Merge pull request #365 from DoubleML/jh-logistic-model

LPLR model

Author: SvenKlaassen

doubleml/__init__.py

@@ -13,6 +13,7 @@
 from .irm.pq import DoubleMLPQ
 from .irm.qte import DoubleMLQTE
 from .irm.ssm import DoubleMLSSM
+from .plm.lplr import DoubleMLLPLR
 from .plm.pliv import DoubleMLPLIV
 from .plm.plr import DoubleMLPLR
 from .utils.blp import DoubleMLBLP
@@ -42,6 +43,7 @@
     "DoubleMLBLP",
     "DoubleMLPolicyTree",
     "DoubleMLSSM",
+    "DoubleMLLPLR",
 ]
 
 __version__ = importlib.metadata.version("doubleml")

doubleml/did/did.py

@@ -70,7 +70,7 @@ class DoubleMLDID(LinearScoreMixin, DoubleML):
     >>> data = make_did_SZ2020(n_obs=500, return_type='DataFrame')
     >>> obj_dml_data = dml.DoubleMLDIDData(data, 'y', 'd')
     >>> dml_did_obj = dml.DoubleMLDID(obj_dml_data, ml_g, ml_m)
-    >>> dml_did_obj.fit().summary
+    >>> dml_did_obj.fit().summary  # doctest: +SKIP
            coef   std err         t     P>|t|     2.5 %    97.5 %
     d -2.840718  1.760386 -1.613691  0.106595 -6.291011  0.609575
 

doubleml/did/did_cs.py

@@ -59,7 +59,8 @@ class DoubleMLDIDCS(LinearScoreMixin, DoubleML):
         Default is ``True``.
 
     Examples
-    --------    >>> import numpy as np
+    --------
+    >>> import numpy as np
     >>> import doubleml as dml
     >>> from doubleml.did.datasets import make_did_SZ2020
     >>> from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
@@ -69,7 +70,7 @@ class DoubleMLDIDCS(LinearScoreMixin, DoubleML):
     >>> data = make_did_SZ2020(n_obs=500, cross_sectional_data=True, return_type='DataFrame')
     >>> obj_dml_data = dml.DoubleMLDIDData(data, 'y', 'd', t_col='t')
     >>> dml_did_obj = dml.DoubleMLDIDCS(obj_dml_data, ml_g, ml_m)
