Merge pull request #114 from DoubleML/p-binary-outcome-irm

Binary outcome variable in IRM and IIVM

Author: PhilippBach

R/datasets.R

@@ -84,14 +84,16 @@ fetch_401k = function(return_type = "DoubleMLData", polynomial_features = FALSE,
     data = data.frame(
       "net_tfa" = data$net_tfa,
       model.matrix(formula_x, data),
-      "p401" = data$p401, "e401" = data$e401)
+      "p401" = data$p401, "e401" = data$e401
+    )
     d_cols = "p401"
     z_cols = "e401"
   } else {
     # see https://github.com/VC2015/DMLonGitHub/blob/b91cbf96c01eccd73367fbd6601ecdd7aa78403b/401K.R#L67
     data = data.frame(
       "net_tfa" = data$net_tfa, model.matrix(formula_x, data),
-      "e401" = data$e401)
+      "e401" = data$e401
+    )
     d_cols = "e401"
   }
   if (return_type == "data.frame") {
@@ -103,7 +105,8 @@ fetch_401k = function(return_type = "DoubleMLData", polynomial_features = FALSE,
     dt = as.data.table(data)
     data = DoubleMLData$new(dt,
       y_col = y_col, d_cols = d_cols, x_cols = x_cols,
-      z_cols = z_cols)
+      z_cols = z_cols
+    )
     return(data)
   }
 }
@@ -199,7 +202,9 @@ fetch_401k = function(return_type = "DoubleMLData", polynomial_features = FALSE,
 #'   x_cols = c(
 #'     "female", "black", "othrace", "dep1", "dep2",
 #'     "q2", "q3", "q4", "q5", "q6", "agelt35", "agegt54",
-#'     "durable", "lusd", "husd"))
+#'     "durable", "lusd", "husd"
+#'   )
+#' )
 #' obj_dml_data_bonus
 #' @export
 fetch_bonus = function(return_type = "DoubleMLData",
@@ -235,7 +240,8 @@ fetch_bonus = function(return_type = "DoubleMLData",
   }
   data = data.frame(
     "inuidur1" = data$inuidur1, model.matrix(formula_x, data),
-    "tg" = data$tg)
+    "tg" = data$tg
+  )
   if (return_type == "data.frame") {
     return(data)
   } else if (return_type == "data.table") {
@@ -315,7 +321,12 @@ make_plr_CCDDHNR2018 = function(n_obs = 500, dim_x = 20, alpha = 0.5,
 
   assert_choice(
     return_type,
-    c("data.table", "matrix", "data.frame", "DoubleMLData"))
+    c("data.table", "matrix", "data.frame", "DoubleMLData")
+  )
+  assert_count(n_obs)
+  assert_count(dim_x)
+  assert_numeric(alpha, len = 1)
+
   cov_mat = toeplitz(0.7^(0:(dim_x - 1)))
   a_0 = 1
   a_1 = 0.25
@@ -405,6 +416,16 @@ make_plr_CCDDHNR2018 = function(n_obs = 500, dim_x = 20, alpha = 0.5,
 make_plr_turrell2018 = function(n_obs = 100, dim_x = 20, theta = 0.5,
   return_type = "DoubleMLData", nu = 0, gamma = 1) {
 
+  assert_choice(
+    return_type,
+    c("data.table", "matrix", "data.frame", "DoubleMLData")
+  )
+  assert_count(n_obs)
+  assert_count(dim_x)
+  assert_numeric(theta, len = 1)
+  assert_numeric(nu, len = 1)
+  assert_numeric(gamma, len = 1)
+
   b = 1 / (1:dim_x)
   sigma = genPositiveDefMat(dim_x)
   x = rmvnorm(n = n_obs, mean = rep(0, dim_x), sigma = sigma$Sigma)
@@ -497,7 +518,12 @@ make_pliv_CHS2015 = function(n_obs, alpha = 1, dim_x = 200, dim_z = 150,
 
   assert_choice(
     return_type,
-    c("data.table", "matrix", "data.frame", "DoubleMLData"))
+    c("data.table", "matrix", "data.frame", "DoubleMLData")
+  )
+  assert_count(n_obs)
+  assert_count(dim_x)
+  assert_count(dim_z)
+  assert_numeric(alpha, len = 1)
   if (dim_x < dim_z) {
     stop("Dimension of X should be greater than dimension of Z.")
   }
@@ -545,14 +571,13 @@ make_pliv_CHS2015 = function(n_obs, alpha = 1, dim_x = 200, dim_z = 150,
       data = DoubleMLData$new(dt,
         y_col = "y", d_cols = "d",
         x_cols = colnames(x),
-        z_cols = colnames(z))
+        z_cols = colnames(z)
+      )
       return(data)
     }
   }
-  return(data)
 }
 
-
 #' @title Generates data from a interactive regression (IRM) model.
 #'
 #' @description
@@ -610,7 +635,16 @@ make_irm_data = function(n_obs = 500, dim_x = 20, theta = 0, R2_d = 0.5,
   # inspired by https://onlinelibrary.wiley.com/doi/abs/10.3982/ECTA12723
   # (see supplement)
 
-  assert_choice(return_type, c("data.table", "matrix", "data.frame", "DoubleMLData"))
+  assert_choice(
+    return_type,
+    c("data.table", "matrix", "data.frame", "DoubleMLData")
+  )
+  assert_count(n_obs)
+  assert_count(dim_x)
+  assert_numeric(theta, len = 1)
+  assert_numeric(R2_d, len = 1)
+  assert_numeric(R2_y, len = 1)
+
   v = runif(n_obs)
   zeta = rnorm(n_obs)
   cov_mat = toeplitz(0.5^(0:(dim_x - 1)))
@@ -702,10 +736,16 @@ make_iivm_data = function(n_obs = 500, dim_x = 20, theta = 1, alpha_x = 0.2,
 
   assert_choice(
     return_type,
-    c("data.table", "matrix", "data.frame", "DoubleMLData"))
+    c("data.table", "matrix", "data.frame", "DoubleMLData")
+  )
+  assert_count(n_obs)
+  assert_count(dim_x)
+  assert_numeric(theta, len = 1)
+  assert_numeric(alpha_x, len = 1)
   xx = rmvnorm(
     n = n_obs, mean = rep(0, 2),
-    sigma = matrix(c(1, 0.3, 0.3, 1), ncol = 2, nrow = 2))
+    sigma = matrix(c(1, 0.3, 0.3, 1), ncol = 2, nrow = 2)
+  )
   u = xx[, 1]
   v = xx[, 2]
 
@@ -739,7 +779,8 @@ make_iivm_data = function(n_obs = 500, dim_x = 20, theta = 1, alpha_x = 0.2,
       data = DoubleMLData$new(dt,
         y_col = "y", d_cols = "d",
         x_cols = colnames(x),
-        z_cols = "z")
+        z_cols = "z"
+      )
       return(data)
     }
   }

