Skip to content

[Repo Assist] fix(estimators): use treatment/control indicators in PropensityScoreWeightingEstimator #1778

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
github-actions wants to merge 1 commit into from
Closed

[Repo Assist] fix(estimators): use treatment/control indicators in PropensityScoreWeightingEstimator #1778

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 38 additions & 69 deletions dowhy/causal_estimators/propensity_score_weighting_estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +148 to +149

# 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]
)
Comment on lines +162 to +164
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":
Expand All @@ -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

Expand Down
48 changes: 48 additions & 0 deletions tests/causal_estimators/test_propensity_score_weighting_estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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})
Comment on lines +105 to +113

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"
Loading