-
Notifications
You must be signed in to change notification settings - Fork 114
Lazy experiment lifecycle: configure → optional prior checks → fit() #1175
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.
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
base: pymc6_and_pymcmarketing1_migration
Are you sure you want to change the base?
Changes from 10 commits
6157ce6
232b8b3
b1503be
e98c1a1
d801d33
c7c2e07
120981d
e8caf65
859df08
021a0af
20351ab
1a5f0f5
5753122
bb15741
04a2d58
410dbc4
2964544
af0830d
b83a015
4b064b8
7e1a63c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -66,7 +66,15 @@ The optional third backend, `PyMCForecastModel` (`causalpy/pymc_forecast_models. | |
|
|
||
| ## Experiment Lifecycle | ||
|
|
||
| Instantiation fits eagerly in `__init__`: `_build_design_matrices()` → `_prepare_data()` → `algorithm()`. There is no separate `.fit()` on the experiment. Each subclass's public `plot(*, ...)` delegates to `_render_plot()`, which calls the subclass's backend-agnostic `_plot()`. Uncertainty rendering keys on data properties of the canonical prediction container (`has_posterior_draws()`), not on backend identity. `effect_summary()` returns `EffectSummary(table, text)` using helpers from `causalpy.reporting`. | ||
| Construction is lazy: `__init__` runs validation and deterministic preprocessing only (`_build_design_matrices()` → `_prepare_data()`); nothing is sampled. The lifecycle is `configure → optional prior checks → fit()`: | ||
|
|
||
| - **`build()`** — public, idempotent, auto-called by both samplers. Merges data-driven priors (defaults → data-derived → user) and constructs the PyMC graph so the spec is inspectable (`pm.model_to_graphviz(exp.model)`, `exp.model.basic_RVs`) before any compute is spent. | ||
| - **`sample_prior_predictive(**kwargs)`** — optional prior phase; populates `exp.prior_result`. Raises `PriorPredictiveNotSupportedException` on backends without a prior phase (sklearn, pymc-forecast, IV, state space). | ||
| - **`fit(**kwargs)`** — posterior phase; runs NUTS + posterior predictive first (forward-sampling machinery mutates shared data nodes, so posterior must sample on freshly armed data), then fills the absent prior phase so `idata` matches the historical eager output. Returns `Self`, so call sites migrate with one appended token: `cp.InterruptedTimeSeries(...).fit()`. | ||
|
|
||
| Draw-derived state lives in per-experiment bundles (`causalpy/experiments/_results.py`) exposed through two raising properties: `exp.result` (posterior group) and `exp.prior_result` (prior group). Nothing derived from draws is assigned to the experiment itself; `is_fitted` / `has_prior_predictive` are predicates over those two slots. Read methods take keyword-only `group: Literal["prior", "posterior"]`; the guard and group→bundle resolution live once in `_render_plot()` / `_resolve_group()`. Missing groups raise `GroupNotSampleedException` naming the call to make; both exceptions are re-exported from `causalpy`. Re-sampling overwrites its own group only (`fit()` again replaces posterior and warns; prior state survives). Assigning a new model — also the documented prior-revision path instead of a `set_priors()` — resets everything, because graph identity *is* the model instance. | ||
|
|
||
| Each subclass's public `plot(*, ...)` delegates to `_render_plot()`, which calls the subclass's backend-agnostic `_plot(bundle=..., group=...)`. Uncertainty rendering keys on data properties of the canonical prediction container (`has_posterior_draws()`), not on backend identity. `plot(group="prior")` renders a reduced panel set (counterfactual vs observed only). `effect_summary(group=...)` returns `EffectSummary(table, text)` using helpers from `causalpy.reporting`; prior-group prose reads as a plausibility check, not a causal claim. | ||
|
|
||
| ## Experiment Inventory | ||
|
|
||
|
|
@@ -98,21 +106,19 @@ Instantiation fits eagerly in `__init__`: `_build_design_matrices()` → `_prepa | |
|
|
||
| | Topic | Detail | | ||
| |-------|--------| | ||
| | **Lazy fitting** | `__init__` never samples. Explicit `fit()` / `sample_prior_predictive()` run the phases; see Experiment Lifecycle. | | ||
| | **Formulas** | Patsy `dmatrices()` for design matrices; `build_design_matrices()` for counterfactual prediction. Bare datetime predictors are encoded as continuous elapsed days from the fitted origin; use `C(date)` for date fixed effects. `PiecewiseITS` uses `step()`/`ramp()` stateful transforms. | | ||
| | **obs_ind** | All experiments set `data.index.name = "obs_ind"`. Canonical xarray/PyMC dimension name. | | ||
| | **treated_units always 2D** | Even single-unit experiments use `treated_units=["unit_0"]`. Never pass 1D y to PyMC. | | ||
| | **Impact uses mu, not y_hat** | The adapter's `predict()` extracts posterior `mu` (conditional expected outcome in observed units), not `y_hat` (with observation noise); impact is `y - predict(X)`. For GLMs, `mu` must be inverse-linked before impact; see `docs/source/knowledgebase/prediction-contract.md`. | | ||
| | **Intercept handling** | Patsy includes intercept by default. sklearn models must use `fit_intercept=False`. | | ||
| | **Eager fitting** | MCMC runs during `__init__`. No lazy `.fit()` on the experiment. | | ||
| | **HDI_PROB** | Project default is 0.94 (ArviZ default), not 0.95. | | ||
| | **create_causalpy_compatible_class** | Applied during `make_model_adapter()` for sklearn backends; clones the user instance before patching. | | ||
|
|
||
| ## Adding New Code | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The That's still true, still surprising (most people assume 0.95), and this was the only place it was written down. Looks like collateral from removing the "Eager fitting" row right next to it. Fix: put the row back. |
||
|
|
||
| Copy the closest existing experiment or model and follow the `BaseExperiment` contract: | ||
|
|
||
| - Declare `supports_ols` / `supports_bayes` (and `supports_pymc_forecast` to opt into the optional pymc-forecast backend); implement a single backend-agnostic `_plot()` (and an explicit `get_plot_data(*, ...)` only where that view is supported) that consumes the canonical prediction container, keying uncertainty rendering on `has_posterior_draws()` rather than backend identity | ||
| - `algorithm()` with the fit/predict/impact flow; every concrete experiment declares its own explicit `effect_summary(...)` contract, using helpers in `causalpy.reporting` where that summary is implemented | ||
| - Implement `_fit_inputs()` returning the `(X, y, coords)` handed to the backend at build time, and `_finalize(group)` computing the experiment's result bundle from the group's draws; every concrete experiment declares its own explicit `effect_summary(*, group=..., ...)` contract, using helpers in `causalpy.reporting` where that summary is implemented | ||
| - Public APIs expose explicit named parameters rather than bare `*args` / `**kwargs`; use keyword-only optional controls for public plotting and plot-data APIs (enforced by `causalpy/tests/test_public_signatures.py` and surveyed by `scripts/audit_public_signatures.py`). A genuine dynamic or third-party forwarder requires an `Other Parameters` contract and a narrow structural-test exemption. For experiments without a unified plot view (e.g. `InversePropensityWeighting`, `InstrumentalVariable`), declare an explicit `plot()` stub that raises `NotImplementedError`. For `hdi_prob` defaults, use ``Defaults to :data:`~causalpy.constants.HDI_PROB` (currently 0.94).`` in the docstring. | ||
|
anevolbap marked this conversation as resolved.
|
||
| - Raise `FormulaException`, `DataException`, or `BadIndexException` from `causalpy.custom_exceptions` for formula, data, and index errors | ||
| - Avoid backwards-compat shims for APIs introduced in the same PR | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| # Copyright 2022 - 2026 The PyMC Labs Developers | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
| """Per-experiment result bundles for the lazy experiment lifecycle. | ||
|
|
||
| Every bundle is *fully populated* whenever it exists: an experiment either has | ||
| no bundle yet (the corresponding lifecycle verb has not run) or a complete one. | ||
| Nothing derived from model draws lives directly on the experiment object; the | ||
| two public slots ``experiment.result`` (posterior group) and | ||
| ``experiment.prior_result`` (prior group) hold these bundles and back the | ||
| ``is_fitted`` / ``has_prior_predictive`` state predicates. | ||
|
|
||
| All prediction/impact fields are typed :class:`xarray.DataArray` on the | ||
| canonical ``("chain", "draw", "obs_ind"[, ...])`` dimensions produced by | ||
| :class:`~causalpy.experiments.model_adapter.ModelAdapter.predict`. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| from dataclasses import dataclass, field | ||
|
|
||
| import pandas as pd | ||
| import xarray as xr | ||
|
|
||
| __all__ = [ | ||
| "CausalResult", | ||
| "CoefficientResult", | ||
| "DiscontinuityResult", | ||
| "GroupComparisonScenario", | ||
| "KinkResult", | ||
| "StaggeredDifferenceInDifferencesResult", | ||
| "SyntheticDifferenceInDifferencesResult", | ||
| ] | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class CausalResult: | ||
| """Result bundle for predict-contrast experiments. | ||
|
|
||
| Used by :class:`~causalpy.experiments.interrupted_time_series.InterruptedTimeSeries`, | ||
| :class:`~causalpy.experiments.synthetic_control.SyntheticControl`, and | ||
| :class:`~causalpy.experiments.piecewise_its.PiecewiseITS`. | ||
|
|
||
| For ``PiecewiseITS`` ``predictions_pre`` carries the fitted expectation over | ||
| the full observation window and ``predictions_post`` / ``impact_post`` / | ||
| ``impact_post_cumulative`` carry the post-first-interruption slices consumed | ||
| by the reporting helpers. | ||
| """ | ||
|
|
||
| predictions_pre: xr.DataArray | ||
| predictions_post: xr.DataArray | ||
| impact_pre: xr.DataArray | ||
| impact_post: xr.DataArray | ||
| impact_post_cumulative: xr.DataArray | ||
| score: pd.Series | None = None | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class SyntheticDifferenceInDifferencesResult(CausalResult): | ||
| """Result bundle for :class:`~causalpy.experiments.synthetic_difference_in_differences.SyntheticDifferenceInDifferences`. | ||
|
|
||
| The ``CausalResult`` prediction/impact fields are reconstructed from the | ||
| synthetic-control imputation; ``tau_posterior`` carries the analytic | ||
| double-difference treatment-effect draws with dimensions | ||
| ``("chain", "draw")``. | ||
| """ | ||
|
|
||
| tau_posterior: xr.DataArray = field(kw_only=True) | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class GroupComparisonScenario: | ||
| """One scenario plotted or summarized by a group-comparison experiment.""" | ||
|
|
||
| inputs: pd.DataFrame | ||
| prediction: xr.DataArray | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class CoefficientResult: | ||
| """Result bundle for coefficient-contrast experiments (DiD, PrePostNEGD). | ||
|
|
||
| ``causal_impact`` holds draws of the treatment-effect coefficient (or its | ||
| algebraically-equivalent prediction contrast) with canonical coefficient | ||
| dimensions. | ||
| """ | ||
|
|
||
| causal_impact: xr.DataArray | ||
| scenario_control: GroupComparisonScenario | ||
| scenario_treated: GroupComparisonScenario | ||
| scenario_counterfactual: GroupComparisonScenario | None = None | ||
| score: pd.Series | None = None | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class DiscontinuityResult: | ||
| """Result bundle for :class:`~causalpy.experiments.regression_discontinuity.RegressionDiscontinuity`.""" | ||
|
|
||
| predictions: xr.DataArray | ||
| discontinuity_at_threshold: xr.DataArray | ||
| score: pd.Series | None = None | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class KinkResult: | ||
| """Result bundle for :class:`~causalpy.experiments.regression_kink.RegressionKink`.""" | ||
|
|
||
| predictions: xr.DataArray | ||
| gradient_change: xr.DataArray | ||
| score: pd.Series | None = None | ||
|
|
||
|
|
||
| @dataclass(frozen=True) | ||
| class StaggeredDifferenceInDifferencesResult: | ||
| """Result bundle for :class:`~causalpy.experiments.staggered_did.StaggeredDifferenceInDifferences`. | ||
|
|
||
| ``att_group_time`` and ``att_event_time`` are aggregated ATT tables; | ||
| ``y_pred`` retains the raw counterfactual draws so placebos and alternate | ||
| HDI levels can re-derive effects without resampling. | ||
| """ | ||
|
|
||
| att_group_time: pd.DataFrame | ||
| att_event_time: pd.DataFrame | ||
| y_pred: xr.DataArray | ||
| hdi_prob: float | ||
| score: pd.Series | None = None |
Uh oh!
There was an error while loading. Please reload this page.