add warnings for multiple treatments

Author: SvenKlaassen

doubleml/double_ml_irm.py

@@ -351,6 +351,10 @@ def cate(self, basis):
             raise ValueError('Invalid score ' + self.score + '. ' +
                              'Valid score ' + ' or '.join(valid_score) + '.')
 
+        if self.n_rep != 1:
+            raise NotImplementedError('Only implemented for one repetition. ' +
+                                      f'Number of repetitions is {str(self.n_rep)}.')
+
         # define the orthogonal signal
         orth_signal = self.psi_elements['psi_b'].reshape(-1)
         # fit the best linear predictor
@@ -378,6 +382,10 @@ def gate(self, groups):
             raise ValueError('Invalid score ' + self.score + '. ' +
                              'Valid score ' + ' or '.join(valid_score) + '.')
 
+        if self.n_rep != 1:
+            raise NotImplementedError('Only implemented for one repetition. ' +
+                                      f'Number of repetitions is {str(self.n_rep)}.')
+
         if not isinstance(groups, pd.DataFrame):
             raise TypeError('Groups must be of DataFrame type. '
                             f'Groups of type {str(type(groups))} was passed.')

doubleml/tests/test_doubleml_exceptions.py

@@ -752,6 +752,19 @@ def test_doubleml_exception_gate():
     with pytest.raises(ValueError, match=msg):
         dml_irm_obj.gate(groups=2)
 
+    dml_irm_obj = DoubleMLIRM(dml_data_irm,
+                              ml_g=Lasso(),
+                              ml_m=LogisticRegression(),
+                              trimming_threshold=0.05,
+                              n_folds=5,
+                              score='ATE',
+                              n_rep=2)
+    dml_irm_obj.fit()
+
+    msg = 'Only implemented for one repetition. Number of repetitions is 2.'
+    with pytest.raises(NotImplementedError, match=msg):
+        dml_irm_obj.gate(groups=2)
+
 
 @pytest.mark.ci
 def test_doubleml_exception_cate():
@@ -766,3 +779,15 @@ def test_doubleml_exception_cate():
     msg = 'Invalid score ATTE. Valid score ATE.'
     with pytest.raises(ValueError, match=msg):
         dml_irm_obj.cate(basis=2)
+
+    dml_irm_obj = DoubleMLIRM(dml_data_irm,
+                              ml_g=Lasso(),
+                              ml_m=LogisticRegression(),
+                              trimming_threshold=0.05,
+                              n_folds=5,
+                              score='ATE',
+                              n_rep=2)
+    dml_irm_obj.fit()
+    msg = 'Only implemented for one repetition. Number of repetitions is 2.'
+    with pytest.raises(NotImplementedError, match=msg):
+        dml_irm_obj.cate(basis=2)