Add cache for outcomeProb
[openmx:openmx.git] / src / omxExpectationBA81.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 <Rmath.h>
19 #include "omxExpectationBA81.h"
20 #include "glue.h"
21 #include "libifa-rpf.h"
22 #include "dmvnorm.h"
23
24 const struct rpf *rpf_model = NULL;
25 int rpf_numModels;
26
27 void pda(const double *ar, int rows, int cols)
28 {
29         std::string buf;
30         for (int rx=0; rx < rows; rx++) {   // column major order
31                 for (int cx=0; cx < cols; cx++) {
32                         buf += string_snprintf("%.6g, ", ar[cx * rows + rx]);
33                 }
34                 buf += "\n";
35         }
36         mxLogBig(buf);
37 }
38
39 void pia(const int *ar, int rows, int cols)
40 {
41         std::string buf;
42         for (int rx=0; rx < rows; rx++) {   // column major order
43                 for (int cx=0; cx < cols; cx++) {
44                         buf += string_snprintf("%d, ", ar[cx * rows + rx]);
45                 }
46                 buf += "\n";
47         }
48         mxLogBig(buf);
49 }
50
51 // state->speQarea[sIndex(state, sx, qx)]
52 OMXINLINE static
53 int sIndex(BA81Expect *state, int sx, int qx)
54 {
55         //if (sx < 0 || sx >= state->numSpecific) error("Out of domain");
56         //if (qx < 0 || qx >= state->quadGridSize) error("Out of domain");
57         return sx * state->quadGridSize + qx;
58 }
59
60 // Depends on item parameters, but not latent distribution
61 void computeRPF(BA81Expect *state, omxMatrix *itemParam, const int *quad,
62                 const bool wantlog, double *out)
63 {
64         omxMatrix *design = state->design;
65         int maxDims = state->maxDims;
66         size_t numItems = state->itemSpec.size();
67
68         double theta[maxDims];
69         pointToWhere(state, quad, theta, maxDims);
70
71         for (size_t ix=0; ix < numItems; ix++) {
72                 const double *spec = state->itemSpec[ix];
73                 int id = spec[RPF_ISpecID];
74                 int dims = spec[RPF_ISpecDims];
75                 double ptheta[dims];
76
77                 for (int dx=0; dx < dims; dx++) {
78                         int ability = (int)omxMatrixElement(design, dx, ix) - 1;
79                         if (ability >= maxDims) ability = maxDims-1;
80                         ptheta[dx] = theta[ability];
81                 }
82
83                 double *iparam = omxMatrixColumn(itemParam, ix);
84                 if (wantlog) {
85                         (*rpf_model[id].logprob)(spec, iparam, ptheta, out);
86                 } else {
87                         (*rpf_model[id].prob)(spec, iparam, ptheta, out);
88                 }
89 #if 0
90                 for (int ox=0; ox < state->itemOutcomes[ix]; ox++) {
91                         if (!isfinite(out[ox]) || out[ox] > 0) {
92                                 mxLog("item param");
93                                 pda(iparam, itemParam->rows, 1);
94                                 mxLog("where");
95                                 pda(ptheta, dims, 1);
96                                 error("RPF returned %20.20f", out[ox]);
97                         }
98                 }
99 #endif
100                 out += state->itemOutcomes[ix];
101         }
102 }
103
104 OMXINLINE static double *
105 getLXKcache(BA81Expect *state, const long qx, const int specific)
106 {
107         long ordinate;
108         if (state->numSpecific == 0) {
109                 ordinate = qx;
110         } else {
111                 ordinate = specific * state->totalQuadPoints + qx;
112         }
113         return state->lxk + state->numUnique * ordinate;
114 }
115
116 OMXINLINE static double *
117 ba81Likelihood(omxExpectation *oo, const int thrId, int specific, const long qx)
118 {
119         BA81Expect *state = (BA81Expect*) oo->argStruct;
120         int numUnique = state->numUnique;
121         std::vector<int> &itemOutcomes = state->itemOutcomes;
122         omxData *data = state->data;
123         size_t numItems = state->itemSpec.size();
124         int *Sgroup = state->Sgroup;
125         double *lxk;
126
127         if (!state->cacheLXK) {
128                 lxk = state->lxk + numUnique * thrId;
129         } else {
130                 lxk = getLXKcache(state, qx, specific);
131         }
132
133         const int *rowMap = state->rowMap;
134         for (int px=0; px < numUnique; px++) {
135                 double lxk1 = 1;
136                 const double *oProb = state->outcomeProb + qx * state->totalOutcomes;
137                 for (size_t ix=0; ix < numItems; ix++) {
138                         int pick = omxIntDataElementUnsafe(data, rowMap[px], ix);
139                         if (specific == Sgroup[ix] && pick != NA_INTEGER) {
140                                 double piece = oProb[pick-1];  // move -1 elsewhere TODO
141                                 lxk1 *= piece;
142                                 //mxLog("%d pick %d piece %.7f", ix, pick-1, piece);
143                         }
144                         oProb += itemOutcomes[ix];
145                 }
146 #if 0
147 #pragma omp critical(ba81LikelihoodDebug1)
148                 if (!isfinite(lxk1) || lxk1 > numItems) {
149                         mxLog("prob");
150                         pda(outcomeProb, 1, state->totalOutcomes);
151                         error("Likelihood of row %d is %f", rowMap[px], lxk1);
152                 }
153 #endif
154                 lxk[px] = lxk1;
155         }
156
157         return lxk;
158 }
159
160 double *ba81LikelihoodFast(omxExpectation *oo, const int thrId, int specific, const long qx)
161 {
162         BA81Expect *state = (BA81Expect*) oo->argStruct;
163         if (!state->cacheLXK) {
164                 double *ret = ba81Likelihood(oo, thrId, specific, qx);
165                 return ret;
166         } else {
167                 return getLXKcache(state, qx, specific);
168         }
169
170 }
171
172 OMXINLINE static void
173 mapLatentSpace(BA81Expect *state, int sgroup, double piece, const double *where,
174                const std::vector<double> &whereGram, double *latentDist)
175 {
176         int maxDims = state->maxDims;
177         int maxAbilities = state->maxAbilities;
178         int pmax = maxDims;
179         if (state->numSpecific) pmax -= 1;
180
181         if (sgroup == 0) {
182                 int gx = 0;
183                 int cx = maxAbilities;
184                 for (int d1=0; d1 < pmax; d1++) {
185                         double piece_w1 = piece * where[d1];
186                         latentDist[d1] += piece_w1;
187                         for (int d2=0; d2 <= d1; d2++) {
188                                 double piece_cov = piece * whereGram[gx];
189                                 latentDist[cx] += piece_cov;
190                                 ++cx; ++gx;
191                         }
192                 }
193         }
194
195         if (state->numSpecific) {
196                 int sdim = pmax + sgroup;
197                 double piece_w1 = piece * where[pmax];
198                 latentDist[sdim] += piece_w1;
199
200                 double piece_var = piece * whereGram[triangleLoc0(pmax)];
201                 int to = maxAbilities + triangleLoc0(sdim);
202                 latentDist[to] += piece_var;
203         }
204 }
205
206 // Eslxk, allElxk (Ei, Eis) depend on the ordinate of the primary dimensions
207 void cai2010(omxExpectation* oo, const int thrId, int recompute, const long primaryQ)
208 {
209         BA81Expect *state = (BA81Expect*) oo->argStruct;
210         int numUnique = state->numUnique;
211         int numSpecific = state->numSpecific;
212         int quadGridSize = state->quadGridSize;
213         double *allElxk = eBase(state, thrId);
214         double *Eslxk = esBase(state, thrId);
215
216         for (int px=0; px < numUnique; px++) {
217                 allElxk[px] = 1;
218                 for (int sx=0; sx < numSpecific; sx++) {
219                         Eslxk[sx * numUnique + px] = 0;
220                 }
221         }
222
223         if (!state->cacheLXK) recompute = TRUE;
224
225         for (int qx=0; qx < quadGridSize; qx++) {
226                 long qloc = primaryQ * quadGridSize + qx;
227
228                 for (int sx=0; sx < numSpecific; sx++) {
229                         double *myEslxk = Eslxk + sx * numUnique;
230                         double *lxk;     // a.k.a. "L_is"
231                         if (recompute) {
232                                 lxk = ba81Likelihood(oo, thrId, sx, qloc);
233                         } else {
234                                 lxk = getLXKcache(state, qloc, sx);
235                         }
236
237                         for (int ix=0; ix < numUnique; ix++) {
238                                 double area = state->speQarea[sIndex(state, sx, qx)];
239                                 double piece = lxk[ix] * area;
240                                 //mxLog("E.is(%d) (%ld,%d) %.6f + %.6f = %.6f",
241                                 //  sx, primaryQ, qx, lxk[ix], area, piece);
242                                 myEslxk[ix] += piece;
243                         }
244                 }
245         }
246
247         for (int sx=0; sx < numSpecific; sx++) {
248                 for (int px=0; px < numUnique; px++) {
249                         //mxLog("E.is(%d) at (%ld) %.6f", sx, primaryQ,
250                         //  Eslxk[sx * numUnique + px]);
251                         allElxk[px] *= Eslxk[sx * numUnique + px];  // allSlxk a.k.a. "E_i"
252                 }
253         }
254 }
255
256 static void ba81OutcomeProb(BA81Expect *state)
257 {
258         int maxDims = state->maxDims;
259         double *qProb = state->outcomeProb =
260                 Realloc(state->outcomeProb, state->totalOutcomes * state->totalQuadPoints, double);
261         for (long qx=0; qx < state->totalQuadPoints; qx++) {
262                 int quad[maxDims];
263                 decodeLocation(qx, maxDims, state->quadGridSize, quad);
264                 double where[maxDims];
265                 pointToWhere(state, quad, where, maxDims);
266                 computeRPF(state, state->EitemParam, quad, FALSE, qProb);
267                 qProb += state->totalOutcomes;
268         }
269 }
270
271 static void ba81Estep1(omxExpectation *oo)
272 {
273         if(OMX_DEBUG) {mxLog("Beginning %s Computation.", oo->name);}
274
275         BA81Expect *state = (BA81Expect*) oo->argStruct;
276         if (state->verbose) {
277                 mxLog("%s: lxk(%d) patternLik ElatentMean ElatentCov",
278                       oo->name, omxGetMatrixVersion(state->EitemParam));
279         }
280
281         int numUnique = state->numUnique;
282         int numSpecific = state->numSpecific;
283         int maxDims = state->maxDims;
284         int maxAbilities = state->maxAbilities;
285         int primaryDims = maxDims;
286
287         state->patternLik = Realloc(state->patternLik, numUnique, double);
288         double *patternLik = state->patternLik;
289         OMXZERO(patternLik, numUnique);
290
291         int numLatents = maxAbilities + triangleLoc1(maxAbilities);
292         int numLatentsPerThread = numUnique * numLatents;
293         double *latentDist = Calloc(numUnique * numLatents * Global->numThreads, double);
294
295         if (numSpecific == 0) {
296 #pragma omp parallel for num_threads(Global->numThreads)
297                 for (long qx=0; qx < state->totalQuadPoints; qx++) {
298                         int thrId = omx_absolute_thread_num();
299                         double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
300                         int quad[maxDims];
301                         decodeLocation(qx, maxDims, state->quadGridSize, quad);
302                         double where[maxDims];
303                         pointToWhere(state, quad, where, maxDims);
304                         std::vector<double> whereGram(triangleLoc1(maxDims));
305                         gramProduct(where, maxDims, whereGram.data());
306
307                         double *lxk = ba81Likelihood(oo, thrId, 0, qx);
308
309                         double area = state->priQarea[qx];
310                         for (int px=0; px < numUnique; px++) {
311                                 double tmp = lxk[px] * area;
312 #if 0
313                                 if (!isfinite(tmp)) {
314                                         mxLog("where");
315                                         pda(where, maxDims, 1);
316                                         error("Row %d lxk %f logArea %f tmp %f",
317                                               state->rowMap[px], lxk[px], logArea, tmp);
318                                 }
319 #endif
320 #pragma omp atomic
321                                 patternLik[px] += tmp;
322                                 mapLatentSpace(state, 0, tmp, where, whereGram,
323                                                thrLatentDist + px * numLatents);
324                         }
325                 }
326         } else {
327                 primaryDims -= 1;
328                 int sDim = primaryDims;
329                 long specificPoints = state->quadGridSize;
330
331 #pragma omp parallel for num_threads(Global->numThreads)
332                 for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
333                         int thrId = omx_absolute_thread_num();
334                         double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
335                         int quad[maxDims];
336                         decodeLocation(qx, primaryDims, state->quadGridSize, quad);
337
338                         cai2010(oo, thrId, TRUE, qx);
339                         double *allElxk = eBase(state, thrId);
340                         double *Eslxk = esBase(state, thrId);
341
342                         for (long sx=0; sx < specificPoints; sx++) {
343                                 long qloc = qx * specificPoints + sx;
344                                 quad[sDim] = sx;
345                                 double where[maxDims];
346                                 pointToWhere(state, quad, where, maxDims);
347                                 std::vector<double> whereGram(triangleLoc1(maxDims));
348                                 gramProduct(where, maxDims, whereGram.data());
349
350                                 for (int sgroup=0; sgroup < numSpecific; sgroup++) {
351                                         double area = areaProduct(state, qx, sx, sgroup);
352                                         double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, qloc);
353                                         for (int px=0; px < numUnique; px++) {
354                                                 double Ei = allElxk[px];
355                                                 double Eis = Eslxk[sgroup * numUnique + px];
356                                                 double tmp = ((Ei / Eis) * lxk[px] * area);
357                                                 mapLatentSpace(state, sgroup, tmp, where, whereGram,
358                                                                thrLatentDist + px * numLatents);
359                                         }
360                                 }
361                         }
362
363                         double priArea = state->priQarea[qx];
364                         for (int px=0; px < numUnique; px++) {
365                                 double Ei = allElxk[px];
366                                 double tmp = (Ei * priArea);
367 #pragma omp atomic
368                                 patternLik[px] += tmp;
369                         }
370                 }
371         }
372
373         int *numIdentical = state->numIdentical;
374
375         //mxLog("raw latent");
376         //pda(latentDist, numLatents, numUnique);
377
378 #pragma omp parallel for num_threads(Global->numThreads) schedule(dynamic)
379         for (int lx=0; lx < maxAbilities + triangleLoc1(primaryDims); ++lx) {
380                 for (int tx=1; tx < Global->numThreads; ++tx) {
381                         double *thrLatentDist = latentDist + tx * numLatentsPerThread;
382                         for (int px=0; px < numUnique; px++) {
383                                 int loc = px * numLatents + lx;
384                                 latentDist[loc] += thrLatentDist[loc];
385                         }
386                 }
387         }
388
389 #pragma omp parallel for num_threads(Global->numThreads)
390         for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
391                 for (int tx=1; tx < Global->numThreads; ++tx) {
392                         double *thrLatentDist = latentDist + tx * numLatentsPerThread;
393                         for (int px=0; px < numUnique; px++) {
394                                 int loc = px * numLatents + maxAbilities + triangleLoc0(sdim);
395                                 latentDist[loc] += thrLatentDist[loc];
396                         }
397                 }
398         }
399
400 #pragma omp parallel for num_threads(Global->numThreads)
401         for (int px=0; px < numUnique; px++) {
402                 if (!isfinite(patternLik[px])) {
403                         omxRaiseErrorf(globalState, "Likelihood of pattern %d is %.3g",
404                                        px, patternLik[px]);
405                 }
406
407                 double *latentDist1 = latentDist + px * numLatents;
408                 double weight = numIdentical[px] / patternLik[px];
409                 int cx = maxAbilities;
410                 for (int d1=0; d1 < primaryDims; d1++) {
411                         latentDist1[d1] *= weight;
412                         for (int d2=0; d2 <= d1; d2++) {
413                                 latentDist1[cx] *= weight;
414                                 ++cx;
415                         }
416                 }
417                 for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
418                         latentDist1[sdim] *= weight;
419                         int loc = maxAbilities + triangleLoc0(sdim);
420                         latentDist1[loc] *= weight;
421                 }
422 #if 0
423                 if (!isfinite(patternLik[px])) {
424                         error("Likelihood of row %d is %f", state->rowMap[px], patternLik[px]);
425                 }
426 #endif
427         }
428
429         //mxLog("raw latent after weighting");
430         //pda(latentDist, numLatents, numUnique);
431
432         std::vector<double> &ElatentMean = state->ElatentMean;
433         std::vector<double> &ElatentCov = state->ElatentCov;
434         
435         ElatentMean.assign(ElatentMean.size(), 0.0);
436         ElatentCov.assign(ElatentCov.size(), 0.0);
437
438 #pragma omp parallel for num_threads(Global->numThreads)
439         for (int d1=0; d1 < maxAbilities; d1++) {
440                 for (int px=0; px < numUnique; px++) {
441                         double *latentDist1 = latentDist + px * numLatents;
442                         int cx = maxAbilities + triangleLoc1(d1);
443                         if (d1 < primaryDims) {
444                                 ElatentMean[d1] += latentDist1[d1];
445                                 for (int d2=0; d2 <= d1; d2++) {
446                                         int cell = d2 * maxAbilities + d1;
447                                         ElatentCov[cell] += latentDist1[cx];
448                                         ++cx;
449                                 }
450                         } else {
451                                 ElatentMean[d1] += latentDist1[d1];
452                                 int cell = d1 * maxAbilities + d1;
453                                 int loc = maxAbilities + triangleLoc0(d1);
454                                 ElatentCov[cell] += latentDist1[loc];
455                         }
456                 }
457         }
458
459         //pda(ElatentMean.data(), 1, state->maxAbilities);
460         //pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);
461
462         omxData *data = state->data;
463         for (int d1=0; d1 < maxAbilities; d1++) {
464                 ElatentMean[d1] /= data->rows;
465         }
466
467         for (int d1=0; d1 < primaryDims; d1++) {
468                 for (int d2=0; d2 <= d1; d2++) {
469                         int cell = d2 * maxAbilities + d1;
470                         int tcell = d1 * maxAbilities + d2;
471                         ElatentCov[tcell] = ElatentCov[cell] =
472                                 ElatentCov[cell] / data->rows - ElatentMean[d1] * ElatentMean[d2];
473                 }
474         }
475         for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
476                 int cell = sdim * maxAbilities + sdim;
477                 ElatentCov[cell] = ElatentCov[cell] / data->rows - ElatentMean[sdim] * ElatentMean[sdim];
478         }
479
480         if (state->cacheLXK) state->LXKcached = TRUE;
481
482         Free(latentDist);
483
484         //mxLog("E-step");
485         //pda(ElatentMean.data(), 1, state->maxAbilities);
486         //pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);
487 }
488
489 static int getLatentVersion(BA81Expect *state)
490 {
491         return omxGetMatrixVersion(state->latentMeanOut) + omxGetMatrixVersion(state->latentCovOut);
492 }
493
494 // Attempt G-H grid? http://dbarajassolano.wordpress.com/2012/01/26/on-sparse-grid-quadratures/
495 static void ba81SetupQuadrature(omxExpectation* oo, int gridsize)
496 {
497         BA81Expect *state = (BA81Expect *) oo->argStruct;
498         if (state->verbose) {
499                 mxLog("%s: quadrature(%d)", oo->name, getLatentVersion(state));
500         }
501         int numUnique = state->numUnique;
502         int numThreads = Global->numThreads;
503         int maxDims = state->maxDims;
504         double Qwidth = state->Qwidth;
505         int numSpecific = state->numSpecific;
506         int priDims = maxDims - (numSpecific? 1 : 0);
507
508         // try starting small and increasing to the cap TODO
509         state->quadGridSize = gridsize;
510
511         state->totalQuadPoints = 1;
512         for (int dx=0; dx < maxDims; dx++) {
513                 state->totalQuadPoints *= state->quadGridSize;
514         }
515
516         state->totalPrimaryPoints = state->totalQuadPoints;
517
518         if (numSpecific) {
519                 state->totalPrimaryPoints /= state->quadGridSize;
520                 state->speQarea.resize(gridsize * numSpecific);
521         }
522
523         state->Qpoint.resize(gridsize);
524         state->priQarea.resize(state->totalPrimaryPoints);
525
526         double qgs = state->quadGridSize-1;
527         for (int px=0; px < state->quadGridSize; px ++) {
528                 state->Qpoint[px] = Qwidth - px * 2 * Qwidth / qgs;
529         }
530
531         //pda(state->latentMeanOut->data, 1, state->maxAbilities);
532         //pda(state->latentCovOut->data, state->maxAbilities, state->maxAbilities);
533
534         double totalArea = 0;
535         for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
536                 int quad[priDims];
537                 decodeLocation(qx, priDims, state->quadGridSize, quad);
538                 double where[priDims];
539                 pointToWhere(state, quad, where, priDims);
540                 state->priQarea[qx] = exp(dmvnorm(priDims, where,
541                                                   state->latentMeanOut->data,
542                                                   state->latentCovOut->data));
543                 totalArea += state->priQarea[qx];
544         }
545         for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
546                 state->priQarea[qx] /= totalArea;
547                 //mxLog("%.5g,", state->priQarea[qx]);
548         }
549
550         for (int sx=0; sx < numSpecific; sx++) {
551                 totalArea = 0;
552                 int covCell = (priDims + sx) * state->maxAbilities + priDims + sx;
553                 double mean = state->latentMeanOut->data[priDims + sx];
554                 double var = state->latentCovOut->data[covCell];
555                 //mxLog("setup[%d] %.2f %.2f", sx, mean, var);
556                 for (int qx=0; qx < state->quadGridSize; qx++) {
557                         double den = dnorm(state->Qpoint[qx], mean, sqrt(var), FALSE);
558                         state->speQarea[sIndex(state, sx, qx)] = den;
559                         totalArea += den;
560                 }
561                 for (int qx=0; qx < state->quadGridSize; qx++) {
562                         state->speQarea[sIndex(state, sx, qx)] /= totalArea;
563                 }
564                 //pda(state->speQarea.data() + sIndex(state, sx, 0), 1, state->quadGridSize);
565         }
566
567         if (!state->cacheLXK) {
568                 state->lxk = Realloc(state->lxk, numUnique * numThreads, double);
569         } else {
570                 int ns = state->numSpecific;
571                 if (ns == 0) ns = 1;
572                 long numOrdinate = ns * state->totalQuadPoints;
573                 state->lxk = Realloc(state->lxk, numUnique * numOrdinate, double);
574         }
575
576         state->expected = Realloc(state->expected, state->totalOutcomes * state->totalQuadPoints, double);
577 }
578
579 static void ba81buildLXKcache(omxExpectation *oo)
580 {
581         BA81Expect *state = (BA81Expect *) oo->argStruct;
582         if (!state->cacheLXK || state->LXKcached) return;
583         
584         ba81Estep1(oo);
585 }
586
587 OMXINLINE static void
588 expectedUpdate(omxData *data, const int *rowMap, const int px, const int item,
589                const double observed, double *out)
590 {
591         int pick = omxIntDataElementUnsafe(data, rowMap[px], item);
592         if (pick != NA_INTEGER) {
593                 out[pick-1] += observed;
594         }
595 }
596
597 OMXINLINE static void
598 ba81Expected(omxExpectation* oo)
599 {
600         BA81Expect *state = (BA81Expect*) oo->argStruct;
601         if (state->verbose) mxLog("%s: EM.expected", oo->name);
602
603         omxData *data = state->data;
604         int numSpecific = state->numSpecific;
605         const int *rowMap = state->rowMap;
606         double *patternLik = state->patternLik;
607         int *numIdentical = state->numIdentical;
608         int numUnique = state->numUnique;
609         int numItems = state->EitemParam->cols;
610         int totalOutcomes = state->totalOutcomes;
611         std::vector<int> &itemOutcomes = state->itemOutcomes;
612
613         OMXZERO(state->expected, totalOutcomes * state->totalQuadPoints);
614
615         if (numSpecific == 0) {
616 #pragma omp parallel for num_threads(Global->numThreads)
617                 for (long qx=0; qx < state->totalQuadPoints; qx++) {
618                         int thrId = omx_absolute_thread_num();
619                         double *lxk = ba81LikelihoodFast(oo, thrId, 0, qx);
620                         for (int px=0; px < numUnique; px++) {
621                                 double *out = state->expected + qx * totalOutcomes;
622                                 double observed = numIdentical[px] * lxk[px] / patternLik[px];
623                                 for (int ix=0; ix < numItems; ix++) {
624                                         const int outcomes = itemOutcomes[ix];
625                                         expectedUpdate(data, rowMap, px, ix, observed, out);
626                                         out += outcomes;
627                                 }
628                         }
629                 }
630         } else {
631                 long specificPoints = state->quadGridSize;
632
633 #pragma omp parallel for num_threads(Global->numThreads)
634                 for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
635                         int thrId = omx_absolute_thread_num();
636
637                         cai2010(oo, thrId, FALSE, qx);
638                         double *allElxk = eBase(state, thrId);
639                         double *Eslxk = esBase(state, thrId);
640
641                         for (long sx=0; sx < specificPoints; sx++) {
642                                 long qloc = qx * specificPoints + sx;
643
644                                 for (int sgroup=0; sgroup < numSpecific; sgroup++) {
645                                         double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, qloc);
646                                         double *myEslxk = Eslxk + sgroup * numUnique;
647
648                                         for (int px=0; px < numUnique; px++) {
649                                                 double *out = state->expected + totalOutcomes * qloc;
650
651                                                 for (int ix=0; ix < numItems; ix++) {
652                                                         const int outcomes = itemOutcomes[ix];
653                                                         if (state->Sgroup[ix] == sgroup) {
654                                                                 double Ei = allElxk[px];
655                                                                 double Eis = myEslxk[px];
656                                                                 double observed = (numIdentical[px] * (Ei / Eis) *
657                                                                                    (lxk[px] / patternLik[px]));
658                                                                 expectedUpdate(data, rowMap, px, ix, observed, out);
659                                                         }
660                                                         out += outcomes;
661                                                 }
662                                         }
663                                 }
664                         }
665                 }
666         }
667
668         if (!state->checkedBadData) {
669                 std::vector<double> byOutcome(totalOutcomes, 0);
670                 for (int ox=0; ox < totalOutcomes; ++ox) {
671                         for (long qx=0; qx < state->totalQuadPoints; qx++) {
672                                 byOutcome[ox] += state->expected[totalOutcomes * qx + ox];
673                         }
674                         if (byOutcome[ox] == 0) {
675                                 int uptoItem = 0;
676                                 for (size_t cx = 0; cx < itemOutcomes.size(); cx++) {
677                                         if (ox < uptoItem + itemOutcomes[cx]) {
678                                                 int bad = ox - uptoItem;
679                                                 omxRaiseErrorf(globalState, "Item %lu outcome %d is never endorsed.\n"
680                                                                "You must collapse categories or omit this item to estimate item parameters.", 1+cx, 1+bad);
681                                                 break;
682                                         }
683                                         uptoItem += itemOutcomes[cx];
684                                 }
685                         }
686                 }
687                 state->checkedBadData = TRUE;
688         }
689         //pda(state->expected, state->totalOutcomes, state->totalQuadPoints);
690 }
691
692 OMXINLINE static void
693 accumulateScores(BA81Expect *state, int px, int sgroup, double piece, const double *where,
694                  int primaryDims, int covEntries, std::vector<double> *mean, std::vector<double> *cov)
695 {
696         int maxDims = state->maxDims;
697         int maxAbilities = state->maxAbilities;
698
699         if (sgroup == 0) {
700                 int cx=0;
701                 for (int d1=0; d1 < primaryDims; d1++) {
702                         double piece_w1 = piece * where[d1];
703                         double &dest1 = (*mean)[px * maxAbilities + d1];
704 #pragma omp atomic
705                         dest1 += piece_w1;
706                         for (int d2=0; d2 <= d1; d2++) {
707                                 double &dest2 = (*cov)[px * covEntries + cx];
708 #pragma omp atomic
709                                 dest2 += where[d2] * piece_w1;
710                                 ++cx;
711                         }
712                 }
713         }
714
715         if (state->numSpecific) {
716                 int sdim = maxDims + sgroup - 1;
717                 double piece_w1 = piece * where[primaryDims];
718                 double &dest3 = (*mean)[px * maxAbilities + sdim];
719 #pragma omp atomic
720                 dest3 += piece_w1;
721
722                 double &dest4 = (*cov)[px * covEntries + triangleLoc0(sdim)];
723 #pragma omp atomic
724                 dest4 += piece_w1 * where[primaryDims];
725         }
726 }
727
728 static void
729 EAPinternalFast(omxExpectation *oo, std::vector<double> *mean, std::vector<double> *cov)
730 {
731         BA81Expect *state = (BA81Expect*) oo->argStruct;
732         if (state->verbose) mxLog("%s: EAP", oo->name);
733
734         int numUnique = state->numUnique;
735         int numSpecific = state->numSpecific;
736         int maxDims = state->maxDims;
737         int maxAbilities = state->maxAbilities;
738         int primaryDims = maxDims;
739         int covEntries = triangleLoc1(maxAbilities);
740
741         mean->assign(numUnique * maxAbilities, 0);
742         cov->assign(numUnique * covEntries, 0);
743
744         if (numSpecific == 0) {
745 #pragma omp parallel for num_threads(Global->numThreads)
746                 for (long qx=0; qx < state->totalQuadPoints; qx++) {
747                         const int thrId = omx_absolute_thread_num();
748                         int quad[maxDims];
749                         decodeLocation(qx, maxDims, state->quadGridSize, quad);
750                         double where[maxDims];
751                         pointToWhere(state, quad, where, maxDims);
752
753                         double *lxk = ba81LikelihoodFast(oo, thrId, 0, qx);
754
755                         double area = state->priQarea[qx];
756                         for (int px=0; px < numUnique; px++) {
757                                 double tmp = lxk[px] * area;
758                                 accumulateScores(state, px, 0, tmp, where, primaryDims, covEntries, mean, cov);
759                         }
760                 }
761         } else {
762                 primaryDims -= 1;
763                 int sDim = primaryDims;
764                 long specificPoints = state->quadGridSize;
765
766 #pragma omp parallel for num_threads(Global->numThreads)
767                 for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
768                         const int thrId = omx_absolute_thread_num();
769                         int quad[maxDims];
770                         decodeLocation(qx, primaryDims, state->quadGridSize, quad);
771
772                         cai2010(oo, thrId, FALSE, qx);
773                         double *allElxk = eBase(state, thrId);
774                         double *Eslxk = esBase(state, thrId);
775
776                         for (int sgroup=0; sgroup < numSpecific; sgroup++) {
777                                 for (long sx=0; sx < specificPoints; sx++) {
778                                         long qloc = qx * specificPoints + sx;
779                                         quad[sDim] = sx;
780                                         double where[maxDims];
781                                         pointToWhere(state, quad, where, maxDims);
782                                         double area = areaProduct(state, qx, sx, sgroup);
783                                         double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, qloc);
784                                         for (int px=0; px < numUnique; px++) {
785                                                 double Ei = allElxk[px];
786                                                 double Eis = Eslxk[sgroup * numUnique + px];
787                                                 double tmp = ((Ei / Eis) * lxk[px] * area);
788                                                 accumulateScores(state, px, sgroup, tmp, where, primaryDims,
789                                                                  covEntries, mean, cov);
790                                         }
791                                 }
792                         }
793                 }
794         }
795
796         double *patternLik = state->patternLik;
797         for (int px=0; px < numUnique; px++) {
798                 double denom = patternLik[px];
799                 for (int ax=0; ax < maxAbilities; ax++) {
800                         (*mean)[px * maxAbilities + ax] /= denom;
801                 }
802                 for (int cx=0; cx < triangleLoc1(primaryDims); ++cx) {
803                         (*cov)[px * covEntries + cx] /= denom;
804                 }
805                 for (int sx=0; sx < numSpecific; sx++) {
806                         (*cov)[px * covEntries + triangleLoc0(primaryDims + sx)] /= denom;
807                 }
808                 int cx=0;
809                 for (int a1=0; a1 < primaryDims; ++a1) {
810                         for (int a2=0; a2 <= a1; ++a2) {
811                                 double ma1 = (*mean)[px * maxAbilities + a1];
812                                 double ma2 = (*mean)[px * maxAbilities + a2];
813                                 (*cov)[px * covEntries + cx] -= ma1 * ma2;
814                                 ++cx;
815                         }
816                 }
817                 for (int sx=0; sx < numSpecific; sx++) {
818                         int sdim = primaryDims + sx;
819                         double ma1 = (*mean)[px * maxAbilities + sdim];
820                         (*cov)[px * covEntries + triangleLoc0(sdim)] -= ma1 * ma1;
821                 }
822         }
823 }
824
825 static void recomputePatternLik(omxExpectation *oo)
826 {
827         BA81Expect *estate = (BA81Expect*) oo->argStruct;
828         if (estate->verbose) mxLog("%s: patternLik", oo->name);
829
830         int numUnique = estate->numUnique;
831         int numSpecific = estate->numSpecific;
832         int maxDims = estate->maxDims;
833         int primaryDims = maxDims;
834         double *patternLik = estate->patternLik;
835         OMXZERO(patternLik, numUnique);
836
837         if (numSpecific == 0) {
838 #pragma omp parallel for num_threads(Global->numThreads)
839                 for (long qx=0; qx < estate->totalQuadPoints; qx++) {
840                         const int thrId = omx_absolute_thread_num();
841                         double area = estate->priQarea[qx];
842                         double *lxk = ba81LikelihoodFast(oo, thrId, 0, qx);
843
844                         for (int px=0; px < numUnique; px++) {
845                                 double tmp = (lxk[px] * area);
846 #pragma omp atomic
847                                 patternLik[px] += tmp;
848                         }
849                 }
850         } else {
851                 primaryDims -= 1;
852
853 #pragma omp parallel for num_threads(Global->numThreads)
854                 for (long qx=0; qx < estate->totalPrimaryPoints; qx++) {
855                         const int thrId = omx_absolute_thread_num();
856
857                         cai2010(oo, thrId, FALSE, qx);
858                         double *allElxk = eBase(estate, thrId);
859
860                         double priArea = estate->priQarea[qx];
861                         for (int px=0; px < numUnique; px++) {
862                                 double Ei = allElxk[px];
863                                 double tmp = (Ei * priArea);
864 #pragma omp atomic
865                                 patternLik[px] += tmp;
866                         }
867                 }
868         }
869 }
870
871 static void
872 ba81compute(omxExpectation *oo, const char *context)
873 {
874         BA81Expect *state = (BA81Expect *) oo->argStruct;
875
876         if (context) {
877                 if (strcmp(context, "EM")==0) {
878                         state->type = EXPECTATION_AUGMENTED;
879                 } else if (context[0] == 0) {
880                         state->type = EXPECTATION_OBSERVED;
881                 } else {
882                         omxRaiseErrorf(globalState, "Unknown context '%s'", context);
883                         return;
884                 }
885         }
886
887         omxRecompute(state->EitemParam);
888
889         bool itemClean = state->itemParamVersion == omxGetMatrixVersion(state->EitemParam);
890         bool latentClean = state->latentParamVersion == getLatentVersion(state);
891
892         if (state->verbose) {
893                 mxLog("%s: Qinit %d itemClean %d latentClean %d (1=clean)",
894                       oo->name, state->Qpoint.size() != 0, itemClean, latentClean);
895         }
896
897         if (state->Qpoint.size() == 0 || !latentClean) {
898                 ba81SetupQuadrature(oo, state->targetQpoints);
899         }
900         if (itemClean) {
901                 ba81buildLXKcache(oo);
902                 if (!latentClean) recomputePatternLik(oo);
903         } else {
904                 ba81OutcomeProb(state);
905                 ba81Estep1(oo);
906         }
907
908         if (state->type == EXPECTATION_AUGMENTED) {
909                 ba81Expected(oo);
910         }
911
912         state->itemParamVersion = omxGetMatrixVersion(state->EitemParam);
913         state->latentParamVersion = getLatentVersion(state);
914 }
915
916 static void
917 copyScore(int rows, int maxAbilities, std::vector<double> &mean,
918           std::vector<double> &cov, const int rx, double *scores, const int dest)
919 {
920         for (int ax=0; ax < maxAbilities; ++ax) {
921                 scores[rows * ax + dest] = mean[maxAbilities * rx + ax];
922         }
923         for (int ax=0; ax < maxAbilities; ++ax) {
924                 scores[rows * (maxAbilities + ax) + dest] =
925                         sqrt(cov[triangleLoc1(maxAbilities) * rx + triangleLoc0(ax)]);
926         }
927         for (int ax=0; ax < triangleLoc1(maxAbilities); ++ax) {
928                 scores[rows * (2*maxAbilities + ax) + dest] =
929                         cov[triangleLoc1(maxAbilities) * rx + ax];
930         }
931 }
932
933 /**
934  * MAP is not affected by the number of items. EAP is. Likelihood can
935  * get concentrated in a single quadrature ordinate. For 3PL, response
936  * patterns can have a bimodal likelihood. This will confuse MAP and
937  * is a key advantage of EAP (Thissen & Orlando, 2001, p. 136).
938  *
939  * Thissen, D. & Orlando, M. (2001). IRT for items scored in two
940  * categories. In D. Thissen & H. Wainer (Eds.), \emph{Test scoring}
941  * (pp 73-140). Lawrence Erlbaum Associates, Inc.
942  */
943 static void
944 ba81PopulateAttributes(omxExpectation *oo, SEXP robj)
945 {
946         BA81Expect *state = (BA81Expect *) oo->argStruct;
947         int maxAbilities = state->maxAbilities;
948
949         SEXP Rmean, Rcov;
950         PROTECT(Rmean = allocVector(REALSXP, maxAbilities));
951         memcpy(REAL(Rmean), state->ElatentMean.data(), maxAbilities * sizeof(double));
952
953         PROTECT(Rcov = allocMatrix(REALSXP, maxAbilities, maxAbilities));
954         memcpy(REAL(Rcov), state->ElatentCov.data(), maxAbilities * maxAbilities * sizeof(double));
955
956         setAttrib(robj, install("empirical.mean"), Rmean);
957         setAttrib(robj, install("empirical.cov"), Rcov);
958
959         if (state->type == EXPECTATION_AUGMENTED) {
960                 int numUnique = state->numUnique;
961                 int totalOutcomes = state->totalOutcomes;
962                 SEXP Rlik;
963                 SEXP Rexpected;
964
965                 PROTECT(Rlik = allocVector(REALSXP, numUnique));
966                 memcpy(REAL(Rlik), state->patternLik, sizeof(double) * numUnique);
967
968                 PROTECT(Rexpected = allocMatrix(REALSXP, totalOutcomes, state->totalQuadPoints));
969                 memcpy(REAL(Rexpected), state->expected, sizeof(double) * totalOutcomes * state->totalQuadPoints);
970
971                 setAttrib(robj, install("patternLikelihood"), Rlik);
972                 setAttrib(robj, install("em.expected"), Rexpected);
973         }
974
975         if (state->scores == SCORES_OMIT || state->type == EXPECTATION_UNINITIALIZED) return;
976
977         // TODO Wainer & Thissen. (1987). Estimating ability with the wrong
978         // model. Journal of Educational Statistics, 12, 339-368.
979
980         /*
981         int numQpoints = state->targetQpoints * 2;  // make configurable TODO
982
983         if (numQpoints < 1 + 2.0 * sqrt(state->itemSpec->cols)) {
984                 // Thissen & Orlando (2001, p. 136)
985                 warning("EAP requires at least 2*sqrt(items) quadrature points");
986         }
987
988         ba81SetupQuadrature(oo, numQpoints, 0);
989         ba81Estep1(oo);
990         */
991
992         std::vector<double> mean;
993         std::vector<double> cov;
994         EAPinternalFast(oo, &mean, &cov);
995
996         int numUnique = state->numUnique;
997         omxData *data = state->data;
998         int rows = state->scores == SCORES_FULL? data->rows : numUnique;
999         int cols = 2 * maxAbilities + triangleLoc1(maxAbilities);
1000         SEXP Rscores;
1001         PROTECT(Rscores = allocMatrix(REALSXP, rows, cols));
1002         double *scores = REAL(Rscores);
1003
1004         const int SMALLBUF = 10;
1005         char buf[SMALLBUF];
1006         SEXP names;
1007         PROTECT(names = allocVector(STRSXP, cols));
1008         for (int nx=0; nx < maxAbilities; ++nx) {
1009                 snprintf(buf, SMALLBUF, "s%d", nx+1);
1010                 SET_STRING_ELT(names, nx, mkChar(buf));
1011                 snprintf(buf, SMALLBUF, "se%d", nx+1);
1012                 SET_STRING_ELT(names, maxAbilities + nx, mkChar(buf));
1013         }
1014         for (int nx=0; nx < triangleLoc1(maxAbilities); ++nx) {
1015                 snprintf(buf, SMALLBUF, "cov%d", nx+1);
1016                 SET_STRING_ELT(names, maxAbilities*2 + nx, mkChar(buf));
1017         }
1018         SEXP dimnames;
1019         PROTECT(dimnames = allocVector(VECSXP, 2));
1020         SET_VECTOR_ELT(dimnames, 1, names);
1021         setAttrib(Rscores, R_DimNamesSymbol, dimnames);
1022
1023         if (state->scores == SCORES_FULL) {
1024 #pragma omp parallel for num_threads(Global->numThreads)
1025                 for (int rx=0; rx < numUnique; rx++) {
1026                         int dups = omxDataNumIdenticalRows(state->data, state->rowMap[rx]);
1027                         for (int dup=0; dup < dups; dup++) {
1028                                 int dest = omxDataIndex(data, state->rowMap[rx]+dup);
1029                                 copyScore(rows, maxAbilities, mean, cov, rx, scores, dest);
1030                         }
1031                 }
1032         } else {
1033 #pragma omp parallel for num_threads(Global->numThreads)
1034                 for (int rx=0; rx < numUnique; rx++) {
1035                         copyScore(rows, maxAbilities, mean, cov, rx, scores, rx);
1036                 }
1037         }
1038
1039         setAttrib(robj, install("scores.out"), Rscores);
1040 }
1041
1042 static void ba81Destroy(omxExpectation *oo) {
1043         if(OMX_DEBUG) {
1044                 mxLog("Freeing %s function.", oo->name);
1045         }
1046         BA81Expect *state = (BA81Expect *) oo->argStruct;
1047         omxFreeAllMatrixData(state->EitemParam);
1048         omxFreeAllMatrixData(state->design);
1049         omxFreeAllMatrixData(state->latentMeanOut);
1050         omxFreeAllMatrixData(state->latentCovOut);
1051         omxFreeAllMatrixData(state->customPrior);
1052         omxFreeAllMatrixData(state->itemParam);
1053         Free(state->numIdentical);
1054         Free(state->rowMap);
1055         Free(state->patternLik);
1056         Free(state->lxk);
1057         Free(state->Eslxk);
1058         Free(state->allElxk);
1059         Free(state->Sgroup);
1060         Free(state->expected);
1061         Free(state->outcomeProb);
1062         delete state;
1063 }
1064
1065 void getMatrixDims(SEXP r_theta, int *rows, int *cols)
1066 {
1067     SEXP matrixDims;
1068     PROTECT(matrixDims = getAttrib(r_theta, R_DimSymbol));
1069     int *dimList = INTEGER(matrixDims);
1070     *rows = dimList[0];
1071     *cols = dimList[1];
1072     UNPROTECT(1);
1073 }
1074
1075 static void ignoreSetVarGroup(omxExpectation*, FreeVarGroup *)
1076 {}
1077
1078 void omxInitExpectationBA81(omxExpectation* oo) {
1079         omxState* currentState = oo->currentState;      
1080         SEXP rObj = oo->rObj;
1081         SEXP tmp;
1082         
1083         if(OMX_DEBUG) {
1084                 mxLog("Initializing %s.", oo->name);
1085         }
1086         if (!rpf_model) {
1087                 if (0) {
1088                         const int wantVersion = 3;
1089                         int version;
1090                         get_librpf_t get_librpf = (get_librpf_t) R_GetCCallable("rpf", "get_librpf_model_GPL");
1091                         (*get_librpf)(&version, &rpf_numModels, &rpf_model);
1092                         if (version < wantVersion) error("librpf binary API %d installed, at least %d is required",
1093                                                          version, wantVersion);
1094                 } else {
1095                         rpf_numModels = librpf_numModels;
1096                         rpf_model = librpf_model;
1097                 }
1098         }
1099         
1100         BA81Expect *state = new BA81Expect;
1101         state->checkedBadData = FALSE;
1102         state->numSpecific = 0;
1103         state->numIdentical = NULL;
1104         state->rowMap = NULL;
1105         state->design = NULL;
1106         state->lxk = NULL;
1107         state->patternLik = NULL;
1108         state->Eslxk = NULL;
1109         state->allElxk = NULL;
1110         state->outcomeProb = NULL;
1111         state->expected = NULL;
1112         state->type = EXPECTATION_UNINITIALIZED;
1113         state->scores = SCORES_OMIT;
1114         state->itemParam = NULL;
1115         state->customPrior = NULL;
1116         state->itemParamVersion = 0;
1117         state->latentParamVersion = 0;
1118         oo->argStruct = (void*) state;
1119
1120         PROTECT(tmp = GET_SLOT(rObj, install("data")));
1121         state->data = omxDataLookupFromState(tmp, currentState);
1122
1123         if (strcmp(omxDataType(state->data), "raw") != 0) {
1124                 omxRaiseErrorf(currentState, "%s unable to handle data type %s", oo->name, omxDataType(state->data));
1125                 return;
1126         }
1127
1128         PROTECT(tmp = GET_SLOT(rObj, install("ItemSpec")));
1129         for (int sx=0; sx < length(tmp); ++sx) {
1130                 SEXP model = VECTOR_ELT(tmp, sx);
1131                 if (!OBJECT(model)) {
1132                         error("Item models must inherit rpf.base");
1133                 }
1134                 SEXP spec;
1135                 PROTECT(spec = GET_SLOT(model, install("spec")));
1136                 state->itemSpec.push_back(REAL(spec));
1137         }
1138
1139         PROTECT(tmp = GET_SLOT(rObj, install("design")));
1140         if (!isNull(tmp)) {
1141                 // better to demand integers and not coerce to real TODO
1142                 state->design = omxNewMatrixFromRPrimitive(tmp, globalState, FALSE, 0);
1143         }
1144
1145         state->latentMeanOut = omxNewMatrixFromSlot(rObj, currentState, "mean");
1146         if (!state->latentMeanOut) error("Failed to retrieve mean matrix");
1147         state->latentCovOut  = omxNewMatrixFromSlot(rObj, currentState, "cov");
1148         if (!state->latentCovOut) error("Failed to retrieve cov matrix");
1149
1150         state->EitemParam =
1151                 omxNewMatrixFromSlot(rObj, currentState, "EItemParam");
1152         if (!state->EitemParam) error("Must supply EItemParam");
1153
1154         state->itemParam =
1155                 omxNewMatrixFromSlot(rObj, globalState, "ItemParam");
1156
1157         if (state->EitemParam->rows != state->itemParam->rows ||
1158             state->EitemParam->cols != state->itemParam->cols) {
1159                 error("ItemParam and EItemParam must be of the same dimension");
1160         }
1161
1162         oo->computeFun = ba81compute;
1163         oo->setVarGroup = ignoreSetVarGroup;
1164         oo->destructFun = ba81Destroy;
1165         oo->populateAttrFun = ba81PopulateAttributes;
1166         
1167         // TODO: Exactly identical rows do not contribute any information.
1168         // The sorting algorithm ought to remove them so we don't waste RAM.
1169         // The following summary stats would be cheaper to calculate too.
1170
1171         int numUnique = 0;
1172         omxData *data = state->data;
1173         if (omxDataNumFactor(data) != data->cols) {
1174                 // verify they are ordered factors TODO
1175                 omxRaiseErrorf(currentState, "%s: all columns must be factors", oo->name);
1176                 return;
1177         }
1178
1179         for (int rx=0; rx < data->rows;) {
1180                 rx += omxDataNumIdenticalRows(state->data, rx);
1181                 ++numUnique;
1182         }
1183         state->numUnique = numUnique;
1184
1185         state->rowMap = Realloc(NULL, numUnique, int);
1186         state->numIdentical = Realloc(NULL, numUnique, int);
1187
1188         state->customPrior =
1189                 omxNewMatrixFromSlot(rObj, globalState, "CustomPrior");
1190         
1191         int numItems = state->EitemParam->cols;
1192         if (data->cols != numItems) {
1193                 error("Data has %d columns for %d items", data->cols, numItems);
1194         }
1195
1196         for (int rx=0, ux=0; rx < data->rows; ux++) {
1197                 if (rx == 0) {
1198                         // all NA rows will sort to the top
1199                         int na=0;
1200                         for (int ix=0; ix < numItems; ix++) {
1201                                 if (omxIntDataElement(data, 0, ix) == NA_INTEGER) { ++na; }
1202                         }
1203                         if (na == numItems) {
1204                                 omxRaiseErrorf(currentState, "Remove rows with all NAs");
1205                                 return;
1206                         }
1207                 }
1208                 int dups = omxDataNumIdenticalRows(state->data, rx);
1209                 state->numIdentical[ux] = dups;
1210                 state->rowMap[ux] = rx;
1211                 rx += dups;
1212         }
1213
1214         int numThreads = Global->numThreads;
1215
1216         int maxSpec = 0;
1217         int maxParam = 0;
1218         state->maxDims = 0;
1219
1220         std::vector<int> &itemOutcomes = state->itemOutcomes;
1221         itemOutcomes.resize(numItems);
1222         int totalOutcomes = 0;
1223         for (int cx = 0; cx < data->cols; cx++) {
1224                 const double *spec = state->itemSpec[cx];
1225                 int id = spec[RPF_ISpecID];
1226                 int dims = spec[RPF_ISpecDims];
1227                 if (state->maxDims < dims)
1228                         state->maxDims = dims;
1229
1230                 int no = spec[RPF_ISpecOutcomes];
1231                 itemOutcomes[cx] = no;
1232                 totalOutcomes += no;
1233
1234                 // TODO this summary stat should be available from omxData
1235                 int dataMax=0;
1236                 for (int rx=0; rx < data->rows; rx++) {
1237                         int pick = omxIntDataElementUnsafe(data, rx, cx);
1238                         if (dataMax < pick)
1239                                 dataMax = pick;
1240                 }
1241                 if (dataMax > no) {
1242                         error("Data for item %d has %d outcomes, not %d", cx+1, dataMax, no);
1243                 } else if (dataMax < no) {
1244                         warning("Data for item %d has only %d outcomes, not %d", cx+1, dataMax, no);
1245                         // promote to error?
1246                         // should complain if an outcome is not represented in the data TODO
1247                 }
1248
1249                 int numSpec = (*rpf_model[id].numSpec)(spec);
1250                 if (maxSpec < numSpec)
1251                         maxSpec = numSpec;
1252
1253                 int numParam = (*rpf_model[id].numParam)(spec);
1254                 if (maxParam < numParam)
1255                         maxParam = numParam;
1256         }
1257
1258         state->totalOutcomes = totalOutcomes;
1259
1260         if (int(state->itemSpec.size()) != data->cols) {
1261                 omxRaiseErrorf(currentState, "ItemSpec must contain %d item model specifications",
1262                                data->cols);
1263                 return;
1264         }
1265         if (state->EitemParam->rows != maxParam) {
1266                 omxRaiseErrorf(currentState, "ItemParam should have %d rows", maxParam);
1267                 return;
1268         }
1269
1270         if (state->design == NULL) {
1271                 state->maxAbilities = state->maxDims;
1272                 state->design = omxInitTemporaryMatrix(NULL, state->maxDims, numItems,
1273                                        TRUE, currentState);
1274                 for (int ix=0; ix < numItems; ix++) {
1275                         const double *spec = state->itemSpec[ix];
1276                         int dims = spec[RPF_ISpecDims];
1277                         for (int dx=0; dx < state->maxDims; dx++) {
1278                                 omxSetMatrixElement(state->design, dx, ix, dx < dims? (double)dx+1 : nan(""));
1279                         }
1280                 }
1281         } else {
1282                 omxMatrix *design = state->design;
1283                 if (design->cols != numItems ||
1284                     design->rows != state->maxDims) {
1285                         omxRaiseErrorf(currentState, "Design matrix should have %d rows and %d columns",
1286                                        state->maxDims, numItems);
1287                         return;
1288                 }
1289
1290                 state->maxAbilities = 0;
1291                 for (int ix=0; ix < design->rows * design->cols; ix++) {
1292                         double got = design->data[ix];
1293                         if (!R_FINITE(got)) continue;
1294                         if (round(got) != (int)got) error("Design matrix can only contain integers"); // TODO better way?
1295                         if (state->maxAbilities < got)
1296                                 state->maxAbilities = got;
1297                 }
1298                 for (int ix=0; ix < design->cols; ix++) {
1299                         const double *idesign = omxMatrixColumn(design, ix);
1300                         int ddim = 0;
1301                         for (int rx=0; rx < design->rows; rx++) {
1302                                 if (isfinite(idesign[rx])) ddim += 1;
1303                         }
1304                         const double *spec = state->itemSpec[ix];
1305                         int dims = spec[RPF_ISpecDims];
1306                         if (ddim > dims) error("Item %d has %d dims but design assigns %d", ix, dims, ddim);
1307                 }
1308         }
1309         if (state->maxAbilities <= state->maxDims) {
1310                 state->Sgroup = Calloc(numItems, int);
1311         } else {
1312                 // Not sure if this is correct, revisit TODO
1313                 int Sgroup0 = -1;
1314                 state->Sgroup = Realloc(NULL, numItems, int);
1315                 for (int dx=0; dx < state->maxDims; dx++) {
1316                         for (int ix=0; ix < numItems; ix++) {
1317                                 int ability = omxMatrixElement(state->design, dx, ix);
1318                                 if (dx < state->maxDims - 1) {
1319                                         if (Sgroup0 <= ability)
1320                                                 Sgroup0 = ability+1;
1321                                         continue;
1322                                 }
1323                                 int ss=-1;
1324                                 if (ability >= Sgroup0) {
1325                                         if (ss == -1) {
1326                                                 ss = ability;
1327                                         } else {
1328                                                 omxRaiseErrorf(currentState, "Item %d cannot belong to more than "
1329                                                                "1 specific dimension (both %d and %d)",
1330                                                                ix, ss, ability);
1331                                                 return;
1332                                         }
1333                                 }
1334                                 if (ss == -1) ss = Sgroup0;
1335                                 state->Sgroup[ix] = ss - Sgroup0;
1336                         }
1337                 }
1338                 state->numSpecific = state->maxAbilities - state->maxDims + 1;
1339                 state->allElxk = Realloc(NULL, numUnique * numThreads, double);
1340                 state->Eslxk = Realloc(NULL, numUnique * state->numSpecific * numThreads, double);
1341         }
1342
1343         if (state->latentMeanOut->rows * state->latentMeanOut->cols != state->maxAbilities) {
1344                 error("The mean matrix '%s' must be 1x%d or %dx1", state->latentMeanOut->name,
1345                       state->maxAbilities, state->maxAbilities);
1346         }
1347         if (state->latentCovOut->rows != state->maxAbilities ||
1348             state->latentCovOut->cols != state->maxAbilities) {
1349                 error("The cov matrix '%s' must be %dx%d",
1350                       state->latentCovOut->name, state->maxAbilities, state->maxAbilities);
1351         }
1352
1353         PROTECT(tmp = GET_SLOT(rObj, install("verbose")));
1354         state->verbose = asLogical(tmp);
1355
1356         PROTECT(tmp = GET_SLOT(rObj, install("cache")));
1357         state->cacheLXK = asLogical(tmp);
1358         state->LXKcached = FALSE;
1359
1360         PROTECT(tmp = GET_SLOT(rObj, install("qpoints")));
1361         state->targetQpoints = asReal(tmp);
1362
1363         PROTECT(tmp = GET_SLOT(rObj, install("qwidth")));
1364         state->Qwidth = asReal(tmp);
1365
1366         PROTECT(tmp = GET_SLOT(rObj, install("scores")));
1367         const char *score_option = CHAR(asChar(tmp));
1368         if (strcmp(score_option, "omit")==0) state->scores = SCORES_OMIT;
1369         if (strcmp(score_option, "unique")==0) state->scores = SCORES_UNIQUE;
1370         if (strcmp(score_option, "full")==0) state->scores = SCORES_FULL;
1371
1372         state->ElatentMean.resize(state->maxAbilities);
1373         state->ElatentCov.resize(state->maxAbilities * state->maxAbilities);
1374
1375         // verify data bounded between 1 and numOutcomes TODO
1376         // hm, looks like something could be added to omxData for column summary stats?
1377 }