simplify if clauses in score_elements of pliv

Author: MalteKurz

R/double_ml_pliv.R

@@ -332,9 +332,19 @@ DoubleMLPLIV = R6Class("DoubleMLPLIV",
       w_hat = d - r_hat
       v_hat = z - m_hat
 
-      if (self$data$n_instr > 1) {
-        v_hat = z - m_hat
-
+      if (self$data$n_instr == 1) {
+        if (is.character(self$score)) {
+            if (self$score == "partialling out") {
+              psi_a = -w_hat * v_hat
+              psi_b = v_hat * u_hat
+            }
+          psis = list(
+            psi_a = psi_a,
+            psi_b = psi_b)
+        } else if (is.function(self$score)) {
+          psis = self$score(y, z, d, g_hat, m_hat, r_hat, smpls)
+        }
+      } else {
         stopifnot(self$apply_cross_fitting)
 
         # Projection: r_hat from projection on m_hat
@@ -348,29 +358,20 @@ DoubleMLPLIV = R6Class("DoubleMLPLIV",
         r_r_tilde = resample(task_r_tilde, ml_r_tilde, resampling_r_tilde,
           store_models = TRUE)
         r_hat_tilde = as.data.table(r_r_tilde$prediction())$response
-      }
+        
       if (is.character(self$score)) {
-        if (self$data$n_instr == 1) {
-          if (self$score == "partialling out") {
-            psi_a = -w_hat * v_hat
-            psi_b = v_hat * u_hat
-          }
-        } else {
           if (self$score == "partialling out") {
             psi_a = -w_hat * r_hat_tilde
             psi_b = r_hat_tilde * u_hat
           }
-        }
         psis = list(
           psi_a = psi_a,
           psi_b = psi_b)
       } else if (is.function(self$score)) {
-        if (self$data$n_instr > 1) {
-          stop(paste(
+        stop(paste(
             "Callable score not implemented for DoubleMLPLIV with",
             "partialX=TRUE and partialZ=FALSE with several instruments."))
-        }
-        psis = self$score(y, z, d, g_hat, m_hat, r_hat, smpls)
+      }
       }
       return(psis)
     },