####### Two-Level Normal independent sampling estimation  #######
#######       SUPPORTS MULTIVARIATE OUTCOMES
####### Copyright 2000, Phil Everson, Swarthmore College. ###########
tlnise_function(Y,V,w=NA,V0=NA,prior=NA,N=1000,seed=NA,
 Tol=1e-6,maxiter=1000,intercept=T,labelY=NA,labelYj=NA,
 labelw=NA,digits=4,brief=1,prnt=T){
#
# This program is free and may be redistributed as long as
# the copyright is preserved. 
#
# This program should be cited if used in published research.
#
# The author will appreciate notification of any comments or
# errors by emailing peverso1@swarthmore.edu
#
# Reference: 
# Everson and Morris (2000). J. R. Statist. Soc. B, 62 prt 2, pp.399-412.
#
# 			!!!IMPORTANT!!! 
# you must change the path in the following line to specify the directory in
# which you stored the Fortran object file tlnise.o.
 dyn.load("/home1/everson/TLNise/tlnisemv1.o")
#          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -> change to agree with your path
#
# TLNise: Two-Level Normal independent sampling estimation 
#
#  Provides inference for 2-level Normal hierarchical models 
#  with univariate (p = 1) or multivariate (p > 1) outcomes, and
#  with known level-1 variances V:
#
#  Level-1:  Yj|thetaj, Vj ~ Np(thetaj, Vj), j = 1,..,J
#
#  Level-2:  thetaj|A, gamma ~ Np(Wjgamma,A)
#
#  Calls Fortran subroutines in tlnise.f (object file tlnise.o).
#
# INPUTS:
#        Y: Jxp (or pxJ) matrix of p-dimensional Normal outcomes
#        V: pxpxJ array of pxp Level-1 covariances (assumed known)
#        w: Jxq (or qxJ) covariate matrix (adds column of 1's if 
#	     not included and intercept = T)
#       V0: "typical" Vj (default is average of Vj's)
#    prior: Prior parameter: p(B0) = |B0|^{(prior-p-1)/2}
#      prior = -(p+1): uniform on level-2 covariance matrix A (default) 
#      prior =      0: Jeffreys' prior
#      prior =  (p+1): uniform on shrinkage matrix B0
#        N: number of Constrained Wishart draws for inference
#      Tol: tolerance for determining modal convergence
#  maxiter: maximum number of EM iterations for finding mode
#intercept: if True, an intercept term is included in the regression
#   labelw: optional names vector for covariates
#   labelY: optional names vector for the J observations
#  labelYj: optional names vector for the p elements of Yj
#   digits: number of significant digits for reporting results
#    brief: level of output, from 0 (minimum) to 2 (maximum)
#     prnt: Controls printing during execution
#
# OUTPUTS:
#  *** brief=0 ***
#    gamma:  matrix of posterior mean and sd estimates of Gamma,
#             and their ratios.
#    theta:  pxJ matrix of posterior mean estimates for thetaj's.
#  SDtheta:  pxJ matrix of posterior SD estimates for thetaj's.
#        A:  pxp estimated posterior mean of variance matrix A.
#      rtA:  p-vector of between group SD estimates.
#
#  *** brief=1 *** 
#   brief=0 outputs, plus
#  Dgamma: rxr estimated posterior covariance matrix for Gamma.
#  Vtheta: pxpxJ array of estimated covariances for thetaj's.
#      B0: pxpxN array of simulated B0 values.
#      lr: N-vector of log density ratios for each B0 value.
#  
#  *** brief=2 *** 
#   brief=1 outputs, plus
#      lf: N-vector of log f(B0|Y) evaluated at each B0.
#     lf0: N-vector of log f0(B0|Y) evaluated at each B0
#           (f0 is the CWish envelope density for f).
#      df: degrees of freedom for f0.
#   Sigma: scale matrix for f0.
#    nvec: number of matrices begun, diagonal and off-diagonal
#           elements simulated to get N CWish matrices.
#    nrej: number of rejections that occurred at each step 1,..,p.
###########
#
# Check inputs:
out.chk_checkcon(Y,V,w,intercept,prior,prnt)
# sets Y: pxJ, w: qxJ, V: pxpxJ
Y_out.chk$Y
V_out.chk$V
w_out.chk$w
J_out.chk$J
p_out.chk$p
q_out.chk$q
r_p*q
prior_out.chk$prior
if(prnt)
 print(paste("******** Prior Parameter =",prior),quote=F)
if(missing(V0)) V0_apply(V,c(1,2),mean)
if(max(abs(V0-diag(p))) < Tol){
 Ys_Y
 Vs_V
 rtV0_diag(p)
 }
else{
# Rotate Y and V by rtV0^{-1}:
 newvars_standard.f(Y,V,V0)
 Ys_newvars$Y
 Vs_newvars$V
 rtV0_newvars$rtVo
# Ys = rtV0^-1%*%Y, Vsj = rtV0^(-1)%*%V%*%rtV0^(-1)
 }
#
# Locate A corresponding to the posterior mode of 
#  B0 = (I+A*)^{-1}  (A* = rtV0^-1%*%A%*%rtV0^-1).
if(prnt){
 cat("\n")
 print("******** Locating Posterior Mode ************ ",quote=F)
 cat("\n")
 }
Astart_diag(p)
out.mode_postmode.f(Ys,Vs,w,rtV0,prior,Astart,Tol, maxiter,J,p,q,r)
modeB0_solve(diag(p)+out.mode$newA)
lfmode_out.mode$lf
# compute A value on scale of data:
modeA_rtV0%*%out.mode$newA%*%rtV0
#
# Set the degrees of freedom (df) and scale parameter (Sigma) for
#  the CWish density f0 so that the mode of f0 is the same as  
#  the mode of f(B0|Y):
df_J-q+prior
Sigma_modeB0/(df-p-1)
# mode(f0) = (df-p-1)*Sigma = modeB0.
#
# Set d as the eigenvalues of Sigma^{-1}, ordered smallest to largest:
eigS_eigen(Sigma,symmetric=T)
d_1/eigS[[1]]
#
# Compute the inverse of Sigma, Siginv, and a square root matrix
#  for Sigma, rtSig ( rtSig%*%t(rtSig) = Sigma).
if(p==1){
 Siginv_matrix(1/Sigma)
 rtSig_matrix(sqrt(Sigma))
 }
else{
 Siginv_eigS[[2]]%*%diag(d)%*%t(eigS[[2]])
 rtSig_eigS[[2]]%*%diag(sqrt(1/d))
 }
# 
# Compute the log of the CWish(df,Sigma;d) density at its mode:
lf0mode_(df-p-1)*ldet(modeB0)/2 - tr(modeB0%*%Siginv)/2
#
# Set adj as the difference in the log-densities at the mode. This
# will be used to adjust the log-density ratio before exponentiating. 
adj_lf0mode-lfmode
iter_out.mode[[16]]
if(prnt){
 if(iter<maxiter){
  print(paste("Converged in",iter,"EM iterations."),quote=F)
  }
 else{
  print(paste("Did not converge in maxiter =",maxiter,"EM steps."),quote=F)
  }
 print("Posterior mode of B0:",quote=F)
 cat("\n")
 print(modeB0,digits)
 cat("\n")
 print(paste(
 "lf(modeB0) =",signif(lfmode,digits),"; lf0(modeB0) =",
  signif(lf0mode,digits),"; adj =", signif(adj,digits)),quote=F)
  }
if(missing(seed)) seed_ceiling(runif(1)*1e8)
cat("\n")
if(prnt) print("******** Drawing Constrained Wisharts ******** ",quote=F)
cat("\n")
# Generate N Constrained-Wishart(df,I;D) draws:
pd_pchisq(d,df)
if(min(pd)> 1-Tol){
 pd_rep(1,p)
 d_rep(0,p)
 }
outU_rscwish(N,p,df,d,pd,seed)
Uvals_outU$Ua
nvec_outU$nvec
nrej_outU$nrej
if(prnt){
 cat("\n")
 print(paste("CWish acceptance rate =",N,"/",nvec[1],"=",
 signif(N/nvec[1],digits)),quote=F)
 }
#
# Rotate Ui's to get B0i's:
B0vals_mammult(rtSig, Uvals, t(rtSig))
#
if(prnt){
 cat("\n")
 print("******** Processing Draws ************** ",quote=F)
 cat("\n")
 }
out.lf_lfB0.f(B0vals,Ys,Vs,w,rtV0,df,Siginv,N,J,p,q,r,prior,adj)
#
lr_out.lf$lrv
lf_out.lf$lfv
lf0_out.lf$lf0v
# changed 7/31/2000 for brief=2 output.
avewt_mean(exp(lr-max(lr)))
if(prnt){
 print(paste("Average scaled importance weight =",
  signif(avewt,digits)),quote=F)
 cat("\n")
 }
meanA_rtV0%*%(out.lf$meanA)%*%rtV0
if(p==1){ rtA_sqrt(meanA)}
else{ rtA_sqrt(diag(meanA))}
#
if(prnt){
 print("Posterior mean estimate of A:",quote=F)
 print(meanA, digits)
 cat("\n")
 print("Between-group SD estimate:", quote=F)
 print(rtA,digits)
 cat("\n")
 }
# Set theta:
if(missing(labelY)) labelY_1:J
if(missing(labelYj)) labelYj_1:p
theta_rtV0%*%out.lf$thetahat
Vtheta_mammult(rtV0,out.lf$Vthetahat,rtV0)
if(p==1){ 
 SDtheta_sqrt(c(Vtheta))
 theta_c(theta)
 names(SDtheta)_names(theta)_labelY
 }
else{
 SDtheta_0*Y
 for(j in 1:J) SDtheta[,j]_sqrt(diag(Vtheta[,,j]))
 dimnames(theta)_dimnames(SDtheta)_list(labelYj,labelY)
 dimnames(Vtheta)_list(labelYj,labelYj,labelY)
 }
#
gamma_out.lf$gamhat
Dgamma_out.lf$Dgamhat
GammaMat_cbind(gamma,sqrt(diag(Dgamma)))
GammaMat_cbind(GammaMat,GammaMat[,1]/GammaMat[,2])
if(missing(labelw)) labelw_seq(0,r-1,1)
dimnames(GammaMat)_list(labelw, c("est","se","est/se"))
names(gamma)_labelw
dimnames(Dgamma)_list(labelw,labelw)
# Output:
if(brief==0)
 out_list(gamma=GammaMat, theta=theta, SDtheta=SDtheta, A=meanA, rtA=rtA)
if(brief==1)
 out_list(gamma=GammaMat, theta=theta, SDtheta=SDtheta, A=meanA, rtA=rtA, 
          Dgamma=Dgamma,Vtheta=Vtheta, B0=B0vals, lr=lr)
if(brief==2)
 out_list(gamma=GammaMat, theta=theta, SDtheta=SDtheta, A=meanA, rtA=rtA, 
          Dgamma=Dgamma,Vtheta=Vtheta, B0=B0vals, lr=lr,lf=lf,
          lf0=lf0,df=df,Sigma=Sigma,nvec=nvec,nrej=nrej)
if(brief>2)
 out_list(out.mode, outU=outU, out.lf=out.lf)
out}

