add apo kwargs for cate and gate

Author: SvenKlaassen

doubleml/irm/apo.py

@@ -389,7 +389,7 @@ def _check_data(self, obj_dml_data):
 
         return
 
-    def capo(self, basis, is_gate=False):
+    def capo(self, basis, is_gate=False, **kwargs):
         """
         Calculate conditional average potential outcomes (CAPO) for a given basis.
 
@@ -398,10 +398,14 @@ def capo(self, basis, is_gate=False):
         basis : :class:`pandas.DataFrame`
             The basis for estimating the best linear predictor. Has to have the shape ``(n_obs, d)``,
             where ``n_obs`` is the number of observations and ``d`` is the number of predictors.
+
         is_gate : bool
             Indicates whether the basis is constructed for GATE/GAPOs (dummy-basis).
             Default is ``False``.
 
+        **kwargs: dict
+            Additional keyword arguments to be passed to :meth:`statsmodels.regression.linear_model.OLS.fit` e.g. ``cov_type``.
+
         Returns
         -------
         model : :class:`doubleML.DoubleMLBLP`
@@ -420,10 +424,10 @@ def capo(self, basis, is_gate=False):
         orth_signal = self.psi_elements['psi_b'].reshape(-1)
         # fit the best linear predictor
         model = DoubleMLBLP(orth_signal, basis=basis, is_gate=is_gate)
-        model.fit()
+        model.fit(**kwargs)
         return model
 
-    def gapo(self, groups):
+    def gapo(self, groups, **kwargs):
         """
         Calculate group average potential outcomes (GAPO) for groups.
 
@@ -434,6 +438,9 @@ def gapo(self, groups):
             Has to be dummy coded with shape ``(n_obs, d)``, where ``n_obs`` is the number of observations
             and ``d`` is the number of groups or ``(n_obs, 1)`` and contain the corresponding groups (as str).
 
+        **kwargs: dict
+            Additional keyword arguments to be passed to :meth:`statsmodels.regression.linear_model.OLS.fit` e.g. ``cov_type``.
+
         Returns
         -------
         model : :class:`doubleML.DoubleMLBLP`
@@ -453,5 +460,5 @@ def gapo(self, groups):
         if any(groups.sum(0) <= 5):
             warnings.warn('At least one group effect is estimated with less than 6 observations.')
 
-        model = self.capo(groups, is_gate=True)
+        model = self.capo(groups, is_gate=True, **kwargs)
         return model

doubleml/irm/tests/test_apo.py

@@ -200,8 +200,14 @@ def test_dml_apo_sensitivity(dml_apo_fixture):
                            rtol=1e-9, atol=1e-4)
 
 
+@pytest.fixture(scope='module',
+                params=["nonrobust", "HC0", "HC1", "HC2", "HC3"])
+def cov_type(request):
+    return request.param
+
+
 @pytest.mark.ci
-def test_dml_apo_capo_gapo(treatment_level):
+def test_dml_apo_capo_gapo(treatment_level, cov_type):
     n = 20
     # collect data
     np.random.seed(42)
@@ -221,25 +227,28 @@ def test_dml_apo_capo_gapo(treatment_level):
     dml_obj.fit()
     # create a random basis
     random_basis = pd.DataFrame(np.random.normal(0, 1, size=(n, 5)))
-    capo = dml_obj.capo(random_basis)
+    capo = dml_obj.capo(random_basis, cov_type=cov_type)
     assert isinstance(capo, dml.utils.blp.DoubleMLBLP)
     assert isinstance(capo.confint(), pd.DataFrame)
+    assert capo.blp_model.cov_type == cov_type
 
     groups_1 = pd.DataFrame(np.column_stack([obj_dml_data.data['X1'] <= -1.0,
                                              obj_dml_data.data['X1'] > 0.2]),
                             columns=['Group 1', 'Group 2'])
     msg = ('At least one group effect is estimated with less than 6 observations.')
     with pytest.warns(UserWarning, match=msg):
-        gapo_1 = dml_obj.gapo(groups_1)
+        gapo_1 = dml_obj.gapo(groups_1, cov_type=cov_type)
     assert isinstance(gapo_1, dml.utils.blp.DoubleMLBLP)
     assert isinstance(gapo_1.confint(), pd.DataFrame)
     assert all(gapo_1.confint().index == groups_1.columns.to_list())
+    assert gapo_1.blp_model.cov_type == cov_type
 
     np.random.seed(42)
     groups_2 = pd.DataFrame(np.random.choice(["1", "2"], n, p=[0.1, 0.9]))
     msg = ('At least one group effect is estimated with less than 6 observations.')
     with pytest.warns(UserWarning, match=msg):
-        gapo_2 = dml_obj.gapo(groups_2)
+        gapo_2 = dml_obj.gapo(groups_2, cov_type=cov_type)
     assert isinstance(gapo_2, dml.utils.blp.DoubleMLBLP)
     assert isinstance(gapo_2.confint(), pd.DataFrame)
     assert all(gapo_2.confint().index == ["Group_1", "Group_2"])
+    assert gapo_2.blp_model.cov_type == cov_type