added clipping mr and tmle, fixed bug on mediator density and tests

Author: houssamzenati

README.md

@@ -12,6 +12,7 @@ Judith Abécassis, Julie Josse and Bertrand Thirion (2022). **Causal mediation a
 med_bench can be installed by executing
 ```
 python setup.py install
+pip install -e .
 ```
 
 Or the package can be directly installed from the GitHub repository using

src/med_bench/estimation/base.py

@@ -14,7 +14,7 @@ class Estimator:
 
     __metaclass__ = ABCMeta
 
-    def __init__(self, verbose: bool = True, mediator_cardinality_threshold: int=10):
+    def __init__(self, verbose: bool = True, mediator_cardinality_threshold: int = 10):
         """Initializes Estimator base class
 
         Parameters
@@ -25,11 +25,11 @@ def __init__(self, verbose: bool = True, mediator_cardinality_threshold: int=10)
         mediator_cardinality_threshold: int
             default 10
             maximal number of categories in a mediator to treat it as discrete or continuous
-            if the mediator is 1-dimensional and 
+            if the mediator is 1-dimensional and
                 if the number of distinct values in the mediator is lower than
                 mediator_cardinality_threshold, the mediator is going to be considered
                 discrete when estimating the mediator probability function, otherwise
-                the mediator is considered continuous, 
+                the mediator is considered continuous,
         """
         self._verbose = verbose
         self._fitted = False
@@ -39,7 +39,6 @@ def __init__(self, verbose: bool = True, mediator_cardinality_threshold: int=10)
     def verbose(self):
         return self._verbose
 
