<commit_before>/*
  Copyright 2012-2013 <NAME> and contributors

  This is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <Rmath.h>
#include "omxExpectationBA81.h"
#include "npsolWrap.h"
#include "libifa-rpf.h"
#include "dmvnorm.h"

const struct rpf *rpf_model = NULL;
int rpf_numModels;
static const double MIN_PATTERNLIK = 1e-100;

void pda(const double *ar, int rows, int cols)
{
	std::string buf;
	for (int rx=0; rx < rows; rx++) {   // column major order
		for (int cx=0; cx < cols; cx++) {
			buf += string_snprintf("%.6g, ", ar[cx * rows + rx]);
		}
		buf += "\n";
	}
	mxLogBig(buf);
}

void pia(const int *ar, int rows, int cols)
{
	std::string buf;
	for (int rx=0; rx < rows; rx++) {   // column major order
		for (int cx=0; cx < cols; cx++) {
			buf += string_snprintf("%d, ", ar[cx * rows + rx]);
		}
		buf += "\n";
	}
	mxLogBig(buf);
}

// state->speQarea[sIndex(state, sx, qx)]
OMXINLINE static
int sIndex(BA81Expect *state, int sx, int qx)
{
	//if (sx < 0 || sx >= state->numSpecific) error("Out of domain");
	//if (qx < 0 || qx >= state->quadGridSize) error("Out of domain");
	return sx * state->quadGridSize + qx;
}

/**
 * \param theta Vector of ability parameters, one per ability
 * \returns A numItems by maxOutcomes colMajor vector of doubles. Caller must Free it.
 */
double *computeRPF(BA81Expect *state, omxMatrix *itemParam, const int *quad, const bool wantlog)
{
	omxMatrix *design = state->design;
	int maxDims = state->maxDims;
	int maxOutcomes = state->maxOutcomes;
	size_t numItems = state->itemSpec.size();

	double theta[maxDims];
	pointToWhere(state, quad, theta, maxDims);

	double *outcomeProb = Realloc(NULL, numItems * maxOutcomes, double);
	//double *outcomeProb = Calloc(numItems * maxOutcomes, double);

	for (size_t ix=0; ix < numItems; ix++) {
		const double *spec = state->itemSpec[ix];
		double *iparam = omxMatrixColumn(itemParam, ix);
		double *out = outcomeProb + ix * maxOutcomes;
		int id = spec[RPF_ISpecID];
		int dims = spec[RPF_ISpecDims];
		double ptheta[dims];

		for (int dx=0; dx < dims; dx++) {
			int ability = (int)omxMatrixElement(design, dx, ix) - 1;
			if (ability >= maxDims) ability = maxDims-1;
			ptheta[dx] = theta[ability];
		}

		if (wantlog) {
			(*rpf_model[id].logprob)(spec, iparam, ptheta, out);
		} else {
			(*rpf_model[id].prob)(spec, iparam, ptheta, out);
		}
#if 0
		for (int ox=0; ox < spec[RPF_ISpecOutcomes]; ox++) {
			if (!isfinite(out[ox]) || out[ox] > 0) {
				mxLog("item param");
				pda(iparam, itemParam->rows, 1);
				mxLog("where");
				pda(ptheta, dims, 1);
				error("RPF returned %20.20f", out[ox]);
			}
		}
#endif
	}

	return outcomeProb;
}

OMXINLINE static double *
getLXKcache(BA81Expect *state, const int *quad, const int specific)
{
	long ordinate;
	if (state->numSpecific == 0) {
		ordinate = encodeLocation(state->maxDims, state->quadGridSize, quad);
	} else {
		ordinate = (specific * state->totalQuadPoints +
			    encodeLocation(state->maxDims, state->quadGridSize, quad));
	}
	return state->lxk + state->numUnique * ordinate;
}

OMXINLINE static double *
ba81Likelihood(omxExpectation *oo, const int thrId, int specific, const int *quad, const double *outcomeProb)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	int numUnique = state->numUnique;
	int maxOutcomes = state->maxOutcomes;
	omxData *data = state->data;
	size_t numItems = state->itemSpec.size();
	int *Sgroup = state->Sgroup;
	double *lxk;

	if (!state->cacheLXK) {
		lxk = state->lxk + numUnique * thrId;
	} else {
		lxk = getLXKcache(state, quad, specific);
	}

	const int *rowMap = state->rowMap;
	for (int px=0; px < numUnique; px++) {
		double lxk1 = 1;
		for (size_t ix=0; ix < numItems; ix++) {
			if (specific != Sgroup[ix]) continue;
			int pick = omxIntDataElementUnsafe(data, rowMap[px], ix);
			if (pick == NA_INTEGER) continue;
			double piece = outcomeProb[ix * maxOutcomes + pick-1];  // move -1 elsewhere TODO
			lxk1 *= piece;
			//mxLog("%d pick %d piece %.3f", ix, pick, piece);
		}
#if 0
#pragma omp critical(ba81LikelihoodDebug1)
		if (!isfinite(lxk1) || lxk1 > numItems) {
			mxLog("where");
			double where[state->maxDims];
			pointToWhere(state, quad, where, state->maxDims);
			pda(where, state->maxDims, 1);
			mxLog("prob");
			pda(outcomeProb, numItems, maxOutcomes);
			error("Likelihood of row %d is %f", rowMap[px], lxk1);
		}
#endif
		lxk[px] = lxk1;
	}

	//mxLog("L.is(%d) at (%d %d) %.2f", specific, quad[0], quad[1], lxk[0]);
	return lxk;
}

double *ba81LikelihoodFast(omxExpectation *oo, const int thrId, int specific, const int *quad)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (!state->cacheLXK) {
		const double *outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
		double *ret = ba81Likelihood(oo, thrId, specific, quad, outcomeProb);
		Free(outcomeProb);
		return ret;
	} else {
		return getLXKcache(state, quad, specific);
	}

}

