Merge pull request #378 from DoubleML/j-static-panel

PLPR Static Panel model

Author: SvenKlaassen

doubleml/__init__.py

@@ -13,6 +13,7 @@
 from .irm.ssm import DoubleMLSSM
 from .plm.lplr import DoubleMLLPLR
 from .plm.pliv import DoubleMLPLIV
+from .plm.plpr import DoubleMLPLPR
 from .plm.plr import DoubleMLPLR
 from .utils.blp import DoubleMLBLP
 from .utils.policytree import DoubleMLPolicyTree
@@ -43,6 +44,7 @@
     "DoubleMLPolicyTree",
     "DoubleMLSSM",
     "DoubleMLLPLR",
+    "DoubleMLPLPR",
 ]
 
 try:

doubleml/data/base_data.py

@@ -199,7 +199,7 @@ def _data_summary_str(self):
 
         if hasattr(self, "is_cluster_data") and self.is_cluster_data:
             data_summary += f"Is cluster data: {self.is_cluster_data}\n"
-        data_summary += f"No. Observations: {self.n_obs}\n"
+        data_summary += f"No. Observations: {self.n_obs}"
         return data_summary
 
     @classmethod

doubleml/data/panel_data.py

@@ -41,6 +41,13 @@ class DoubleMLPanelData(DoubleMLData):
         The instrumental variable(s).
         Default is ``None``.
 
+    static_panel : bool
+        Indicates whether the data model corresponds to a static
+        panel data approach (``True``) or to staggered adoption panel data
+        (``False``). In the latter case, the treatment groups/values are defined in terms of the first time of
+        treatment exposure.
+        Default is ``False``.
+
     use_other_treat_as_covariate : bool
         Indicates whether in the multiple-treatment case the other treatment variables should be added as covariates.
         Default is ``True``.
@@ -82,31 +89,40 @@ def __init__(
         id_col,
         x_cols=None,
         z_cols=None,
+        static_panel=False,
         use_other_treat_as_covariate=True,
         force_all_x_finite=True,
         datetime_unit="M",
     ):
         DoubleMLBaseData.__init__(self, data)
 
+        self._static_panel = static_panel
+
         # we need to set id_col (needs _data) before call to the super __init__ because of the x_cols setter
         self.id_col = id_col
-        self._datetime_unit = _is_valid_datetime_unit(datetime_unit)
         self._set_id_var()
-
         # Set time column before calling parent constructor
         self.t_col = t_col
+        self._datetime_unit = _is_valid_datetime_unit(datetime_unit)
+
+        if not self.static_panel:
+            cluster_cols = None
+            force_all_d_finite = False
+        else:
+            cluster_cols = id_col
+            force_all_d_finite = True
 
-        # Call parent constructor
         DoubleMLData.__init__(
             self,
             data=data,
             y_col=y_col,
             d_cols=d_cols,
             x_cols=x_cols,
             z_cols=z_cols,
+            cluster_cols=cluster_cols,
             use_other_treat_as_covariate=use_other_treat_as_covariate,
             force_all_x_finite=force_all_x_finite,
-            force_all_d_finite=False,
+            force_all_d_finite=force_all_d_finite,
         )
 
         # reset index to ensure a simple RangeIndex
@@ -115,15 +131,15 @@ def __init__(
         # Set time variable array after data is loaded
         self._set_time_var()
 
-        if self.n_treat != 1:
-            raise ValueError("Only one treatment column is allowed for panel data.")
-
         self._check_disjoint_sets_id_col()
 
         # intialize the unique values of g and t
         self._g_values = np.sort(np.unique(self.d))  # unique values of g
         self._t_values = np.sort(np.unique(self.t))  # unique values of t
 
+        if self.n_treat != 1:
+            raise ValueError("Only one treatment column is allowed for panel data.")
+
     def __str__(self):
         data_summary = self._data_summary_str()
         buf = io.StringIO()
@@ -146,9 +162,10 @@ def _data_summary_str(self):
             f"Instrument variable(s): {self.z_cols}\n"
             f"Time variable: {self.t_col}\n"
             f"Id variable: {self.id_col}\n"
+            f"Static panel data: {self.static_panel}\n"
         )
 
-        data_summary += f"No. Unique Ids: {self.n_ids}\n" f"No. Observations: {self.n_obs}\n"
+        data_summary += f"No. Unique Ids: {self.n_ids}\n" f"No. Observations: {self.n_obs}"
         return data_summary
 
     @classmethod