-
     @abstractmethod
     def fit(self, t, m, x, y):
         """Fits nuisance parameters to data
@@ -231,7 +230,9 @@ def _resize(self, t, m, x, y):
 
     def _fit_mediator_discretizer(self, m):
         """Fits the discretization procedure of mediators"""
-        if (is_array_integer(m)) and (len(np.unique(m)) <= self.mediator_cardinality_threshold):
+        if (is_array_integer(m)) and (
+            len(np.unique(m)) <= self.mediator_cardinality_threshold
+        ):
             self.discretizer = LabelEncoder()
             self.discretizer.fit(m.ravel())
             self.mediator_bins = self.discretizer.classes_
@@ -362,12 +363,15 @@ def _estimate_discrete_mediator_probability(self, x, m_label):
         t0 = np.zeros((n, 1))
         t1 = np.ones((n, 1))
 
-
         t0_x = np.hstack([t0.reshape(-1, 1), x])
         t1_x = np.hstack([t1.reshape(-1, 1), x])
 
-        f_m0x = self._classifier_m.predict_proba(t0_x)[np.arange(m_label.shape[0]), m_label]
-        f_m1x = self._classifier_m.predict_proba(t1_x)[np.arange(m_label.shape[0]), m_label]
+        f_m0x = self._classifier_m.predict_proba(t0_x)[
+            np.arange(m_label.shape[0]), m_label
+        ]
+        f_m1x = self._classifier_m.predict_proba(t1_x)[
+            np.arange(m_label.shape[0]), m_label
+        ]
 
         return f_m0x, f_m1x
 
@@ -392,8 +396,8 @@ def _estimate_mediator_density(self, x, m):
         t0_x = np.hstack([t0.reshape(-1, 1), x])
         t1_x = np.hstack([t1.reshape(-1, 1), x])
 
-        f_m0x = self._density_m.pdf(t0_x, m)
-        f_m1x = self._density_m.pdf(t1_x, m)
+        f_m0x = self._density_m.pdf(t0_x, m.squeeze())
+        f_m1x = self._density_m.pdf(t1_x, m.squeeze())
 
         return f_m0x, f_m1x
 

src/med_bench/estimation/mediation_mr.py

@@ -13,6 +13,8 @@ def __init__(
         self,
         regressor,
         classifier,
+        clip: float,
+        trim: float,
         prop_ratio="treatment",
         integration="implicit",
         normalized=True,
@@ -26,6 +28,10 @@ def __init__(
             Regressor used for mu estimation, can be any object with a fit and predict method
         classifier
             Classifier used for propensity estimation, can be any object with a fit and predict_proba method
+        clip : float
+            Clipping value for propensity scores
+        trim : float
+            Trimming value for propensity scores
         prop_ratio : str
             prop_ratio to use for estimation, can be either 'mediator' or 'treatment'
         integration : str
@@ -47,6 +53,8 @@ def __init__(
         self.regressor = regressor
         self.classifier = classifier
 
+        self._clip = clip
+        self._trim = trim
         assert prop_ratio in ["mediator", "treatment"]
         assert integration in ["implicit", "explicit"]
         self._prop_ratio = prop_ratio
@@ -96,15 +104,19 @@ def _pointwise_estimate(self, t, m, x, y):
         # Format checking
         t, m, x, y = self._resize(t, m, x, y)
 
+        p_x = self._estimate_treatment_propensity_x(x)
+        p_x = np.clip(p_x, self._clip, 1 - self._clip)
+
         if self._prop_ratio == "mediator":
             f_m0x, f_m1x = self._estimate_mediator_probability(x, m)
-            p_x = self._estimate_treatment_propensity_x(x)
+            f_m0x = np.clip(f_m0x, self._clip, None)
+            f_m1x = np.clip(f_m1x, self._clip, None)
             prop_ratio_t1_m0 = f_m0x / (p_x * f_m1x)
             prop_ratio_t0_m1 = f_m1x / ((1 - p_x) * f_m0x)
 
         elif self._prop_ratio == "treatment":
-            p_x = self._estimate_treatment_propensity_x(x)
             p_xm = self._estimate_treatment_propensity_xm(m, x)
+            p_xm = np.clip(p_xm, self._clip, 1 - self._clip)
             prop_ratio_t1_m0 = (1 - p_xm) / ((1 - p_x) * p_xm)
             prop_ratio_t0_m1 = p_xm / ((1 - p_xm) * p_x)
 
@@ -133,13 +145,28 @@ def _pointwise_estimate(self, t, m, x, y):
             y1m0 = y1m0.mean()
             y0m1 = y0m1.mean()
 
-            mu_0mx, mu_1mx = self._estimate_conditional_mean_outcome(x, m_discrete_value)
+            mu_0mx, mu_1mx = self._estimate_conditional_mean_outcome(
+                x, m_discrete_value
+            )
 
         elif self._integration == "implicit":
             mu_0mx, mu_1mx, y0m0, y0m1, y1m0, y1m1 = (
                 self._estimate_cross_conditional_mean_outcome(m, x)
             )
 
+        ind = (p_x > self._trim) & (p_x < (1 - self._trim))
+        y, t, p_x, prop_ratio_t1_m0, prop_ratio_t0_m1 = (
+            y[ind],
+            t[ind],
+            p_x[ind],
+            prop_ratio_t1_m0[ind],
+            prop_ratio_t0_m1[ind],
+        )
+        mu_0mx, mu_1mx = (
+            mu_0mx[ind],
+            mu_1mx[ind],
+        )
+
         # score computing
         if self._normalized:
             sum_score_m1 = np.mean(t / p_x)

src/med_bench/estimation/mediation_tmle.py

@@ -12,7 +12,9 @@
 class TMLE(Estimator):
     """Implementation of targeted maximum likelihood estimation method class"""
 
-    def __init__(self, regressor, classifier, prop_ratio, **kwargs):
+    def __init__(
+        self, regressor, classifier, prop_ratio, clip: float, trim: float, **kwargs
+    ):
         """_summary_
 
         Parameters
@@ -23,6 +25,10 @@ def __init__(self, regressor, classifier, prop_ratio, **kwargs):
             Classifier used for propensity estimation, can be any object with a fit and predict_proba method
         prop_ratio : str
             prop_ratio to use for estimation, can be either 'mediator' or 'treatment'