OMXINLINE static void
mapLatentSpace(BA81Expect *state, int sgroup, double piece, const double *where,
	       const std::vector<double> &whereGram, double *latentDist)
{
	int maxDims = state->maxDims;
	int maxAbilities = state->maxAbilities;
	int pmax = maxDims;
	if (state->numSpecific) pmax -= 1;

	if (sgroup == 0) {
		int gx = 0;
		int cx = maxAbilities;
		for (int d1=0; d1 < pmax; d1++) {
			double piece_w1 = piece * where[d1];
			latentDist[d1] += piece_w1;
			for (int d2=0; d2 <= d1; d2++) {
				double piece_cov = piece * whereGram[gx];
				latentDist[cx] += piece_cov;
				++cx; ++gx;
			}
		}
	}

	if (state->numSpecific) {
		int sdim = pmax + sgroup;
		double piece_w1 = piece * where[pmax];
		latentDist[sdim] += piece_w1;

		double piece_var = piece * whereGram[triangleLoc0(pmax)];
		int to = maxAbilities + triangleLoc0(sdim);
		latentDist[to] += piece_var;
	}
}

// Eslxk, allElxk (Ei, Eis) depend on the ordinate of the primary dimensions
void cai2010(omxExpectation* oo, const int thrId, int recompute, const int *primaryQuad)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	int numUnique = state->numUnique;
	int numSpecific = state->numSpecific;
	int maxDims = state->maxDims;
	int sDim = maxDims-1;
	int quadGridSize = state->quadGridSize;
	int quad[maxDims];
	memcpy(quad, primaryQuad, sizeof(int)*sDim);
	double *allElxk = eBase(state, thrId);
	double *Eslxk = esBase(state, thrId);

	for (int px=0; px < numUnique; px++) {
		allElxk[px] = 1;
		for (int sx=0; sx < numSpecific; sx++) {
			Eslxk[sx * numUnique + px] = 0;
		}
	}

	if (!state->cacheLXK) recompute = TRUE;

	for (int qx=0; qx < quadGridSize; qx++) {
		quad[sDim] = qx;
		double *outcomeProb = NULL;
		if (recompute) {
			outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
		}
		for (int sx=0; sx < numSpecific; sx++) {
			double *lxk;     // a.k.a. "L_is"
			if (recompute) {
				lxk = ba81Likelihood(oo, thrId, sx, quad, outcomeProb);
			} else {
				lxk = getLXKcache(state, quad, sx);
			}

			for (int ix=0; ix < numUnique; ix++) {
				double area = state->speQarea[sIndex(state, sx, qx)];
				double piece = lxk[ix] * area;
				//mxLog("E.is(%d) at (%d, %d) %.2f + %.2f = %.2f",
				//  sx, primaryQuad[0], qx, lxk[ix], area, piece);
				Eslxk[sx * numUnique + ix] += piece;
			}
		}
		Free(outcomeProb);
	}

	for (int sx=0; sx < numSpecific; sx++) {
		for (int px=0; px < numUnique; px++) {
			//mxLog("E.is(%d) at (%d) %.2f", sx, primaryQuad[0], state->Eslxk[esIndex(state, sx, 0)]);
			allElxk[px] *= Eslxk[sx * numUnique + px];  // allSlxk a.k.a. "E_i"
		}
	}
}

static void ba81Estep1(omxExpectation *oo)
{
	if(OMX_DEBUG) {mxLog("Beginning %s Computation.", oo->name);}

	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (state->verbose) {
		mxLog("%s: lxk(%d) patternLik ElatentMean ElatentCov",
		      oo->name, omxGetMatrixVersion(state->EitemParam));
	}

	int numUnique = state->numUnique;
	int numSpecific = state->numSpecific;
	int maxDims = state->maxDims;
	int maxAbilities = state->maxAbilities;
	int primaryDims = maxDims;

	state->patternLik = Realloc(state->patternLik, numUnique, double);
	double *patternLik = state->patternLik;
	OMXZERO(patternLik, numUnique);

	int numLatents = maxAbilities + triangleLoc1(maxAbilities);
	int numLatentsPerThread = numUnique * numLatents;
	double *latentDist = Calloc(numUnique * numLatents * Global->numThreads, double);

	// E-step, marginalize person ability
	//
	// Note: In the notation of <NAME> (1981) and
	// Cai~(2010), these loops are reversed.  That is, the inner
	// loop is over quadrature points and the outer loop is over
	// all response patterns.
	//
	if (numSpecific == 0) {
#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalQuadPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);
			double where[maxDims];
			pointToWhere(state, quad, where, maxDims);
			std::vector<double> whereGram(triangleLoc1(maxDims));
			gramProduct(where, maxDims, whereGram.data());

			const double *outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
			double *lxk = ba81Likelihood(oo, thrId, 0, quad, outcomeProb);
			Free(outcomeProb);

			double area = state->priQarea[qx];
			for (int px=0; px < numUnique; px++) {
				double tmp = lxk[px] * area;
#if 0
				if (!isfinite(tmp)) {
					mxLog("where");
					pda(where, maxDims, 1);
					error("Row %d lxk %f logArea %f tmp %f",
					      state->rowMap[px], lxk[px], logArea, tmp);
				}
#endif
#pragma omp atomic
				patternLik[px] += tmp;
				mapLatentSpace(state, 0, tmp, where, whereGram,
					       thrLatentDist + px * numLatents);
			}
		}
	} else {
		primaryDims -= 1;
		int sDim = primaryDims;
		long specificPoints = state->quadGridSize;

#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
			int quad[maxDims];
			decodeLocation(qx, primaryDims, state->quadGridSize, quad);

			cai2010(oo, thrId, TRUE, quad);
			double *allElxk = eBase(state, thrId);
			double *Eslxk = esBase(state, thrId);

			for (long sx=0; sx < specificPoints; sx++) {
				quad[sDim] = sx;
				double where[maxDims];
				pointToWhere(state, quad, where, maxDims);
				std::vector<double> whereGram(triangleLoc1(maxDims));
				gramProduct(where, maxDims, whereGram.data());

				for (int sgroup=0; sgroup < numSpecific; sgroup++) {
					double area = areaProduct(state, quad, sgroup);
					double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, quad);
					for (int px=0; px < numUnique; px++) {
						double Ei = allElxk[px];
						double Eis = Eslxk[sgroup * numUnique + px];
						double tmp = ((Ei / Eis) * lxk[px] * area);
						mapLatentSpace(state, sgroup, tmp, where, whereGram,
							       thrLatentDist + px * numLatents);
					}
				}
			}

			double priArea = state->priQarea[qx];
			for (int px=0; px < numUnique; px++) {
				double Ei = allElxk[px];
				double tmp = (Ei * priArea);
#pragma omp atomic
				patternLik[px] += tmp;
			}
		}
	}

	int *numIdentical = state->numIdentical;

	//mxLog("raw latent");
	//pda(latentDist, numLatents, numUnique);

