Fix honest sample splitting

R/NetworkCausalTree.R

@@ -157,17 +157,29 @@ NetworkCausalTree <- function(X, Y, W,
   tree$FILTER <- as.character(tree$FILTER)
   tree$FILTER[is.na(tree$FILTER)] <- "NA"
 
+  clusters_est <- setdiff(1:k, clusters_disc)
+  est_units <- which(K %in% clusters_est)
+
+  A_est <- A[est_units, est_units]
+  Ne_est <- rowSums(A_est)
+  Ne_treated_est <- as.vector(A_est %*% W[est_units])
+  G_est <- rep(1, length(est_units))
+  G_est[Ne_treated_est == 0] <- 0
+  Ne_list_est <- lapply(seq_along(est_units), function(i) {
+    which(A_est[i, ] > 0)
+  })
+
   # Compute estimates on the full tree
   tree_with_estimates <- compute_effects_nct(output = output,
                                              nct_partition = tree,
-                                             N = N,
-                                             W = W,
-                                             G = G,
-                                             Y = Y,
-                                             X = X,
-                                             p = p,
-                                             Ne = Ne,
-                                             Ne_list = Ne_list,
+                                             N = length(est_units),
+                                             W = W[est_units],
+                                             G = G_est,
+                                             Y = Y[est_units],
+                                             X = X[est_units, , drop = FALSE],
+                                             p = p[est_units],
+                                             Ne = Ne_est,
+                                             Ne_list = Ne_list_est,
                                              minsize = minsize)
 
   return(tree_with_estimates)

tests/testthat/test-plot_NCT.R

@@ -2,11 +2,13 @@ test_that("plot_NCT runs on a real tiny NCT object", {
   skip_on_cran()
 
   set.seed(100)
-  d <- data_generator_direct_indirect(N = 100, k = 5, M = 2, p = rep(0.3,100), remove_isolates = TRUE)
+  d <- data_generator_direct_indirect(N = 500, k = 10, M = 2, p = rep(0.3, 500), remove_isolates = TRUE)
 
   fit <- NetworkCausalTree(
     Y = d$Y, W = d$W, X = d$X, A = d$A, p = d$p,
-    method = "OF", depth = 1, minsize = 20,
+    K = d$K,
+    effect_weights = c(1,0,0,0),
+    method = "singular", depth = 1, minsize = 20,
     ratio_disc = 0.5
   )
 
@@ -22,11 +24,13 @@ test_that("plot_NCT works in detection mode", {
   skip_on_cran()
 
   set.seed(100)
-  d <- data_generator_direct_indirect(N = 100, k = 5, M = 2, p = rep(0.3,100), remove_isolates = TRUE)
+  d <- data_generator_direct_indirect(N = 500, k = 10, M = 2, p = rep(0.3, 500), remove_isolates = TRUE)
 
   fit <- NetworkCausalTree(
     Y = d$Y, W = d$W, X = d$X, A = d$A, p = d$p,
-    method = "OF", depth = 1, minsize = 20,
+    K = d$K,
+    effect_weights = c(1,0,0,0),
+    method = "singular", depth = 1, minsize = 20,
     ratio_disc = 0.5,
     output = "detection"
   )