Merge pull request #108 from DoubleML/m-check-01-binary-for-classif

Extended exception handling for classification

Author: MalteKurz

R/helper.R

@@ -229,6 +229,8 @@ initiate_task = function(id, data, target, select_cols, learner_class) {
     task = TaskRegr$new(id = id, backend = data, target = target)
   } else if (learner_class == "LearnerClassif") {
     data[[target]] = factor(data[[target]])
+    assert_set_equal(levels(data[[target]]),
+                     c("0", "1"))
     task = TaskClassif$new(
       id = id, backend = data, target = target,
       positive = "1")

tests/testthat/test-double_ml_plr_classifier.R

@@ -89,3 +89,33 @@ patrick::with_parameters_test_that("Unit tests for PLR with classifier for ml_m:
     }
   }
 )
+
+test_that("Unit tests for exception handling of PLR with classifier for ml_m:", {
+  # Only binary outcome with values 0 and 1 is allowed when ml_m is a classifier
+  
+  # Test with 0 and 2
+  df = data_irm$df
+  df['d'] = df['d']*2
+  dml_data = double_ml_data_from_data_frame(df, y_col = 'y', d_cols = 'd')
+  double_mlplr_obj = DoubleMLPLR$new(data = dml_data,
+                                     ml_g = mlr3::lrn('regr.rpart'),
+                                     ml_m = mlr3::lrn('classif.rpart'))
+  msg = paste("Assertion on 'levels\\(data\\[\\[target\\]\\])' failed:",
+              "Must be equal to set \\{'0','1'\\}, but is \\{'0','2'\\}.")
+  expect_error(double_mlplr_obj$fit(),
+               regexp = msg)
+  
+  # Test with 0.5 and 1
+  df = data_irm$df
+  df['d'] = (df['d']+2)/2
+  dml_data = double_ml_data_from_data_frame(df, y_col = 'y', d_cols = 'd')
+  double_mlplr_obj = DoubleMLPLR$new(data = dml_data,
+                                     ml_g = mlr3::lrn('regr.rpart'),
+                                     ml_m = mlr3::lrn('classif.rpart'))
+  msg = paste("Assertion on 'levels\\(data\\[\\[target\\]\\])' failed:",
+              "Must be equal to set \\{'0','1'\\}, but is \\{'1','1.5'\\}.")
+  expect_error(double_mlplr_obj$fit(),
+               regexp = msg)
+
+  }
+)