#pragma omp parallel for num_threads(Global->numThreads) schedule(dynamic)
	for (int lx=0; lx < maxAbilities + triangleLoc1(primaryDims); ++lx) {
		for (int tx=1; tx < Global->numThreads; ++tx) {
			double *thrLatentDist = latentDist + tx * numLatentsPerThread;
			for (int px=0; px < numUnique; px++) {
				int loc = px * numLatents + lx;
				latentDist[loc] += thrLatentDist[loc];
			}
		}
	}

#pragma omp parallel for num_threads(Global->numThreads)
	for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
		for (int tx=1; tx < Global->numThreads; ++tx) {
			double *thrLatentDist = latentDist + tx * numLatentsPerThread;
			for (int px=0; px < numUnique; px++) {
				int loc = px * numLatents + maxAbilities + triangleLoc0(sdim);
				latentDist[loc] += thrLatentDist[loc];
			}
		}
	}

#pragma omp parallel for num_threads(Global->numThreads)
	for (int px=0; px < numUnique; px++) {
		if (patternLik[px] < MIN_PATTERNLIK) {
			patternLik[px] = MIN_PATTERNLIK;
			warning("Likelihood of pattern %d is 0, forcing to %.3g",
				px, MIN_PATTERNLIK);
		}

		double *latentDist1 = latentDist + px * numLatents;
		double weight = numIdentical[px] / patternLik[px];
		int cx = maxAbilities;
		for (int d1=0; d1 < primaryDims; d1++) {
			latentDist1[d1] *= weight;
			for (int d2=0; d2 <= d1; d2++) {
				latentDist1[cx] *= weight;
				++cx;
			}
		}
		for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
			latentDist1[sdim] *= weight;
			int loc = maxAbilities + triangleLoc0(sdim);
			latentDist1[loc] *= weight;
		}
#if 0
		if (!isfinite(patternLik[px])) {
			error("Likelihood of row %d is %f", state->rowMap[px], patternLik[px]);
		}
#endif
	}

	//mxLog("raw latent after weighting");
	//pda(latentDist, numLatents, numUnique);

	std::vector<double> &ElatentMean = state->ElatentMean;
	std::vector<double> &ElatentCov = state->ElatentCov;
	
	ElatentMean.assign(ElatentMean.size(), 0.0);
	ElatentCov.assign(ElatentCov.size(), 0.0);

#pragma omp parallel for num_threads(Global->numThreads)
	for (int d1=0; d1 < maxAbilities; d1++) {
		for (int px=0; px < numUnique; px++) {
			double *latentDist1 = latentDist + px * numLatents;
			int cx = maxAbilities + triangleLoc1(d1);
			if (d1 < primaryDims) {
				ElatentMean[d1] += latentDist1[d1];
				for (int d2=0; d2 <= d1; d2++) {
					int cell = d2 * maxAbilities + d1;
					ElatentCov[cell] += latentDist1[cx];
					++cx;
				}
			} else {
				ElatentMean[d1] += latentDist1[d1];
				int cell = d1 * maxAbilities + d1;
				int loc = maxAbilities + triangleLoc0(d1);
				ElatentCov[cell] += latentDist1[loc];
			}
		}
	}

	//pda(ElatentMean.data(), 1, state->maxAbilities);
	//pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);

	omxData *data = state->data;
	for (int d1=0; d1 < maxAbilities; d1++) {
		ElatentMean[d1] /= data->rows;
	}

	for (int d1=0; d1 < primaryDims; d1++) {
		for (int d2=0; d2 <= d1; d2++) {
			int cell = d2 * maxAbilities + d1;
			int tcell = d1 * maxAbilities + d2;
			ElatentCov[tcell] = ElatentCov[cell] =
				ElatentCov[cell] / data->rows - ElatentMean[d1] * ElatentMean[d2];
		}
	}
	for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
		int cell = sdim * maxAbilities + sdim;
		ElatentCov[cell] = ElatentCov[cell] / data->rows - ElatentMean[sdim] * ElatentMean[sdim];
	}

	if (state->cacheLXK) state->LXKcached = TRUE;

	Free(latentDist);

	//mxLog("E-step");
	//pda(ElatentMean.data(), 1, state->maxAbilities);
	//pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);
}

static int getLatentVersion(BA81Expect *state)
{
	return omxGetMatrixVersion(state->latentMeanOut) + omxGetMatrixVersion(state->latentCovOut);
}