# source("/home/everson/TLNise/tlnise.mv1.src")
# out_tlnise(Y,V,w,seed=123456789)



postmode.f_function(Y,V,w,rtV0,prior,Astart,EPS,MAXIT,K,p,q,r){
# Calls Fortran subroutine postmode to locate the A value
# corresponding to the posterior mode of f(B0|Y).
# Y: pxJ, V: pxpxJ, w: qxJ
indxr_gamma_rep(0,r)
indxp_rep(0,p)
indxpk_0*Y
H_newA_0*Astart
Dgam_matrix(0,ncol=r,nrow=r)
lf_llik_0
Wi_matrix(0,ncol=p,nrow=r)
out_.Fortran("postmode",
 Astart=Astart,as.double(Y), as.double(w),as.double(V),
 as.double(rtV0),as.double(prior), # 6
 as.integer(p), as.integer(q),as.integer(r), as.integer(K),
 as.double(EPS), as.integer(MAXIT),    # 12
 newA=newA,
 lf=lf,
 llik=llik,
 as.integer(1), # 16 iter
 as.double(gamma),as.double(indxpk),
 as.double(Dgam),as.double(Dgam),
 as.double(0*V), 
 as.double(indxp),  # Yi
 as.double(Wi),   # Wi  23
 as.double(t(Wi)), 
 as.double(H),  # Di   
 as.double(H),  # Si
 H=H,  # H
 as.double(H),  # Hi
 as.integer(indxp), 
 as.integer(indxp), # indxp2 
 as.integer(indxpk),   # 30
 as.integer(indxr))
out
}


