R/05a.Grow_continuous.R

In btrm: Bayesian Treed Regression Model for Personalized Prediction and Precision Diagnostics

Grow_continuous=function(ET,STN,minsample){

  ####
  #1. Initial value
  ####
  p=ET$p
  q=ET$q
  y=ET$y
  x=ET$x

  splitVariable=ET$splitVariable
  cutoff=ET$cutoff
  marker=ET$marker

  node.hat=ET$node.hat
  marker.hat=ET$marker.hat

  internal=ET$internal
  terminal=ET$terminal

  eta=ET$eta

  ####
  # 2. Update tree structure
  ####
  #2.1. select feature
  splitVariable[STN]=sample(1:q,1,prob = ET$dir.predictor)

  idx=(node.hat==STN)                            # idx, i.e. those who belong to terminal node=STN
  x.sel=unique(x[which(idx),splitVariable[STN]]) # selected x for those who belong to terminal node=STN
  cutoff[STN]=sample(x.sel,1)
  eta[STN]=length(x.sel)                         # the number of available values which could be select to split the chosen terminal node

  #2.2. Split STN into left & right
  left=which(idx & x[,splitVariable[STN]]<=cutoff[STN])
  right=which(idx & x[,splitVariable[STN]]>cutoff[STN])

  LC.STN=2*STN   #left child node number for STN
  RC.STN=2*STN+1 #right child node number

  node.hat[left]=LC.STN
  node.hat[right]=RC.STN

  cond2=length(left)>minsample
  cond3=length(right)>minsample
  size.cond=cond2&cond3

  if(size.cond){
    #2.3. select marker
    sel.marker.L=sample(1:p,1,prob = ET$dir.marker)
    sel.marker.R=sample(1:p,1,prob = ET$dir.marker)

    marker[LC.STN]=sel.marker.L
    marker[RC.STN]=sel.marker.R

    marker.hat[left]=marker[LC.STN]               # selected marker for each subj
    marker.hat[right]=marker[RC.STN]

    #2.4. updated internal & terminal nodes
    internal=sort(c(internal,STN))               # NA is removed by sort, if internal is NA
    terminal=sort(c(setdiff(terminal,internal),LC.STN,RC.STN))
            #setdiff(sort(c(terminal,LC.STN,RC.STN)),internal)

  	####
    # 3. summary
    ####
    ET$eta=eta
    ET$splitVariable=splitVariable
    ET$cutoff=cutoff

    ET$marker=marker
    ET$node.hat=node.hat
    ET$marker.hat=marker.hat

    ET$internal=internal
    ET$terminal=terminal

    ET$numNodes=length(ET$terminal)
  }

  ET$size.cond=size.cond
  ET$STN=STN

  return(ET)
}