R/05d.Swap.R

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

Swap=function(ET,SIN,minsample){

  ####
  #1. Initial value
  ####
  n=ET$n
  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=swapSplit(splitVariable, SIN, parentNode(SIN))
  cutoff=swapSplit(cutoff, SIN, parentNode(SIN))

  eta=eta[1]
  node.hat=rep(1,n)
  reverse=FALSE
  for(i in internal){
    # define index number and eta number
    idx=(node.hat==i)
    x.sel=unique(x[which(idx),splitVariable[i]])
    eta[i]=length(x.sel)

    # Split node i into left & right
    left=which(idx & x[,splitVariable[i]]<=cutoff[i])
    right=which(idx & x[,splitVariable[i]]>cutoff[i])

    node.hat[left]=2*i
    node.hat[right]=2*i+1

    cond1=min(table(y,node.hat))>minsample         # subgroup sample size is large enough
    cond2=length(left)>0                           # table above does not work if left or right is empty
    cond3=length(right)>0
    cond4=length(table(y,node.hat))>1              # exclude the cases table contains NA values
    size.cond=cond1&cond2&cond3&cond4

    if(!size.cond){
      reverse=TRUE
      break
    }else{
      marker.hat[left]=marker[2*i]               # selected marker for each subj
      marker.hat[right]=marker[2*i+1]
    }
  }

  ####
  # 3. summary
  ####
  if(reverse){                        # if we did not do "SWAP"
    ET$size.cond=FALSE
  }else{
    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=TRUE
  }

  return(ET)
}