[Repo Assist] fix(estimators): use treatment/control indicators in PropensityScoreWeightingEstimator - #1778
Closed
github-actions[bot] wants to merge 1 commit into
Conversation
…eightingEstimator
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] <github-actions[bot]@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Fixes PropensityScoreWeightingEstimator weight computations for non-{0,1} binary treatment encodings by replacing arithmetic on the raw treatment column with an explicit treated/control indicator, and adds a regression test to prevent future regressions.
Changes:
- Introduce a binary treated indicator (
d) after propensity-score trimming and use it across all IPS/Hajek/stabilized weight formulas. - Update effect calculation to use the treated indicator instead of the raw treatment column.
- Add a regression test covering treatment encoded as
{1,2}.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
dowhy/causal_estimators/propensity_score_weighting_estimator.py |
Replaces hardcoded {0,1} treatment arithmetic with a derived treated indicator in all weighting formulas. |
tests/causal_estimators/test_propensity_score_weighting_estimator.py |
Adds a regression test ensuring PSW works with treatment encoded as {1,2}. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+148
to
+149
| # 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) |
| 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])) |
Comment on lines
+162
to
+164
| data["ips_weight"] = d / data[self.propensity_score_column] + (1 - d) / ( | ||
| 1 - 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
+105
to
+113
| 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}) |
This was referenced Sep 1, 2026
Member
|
superseded by #1795 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🤖 This PR was created by Repo Assist, an automated AI assistant.
Problem
PropensityScoreWeightingEstimatorhardcoded the assumption that treatment is encoded as{0, 1}in every IPW weight formula. The method useddata[treatment_col]directly in arithmetic:For any binary treatment encoding other than
{0, 1}(e.g.{1, 2},{-1, 1}, boolean), all weight variants (vanilla IPS, Hajek, stabilized — and theirtips/cipsATT/ATC variants) are computed incorrectly. The numerators and denominators of every weight formula includedata[T]as if it equals 1 for treated and 0 for control.Fix
Derive a float binary indicator immediately after propensity-score trimming:
Then replace every
data[T]withdand every(1 - data[T])with(1 - d)throughout all weight formulas. This is correct for any binary encoding.Relationship to Other PRs
PropensityScoreMatchingEstimator) and PR [Repo Assist] fix(estimators): use treatment/control indicators in PropensityScoreStratificationEstimator #1766 (PropensityScoreStratificationEstimator+ base-class binary check).PropensityScoreEstimator.fit()binary check (relaxed from{0,1}-only to "exactly 2 distinct values" in PR [Repo Assist] fix(estimators): use treatment/control indicators in PropensityScoreStratificationEstimator #1766) is a prerequisite for calling PSW with non-{0,1}treatment, but that change is independent.Test Status
✅ New regression test
test_psw_non_zero_one_treatment_encoding— treatment{1, 2}(control=1, treated=2) now produces a finite ATE estimate close to the true value of 3.0.✅ Syntax check passes on both modified files.
✅
blackandisortformatting verified (imports sorted correctly, line length ≤ 120).TestPropensityScoreWeightingEstimator) not run due to environment limitations; the change is a mechanical substitution ofdata[T]→dthroughout weight formulas.