fix tests

judithabk6 · Nov 14, 2024 · ac77369 · ac77369
1 parent 9da25c0
commit ac77369
Show file tree

Hide file tree

Showing 11 changed files with 75 additions and 928 deletions.
diff --git a/src/med_bench/estimation/mediation_tmle.py b/src/med_bench/estimation/mediation_tmle.py
@@ -3,7 +3,6 @@
 from sklearn.linear_model import LinearRegression
 
 from med_bench.estimation.base import Estimator
-from med_bench.nuisances.utils import _get_regressor
 from med_bench.utils.decorators import fitted
 
 ALPHA = 10
@@ -85,7 +84,7 @@ def _one_step_correction_direct(self, t, m, x, y):
         mu_t0_mx_star = mu_t0_mx + epsilon_h * h_corrector_t0
         mu_t1_mx_star = mu_t1_mx + epsilon_h * h_corrector_t1
 
-        regressor_y = _get_regressor(self._regularize, self._use_forest)
+        regressor_y = self.regressor
         reg_cross = clone(regressor_y)
         reg_cross.fit(
             x[t == 0], (mu_t1_mx_star[t == 0] -
@@ -143,8 +142,7 @@ def _one_step_correction_indirect(self, t, m, x, y):
         h_corrector_t1 = t1 / p_x - t1 * ratio
         mu_t1_mx_star = mu_t1_mx + epsilon_h * h_corrector_t1
 
-        regressor_y = _get_regressor(self._regularize,
-                                     self._use_forest)
+        regressor_y = self.regressor
 
         reg_cross = clone(regressor_y)
         reg_cross.fit(x[t == 0], mu_t1_mx_star[t == 0])

diff --git a/src/med_bench/get_estimation_results.py b/src/med_bench/get_estimation_results.py
@@ -15,7 +15,7 @@
 from med_bench.estimation.mediation_g_computation import GComputation
 from med_bench.estimation.mediation_ipw import InversePropensityWeighting
 from med_bench.estimation.mediation_mr import MultiplyRobust
-from med_bench.nuisances.utils import _get_regularization_parameters
+from med_bench.utils.utils import _get_regularization_parameters
 from med_bench.utils.constants import CV_FOLDS
 
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor

diff --git a/src/med_bench/mediation.py b/src/med_bench/mediation.py
@@ -9,7 +9,6 @@
 import numpy as np
 import pandas as pd
 from sklearn.base import clone
-from sklearn.linear_model import RidgeCV
 
 
 from .utils.nuisances import (

diff --git a/src/med_bench/nuisances/conditional_outcome.py b/src/med_bench/nuisances/conditional_outcome.py
diff --git a/src/med_bench/nuisances/cross_conditional_outcome.py b/src/med_bench/nuisances/cross_conditional_outcome.py