Change model name to DoubleMLSSM and remove line redefining pi_hat

doubleml/__init__.py

@@ -13,7 +13,7 @@
 from .double_ml_lpq import DoubleMLLPQ
 from .double_ml_cvar import DoubleMLCVAR
 from .double_ml_policytree import DoubleMLPolicyTree
-from .double_ml_selection import DoubleMLS
+from .double_ml_selection import DoubleMLSSM
 
 __all__ = ['DoubleMLPLR',
            'DoubleMLPLIV',
@@ -29,6 +29,6 @@
            'DoubleMLLPQ',
            'DoubleMLCVAR',
            'DoubleMLPolicyTree',
-           'DoubleMLS']
+           'DoubleMLSSM']
 
 __version__ = get_distribution('doubleml').version

doubleml/double_ml_selection.py

@@ -18,7 +18,7 @@
 from .double_ml_score_mixins import LinearScoreMixin
 
 
-class DoubleMLS(LinearScoreMixin, DoubleML):
+class DoubleMLSSM(LinearScoreMixin, DoubleML):
     """Double machine learning for sample selection models
 
     Parameters
@@ -279,13 +279,12 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             y_1 = y
 
         else:
-            # initialize nuisance predictions, targets and models
+            # initialize nuisance predictions, targets and models -- will be overwritten in each iteration
             mu_hat_d1 = {'models': None,
                          'targets': np.full(shape=self._dml_data.n_obs, fill_value=np.nan),
                          'preds': np.full(shape=self._dml_data.n_obs, fill_value=np.nan)
                          }
             mu_hat_d0 = copy.deepcopy(mu_hat_d1)
-            pi_hat = copy.deepcopy(mu_hat_d1)
             pi_hat_d1 = copy.deepcopy(mu_hat_d1)
             pi_hat_d0 = copy.deepcopy(mu_hat_d1)
             p_hat_d1 = copy.deepcopy(mu_hat_d1)
@@ -295,6 +294,7 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             strata = self._dml_data.d.reshape(-1, 1) + 2 * self._dml_data.t.reshape(-1, 1)
 
             # initialize nuisance predictions, targets and models
+            # pi_hat - used for preliminary estimation of propensity score pi, overwritten in each iteration
             pi_hat = {'models': None,
                       'targets': [],
                       'preds': []

doubleml/tests/test_selection.py

@@ -61,19 +61,19 @@ def dml_selection_fixture(generate_data_selection, learner, score, dml_procedure
 
     if score == 'mar':
         obj_dml_data = dml.DoubleMLData.from_arrays(x, y, d, z=None, t=s)
-        dml_sel_obj = dml.DoubleMLS(obj_dml_data,
-                                    ml_mu, ml_pi, ml_p,
-                                    n_folds=n_folds,
-                                    score=score,
-                                    dml_procedure=dml_procedure)
+        dml_sel_obj = dml.DoubleMLSSM(obj_dml_data,
+                                      ml_mu, ml_pi, ml_p,
+                                      n_folds=n_folds,
+                                      score=score,
+                                      dml_procedure=dml_procedure)
     else:
         assert score == 'nonignorable'
         obj_dml_data = dml.DoubleMLData.from_arrays(x, y, d, z=z, t=s)
-        dml_sel_obj = dml.DoubleMLS(obj_dml_data,
-                                    ml_mu, ml_pi, ml_p,
-                                    n_folds=n_folds,
-                                    score=score,
-                                    dml_procedure=dml_procedure)
+        dml_sel_obj = dml.DoubleMLSSM(obj_dml_data,
+                                      ml_mu, ml_pi, ml_p,
+                                      n_folds=n_folds,
+                                      score=score,
+                                      dml_procedure=dml_procedure)
 
     n_obs = len(y)
     all_smpls = draw_smpls(n_obs, n_folds)