#Splus code for hierarchical Bayes seedling growth data 
#  J.S. Clark 2002

plotfile_"acruplot"        # or "lituplot"
plantfile_"acruplant"      # or "lituplant"
b_read.table(plotfile,header=F)
  lightplot_b[,1]; varlight_b[,2]; dens_b[,3]
  #lightplot = obs light level at each plot
  #varlight  = variance in light at each plot
  #dens      = number of seedlings on each plot

b_read.table(plantfile,header=F)
  plot_b[,1]; lightplant_b[,2]; grow_b[,3]
  #lightplant = light estimate for the plot this plant occupies
  #grow       = growth rate of each plant (cm per yr)
  #plot       = plot number for each plant

#beta parameters describing light distributions (appendix A)
   a.x_lightplot*(lightplot*(1-lightplot)/varlight - 1)
   b.x_(1-lightplot)*(lightplot*(1-lightplot)/varlight - 1)

#frequentist fit - used to set initial parameter values and for comparison
 fitreg_function(param){
  ap_param[1]
  bp_param[2]
  th1_param[3]
  sig_param[4]
  fx_ap + bp*(lightplant/(th1 + lightplant))
  llik_-log(2*3.14159)/2 - log(sig) - .5/sig/sig * (grow - fx)^2
  return(-sum(llik))
 }
  plot(lightplant,grow)
  param0_c(-1,20,.2,1)
  out_nlminb(param0,fitreg,lower=c(-100,.1,.00001,.00001),upper=c(10,300,20,50))
    lseq_seq(0,.8,length=50)
    mua_out$parameters[1]
    mub_out$parameters[2]
    tt_out$parameters[3]
    pred_mua + mub*lseq/(tt + lseq)
    lines(lseq,pred)

###############################################
#Functions for Gibbs sampler

#sum of squares
rfunc_function(tt){
  rvar_sum((grow - a - b*lplant.t/(tt+lplant.t))^2)
  return(rvar)
}
#light probability-'lite' has one light value for each plot
truelight_function(lite){
   pp_rep(0,nplot)
   for(k in 1:nplot){
     gvec_grow[plot[] == k]
     aa_a[plot[] == k]
     bb_b[plot[] == k]
     fx_aa + bb*(lite[k]/(th + lite[k]))
     pp[k]_sum(log(dnorm(gvec,fx,sqrt(ssv))))
     pp[pp[] < -1000]_-1000
   }
  return(pp)  #return log probability
}
#########################################
#Gibbs parameters and initialization
nstep_20000                             #number of Gibbs steps
options(object.size=5e20)               #make space
    as_1; bs_1                          #prior for sigma
    sga_100; sgb_100                    #initial var on alpha, beta
    vv_c(sga,sgb,0); vinvert_solve(diag(vv[1:2]))
      v.out_matrix(NA,nrow=nstep,ncol=3); v.out[1,]_vv
    a0_.1; b0_.1                       #Inv Gamma prior parameters for variance
    svar_rep(0,nstep);svar[1]_1            #variance regression
    ath_100;  bth_ath*(1/tt - 1)     #prior theta-center on MLE
      theta_rep(0,nstep); theta[1]_tt
    w_round(nplant/20)               #Wishart prior
      r11_(20/8)^2; r22_(80/8)^2; rmat_diag(c(r11,r22)) #inv rmat is prior precision
      w1_matrix(0,nrow=2,ncol=2)
      ww_matrix(0,nrow=2,ncol=2);wpar_ww*0
    dmat_rbind(cbind(800,0),cbind(0,1000)); dinvert_solve(dmat) #parameter covariance
    hprior_c(mua,mub)                      #prior hyperparameters
    hypar_matrix(NA,nrow=nstep,ncol=2); hypar[1,]_hprior
    zmat_as.matrix(cbind(rep(1,nplant),cbind(rep(0,nplant)))) #design matrix
    beta0_matrix(0,nrow=nstep,ncol=nplant)  #individual regression parameters
      beta1_matrix(0,nrow=nstep,ncol=nplant)
    mvec_matrix(0,nrow=nplant,ncol=2)      #arrays needed for regression
      b.sd_matrix(0,nrow=nplant,ncol=2)
      mv_rep(0,2)
      vi.big_matrix(0,nrow=2,ncol=2)
      bpar_matrix(0,nrow=nplant,ncol=2)  
    plotlight_matrix(NA,nrow=nstep,ncol=nplot)  #imputed light values
      lightplot.t_lightplot; plotlight[1,]_lightplot.t