-    >>> dml_did_obj.fit().summary
+    >>> dml_did_obj.fit().summary  # doctest: +SKIP
          coef   std err         t     P>|t|      2.5 %    97.5 %
     d -4.9944  7.561785 -0.660479  0.508947 -19.815226  9.826426
     """

doubleml/did/did_multi.py

@@ -140,7 +140,7 @@ class DoubleMLDIDMulti:
     ...     gt_combinations="standard",
     ...     control_group="never_treated",
     ... )
-    >>> print(dml_did_obj.fit().summary)
+    >>> print(dml_did_obj.fit().summary)  # doctest: +SKIP
                                       coef   std err  ...     2.5 %    97.5 %
     ATT(2025-03,2025-01,2025-02) -0.797617  0.459617  ... -1.698450  0.103215
     ATT(2025-03,2025-02,2025-03)  0.270311  0.456453  ... -0.624320  1.164941

doubleml/double_ml.py

@@ -22,7 +22,7 @@
 class DoubleML(SampleSplittingMixin, ABC):
     """Double Machine Learning."""
 
-    def __init__(self, obj_dml_data, n_folds, n_rep, score, draw_sample_splitting):
+    def __init__(self, obj_dml_data, n_folds, n_rep, score, draw_sample_splitting, double_sample_splitting=False):
         # check and pick up obj_dml_data
         if not isinstance(obj_dml_data, DoubleMLBaseData):
             raise TypeError(
@@ -34,18 +34,10 @@ def __init__(self, obj_dml_data, n_folds, n_rep, score, draw_sample_splitting):
             if obj_dml_data.n_cluster_vars > 2:
                 raise NotImplementedError("Multi-way (n_ways > 2) clustering not yet implemented.")
             self._is_cluster_data = True
-        self._is_panel_data = False
-        if isinstance(obj_dml_data, DoubleMLPanelData):
-            self._is_panel_data = True
-        self._is_did_data = False
-        if isinstance(obj_dml_data, DoubleMLDIDData):
-            self._is_did_data = True
-        self._is_ssm_data = False
-        if isinstance(obj_dml_data, DoubleMLSSMData):
-            self._is_ssm_data = True
-        self._is_rdd_data = False
-        if isinstance(obj_dml_data, DoubleMLRDDData):
-            self._is_rdd_data = True
+        self._is_panel_data = isinstance(obj_dml_data, DoubleMLPanelData)
+        self._is_did_data = isinstance(obj_dml_data, DoubleMLDIDData)
+        self._is_ssm_data = isinstance(obj_dml_data, DoubleMLSSMData)
+        self._is_rdd_data = isinstance(obj_dml_data, DoubleMLRDDData)
 
         self._dml_data = obj_dml_data
         self._n_obs = self._dml_data.n_obs
@@ -108,6 +100,9 @@ def __init__(self, obj_dml_data, n_folds, n_rep, score, draw_sample_splitting):
         self._smpls = None
         self._smpls_cluster = None
         self._n_obs_sample_splitting = self.n_obs
+        self._double_sample_splitting = double_sample_splitting
+        if self._double_sample_splitting:
+            self._smpls_inner = None
         if draw_sample_splitting:
             self.draw_sample_splitting()
         self._score_dim = (self._dml_data.n_obs, self.n_rep, self._dml_data.n_coefs)
@@ -359,6 +354,21 @@ def smpls(self):
             raise ValueError(err_msg)
         return self._smpls
 
+    @property
+    def smpls_inner(self):
+        """
+        The partition used for cross-fitting.
+        """
+        if not self._double_sample_splitting:
+            raise ValueError("smpls_inner is only available for double sample splitting.")
+        if self._smpls_inner is None:
+            err_msg = (
+                "Sample splitting not specified. Either draw samples via .draw_sample splitting() "
+                + "or set external samples via .set_sample_splitting()."
+            )
+            raise ValueError(err_msg)
+        return self._smpls_inner
+
     @property
     def smpls_cluster(self):
         """
@@ -507,6 +517,18 @@ def summary(self):
     def __smpls(self):
         return self._smpls[self._i_rep]
 
+    @property
+    def __smpls__inner(self):
+        if not self._double_sample_splitting:
+            raise ValueError("smpls_inner is only available for double sample splitting.")
+        if self._smpls_inner is None:
+            err_msg = (
+                "Sample splitting not specified. Either draw samples via .draw_sample splitting() "
+                + "or set external samples via .set_sample_splitting()."
+            )
+            raise ValueError(err_msg)
+        return self._smpls_inner[self._i_rep]
+
     @property
     def __smpls_cluster(self):
         return self._smpls_cluster[self._i_rep]
@@ -1081,7 +1103,10 @@ def _initalize_fit(self, store_predictions, store_models):
 
     def _fit_nuisance_and_score_elements(self, n_jobs_cv, store_predictions, external_predictions, store_models):
         ext_prediction_dict = _set_external_predictions(
-            external_predictions, learners=self.params_names, treatment=self._dml_data.d_cols[self._i_treat], i_rep=self._i_rep
+            external_predictions,
+            learners=self.params_names,
+            treatment=self._dml_data.d_cols[self._i_treat],
+            i_rep=self._i_rep,
         )
 
         # ml estimation of nuisance models and computation of score elements
@@ -1230,7 +1255,7 @@ def evaluate_learners(self, learners=None, metric=_rmse):
         >>> def mae(y_true, y_pred):
         ...     subset = np.logical_not(np.isnan(y_true))
         ...     return mean_absolute_error(y_true[subset], y_pred[subset])
