R/post_AsContsimmap.R

In bstaggmartin/contSimmap: Stochastic Character Maps of Continuous Trait Data on Phylogenies

#based on these initial tests, coming up with some kind of built-in "shortcut" for processing inputs might be good
#either based on things be univariate and/or some argument like check.inputs=FALSE

#' @export
as.contsimmap<-function(fit,res=100,tol=0.001,max.iter=1e5,verbose=FALSE){
  chains<-fit[['chains']]
  dims<-dim(chains)
  if(length(dims)<3){
    chains<-array(chains,c(dims,1),c(dimnames(chains),chains="chain 1"))
  }
  R0<-as.list(chains[,'R_0',,drop=FALSE])
  nsims<-length(R0)
  Rsig2<-if(any(dimnames(chains)[[2]]=='R_sig2')) as.list(chains[,'R_sig2',,drop=FALSE]) else 0
  Rmu<-if(any(dimnames(chains)[[2]]=='R_mu')) as.list(chains[,'R_mu',,drop=FALSE]) else 0
  #set up "dummy" contsimmap
  list2env(.prep.contsimmap(tree=fit[['call']][['tree']],
                                         nsims=nsims,
                                         res=res,
                                         trait.data=R0,
                                         Xsig2=Rsig2,
                                         Ysig2=0,
                                         mu=Rmu,
                                         conditional=FALSE,
                                         ntraits=1,
                                         traits='R',
                                         nobs=0),envir=environment())
  
  #have some "shortcut" for when Rsig2 is 0...
  #some initialization...
  elen<-c(0,tree[[1]]$edge.length)
  nhgt<-c(0,edge.ranges(tree[[1]])[,2])
  Rsig2<-unlist(Rsig2,use.names=FALSE)
  Rsig<-sqrt(Rsig2)
  Rp<-1/Rsig2
  Rmu<-unlist(Rmu,use.names=FALSE)
  dt.Rs<-sweep(cbind(0,matrix(aperm(remove.trend(fit,simplify=FALSE),c(1,3,2)),nsims,Nedge(fit))),2,elen,'/',check.margin=FALSE)
  t.Rs<-sweep(cbind(0,matrix(aperm(get.R(fit,simplify=FALSE),c(1,3,2)),nsims,Nedge(fit))),2,log(elen),'+',check.margin=FALSE)
  
  #helper functions...
  get.ndes<-function(){
    ndes[e]
  }
  calc.v<-function(){
    1/(2/elen[e]+3*sum(1/elen[des[[e]]]))
  }
  calc.r<-function(){
    tmp.v*(2*dt.Rs[,e]+3*.rowSums(dt.Rs[,des[[e]],drop=FALSE],nsims,tmp.ndes))+Rmu*nhgt[e]
  }
  get.nt<-function(){
    nts[e,1]
  }
  get.dt<-function(){
    maps[[e-1,1,'dts']][1]
  }
  get.chol<-function(){
    tmp<-.col(c(tmp.nt,tmp.nt))
    ttmp<-t(tmp)
    tmp.inds<-tmp>ttmp
    tmp[tmp.inds]<-ttmp[tmp.inds]
    tmp.seq<-tmp.nt:1
    ttmp<-tmp[tmp.seq,tmp.seq]
    tmp.seq<-seq(tmp.dt,length.out=tmp.nt,by=tmp.dt)/sqrt(elen[e])
    chol(matrix(tmp.seq[tmp]*tmp.seq[ttmp],tmp.nt,tmp.nt)+tcrossprod(seq(0,sqrt(tmp.v),length.out=tmp.nt+1)[-1]))
  }
  get.mu<-function(){
    anc.r<-seed[[anc[[e]],1]][nts[anc[[e]],1],,]
    tcrossprod(tmp.r-anc.r,seq(0,1,length.out=tmp.nt+2)[-c(1,tmp.nt+2)])+anc.r
  }
  
  scores<-matrix(0,nsims,dim(seed)[1])
  
  #main traversal...
  #first do R0!
  seed[[1,1]][1,,]<-X0[,lookup[[1]][['nobs.X0']][,2],drop=FALSE]
  if(verbose){
    counter<-cur.prog<-0
    tot<-length(prune.seq)-1
    cat("Sampling rate maps...\n")
    cat(.prog(cur.prog))
  }
  for(e in rev(prune.seq)[-1]){
    #sample rates at edge's descendant node...
    tmp.ndes<-get.ndes()
    if(tmp.ndes){
      tmp.v<-calc.v()
      tmp.r<-calc.r()
    }else{
      tmp.v<-elen[e]/2
      tmp.r<-dt.Rs[,e]*elen[e]+Rmu*nhgt[e]
    }
    #"quick" MCMC?
    #likelihood is simple enough--just the sum of squares of the raw seed + squared deviation between resulting integral and target divided by some small sd, I think...
    #maybe use a standard normal proposal distribution with sd of 0.1ish?
    #pseudo-MCMC actually seems remarkably quick to converge and results in much better approximation
    #able to use MUCH lower tolerances and higher number of iterations!!!
    tmp.nt<-get.nt()
    tmp.dt<-get.dt()
    log.dt<-log(tmp.dt)
    tmp.chol<-get.chol()
    tmp.mu<-get.mu()
    targ<-t.Rs[,e]
    tol2<-tol^2
    cur.seed<-matrix(seed[[e,1]],nsims,tmp.nt,byrow=TRUE)
    tmp.rs<-tmp.mu+sweep(cur.seed%*%tmp.chol,1,Rsig,'*',check.margin=FALSE)
    tmp.score<-(log(.rowSums(exp(tmp.rs),nsims,tmp.nt))+log.dt-targ)^2/tol2
    cur.lik<- -0.5*(.rowSums(cur.seed^2,nsims,tmp.nt)+tmp.score)
    inds<-which(tmp.score>1)
    n.inds<-length(inds)
    for(i in seq_len(max.iter)){
      prop.seed<-cur.seed[inds,,drop=FALSE]+rnorm(n.inds*tmp.nt,0,0.1)
      tmp.rs<-tmp.mu[inds,,drop=FALSE]+sweep(prop.seed%*%tmp.chol,1,Rsig[inds],'*',check.margin=FALSE)
      tmp.score<-(log(.rowSums(exp(tmp.rs),n.inds,tmp.nt))+log.dt-targ[inds])^2/tol2
      prop.lik<- -0.5*(.rowSums(prop.seed^2,n.inds,tmp.nt)+tmp.score)
      tmp<-tmp.score>1
      new.inds<-inds[tmp]
      n.inds<-sum(tmp)
      tmp[tmp]<-(prop.lik[tmp]-cur.lik[inds[tmp]])< -rexp(n.inds)
      tmp<-!tmp
      cur.seed[inds[tmp],]<-prop.seed[tmp,,drop=FALSE]
      cur.lik[inds[tmp]]<-prop.lik[tmp]
      inds<-new.inds
      if(!n.inds) break
    }
    tmp.rs<-tmp.mu+sweep(cur.seed%*%tmp.chol,1,Rsig,'*',check.margin=FALSE)
    seed[[e,1]][]<-as.vector(t(tmp.rs))
    scores[,e]<-abs(log(.rowSums(exp(tmp.rs),nsims,tmp.nt))+log.dt-targ)
    if(verbose){
      counter<-counter+1
      prop.prog<-floor(100*counter/tot)
      if(prop.prog>cur.prog){
        cur.prog<-prop.prog
        cat(.prog(cur.prog))
      }
    }
  }
  
  #format output
  x<-seed
  attr(x,'ts')<-ts
  attr(x,'tree')<-tree
  attr(x,'maps')<-maps
  attr(x,'treeID')<-treeID
  traits<-colnames(Xsig2[[1]])
  attr(x,'traits')<-setNames(rep(1,length(traits)),traits)
  attr(x,'params')<-matrix(list(NULL,Xsig2,Ysig2,mu,lookup,list(fxn="as.contsimmap.evorates_fit")),
                           6,1,
                           dimnames=list(c('trait.data','Xsig2','Ysig2','mu','lookup','call_info'),NULL))
  out<-.uncompress(x)
  #now for the million-dollar test...
  Ysig2<-if(any(dimnames(chains)[[2]]=='Y_sig2')) as.list(chains[,'Y_sig2',,drop=FALSE]) else NULL
  tmp<-fit[['call']][['trait.se']]
  not.nas<-!is.na(tmp)&!is.infinite(tmp)
  Ysig2<-c(list(Ysig2),setNames(as.list(tmp[not.nas]),rownames(tmp[not.nas,,drop=FALSE])))
  dat<-fit[['call']][['trait.data']]
  form<-as.formula(paste0(colnames(dat),"~diffusion(",colnames(dat),"_Xsig2='exp_R',Ysig2=Ysig2,trait.data=dat",if(verbose) ",verbose=TRUE",")"))
  out<-make.traits(out,exp_R~exp(R),form)
}