From 80ead377857741bfe624ece5a8877457a61512e0 Mon Sep 17 00:00:00 2001 From: charliewang123 Date: Thu, 25 Jun 2026 09:03:37 -0400 Subject: [PATCH] Update NetworkCausalTree.R --- R/NetworkCausalTree.R | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/R/NetworkCausalTree.R b/R/NetworkCausalTree.R index 08a91fa..1873259 100644 --- a/R/NetworkCausalTree.R +++ b/R/NetworkCausalTree.R @@ -128,9 +128,7 @@ NetworkCausalTree <- function(X, Y, W, # Sample the clusters to be assigned to the discovery set nclusters_disc = round(k * ratio_disc) - clusters_disc <- sample(1 : k, - size = nclusters_disc, - replace = FALSE) + clusters_disc <- sample(unique(K), size = nclusters_disc, replace = FALSE) # Generate the tree on the discovery set tree <- sprout_nct(method = method, @@ -157,7 +155,7 @@ 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) + clusters_est <- setdiff(unique(K), clusters_disc) est_units <- which(K %in% clusters_est) A_est <- A[est_units, est_units]