fix clipping in IPW

Author: judithabk6

src/med_bench/estimation/mediation_ipw.py

@@ -59,14 +59,18 @@ def _pointwise_estimate(self, t, m, x, y):
 
         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 == "treatment":
             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)
 
         elif self._prop_ratio == "mediator":
             f_m0x, f_m1x = self._estimate_mediator_probability(x, m)
+            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)
 
@@ -79,8 +83,6 @@ def _pointwise_estimate(self, t, m, x, y):
             prop_ratio_t0_m1[ind],
         )
 
-        p_x = np.clip(p_x, self._clip, 1 - self._clip)
-
         # importance weighting
         y1m1 = (y * t / p_x) / np.mean(t / p_x)
         y1m0 = (y * t * prop_ratio_t1_m0) / np.mean(t * prop_ratio_t1_m0)