// Attempt G-H grid? http://dbarajassolano.wordpress.com/2012/01/26/on-sparse-grid-quadratures/
static void ba81SetupQuadrature(omxExpectation* oo, int gridsize)
{
	BA81Expect *state = (BA81Expect *) oo->argStruct;
	if (state->verbose) {
		mxLog("%s: quadrature(%d)", oo->name, getLatentVersion(state));
	}
	int numUnique = state->numUnique;
	int numThreads = Global->numThreads;
	int maxDims = state->maxDims;
	int Qwidth = state->Qwidth;
	int numSpecific = state->numSpecific;
	int priDims = maxDims - (numSpecific? 1 : 0);

	// try starting small and increasing to the cap TODO
	state->quadGridSize = gridsize;

	state->totalQuadPoints = 1;
	for (int dx=0; dx < maxDims; dx++) {
		state->totalQuadPoints *= state->quadGridSize;
	}

	state->totalPrimaryPoints = state->totalQuadPoints;

	if (numSpecific) {
		state->totalPrimaryPoints /= state->quadGridSize;
		state->speQarea.resize(gridsize * numSpecific);
	}

	state->Qpoint.resize(gridsize);
	state->priQarea.resize(state->totalPrimaryPoints);

	double qgs = state->quadGridSize-1;
	for (int px=0; px < state->quadGridSize; px ++) {
		state->Qpoint[px] = Qwidth - px * 2 * Qwidth / qgs;
	}

	//pda(state->latentMeanOut->data, 1, state->maxAbilities);
	//pda(state->latentCovOut->data, state->maxAbilities, state->maxAbilities);

	double totalArea = 0;
	for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
		int quad[priDims];
		decodeLocation(qx, priDims, state->quadGridSize, quad);
		double where[priDims];
		pointToWhere(state, quad, where, priDims);
		state->priQarea[qx] = exp(dmvnorm(priDims, where,
						  state->latentMeanOut->data,
						  state->latentCovOut->data));
		totalArea += state->priQarea[qx];
	}
	for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
		state->priQarea[qx] /= totalArea;
		//mxLog("%.5g,", state->priQarea[qx]);
	}

	for (int sx=0; sx < numSpecific; sx++) {
		totalArea = 0;
		int covCell = (priDims + sx) * state->maxAbilities + priDims + sx;
		double mean = state->latentMeanOut->data[priDims + sx];
		double var = state->latentCovOut->data[covCell];
		//mxLog("setup[%d] %.2f %.2f", sx, mean, var);
		for (int qx=0; qx < state->quadGridSize; qx++) {
			double den = dnorm(state->Qpoint[qx], mean, sqrt(var), FALSE);
			state->speQarea[sIndex(state, sx, qx)] = den;
			totalArea += den;
		}
		for (int qx=0; qx < state->quadGridSize; qx++) {
			state->speQarea[sIndex(state, sx, qx)] /= totalArea;
		}
		//pda(state->speQarea.data() + sIndex(state, sx, 0), 1, state->quadGridSize);
	}

	if (!state->cacheLXK) {
		state->lxk = Realloc(state->lxk, numUnique * numThreads, double);
	} else {
		int ns = state->numSpecific;
		if (ns == 0) ns = 1;
		long numOrdinate = ns * state->totalQuadPoints;
		state->lxk = Realloc(state->lxk, numUnique * numOrdinate, double);
	}

	state->expected = Realloc(state->expected, state->totalOutcomes * state->totalQuadPoints, double);
}

static void ba81buildLXKcache(omxExpectation *oo)
{
	BA81Expect *state = (BA81Expect *) oo->argStruct;
	if (!state->cacheLXK || state->LXKcached) return;
	
	ba81Estep1(oo);
}

OMXINLINE static void
expectedUpdate(omxData *data, const int *rowMap, const int px, const int item,
	       const double observed, const int outcomes, double *out)
{
	int pick = omxIntDataElementUnsafe(data, rowMap[px], item);
	if (pick == NA_INTEGER) {
		double slice = observed / outcomes;
		for (int ox=0; ox < outcomes; ox++) {
			out[ox] += slice;
		}
	} else {
		out[pick-1] += observed;
	}
}

OMXINLINE static void
ba81Expected(omxExpectation* oo)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (state->verbose) mxLog("%s: EM.expected", oo->name);

	omxData *data = state->data;
	int numSpecific = state->numSpecific;
	const int *rowMap = state->rowMap;
	double *patternLik = state->patternLik;
	int *numIdentical = state->numIdentical;
	int numUnique = state->numUnique;
	int maxDims = state->maxDims;
	int numItems = state->EitemParam->cols;
	int totalOutcomes = state->totalOutcomes;

	OMXZERO(state->expected, totalOutcomes * state->totalQuadPoints);

	if (numSpecific == 0) {
#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalQuadPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);
			double *lxk = ba81LikelihoodFast(oo, thrId, 0, quad);
			for (int px=0; px < numUnique; px++) {
				double *out = state->expected + qx * totalOutcomes;
				double observed = numIdentical[px] * lxk[px] / patternLik[px];
				for (int ix=0; ix < numItems; ix++) {
					const double *spec = state->itemSpec[ix];
					int outcomes = spec[RPF_ISpecOutcomes];
					expectedUpdate(data, rowMap, px, ix, observed, outcomes, out);
					out += outcomes;
				}
			}
		}
	} else {
		int sDim = state->maxDims-1;
		long specificPoints = state->quadGridSize;

#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);

			cai2010(oo, thrId, FALSE, quad);
			double *allElxk = eBase(state, thrId);
			double *Eslxk = esBase(state, thrId);

			for (long sx=0; sx < specificPoints; sx++) {
				quad[sDim] = sx;
				long qloc = encodeLocation(state->maxDims, state->quadGridSize, quad);

				for (int sgroup=0; sgroup < numSpecific; sgroup++) {
					double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, quad);

					for (int px=0; px < numUnique; px++) {
						double *out = state->expected + totalOutcomes * qloc;

						for (int ix=0; ix < numItems; ix++) {
							const double *spec = state->itemSpec[ix];
							int outcomes = spec[RPF_ISpecOutcomes];
							if (state->Sgroup[ix] == sgroup) {
								double Ei = allElxk[px];
								double Eis = Eslxk[sgroup * numUnique + px];
								double observed = (numIdentical[px] * (Ei / Eis) *
										   (lxk[px] / patternLik[px]));
								expectedUpdate(data, rowMap, px, ix, observed, outcomes, out);
							}
							out += outcomes;
						}
					}
				}
			}
		}
	}
	//pda(state->expected, state->totalOutcomes, state->totalQuadPoints);
}

