Consolidate experiment design-matrix attributes into xr.Dataset#849
Conversation
Add _predictor_data_name and _target_data_name class attributes to PyMCModel so subclasses using non-default data node names can customize without re-implementing _data_setter. Validate at predict time and raise a clear ValueError if expected nodes are missing. Also fix pre-existing mypy type: ignore codes in panel_regression.py (attr-defined -> union-attr). Made-with: Cursor
Remove _predictor_data_name / _target_data_name class attributes (added complexity for a case no existing subclass needs). Keep the validation that raises a clear ValueError when expected data nodes are missing. Revert undeclared y_dtype behavioral change. Improve test coverage with separate X-missing and y-missing error paths. Made-with: Cursor
Bundle loose xr.DataArray attributes on experiment classes into xr.Dataset objects to reduce attribute sprawl. Pre/post classes (ITS, SC) use pre_design/post_design; formula-based classes use a single design Dataset. Deprecated @Property accessors preserve backward compatibility. Closes #199. Made-with: Cursor
github-actions
Bot
added
the
refactor
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #849 +/- ##
==========================================
+ Coverage 95.14% 95.16% +0.02%
==========================================
Files 92 93 +1
Lines 14860 15030 +170
Branches 890 896 +6
==========================================
+ Hits 14138 14304 +166
- Misses 505 507 +2
- Partials 217 219 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Adversarial review from gpt-5.4-xhigh.
-
causalpy/checks/convex_hull.pystill calls deprecated aliases.ConvexHullCheck.run()readssc.datapre_controlandsc.datapre_treated, which on this branch now route through the new deprecation shims. I reproduced this locally and the check emits:datapre_control is deprecated, use pre_design['control']datapre_treated is deprecated, use pre_design['treated']
That means a normal synthetic-control workflow now makes CausalPy warn against itself. Immediate fix direction: do a repo-wide sweep for internal uses of the deprecated names, switch those callers to
design[...]/pre_design[...]/post_design[...], and add a warning-focused test incausalpy/tests/test_cross_cutting_checks.pyasserting thatConvexHullCheck.run()stays warning-free. -
The refactor still has a lot of duplicate migration boilerplate, and that is making the change less elegant than it could be. Right now the same migration pattern is hand-written across the experiment classes:
- build one or two
xr.Datasets with near-identicalX/yorcontrol/treatedwrapping logic, - preserve temporary raw arrays just long enough to build those datasets,
- add near-identical deprecated properties that warn and then forward to the new dataset-backed location.
You can see this pattern repeated in
diff_in_diff.py,panel_regression.py,piecewise_its.py,prepostnegd.py,regression_discontinuity.py,regression_kink.py, plus the split pre/post variants ininterrupted_time_series.pyandsynthetic_control.py. The practical problem is not just aesthetics: every repeated shim and dataset wrapper is another place to miss an internal migration, drift in warning text, or leave coverage behind. TheConvexHullCheckmiss feels like a direct symptom of that copy-paste surface area.Suggested fix direction: centralize more of this in
BaseExperiment(or a small mixin/helper module) so subclasses only describe what is experiment-specific. For example:- one helper for the common single-design case that builds the standard
{"X", "y"}dataset given raw arrays, observation index, labels, and treated-unit labels; - one helper for split designs so ITS-style classes can reuse the same builder for pre/post periods;
- one shared deprecation-forwarding helper so the compatibility properties become one-liners instead of each class repeating its own
warnings.warn(...); return self.design[...]block.
I would prefer that over adding more bespoke per-class migration code. If you want to avoid magic, you do not need a dynamic
__getattr__; even explicit properties backed by a shared helper would already remove most of the duplication. - build one or two
-
The backward-compatibility layer is still under-tested, and the remote results reflect that. Every non-Codecov check is green, but
codecov/patchis failing at84.71%with37missing lines, concentrated in the new deprecated properties across the experiment classes. Since preserving old attribute access is one of the main claims of the PR, I’d expect targeted tests that exercise those legacy aliases and assert both the warning and the data equivalence to the new dataset-backed API.Suggested fix direction: add a small parametrized compatibility test suite rather than lots of bespoke tests. I would cover:
- representative single-design aliases (
X,y) and split-design aliases (pre_X,pre_y,post_X,post_y); - the synthetic-control aliases (
datapre_control,datapre_treated,datapost_control,datapost_treated); - equality of the returned object/data with the new dataset-backed access path;
- correct deprecation behavior via
pytest.deprecated_call()or explicit warning capture.
If you centralize the alias metadata, you can drive these tests from the same mapping and shrink both implementation duplication and test duplication at the same time.
- representative single-design aliases (
Net: I like the direction of moving related arrays into xr.Datasets, but I don’t think this is ready yet. The current version is functionally plausible, but not yet simple or especially elegant because the migration logic is still spread across too many classes and one real regression has already slipped through. I’d first eliminate internal uses of deprecated aliases, then centralize the migration helpers, then land this with focused compatibility coverage.
Move _build_design_dataset helper and __getattr__ deprecation forwarding into BaseExperiment, replacing per-class @Property blocks with a declarative _deprecated_design_aliases dict. Fix convex_hull.py and maketables_adapters.py to use the new API directly. Add parametrized backward-compatibility tests covering all deprecated aliases. Made-with: Cursor
drbenvincent
left a comment
drbenvincent
left a comment
There was a problem hiding this comment.
Follow-up adversarial review from gpt-5.4-xhigh.
No new blocking findings from me on the latest revision.
What satisfies my earlier review:
- The duplicate migration boilerplate is now materially reduced. Moving deprecated alias forwarding into
BaseExperiment.__getattr__with a declarative_deprecated_design_aliasesmapping is the kind of centralization I was asking for, and_build_design_dataset()usefully consolidates the commondesign["X"]/design["y"]wrapping path. - The self-warning regression is fixed.
causalpy/checks/convex_hull.pynow readspre_design[...]directly, andcausalpy/maketables_adapters.pyalso no longer relies on deprecated design aliases internally. - The backward-compatibility contract is now tested in the right way.
causalpy/tests/test_deprecated_design_aliases.pychecks both deprecation behavior and data identity, and it adds a warning-free test for the convex-hull path. - The remote results now support the implementation:
codecov/patchis green, bothtestjobs are green, notebooks are green, docs are green, andprekis green.
I also ran the two most relevant local test targets from this follow-up:
pytest causalpy/tests/test_deprecated_design_aliases.py -qpytest causalpy/tests/test_maketables_plugin.py -q
Both passed locally. The standalone invocations tripped the repo-wide coverage floor, which is expected for narrow pytest runs here and not a concern for the PR itself.
Residual note, not a blocker: the shared dataset helper currently covers the standard X/y layout but not the SyntheticControl control/treated layout, so there is still some experiment-specific wrapping code there. That no longer looks like problematic duplication to me; the important compatibility and warning-forwarding boilerplate is now centralized.
This satisfies my earlier objections.
|
|
||
| _default_model_class: type[PyMCModel] | None = None | ||
|
|
||
| _deprecated_design_aliases: dict[str, tuple[str, str]] = {} |
There was a problem hiding this comment.
Quick one — InversePropensityWeighting, InstrumentalVariable, and StaggeredDiD didn't get migrated to design (still using numpy self.X / self.y, see e.g. causalpy/experiments/inverse_propensity_weighting.py:111, instrumental_variable.py:155, staggered_did.py:332-338) and accordingly don't show up in _deprecated_design_aliases here. Assuming intentional given their non-standard layouts (IPW has t + outcome, IV has Z + endogenous treatment, staggered has the train/full split)? If so, maybe worth a sentence in the PR body so it's not mistaken for an oversight.
There was a problem hiding this comment.
Yes, intentional — IPW, IV, and StaggeredDiD have non-standard design layouts (IPW pairs covariates with a treatment vector, IV has an instrument matrix plus endogenous treatment, StaggeredDiD keeps a train/full split) that don't fit the shared X/y Dataset shape. I've added a sentence to the PR body making this explicit. Note that SyntheticDifferenceInDifferences (added on main in #823 after this branched) did fit the SC-style four-quadrant pattern, so it has now been migrated to pre_design/post_design with the same deprecated aliases.
| self.pre_design["treated"].isel(treated_units=i).values, | ||
| self.pre_design["control"].values, |
There was a problem hiding this comment.
Heads-up on the input convention: this call site converts to numpy via .values, but the new ConvexHullCheck.run in causalpy/checks/convex_hull.py:60-62 hands the same data to check_convex_hull_violation as xarray.DataArray directly (no .values). Both work — I checked end-to-end with violation cases and got identical results — but it's worth normalizing on one style across the two call sites and adding a one-liner in the helper's docstring saying it accepts numpy or xarray. Otherwise this becomes a quiet footgun the first time someone changes the helper to assume one shape.
There was a problem hiding this comment.
Good catch — normalized both call sites on xarray: SyntheticControl._check_convex_hull now passes the DataArrays directly (no .values), matching ConvexHullCheck.run, and check_convex_hull_violation's docstring and type hints now state it accepts np.ndarray or xr.DataArray.
| datapre_control = sc.pre_design["control"] # type: ignore[attr-defined] | ||
| datapre_treated = sc.pre_design["treated"] # type: ignore[attr-defined] |
There was a problem hiding this comment.
Tiny nit: since applicable_methods = {SyntheticControl} and validate() already enforces the type, you could drop both # type: ignore[attr-defined] markers by adding assert isinstance(experiment, SyntheticControl) (or just calling self.validate(experiment)) at the top of run() — that narrows sc to SyntheticControl and pre_design becomes a known attribute for mypy. Same effect, no escape hatches.
There was a problem hiding this comment.
Done — run() now calls self.validate(experiment) followed by an isinstance assert to narrow the type, and both type: ignore markers are gone. mypy passes via prek.
| for name in ("X", "y"): | ||
| if name not in self.named_vars: | ||
| raise ValueError( | ||
| f"Data node '{name}' not found in model. " |
There was a problem hiding this comment.
Optional polish: the class docstring above already points users at BayesianBasisExpansionTimeSeries as a concrete override example. Worth echoing that in the runtime message itself so users hitting this in a notebook don't have to dig — e.g. f"... override _data_setter() (see BayesianBasisExpansionTimeSeries for an example).". Pure DX nicety, not blocking.
There was a problem hiding this comment.
Done — the runtime message now ends with "override _data_setter() (see BayesianBasisExpansionTimeSeries for an example)."
…y-datasets Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # causalpy/experiments/base.py # causalpy/experiments/diff_in_diff.py # causalpy/experiments/panel_regression.py # causalpy/experiments/piecewise_its.py # causalpy/pymc_models.py
- Migrate SyntheticDifferenceInDifferences (added on main in #823) to the pre_design/post_design Dataset pattern with deprecated aliases, matching SyntheticControl, and extend the alias test suite to cover it. - Normalize convex-hull call sites on xarray inputs and document that check_convex_hull_violation accepts numpy or xarray. - Replace type: ignore escape hatches in ConvexHullCheck.run with an isinstance assertion that narrows the type for mypy. - Point users at BayesianBasisExpansionTimeSeries in the _data_setter error message. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2 participants
Summary
xr.DataArrayattributes on 9 experiment classes intoxr.Datasetobjects, reducing attribute sprawl while keeping related design matrices together.pre_design/post_designDatasets.designDataset.BaseExperiment.__getattr__driven by a declarative_deprecated_design_aliasesmapping) that emitsDeprecationWarningfor backward compatibility.reporting.py,maketables_adapters.py,checks/convex_hull.py, and tests to use the new Dataset access patterns internally, so CausalPy never warns against itself.InversePropensityWeighting,InstrumentalVariable, andStaggeredDiDare intentionally not migrated: their design matrices have non-standard layouts (IPW pairs covariates with a treatment vectort, IV has an instrument matrixZplus an endogenous treatment, StaggeredDiD maintains a train/full split) that don't fit the sharedX/yDataset shape, so they keep their existing attributes with no deprecation shims.docs/source/_static/classes.png/packages.png) regenerated viamake umlto reflect the consolidated attributes.Closes #199. Follow-up notebook migration tracked in #848. Removal of the deprecation shims is tracked in draft PR #958 (do not merge that until a release has shipped with the warnings in place).
Test plan
test_deprecated_design_aliases.py) covering every alias on every migrated classprek run --all-filespasses (lint, format, mypy, codespell, notebook validation)Made with Cursor