#Gibbs loop#####################################
for(j in 1:(nstep-1)){
  lightplot.t_plotlight[j,]  #current light value for each plot
  th_theta[j]                #current theta value
  ssv_svar[j]                #current residual variance
  lplant.t_lightplot.t[plot] #assign latent (true) light level to each plant

  #regression parameters- algebra to avoid matrix inversions
    zz_lplant.t/(th+lplant.t)
    zmat[,2]_zz
    v1 _1/ssv + vinvert[1,1]
    v12_zz/ssv + vinvert[1,2]
    v2 _ zz^2/ssv + vinvert[2,2]
    vdenom_1/(v1*v2 - v12^2)
    vi.b1 _v2*vdenom
    vi.b12_-v12*vdenom
    vi.b2 _v1*vdenom
    vi.s1 _grow/ssv + hypar[j,1]*vinvert[1,1] + hypar[j,2]*vinvert[1,2]
    vi.s2 _zz*grow/ssv + hypar[j,1]*vinvert[1,2] + hypar[j,2]*vinvert[2,2]
    b.sd[,1]_sqrt(vi.b1)
    b.sd[,2]_sqrt(vi.b2)
    bcor_vi.b12/b.sd[,1]/b.sd[,2]
    mvec[,1]_vi.b1*vi.s1 + vi.b12*vi.s2
    mvec[,2]_vi.b12*vi.s1 + vi.b2*vi.s2
    bpar[,1:2]_rmvnorm(nplant,mean=mvec,sd=b.sd,rho=bcor,d=2)
    a_bpar[,1]
    b_bpar[,2]
   #regression hyperparameters
    v1 _nplant*vinvert[1,1] + dinvert[1,1]
      v12_nplant*vinvert[1,2] + dinvert[1,2]
      v2 _nplant*vinvert[2,2] + dinvert[2,2]
      vdenom_1/(v1*v2 - v12^2)
    vi.b1 _v2*vdenom; vi.b12_-v12*vdenom; vi.b2 _v1*vdenom
    bs1_sum(bpar[,1]); bs2_sum(bpar[,2])
    vi.s1_vinvert[1,1]*bs1 + vinvert[1,2]*bs2 + dinvert[1,1]*hprior[1]
    vi.s2_vinvert[1,2]*bs1 + vinvert[2,2]*bs2 + dinvert[2,2]*hprior[2]
    mv[1]_vi.b1*vi.s1 + vi.b12*vi.s2
    mv[2]_vi.b12*vi.s1 + vi.b2*vi.s2
    hypar[j+1,]_rmvnorm(1,mean=mv,cov=rbind(c(vi.b1,vi.b12),c(vi.b12,vi.b2)))
   
   #Wishart covariance matrix
     w1[1,1]_sum((bpar[,1] - hypar[j+1,1])^2)
       w1[2,2]_sum((bpar[,2] - hypar[j+1,2])^2)
       w1[1,2]_sum((bpar[,1] - hypar[j+1,1]) * (bpar[,2] - hypar[j+1,2]))
       w1[2,1]_w1[1,2]
     zp_w1 + w*rmat; wdenom_zp[1,1]*zp[2,2] - zp[1,2]*zp[2,1]
     wpar[1,1]_zp[2,2]; wpar[1,2]_-zp[1,2]; wpar[2,1]_-zp[2,1];wpar[2,2]_zp[1,1]
     wpp_wpar/wdenom
     alpha_rmvnorm((nplant + w),mean=rep(0,2),cov=wpp)
     ww[1,1]_sum(alpha[,1]^2)
       ww[2,2]_sum(alpha[,2]^2)
       ww[1,2]_sum(alpha[,1] * alpha[,2])
       ww[2,1]_ww[1,2]
       vinvert_ww
     
  #theta: Met Hast step--beta jump density with mean at current theta value
    a.j_abs(th*(th*(1-th)/.00001 - 1))
    b.j_abs((1-th)*(th*(1-th)/.00001 - 1))
    st_rbeta(1,a.j,b.j)                      #proposal
      a.s_abs(st*(st*(1-st)/.00001 - 1))
      b.s_abs((1-st)*(st*(1-st)/.00001 - 1))
    predmean1_a + b*zz
      lb1_sum(log(dnorm(grow,predmean1,sqrt(ssv)))) + log(dbeta(th,ath,bth))
      lb2_log(dbeta(th,a.s,b.s))
    zzs_lplant.t/(st+lplant.t)
    predmean2_a + b*zzs        #predicted growth rate for proposal
      lt1_sum(log(dnorm(grow,predmean2,sqrt(ssv)))) + log(dbeta(st,ath,bth))
      lt2_log(dbeta(st,a.j,b.j))
    ltt_lt1 - lt2 - lb1 + lb2
    r_exp(min(ltt,20))
    z_runif(1,0,1)
    pr_min(r,1)
    theta[j+1]_theta[j]
    if(pr > z)theta[j+1]_st
    th_theta[j+1]

  #variance: Inv Gamma
    apost_a0 + nplant/2
    bpost_b0 + .5*rfunc(theta[j+1])
    svar[j+1]_1/rgamma(1,apost,bpost); ssv_svar[j+1]

   #true light value: metropolis, beta jump density
    pss_rbeta(nplot,a.x,b.x)      #proposal light for each plot
     l.temp_rep(pss,times=dens)    #for each plant
     a.j_abs(pss*(pss*(1-pss)/varlight - 1))
     b.j_abs((1-pss)*(pss*(1-pss)/varlight - 1))
    lt1_truelight(pss) + log(dbeta(pss,a.x,b.x))
    lt2_log(dbeta(pss,a.x,b.x))
    lb1_truelight(lightplot.t) + log(dbeta(lightplot.t,a.x,b.x))
    lb2_log(dbeta(lightplot.t,a.j,b.j))
     r_exp(lt1 - lt2 - lb1 + lb2)
     r[r[] > 1]_1
     z_runif(nplot,0,1)
     lightplot.t[r[] > z[]]_pss[r[] > z[]]
     plotlight[j+1,]_lightplot.t

   vtt_solve(vinvert)             #updates
   v.out[j+1,]_c(diag(vtt),vtt[2,1])
   beta0[j,]_bpar[,1]; beta1[j,]_bpar[,2]
   lplant.t_lightplot.t[plot]
   print(c(j,hypar[j,],th,ssv))
}

