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