rscwish_function(N,p,df,d,pd,seed){
# Calls Fortran subroutine "rscwish" to generate N pxp
# standard constrained Wishart matrices with df degrees of
# freedom and a p-vector of constraints d. pd is the p-vector
# of the Chi-square(df) CDF evaluated at d. If all elements of 
# d are 0 and all elements of pd are 1, then rscwish returns N 
# unconstrained Wisharts.
#
 Ua_array(0.0,c(p,p,N))
 Tmat_matrix(0.0,p,p)
 Z_rep(0,p)
 chitab_c(qchisq(seq(.001,.999,.001), df))
# chitab is a table of 999 Chi-square(df) quantiles
#  to pass to Fortran. 
 out_.Fortran("rscwish",
        Ua=Ua,
        as.integer(N), 
        as.integer(p), 
        as.double(df),
	as.double(d), 
        as.double(pd), 
        as.double(chitab),
        as.integer(seed),
	as.double(Tmat),     # Tmat
        as.double(Tmat),     # Deltinv
        as.double(Tmat),     # H11
	as.double(Z),        # H12
        as.double(Z),        # Zt
        as.double(Z),        # Zts
        as.double(Z),        # tempp
	as.integer(0),       # mcnt
        as.integer(0),       # drvcnt
        as.integer(0),       # orvcnt
        as.integer(Z)        # rejcnt
        )
 list(Ua=out[[1]],nvec=c(out[[16]],out[[17]],out[[18]]),nrej=c(out[[19]]))
 }



