Invert Hessian piecewise
[openmx:openmx.git] / src / omxFitFunctionBA81.cpp
1 /*
2   Copyright 2012-2013 Joshua Nathaniel Pritikin and contributors
3
4   This is free software: you can redistribute it and/or modify
5   it under the terms of the GNU General Public License as published by
6   the Free Software Foundation, either version 3 of the License, or
7   (at your option) any later version.
8
9   This program is distributed in the hope that it will be useful,
10   but WITHOUT ANY WARRANTY; without even the implied warranty of
11   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12   GNU General Public License for more details.
13
14   You should have received a copy of the GNU General Public License
15   along with this program.  If not, see <http://www.gnu.org/licenses/>.
16 */
17
18 #include "omxFitFunction.h"
19 #include "omxExpectationBA81.h"
20 #include "omxOpenmpWrap.h"
21 #include "libifa-rpf.h"
22
23 struct BA81FitState {
24
25         bool haveLatentMap;
26         std::vector<int> latentMap;
27
28         bool haveItemMap;
29         int itemDerivPadSize;     // maxParam + maxParam*(1+maxParam)/2
30         std::vector<int> paramMap;           // itemParam->cols * itemDerivPadSize -> index of free parameter
31         std::vector<int> paramLocations;     // param# -> count of appearances in ItemParam
32         std::vector<int> itemParamFree;      // itemParam->cols * itemParam->rows
33         std::vector<int> ihessDivisor;       // freeParam * freeParam
34
35         omxMatrix *cholCov;
36         int choleskyError;
37         double *tmpLatentMean;    // maxDims
38         double *tmpLatentCov;     // maxDims * maxDims ; only lower triangle is used
39         omxMatrix *icov;          // inverse covariance matrix
40
41         std::vector< FreeVarGroup* > varGroups;
42         size_t numItemParam;
43
44         BA81FitState();
45         ~BA81FitState();
46 };
47
48 BA81FitState::BA81FitState()
49 {
50         tmpLatentMean = NULL;
51         tmpLatentCov = NULL;
52         haveItemMap = false;
53         haveLatentMap = false;
54 }
55
56 static void buildLatentParamMap(omxFitFunction* oo, FitContext *fc)
57 {
58         FreeVarGroup *fvg = fc->varGroup;
59         BA81FitState *state = (BA81FitState *) oo->argStruct;
60         std::vector<int> &latentMap = state->latentMap;
61         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
62         int meanNum = estate->latentMeanOut->matrixNumber;
63         int covNum = estate->latentCovOut->matrixNumber;
64         int itemNum = estate->itemParam->matrixNumber;
65         int maxAbilities = estate->maxAbilities;
66         int numLatents = maxAbilities + triangleLoc1(maxAbilities);
67
68         latentMap.assign(numLatents, -1);
69
70         int numParam = int(fvg->vars.size());
71         for (int px=0; px < numParam; px++) {
72                 omxFreeVar *fv = fvg->vars[px];
73                 for (size_t lx=0; lx < fv->locations.size(); lx++) {
74                         omxFreeVarLocation *loc = &fv->locations[lx];
75                         int matNum = ~loc->matrix;
76                         if (matNum == meanNum) {
77                                 latentMap[loc->row + loc->col] = px;
78                         } else if (matNum == covNum) {
79                                 int a1 = loc->row;
80                                 int a2 = loc->col;
81                                 if (a1 < a2) std::swap(a1, a2);
82                                 int cell = maxAbilities + triangleLoc1(a1) + a2;
83                                 if (latentMap[cell] == -1) {
84                                         latentMap[cell] = px;
85
86                                         if (a1 == a2 && fv->lbound == NEG_INF) {
87                                                 fv->lbound = 1e-6;  // variance must be positive
88                                                 if (fc->est[px] < fv->lbound) {
89                                                         error("Starting value for variance %s is negative", fv->name);
90                                                 }
91                                         }
92                                 } else if (latentMap[cell] != px) {
93                                         // doesn't work for multigroup constraints TODO
94                                         error("In covariance matrix, %s and %s must be constrained equal to preserve symmetry",
95                                               fvg->vars[latentMap[cell]]->name, fv->name);
96                                 }
97                         } else if (matNum == itemNum) {
98                                 omxRaiseErrorf(globalState, "The fitfunction free.set should consist of "
99                                                "latent distribution parameters, excluding item parameters");
100                         }
101                 }
102         }
103         state->haveLatentMap = TRUE;
104 }
105
106 static void buildItemParamMap(omxFitFunction* oo, FitContext *fc)
107 {
108         FreeVarGroup *fvg = fc->varGroup;
109         BA81FitState *state = (BA81FitState *) oo->argStruct;
110         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
111         omxMatrix *itemParam = estate->itemParam;
112         int size = itemParam->cols * state->itemDerivPadSize;
113         state->paramMap.assign(size, -1);  // matrix location to free param index
114         state->itemParamFree.assign(itemParam->rows * itemParam->cols, FALSE);
115
116         size_t numFreeParams = state->numItemParam = fvg->vars.size();
117         state->paramLocations.assign(numFreeParams, 0);
118
119         for (size_t px=0; px < numFreeParams; px++) {
120                 omxFreeVar *fv = fvg->vars[px];
121                 state->paramLocations[px] = int(fv->locations.size());
122                 for (size_t lx=0; lx < fv->locations.size(); lx++) {
123                         omxFreeVarLocation *loc = &fv->locations[lx];
124                         int matNum = ~loc->matrix;
125                         if (matNum == itemParam->matrixNumber) {
126                                 int at = loc->col * state->itemDerivPadSize + loc->row;
127                                 state->paramMap[at] = px;
128                                 state->itemParamFree[loc->col * itemParam->rows + loc->row] = TRUE;
129
130                                 const double *spec = estate->itemSpec[loc->col];
131                                 int id = spec[RPF_ISpecID];
132                                 double upper, lower;
133                                 (*rpf_model[id].paramBound)(spec, loc->row, &upper, &lower);
134                                 if (fv->lbound == NEG_INF && isfinite(lower)) {
135                                         fv->lbound = lower;
136                                         if (fc->est[px] < fv->lbound) {
137                                                 error("Starting value %s %f less than lower bound %f",
138                                                       fv->name, fc->est[px], lower);
139                                         }
140                                 }
141                                 if (fv->ubound == INF && isfinite(upper)) {
142                                         fv->ubound = upper;
143                                         if (fc->est[px] > fv->ubound) {
144                                                 error("Starting value %s %f greater than upper bound %f",
145                                                       fv->name, fc->est[px], upper);
146                                         }
147                                 }
148                         }
149                         // guard against latent parameters here? TODO
150                 }
151         }
152
153         state->ihessDivisor.resize(size);
154
155         for (int cx=0; cx < itemParam->cols; ++cx) {
156                 const double *spec = estate->itemSpec[cx];
157                 int id = spec[RPF_ISpecID];
158                 int numParam = (*rpf_model[id].numParam)(spec);
159
160                 for (int p1=0; p1 < numParam; p1++) {
161                         int at1 = state->paramMap[cx * state->itemDerivPadSize + p1];
162                         if (at1 < 0) continue;
163
164                         for (int p2=0; p2 <= p1; p2++) {
165                                 int at2 = state->paramMap[cx * state->itemDerivPadSize + p2];
166                                 if (at2 < 0) continue;
167
168                                 if (at1 < at2) std::swap(at1, at2);  // lower triangle
169
170                                 //mxLog("Item %d param(%d,%d) -> H[%d,%d]", cx, p1, p2, at1, at2);
171                                 int at = cx * state->itemDerivPadSize + numParam + triangleLoc1(p1) + p2;
172                                 state->paramMap[at] = numFreeParams + at1 * numFreeParams + at2;
173
174                                 state->ihessDivisor[at] =
175                                         state->paramLocations[at1] * state->paramLocations[at2];
176                         }
177                 }
178         }
179
180         state->haveItemMap = TRUE;
181         //pia(state->paramMap.data(), state->itemDerivPadSize, itemParam->cols);
182 }
183
184 OMXINLINE static double
185 ba81Fit1Ordinate(omxFitFunction* oo, const long qx, const int *quad,
186                  const double *weight, int want, double *myDeriv)
187 {
188         BA81FitState *state = (BA81FitState*) oo->argStruct;
189         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
190         omxMatrix *itemParam = estate->itemParam;
191         int numItems = itemParam->cols;
192         int maxDims = estate->maxDims;
193         int do_fit = want & FF_COMPUTE_FIT;
194         int do_deriv = want & (FF_COMPUTE_GRADIENT | FF_COMPUTE_HESSIAN | FF_COMPUTE_IHESSIAN);
195
196         double where[maxDims];
197         pointToWhere(estate, quad, where, maxDims);
198
199         double *outcomeProb = NULL;
200         if (do_fit) {
201                 outcomeProb = Realloc(NULL, estate->totalOutcomes, double); // avoid malloc/free? TODO
202                 computeRPF(estate, itemParam, quad, TRUE, outcomeProb);
203         }
204
205         double thr_ll = 0;
206         const double *oProb = outcomeProb;
207         for (int ix=0; ix < numItems; ix++) {
208                 const double *spec = estate->itemSpec[ix];
209                 int id = spec[RPF_ISpecID];
210                 int iOutcomes = estate->itemOutcomes[ix];
211
212                 if (do_fit) {
213                         for (int ox=0; ox < iOutcomes; ox++) {
214                                 double got = weight[ox] * oProb[ox];
215                                 thr_ll += got;
216                         }
217                 }
218
219                 if (do_deriv) {
220                         double *iparam = omxMatrixColumn(itemParam, ix);
221                         double *pad = myDeriv + ix * state->itemDerivPadSize;
222                         (*rpf_model[id].dLL1)(spec, iparam, where, weight, pad);
223                 }
224                 oProb += iOutcomes;
225                 weight += iOutcomes;
226         }
227
228         Free(outcomeProb);
229
230         return thr_ll;
231 }
232
233 static double
234 ba81ComputeMFit1(omxFitFunction* oo, int want, FitContext *fc)
235 {
236         BA81FitState *state = (BA81FitState*) oo->argStruct;
237         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
238         if (estate->verbose) mxLog("%s: em.fit-%d", oo->matrix->name, want);
239
240         omxMatrix *customPrior = estate->customPrior;
241         omxMatrix *itemParam = estate->itemParam;
242         std::vector<const double*> &itemSpec = estate->itemSpec;
243         int maxDims = estate->maxDims;
244         const int totalOutcomes = estate->totalOutcomes;
245
246         double *thrDeriv = Calloc(itemParam->cols * state->itemDerivPadSize * Global->numThreads, double);
247
248         double ll = 0;
249         if (customPrior) {
250                 omxRecompute(customPrior);
251                 ll = customPrior->data[0];
252                 // need deriv adjustment TODO
253         }
254
255         if (!isfinite(ll)) {
256                 omxPrint(itemParam, "item param");
257                 error("Bayesian prior returned %g; do you need to add a lbound/ubound?", ll);
258         }
259
260 #pragma omp parallel for num_threads(Global->numThreads) schedule(static,4)
261         for (long qx=0; qx < estate->totalQuadPoints; qx++) {
262                 int quad[maxDims];
263                 decodeLocation(qx, maxDims, estate->quadGridSize, quad);
264                 double *weight = estate->expected + qx * totalOutcomes;
265                 double *myDeriv = thrDeriv + itemParam->cols * state->itemDerivPadSize * omx_absolute_thread_num();
266                 double thr_ll = ba81Fit1Ordinate(oo, qx, quad, weight, want, myDeriv);
267                 
268 #pragma omp atomic
269                 ll += thr_ll;
270         }
271
272         if (want & (FF_COMPUTE_GRADIENT | FF_COMPUTE_HESSIAN | FF_COMPUTE_IHESSIAN)) {
273                 double *deriv0 = thrDeriv;
274
275                 int perThread = itemParam->cols * state->itemDerivPadSize;
276                 for (int th=1; th < Global->numThreads; th++) {
277                         double *thrD = thrDeriv + th * perThread;
278                         for (int ox=0; ox < perThread; ox++) deriv0[ox] += thrD[ox];
279                 }
280
281                 int numItems = itemParam->cols;
282                 for (int ix=0; ix < numItems; ix++) {
283                         const double *spec = itemSpec[ix];
284                         int id = spec[RPF_ISpecID];
285                         double *iparam = omxMatrixColumn(itemParam, ix);
286                         double *pad = deriv0 + ix * state->itemDerivPadSize;
287                         (*rpf_model[id].dLL2)(spec, iparam, pad);
288                 }
289
290                 int numFreeParams = int(state->numItemParam);
291                 int numParams = itemParam->cols * state->itemDerivPadSize;
292                 for (int ox=0; ox < numParams; ox++) {
293                         int to = state->paramMap[ox];
294                         if (to == -1) continue;
295
296                         // Need to check because this can happen if
297                         // lbounds/ubounds are not set appropriately.
298                         if (0 && !isfinite(deriv0[ox])) {
299                                 int item = ox / itemParam->rows;
300                                 mxLog("item parameters:\n");
301                                 const double *spec = itemSpec[item];
302                                 int id = spec[RPF_ISpecID];
303                                 int numParam = (*rpf_model[id].numParam)(spec);
304                                 double *iparam = omxMatrixColumn(itemParam, item);
305                                 pda(iparam, numParam, 1);
306                                 // Perhaps bounds can be pulled in from librpf? TODO
307                                 error("Deriv %d for item %d is %f; are you missing a lbound/ubound?",
308                                       ox, item, deriv0[ox]);
309                         }
310
311                         if (to < numFreeParams) {
312                                 if (want & FF_COMPUTE_GRADIENT) {
313                                         fc->grad[to] -= deriv0[ox];
314                                 }
315                         } else {
316                                 if (want & FF_COMPUTE_HESSIAN) {
317                                         int Hto = to - numFreeParams;
318                                         fc->hess[Hto] -= deriv0[ox];
319                                 }
320                         }
321                 }
322
323                 if (want & FF_COMPUTE_IHESSIAN) {
324                         for (int ix=0; ix < numItems; ix++) {
325                                 const double *spec = itemSpec[ix];
326                                 int id = spec[RPF_ISpecID];
327                                 int iParams = (*rpf_model[id].numParam)(spec);
328                                 double *pad = deriv0 + ix * state->itemDerivPadSize + iParams;
329                                 int *mask = state->itemParamFree.data() + ix * itemParam->rows;
330                                 double stress;
331                                 omxApproxInvertPackedPosDefTriangular(iParams, mask, pad, &stress);
332                         }
333
334                         for (int ox=0; ox < numParams; ox++) {
335                                 int to = state->paramMap[ox];
336                                 if (to == -1) continue;
337                                 if (to >= numFreeParams) {
338                                         int Hto = to - numFreeParams;
339                                         fc->ihess[Hto] += deriv0[ox] / state->ihessDivisor[ox];
340                                 }
341                         }
342                 }
343         }
344
345         Free(thrDeriv);
346
347         return -ll;
348 }
349
350 static int
351 moveLatentDistribution(omxFitFunction *oo, FitContext *fc,
352                        double *ElatentMean, double *ElatentCov)
353 {
354         BA81FitState *state = (BA81FitState*) oo->argStruct;
355         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
356         std::vector<const double*> &itemSpec = estate->itemSpec;
357         omxMatrix *itemParam = estate->itemParam;
358         omxMatrix *design = estate->design;
359         double *tmpLatentMean = state->tmpLatentMean;
360         double *tmpLatentCov = state->tmpLatentCov;
361         int maxDims = estate->maxDims;
362         int maxAbilities = estate->maxAbilities;
363         int moveCount = 0;
364
365         int numItems = itemParam->cols;
366         for (int ix=0; ix < numItems; ix++) {
367                 const double *spec = itemSpec[ix];
368                 int id = spec[RPF_ISpecID];
369                 const double *rawDesign = omxMatrixColumn(design, ix);
370                 int idesign[design->rows];
371                 int idx = 0;
372                 for (int dx=0; dx < design->rows; dx++) {
373                         if (isfinite(rawDesign[dx])) {
374                                 idesign[idx++] = rawDesign[dx]-1;
375                         } else {
376                                 idesign[idx++] = -1;
377                         }
378                 }
379                 for (int d1=0; d1 < idx; d1++) {
380                         if (idesign[d1] == -1) {
381                                 tmpLatentMean[d1] = 0;
382                         } else {
383                                 tmpLatentMean[d1] = ElatentMean[idesign[d1]];
384                         }
385                         for (int d2=0; d2 <= d1; d2++) {
386                                 int cell = idesign[d2] * maxAbilities + idesign[d1];
387                                 if (idesign[d1] == -1 || idesign[d2] == -1) {
388                                         tmpLatentCov[d2 * maxDims + d1] = d1==d2? 1 : 0;
389                                 } else {
390                                         tmpLatentCov[d2 * maxDims + d1] = ElatentCov[cell];
391                                 }
392                         }
393                 }
394                 if (1) {  // ease debugging, make optional TODO
395                         for (int d1=idx; d1 < maxDims; d1++) tmpLatentMean[d1] = nan("");
396                         for (int d1=0; d1 < maxDims; d1++) {
397                                 for (int d2=0; d2 < maxDims; d2++) {
398                                         if (d1 < idx && d2 < idx) continue;
399                                         tmpLatentCov[d2 * maxDims + d1] = nan("");
400                                 }
401                         }
402                 }
403                 double *iparam = omxMatrixColumn(itemParam, ix);
404                 int *mask = state->paramMap.data() + state->itemDerivPadSize * ix;
405                 rpf_model[id].rescale(spec, iparam, mask, tmpLatentMean, tmpLatentCov);
406         }
407
408         int numFreeParams = int(fc->varGroup->vars.size());
409         for (int rx=0; rx < itemParam->rows; rx++) {
410                 for (int cx=0; cx < itemParam->cols; cx++) {
411                         int px = cx * state->itemDerivPadSize + rx;
412                         int vx = state->paramMap[px];
413                         if (vx >= 0 && vx < numFreeParams) {
414                                 fc->est[vx] = omxMatrixElement(itemParam, rx, cx);
415                                 ++moveCount;
416                         }
417                 }
418         }
419         return moveCount;
420 }
421
422 static void
423 schilling_bock_2005_rescale(omxFitFunction *oo, FitContext *fc)
424 {
425         BA81FitState *state = (BA81FitState*) oo->argStruct;
426         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
427         omxMatrix *cholCov = state->cholCov;
428         int maxAbilities = estate->maxAbilities;
429
430         //pda(ElatentMean, maxAbilities, 1);
431         //pda(estate->ElatentCov.data(), maxAbilities, maxAbilities);
432         //omxPrint(design, "design");
433
434         memcpy(cholCov->data, estate->ElatentCov.data(), sizeof(double) * maxAbilities * maxAbilities);
435
436         const char triangle = 'L';
437         F77_CALL(dpotrf)(&triangle, &maxAbilities, cholCov->data, &maxAbilities, &state->choleskyError);
438         if (state->choleskyError != 0) {
439                 warning("Cholesky failed with %d; rescaling disabled", state->choleskyError); // make error TODO?
440                 return;
441         }
442
443         //pda(cholCov->data, maxAbilities, maxAbilities);
444
445         int moveCount = moveLatentDistribution(oo, fc, estate->ElatentMean.data(), cholCov->data);
446         if (estate->verbose) mxLog("%s: schilling-bock (%d param)", oo->matrix->name, moveCount);
447 }
448
449 void ba81SetFreeVarGroup(omxFitFunction *oo, FreeVarGroup *fvg)
450 {}
451
452 // can use same reorganization to avoid the gram product of where for every pattern TODO
453 static void mapLatentDeriv(BA81FitState *state, BA81Expect *estate, int sgroup, double piece,
454                            const std::vector<double> &derivCoef,
455                            double *derivOut)
456 {
457         int maxAbilities = estate->maxAbilities;
458         int maxDims = estate->maxDims;
459         int pmax = maxDims;
460         if (estate->numSpecific) pmax -= 1;
461
462         if (sgroup == 0) {
463                 int cx = 0;
464                 for (int d1=0; d1 < pmax; ++d1) {
465                         double amt1 = piece * derivCoef[d1];
466 #pragma omp atomic
467                         derivOut[d1] += amt1;
468                         for (int d2=0; d2 <= d1; ++d2) {
469                                 int to = maxAbilities + cx;
470                                 double amt2 = piece * derivCoef[maxDims + cx];
471 #pragma omp atomic
472                                 derivOut[to] += amt2;
473                                 ++cx;
474                         }
475                 }
476         }
477
478         if (estate->numSpecific) {
479                 int sdim = pmax + sgroup;
480                 double amt3 = piece * derivCoef[pmax];
481 #pragma omp atomic
482                 derivOut[sdim] += amt3;
483
484                 double amt4 = piece * derivCoef[maxDims + triangleLoc0(pmax)];
485                 int to = maxAbilities + triangleLoc0(sdim);
486 #pragma omp atomic
487                 derivOut[to] += amt4;
488         }
489 }
490
491 static double *reduceForSpecific(omxMatrix *mat, int maxDims, int sgroup)
492 {
493         double *out = Calloc(maxDims * maxDims, double);
494         for (int d1=0; d1 < maxDims-1; ++d1) {
495                 int cell = d1 * maxDims;
496                 for (int d2=0; d2 < maxDims-1; ++d2) {
497                         out[cell + d2] = omxMatrixElement(mat, d1, d2);
498                 }
499         }
500         int sloc = maxDims-1 + sgroup;
501         out[maxDims*maxDims - 1] = omxMatrixElement(mat, sloc, sloc);
502         return out;
503 }
504
505 static void calcDerivCoef(BA81FitState *state, BA81Expect *estate,
506                           double *where, int sgroup, std::vector<double> *derivCoef)
507 {
508         omxMatrix *mean = estate->latentMeanOut;
509         omxMatrix *cov = estate->latentCovOut;
510         omxMatrix *icov = state->icov;
511         double *covData = cov->data;
512         double *icovData = icov->data;
513         int maxDims = estate->maxDims;
514         const char R='R';
515         const char L='L';
516         const char U='U';
517         const double alpha = 1;
518         const double beta = 0;
519         const int one = 1;
520
521         double *scov = NULL;
522         double *sicov = NULL;
523         if (estate->numSpecific) {
524                 scov = reduceForSpecific(cov, maxDims, sgroup);
525                 covData = scov;
526
527                 sicov = reduceForSpecific(icov, maxDims, sgroup);
528                 icovData = sicov;
529         }
530
531         std::vector<double> whereDiff(maxDims);
532         std::vector<double> whereGram(triangleLoc1(maxDims));
533         for (int d1=0; d1 < maxDims; ++d1) {
534                 whereDiff[d1] = where[d1] - omxVectorElement(mean, d1);
535         }
536         gramProduct(whereDiff.data(), whereDiff.size(), whereGram.data());
537
538         F77_CALL(dsymv)(&U, &maxDims, &alpha, icovData, &maxDims, whereDiff.data(), &one,
539                         &beta, derivCoef->data(), &one);
540
541         std::vector<double> covGrad1(maxDims * maxDims);
542         std::vector<double> covGrad2(maxDims * maxDims);
543
544         int cx=0;
545         for (int d1=0; d1 < maxDims; ++d1) {
546                 for (int d2=0; d2 <= d1; ++d2) {
547                         covGrad1[d2 * maxDims + d1] = covData[d2 * maxDims + d1] - whereGram[cx];
548                         ++cx;
549                 }
550         }
551
552         F77_CALL(dsymm)(&R, &L, &maxDims, &maxDims, &alpha, covGrad1.data(), &maxDims, icovData,
553                         &maxDims, &beta, covGrad2.data(), &maxDims);
554         F77_CALL(dsymm)(&R, &L, &maxDims, &maxDims, &alpha, icovData, &maxDims, covGrad2.data(),
555                         &maxDims, &beta, covGrad1.data(), &maxDims);
556
557         for (int d1=0; d1 < maxDims; ++d1) {
558                 covGrad1[d1 * maxDims + d1] /= 2.0;
559         }
560
561         cx = maxDims;
562         for (int d1=0; d1 < maxDims; ++d1) {
563                 int cell = d1 * maxDims;
564                 for (int d2=0; d2 <= d1; ++d2) {
565                         (*derivCoef)[cx] = -covGrad1[cell + d2];
566                         ++cx;
567                 }
568         }
569
570         Free(scov);
571         Free(sicov);
572 }
573
574 static bool latentDeriv(omxFitFunction *oo, double *gradient)
575 {
576         omxExpectation *expectation = oo->expectation;
577         BA81FitState *state = (BA81FitState*) oo->argStruct;
578         BA81Expect *estate = (BA81Expect*) expectation->argStruct;
579         if (estate->verbose) mxLog("%s: latentDeriv", oo->matrix->name);
580
581         int numUnique = estate->numUnique;
582         int numSpecific = estate->numSpecific;
583         int maxDims = estate->maxDims;
584         int maxAbilities = estate->maxAbilities;
585         int primaryDims = maxDims;
586         omxMatrix *cov = estate->latentCovOut;
587         int *numIdentical = estate->numIdentical;
588         double *patternLik = estate->patternLik;
589
590         OMXZERO(patternLik, numUnique);
591
592         omxCopyMatrix(state->icov, cov);
593
594         int info;
595         omxDPOTRF(state->icov, &info);
596         if (info != 0) {
597                 if (info < 0) error("dpotrf invalid argument %d", -info);
598                 return FALSE;
599         }
600         omxDPOTRI(state->icov, &info);
601         if (info != 0) {
602                 if (info < 0) error("dpotri invalid argument %d", -info);
603                 return FALSE;
604         }
605         // fill in rest from upper triangle
606         for (int rx=1; rx < maxAbilities; ++rx) {
607                 for (int cx=0; cx < rx; ++cx) {
608                         omxSetMatrixElement(state->icov, rx, cx, omxMatrixElement(state->icov, cx, rx));
609                 }
610         }
611
612         int maxDerivCoef = maxDims + triangleLoc1(maxDims);
613         int numLatents = maxAbilities + triangleLoc1(maxAbilities);
614         double *uniqueDeriv = Calloc(numUnique * numLatents, double);
615
616         if (numSpecific == 0) {
617 #pragma omp parallel for num_threads(Global->numThreads)
618                 for (long qx=0; qx < estate->totalQuadPoints; qx++) {
619                         const int thrId = omx_absolute_thread_num();
620                         int quad[maxDims];
621                         decodeLocation(qx, maxDims, estate->quadGridSize, quad);
622                         double where[maxDims];
623                         pointToWhere(estate, quad, where, maxDims);
624                         std::vector<double> derivCoef(maxDerivCoef);
625                         calcDerivCoef(state, estate, where, 0, &derivCoef);
626                         double area = estate->priQarea[qx];
627                         double *lxk = ba81LikelihoodFast1(expectation, thrId, 0, qx);
628
629                         for (int px=0; px < numUnique; px++) {
630                                 double tmp = (lxk[px] * area);
631 #pragma omp atomic
632                                 patternLik[px] += tmp;
633                                 mapLatentDeriv(state, estate, 0, tmp, derivCoef,
634                                                uniqueDeriv + px * numLatents);
635                         }
636                 }
637         } else {
638                 primaryDims -= 1;
639                 int sDim = primaryDims;
640                 long specificPoints = estate->quadGridSize;
641
642 #pragma omp parallel for num_threads(Global->numThreads)
643                 for (long qx=0; qx < estate->totalPrimaryPoints; qx++) {
644                         const int thrId = omx_absolute_thread_num();
645                         int quad[maxDims];
646                         decodeLocation(qx, primaryDims, estate->quadGridSize, quad);
647
648                         cai2010(expectation, thrId, FALSE, qx);
649                         double *allElxk = eBase(estate, thrId);
650                         double *Eslxk = esBase(estate, thrId);
651
652                         for (long sx=0; sx < specificPoints; sx++) {
653                                 long qloc = qx * specificPoints + sx;
654                                 quad[sDim] = sx;
655                                 double where[maxDims];
656                                 pointToWhere(estate, quad, where, maxDims);
657                                 for (int sgroup=0; sgroup < numSpecific; sgroup++) {
658                                         std::vector<double> derivCoef(maxDerivCoef);
659                                         calcDerivCoef(state, estate, where, sgroup, &derivCoef);
660                                         double area = areaProduct(estate, qx, sx, sgroup);
661                                         double *lxk = ba81LikelihoodFast1(expectation, thrId, sgroup, qloc);
662                                         for (int px=0; px < numUnique; px++) {
663                                                 double Ei = allElxk[px];
664                                                 double Eis = Eslxk[sgroup * numUnique + px];
665                                                 double tmp = ((Ei / Eis) * lxk[px] * area);
666                                                 mapLatentDeriv(state, estate, sgroup, tmp, derivCoef,
667                                                                uniqueDeriv + px * numLatents);
668                                         }
669                                 }
670                         }
671
672                         double priArea = estate->priQarea[qx];
673                         for (int px=0; px < numUnique; px++) {
674                                 double Ei = allElxk[px];
675                                 double tmp = (Ei * priArea);
676 #pragma omp atomic
677                                 patternLik[px] += tmp;
678                         }
679                 }
680         }
681
682         /*
683         std::vector<double> hess1(triangleLoc1(numLatents));
684         std::vector<double> hessSum(triangleLoc1(numLatents));
685
686         // could run nicely in parallel with numUnique * triangleLoc(numLatents) buffer
687         for (int px=0; px < numUnique; ++px) {
688                 gramProduct(uniqueDeriv + px * numLatents, numLatents, hess1.data());
689                 double dups = numIdentical[px];
690                 for (int rx=0; rx < triangleLoc1(numLatents); ++rx) {
691                         hessSum[rx] += hess1[rx] * dups;
692                 }
693         }
694         */
695
696 #pragma omp parallel for num_threads(Global->numThreads)
697         for (int px=0; px < numUnique; ++px) {
698                 double weight = numIdentical[px] / patternLik[px];
699                 for (int rx=0; rx < numLatents; ++rx) {
700                         uniqueDeriv[px * numLatents + rx] *= weight;
701                 }
702         }
703
704 #pragma omp parallel for num_threads(Global->numThreads)
705         for (int rx=0; rx < numLatents; ++rx) {
706                 for (int px=1; px < numUnique; ++px) {
707                         uniqueDeriv[rx] += uniqueDeriv[px * numLatents + rx];
708                 }
709         }
710
711         for (int l1=0; l1 < numLatents; ++l1) {
712                 int t1 = state->latentMap[l1];
713                 if (t1 < 0) continue;
714                 gradient[t1] -= 2 * uniqueDeriv[l1];
715
716                 /*
717                 for (int l2=0; l2 <= l1; ++l2) {
718                         int t2 = state->latentMap[l2];
719                         if (t2 < 0) continue;
720                         hessian[numLatents * t1 + t2] -= 2 * hessSum[triangleLoc1(l1) + l2];
721                 }
722                 */
723         }
724
725         Free(uniqueDeriv);
726
727         return TRUE;
728 }
729
730 static void setLatentStartingValues(omxFitFunction *oo, FitContext *fc)
731 {
732         BA81FitState *state = (BA81FitState*) oo->argStruct;
733         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
734         std::vector<int> &latentMap = state->latentMap;
735         std::vector<double> &ElatentMean = estate->ElatentMean;
736         std::vector<double> &ElatentCov = estate->ElatentCov;
737         int maxAbilities = estate->maxAbilities;
738
739         for (int a1 = 0; a1 < maxAbilities; ++a1) {
740                 if (latentMap[a1] >= 0) {
741                         int to = latentMap[a1];
742                         fc->est[to] = ElatentMean[a1];
743                 }
744
745                 for (int a2 = 0; a2 <= a1; ++a2) {
746                         int to = latentMap[maxAbilities + triangleLoc1(a1) + a2];
747                         if (to < 0) continue;
748                         fc->est[to] = ElatentCov[a1 * maxAbilities + a2];
749                 }
750         }
751
752         //fc->log("setLatentStartingValues", FF_COMPUTE_ESTIMATE);
753 }
754
755 static double
756 ba81ComputeFit(omxFitFunction* oo, int want, FitContext *fc)
757 {
758         BA81FitState *state = (BA81FitState*) oo->argStruct;
759         BA81Expect *estate = (BA81Expect*) oo->expectation->argStruct;
760
761         if (want & FF_COMPUTE_POSTOPTIMIZE) return 0;
762
763         if (estate->type == EXPECTATION_AUGMENTED) {
764                 if (!state->haveItemMap) buildItemParamMap(oo, fc);
765
766                 if (want & FF_COMPUTE_PREOPTIMIZE) {
767                         // schilling_bock_2005_rescale(oo, fc); seems counterproductive
768                         return 0;
769                 }
770
771                 if (state->numItemParam != fc->varGroup->vars.size()) error("mismatch"); // remove TODO
772                 double got = ba81ComputeMFit1(oo, want, fc);
773                 return got;
774         } else if (estate->type == EXPECTATION_OBSERVED) {
775                 if (!state->haveLatentMap) buildLatentParamMap(oo, fc);
776
777                 if (want & FF_COMPUTE_PREOPTIMIZE) {
778                         setLatentStartingValues(oo, fc);
779                         return 0;
780                 }
781
782                 omxExpectationCompute(oo->expectation, NULL);
783
784                 if (want & FF_COMPUTE_GRADIENT) {
785                         if (!latentDeriv(oo, fc->grad)) {
786                                 return INFINITY;
787                         }
788                 }
789
790                 if (want & FF_COMPUTE_HESSIAN) {
791                         warning("%s: Hessian is not available for latent distribution parameters", oo->matrix->name);
792                 }
793
794                 if (want & FF_COMPUTE_FIT) {
795                         if (estate->verbose) mxLog("%s: fit", oo->matrix->name);
796                         double *patternLik = estate->patternLik;
797                         int *numIdentical = estate->numIdentical;
798                         int numUnique = estate->numUnique;
799                         double got = 0;
800 #pragma omp parallel for num_threads(Global->numThreads) schedule(static,64) reduction(+:got)
801                         for (int ux=0; ux < numUnique; ux++) {
802                                 got += numIdentical[ux] * log(patternLik[ux]);
803                         }
804                         //mxLog("fit %.2f", -2 * got);
805                         return -2 * got;
806                 }
807
808                 // if (want & FF_COMPUTE_POSTOPTIMIZE)  discard lxk cache? TODO
809
810                 return 0;
811         } else {
812                 error("Confused");
813         }
814 }
815
816 static void ba81Compute(omxFitFunction *oo, int want, FitContext *fc)
817 {
818         if (!want) return;
819         double got = ba81ComputeFit(oo, want, fc);
820         if (got) oo->matrix->data[0] = got;
821 }
822
823 BA81FitState::~BA81FitState()
824 {
825         Free(tmpLatentMean);
826         Free(tmpLatentCov);
827         omxFreeAllMatrixData(icov);
828         omxFreeAllMatrixData(cholCov);
829 }
830
831 static void ba81Destroy(omxFitFunction *oo) {
832         BA81FitState *state = (BA81FitState *) oo->argStruct;
833         delete state;
834 }
835
836 void omxInitFitFunctionBA81(omxFitFunction* oo)
837 {
838         if (!oo->argStruct) { // ugh!
839                 BA81FitState *state = new BA81FitState;
840                 oo->argStruct = state;
841         }
842
843         BA81FitState *state = (BA81FitState*) oo->argStruct;
844
845         omxExpectation *expectation = oo->expectation;
846         BA81Expect *estate = (BA81Expect*) expectation->argStruct;
847
848         //newObj->data = oo->expectation->data;
849
850         oo->computeFun = ba81Compute;
851         oo->setVarGroup = ba81SetFreeVarGroup;
852         oo->destructFun = ba81Destroy;
853         oo->gradientAvailable = TRUE;
854         oo->hessianAvailable = TRUE;
855
856         int maxParam = estate->itemParam->rows;
857         state->itemDerivPadSize = maxParam + triangleLoc1(maxParam);
858
859         int maxAbilities = estate->maxAbilities;
860
861         state->tmpLatentMean = Realloc(NULL, estate->maxDims, double);
862         state->tmpLatentCov = Realloc(NULL, estate->maxDims * estate->maxDims, double);
863
864         int numItems = estate->itemParam->cols;
865         for (int ix=0; ix < numItems; ix++) {
866                 const double *spec = estate->itemSpec[ix];
867                 int id = spec[RPF_ISpecID];
868                 if (id < 0 || id >= rpf_numModels) {
869                         error("ItemSpec %d has unknown item model %d", ix, id);
870                 }
871         }
872
873         state->icov = omxInitMatrix(NULL, maxAbilities, maxAbilities, TRUE, globalState);
874         state->cholCov = omxInitMatrix(NULL, maxAbilities, maxAbilities, TRUE, globalState);
875 }