Update prescriptive warning message (grf-labs#1040)

vallader · Aug 25, 2021 · b2d4575 · b2d4575
1 parent 6684b14
commit b2d4575
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/r-package/grf/R/multi_arm_causal_forest.R b/r-package/grf/R/multi_arm_causal_forest.R
@@ -202,8 +202,8 @@ multi_arm_causal_forest <- function(X, Y, W,
   if (nlevels(W) == 1) {
     stop("Can not compute contrasts from a single treatment.")
   }
-  if (!identical(levels(W), levels(droplevels(W)))) {
-    warning("The treatment vector W contains unused levels.")
+  if (nlevels(W) != nlevels(droplevels(W))) {
+    warning("The treatment vector W contains unused levels (see `droplevels()` to drop unused levels).")
   }
 
   args.orthog <- list(X = X,

diff --git a/r-package/grf/R/probability_forest.R b/r-package/grf/R/probability_forest.R
@@ -109,8 +109,8 @@ probability_forest <- function(X, Y,
   if (!is.factor(Y)) {
     stop("The class labels must be a factor vector.")
   }
-  if (!identical(levels(Y), levels(droplevels(Y)))) {
-    warning("The class label vector Y contains unused levels.")
+  if (nlevels(Y) != nlevels(droplevels(Y))) {
+    warning("The class label vector Y contains unused levels (see `droplevels()` to drop unused levels).")
   }
   Y.relabeled <- as.numeric(Y) - 1 # convert to integers between 0 and num_classes
   class.names <- levels(Y)