OMXINLINE static void
accumulateScores(BA81Expect *state, int px, int sgroup, double piece, const double *where,
		 int primaryDims, int covEntries, std::vector<double> *mean, std::vector<double> *cov)
{<commit_msg>A few more minor optimizations<commit_after>/*
  Copyright 2012-2013 <NAME> and contributors

  This is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <Rmath.h>
#include "omxExpectationBA81.h"
#include "npsolWrap.h"
#include "libifa-rpf.h"
#include "dmvnorm.h"

const struct rpf *rpf_model = NULL;
int rpf_numModels;
static const double MIN_PATTERNLIK = 1e-100;

void pda(const double *ar, int rows, int cols)
{
	std::string buf;
	for (int rx=0; rx < rows; rx++) {   // column major order
		for (int cx=0; cx < cols; cx++) {
			buf += string_snprintf("%.6g, ", ar[cx * rows + rx]);
		}
		buf += "\n";
	}
	mxLogBig(buf);
}

void pia(const int *ar, int rows, int cols)
{
	std::string buf;
	for (int rx=0; rx < rows; rx++) {   // column major order
		for (int cx=0; cx < cols; cx++) {
			buf += string_snprintf("%d, ", ar[cx * rows + rx]);
		}
		buf += "\n";
	}
	mxLogBig(buf);
}

// state->speQarea[sIndex(state, sx, qx)]
OMXINLINE static
int sIndex(BA81Expect *state, int sx, int qx)
{
	//if (sx < 0 || sx >= state->numSpecific) error("Out of domain");
	//if (qx < 0 || qx >= state->quadGridSize) error("Out of domain");
	return sx * state->quadGridSize + qx;
}

/**
 * \param theta Vector of ability parameters, one per ability
 * \returns A numItems by maxOutcomes colMajor vector of doubles. Caller must Free it.
 */
double *computeRPF(BA81Expect *state, omxMatrix *itemParam, const int *quad, const bool wantlog)
{
	omxMatrix *design = state->design;
	int maxDims = state->maxDims;
	int maxOutcomes = state->maxOutcomes;
	size_t numItems = state->itemSpec.size();

	double theta[maxDims];
	pointToWhere(state, quad, theta, maxDims);

	double *outcomeProb = Realloc(NULL, numItems * maxOutcomes, double);
	//double *outcomeProb = Calloc(numItems * maxOutcomes, double);

	for (size_t ix=0; ix < numItems; ix++) {
		const double *spec = state->itemSpec[ix];
		double *iparam = omxMatrixColumn(itemParam, ix);
		double *out = outcomeProb + ix * maxOutcomes;
		int id = spec[RPF_ISpecID];
		int dims = spec[RPF_ISpecDims];
		double ptheta[dims];

		for (int dx=0; dx < dims; dx++) {
			int ability = (int)omxMatrixElement(design, dx, ix) - 1;
			if (ability >= maxDims) ability = maxDims-1;
			ptheta[dx] = theta[ability];
		}

		if (wantlog) {
			(*rpf_model[id].logprob)(spec, iparam, ptheta, out);
		} else {
			(*rpf_model[id].prob)(spec, iparam, ptheta, out);
		}
#if 0
		for (int ox=0; ox < spec[RPF_ISpecOutcomes]; ox++) {
			if (!isfinite(out[ox]) || out[ox] > 0) {
				mxLog("item param");
				pda(iparam, itemParam->rows, 1);
				mxLog("where");
				pda(ptheta, dims, 1);
				error("RPF returned %20.20f", out[ox]);
			}
		}
#endif
	}

	return outcomeProb;
}

OMXINLINE static double *
getLXKcache(BA81Expect *state, const int *quad, const int specific)
{
	long ordinate;
	if (state->numSpecific == 0) {
		ordinate = encodeLocation(state->maxDims, state->quadGridSize, quad);
	} else {
		ordinate = (specific * state->totalQuadPoints +
			    encodeLocation(state->maxDims, state->quadGridSize, quad));
	}
	return state->lxk + state->numUnique * ordinate;
}

OMXINLINE static double *
ba81Likelihood(omxExpectation *oo, const int thrId, int specific, const int *quad, const double *outcomeProb)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	int numUnique = state->numUnique;
	int maxOutcomes = state->maxOutcomes;
	omxData *data = state->data;
	size_t numItems = state->itemSpec.size();
	int *Sgroup = state->Sgroup;
	double *lxk;

	if (!state->cacheLXK) {
		lxk = state->lxk + numUnique * thrId;
	} else {
		lxk = getLXKcache(state, quad, specific);
	}

	const int *rowMap = state->rowMap;
	for (int px=0; px < numUnique; px++) {
		double lxk1 = 1;
		for (size_t ix=0; ix < numItems; ix++) {
			if (specific != Sgroup[ix]) continue;
			int pick = omxIntDataElementUnsafe(data, rowMap[px], ix);
			if (pick == NA_INTEGER) continue;
			double piece = outcomeProb[ix * maxOutcomes + pick-1];  // move -1 elsewhere TODO
			lxk1 *= piece;
			//mxLog("%d pick %d piece %.3f", ix, pick, piece);
		}
#if 0
#pragma omp critical(ba81LikelihoodDebug1)
		if (!isfinite(lxk1) || lxk1 > numItems) {
			mxLog("where");
			double where[state->maxDims];
			pointToWhere(state, quad, where, state->maxDims);
			pda(where, state->maxDims, 1);
			mxLog("prob");
			pda(outcomeProb, numItems, maxOutcomes);
			error("Likelihood of row %d is %f", rowMap[px], lxk1);
		}
#endif
		lxk[px] = lxk1;
	}

	//mxLog("L.is(%d) at (%d %d) %.2f", specific, quad[0], quad[1], lxk[0]);
	return lxk;
}

double *ba81LikelihoodFast(omxExpectation *oo, const int thrId, int specific, const int *quad)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (!state->cacheLXK) {
		const double *outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
		double *ret = ba81Likelihood(oo, thrId, specific, quad, outcomeProb);
		Free(outcomeProb);
		return ret;
	} else {
		return getLXKcache(state, quad, specific);
	}

}

OMXINLINE static void
mapLatentSpace(BA81Expect *state, int sgroup, double piece, const double *where,
	       const std::vector<double> &whereGram, double *latentDist)
{
	int maxDims = state->maxDims;
	int maxAbilities = state->maxAbilities;
	int pmax = maxDims;
	if (state->numSpecific) pmax -= 1;

	if (sgroup == 0) {
		int gx = 0;
		int cx = maxAbilities;
		for (int d1=0; d1 < pmax; d1++) {
			double piece_w1 = piece * where[d1];
			latentDist[d1] += piece_w1;
			for (int d2=0; d2 <= d1; d2++) {
				double piece_cov = piece * whereGram[gx];
				latentDist[cx] += piece_cov;
				++cx; ++gx;
			}
		}
	}

	if (state->numSpecific) {
		int sdim = pmax + sgroup;
		double piece_w1 = piece * where[pmax];
		latentDist[sdim] += piece_w1;

		double piece_var = piece * whereGram[triangleLoc0(pmax)];
		int to = maxAbilities + triangleLoc0(sdim);
		latentDist[to] += piece_var;
	}
}

