R/simulation_functions.R

In cfhammill/hierarchyTrees: Code for testing taxonomy tree models

#' Recompute volumes for non-leaf nodes
#'
#' Destructively aggregate volumes from the leaves to the
#' root. This sets each non-leaf node to be the sum of its
#' children
#' @param tree the tree of interest
#' @return the tree invisibly
#' @export
aggregate_volsD <-
  function(tree){

    tree$Do(function(n){
      n$volumes <-
        `if`(!isLeaf(n)
           , Reduce(function(acc, x) cbind(acc, x$volumes), n$children, NULL) %>%
             rowSums()
           , n$volumes)
      
      n$meanVolume <- mean(n$volumes)      
    }, traversal = "post-order")

    invisible(tree)
  }

#' Check parent
#'
#' Check if node descends from node "name"
#'
#' @param node The node of interest
#' @param name The name to look for in the nodes
#' ancestry
#' @return boolean whether one of nodes ancestors
#' has the specified `name`
#' @export
child_of <- function(node, name){
  if(is.null(node$parent))
    return(FALSE)

  if(node$parent$name == name)
    return(TRUE)

  child_of(node$parent, name)
}

#' Increment volumes
#'
#' Destructively add effect to the volume
#' of all leaves at or below node "name"
#'
#' @param tree The tree of interest
#' @param name The node name to increment at or below
#' @param effect a scalar or vector corresponding to the
#' amount to add to each leaf.
#' @return the tree invisibly
#' @export
add_effectD <-
  function(tree, effect, name = tree$name){
    tree$Do(function(n){
      if(n$name == name || child_of(n, name))
         n$volumes <- n$volumes + effect
    })

    invisible(tree)
  }

#' Scale node volumes up
#'
#' Destructively apply a scale factor to the volumes of all leaf
#' nodes at or below node "name"
#'
#' @param tree the tree of interest
#' @param scale a scalar or vector of how much to scale each
#' volume
#' @param name the node name to increment at or below
#' @return the tree invisibly
#' @export
scale_volumesD <-
  function(tree, scale, name = tree$name){
    tree$Do(function(n){
      if(n$name == name || child_of(n, name))
        n$volumes <- n$volumes * scale
    })

    invisible(tree)    
  }

#' Set leaves
#'
#' Set the volumes at the leaves of a tree to a specific value.
#'
#' @param tree the tree of interest
#' @param vals a scalar or vector of replacement values
#' @param name the node name to set at or below
#' @return the tree invisibly
#' @export
set_leavesD <-
  function(tree, vals, name = tree$name){
    tree$Do(function(n){
      if(n$name == name || child_of(n, name))
        n$volumes[] <- vals
    })

    invisible(tree)    
  }


#' Generate a simulated tree
#'
#' Create a simulated tree from a noise sd,
#' named vectors of leaf effects, volumes, individual
#' effects, metadata, and an initial tree
#'
#' @param noise_sd What should the sd of the gaussian noise
#' added to observations be
#' @param leaf_effects A named vector of effects to apply to
#' the leaves
#' @param leaf_volume A named vector of scaling constants to
#' be applied to the leaf nodes, converting them to `volume` scale.
#' @param indiv_effects A vector of individual level effects
#' @param base_tree A tree indicating the ontology, should contain
#' volumes already to be overwritten.
#' @param extra_effects A named vector of extra effects to be applied
#' at nodes. Default `NA` for no extra effects
#' @param metadata A vector containing a covariate
#' with elements matching the elements of the indiv effects / tree
#' volumes.
#' @return A two element list containing
#' 1. The simulated tree
#' 2. A vector with the composite `leaf_effects` generated by adding
#'  `extra_effects` to `leaf_effects` where applicable.
#' @export
create_tree <-
  function(noise_sd
         , leaf_effects
         , leaf_volumes
         , indiv_effects
         , metadata
         , base_tree
         , intercept_sd = 0
         , extra_effects = NA){
    
    simulated_tree <- Clone(base_tree)
    nodes <- simulated_tree$Get("name")

    ## Zero the leaves
    set_leavesD(simulated_tree, 0)

    ## Add effects to tree
    walk2(leaf_effects, names(leaf_effects)
        , function(b, name){
          add_effectD(simulated_tree, b, name = name)
        })

    ## Optionally Add extra effects
    if(!is.na(extra_effects[[1]])){
      walk2(extra_effects, names(extra_effects),
            function(b, name){
              add_effectD(simulated_tree, b, name)
            })
    }

    ## Aggregate effects to see parent effects
    aggregate_volsD(simulated_tree)

    ## Compute composite effects
    composite_effects <-
      map_dbl(nodes, function(l)
        simulated_tree$Get(function(n) n$volumes[1]
                         , filterFun = function(n) n$name == l)
        )
    
    ## Multiply in covariate
    scale_volumesD(simulated_tree, metadata)

    ## Add some intercept jittercx
    ## intercepts <- rnorm(length(leaf_effects), 0, sd = intercept_sd)
    ## names(intercepts) <- names(leaf_effects)
    ## walk2(intercepts, names(intercepts)
    ##       , function(i, name){
    ##         add_effectD(simulated_tree, i * (metadata == 0), name)
    ##       })

    ## Add individual effects
    add_effectD(simulated_tree, indiv_effects)

    ## Add some jitter
    walk(names(leaf_effects)
       , function(name)
         add_effectD(simulated_tree, rnorm(length(metadata), 0, noise_sd), name = name))



    ## Add an intercept and scale the effects
    ## walk2(leaf_volumes, names(leaf_volumes)
    ##     , function(v, name){
    ##       add_effectD(simulated_tree, 1, name = name)
    ##       scale_volumesD(simulated_tree, v, name = name)
    ##     })

    ## Aggregate up the tree
    aggregate_volsD(simulated_tree)

    list(effects = composite_effects
       , tree = simulated_tree)
  }