/*
************************************************************************
*   (C) Copyright 1999, Peter Lenk. All Rights Reserved.
*	HBREG2.GSS
*	HB Linear Regression Model.
*	Allows for subject level design matrices.
*	
*		Y_i = X_i*beta_i + epsilon_i for i = 1,..., nsub
*		epsilon_i is N(0,sigma^2 I)
*
*		beta_i = Theta'z_i + delta_i
*		delta_i is N(0,Lambda)
*		B		= Z*Theta + Delta
*
*	PRIORS
*		sigma^2 is Inverted Gamma(r0/2,s0/2)
*		Theta is maxtrix normal(u0,v0).
*		That is, vec(Theta) is N(vec(u0),v0).
*		vec(theta) stacks the columns of theta.
*		Lambda is Inverted Wishart(f0, g0) 
***********************************************************************
*/

new;
outfile		= "result.res";		@ Specify output file for saving results 					@
								@ outfile is a string variable that contains a file name 	@
inxy 		= "xydata";			@ Name of Gauss file with X,Y data							@
inz	 		= "zdata";			@ Name of Gauss file with Z data							@
flagtrue	= 1;				@ 1 -> knows true parameters from simulation				@
/*
********************************************************************
*	Initialize parameters for MCMC 
********************************************************************
*/
smcmc		= 100;		@ number of iterations to save for analysis 				@
skip		= 1;		@ Save every skip iterations								@
nblow		= 100;		@ Initial transition iterations 							@
nmcmc		= nblow + skip*smcmc;	@ total number of iterations					@

/*
********************************************************************
*	Get data
********************************************************************
*/
@ Input Gauss files @
open f1 	= ^inxy;	@ Get Gauss file for X, Y data 				@
				@ Opens Gauss file & assigns file handle f1 @
xydata		= readr(f1,rowsf(f1));		
			@ readr reads in Gauss file with file handle f1.			@
			@ rowsf(f1) returns the number of rows in the Gauss file. 	@
			@ readr reads rowsf(f1) rows, which is the entir dataset.	@
ci		= close(f1);
xynames		= setvars(inxy);			@ Get the variable names that accompnay X, Y data @
ynames		= xynames[rows(xynames)];	
xnames		= "CS"|xynames[1:rows(xynames)-1];

@ Last row of xydata is Y. @
xdata		= xydata[.,1:cols(xydata)-1];		@ Get independent variables 	@
ydata		= xydata[.,cols(xydata)];		@ Get dependent variable	@

open f1		= ^inz;
zdata		= readr(f1,rowsf(f1));	
ci		= close(f1);
znames		= setvars(inz);
znames		= "CS"|znames;

loadm	yrows	= yrows;	@ yrows gives the number of observations per subject @
nsub	= rows(yrows);		@ Number of subjects. @
ntot	= sumc(yrows);		@ Total number of observations. @
@ Create pointer based on yrows to access xdata and ydata @
b	= cumsumc(yrows);	@ cumsumc is the cumulative sum of yrows @
a	= 1|(1+b[1:nsub-1]);
iptxy	= a~b;
@ To use iptxy to get the design matrix and data for subject i: @
@ x_i = xdata[iptxy[i,1]:iptxy[i,2],.] @
@ y_i = ydata[iptxy[i,1]:iptxy[i,2]] @

xdim	= cols(xdata);
zdim	= cols(zdata);
@ add intercepts to xdata and zdata @
xdata 	= ones(ntot,1)~xdata;
zdata	= ones(nsub,1)~zdata;
rankx	= cols(xdata);
rankz	= cols(zdata);
thdim	= rankx*rankz;		@ dimension of vec(theta) @