// Eslxk, allElxk (Ei, Eis) depend on the ordinate of the primary dimensions
void cai2010(omxExpectation* oo, const int thrId, int recompute, const int *primaryQuad)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	int numUnique = state->numUnique;
	int numSpecific = state->numSpecific;
	int maxDims = state->maxDims;
	int sDim = maxDims-1;
	int quadGridSize = state->quadGridSize;
	int quad[maxDims];
	memcpy(quad, primaryQuad, sizeof(int)*sDim);
	double *allElxk = eBase(state, thrId);
	double *Eslxk = esBase(state, thrId);

	for (int px=0; px < numUnique; px++) {
		allElxk[px] = 1;
		for (int sx=0; sx < numSpecific; sx++) {
			Eslxk[sx * numUnique + px] = 0;
		}
	}

	if (!state->cacheLXK) recompute = TRUE;

	for (int qx=0; qx < quadGridSize; qx++) {
		quad[sDim] = qx;
		double *outcomeProb = NULL;
		if (recompute) {
			outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
		}
		for (int sx=0; sx < numSpecific; sx++) {
			double *myEslxk = Eslxk + sx * numUnique;
			double *lxk;     // a.k.a. "L_is"
			if (recompute) {
				lxk = ba81Likelihood(oo, thrId, sx, quad, outcomeProb);
			} else {
				lxk = getLXKcache(state, quad, sx);
			}

			for (int ix=0; ix < numUnique; ix++) {
				double area = state->speQarea[sIndex(state, sx, qx)];
				double piece = lxk[ix] * area;
				//mxLog("E.is(%d) at (%d, %d) %.2f + %.2f = %.2f",
				//  sx, primaryQuad[0], qx, lxk[ix], area, piece);
				myEslxk[ix] += piece;
			}
		}
		Free(outcomeProb);
	}

	for (int sx=0; sx < numSpecific; sx++) {
		for (int px=0; px < numUnique; px++) {
			//mxLog("E.is(%d) at (%d) %.2f", sx, primaryQuad[0], state->Eslxk[esIndex(state, sx, 0)]);
			allElxk[px] *= Eslxk[sx * numUnique + px];  // allSlxk a.k.a. "E_i"
		}
	}
}

static void ba81Estep1(omxExpectation *oo)
{
	if(OMX_DEBUG) {mxLog("Beginning %s Computation.", oo->name);}

	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (state->verbose) {
		mxLog("%s: lxk(%d) patternLik ElatentMean ElatentCov",
		      oo->name, omxGetMatrixVersion(state->EitemParam));
	}

	int numUnique = state->numUnique;
	int numSpecific = state->numSpecific;
	int maxDims = state->maxDims;
	int maxAbilities = state->maxAbilities;
	int primaryDims = maxDims;

	state->patternLik = Realloc(state->patternLik, numUnique, double);
	double *patternLik = state->patternLik;
	OMXZERO(patternLik, numUnique);

	int numLatents = maxAbilities + triangleLoc1(maxAbilities);
	int numLatentsPerThread = numUnique * numLatents;
	double *latentDist = Calloc(numUnique * numLatents * Global->numThreads, double);

	// E-step, marginalize person ability
	//
	// Note: In the notation of Bock & Aitkin (1981) and
	// Cai~(2010), these loops are reversed.  That is, the inner
	// loop is over quadrature points and the outer loop is over
	// all response patterns.
	//
	if (numSpecific == 0) {
#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalQuadPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);
			double where[maxDims];
			pointToWhere(state, quad, where, maxDims);
			std::vector<double> whereGram(triangleLoc1(maxDims));
			gramProduct(where, maxDims, whereGram.data());

			const double *outcomeProb = computeRPF(state, state->EitemParam, quad, FALSE);
			double *lxk = ba81Likelihood(oo, thrId, 0, quad, outcomeProb);
			Free(outcomeProb);

			double area = state->priQarea[qx];
			for (int px=0; px < numUnique; px++) {
				double tmp = lxk[px] * area;
#if 0
				if (!isfinite(tmp)) {
					mxLog("where");
					pda(where, maxDims, 1);
					error("Row %d lxk %f logArea %f tmp %f",
					      state->rowMap[px], lxk[px], logArea, tmp);
				}
#endif
#pragma omp atomic
				patternLik[px] += tmp;
				mapLatentSpace(state, 0, tmp, where, whereGram,
					       thrLatentDist + px * numLatents);
			}
		}
	} else {
		primaryDims -= 1;
		int sDim = primaryDims;
		long specificPoints = state->quadGridSize;

#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			double *thrLatentDist = latentDist + thrId * numLatentsPerThread;
			int quad[maxDims];
			decodeLocation(qx, primaryDims, state->quadGridSize, quad);

			cai2010(oo, thrId, TRUE, quad);
			double *allElxk = eBase(state, thrId);
			double *Eslxk = esBase(state, thrId);

			for (long sx=0; sx < specificPoints; sx++) {
				quad[sDim] = sx;
				double where[maxDims];
				pointToWhere(state, quad, where, maxDims);
				std::vector<double> whereGram(triangleLoc1(maxDims));
				gramProduct(where, maxDims, whereGram.data());

				for (int sgroup=0; sgroup < numSpecific; sgroup++) {
					double area = areaProduct(state, quad, sgroup);
					double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, quad);
					for (int px=0; px < numUnique; px++) {
						double Ei = allElxk[px];
						double Eis = Eslxk[sgroup * numUnique + px];
						double tmp = ((Ei / Eis) * lxk[px] * area);
						mapLatentSpace(state, sgroup, tmp, where, whereGram,
							       thrLatentDist + px * numLatents);
					}
				}
			}

			double priArea = state->priQarea[qx];
			for (int px=0; px < numUnique; px++) {
				double Ei = allElxk[px];
				double tmp = (Ei * priArea);
#pragma omp atomic
				patternLik[px] += tmp;
			}
		}
	}

	int *numIdentical = state->numIdentical;

	//mxLog("raw latent");
	//pda(latentDist, numLatents, numUnique);