#########################################################################

burnin_2000             #set burnin
beta0_beta0[burnin:(nstep-1),]
beta1_beta1[burnin:(nstep-1),]
theta_theta[burnin:(nstep-1)]
dev_dev[burnin:(nstep-1)]
svar_svar[burnin:(nstep-1)]
v.out_v.out[burnin:(nstep-1),]
hypar_hypar[burnin:(nstep-1),]
plotlight_plotlight[burnin:(nstep-1),]

#plot Markov chains
  par(mfrow=c(2,2))
  plot(hypar[,1],type='l',ylim=range(hypar)); lines(hypar[,2])
  plot(theta,type='l')
  plot(svar,type='l')
  plot(v.out[,1],type='l',ylim=c(-10000,20000));lines(v.out[,2],lty=2)
     lines(v.out[,3],lty=3)

#extract positive values, calculate transformed parameters
  ndim_length(theta)
  trep_rep(theta,each=nplant)
  theta.mat_matrix(trep,nrow=ndim,byrow=T)
  b0.vec_as.vector(beta0); b1.vec_as.vector(beta1)
  th.vec_as.vector(theta.mat)
  b3_b0.vec + b1.vec
  b0.tmp_b0.vec[b0.vec < 0 & b1.vec > 0 & b3 > 0]
  b1.tmp_b1.vec[b0.vec < 0 & b1.vec > 0 & b3 > 0]
  th.tmp_th.vec[b0.vec < 0 & b1.vec > 0 & b3 > 0]
  mg_b0.tmp + b1.tmp; mx_-b0.tmp/mg * th.tmp
  xm_rep(0,nplant);gm_rep(0,nplant)
  gm2_beta0*0; xm2_beta0*0; gsd_xm*0;xsd_xm*0
  ttt_mean(theta)
  for(i in 1:nplant){
   b00_beta0[,i]; b11_beta1[,i]
   b0_b00[b00 < 0 & b11 > 0]
   b1_b11[b00 < 0 & b11 > 0]
   gm[i]_mean(b0 + b1)
   gsd[i]_sqrt(var(b0 + b1))
   xm[i]_mean(-b0*ttt/(b0 + b1))
   xsd[i]_sqrt(var(-b0*ttt/(b0 + b1)))
   b00[b00 > 0 | b11 < 0]_NA; b11[b00 > 0 | b11 < 0]_NA
   gm2[,i]_b00 + b11
   xm2[,i]_-b00 *ttt/(b00+b11)
 }

