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