@ Compute some sufficient statistics @
@ Get point to access stacked matrices of xtx, xtxi, and xty @
b	= seqa(rankx,rankx,nsub);
a	= 1|(1+b[1:nsub-1]);
iptxt	= a~b;
bpool	= invpd(xdata'xdata)*xdata'ydata;
xtx		= zeros(rankx*nsub,rankx);
xty		= zeros(rankx*nsub,1);
bhat	= zeros(nsub,rankx);		@ MLE of beta_i @
sse	= 0;
for i0 (1,nsub,1);	i = i0;
	xi	= xdata[iptxy[i,1]:iptxy[i,2],.];
	yi	= ydata[iptxy[i,1]:iptxy[i,2],.];
	xitxi	= xi'xi;
	xityi	= xi'yi;
	xtx[iptxt[i,1]:iptxt[i,2],.] = xitxi;
	xty[iptxt[i,1]:iptxt[i,2],.] = xityi;
	if rank(xitxi) >= rankx;		@ Got an inverse @
		xitxii	= invpd(xitxi);
		bhat[i,.] = (xitxii*xityi)';
	else;							@ Use the pooled estimate @
		bhat[i,.] = bpool';
	endif;
	resid	= yi - xi*(bhat[i,.]');
	sse	= sse + resid'resid;
endfor;
s2hat	= sse/ntot;		@ MLE of sigma2 @
sighat	= sqrt(s2hat);
ztz	= zdata'zdata;



/*
********************************************************************
*	Initialize Priors 
********************************************************************
*/

@ Prior for sigma2 is IG(r0/2, s0/2) @
r0 	= 2; s0 = 2;
rn 	= r0 + ntot;


@ Prior for theta is N(u0,v0) @

u0 	= zeros(thdim,1);
v0   	= 100*eye(thdim);		@ thdim = rankx*rankz    @
v0i  	= invpd(v0);    		@ used in updating theta @
v0iu0 	= v0i*vec(u0);          @ used in updating theta @

@ Lambda^{-1} is W_rankx(f0,g0 ) @
@ f0 = prior df, g0 = prior scale matrix @
f0 		= rankx+2;  f0n = f0 + nsub;
g0i 	= eye(rankx);   @ g0^{-1} @

/*
*******************************************************************
*	Initialize MCMC
******************************************************************
*/

beta 	= zeros(nsub, rankx);
sigma 	= 1;
sigma2	= s2hat;
theta	= zeros(rankz,rankx);
lambda	= eye(rankx);
lambdai	= invpd(lambda);

@ Define data structures for saving iterates & computing posterior means & std @
betam	= zeros(nsub,rankx);	@ posterior mean of beta 	@
betas	= zeros(nsub,rankx);	@ posterior std of beta 	@
sigmag	= zeros(smcmc,1);
thetag	= zeros(smcmc,thdim);
c	= rankx*(rankx+1)/2;
lambdag = zeros(smcmc,c);	@ save unique elements of lambda @


/*
********************************************************************
*	Do MCMC
********************************************************************
*/
@ Do the initial transition period @
for i1 (1,nblow,1);	imcmc = i1;
	call gethbreg;
endfor;

for i1 (1,smcmc,1);	imcmc = i1;		@ Save smcmc iterations 			@
	for i2 (1,skip,1); jmcmc = i2;		@ Save every skip iterations 	@
		call gethbreg;
	endfor;
	sigmag[imcmc]		= sigma;
	thetag[imcmc,.]		= vecr(theta)';
	betam				= betam + beta;
	betas				= betas + beta^2;
	lambdag[imcmc,.]	= vech(lambda)'; @ vech gets unique elements of symmetric matrix @
endfor;


/*
******************************************************************
*	Compute Posterior Means and STD
******************************************************************
*/
@ Compute Posterior Means @
betam		= betam/smcmc;
sigmam		= meanc(sigmag);
thetam		= meanc(thetag);
thetam		= reshape(thetam,rankz,rankx);
lambdam		= xpnd(meanc(lambdag));	@ xpnd is opposite of vech @

@ Compute Posterior STD @
betas		= sqrt( abs(betas - smcmc*betam^2)/smcmc);
sigmas		= stdc(sigmag);
thetas		= stdc(thetag);
thetas		= reshape(thetas, rankz, rankx);
lambdas		= xpnd(stdc(lambdag));

@ Predict yi @
yhhb	= zeros(ntot,1);
yhml	= yhhb;
for i0 (1, nsub, 1); i = i0;
	xi	= xdata[iptxy[i,1]:iptxy[i,2],.];
	yhhb[iptxy[i,1]:iptxy[i,2]] = xi*(betam[i,.]');
	yhml[iptxy[i,1]:iptxy[i,2]] = xi*(bhat[i,.]');
endfor;
cr		= corrx(ydata~yhhb);
hbr		= cr[1,2];
estd	= ydata - yhhb;
estd	= sqrt(estd'estd/ntot);
cr		= corrx(ydata~yhml);
mlr		= cr[1,2];
estdml 	=  ydata - yhml;
estdml 	= sqrt(estdml'estdml/ntot);


/*
****************************************************************
*	Do some output
****************************************************************
*/
call outputanal;


@ Plot saved iterations against iterations number @
t 	= seqa(nblow+skip,skip,smcmc);		@ saved iteration number @
title("Error STD versus Iteration");
xy(t,sigmag);
title("Theta versus Iteration");
xy(t,thetag);
title("Lambda versus Iteration");
xy(t,lambdag);
title("MLE & HB for " $+ xnames[2]$+ " versus " $+ xnames[1]);
_plctrl = -1;
xy(bhat[.,1]~betam[.,1],bhat[.,2]~betam[.,2]);
graphset;			@ Return to default settings @


end;

/*
****************************************************************
* GETHBREG
*	Does one iteration of the HB regression model.
*	INPUT
*		Global Variables
*	OUTPUT
*		Global Variables
****************************************************************
*/
PROC (0) = gethbreg;
local vibn, vibn12, ebin, yhat, sse, sn, resid, gni, gn, gn12, w, i0, i, limub, yhati,
betai;

	limub	= zdata*theta*lambdai;		@ Used in posterior mean of beta @
	/*
	***********************************************************
	* Generate beta
	* beta_i is N(mbin, vbn)
	* vbn 	= ( X_i'X_i/sigma2 + Lambda^{-1} }^{-1}
	* mbin	= vbn*( X_i'Y_i/sigma2 + Lambda^{-1}*Theta*Z_i)
	**********************************************************
	*/
	sse	= 0;
	for i0 (1,nsub,1); i = i0;
		xi 			= xdata[iptxy[i,1]:iptxy[i,2],.];
		yi 			= ydata[iptxy[i,1]:iptxy[i,2],.];
		xitxi 		= xtx[iptxt[i,1]:iptxt[i,2],.];
		xityi 		= xty[iptxt[i,1]:iptxt[i,2],.];
		vibn		= xitxi/sigma2 + lambdai;
		vibn12		= chol(vibn);
		ebin		= xityi/sigma2 + limub[i,.]';
		betai		= cholsol(ebin + vibn12'rndn(rankx,1), vibn12);
		beta[i,.]	= betai';
		yhati		= xi*betai;
		resid		= yi - yhati;
		sse			= sse + resid'resid;
	endfor;

	/*
	***************************************************************
	* Generate sigma2
	* sigma2 is IG(rn/2, sn/2)
	* rn = r0 + ntot
	* sn = s0 + sum_{i=1}^{nsub} (Y_i - X*beta_i)'(Y_i - X*beta_i)
	****************************************************************
	*/
	sn		= s0 + sse;
	sigma2	= sn/(2*rndgam(1,1,rn/2)); 
	sigma	= sqrt(sigma2);

	/*
	***********************************************************
	* Generate Theta and Lambda from multivariate model:
	*	B = Z*Theta + N(0,Lambda)
	************************************************************
	*/
	{theta, lambda, lambdai} = 
	getmulreg(beta,zdata,ztz,theta,lambda,lambdai,v0i,v0iu0,f0n,g0i);


endp;



/*
****************************************************************
* GETMULREG
*	Generate multivariate regression parameters.
*		Yd		= Xd*parmat + epsilon
*	
*	INPUT
*		yd		= dependent variables
*		xd		= independet variables
*		xdtxd		= xd'xd
*		
*		parmat	= current value of coefficient matrix
*		var		= current value of covariance matrix
*		vari		= its inverse
*		v0i		= prior precisions for bmat
*		v0iu0		= prior precision*prior mean for bmat
*		f0n		= posterior df for sigma
*		g0i		= prior scaling matrix inverse for sigma