Enhance learner evaluation checks and handle NaN targets in DoubleML class

Author: SvenKlaassen

doubleml/double_ml.py

@@ -1274,12 +1274,19 @@ def evaluate_learners(self, learners=None, metric=_rmse):
             for learner in learners:
                 for rep in range(self.n_rep):
                     for coef_idx in range(self._dml_data.n_coefs):
-                        res = metric(
-                            y_pred=self.predictions[learner][:, rep, coef_idx].reshape(1, -1),
-                            y_true=self.nuisance_targets[learner][:, rep, coef_idx].reshape(1, -1),
-                        )
-                        if not np.isfinite(res):
-                            raise ValueError(f"Evaluation from learner {str(learner)} is not finite.")
+                        targets = self.nuisance_targets[learner][:, rep, coef_idx].reshape(1, -1)
+
+                        if np.all(np.isnan(targets)):
+                            res = np.nan
+                        else:
+                            predictions = self.predictions[learner][:, rep, coef_idx].reshape(1, -1)
+                            res = metric(
+                                y_pred=predictions,
+                                y_true=targets,
+                            )
+                            if not np.isfinite(res):
+                                raise ValueError(f"Evaluation from learner {str(learner)} is not finite.")
+
                         dist[learner][rep, coef_idx] = res
             return dist
         else:

doubleml/utils/_check_return_types.py

@@ -113,6 +113,12 @@ def check_basic_predictions_and_targets(dml_obj, n_obs, n_treat, n_rep):
         assert isinstance(dml_obj.nuisance_loss[key], np.ndarray)
         assert dml_obj.nuisance_loss[key].shape == (n_rep, n_treat)
 
+    learner_eval = dml_obj.evaluate_learners()
+    assert isinstance(learner_eval, dict)
+    for key in expected_keys:
+        assert key in learner_eval
+        assert isinstance(learner_eval[key], np.ndarray)
+        assert learner_eval[key].shape == (n_rep, n_treat)
     return