From 99411c8ea726bc9f0ae6139bbb5cd0adb00ce66c Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:54:16 +0000 Subject: [PATCH] fix(estimators): use treatment/control indicators in PropensityScoreWeightingEstimator PSW weight formulas used data[treatment_col] directly in arithmetic (e.g. data[T] / ps, 1 - data[T]), which only gives correct results when treatment is encoded as {0, 1}. For any other binary encoding ({1, 2}, {-1, 1}, etc.) all IPW weight variants (vanilla, Hajek, stabilized) were computed incorrectly. Fix: derive a float indicator d = (data[T] == treatment_value).astype(float) immediately after propensity-score trimming, then use d / (1-d) throughout every weight formula. This mirrors the fix applied to PSS in PR #1766 and to PSM in PR #1762, completing the trilogy for all three propensity-score estimators. Adds regression test: treatment encoded as {1, 2} now produces a finite, sensible ATE estimate (was giving wrong weights before this fix). Note: the base-class binary check in PropensityScoreEstimator.fit() still validates that treatment is binary via two distinct values (relaxed in #1766); this commit focuses only on the weight-calculation fix in PSW. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: github-actions[bot] --- .../propensity_score_weighting_estimator.py | 107 +++++++----------- ...st_propensity_score_weighting_estimator.py | 48 ++++++++ 2 files changed, 86 insertions(+), 69 deletions(-) diff --git a/dowhy/causal_estimators/propensity_score_weighting_estimator.py b/dowhy/causal_estimators/propensity_score_weighting_estimator.py index 7ffecf9156..b2135b3473 100755 --- a/dowhy/causal_estimators/propensity_score_weighting_estimator.py +++ b/dowhy/causal_estimators/propensity_score_weighting_estimator.py @@ -145,82 +145,61 @@ def estimate_effect( data[self.propensity_score_column] = np.minimum(self.max_ps_score, data[self.propensity_score_column]) data[self.propensity_score_column] = np.maximum(self.min_ps_score, data[self.propensity_score_column]) + # Binary indicator for treated units: works for any binary encoding (e.g. {1,2}, {-1,1}, {True,False}) + d = (data[self._target_estimand.treatment_variable[0]] == treatment_value).astype(float) + # ips ==> (isTreated(y)/ps(y)) + ((1-isTreated(y))/(1-ps(y))) # nips ==> ips / (sum of ips over all units) # icps ==> ps(y)/(1-ps(y)) / (sum of (ps(y)/(1-ps(y))) over all control units) # itps ==> ps(y)/(1-ps(y)) / (sum of (ps(y)/(1-ps(y))) over all treatment units) - ipst_sum = sum(data[self._target_estimand.treatment_variable[0]] / data[self.propensity_score_column]) - ipsc_sum = sum( - (1 - data[self._target_estimand.treatment_variable[0]]) / (1 - data[self.propensity_score_column]) - ) - num_units = len(data[self._target_estimand.treatment_variable[0]]) - num_treatment_units = sum(data[self._target_estimand.treatment_variable[0]]) + ipst_sum = sum(d / data[self.propensity_score_column]) + ipsc_sum = sum((1 - d) / (1 - data[self.propensity_score_column])) + num_units = len(d) + num_treatment_units = sum(d) num_control_units = num_units - num_treatment_units # Vanilla IPS estimator - data["ips_weight"] = data[self._target_estimand.treatment_variable[0]] / data[self.propensity_score_column] + ( - 1 - data[self._target_estimand.treatment_variable[0]] - ) / (1 - data[self.propensity_score_column]) - data["tips_weight"] = data[self._target_estimand.treatment_variable[0]] + ( - 1 - data[self._target_estimand.treatment_variable[0]] - ) * data[self.propensity_score_column] / (1 - data[self.propensity_score_column]) - data["cips_weight"] = data[self._target_estimand.treatment_variable[0]] * ( + data["ips_weight"] = d / data[self.propensity_score_column] + (1 - d) / ( + 1 - data[self.propensity_score_column] + ) + data["tips_weight"] = d + (1 - d) * data[self.propensity_score_column] / ( 1 - data[self.propensity_score_column] - ) / data[self.propensity_score_column] + (1 - data[self._target_estimand.treatment_variable[0]]) + ) + data["cips_weight"] = d * (1 - data[self.propensity_score_column]) / data[self.propensity_score_column] + ( + 1 - d + ) # The Hajek estimator (or the self-normalized estimator) data["ips_normalized_weight"] = ( - data[self._target_estimand.treatment_variable[0]] / data[self.propensity_score_column] / ipst_sum - + (1 - data[self._target_estimand.treatment_variable[0]]) - / (1 - data[self.propensity_score_column]) - / ipsc_sum - ) - ipst_for_att_sum = sum(data[self._target_estimand.treatment_variable[0]]) - ipsc_for_att_sum = sum( - (1 - data[self._target_estimand.treatment_variable[0]]) - / (1 - data[self.propensity_score_column]) - * data[self.propensity_score_column] + d / data[self.propensity_score_column] / ipst_sum + + (1 - d) / (1 - data[self.propensity_score_column]) / ipsc_sum ) + ipst_for_att_sum = sum(d) + ipsc_for_att_sum = sum((1 - d) / (1 - data[self.propensity_score_column]) * data[self.propensity_score_column]) data["tips_normalized_weight"] = ( - data[self._target_estimand.treatment_variable[0]] / ipst_for_att_sum - + (1 - data[self._target_estimand.treatment_variable[0]]) - * data[self.propensity_score_column] - / (1 - data[self.propensity_score_column]) - / ipsc_for_att_sum - ) - ipst_for_atc_sum = sum( - data[self._target_estimand.treatment_variable[0]] - / data[self.propensity_score_column] - * (1 - data[self.propensity_score_column]) + d / ipst_for_att_sum + + (1 - d) * data[self.propensity_score_column] / (1 - data[self.propensity_score_column]) / ipsc_for_att_sum ) - ipsc_for_atc_sum = sum((1 - data[self._target_estimand.treatment_variable[0]])) + ipst_for_atc_sum = sum(d / data[self.propensity_score_column] * (1 - data[self.propensity_score_column])) + ipsc_for_atc_sum = sum(1 - d) data["cips_normalized_weight"] = ( - data[self._target_estimand.treatment_variable[0]] - * (1 - data[self.propensity_score_column]) - / data[self.propensity_score_column] - / ipst_for_atc_sum - + (1 - data[self._target_estimand.treatment_variable[0]]) / ipsc_for_atc_sum + d * (1 - data[self.propensity_score_column]) / data[self.propensity_score_column] / ipst_for_atc_sum + + (1 - d) / ipsc_for_atc_sum ) # Stabilized weights (from Robins, Hernan, Brumback (2000)) # Paper: Marginal Structural Models and Causal Inference in Epidemiology - p_treatment = sum(data[self._target_estimand.treatment_variable[0]]) / num_units - data["ips_stabilized_weight"] = data[self._target_estimand.treatment_variable[0]] / data[ - self.propensity_score_column - ] * p_treatment + (1 - data[self._target_estimand.treatment_variable[0]]) / ( - 1 - data[self.propensity_score_column] - ) * ( - 1 - p_treatment + p_treatment = sum(d) / num_units + data["ips_stabilized_weight"] = ( + d / data[self.propensity_score_column] * p_treatment + + (1 - d) / (1 - data[self.propensity_score_column]) * (1 - p_treatment) ) - data["tips_stabilized_weight"] = data[self._target_estimand.treatment_variable[0]] * p_treatment + ( - 1 - data[self._target_estimand.treatment_variable[0]] - ) * data[self.propensity_score_column] / (1 - data[self.propensity_score_column]) * (1 - p_treatment) - data["cips_stabilized_weight"] = data[self._target_estimand.treatment_variable[0]] * ( + data["tips_stabilized_weight"] = d * p_treatment + (1 - d) * data[self.propensity_score_column] / ( 1 - data[self.propensity_score_column] - ) / data[self.propensity_score_column] * p_treatment + ( - 1 - data[self._target_estimand.treatment_variable[0]] - ) * ( - 1 - p_treatment + ) * (1 - p_treatment) + data["cips_stabilized_weight"] = ( + d * (1 - data[self.propensity_score_column]) / data[self.propensity_score_column] * p_treatment + + (1 - d) * (1 - p_treatment) ) if isinstance(target_units, pd.DataFrame) or target_units == "ate": @@ -233,20 +212,10 @@ def estimate_effect( raise ValueError(f"Target units value {target_units} not supported") # Calculating the effect - data["d_y"] = ( - data[weighting_scheme_name] - * data[self._target_estimand.treatment_variable[0]] - * data[self._target_estimand.outcome_variable[0]] - ) - data["dbar_y"] = ( - data[weighting_scheme_name] - * (1 - data[self._target_estimand.treatment_variable[0]]) - * data[self._target_estimand.outcome_variable[0]] - ) - sum_dy_weights = np.sum(data[self._target_estimand.treatment_variable[0]] * data[weighting_scheme_name]) - sum_dbary_weights = np.sum( - (1 - data[self._target_estimand.treatment_variable[0]]) * data[weighting_scheme_name] - ) + data["d_y"] = data[weighting_scheme_name] * d * data[self._target_estimand.outcome_variable[0]] + data["dbar_y"] = data[weighting_scheme_name] * (1 - d) * data[self._target_estimand.outcome_variable[0]] + sum_dy_weights = np.sum(d * data[weighting_scheme_name]) + sum_dbary_weights = np.sum((1 - d) * data[weighting_scheme_name]) # Subtracting the weighted means est = data["d_y"].sum() / sum_dy_weights - data["dbar_y"].sum() / sum_dbary_weights diff --git a/tests/causal_estimators/test_propensity_score_weighting_estimator.py b/tests/causal_estimators/test_propensity_score_weighting_estimator.py index 54d239821e..ffd6fe0101 100755 --- a/tests/causal_estimators/test_propensity_score_weighting_estimator.py +++ b/tests/causal_estimators/test_propensity_score_weighting_estimator.py @@ -1,6 +1,11 @@ from pytest import mark +import numpy as np +import pandas as pd + +from dowhy import EstimandType, identify_effect_auto from dowhy.causal_estimators.propensity_score_weighting_estimator import PropensityScoreWeightingEstimator +from dowhy.graph import build_graph_from_str from .base import SimpleEstimator @@ -88,3 +93,46 @@ def test_average_treatment_effect( ], method_params={"num_simulations": 1, "num_null_simulations": 1}, ) + + +def test_psw_non_zero_one_treatment_encoding(): + """Regression test: PSW must use treatment_value/control_value indicators, not raw data[treatment]. + + When treatment is encoded as {1, 2} instead of {0, 1}, the old implementation used + ``data[T] / ps`` and ``1 - data[T]`` directly in weight formulas, which is only + correct for {0, 1} encodings and gives wrong results for any other binary encoding. + """ + rng = np.random.default_rng(42) + n = 2000 + X = rng.standard_normal(n) + # Treatment encoded as {1, 2} (control=1, treated=2) + ps = 1 / (1 + np.exp(-X)) + T = (rng.random(n) < ps).astype(int) + 1 # 1 = control, 2 = treated + Y = 3.0 * (T == 2) + 0.5 * X + rng.standard_normal(n) + + df = pd.DataFrame({"X": X, "T": T, "Y": Y}) + + gml = """graph [ + directed 1 + node [id "T" label "T"] + node [id "X" label "X"] + node [id "Y" label "Y"] + edge [source "X" target "T"] + edge [source "X" target "Y"] + edge [source "T" target "Y"] + ]""" + graph = build_graph_from_str(gml) + estimand = identify_effect_auto( + graph, + observed_nodes=["T", "X", "Y"], + action_nodes=["T"], + outcome_nodes=["Y"], + estimand_type=EstimandType.NONPARAMETRIC_ATE, + ) + + estimator = PropensityScoreWeightingEstimator(identified_estimand=estimand) + estimator.fit(df) + estimate = estimator.estimate_effect(df, treatment_value=2, control_value=1, target_units="ate") + + assert np.isfinite(estimate.value), f"Expected a finite estimate, got {estimate.value}" + assert abs(estimate.value - 3.0) < 1.0, f"ATE estimate {estimate.value} too far from true value 3.0"