hierarchical within-subjects gaussian with sufficient stats and reduce_sum

hwg_sufficient.stan

functions{
	real partial_lpdf( 
		array[] real real_array_to_slice
		, int start
		, int end
		, real sigma
		, vector Zc
		, matrix iZc_chol_mat
		, matrix iZc
		, array[] matrix iZc_
		, array[] int iXc
		, matrix tXc
		, vector Y_mean
		, vector Y_var
		, int rXc_minus_1
		, vector sqrt_Y_n	
	){
		int n = end - start;
		real lp = 0 ;
		// multivariate normal for individuals' coefficients ----
		matrix[nXc,n] iZc_mat ;
		if(centered==1){
			lp += multi_normal_cholesky_lpdf( iZc[start:end] |	Zc , iZc_chol_mat ) ;
			iZc_mat = to_matrix(iZc[start:end]) ;
		}else{
			lp += std_normal_lupdf(to_vector(iZc_[1][,start:end])) ;
			iZc_mat = (
				rep_matrix(Zc,n)
				+ (
					iZc_chol_mat
					* (iZc_[1][,start:end])
				)
			);

		}
		return(
			lp
			+ (
				(
					chi_square_lupdf(
						Y_var[start:end] / pow(sigma,2)
						| 1
					)
					- 2*(n*log(sigma))
				)
				* (n-1)
			)
			+ normal_lupdf(
				Y_mean
				|
				columns_dot_product( iZc_mat[,(iXc[start:end]-iXc[start])] , tXc[,start:end] )
				, sigma / sqrt_Y_n[start:end]
			)
		);
	}

}

data{

	// nI: number of individuals
	int<lower=2> nI ;

	// nXc: number of condition-level predictors
	int<lower=1> nXc ;

	// rXc: number of rows in the condition-level predictor matrix
	int<lower=nXc> rXc ;

	// Xc: condition-level predictor matrix
	matrix[rXc,nXc] Xc ;

	// iXc: which individual is associated with each row in Xc
	array[rXc] int<lower=1,upper=nI> iXc ;

	// Y_mean
	vector[rXc] Y_mean ;

	// Y_sd
	vector[rXc] Y_sd ;

	// Y_n
	vector[rXc] Y_n ;

	// centered: whether to monolithically-center (1) or non-center (2) mid-hierarchy parameters
	int<lower=0,upper=1> centered ;

	// use_reduce_sum: whether to use reeduce_sum for parallel computation of the likelihood
	int<lower=0,upper=1> use_reduce_sum ;
}

transformed data{

	// tXc: transposed copy of Xc
	matrix[nXc,rXc] tXc = transpose(Xc) ;

	array[rXc] real real_array_to_slice = rep_array(0.0,rXc) ;

	vector[rXc] Y_var = pow(Y_sd,2) ;

	vector[rXc] sqrt_Y_n = sqrt(Y_n) ;

	int rXc_minus_1 = rXc - 1 ;

}

parameters{

	// sigma: magnitude of observation-level variability
	real<lower=0> sigma ;

	// Z: coefficients associated with predictors (group-level, condition-level, & interactions)
	vector[nXc] Zc ;

	// iZc_sd: magnitude of variability among individuals within a group
	vector<lower=0>[nXc] iZc_sd ;

	// iZc_cholcorr: cholesky-factor of correlation structure associated with variability among individuals on influence of within-individual predictors
	cholesky_factor_corr[nXc] iZc_cholcorr ;

	//next two parameters represent the same quantity, just with structures suitable for either
	//  centered or non-centered parameterization. The `centered` data variable will make one or the other zero-length.
	//  Hopefully Stan will soon allow matrix arguments to multi_normal_*(), which would obviate this hack.

	// iZc: by-individual coefficients (centered parameterization)
	array[nI*centered] vector[nXc] iZc ;

	// iZc_: helper-variable (note underscore suffix) for non-centered parameterization of iZc
	array[(1-centered)] matrix[nXc,nI] iZc_ ;

}

model{

	// priors ----

	// standard-normal priors on all group-level coefficients
	Zc ~ std_normal() ;

	// prior peaked around .8 for magnitude of observation-level variability
	sigma ~ weibull(2,1) ;

	// positive-standard-normal priors on magnitude of variability among individuals within a group
	iZc_sd ~ std_normal() ;

	// flat prior on correlations
	iZc_cholcorr ~ lkj_corr_cholesky(1) ;


	// Dot-products & likelihood ----

	// observations as normal with common error variability and varying mean
	if(use_reduce_sum==0){
		// multivariate normal for individuals' coefficients ----
		matrix[nXc,nI] iZc_mat ;
		if(centered==1){

			// multi-normal structure for iZc
			iZc ~ multi_normal_cholesky(
				Zc
				, diag_pre_multiply(iZc_sd, iZc_cholcorr)
			) ;

			//convert iZc from array of vectors to matrix
			iZc_mat = to_matrix(iZc) ;

		}else{
			// iZc_ must have a standard-normal prior for non-centered parameterization
			to_vector(iZc_[1]) ~ std_normal() ;

			// compute values for individuals given group associated group values and the individuals'
			//   deviations therefrom
			iZc_mat = (
				rep_matrix(Zc,nI)
				+ (
					diag_pre_multiply(
						iZc_sd
						, iZc_cholcorr
					)
					* (iZc_[1])
				)
			);

		}
		target += (
			(
				(
					chi_square_lupdf(
						Y_var / pow(sigma,2)
						| 1
					)
					- 2*(rXc*log(sigma))
				)
				* rXc_minus_1
			)
			+ normal_lupdf(
				Y_mean
				|
				columns_dot_product( iZc_mat[,iXc] , tXc )
				, sigma / sqrt_Y_n
			)
		) ;
	}else{
		matrix[nXc,nXc] iZc_chol_mat = diag_pre_multiply(iZc_sd, iZc_cholcorr) ;
		target += reduce_sum(
			partial_lupdf
			, real_array_to_slice
			, 1 //grain-size (1=auto)
			, sigma 
			, Zc
			, iZc_chol_mat
			, iZc
			, iZc_
			, iXc
			, tXc
			, Y_mean
			, Y_var
			, rXc_minus_1
			, sqrt_Y_n
		) ;
	}

}