lfB0.f_function(B0vals, Y, V, w, rtV0, df,Siginv,N,J,p,q,r, 
  prior=NA,adj=0){
# Returns the log posterior density of B0 = (I+A)^{-1},
# assuming prior distribution:
#   p(B0)dB0 propto |B0|^{(prior-p-1)/2} dB0, 0 < B0 < I.
# Assumes data have been rotated by rtV0^-1.
if(missing(prior))
# Uniform prior on B0 -> reports Likelihood of B0.
 prior_p+1
lfv_lf0v_lrv_rep(0,N)
gamma_rep(0,r)
Dgam_matrix(0,ncol=r,nrow=r)
theta_0*Y
Vtheta_0*V
Yj_rep(0,p)
Wj_matrix(0,ncol=r,nrow=p)
A_matrix(0,ncol=p,nrow=p)
resid_0*Y
sumwt_0
out_.Fortran("lfb0", B0vals=B0vals,as.double(Y),as.double(V), 
 as.double(w), as.double(rtV0), 
 as.integer(N),as.integer(p),as.integer(q),as.integer(r),   # 9
 as.integer(J),as.double(df),as.double(Siginv), #12
 as.double(prior),as.double(adj),lfv=lfv,lf0v=lf0v,lrv=lrv,  #17
 gamhat=gamma,Dgamhat=Dgam, thetahat=theta, Vthetahat=Vtheta, 
 meanA=A,sumwt=sumwt,
 as.double(gamma),as.double(Dgam), as.double(Dgam), as.double(Yj),  #27
 as.double(Vtheta),as.double(resid),as.double(A),as.double(A),  
 as.double(A),as.double(Dgam), as.double(Vtheta),as.double(Yj), #35
 as.double(A),as.double(Wj), as.double(t(Wj)),as.double(A),
 as.integer(Yj),as.integer(Yj),as.integer(resid),as.integer(gamma)) #43
out}