R/double_ml_data.R

@@ -289,10 +289,10 @@ DoubleMLData = R6Class("DoubleMLData",
     },
 
     #' @description
-    #' Setter function for `data_model`. The function implements the causal model
-    #' as specified by the user via `y_col`, `d_cols`, `x_cols` and `z_cols` and
-    #' assigns the role for the treatment variables in the multiple-treatment
-    #' case.
+    #' Setter function for `data_model`. The function implements the causal
+    #' model as specified by the user via `y_col`, `d_cols`, `x_cols` and
+    #' `z_cols` and assigns the role for the treatment variables in the
+    #' multiple-treatment case.
     #' @param treatment_var (`character()`)\cr
     #' Active treatment variable that will be set to `treat_col`.
     set_data_model = function(treatment_var) {

R/double_ml_iivm.R

@@ -139,16 +139,19 @@ DoubleMLIIVM = R6Class("DoubleMLIIVM",
     #' of the causal model.
     #'
     #' @param ml_g ([`LearnerRegr`][mlr3::LearnerRegr],
-    #' [`Learner`][mlr3::Learner], `character(1)`) \cr
+    #' [`LearnerClassif`][mlr3::LearnerClassif], [`Learner`][mlr3::Learner],
+    #' `character(1)`) \cr
     #' A learner of the class [`LearnerRegr`][mlr3::LearnerRegr], which is
     #' available from [mlr3](https://mlr3.mlr-org.com/index.html) or its
     #' extension packages [mlr3learners](https://mlr3learners.mlr-org.com/) or
     #' [mlr3extralearners](https://mlr3extralearners.mlr-org.com/).
+    #' For binary treatment outcomes, an object of the class
+    #' [`LearnerClassif`][mlr3::LearnerClassif] can be passed, for example
+    #' `lrn("classif.cv_glmnet", s = "lambda.min")`.
     #' Alternatively, a [`Learner`][mlr3::Learner] object with public field
-    #' `task_type = "regr"` can be passed, for example of class
-    #' [`GraphLearner`][mlr3pipelines::GraphLearner]. The learner can possibly
-    #' be passed with specified parameters, for example
-    #' `lrn("regr.cv_glmnet", s = "lambda.min")`. \cr
+    #' `task_type = "regr"` or `task_type = "classif"` can be passed,
+    #' respectively, for example of class
+    #' [`GraphLearner`][mlr3pipelines::GraphLearner]. \cr
     #' `ml_g` refers to the nuisance function \eqn{g_0(Z,X) = E[Y|X,Z]}.
     #'
     #' @param ml_m ([`LearnerClassif`][mlr3::LearnerClassif],
@@ -246,7 +249,7 @@ DoubleMLIIVM = R6Class("DoubleMLIIVM",
         "ml_g" = NULL,
         "ml_m" = NULL,
         "ml_r" = NULL)
-      ml_g = private$assert_learner(ml_g, "ml_g", Regr = TRUE, Classif = FALSE)
+      ml_g = private$assert_learner(ml_g, "ml_g", Regr = TRUE, Classif = TRUE)
       ml_m = private$assert_learner(ml_m, "ml_m", Regr = FALSE, Classif = TRUE)
       ml_r = private$assert_learner(ml_r, "ml_r", Regr = FALSE, Classif = TRUE)
 

R/double_ml_irm.R

@@ -108,16 +108,19 @@ DoubleMLIRM = R6Class("DoubleMLIRM",
     #' of the causal model.
     #'
     #' @param ml_g ([`LearnerRegr`][mlr3::LearnerRegr],
-    #' [`Learner`][mlr3::Learner], `character(1)`) \cr
+    #' [`LearnerClassif`][mlr3::LearnerClassif], [`Learner`][mlr3::Learner],
+    #' `character(1)`) \cr
     #' A learner of the class [`LearnerRegr`][mlr3::LearnerRegr], which is
     #' available from [mlr3](https://mlr3.mlr-org.com/index.html) or its
     #' extension packages [mlr3learners](https://mlr3learners.mlr-org.com/) or
     #' [mlr3extralearners](https://mlr3extralearners.mlr-org.com/).
+    #' For binary treatment outcomes, an object of the class
+    #' [`LearnerClassif`][mlr3::LearnerClassif] can be passed, for example
+    #' `lrn("classif.cv_glmnet", s = "lambda.min")`.
     #' Alternatively, a [`Learner`][mlr3::Learner] object with public field
-    #' `task_type = "regr"` can be passed, for example of class
-    #' [`GraphLearner`][mlr3pipelines::GraphLearner]. The learner can possibly
-    #' be passed with specified parameters, for example
-    #' `lrn("regr.cv_glmnet", s = "lambda.min")`. \cr
+    #' `task_type = "regr"` or `task_type = "classif"` can be passed,
+    #' respectively, for example of class
+    #' [`GraphLearner`][mlr3pipelines::GraphLearner]. \cr
     #' `ml_g` refers to the nuisance function \eqn{g_0(X) = E[Y|X,D]}.
     #'
     #' @param ml_m ([`LearnerClassif`][mlr3::LearnerClassif],
@@ -189,7 +192,7 @@ DoubleMLIRM = R6Class("DoubleMLIRM",
       private$task_type = list(
         "ml_g" = NULL,
         "ml_m" = NULL)
-      ml_g = private$assert_learner(ml_g, "ml_g", Regr = TRUE, Classif = FALSE)
+      ml_g = private$assert_learner(ml_g, "ml_g", Regr = TRUE, Classif = TRUE)
       ml_m = private$assert_learner(ml_m, "ml_m", Regr = FALSE, Classif = TRUE)
 
       private$learner_ = list(

R/helper.R

@@ -190,7 +190,9 @@ extract_prediction = function(obj_resampling, task_type, n_obs,
     }
   } else {
     preds = rep(NA_real_, n_obs)
-    if (testR6(obj_resampling, classes = "ResampleResult")) obj_resampling = list(obj_resampling)
+    if (testR6(obj_resampling, classes = "ResampleResult")) {
+      obj_resampling = list(obj_resampling)
+    }
     n_obj_rsmp = length(obj_resampling)
     for (i_obj_rsmp in 1:n_obj_rsmp) {
       f_hat = as.data.table(obj_resampling[[i_obj_rsmp]]$prediction("test"))
@@ -205,7 +207,9 @@ initiate_learner = function(learner, task_type, params, return_train_preds = FAL
   ml_learner = learner$clone()
 
   if (!is.null(params)) {
-    ml_learner$param_set$values = insert_named(ml_learner$param_set$values, params)
+    ml_learner$param_set$values = insert_named(
+      ml_learner$param_set$values,
+      params)
   } # else if (is.null(params) | length(params) == 0) {
   # message("No parameters provided for learners. Default values are used.")
   # }