#light parameters: plot densities, compute quantiles
  ci.light_matrix(NA,nrow=nplot,ncol=3)    #95% ci for each plot
  plot(density(plotlight[,1],width=.05),type='l',xlim=c(0,1.4),ylim=c(0,30))
    for(k in 1:nplot){
      lines(density(plotlight[,k],width=.05))
      ci.light[k,]_quantile(plotlight[,k],c(.5,.025,.975))
    }

#output for plot parameters
  sdl_sqrt(apply(plotlight,2,var))
  outplot_signif(cbind(c(1:nplot),lightplot,ci.light[,2],sdl,ci.light[,2:3]),4)
  dimnames(outplot)[2]_list(c("plot","x - obs","xT - latent","se","0.025","0.975"))
  par(mfrow=c(4,1))

#regression parameters
  ci.x0_matrix(NA,nrow=nplant,ncol=3)
  #plot x0 distributions
    plot(density(xm2[,1],width=.05,na.rm=T),type='l',
      xlim=c(0,1.4),ylim=c(0,50))
    for(i in 1:nplant){
      xtmp_xm2[,i]
      xi_xtmp[xtmp < 1]
      lines(density(xi,width=.05,na.rm=T))
      ci.x0[i,]_quantile(xi,c(.5,.025,.975),na.rm=T)
    }
   xtmp_xm2[xm2 < .2 & xm2 > 0]
   plot(density(xtmp,width=.002,na.rm=T),type='l',xlim=c(0,1.4))
   ci.xall_quantile(xtmp,c(.5,.025,.975),na.rm=T)
   sdx0_stdev(xtmp,na.rm=T)
   abline(v=ci.xall,lty=2)

#asymptotes
  ci.g_matrix(NA,nrow=nplant,ncol=3)
    for(i in 1:nplant)ci.g[i,]_quantile(gm2[,i],c(.025,.5,.975),na.rm=T)
    ci.gall_quantile(gm2,c(.5,.025,.975),na.rm=T)
    sdg_stdev(gm2,na.rm=T)

#output for plant parameters
   outplant_cbind(c(1:nplant),grow,ci.x0[,1],xsd,ci.x0[,2:3],ci.g[,1],gsd,ci.g[,2:3])
   dimnames(outplant)[2]_list(c("plant","growth","x0","x0.se","x0.025","x0.975",
      "G","G.se","G.025","G.975"))
   all.p_c(0,mean(grow),ci.xall[1],sdx0,ci.xall[2:3],ci.gall[1],sdg,ci.gall[2:3])
   outnew_signif(rbind(all.p,outplant),3)

