Use parent forest seed in auxiliary forest (grf-labs#1070)

vallader · Nov 25, 2021 · 0dd9537 · 0dd9537
1 parent 3b5f89b
commit 0dd9537
Show file tree

Hide file tree

Showing 11 changed files with 14 additions and 2 deletions.
diff --git a/r-package/grf/R/causal_forest.R b/r-package/grf/R/causal_forest.R
@@ -260,6 +260,7 @@ causal_forest <- function(X, Y, W,
 
   forest <- do.call.rcpp(causal_train, c(data, args))
   class(forest) <- c("causal_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["ci.group.size"]] <- ci.group.size
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/causal_survival_forest.R b/r-package/grf/R/causal_survival_forest.R
@@ -327,6 +327,7 @@ causal_survival_forest <- function(X, Y, W, D,
 
   forest <- do.call.rcpp(causal_survival_train, c(data, args))
   class(forest) <- c("causal_survival_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["_eta"]] <- eta
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/get_scores.R b/r-package/grf/R/get_scores.R
@@ -78,7 +78,8 @@ get_scores.causal_forest <- function(forest,
                                            clusters = clusters,
                                            sample.weights = forest$sample.weights,
                                            num.trees = num.trees.for.weights,
-                                           ci.group.size = 1)
+                                           ci.group.size = 1,
+                                           seed = forest$seed)
       V.hat <- predict(variance_forest)$predictions
       debiasing.weights.all <- (forest$W.orig - forest$W.hat) / V.hat
       debiasing.weights <- debiasing.weights.all[subset]
@@ -176,7 +177,8 @@ get_scores.instrumental_forest <- function(forest,
                                          W.hat = forest$Z.hat,
                                          sample.weights = forest$sample.weights,
                                          clusters = clusters,
-                                         num.trees = num.trees.for.weights)
+                                         num.trees = num.trees.for.weights,
+                                         seed = forest$seed)
       compliance.score <- predict(compliance.forest)$predictions
       compliance.score <- compliance.score[subset]
     } else if (length(compliance.score) == length(forest$Y.orig)) {

diff --git a/r-package/grf/R/instrumental_forest.R b/r-package/grf/R/instrumental_forest.R
@@ -238,6 +238,7 @@ instrumental_forest <- function(X, Y, W, Z,
 
   forest <- do.call.rcpp(instrumental_train, c(data, args))
   class(forest) <- c("instrumental_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["ci.group.size"]] <- ci.group.size
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/ll_regression_forest.R b/r-package/grf/R/ll_regression_forest.R
@@ -197,6 +197,7 @@ ll_regression_forest <- function(X, Y,
   }
 
   class(forest) <- c("ll_regression_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["ci.group.size"]] <- ci.group.size
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/multi_arm_causal_forest.R b/r-package/grf/R/multi_arm_causal_forest.R
@@ -271,6 +271,7 @@ multi_arm_causal_forest <- function(X, Y, W,
 
   forest <- do.call.rcpp(multi_causal_train, c(data, args))
   class(forest) <- c("multi_arm_causal_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["ci.group.size"]] <- ci.group.size
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/multi_regression_forest.R b/r-package/grf/R/multi_regression_forest.R
@@ -109,6 +109,7 @@ multi_regression_forest <- function(X, Y,
 
   forest <- do.call.rcpp(multi_regression_train, c(data, args))
   class(forest) <- c("multi_regression_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y
   forest[["sample.weights"]] <- sample.weights

diff --git a/r-package/grf/R/probability_forest.R b/r-package/grf/R/probability_forest.R
@@ -136,6 +136,7 @@ probability_forest <- function(X, Y,
 
   forest <- do.call.rcpp(probability_train, c(data, args))
   class(forest) <- c("probability_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y
   forest[["Y.relabeled"]] <- Y.relabeled

diff --git a/r-package/grf/R/quantile_forest.R b/r-package/grf/R/quantile_forest.R
@@ -128,6 +128,7 @@ quantile_forest <- function(X, Y,
 
   forest <- do.call.rcpp(quantile_train, c(data, args))
   class(forest) <- c("quantile_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y
   forest[["quantiles.orig"]] <- quantiles

diff --git a/r-package/grf/R/regression_forest.R b/r-package/grf/R/regression_forest.R
@@ -167,6 +167,7 @@ regression_forest <- function(X, Y,
 
   forest <- do.call.rcpp(regression_train, c(data, args))
   class(forest) <- c("regression_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["ci.group.size"]] <- ci.group.size
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y

diff --git a/r-package/grf/R/survival_forest.R b/r-package/grf/R/survival_forest.R
@@ -179,6 +179,7 @@ survival_forest <- function(X, Y, D,
 
   forest <- do.call.rcpp(survival_train, c(data, args))
   class(forest) <- c("survival_forest", "grf")
+  forest[["seed"]] <- seed
   forest[["X.orig"]] <- X
   forest[["Y.orig"]] <- Y
   forest[["Y.relabeled"]] <- Y.relabeled