Merge pull request #3 from DoubleML/m-iivm-subgroups

implemented the subgroups for the IIVM model for the serverless class

Author: MalteKurz

doubleml_serverless/double_ml_iivm_aws_lambda.py

@@ -74,32 +74,49 @@ def _ml_nuisance_aws_lambda(self, cv_params):
                         self._dml_data.z_cols[0], self._dml_data.x_cols,
                         method='predict_proba')
 
-        _attach_learner(payload_ml_r0,
-                        'ml_r0', self.learner['ml_r'],
-                        self._dml_data.d_cols[0], self._dml_data.x_cols,
-                        method='predict_proba')
-
-        _attach_learner(payload_ml_r1,
-                        'ml_r1', self.learner['ml_r'],
-                        self._dml_data.d_cols[0], self._dml_data.x_cols,
-                        method='predict_proba')
-
-        all_payloads = [payload_ml_g0, payload_ml_g1, payload_ml_m, payload_ml_r0, payload_ml_r1]
-        all_smpls = [smpls_z0, smpls_z1, self.smpls, smpls_z0, smpls_z1]
+        all_payloads = [payload_ml_g0, payload_ml_g1, payload_ml_m]
+        all_smpls = [smpls_z0, smpls_z1, self.smpls]
+        send_train_ids = [True, True, False]
+        params_names = ['ml_g0', 'ml_g1', 'ml_m']
+
+        if self.subgroups['always_takers']:
+            _attach_learner(payload_ml_r0,
+                            'ml_r0', self.learner['ml_r'],
+                            self._dml_data.d_cols[0], self._dml_data.x_cols,
+                            method='predict_proba')
+            all_payloads.append(payload_ml_r0)
+            all_smpls.append(smpls_z0)
+            send_train_ids.append(True)
+            params_names.append('ml_r0')
+
+        if self.subgroups['never_takers']:
+            _attach_learner(payload_ml_r1,
+                            'ml_r1', self.learner['ml_r'],
+                            self._dml_data.d_cols[0], self._dml_data.x_cols,
+                            method='predict_proba')
+            all_payloads.append(payload_ml_r1)
+            all_smpls.append(smpls_z1)
+            send_train_ids.append(True)
+            params_names.append('ml_r1')
 
         payloads = _attach_smpls(all_payloads,
                                  all_smpls,
                                  self.n_folds,
                                  self.n_rep,
                                  self._dml_data.n_obs,
                                  cv_params['n_lambdas_cv'],
-                                 [True, True, False, True, True],
+                                 send_train_ids,
                                  cv_params['seed'])
 
-        preds = self.invoke_lambdas(payloads, self.smpls, self.params_names,
+        preds = self.invoke_lambdas(payloads, self.smpls, params_names,
                                     self._dml_data.n_obs, self.n_rep,
                                     cv_params['n_lambdas_cv'])
 
+        if not self.subgroups['always_takers']:
+            preds['ml_r0'] = np.zeros_like(preds['ml_g0'])
+        if not self.subgroups['never_takers']:
+            preds['ml_r1'] = np.ones_like(preds['ml_g1'])
+
         for i_rep in range(self.n_rep):
             # compute score elements
 

doubleml_serverless/tests/test_iivm.py

@@ -51,8 +51,16 @@ def trimming_threshold(request):
     return request.param
 
 
+@pytest.fixture(scope='module',
+                params=[{'always_takers': True, 'never_takers': True},
+                        {'always_takers': False, 'never_takers': True},
+                        {'always_takers': True, 'never_takers': False}])
+def subgroups(request):
+    return request.param
+
+
 @pytest.fixture(scope="module")
-def dml_iivm_fixture(generate_data_iivm, idx, learner, score, dml_procedure, trimming_threshold):
+def dml_iivm_fixture(generate_data_iivm, idx, learner, score, dml_procedure, trimming_threshold, subgroups):
     boot_methods = ['normal']
     n_folds = 4
     n_rep_boot = 502
@@ -77,6 +85,7 @@ def dml_iivm_fixture(generate_data_iivm, idx, learner, score, dml_procedure, tri
                                                   ml_g, ml_m, ml_r,
                                                   n_folds,
                                                   score=score,
+                                                  subgroups=subgroups,
                                                   dml_procedure=dml_procedure)
 
     dml_iivm_lambda.fit_aws_lambda()
@@ -87,6 +96,7 @@ def dml_iivm_fixture(generate_data_iivm, idx, learner, score, dml_procedure, tri
                                 ml_g, ml_m, ml_r,
                                 n_folds,
                                 score=score,
+                                subgroups=subgroups,
                                 dml_procedure=dml_procedure)
 
     dml_iivm.fit()