-        >>> dml_irm_obj.evaluate_learners(metric=mae)
+        >>> dml_irm_obj.evaluate_learners(metric=mae)  # doctest: +SKIP
         {'ml_g0': array([[0.88173585]]), 'ml_g1': array([[0.83854057]]), 'ml_m': array([[0.35871235]])}
         """
         # if no learners are provided try to evaluate all learners
@@ -1249,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/double_ml_sampling_mixins.py

@@ -1,7 +1,7 @@
 from abc import abstractmethod
 
 from doubleml.utils._checks import _check_sample_splitting
-from doubleml.utils.resampling import DoubleMLClusterResampling, DoubleMLResampling
+from doubleml.utils.resampling import DoubleMLClusterResampling, DoubleMLDoubleResampling, DoubleMLResampling
 
 
 class SampleSplittingMixin:
@@ -17,6 +17,8 @@ class SampleSplittingMixin:
     `sample splitting <https://docs.doubleml.org/stable/guide/resampling.html>`_ in the DoubleML user guide.
     """
 
+    _double_sample_splitting = False
+
     def draw_sample_splitting(self):
         """
         Draw sample splitting for DoubleML models.
@@ -29,6 +31,8 @@ def draw_sample_splitting(self):
         self : object
         """
         if self._is_cluster_data:
+            if self._double_sample_splitting:
+                raise ValueError("Cluster data not supported for double sample splitting.")
             obj_dml_resampling = DoubleMLClusterResampling(
                 n_folds=self._n_folds_per_cluster,
                 n_rep=self.n_rep,
@@ -38,10 +42,20 @@ def draw_sample_splitting(self):
             )
             self._smpls, self._smpls_cluster = obj_dml_resampling.split_samples()
         else:
-            obj_dml_resampling = DoubleMLResampling(
-                n_folds=self.n_folds, n_rep=self.n_rep, n_obs=self._n_obs_sample_splitting, stratify=self._strata
-            )
-            self._smpls = obj_dml_resampling.split_samples()
+            if self._double_sample_splitting:
+                obj_dml_resampling = DoubleMLDoubleResampling(
+                    n_folds=self.n_folds,
+                    n_folds_inner=self.n_folds_inner,
+                    n_rep=self.n_rep,
+                    n_obs=self._dml_data.n_obs,
+                    stratify=self._strata,
+                )
+                self._smpls, self._smpls_inner = obj_dml_resampling.split_samples()
+            else:
+                obj_dml_resampling = DoubleMLResampling(
+                    n_folds=self.n_folds, n_rep=self.n_rep, n_obs=self._n_obs_sample_splitting, stratify=self._strata
+                )
+                self._smpls = obj_dml_resampling.split_samples()
 
         return self
 