@@ -296,6 +313,11 @@ def n_t_periods(self):
         """
         return len(self.t_values)
 
+    @property
+    def static_panel(self):
+        """Indicates whether the data model corresponds to a static panel data approach."""
+        return self._static_panel
+
     def _get_optional_col_sets(self):
         base_optional_col_sets = super()._get_optional_col_sets()
         id_col_set = {self.id_col}
@@ -370,4 +392,9 @@ def _set_id_var(self):
     def _set_time_var(self):
         """Set the time variable array."""
         if hasattr(self, "_data") and self.t_col in self.data.columns:
-            self._t = self.data.loc[:, self.t_col]
+            t_values = self.data.loc[:, self.t_col]
+            expected_dtypes = (np.integer, np.floating, np.datetime64)
+            if not any(np.issubdtype(t_values.dtype, dt) for dt in expected_dtypes):
+                raise ValueError(f"Invalid data type for time variable: expected one of {expected_dtypes}.")
+            else:
+                self._t = t_values

doubleml/data/tests/test_panel_data.py

@@ -157,14 +157,26 @@ def test_panel_data_str():
     assert "Time variable: t" in dml_str
     assert "Id variable: id" in dml_str
     assert "No. Observations:" in dml_str
+    assert "Static panel data:" in dml_str
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def static_panel(request):
+    return request.param
 
 
 @pytest.mark.ci
-def test_panel_data_properties():
+def test_panel_data_properties(static_panel):
     np.random.seed(3141)
     df = make_did_SZ2020(n_obs=100, return_type="DoubleMLPanelData")._data
     dml_data = DoubleMLPanelData(
-        data=df, y_col="y", d_cols="d", t_col="t", id_col="id", x_cols=[f"Z{i + 1}" for i in np.arange(4)]
+        data=df,
+        y_col="y",
+        d_cols="d",
+        t_col="t",
+        id_col="id",
+        x_cols=[f"Z{i + 1}" for i in np.arange(4)],
+        static_panel=static_panel,
     )
 
     assert np.array_equal(dml_data.id_var, df["id"].values)
@@ -176,3 +188,10 @@ def test_panel_data_properties():
     assert dml_data.n_groups == len(np.unique(df["d"].values))
     assert np.array_equal(dml_data.t_values, np.sort(np.unique(df["t"].values)))
     assert dml_data.n_t_periods == len(np.unique(df["t"].values))
+
+    if static_panel:
+        assert dml_data.static_panel is True
+        assert dml_data.cluster_cols == ["id"]
+    else:
+        assert dml_data.static_panel is False
+        assert dml_data.cluster_cols is None

doubleml/data/tests/test_panel_data_exceptions.py

@@ -61,6 +61,19 @@ def test_time_col_none_exception(sample_data):
         )
 
 
+@pytest.mark.ci
+def test_time_var_data_type_exception(sample_data):
+    # Test exception when time var is not int, float or datetime
+    msg = (
+        r"Invalid data type for time variable: expected one of \(<class 'numpy.integer'>, <class 'numpy.floating'>, "
+        r"<class 'numpy.datetime64'>\)"
+    )
+    with pytest.raises(ValueError, match=msg):
+        data_time_type = sample_data.copy()
+        data_time_type["time"] = data_time_type["time"].astype(str)
+        DoubleMLPanelData(data=data_time_type, y_col="y", d_cols="treatment", t_col="time", id_col="id")
+
+
 @pytest.mark.ci
 def test_overlapping_variables_exception(sample_data):
     # Test exception when id_col overlaps with another variable
@@ -80,9 +93,10 @@ def test_overlapping_variables_exception(sample_data):
         DoubleMLPanelData(data=sample_data, y_col="y", d_cols="id", t_col="time", id_col="id")  # Using id as treatment
 
     # Test time variable overlapping
+    # using t_col="id", id_col="id" gives invalid data type for time variable exception first
     msg = r"At least one variable/column is set as time variable \(``t_col``\) and identifier variable \(``id_col``\)."
     with pytest.raises(ValueError, match=msg):
-        DoubleMLPanelData(data=sample_data, y_col="y", d_cols="treatment", t_col="id", id_col="id")  # Using id as time
+        DoubleMLPanelData(data=sample_data, y_col="y", d_cols="treatment", t_col="time", id_col="time")  # Using time as id
 
 
 @pytest.mark.ci

doubleml/double_ml.py

@@ -192,7 +192,7 @@ def __str__(self):
 
         additional_info = self._format_additional_info_str()
         if additional_info:
-            representation += f"\n\n------------------ Additional Information ------------------\n" f"{additional_info}"
+            representation += f"\n\n------------------ Additional Information -------------\n" f"{additional_info}"
         return representation
 
     @property
@@ -1177,7 +1177,7 @@ def _initialize_ml_nuisance_params(self):
         pass
 
     @abstractmethod
-    def _nuisance_est(self, smpls, n_jobs_cv, return_models, external_predictions):
+    def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models):
         pass
 
     @abstractmethod

doubleml/plm/__init__.py

@@ -4,6 +4,7 @@
 
 from .lplr import DoubleMLLPLR
 from .pliv import DoubleMLPLIV
+from .plpr import DoubleMLPLPR
 from .plr import DoubleMLPLR
 
-__all__ = ["DoubleMLPLR", "DoubleMLPLIV", "DoubleMLLPLR"]
+__all__ = ["DoubleMLPLR", "DoubleMLPLIV", "DoubleMLLPLR", "DoubleMLPLPR"]

doubleml/plm/datasets/__init__.py

@@ -7,6 +7,7 @@
 from .dgp_lplr_LZZ2020 import make_lplr_LZZ2020
 from .dgp_pliv_CHS2015 import make_pliv_CHS2015
 from .dgp_pliv_multiway_cluster_CKMS2021 import make_pliv_multiway_cluster_CKMS2021
+from .dgp_plpr_CP2025 import make_plpr_CP2025
 from .dgp_plr_CCDDHNR2018 import make_plr_CCDDHNR2018
 from .dgp_plr_turrell2018 import make_plr_turrell2018
 
@@ -18,4 +19,5 @@
     "make_pliv_multiway_cluster_CKMS2021",
     "make_lplr_LZZ2020",
     "_make_pliv_data",
+    "make_plpr_CP2025",
 ]