refactor BLP, CATE and GATE

Author: SvenKlaassen

.coverage

@@ -5,14 +5,14 @@
 from .double_ml_irm import DoubleMLIRM
 from .double_ml_iivm import DoubleMLIIVM
 from .double_ml_data import DoubleMLData, DoubleMLClusterData
-from .double_ml_blp import DoubleMLIRMBLP
+from .double_ml_blp import DoubleMLBLP
 
 __all__ = ['DoubleMLPLR',
            'DoubleMLPLIV',
            'DoubleMLIRM',
            'DoubleMLIIVM',
            'DoubleMLData',
            'DoubleMLClusterData',
-           'DoubleMLIRMBLP']
+           'DoubleMLBLP']
 
 __version__ = get_distribution('doubleml').version

doubleml/__init__.py

@@ -6,8 +6,8 @@
 from scipy.linalg import sqrtm
 
 
-class DoubleMLIRMBLP:
-    """Best linear predictor for DoubleML IRM models
+class DoubleMLBLP:
+    """Best linear predictor (BLP) for DoubleML with orthogonal signals.
 
         Parameters
         ----------
@@ -17,11 +17,16 @@ class DoubleMLIRMBLP:
         basis : :class:`pandas.DataFrame`
             The basis for estimating the best linear predictor. Has to have the shape (n,d),
             where d is the number of predictors.
+
+        is_gate : bool
+            Indicates whether the basis is constructed for GATEs (dummy-basis).
+            Default is ``False``.
     """
 
     def __init__(self,
                  orth_signal,
-                 basis):
+                 basis,
+                 is_gate=False):
 
         if not isinstance(orth_signal, np.ndarray):
             raise TypeError('The signal must be of np.ndarray type. '
@@ -41,6 +46,7 @@ def __init__(self,
 
         self._orth_signal = orth_signal
         self._basis = basis
+        self._is_gate = is_gate
 
         # initialize the score and the covariance
         self._blp_model = None
@@ -89,15 +95,16 @@ def fit(self):
 
         return self
 
-    def confint(self, basis, joint=False, level=0.95, n_rep_boot=500):
+    def confint(self, basis=None, joint=False, level=0.95, n_rep_boot=500):
         """
-        Confidence intervals for BLP for DoubleML IRM.
+        Confidence intervals for the BLP model.
 
         Parameters
         ----------
         basis : :class:`pandas.DataFrame`
             The basis for constructing the confidence interval. Has to have the same form as the basis from
-            the construction.
+            the construction. If ``None`` the basis for the construction of the model is used.
+            Default is ``None``
 
         joint : bool
             Indicates whether joint confidence intervals are computed.
@@ -138,6 +145,20 @@ def confint(self, basis, joint=False, level=0.95, n_rep_boot=500):
             raise ValueError('Apply fit() before confint().')
 
         alpha = 1 - level
+        gate_names = None
+        # define basis if none is supplied
+        if basis is None:
+            if self._is_gate:
+                # reduce to unique groups
+                basis = pd.DataFrame(np.diag(v=np.full((self._basis.shape[1]), True)))
+                gate_names = list(self._basis.columns.values)
+            else:
+                basis = self._basis
+        elif not (basis.shape[1] == self._basis.shape[1]):
+            raise ValueError('Invalid basis: DataFrame has to have the exact same number and ordering of columns.')
+        elif not list(basis.columns.values) == list(self._basis.columns.values):
+            raise ValueError('Invalid basis: DataFrame has to have the exact same number and ordering of columns.')
+
         # blp of the orthogonal signal
         g_hat = self._blp_model.predict(basis)
 
@@ -167,4 +188,8 @@ def confint(self, basis, joint=False, level=0.95, n_rep_boot=500):
         df_ci = pd.DataFrame(ci,
                              columns=['{:.1f} %'.format(alpha/2 * 100), 'effect', '{:.1f} %'.format((1-alpha/2) * 100)],
                              index=basis.index)
+
+        if self._is_gate and gate_names is not None:
+            df_ci.index = gate_names
+
         return df_ci

doubleml/double_ml_blp.py

@@ -1,11 +1,11 @@
 import numpy as np
 import pandas as pd
+import warnings
 from sklearn.utils import check_X_y
 from sklearn.utils.multiclass import type_of_target
 
 from .double_ml import DoubleML