#pragma omp parallel for num_threads(Global->numThreads) schedule(dynamic)
	for (int lx=0; lx < maxAbilities + triangleLoc1(primaryDims); ++lx) {
		for (int tx=1; tx < Global->numThreads; ++tx) {
			double *thrLatentDist = latentDist + tx * numLatentsPerThread;
			for (int px=0; px < numUnique; px++) {
				int loc = px * numLatents + lx;
				latentDist[loc] += thrLatentDist[loc];
			}
		}
	}

#pragma omp parallel for num_threads(Global->numThreads)
	for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
		for (int tx=1; tx < Global->numThreads; ++tx) {
			double *thrLatentDist = latentDist + tx * numLatentsPerThread;
			for (int px=0; px < numUnique; px++) {
				int loc = px * numLatents + maxAbilities + triangleLoc0(sdim);
				latentDist[loc] += thrLatentDist[loc];
			}
		}
	}

#pragma omp parallel for num_threads(Global->numThreads)
	for (int px=0; px < numUnique; px++) {
		if (patternLik[px] < MIN_PATTERNLIK) {
			patternLik[px] = MIN_PATTERNLIK;
			warning("Likelihood of pattern %d is 0, forcing to %.3g",
				px, MIN_PATTERNLIK);
		}

		double *latentDist1 = latentDist + px * numLatents;
		double weight = numIdentical[px] / patternLik[px];
		int cx = maxAbilities;
		for (int d1=0; d1 < primaryDims; d1++) {
			latentDist1[d1] *= weight;
			for (int d2=0; d2 <= d1; d2++) {
				latentDist1[cx] *= weight;
				++cx;
			}
		}
		for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
			latentDist1[sdim] *= weight;
			int loc = maxAbilities + triangleLoc0(sdim);
			latentDist1[loc] *= weight;
		}
#if 0
		if (!isfinite(patternLik[px])) {
			error("Likelihood of row %d is %f", state->rowMap[px], patternLik[px]);
		}
#endif
	}

	//mxLog("raw latent after weighting");
	//pda(latentDist, numLatents, numUnique);

	std::vector<double> &ElatentMean = state->ElatentMean;
	std::vector<double> &ElatentCov = state->ElatentCov;
	
	ElatentMean.assign(ElatentMean.size(), 0.0);
	ElatentCov.assign(ElatentCov.size(), 0.0);

#pragma omp parallel for num_threads(Global->numThreads)
	for (int d1=0; d1 < maxAbilities; d1++) {
		for (int px=0; px < numUnique; px++) {
			double *latentDist1 = latentDist + px * numLatents;
			int cx = maxAbilities + triangleLoc1(d1);
			if (d1 < primaryDims) {
				ElatentMean[d1] += latentDist1[d1];
				for (int d2=0; d2 <= d1; d2++) {
					int cell = d2 * maxAbilities + d1;
					ElatentCov[cell] += latentDist1[cx];
					++cx;
				}
			} else {
				ElatentMean[d1] += latentDist1[d1];
				int cell = d1 * maxAbilities + d1;
				int loc = maxAbilities + triangleLoc0(d1);
				ElatentCov[cell] += latentDist1[loc];
			}
		}
	}

	//pda(ElatentMean.data(), 1, state->maxAbilities);
	//pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);

	omxData *data = state->data;
	for (int d1=0; d1 < maxAbilities; d1++) {
		ElatentMean[d1] /= data->rows;
	}

	for (int d1=0; d1 < primaryDims; d1++) {
		for (int d2=0; d2 <= d1; d2++) {
			int cell = d2 * maxAbilities + d1;
			int tcell = d1 * maxAbilities + d2;
			ElatentCov[tcell] = ElatentCov[cell] =
				ElatentCov[cell] / data->rows - ElatentMean[d1] * ElatentMean[d2];
		}
	}
	for (int sdim=primaryDims; sdim < maxAbilities; sdim++) {
		int cell = sdim * maxAbilities + sdim;
		ElatentCov[cell] = ElatentCov[cell] / data->rows - ElatentMean[sdim] * ElatentMean[sdim];
	}

	if (state->cacheLXK) state->LXKcached = TRUE;

	Free(latentDist);

	//mxLog("E-step");
	//pda(ElatentMean.data(), 1, state->maxAbilities);
	//pda(ElatentCov.data(), state->maxAbilities, state->maxAbilities);
}

static int getLatentVersion(BA81Expect *state)
{
	return omxGetMatrixVersion(state->latentMeanOut) + omxGetMatrixVersion(state->latentCovOut);
}

