[Repo Assist] fix(pss): stop _get_strata from mutating the caller's DataFrame

[github-actions]

🤖 This PR was created by Repo Assist, an automated AI assistant.

Problem

PropensityScoreStratificationEstimator._get_strata directly mutates the DataFrame passed to it by adding four internal bookkeeping columns:

data["strata"] = ...
data["dbar"] = ...
data["d_y"] = ...
data["dbar_y"] = ...

Because estimate_effect() passes the caller's data argument straight into _get_strata, any code that reuses the same DataFrame after calling estimate_effect() finds four unexpected columns polluting its data.

Example of the bug:

model.estimate_effect(estimand, method_name="backdoor.propensity_score_stratification", ...)
print(data.columns)
# Index(['W', 'T', 'Y', 'propensity_score', 'strata', 'dbar', 'd_y', 'dbar_y'], ...)
#                                              ^^^^^^  unexpected internal columns ^

Fix

Add data = data.copy() at the start of _get_strata so all mutations are local to the method. The copy is a shallow/view-friendly pandas copy and adds negligible overhead since _get_strata is called at most a handful of times per estimate_effect invocation.

The clipped subset returned by the method still carries the newly added columns (required by the aggregation in estimate_effect()), so correctness is unaffected.

Test

Added test_estimate_effect_does_not_mutate_input_dataframe() to tests/causal_estimators/test_propensity_score_stratification_estimator.py. The test:

1. Creates a simple linear dataset.
2. Calls estimate_effect with PSS.
3. Asserts that none of the four internal column names (strata, dbar, d_y, dbar_y) are present on the original DataFrame.

Test Status

• ✅ black --check passes on all modified files
• ✅ flake8 — no new violations introduced
• ⚠️ Full test suite requires poetry install (causallearn not available in the CI-lite environment); the logic of _get_strata is unchanged — only a .copy() is prepended

Generated by 🌈 Repo Assist, see workflow run. Learn more.

Generated by 🌈 Repo Assist, see workflow run. Learn more.

To install this agentic workflow, run

gh aw add githubnext/agentics/workflows/repo-assist.md@11c9a2c442e519ff2b427bf58679f5a525353f76