+        clip : float
+            Clipping value for propensity scores
+        trim : float
+            Trimming value for propensity scores
         """
         super().__init__(**kwargs)
 
@@ -37,6 +43,8 @@ def __init__(self, regressor, classifier, prop_ratio, **kwargs):
         self.regressor = regressor
         self.classifier = classifier
 
+        self._clip = clip
+        self._trim = trim
         assert prop_ratio in ["mediator", "treatment"]
         self._prop_ratio = prop_ratio
 
@@ -65,7 +73,7 @@ def fit(self, t, m, x, y):
             print("Nuisance models fitted")
 
         return self
-        
+
     def _one_step_correction_direct(self, t, m, x, y):
         """Implements the one step correction for the estimation of the natural
         direct effect with the prop_ratio of mediator densities or treatment
@@ -77,17 +85,31 @@ def _one_step_correction_direct(self, t, m, x, y):
         t0 = np.zeros((n))
         t1 = np.ones((n))
 
+        p_x = self._estimate_treatment_propensity_x(x)
+        p_x = np.clip(p_x, self._clip, 1 - self._clip)
+
         # estimate mediator densities
         if self._prop_ratio == "mediator":
             f_m0x, f_m1x = self._estimate_mediator_probability(x, m)
-            p_x = self._estimate_treatment_propensity_x(x)
+            f_m0x = np.clip(f_m0x, self._clip, None)
+            f_m1x = np.clip(f_m1x, self._clip, None)
             prop_ratio = f_m0x / (p_x * f_m1x)
 
         elif self._prop_ratio == "treatment":
-            p_x = self._estimate_treatment_propensity_x(x)
             p_xm = self._estimate_treatment_propensity_xm(m, x)
+            p_xm = np.clip(p_xm, self._clip, 1 - self._clip)
             prop_ratio = (1 - p_xm) / ((1 - p_x) * p_xm)
 
+        ind = (p_x > self._trim) & (p_x < (1 - self._trim))
+        y, t, m, x, p_x, prop_ratio = (
+            y[ind],
+            t[ind],
+            m[ind],
+            x[ind],
+            p_x[ind],
+            prop_ratio[ind],
+        )
+
         # estimation of corrective features for the conditional mean outcome
         h_corrector = t * prop_ratio - (1 - t) / (1 - p_x)
 
@@ -151,17 +173,30 @@ def _one_step_correction_indirect(self, t, m, x, y):
         t0 = np.zeros((n))
         t1 = np.ones((n))
 
+        p_x = self._estimate_treatment_propensity_x(x)
+        p_x = np.clip(p_x, self._clip, 1 - self._clip)
+
         # estimate mediator densities
         if self._prop_ratio == "mediator":
             f_m0x, f_m1x = self._estimate_mediator_probability(x, m)
-            p_x = self._estimate_treatment_propensity_x(x)
+            f_m0x = np.clip(f_m0x, self._clip, None)
+            f_m1x = np.clip(f_m1x, self._clip, None)
             prop_ratio = f_m0x / (p_x * f_m1x)
 
         elif self._prop_ratio == "treatment":
-            p_x = self._estimate_treatment_propensity_x(x)
             p_xm = self._estimate_treatment_propensity_xm(m, x)
+            p_xm = np.clip(p_xm, self._clip, 1 - self._clip)
             prop_ratio = (1 - p_xm) / ((1 - p_x) * p_xm)
 
+        ind = (p_x > self._trim) & (p_x < (1 - self._trim))
+        y, t, m, x, p_x, prop_ratio = (
+            y[ind],
+            t[ind],
+            m[ind],
+            x[ind],
+            p_x[ind],
+            prop_ratio[ind],
+        )
         # estimation of corrective features for the conditional mean outcome
         h_corrector = t / p_x - t * prop_ratio
 
@@ -220,7 +255,6 @@ def _one_step_correction_indirect(self, t, m, x, y):
 
         return delta_1
 
-
     @fitted
     def estimate(self, t, m, x, y):
         """Estimates causal effect on data"""