// Attempt G-H grid? http://dbarajassolano.wordpress.com/2012/01/26/on-sparse-grid-quadratures/
static void ba81SetupQuadrature(omxExpectation* oo, int gridsize)
{
	BA81Expect *state = (BA81Expect *) oo->argStruct;
	if (state->verbose) {
		mxLog("%s: quadrature(%d)", oo->name, getLatentVersion(state));
	}
	int numUnique = state->numUnique;
	int numThreads = Global->numThreads;
	int maxDims = state->maxDims;
	int Qwidth = state->Qwidth;
	int numSpecific = state->numSpecific;
	int priDims = maxDims - (numSpecific? 1 : 0);

	// try starting small and increasing to the cap TODO
	state->quadGridSize = gridsize;

	state->totalQuadPoints = 1;
	for (int dx=0; dx < maxDims; dx++) {
		state->totalQuadPoints *= state->quadGridSize;
	}

	state->totalPrimaryPoints = state->totalQuadPoints;

	if (numSpecific) {
		state->totalPrimaryPoints /= state->quadGridSize;
		state->speQarea.resize(gridsize * numSpecific);
	}

	state->Qpoint.resize(gridsize);
	state->priQarea.resize(state->totalPrimaryPoints);

	double qgs = state->quadGridSize-1;
	for (int px=0; px < state->quadGridSize; px ++) {
		state->Qpoint[px] = Qwidth - px * 2 * Qwidth / qgs;
	}

	//pda(state->latentMeanOut->data, 1, state->maxAbilities);
	//pda(state->latentCovOut->data, state->maxAbilities, state->maxAbilities);

	double totalArea = 0;
	for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
		int quad[priDims];
		decodeLocation(qx, priDims, state->quadGridSize, quad);
		double where[priDims];
		pointToWhere(state, quad, where, priDims);
		state->priQarea[qx] = exp(dmvnorm(priDims, where,
						  state->latentMeanOut->data,
						  state->latentCovOut->data));
		totalArea += state->priQarea[qx];
	}
	for (int qx=0; qx < state->totalPrimaryPoints; qx++) {
		state->priQarea[qx] /= totalArea;
		//mxLog("%.5g,", state->priQarea[qx]);
	}

	for (int sx=0; sx < numSpecific; sx++) {
		totalArea = 0;
		int covCell = (priDims + sx) * state->maxAbilities + priDims + sx;
		double mean = state->latentMeanOut->data[priDims + sx];
		double var = state->latentCovOut->data[covCell];
		//mxLog("setup[%d] %.2f %.2f", sx, mean, var);
		for (int qx=0; qx < state->quadGridSize; qx++) {
			double den = dnorm(state->Qpoint[qx], mean, sqrt(var), FALSE);
			state->speQarea[sIndex(state, sx, qx)] = den;
			totalArea += den;
		}
		for (int qx=0; qx < state->quadGridSize; qx++) {
			state->speQarea[sIndex(state, sx, qx)] /= totalArea;
		}
		//pda(state->speQarea.data() + sIndex(state, sx, 0), 1, state->quadGridSize);
	}

	if (!state->cacheLXK) {
		state->lxk = Realloc(state->lxk, numUnique * numThreads, double);
	} else {
		int ns = state->numSpecific;
		if (ns == 0) ns = 1;
		long numOrdinate = ns * state->totalQuadPoints;
		state->lxk = Realloc(state->lxk, numUnique * numOrdinate, double);
	}

	state->expected = Realloc(state->expected, state->totalOutcomes * state->totalQuadPoints, double);
}

static void ba81buildLXKcache(omxExpectation *oo)
{
	BA81Expect *state = (BA81Expect *) oo->argStruct;
	if (!state->cacheLXK || state->LXKcached) return;
	
	ba81Estep1(oo);
}

OMXINLINE static void
expectedUpdate(omxData *data, const int *rowMap, const int px, const int item,
	       const double observed, const int outcomes, double *out)
{
	int pick = omxIntDataElementUnsafe(data, rowMap[px], item);
	if (pick == NA_INTEGER) {
		double slice = observed / outcomes;
		for (int ox=0; ox < outcomes; ox++) {
			out[ox] += slice;
		}
	} else {
		out[pick-1] += observed;
	}
}

OMXINLINE static void
ba81Expected(omxExpectation* oo)
{
	BA81Expect *state = (BA81Expect*) oo->argStruct;
	if (state->verbose) mxLog("%s: EM.expected", oo->name);

	omxData *data = state->data;
	int numSpecific = state->numSpecific;
	const int *rowMap = state->rowMap;
	double *patternLik = state->patternLik;
	int *numIdentical = state->numIdentical;
	int numUnique = state->numUnique;
	int maxDims = state->maxDims;
	int numItems = state->EitemParam->cols;
	int totalOutcomes = state->totalOutcomes;

	std::vector<int> itemOutcomes(numItems);
	for (int ix=0; ix < numItems; ix++) {
		const double *spec = state->itemSpec[ix];
		itemOutcomes[ix] = spec[RPF_ISpecOutcomes];
	}

	OMXZERO(state->expected, totalOutcomes * state->totalQuadPoints);

	if (numSpecific == 0) {
#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalQuadPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);
			double *lxk = ba81LikelihoodFast(oo, thrId, 0, quad);
			for (int px=0; px < numUnique; px++) {
				double *out = state->expected + qx * totalOutcomes;
				double observed = numIdentical[px] * lxk[px] / patternLik[px];
				for (int ix=0; ix < numItems; ix++) {
					const int outcomes = itemOutcomes[ix];
					expectedUpdate(data, rowMap, px, ix, observed, outcomes, out);
					out += outcomes;
				}
			}
		}
	} else {
		int sDim = state->maxDims-1;
		long specificPoints = state->quadGridSize;

#pragma omp parallel for num_threads(Global->numThreads)
		for (long qx=0; qx < state->totalPrimaryPoints; qx++) {
			int thrId = omx_absolute_thread_num();
			int quad[maxDims];
			decodeLocation(qx, maxDims, state->quadGridSize, quad);

			cai2010(oo, thrId, FALSE, quad);
			double *allElxk = eBase(state, thrId);
			double *Eslxk = esBase(state, thrId);

			for (long sx=0; sx < specificPoints; sx++) {
				quad[sDim] = sx;
				long qloc = encodeLocation(state->maxDims, state->quadGridSize, quad);

				for (int sgroup=0; sgroup < numSpecific; sgroup++) {
					double *lxk = ba81LikelihoodFast(oo, thrId, sgroup, quad);
					double *myEslxk = Eslxk + sgroup * numUnique;

					for (int px=0; px < numUnique; px++) {
						double *out = state->expected + totalOutcomes * qloc;

						for (int ix=0; ix < numItems; ix++) {
							const int outcomes = itemOutcomes[ix];
							if (state->Sgroup[ix] == sgroup) {
								double Ei = allElxk[px];
								double Eis = myEslxk[px];
								double observed = (numIdentical[px] * (Ei / Eis) *
										   (lxk[px] / patternLik[px]));
								expectedUpdate(data, rowMap, px, ix, observed, outcomes, out);
							}
							out += outcomes;
						}
					}
				}
			}
		}
	}
	//pda(state->expected, state->totalOutcomes, state->totalQuadPoints);
}

OMXINLINE static void
accumulateScores(BA81Expect *state, int px, int sgroup, double piece, const double *where,
		 int primaryDims, int covEntries, std::vector<double> *mean, std::vector<double> *cov)
{
{'fraction_non_alphanumeric': 0.1521436239464608, 'fraction_numerical': 0.014732418172266509, 'mean_word_length': 3.8978632956706565, 'pattern_counts': {'":': 0, '<': 150, '<?xml version=': 0, '>': 422, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2236959', 'n_tokens_mistral': 18612, 'n_tokens_neox': 16654, 'n_words': 4081}