Return offspring as part of object

R/epichains.R

@@ -19,13 +19,15 @@
 new_epichains_tree <- function(tree_df,
                                ntrees = integer(),
                                statistic = character(),
+                               offspring_dist = character(),
                                stat_max = integer(),
                                track_pop = logical()) {
   # Assemble the elements of the object
   obj <- structure(
     tree_df,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max,
     track_pop = track_pop,
     class = c("epichains_tree", class(tree_df))
@@ -58,12 +60,15 @@ new_epichains_tree <- function(tree_df,
 epichains_tree <- function(tree_df,
                            ntrees = integer(),
                            statistic = character(),
+                           offspring_dist = character(),
                            stat_max = integer(),
                            track_pop = logical()) {
   # Check that inputs are well specified
   checkmate::assert_data_frame(tree_df)
   checkmate::assert_integerish(ntrees, null.ok = TRUE)
   checkmate::assert_character(statistic, null.ok = TRUE)
+  check_offspring_valid(offspring_dist)
+  check_offspring_func_valid(paste0("r", offspring_dist))
   checkmate::assert_logical(track_pop)
   checkmate::assert_number(stat_max, null.ok = TRUE)
 
@@ -72,6 +77,7 @@ epichains_tree <- function(tree_df,
     tree_df = tree_df,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max,
     track_pop = track_pop
   )
@@ -103,12 +109,14 @@ epichains_tree <- function(tree_df,
 new_epichains_summary <- function(chains_summary,
                                   ntrees = integer(),
                                   statistic = character(),
+                                  offspring_dist = character(),
                                   stat_max = integer()) {
   # Assemble the elements of the object
   obj <- structure(
     chains_summary,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max,
     class = c("epichains_summary", class(chains_summary))
   )
@@ -133,18 +141,22 @@ new_epichains_summary <- function(chains_summary,
 epichains_summary <- function(chains_summary,
                               ntrees = integer(),
                               statistic = character(),
+                              offspring_dist = character(),
                               stat_max = integer()) {
   # Check that inputs are well specified
   checkmate::assert_vector(chains_summary)
   checkmate::assert_integerish(ntrees, null.ok = TRUE)
   checkmate::assert_character(statistic)
+  check_offspring_valid(offspring_dist)
+  check_offspring_func_valid(paste0("r", offspring_dist))
   checkmate::assert_number(stat_max, null.ok = TRUE)
 
   # Create <epichains_summary> object
   epichains_summary <- new_epichains_summary(
     chains_summary,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max
   )
 

R/simulate.r

@@ -242,6 +242,7 @@ simulate_tree <- function(ntrees, statistic = c("size", "length"),
     tree_df = tree_df,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max,
     track_pop = FALSE
   )
@@ -339,6 +340,7 @@ simulate_summary <- function(ntrees, statistic = c("size", "length"),
     chains_summary = stat_track,
     ntrees = ntrees,
     statistic = statistic,
+    offspring_dist = offspring_dist,
     stat_max = stat_max
   )
 
@@ -552,6 +554,7 @@ simulate_tree_from_pop <- function(pop,
     tree_df,
     ntrees = NULL,
     statistic = NULL,
+    offspring_dist = offspring_dist,
     stat_max = NULL,
     track_pop = TRUE
   )