add return_type tests for pred, targets and rmses

Author: SvenKlaassen

.gitignore

@@ -25,5 +25,6 @@ share/python-wheels/
 *.egg-info/
 .installed.cfg
 *.egg
+*.vscode
 MANIFEST
 *.idea

doubleml/double_ml.py

@@ -47,7 +47,7 @@ def __init__(self,
 
         # initialize predictions and target to None which are only stored if method fit is called with store_predictions=True
         self._predictions = None
-        self._nuisance_target = None
+        self._nuisance_targets = None
         self._rmses = None
 
         # initialize models to None which are only stored if method fit is called with store_models=True
@@ -231,11 +231,11 @@ def predictions(self):
         return self._predictions
 
     @property
-    def nuisance_target(self):
+    def nuisance_targets(self):
         """
         The outcome of the nuisance models.
         """
-        return self._nuisance_target
+        return self._nuisance_targets
 
     @property
     def rmses(self):
@@ -249,7 +249,7 @@ def models(self):
         """
         The fitted nuisance models.
         """
-        return self._model
+        return self._models
 
     def get_params(self, learner):
         """
@@ -1010,8 +1010,8 @@ def _initialize_boot_arrays(self, n_rep_boot):
     def _initialize_predictions_and_targets(self):
         self._predictions = {learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan)
                              for learner in self.params_names}
-        self._nuisance_target = {learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan)
-                                 for learner in self.params_names}
+        self._nuisance_targets = {learner: np.full((self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs), np.nan) 
+                                  for learner in self.params_names}
 
     def _initialize_rmses(self):
         self._rmses = {learner: np.full((self.n_rep, self._dml_data.n_coefs), np.nan)
@@ -1024,7 +1024,7 @@ def _initialize_models(self):
     def _store_predictions_and_targets(self, preds, targets):
         for learner in self.params_names:
             self._predictions[learner][:, self._i_rep, self._i_treat] = preds[learner]
-            self._nuisance_target[learner][:, self._i_rep, self._i_treat] = targets[learner]
+            self._nuisance_targets[learner][:, self._i_rep, self._i_treat] = targets[learner]
 
     def _calc_rmses(self, preds, targets):
         for learner in self.params_names:

doubleml/tests/test_doubleml_return_types.py

@@ -165,3 +165,63 @@ def test_stored_models():
 
     assert np.all([isinstance(mdl, plr_dml1.learner['ml_l'].__class__) for mdl in plr_dml1.models['ml_l']['d'][0]])
     assert np.all([isinstance(mdl, plr_dml1.learner['ml_m'].__class__) for mdl in plr_dml1.models['ml_m']['d'][0]])
+
+
+@pytest.mark.ci
+def test_stored_predictions():
+    assert plr_dml1.predictions['ml_l'].shape == (n_obs, n_rep, n_treat)
+    assert plr_dml1.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
+
+    assert pliv_dml1.predictions['ml_l'].shape == (n_obs, n_rep, n_treat)
+    assert pliv_dml1.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
+    assert pliv_dml1.predictions['ml_r'].shape == (n_obs, n_rep, n_treat)
+
+    assert irm_dml1.predictions['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert irm_dml1.predictions['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert irm_dml1.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
+
+    assert iivm_dml1.predictions['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.predictions['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.predictions['ml_r0'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.predictions['ml_r1'].shape == (n_obs, n_rep, n_treat)
+
+
+@pytest.mark.ci
+def test_stored_nuisance_targets():
+    assert plr_dml1.nuisance_targets['ml_l'].shape == (n_obs, n_rep, n_treat)
+    assert plr_dml1.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
+
+    assert pliv_dml1.nuisance_targets['ml_l'].shape == (n_obs, n_rep, n_treat)
+    assert pliv_dml1.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
+    assert pliv_dml1.nuisance_targets['ml_r'].shape == (n_obs, n_rep, n_treat)
+
+    assert irm_dml1.nuisance_targets['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert irm_dml1.nuisance_targets['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert irm_dml1.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
+
+    assert iivm_dml1.nuisance_targets['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.nuisance_targets['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.nuisance_targets['ml_r0'].shape == (n_obs, n_rep, n_treat)
+    assert iivm_dml1.nuisance_targets['ml_r1'].shape == (n_obs, n_rep, n_treat)
+
+
+@pytest.mark.ci
+def test_rmses():
+    assert plr_dml1.rmses['ml_l'].shape == (n_rep, n_treat)
+    assert plr_dml1.rmses['ml_m'].shape == (n_rep, n_treat)
+
+    assert pliv_dml1.rmses['ml_l'].shape == (n_rep, n_treat)
+    assert pliv_dml1.rmses['ml_m'].shape == (n_rep, n_treat)
+    assert pliv_dml1.rmses['ml_r'].shape == (n_rep, n_treat)
+
+    assert irm_dml1.rmses['ml_g0'].shape == (n_rep, n_treat)
+    assert irm_dml1.rmses['ml_g1'].shape == (n_rep, n_treat)
+    assert irm_dml1.rmses['ml_m'].shape == (n_rep, n_treat)
+
+    assert iivm_dml1.rmses['ml_g0'].shape == (n_rep, n_treat)
+    assert iivm_dml1.rmses['ml_g1'].shape == (n_rep, n_treat)
+    assert iivm_dml1.rmses['ml_m'].shape == (n_rep, n_treat)
+    assert iivm_dml1.rmses['ml_r0'].shape == (n_rep, n_treat)
+    assert iivm_dml1.rmses['ml_r1'].shape == (n_rep, n_treat)