-from .double_ml_blp import DoubleMLIRMBLP
-
+from .double_ml_blp import DoubleMLBLP
 from ._utils import _dml_cv_predict, _get_cond_smpls, _dml_tune, _check_finite_predictions
 
 
@@ -334,43 +334,40 @@ def cate(self, basis):
 
         Returns
         -------
-        model : :class:`doubleML.DoubleMLIRMBLP`
+        model : :class:`doubleML.DoubleMLBLP`
             Best linear Predictor model.
         """
+        valid_score = ['ATE']
+        if self.score not in valid_score:
+            raise ValueError('Invalid score ' + self.score + '. ' +
+                             'Valid score ' + ' or '.join(valid_score) + '.')
+
         # define the orthogonal signal
         orth_signal = self.psi_b.reshape(-1)
         # fit the best linear predictor
-        model = DoubleMLIRMBLP(orth_signal, basis=basis).fit()
+        model = DoubleMLBLP(orth_signal, basis=basis).fit()
 
         return model
 
-    def gate(self, groups, joint=False, level=0.95, n_rep_boot=500):
+    def gate(self, groups):
         """
-        Calculate group average treatment effects (GATE) for a given basis.
+        Calculate group average treatment effects (GATE) for mutually exclusive groups.
 
         Parameters
         ----------
         groups : :class:`pandas.DataFrame`
-            The group indicator for estimating the best linear predictor. Has to have the shape (n,d),
+            The group indicator for estimating the best linear predictor. Has to be dummy coded with shape (n,d),
             where d is the number of groups or (n,1) and contain the corresponding groups.
 
-        joint : bool
-            Indicates whether joint confidence intervals are computed.
-            Default is ``False``
-
-        level : float
-            The confidence level.
-            Default is ``0.95``.
-
-        n_rep_boot : int
-            Number of bootstrap samples for joint confidence interval.
-            Default is ``500``.
-
         Returns
         -------