@@ -104,6 +118,9 @@ def set_sample_splitting(self, all_smpls, all_smpls_cluster=None):
         >>> dml_plr_obj.set_sample_splitting(smpls) # doctest: +ELLIPSIS
         <doubleml.plm.plr.DoubleMLPLR object at 0x...>
         """
+        if self._double_sample_splitting:
+            raise ValueError("set_sample_splitting not supported for double sample splitting.")
+
         self._smpls, self._smpls_cluster, self._n_rep, self._n_folds = _check_sample_splitting(
             all_smpls, all_smpls_cluster, self._dml_data, self._is_cluster_data, n_obs=self._n_obs_sample_splitting
         )

doubleml/double_ml_score_mixins.py

@@ -86,6 +86,7 @@ class NonLinearScoreMixin:
     _score_type = "nonlinear"
     _coef_start_val = np.nan
     _coef_bounds = None
+    _error_on_convergence_failure = False
 
     @property
     @abstractmethod
@@ -149,12 +150,16 @@ def score_deriv(theta):
             theta_hat = root_res.root
             if not root_res.converged:
                 score_val = score(theta_hat)
-                warnings.warn(
+                msg = (
                     "Could not find a root of the score function.\n "
                     f"Flag: {root_res.flag}.\n"
                     f"Score value found is {score_val} "
                     f"for parameter theta equal to {theta_hat}."
                 )
+                if self._error_on_convergence_failure:
+                    raise ValueError(msg)
+                else:
+                    warnings.warn(msg)
         else:
             signs_different, bracket_guess = _get_bracket_guess(score, self._coef_start_val, self._coef_bounds)
 
@@ -186,12 +191,16 @@ def score_squared(theta):
                             score, self._coef_start_val, approx_grad=True, bounds=[self._coef_bounds]
                         )
                         theta_hat = theta_hat_array.item()
-                        warnings.warn(
+                        msg = (
                             "Could not find a root of the score function.\n "
                             f"Minimum score value found is {score_val} "
                             f"for parameter theta equal to {theta_hat}.\n "
                             "No theta found such that the score function evaluates to a negative value."
                         )
+                        if self._error_on_convergence_failure:
+                            raise ValueError(msg)
+                        else:
+                            warnings.warn(msg)
                     else:
 
                         def neg_score(theta):
@@ -202,11 +211,15 @@ def neg_score(theta):
                             neg_score, self._coef_start_val, approx_grad=True, bounds=[self._coef_bounds]
                         )
                         theta_hat = theta_hat_array.item()
-                        warnings.warn(
+                        msg = (
                             "Could not find a root of the score function. "
                             f"Maximum score value found is {-1 * neg_score_val} "
                             f"for parameter theta equal to {theta_hat}. "
                             "No theta found such that the score function evaluates to a positive value."
                         )
+                        if self._error_on_convergence_failure:
+                            raise ValueError(msg)
+                        else:
+                            warnings.warn(msg)
 
         return theta_hat

doubleml/irm/cvar.py

@@ -97,7 +97,7 @@ class DoubleMLCVAR(LinearScoreMixin, DoubleML):
     >>> data = make_irm_data(theta=0.5, n_obs=500, dim_x=20, return_type='DataFrame')
     >>> obj_dml_data = dml.DoubleMLData(data, 'y', 'd')
     >>> dml_cvar_obj = dml.DoubleMLCVAR(obj_dml_data, ml_g, ml_m, treatment=1, quantile=0.5)
-    >>> dml_cvar_obj.fit().summary
+    >>> dml_cvar_obj.fit().summary  # doctest: +SKIP
            coef   std err         t         P>|t|     2.5 %    97.5 %
     d  1.588364  0.096616  16.43989  9.909942e-61  1.398999  1.777728
 

doubleml/irm/iivm.py

@@ -95,7 +95,7 @@ class DoubleMLIIVM(LinearScoreMixin, DoubleML):
     >>> data = make_iivm_data(theta=0.5, n_obs=1000, dim_x=20, alpha_x=1.0, return_type='DataFrame')
     >>> obj_dml_data = dml.DoubleMLData(data, 'y', 'd', z_cols='z')
     >>> dml_iivm_obj = dml.DoubleMLIIVM(obj_dml_data, ml_g, ml_m, ml_r)
-    >>> dml_iivm_obj.fit().summary
+    >>> dml_iivm_obj.fit().summary  # doctest: +SKIP
            coef   std err         t     P>|t|     2.5 %    97.5 %
     d  0.362398  0.191578  1.891649  0.058538 -0.013088  0.737884
 

doubleml/irm/irm.py

@@ -96,7 +96,7 @@ class DoubleMLIRM(LinearScoreMixin, DoubleML):
     >>> data = make_irm_data(theta=0.5, n_obs=500, dim_x=20, return_type='DataFrame')
     >>> obj_dml_data = dml.DoubleMLData(data, 'y', 'd')
     >>> dml_irm_obj = dml.DoubleMLIRM(obj_dml_data, ml_g, ml_m)
-    >>> dml_irm_obj.fit().summary
+    >>> dml_irm_obj.fit().summary  # doctest: +SKIP
            coef   std err         t     P>|t|     2.5 %    97.5 %
     d  0.371972  0.206802  1.798685  0.072069 -0.033353  0.777297