checkcon_function(Y,V,w,intercept,prior,prnt){
if(length(Y)==length(V)){
# univariate problem:
 J_length(Y)
 Y_matrix(Y,ncol=J,nrow=1)
 V_array(V,c(1,1,J))
 }
dmV_dim(V)
p_dmV[1]
if(p!=dmV[2]) stop("error in dimension of V")
J_dmV[3]
dmY_dim(Y)
if(dmY[1]==J){
 if(dmY[2]!=p) stop("error in dimension of Y or V")
# Set Y so it is pxJ:
 Y_t(Y)
 }
else{
 if(dmY[1]!=p|dmY[2]!=J) stop("error in dimension of Y or V")
 }
if(w[1]=="NA"){
 w_matrix(rep(1,J),nrow=1)
 q_1
 }
else{
 if(length(w)==J)
  w_matrix(w,nrow=1)
 dmw_dim(w)
 if(dmw[1]==J){
  q_dmw[2]
  w_t(w)
  }
 else{
  if(dmw[2]!=J) stop("error in dimension of w or V")
  q_dmw[1]
  }
 }
if(intercept&&max(abs(w[1,]-1))!=0){
#  add a row of 1's for an intercept:
  w_rbind(rep(1,J),w)
  q_q+1
  }
# Returns w: qxJ
if(J-q < 1) 
 stop("too many covariates for the number of observations")
if(prior!="NA"){
 if(J-q+prior < 1){
  if(prnt){
   print(paste("insufficient data for prior parameter =",prior,"."),quote=F)
   print("need J-q-p-1 + prior > 0", quote=F)
   }
  prior_max(prior,-(p+1))
  }
 }
else{
 prior_-(p+1)
 }
# Assuming a Uniform prior on A; p(A)dA propto dA.
# For B0 = (I+A*)^{-1}, this implies p(B0)dB0 \propto |B0|^{-(p+1)}dB0.
if(J-q+prior-p-1 <= 0)
 prior_0
if(J-q+prior-p-1 <= 0)
 prior_p+1
#
list(Y=Y,V=V,w=w,J=J,p=p,q=q,prior=prior)
}


standard.f_function(Y,V,Vo){
# Y is pxJ, V is pxpxJ and Vo is pxp. Generates rtVo, the
# symmetric square root matrix for Vo. Returns rtVo^{-1}%*%Y, 
#  rtVo^{-1}%*%V%*%rtVo^{-1}, and rtVo.
 p_dim(Vo)[1]
 J_dim(V)[3]
 Yj_rep(0,p)
 Vj_rtVo_0*Vo
 indxp_rep(0,p)
 out_.Fortran("standardize", Y=Y, V=V, Vo=Vo,as.integer(p),as.integer(J), 
        as.double(Yj), Vj=Vj,rtVo=rtVo,as.integer(indxp))
 list(Y=out$Y, V=out$V, rtVo=out$rtVo)}


mammult_function(M,A,tM){
# Calls Fortran subroutine "mammult" to pre-multiply by M
# (pxp) and post-multiply by tM (pxp) each of the N pxp  
# elements of A (pxpxN).
#
 pp_dim(A)
 p_pp[1]
 N_pp[3]
 MAM_A*0.0
 out_.Fortran("mammult",
       as.double(M),
       as.double(tM),
       as.double(A),
       as.integer(p),
       as.integer(N),
       as.double(0*M), # U
       as.double(0*M), # V
       MAM=MAM
       )
 out$MAM
 }


tr_function(A){
# calculate trace(A):
 sum(diag(A))
 }

ldet_function(A){
# calculate log(det(A)):
	sum(log(abs(diag(qr(A)[[1]]))))
}