-        df_ci : pd.DataFrame
-            A data frame with the confidence interval(s).
+        model : :class:`doubleML.DoubleMLBLPGATE`
+            Best linear Predictor model for Group Effects.
         """
+        valid_score = ['ATE']
+        if self.score not in valid_score:
+            raise ValueError('Invalid score ' + self.score + '. ' +
+                             'Valid score ' + ' or '.join(valid_score) + '.')
 
         if not isinstance(groups, pd.DataFrame):
             raise TypeError('Groups must be of DataFrame type. '
@@ -380,17 +377,15 @@ def gate(self, groups, joint=False, level=0.95, n_rep_boot=500):
             if groups.shape[1] == 1:
                 groups = pd.get_dummies(groups, prefix='Group', prefix_sep='_')
             else:
-                raise TypeError('Columns must be of of bool or int type or the data frame only should contain '
-                                'one column.')
+                raise TypeError('Columns of groups must be of bool type or int type (dummy coded). '
+                                'Alternatively, groups should only contain one column.')
+
+        if any(groups.sum(0) <= 5):
+            warnings.warn('At least one group effect is estimated with less than 6 observations.')
 
         # define the orthogonal signal
         orth_signal = self.psi_b.reshape(-1)
-        # fit the best linear predictor
-        model = DoubleMLIRMBLP(orth_signal, basis=groups).fit()
+        # fit the best linear predictor for GATE (different confint() method)
+        model = DoubleMLBLP(orth_signal, basis=groups, is_gate=True).fit()
 
-        # reduce to unique groups and create confidence interval
-        unique_groups = pd.DataFrame(np.diag(v=np.full((groups.shape[1]), True)))
-        df_ci = model.confint(unique_groups, joint, level, n_rep_boot)
-        df_ci.index = groups.columns.values
-
-        return df_ci
+        return model

doubleml/double_ml_irm.py

@@ -26,11 +26,13 @@ def dml_blp_fixture(ci_joint, ci_level):
     random_basis = pd.DataFrame(np.random.normal(0, 1, size=(n, 3)))
     random_signal = np.random.normal(0, 1, size=(n, ))
 
-    blp = dml.DoubleMLIRMBLP(random_signal, random_basis).fit()
+    blp = dml.DoubleMLBLP(random_signal, random_basis).fit()
     blp_manual = fit_blp(random_signal, random_basis)
 
     np.random.seed(42)
-    ci = blp.confint(random_basis, joint=ci_joint, level=ci_level, n_rep_boot=1000)
+    ci_1 = blp.confint(random_basis, joint=ci_joint, level=ci_level, n_rep_boot=1000)
+    np.random.seed(42)
+    ci_2 = blp.confint(joint=ci_joint, level=ci_level, n_rep_boot=1000)
     np.random.seed(42)
     ci_manual = blp_confint(blp_manual, random_basis, joint=ci_joint, level=ci_level, n_rep_boot=1000)
 
@@ -40,7 +42,8 @@ def dml_blp_fixture(ci_joint, ci_level):
                 'values_manual':  blp_manual.fittedvalues,
                 'omega': blp.blp_omega,
                 'omega_manual': blp_manual.cov_HC0,
-                'ci': ci,
+                'ci_1': ci_1,
+                'ci_2': ci_2,
                 'ci_manual': ci_manual}
 
     return res_dict
@@ -68,7 +71,14 @@ def test_dml_blp_omega(dml_blp_fixture):
 
 
 @pytest.mark.ci
-def test_dml_blp_ci(dml_blp_fixture):
-    assert np.allclose(dml_blp_fixture['ci'],
+def test_dml_blp_ci_1(dml_blp_fixture):
+    assert np.allclose(dml_blp_fixture['ci_1'],
+                       dml_blp_fixture['ci_2'],
+                       rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_blp_ci_2(dml_blp_fixture):
+    assert np.allclose(dml_blp_fixture['ci_1'],
                        dml_blp_fixture['ci_manual'],
                        rtol=1e-9, atol=1e-4)

doubleml/tests/test_blp.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import numpy as np
 
-from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV, DoubleMLData, DoubleMLClusterData, DoubleMLIRMBLP
+from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV, DoubleMLData, DoubleMLClusterData, DoubleMLBLP
 from doubleml.datasets import make_plr_CCDDHNR2018, make_irm_data, make_pliv_CHS2015, make_iivm_data, \
     make_pliv_multiway_cluster_CKMS2021
 
@@ -664,19 +664,24 @@ def test_doubleml_exception_blp():
 
     msg = "The signal must be of np.ndarray type. Signal of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
-        DoubleMLIRMBLP(orth_signal=1, basis=random_basis)
+        DoubleMLBLP(orth_signal=1, basis=random_basis)
     msg = 'The signal must be of one dimensional. Signal of dimensions 2 was passed.'
     with pytest.raises(ValueError, match=msg):
-        DoubleMLIRMBLP(orth_signal=np.array([[1], [2]]), basis=random_basis)
+        DoubleMLBLP(orth_signal=np.array([[1], [2]]), basis=random_basis)
     msg = "The basis must be of DataFrame type. Basis of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
-        DoubleMLIRMBLP(orth_signal=signal, basis=1)
+        DoubleMLBLP(orth_signal=signal, basis=1)
     msg = 'Invalid pd.DataFrame: Contains duplicate column names.'
     with pytest.raises(ValueError, match=msg):
-        DoubleMLIRMBLP(orth_signal=signal, basis=pd.DataFrame(np.array([[1, 2], [4, 5]]),
-                                                              columns=['x_1', 'x_1']))
+        DoubleMLBLP(orth_signal=signal, basis=pd.DataFrame(np.array([[1, 2], [4, 5]]),
+                                                           columns=['a_1', 'a_1']))
 
-    dml_blp_confint = DoubleMLIRMBLP(orth_signal=signal, basis=random_basis)
+    dml_blp_confint = DoubleMLBLP(orth_signal=signal, basis=random_basis)
+    msg = r'Apply fit\(\) before confint\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_blp_confint.confint(random_basis)
+
+    dml_blp_confint.fit()
     msg = 'joint must be True or False. Got 1.'
     with pytest.raises(TypeError, match=msg):
         dml_blp_confint.confint(random_basis, joint=1)
@@ -692,9 +697,12 @@ def test_doubleml_exception_blp():
     msg = 'The number of bootstrap replications must be positive. 0 was passed.'
     with pytest.raises(ValueError, match=msg):
         dml_blp_confint.confint(random_basis, n_rep_boot=0)
-    msg = r'Apply fit\(\) before confint\(\).'
+    msg = 'Invalid basis: DataFrame has to have the exact same number and ordering of columns.'
     with pytest.raises(ValueError, match=msg):
-        dml_blp_confint.confint(random_basis)
+        dml_blp_confint.confint(basis=pd.DataFrame(np.array([[1], [4]]), columns=['a_1']))
+    msg = 'Invalid basis: DataFrame has to have the exact same number and ordering of columns.'
+    with pytest.raises(ValueError, match=msg):
+        dml_blp_confint.confint(basis=pd.DataFrame(np.array([[1, 2, 3], [4, 5, 6]]), columns=['x_1', 'x_2', 'x_3']))
 
 
 @pytest.mark.ci
@@ -709,6 +717,34 @@ def test_doubleml_exception_gate():
     msg = "Groups must be of DataFrame type. Groups of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
         dml_irm_obj.gate(groups=2)
-    msg = 'Columns must be of of bool or int type or the data frame only should contain one column.'
+    msg = (r'Columns of groups must be of bool type or int type \(dummy coded\). '
+           'Alternatively, groups should only contain one column.')
     with pytest.raises(TypeError, match=msg):
-        dml_irm_obj.gate(groups=pd.DataFrame(np.random.normal(0, 1, size=(50, 3))))
+        dml_irm_obj.gate(groups=pd.DataFrame(np.random.normal(0, 1, size=(dml_data_irm.n_obs, 3))))
+
+    dml_irm_obj = DoubleMLIRM(dml_data_irm,
+                              ml_g=Lasso(),
+                              ml_m=LogisticRegression(),
+                              trimming_threshold=0.05,
+                              n_folds=5,
+                              score='ATTE')
+    dml_irm_obj.fit()
+
+    msg = 'Invalid score ATTE. Valid score ATE.'
+    with pytest.raises(ValueError, match=msg):
+        dml_irm_obj.gate(groups=2)
+
+
+@pytest.mark.ci
+def test_doubleml_exception_cate():
+    dml_irm_obj = DoubleMLIRM(dml_data_irm,
+                              ml_g=Lasso(),
+                              ml_m=LogisticRegression(),
+                              trimming_threshold=0.05,
+                              n_folds=5,
+                              score='ATTE')
+    dml_irm_obj.fit()
+
+    msg = 'Invalid score ATTE. Valid score ATE.'
+    with pytest.raises(ValueError, match=msg):
+        dml_irm_obj.cate(basis=2)

doubleml/tests/test_doubleml_exceptions.py

@@ -124,7 +124,7 @@ def test_dml_irm_boot(dml_irm_fixture):
 
 @pytest.mark.ci
 def test_dml_irm_cate_gate():
-    n = 50
+    n = 9
     # collect data
     np.random.seed(42)
     obj_dml_data = make_irm_data(n_obs=n, dim_x=2)
@@ -143,12 +143,24 @@ def test_dml_irm_cate_gate():
     # create a random basis
     random_basis = pd.DataFrame(np.random.normal(0, 1, size=(n, 5)))
     cate = dml_irm_obj.cate(random_basis)
-    assert isinstance(cate, dml.double_ml_blp.DoubleMLIRMBLP)
+    assert isinstance(cate, dml.double_ml_blp.DoubleMLBLP)
+    assert isinstance(cate.confint(), pd.DataFrame)
 
     groups_1 = pd.DataFrame(np.column_stack([obj_dml_data.data['X1'] <= 0,
                                              obj_dml_data.data['X1'] > 0.2]),
                             columns=['Group 1', 'Group 2'])
-    assert isinstance(dml_irm_obj.gate(groups_1), pd.DataFrame)
+    msg = ('At least one group effect is estimated with less than 6 observations.')
+    with pytest.warns(UserWarning, match=msg):
+        gate_1 = dml_irm_obj.gate(groups_1)
+    assert isinstance(gate_1, dml.double_ml_blp.DoubleMLBLP)
+    assert isinstance(gate_1.confint(), pd.DataFrame)
+    assert all(gate_1.confint().index == groups_1.columns)
 
+    np.random.seed(42)
     groups_2 = pd.DataFrame(np.random.choice(["1", "2"], n))
-    assert isinstance(dml_irm_obj.gate(groups_2), pd.DataFrame)
+    msg = ('At least one group effect is estimated with less than 6 observations.')
+    with pytest.warns(UserWarning, match=msg):
+        gate_2 = dml_irm_obj.gate(groups_2)
+    assert isinstance(gate_2, dml.double_ml_blp.DoubleMLBLP)
+    assert isinstance(gate_2.confint(), pd.DataFrame)
+    assert all(gate_2.confint().index == ["Group_1", "Group_2"])