Add pymc_forecast model-provider adapter for InterruptedTimeSeries

[twiecki]

Closes #1013.

Implements the ITS-only first pass: a thin adapter that lets a pymc_forecast forecasting model act as a model provider behind CausalPy's existing experiment API, in the same spirit as the sklearn wrapper. CausalPy keeps identification, counterfactual construction, and placebo methods; pymc_forecast provides the fitted forecasting model.

What's in here

• causalpy/pymc_forecast_models.py — PyMCForecastModel, a thin wrapper (comparable in size to skl_models.py) around a pymc_forecast model function / ForecastingModel plus a forecaster class (NUTS HMCForecaster by default, matching CausalPy's native backends; ADVI/Pathfinder are explicit opt-ins). Maps the experiment protocol onto the upstream 0.1.0 API exactly as scoped in the issue:
  • fit(X, y) → construct + fit the forecaster on the pre-period; the patsy design matrix rides along as covariates (obs_ind → time, coeffs → covariate).
  • predict(X) (pre) → predict_in_sample(), posterior_predictive group.
  • predict(X, out_of_sample=True) (post, "as if untreated") → forecast(future_covariates=post_X) with exogenous regressors, or forecast(future_index=post_index) for a covariate-free trend model.
  • Draw-level samples come out of prediction_samples(result) and the documented schema dims are renamed onto CausalPy coords (time_future → obs_ind, series → treated_units).
  • Priors flow through the model object (pymc-extras Prior), keeping the transparent-prior ethos.
• PyMCForecastAdapter in experiments/model_adapter.py, gated by a new supports_pymc_forecast flag — enabled for InterruptedTimeSeries only (SyntheticControl is a tracked follow-up). No change to the public API/UX: users just pass model=PyMCForecastModel(...) to cp.InterruptedTimeSeries.
• Optional extra, not a hard dep: pip install causalpy[forecast] → pymc-forecast>=0.1.0,<0.2 (minor pinned while upstream is 0.x; the output schema is the surface upstream contract-tests). Also added to the test extra so CI exercises the adapter; the test module importorskips it.
• Round-trip tests (test_pymc_forecast_adapter.py, 12 tests): fit pre / forecast post-as-untreated / calculate_impact on draw-level samples, with the output contract and recovered effect checked against the existing PyMC LinearRegression path on the same synthetic data; plus the covariate-free future_index path, three-period designs, plot/summary()/effect_summary()/get_plot_data_bayesian() smoke, single-unit validation, and adapter resolution/gating.
• Docs: module docstring carries the "when to reach for this backend vs the existing PyMCModel classes" guidance (referenced from the ITS docstring), and the module is added to the API reference index.

Out of scope (per the issue)

No outsourcing of placebo-in-time, power analysis, or power curves — the boundary stays at model-provider + forecasting metrics.

Caveats

Coefficient printing reports posterior summaries of the model's scalar parameters instead of design-matrix coefficients.

pymc_forecast exposes the draw-level posterior predictive of the observed variable rather than a separate noise-free mu ... Resolved with pymc-forecast 0.2: mu now carries the upstream noise-free latent predictor (mu/mu_future), so causal impact excludes observation noise exactly like the native backends, one posterior subsample is shared by the pre-period fit and the counterfactual (draw coherence via posterior=), and the forecaster uses the deferred-fit lifecycle. Pinned pymc-forecast>=0.2,<0.3.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 causalpy | 🛠️ Build #33644297 | 📁 Comparing 780bb1d against latest (71be9df)

  🔍 Preview build  

56 files changed · + 15 added · ± 41 modified

+ Added

• _modules/causalpy/pymc_forecast_models.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.coefficients.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.fit.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.predict.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.print_coefficients.html
• api/generated/causalpy.experiments.model_adapter.PyMCForecastAdapter.score.html
• api/generated/causalpy.pymc_forecast_models.PyMCForecastModel.calculate_cumulative_impact.html
• api/generated/causalpy.pymc_forecast_models.PyMCForecastModel.calculate_impact.html
• api/generated/causalpy.pymc_forecast_models.PyMCForecastModel.fit.html
• and 5 more...

± Modified

• 404.html
• genindex.html
• py-modindex.html
• _modules/index.html
• api/index.html
• api/generated/causalpy.checks.placebo_in_time.PlaceboInTime.validate.html
• api/generated/causalpy.checks.placebo_in_time.html
• api/generated/causalpy.experiments.base.BaseExperiment.html
• api/generated/causalpy.experiments.diff_in_diff.DifferenceInDifferences.html
• api/generated/causalpy.experiments.diff_in_diff.DifferenceInDifferences.plot.html
• and 31 more...

[drbenvincent]

Automated triage

Recommendation: review:high — no decision gate identified.

Why:

• Public API + behaviour change: Adds a new first-class causal backend (PyMCForecastModel) exposed via causalpy.__init__ and integrated into InterruptedTimeSeries (9 files, ~820 lines). This is a new integration surface, not a contained internal refactor.
• Dependency change: Introduces pymc-forecast>=0.1.0,<0.2 as an optional extra in pyproject.toml and environment.yml, changing the install/CI-surface for the project.
• CI failure: The test (3.11) job failed and test (3.14) was cancelled. The cause is not obvious from the metadata — needs confirmation whether this is flaky or a genuine regression introduced by the new adapter.

Review focus:

1. Confirm the failing CI (test (3.11)) — is it related to the new adapter or pre-existing?
2. Verify that the adapter boundary (fit/predict schema mapping) correctly isolates CausalPy identification logic from pymc_forecast model fitting, especially the mu/observation-noise caveat documented in the module docstring.
3. Check that the optional-extra dependency pin (>=0.1.0,<0.2) is appropriate for the upstream's pre-1.0 API stability.

Confidence: high

[twiecki]

The upstream changes this adapter requested are now in: pymc-labs/pymc-forecast#44 implemented all four feature requests (pymc-labs/pymc-forecast#36 noise-free mu/mu_future in the output schema, pymc-labs/pymc-forecast#37 posterior= passthrough for draw-coherent predictions, pymc-labs/pymc-forecast#38 ADVI non-convergence warning, pymc-labs/pymc-forecast#39 uniform progressbar= + deferred fit) and is merged; the 0.2.0 release PR (pymc-labs/pymc-forecast#45) is green and awaiting merge + PyPI publish.

This PR is updated accordingly (latest commit):

• Noise-free impact: mu is now the upstream latent predictor, y_hat the posterior predictive — the documented observation-noise caveat is gone, and impact HDIs match the native backends' convention. A new test locks std(mu) < std(y_hat) and std(impact) == std(mu).
• Draw coherence: one posterior subsample is drawn at fit time and passed to both predict_in_sample and forecast via posterior=; a new test verifies pre/post mu reproduce X @ beta from the shared draws draw-for-draw and that repeated predict() calls are bit-identical.
• Deferred fit: PyMCForecastModel now constructs the (unfitted) forecaster at construction and calls fit(data, covariates), dropping the carried (class, kwargs) recipe; invalid options fail early. The default forecaster resolves to StatespaceForecaster for StatespaceModel definitions.
• Pins bumped to pymc-forecast>=0.2,<0.3.

Note: CI will stay red on the dependency resolve until pymc-forecast==0.2.0 is published to PyPI (release PR pymc-labs/pymc-forecast#45 is ready). All 14 adapter tests pass locally against 0.2.0.

[codecov]

Codecov Report

❌ Patch coverage is 97.44318% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.58%. Comparing base (71be9df) to head (780bb1d).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
causalpy/pymc_forecast_models.py	92.17%	5 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1014      +/-   ##
==========================================
+ Coverage   95.51%   95.58%   +0.06%     
==========================================
  Files          98      102       +4     
  Lines       15870    16324     +454     
  Branches      931      955      +24     
==========================================
+ Hits        15159    15604     +445     
- Misses        504      509       +5     
- Partials      207      211       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[drbenvincent]

Thanks for this — the core direction looks right. pymc-forecast should remain a distinct, experiment-gated backend rather than inherit from PyMCModel: it owns a forecaster with a different fit/forecast lifecycle, coefficient contract, and inference-result shape. Once its outputs are normalized, treating it as Bayesian through is_bayesian is appropriate. supports_pymc_forecast is a reasonable incremental gate for the ITS-only first pass; before adding further providers, we should replace the parallel booleans with an explicit set of supported backend kinds.

Three areas need addressing before merge:

1. Make the existing CausalPy placebo-in-time path work

The PR/issue says that placebo methods remain CausalPy's responsibility, but PlaceboInTime.validate() currently rejects PyMCForecastModel because it checks isinstance(model, PyMCModel). Refitting also fails because PyMCForecastModel has no _clone() and its fallback deepcopy() raises TypeError: cannot pickle 'module' object.

This is a small CausalPy-only fix, not a new upstream feature: add an _clone() that reconstructs an unfitted wrapper from its stored configuration, validate against the backend's Bayesian capability, and add coverage that a forecast-backed ITS can be cloned and accepted/refitted by PlaceboInTime. Track broader support for other sensitivity checks separately.

2. Do not alias noisy statespace predictions to mu

For StatespaceForecaster, the adapter currently falls back from a missing noise-free expected observation to y_hat, then exposes that as mu. ITS consequently computes causal impact with observation noise included, contrary to CausalPy's mu-based impact contract.

This requires a new upstream pymc-forecast capability: statespace prediction currently exposes latent-state draws and noisy observed draws, but not the expected observation in observation space with measurement error removed. Please open and link an upstream issue requesting draw-level in-sample and future expected-observation outputs that:

• apply the state-space observation equation rather than returning the latent state vector;
• exclude measurement error;
• remain conditioned on the same supplied posterior= draws;
• preserve matching chain/draw/time/series coordinates; and
• align with the core forecaster's mu / mu_future schema where possible.

Until that capability exists, please reject or stop advertising the statespace path here rather than silently substituting y_hat for mu. Core pymc-forecast models built with predict() already provide proper mu / mu_future and can remain supported.

3. Complete the new backend's central contracts

• Update ARCHITECTURE.md; it still describes two backends and its maintenance rule explicitly requires updates when backend dispatch changes.
• Update the central BaseExperiment and ModelAdapter.model annotations/docstrings, which still exclude PyMCForecastModel.
• Align the optional dependency with the documented API. Upstream places statespace/Pathfinder support behind pymc-forecast[extras] with pymc-extras>=0.10, while this PR installs bare pymc-forecast and CausalPy permits pymc-extras>=0.3.0.
• Make inference diagnostics discoverable. The full HMC/Pathfinder fit remains available at model.forecaster.idata while result.idata contains the draw-coherent prediction subset; document that distinction or expose a direct fit_idata/fit_result property. No result.idata redesign is required in this PR.

Maketables support can be a separate CausalPy issue. Forecasting parameters do not necessarily map to Patsy coefficients, so that needs an explicit table-semantics decision rather than just another dispatch branch. For this PR, documenting the limitation and raising an intentional error is enough.

The distribution/import naming (pymc-forecast / pymc_forecast), optional minor pin, draw-coherence work, output normalization, and ITS integration tests all look appropriate. CI is green. With the placebo/refit path, statespace semantic boundary, and central architecture contracts addressed, this will establish a sound direction for future forecast-backed experiments.

[twiecki]

Thanks for the review @drbenvincent — all four points addressed in 9f09035:

1. PlaceboInTime: validate() now checks the backend capability (_model_backend.is_bayesian) instead of isinstance(model, PyMCModel), and PyMCForecastModel gained _clone() so clone_model() reconstructs an unfitted wrapper from its configuration. Covered by a real 2-fold PlaceboInTime.run() test asserting each fold received a fresh, refitted PyMCForecastModel.

2. Statespace mu aliasing: went with reject-plus-upstream-issue rather than silently substituting the noisy predictive. StatespaceModel / StatespaceForecaster now raise NotImplementedError at construction, the mu/mu_future lookups are strict (a missing latent fails loudly), and the request for draw-level expected observations excluding measurement error is filed as StatespaceForecaster: expose noise-free expected observations (mu/mu_future) excluding measurement error pymc-forecast#50 — once that ships, the rejection lifts.

3. Docs/annotations: ARCHITECTURE.md now covers the third backend (intro, module map, a dedicated backend paragraph, and the supports_pymc_forecast note in the authoring checklist); BaseExperiment / ModelAdapter.model annotations and docstrings include PyMCForecastModel. Dependency pins stand at pymc-forecast>=0.2,<0.3 in both the forecast and test extras, matching the API the module documents.

4. Diagnostics discoverability: added a fit_idata property exposing the full forecaster fit result (complete NUTS InferenceData with sample stats), documented against idata (the thinned, draw-coherent posterior subsample used for prediction). Raises informatively when unfitted or when the forecaster doesn't retain an InferenceData (variational fits point to .forecaster.approx / .losses).

118 tests green locally across the adapter, placebo-in-time, model-adapter, panel-regression, and guard-script suites, plus doctests and the full hook set (mypy, numpydoc, ruff, architecture/exports).

🤖 Generated with Claude Code

[drbenvincent]

Thanks @twiecki — I rechecked the follow-up commit and the substantive review points are addressed: placebo refitting is covered, statespace now fails safely with pymc-forecast#50 tracking the missing capability, and the architecture, typing, diagnostics, tests, and current CI look good. The upstream statespace issue does not need to block this PR.

The one remaining item is the optional dependency contract: the adapter advertises Pathfinder, but the extras still install bare pymc-forecast while allowing pymc-extras>=0.3; upstream declares these integrations through pymc-forecast[extras] with pymc-extras>=0.10. Please either use pymc-forecast[extras]>=0.2,<0.3 in the forecast/test extras or confirm and document why the lower floor is supported. Once that is resolved and the updated CI remains green, I’m happy to approve this for merge.

[twiecki]

@drbenvincent done.

[drbenvincent]

One final scope clarification before approval, following a broader look at CausalPy's prediction semantics.

Current behavior: CausalPy's calculate_impact() interprets mu as the conditional expected outcome in observed outcome units and computes observed - mu. Most existing built-in models satisfy that convention because they use Gaussian identity links. However, this is currently an implicit convention rather than an enforced backend contract.

Boundary exposed by this PR: pymc-forecast documents mu / mu_future as the latent passed to predict(), which may be a link-scale linear predictor for a GLM. The Gaussian identity-link path implemented and tested here is correct. But for a linked model, directly consuming a link-scale mu would mix units—for example, observed counts minus a log-rate.

No adapter redesign or upstream release should block this PR. Please add a short note documenting the current requirement: models using a link function must pass the inverse-linked conditional expectation in observed outcome units to pymc_forecast.predict(). Once that scope is explicit and CI remains green, I am happy to approve.

We have opened follow-up issues so the temporary convention and the intended long-term design are historically connected:

• Define outcome-scale prediction semantics across model backends #1016 — define outcome-scale prediction semantics across CausalPy backends, audit current behavior, and add contract tests.
• Analyze customer-support ticket counts with CausalPy InterruptedTimeSeries #1017 — add a g-computation tutorial showing correct and incorrect causal contrasts with link functions.
• Expose outcome-scale expected observations separately from latent predictors pymc-forecast#52 — expose an explicit outcome-scale expected observation separately from the generic latent predictor.
• Require sorted, unique time indexes for InterruptedTimeSeries #1018 — require sorted, unique ITS indexes; this is independent of the link-function contract but was exposed while reviewing the forecasting integration.

The existing statespace limitation remains tracked in pymc-labs/pymc-forecast#50.

[drbenvincent]

FYI: No action needed from @twiecki. I'm going to get an agent to blast through this and the issues it generated