#theta
  ci.t_quantile(theta,c(.5,0.025,0.975))
  tdens_density(theta,width=.01)
  plot(tdens,type='l',xlim=c(0,1.4)); abline(v=ci.t,lty=2)
  theta.est_signif(c(mean(theta),stdev(theta),ci.t[2:3]),3)

#sigma
  ci.s_quantile(svar,c(.5,0.025,0.975))
  sdens_density(svar,width=.2)
  svar.est_signif(c(mean(svar),stdev(svar),ci.s[2:3]),3)

#hyperparameters
  ghypar_hypar[,1] + hypar[,2]
  ci.gh_quantile(ghypar,c(.5,0.025,0.975))
  ghdens_density(ghypar,width=5)
  xhypar_-hypar[,1]*theta/ghypar
  ci.xh_quantile(xhypar,c(.5,0.025,0.975))
  xhdens_density(xhypar,width=.004)
  plot(xhdens,type='l',xlim=c(0,1.4)); abline(v=ci.xh,lty=2)
  hyp.mean_apply(hypar,2,mean); hyp.var_var(hypar)
  b0h.est_signif(c(hyp.mean[1],sqrt(hyp.var[1,1]),quantile(hypar[,1],c(.025,.975))),3)
  b1h.est_signif(c(hyp.mean[2],sqrt(hyp.var[2,2]),quantile(hypar[,2],c(.025,.975))),3)
  gh.est_signif(c(mean(ghypar),stdev(ghypar),ci.gh[2:3]),3)
  x0h.est_signif(c(mean(xhypar),stdev(xhypar),ci.xh[2:3]),3)

#parameter covariances
  vt_v.out
  x11.est_signif(c(mean(vt[,1]),stdev(vt[,1]),quantile(vt[,1],c(.025,.975))),3)
  x22.est_signif(c(mean(vt[,2]),stdev(vt[,2]),quantile(vt[,2],c(.025,.975))),3)
  x12.est_signif(c(mean(vt[,3]),stdev(vt[,3]),quantile(vt[,3],c(.025,.975))),3)

#output for regression parameters
  outpar_rbind(b0h.est,b1h.est,gh.est,x0h.est,theta.est,
    svar.est,x11.est,x22.est,x12.est)
  dimnames(outpar)[2]_list(c("estimate","se","2.5%","97.5%"))

#individual transformed parameters
  par(mfrow=c(1,1))
  plot(lightplant,grow,xlim=c(0,1.4),ylim=c(-50,200))
    for(i in 1:nplant){
      gdens_density(gm2[,i],width=8,na.rm=T)
      lines((1-.5*gdens$y),gdens$x)
    }
     gdens_density(gm2,width=30,na.rm=T)
      lines((1.2-8*gdens$y),gdens$x)
      abline(h=c(ci.gall[2],ci.gall[3]),lty=2)

#hypar asymptote
  lines((1.3-ghdens$y),ghdens$x); abline(h=ci.gh[2:3],lty=2)
 
#sigma
  sss_seq(.7,1.3,length=100)
  sseq_ci.gh[1]*sss
  sigma_sqrt(svar)
  lines((1.4-dnorm(sseq,ci.gh[1],mean(sigma))),sseq)
  abline(h=qnorm(c(.025,.975),ci.gh[1],mean(sigma)),lty=2)

#confidence intervals for regression
  nrep_10; ntot_nrep*nplant; npar_length(theta)
  mdim_length(b0.tmp)
  par_sample((1:mdim),npar,replace=T)
  regmat_matrix(NA,nrow=npar,ncol=50)
  ki_0
  for(ki in 1:npar){
    pm_par[ki]
    regmat[ki,]_b0.tmp[pm] + b1.tmp[pm]*(lseq/(lseq + th.tmp[pm]))
  }
  cireg_matrix(NA,nrow=50,ncol=3)
  for(j in 1:50)cireg[j,]_quantile(regmat[,j],c(.025,.5,.975))
  lines(lseq,cireg[,1],lty=2)
  lines(lseq,cireg[,3],lty=2)
  lines(lseq,cireg[,2],lty=1)