`mapi_pymc`¶

This package is for PyMC models. You write build(). The package gives engines, diagnostics, new levels and counterfactuals.

The base class¶

PyMCModel: authors write build(); the platform fits, predicts and checkpoints (§5).

.. code-block:: python

class Radon(PyMCModel[RadonSettings]):
    id = "demo.radon"
    version = "1.0.0"
    Settings = RadonSettings
    schema_template = {...}
    outputs = {"y": {"dims": ["obs"], "units_of": "y"}}
    expected = {"y": "mu"}             # output -> Deterministic holding E[y]

    def build(self, ds, s, ctx):
        with pm.Model(coords=ctx.coords()) as m:
            floor = ctx.data("floor")
            county = ctx.data("county_idx")
            a = ctx.hier_effect("a", "county")
            b = pm.Normal("b", 0, 1)
            mu = pm.Deterministic("mu", a[county] + b * floor, dims="obs")
            pm.Normal("y", mu, pm.HalfNormal("sigma", 1), observed=ctx.observed("y"),
                      dims="obs")
        return m

fit dispatches to the engine in fit.engine (see :mod:mapi_pymc.engines), adds the opt-in groups (prior/prior_predictive, log_likelihood, a thinned posterior_predictive) and turns r-hat, ESS and divergences into diagnostics.

score rebuilds the model on the canonicalised score data and runs pm.sample_posterior_predictive(..., predictions=True) against the posterior:

output.kind = expected returns the Deterministic named in :attr:PyMCModel.expected (no observation noise; counterfactual deltas on it are exact per draw);
output.kind = predictive returns draws of the observed variable through icdf(dist, U) with frozen uniforms U seeded by (seed, name, shape). A counterfactual's scenario and baseline therefore share their noise (common random numbers) and the delta's variance is not inflated. Discrete distributions without a symbolic inverse CDF (Poisson, negative binomial, binomial) sample their parameters and push the same frozen uniforms through a numeric inverse CDF (noise = crn_ppf), so a change in one cell's mean never reshuffles the noise of the others. Only distributions with neither fall back to the seeded sampler stream (seeded_stream), whose draws in unchanged cells can still differ between a scenario and its baseline - prefer kind: expected deltas for those.

PyMCModel ¶

Bases: CustomModel[SettingsT]

Base class for PyMC plugins: implement :meth:build; declare outputs and expected.

Class attributes beyond :class:~mapi_sdk.models.CustomModel:

expected - {output: Deterministic name} for kind=expected predictions;
predictive - {output: observed RV name} (default: the output's own name);
default_engine - engine when the spec names none (default nutpie, else pymc);
fast_score_spec - a :class:~mapi_pymc.frozen.FrozenPredictorSpec enabling :meth:frozen_predictor (the manifest then declares fast_score);
prior_overlap_vars - free parameters whose prior draws (prior_overlap_draws, forward sampling only, so cheap) are kept as the prior group when the spec keeps no prior itself, so the sanity scorecard can measure how much the data moved them (M3, plausibility.prior_overlap). Empty (the default) keeps nothing.

build ¶

build(
    ds: Dataset, s: SettingsT, ctx: BuildContext
) -> Model

Return a pm.Model for ds (called at fit and again at every score).

build_model ¶

build_model(
    ds: Dataset, s: SettingsT, ctx: BuildContext
) -> Model

Call :meth:build and check that it returned a model with the declared outputs.

check_fast_score ¶

check_fast_score(
    fitted: Fitted, ds: Dataset, *, seed: int = 0
) -> str

Contract-kit hook (check_model's fast_score_matches_exact): the frozen predictor agrees with exact scoring on model-stage ds within MCSE.

Raises AssertionError when it does not (see :func:mapi_pymc.testing. check_fast_score); returns a one-line summary otherwise.

frozen_predictor ¶

frozen_predictor(
    fitted: Fitted,
    ds: Dataset,
    spec: FrozenPredictorSpec | None = None,
    *,
    num_samples: int | None = None,
    seed: int | None = None,
) -> FrozenPredictor

Compile a :class:~mapi_pymc.frozen.FrozenPredictor from fast_score_spec.

ds is a model-stage dataset fixing the non-input data (canonicalised score data or the fit data); the predictor then varies only spec.input_vars.

Build context¶

BuildContext: what a PyMC model's build() uses instead of raw coordinates (plan §5).

build(ds, settings, ctx) runs twice in a model's life: at fit on the training data, and at score on canonicalised score data (fit levels in fit order, then new levels, padded extend dims). The helpers hide the differences so the same build serves both:

:meth:BuildContext.coords - the pm.Model(coords=...) mapping; adds <dim>__fit and <dim>__new when the score data holds unseen levels;
:meth:BuildContext.data - pm.Data over the current values (rebuild-and-predict);
:meth:BuildContext.observed - the target at fit, a shape-only placeholder at score;
:meth:BuildContext.hier_effect - a hierarchical effect over a new_levels dim whose new levels are filled according to the dim's strategy (hyperprior draws them from the fitted hyperprior during prediction, with frozen noise so scenario and baseline agree);
:meth:BuildContext.position / :attr:BuildContext.anticipated_coords - stable positions along extend dims, so trends and HSGPs see the same axis at fit and at forecast time.

BuildContext.tree (M4) is the bound tree dataset for models whose capabilities declare input="tree" (FitContext.tree / ScoreContext.tree), else None.

BuildContext ¶

BuildContext(
    phase: Phase,
    ds: Dataset,
    schema: DataSchema | None = None,
    policies: Mapping[str, DimPolicy] = dict(),
    fit_coords: Mapping[str, ndarray] = dict(),
    seed: int = 0,
    mode: str | None = None,
    kind: Literal["expected", "predictive"] | None = None,
    data_vars: dict[str, tuple[str, ...]] = dict(),
    observed_vars: dict[str, tuple[str, ...]] = dict(),
    data_values: dict[str, ndarray] = dict(),
    observed_values: dict[str, ndarray] = dict(),
    hier_effects: dict[str, HierEffect] = dict(),
    frozen: dict[str, FrozenNoise] = dict(),
    tree: DataTree | None = None,
)

Platform helpers for PyMCModel.build (see the module docstring).

rng ¶

rng: Generator

A generator seeded from the run seed (fresh on each access, so reproducible).

anticipated_coords ¶

anticipated_coords: dict[str, ndarray]

For each extend dim: the fit labels plus the max_extension labels a forecast may reach. Use it to size HSGP domains or to normalise time at fit so the normalisation is identical at forecast time.

labels ¶

labels(dim: str) -> ndarray

Current labels along dim (positions when the dim has no coordinate).

new_levels ¶

new_levels(dim: str) -> ndarray

Levels of dim unseen at fit (always empty at fit).

fit_levels ¶

fit_levels(dim: str) -> ndarray

Levels of dim known at fit, in fit order.

strategy ¶

strategy(dim: str) -> str

The dim's new-level strategy (hyperprior when no schema says otherwise).

coords ¶

coords(ds: Dataset | None = None) -> dict[str, Any]

pm.Model coords for ds (default: the build data).

Every dim gets its labels. At score time a dim with unseen levels also gets <dim>__fit (the fit levels, which posterior arrays are indexed by) and <dim>__new (the unseen levels, appended after the fit levels).

data ¶

data(
    name: str,
    values: DataArray
    | ndarray
    | Sequence[Any]
    | float
    | None = None,
    *,
    dims: Sequence[str] | None = None,
) -> TensorVariable

Register pm.Data(name, ...) over the current values.

values defaults to ds[name]; dims defaults to its dims. At score time the model is rebuilt with the score values, so authors never call pm.set_data. Missing values are rejected (pm.Data cannot hold NaN): impute them in the pipeline or pass explicit values.

observed ¶

observed(
    name: str, *, dims: Sequence[str] | None = None
) -> ndarray

Observed values for a likelihood: the target at fit, zeros of the right shape at score (prediction never conditions on them; missing targets are not imputed).

align ¶

align(da: DataArray) -> DataArray

Put dims the model schema does not declare (channel in a panel whose template only names geo and time) back into their fit order.

The platform canonicalises declared dims; undeclared ones arrive in payload order. Positional models must see them in fit order, and with exactly the fit labels.

hier_effect ¶

hier_effect(
    name: str,
    dim: str,
    *,
    mu: Any = None,
    sigma: Any = None,
    centered: bool = False,
    mu_sigma: float = 1.0,
    sigma_scale: float = 1.0,
) -> TensorVariable

A partially pooled effect over dim: name[dim] ~ Normal(mu, sigma).

Creates name (dims dim) - non-centred by default via name__offset - plus hyperpriors name__mu ~ Normal(0, mu_sigma) and name__sigma ~ HalfNormal (sigma_scale) unless mu/sigma are given (mu may be a scalar or a vector over the full dim). When the score data holds unseen levels, name__new covers them according to the dim's new-level strategy:

hyperprior - mu + sigma * z with z frozen standard-normal noise, i.e. a fresh level drawn from the fitted hyperprior for every posterior draw;
pooled - the population mean mu;
zero - 0;
error - raises UNSEEN_LEVELS.

Returns the effect over the full current dim (fit levels, then new levels).

frozen_noise ¶

frozen_noise(
    name: str,
    *,
    dims: Sequence[str],
    kind: NoiseKind = "normal",
) -> TensorVariable

A latent filled with seeded noise at score time (standard normal or uniform).

Only meaningful when predicting; the platform supplies its values per posterior draw from (seed, name, shape), so repeated scores agree exactly.

position ¶

position(dim: str, *, scale: bool = False) -> ndarray

Positions of the current labels along dim relative to the first fit label.

Datetimes count periods of the dim's frequency, integers count steps of the fit spacing, categories use the fit order (new levels continue after it). With scale=True positions are divided by the last fit position so the fit range maps onto [0, 1] - identically at fit and at forecast time.

HierEffect ¶

HierEffect(
    name: str,
    dim: str,
    strategy: str,
    n_fit: int,
    n_new: int,
    centered: bool,
)

Book-keeping for one :meth:BuildContext.hier_effect call.

FrozenNoise ¶

FrozenNoise(
    name: str, kind: NoiseKind, dims: tuple[str, ...]
)

A latent the platform fills with seeded draws at score time (never sampled by MCMC).

Used for new-level offsets and for common-random-number predictive noise: the values depend only on the seed, the variable name and the shape, so a counterfactual's scenario and baseline share them exactly.

Engines¶

Inference engines for :class:~mapi_pymc.PyMCModel (plan §5).

fit.engine picks one of:

============== ========================================================================== nutpie nutpie NUTS (default when installed). Fast, parallel chains. Not bit-for-bit reproducible from a seed. Checkpoints are snapshots (draws so far): nutpie cannot export its adaptation, so a retry samples again from the start. pymc PyMC NUTS with segmented sampling: tune once, then draw in chunks of fit.checkpoint.every_draws. After each chunk a resumable checkpoint holds the draws, every chain's position and the full sampler state; a retry resumes exactly where the last chunk ended. Chains run in this process one after another (models with discrete variables fall back to pm.sample without segments). numpyro PyMC's external NUTS bridges, when the library is installed. No segments. blackjax map pm.find_MAP - one "draw" (chain 1, draw 1). advi mean-field ADVI; fit.tune is the number of optimisation steps (default 10,000) and fit.draws the number of draws from the approximation. pathfinder pymc-extras' Pathfinder when pymc-extras is installed (extra pathfinder: mapi-pymc[pathfinder]). One chain; no progress table on stdout. ============== ==========================================================================

Every engine reports progress through one callback (fraction, message) where the message names the chain ("chain 2/4: sampling 150/500, 0 divergences").

Tracking. When :attr:EngineHooks.tracker is set (:meth:EngineHooks.for_context sets it when the run records metrics), the engines also report every chain's progress (nutpie, pymc) and step metrics: step_size, num_steps and cumulative divergences per chain (label chain; step = draws done, tuning included; the pymc engine adds tree_depth and logs every k-th draw so that a chain has at most :data:MAX_HISTORY_POINTS points per metric), and the ADVI loss (negative ELBO) at most :data:MAX_HISTORY_POINTS times per fit. :func:pymc_draw_callback gives the same for a library that calls pm.sample itself. numpyro, blackjax, map and pathfinder report progress only.

Start points. By default every MCMC chain starts at the model's initial point with a random jitter. A plugin whose posterior has a trap near the default start (a mixture whose components all switch off) sets :attr:EnginePlan.initvals (values on the constrained scale, by free variable name, as pm.sample(initvals=...)) and jitter=False with dataclasses.replace(plan, ...): every chain then starts at exactly that point (pymc, nutpie and the external NUTS bridges; map, advi and pathfinder ignore both).

EnginePlan ¶

EnginePlan(
    engine: str,
    tune: int = DEFAULT_TUNE,
    draws: int = DEFAULT_DRAWS,
    chains: int = DEFAULT_CHAINS,
    seed: int = 0,
    target_accept: float | None = None,
    checkpoint_every: int | None = None,
    initvals: Mapping[str, Any] | None = None,
    jitter: bool = True,
)

Resolved sampler settings (FitConfig with defaults applied).

initvals ¶

initvals: Mapping[str, Any] | None = None

Start values by free variable name (constrained scale); None: the model's own.

jitter ¶

jitter: bool = True

Jitter the start of each chain (False: every chain starts at the same point).

EngineResult ¶

EngineResult(
    idata: DataTree,
    plan: EnginePlan,
    resumed_from: int = 0,
    checkpoints: int = 0,
    notes: list[str] = list(),
)

What an engine produced: ArviZ groups plus bookkeeping.

EngineHooks ¶

EngineHooks(
    progress: Callable[[float, str], None] | None = None,
    cancel: CancelToken | None = None,
    log: Callable[[str], None] | None = None,
    sink: CheckpointSink | None = None,
    key: str | None = None,
    min_interval: float = _PROGRESS_EVERY_S,
    tracker: Tracker | None = None,
    _last: float = float("-inf"),
    _chains_last: float = float("-inf"),
)

Callbacks an engine run uses: progress, cancellation, logging, checkpoints and tracking (chains and step metrics; tracker=None skips building them).

for_context ¶

for_context(
    ctx: FitContext,
    *,
    progress: Callable[[float, str], None] | None,
    sink: CheckpointSink | None = None,
    key: str | None = None,
) -> EngineHooks

Hooks wired to a plugin's FitContext (its cancel token, log and tracker).

report ¶

report(
    fraction: float, message: str, *, force: bool = False
) -> None

Forward progress, at most once per min_interval seconds unless forced.

metrics ¶

metrics(
    values: Mapping[str, float],
    *,
    step: int,
    labels: Mapping[str, str] | None = None,
) -> None

Log step metrics (no-op without a tracker).

chains ¶

chains(chains: Sequence[ChainProgress]) -> None

Report every chain's progress (no-op without a tracker).

chains_every ¶

chains_every(
    build: Callable[[], Sequence[ChainProgress]],
    *,
    force: bool = False,
) -> None

:meth:chains, at most once per min_interval seconds unless forced.

ChainProgress ¶

ChainProgress(
    chain: int,
    chains: int,
    done: int,
    total: int,
    tuning: bool,
    divergences: int = 0,
    step_size: float | None = None,
    num_steps: int | None = None,
)

Progress of one chain.

to_spec ¶

to_spec() -> ChainProgress

The run-protocol form (progress event chains).

run_engine ¶

run_engine(
    model: Model,
    plan: EnginePlan,
    hooks: EngineHooks | None = None,
) -> EngineResult

Fit model with plan.engine.

resolve_engine ¶

resolve_engine(
    name: str | None, default: str | None = None
) -> str

The engine to use: the spec's, else the plugin default, else nutpie (or pymc).

engine_available ¶

engine_available(name: str) -> bool

Whether name can run in this environment.

Fast scoring¶

FrozenPredictor: one compiled PyTensor function for fast posterior predictions (plan §5).

Ported from pymc-model-interactivity/frozen_predictor.py to PyMC 6 / PyTensor 3 and ArviZ 1.x (posteriors are xr.DataTree / xr.Dataset). Given a model, posterior samples, response expressions and the pm.Data inputs to vary, it compiles a single function vectorised over posterior samples and scenarios:

only the free random variables that are ancestors of the responses are bound, to constant posterior samples (so compilation stays small);
observed random variables upstream of a response are replaced by their observed values when they do not depend on the inputs, and by icdf(dist, U) with frozen uniform quantiles U (one per sample) when they do - common random numbers across scenarios;
SpecifyShape constraints on dims that only the inputs use are removed, so inputs may change length (dims shared with other model variables stay fixed).

Inputs may have the data variable's own shape (one scenario) or an extra leading scenario axis. Outputs have shape (sample, [scenario,] *dims) before the optional reducer.

The plan's FrozenPredictorSpec(response_exprs, input_vars, reducer) is the JSON form a plugin declares when capabilities.fast_score is true.

FrozenPredictorSpec ¶

Bases: BaseModel

Declarative fast-score recipe: {output: model variable}, the pm.Data inputs to vary, and how to reduce over posterior samples.

FrozenPredictor ¶

FrozenPredictor(
    fn: Callable[..., list[ndarray]],
    *,
    output_names: list[str],
    output_dims: dict[str, tuple[str, ...]],
    input_vars: list[str],
    input_dims: dict[str, tuple[str, ...]],
    input_ndim: dict[str, int],
    post_fns: list[Callable[[ndarray], ndarray] | None],
    samples: Dataset,
    symbolic: list[bool] | None = None,
    reducer: Callable[[ndarray], ndarray] | None = None,
)

A compiled predictor; call it with input arrays or use :meth:predict with xarray.

predict ¶

predict(
    inputs: Mapping[str, DataArray] | Dataset,
    *,
    scenario_dim: str = "scenario",
) -> Dataset

xarray in, xarray out: outputs as (sample, [scenario,] *dims) with chain/draw coordinates on sample (no reducers - use the raw call for those).

Reducer ¶

Bases: StrEnum

How to collapse the sample axis of each response.

create_frozen_predictor ¶

create_frozen_predictor(
    model: Model,
    posterior: DataTree | Dataset,
    response_exprs: Mapping[str, ExprSpec],
    input_vars: Sequence[str],
    *,
    group: str = "posterior",
    num_samples: int | None = None,
    random_seed: int | None = None,
    reducer: Reducer | str = NONE,
) -> FrozenPredictor

Compile a predictor of response_exprs at new values of input_vars.

response_exprs maps output names to a model variable name, a PyTensor expression of the model, or (expr, post_fn) where post_fn post-processes the numpy result. num_samples subsamples the posterior (seeded by random_seed); random_seed also seeds the frozen uniforms for input-dependent observed variables.

Checkpoints¶

Checkpoints for segmented sampling (plan §5, §8.3).

A long MCMC fit tunes once and then draws in chunks. After each chunk the engine hands a :class:SamplingCheckpoint to a :class:CheckpointSink; a retried run with the same fit key resumes from the last checkpoint instead of starting again.

Two kinds of checkpoint exist:

resumable (engine pymc): the draws so far, every chain's current position and the full NUTS sampling state (adapted step size, mass matrix, RNG state). Resuming continues the same chains exactly; the result equals an uninterrupted run with the same seed.
snapshot (engine nutpie): the draws so far only. nutpie cannot export or restore its adaptation state, so a snapshot feeds a PARTIAL bundle but a retry starts again.

Nothing here pickles. Sampling states are encoded to JSON by :func:encode_state; arrays travel base64-encoded and dataclasses are rebuilt only from classes defined in an allow-list of PyMC step-method modules (and this package).

CheckpointSink ¶

Bases: Protocol

Where engines put checkpoints. The runner provides one backed by the run prefix.

MemoryCheckpointSink ¶

MemoryCheckpointSink()

Keeps the latest checkpoint per key in memory (tests, notebooks).

DirectoryCheckpointSink ¶

DirectoryCheckpointSink(root: str | PathLike[str])

One directory per key: meta.json plus draws.zarr.zip (written atomically).

SamplingCheckpoint ¶

SamplingCheckpoint(
    key: str,
    engine: str,
    resumable: bool,
    chains: int,
    tune: int,
    draws: int,
    completed: int,
    data: Dataset = Dataset(),
    state: dict[str, Any] = dict(),
)

The state of a segmented fit after completed post-tuning draws per chain.

fraction ¶

fraction: float

Share of the post-tuning draws already taken.

Diagnostics¶

Convergence diagnostics as scorecard-ready :class:~mapi_spec.capabilities.Diagnostic objects.

Every check reports something - info when it passes - so the sanity scorecard can list it. Thresholds follow current practice (Vehtari et al. 2021): rank-normalised split r-hat below 1.01 and at least 100 effective draws per chain for bulk and tail ESS.

convergence_diagnostics ¶

convergence_diagnostics(
    posterior: Dataset,
    sample_stats: Dataset | None = None,
    *,
    var_names: Sequence[str] | None = None,
    rhat_max: float = 1.01,
    ess_per_chain: int = 100,
) -> list[Diagnostic]

r-hat, bulk/tail ESS, divergences and tree-depth saturation for an MCMC posterior.

var_names limits the checks (pass the free random variables; deterministics over data dims are large and add nothing). Needs at least 2 chains and 4 draws for r-hat/ESS.

mapi_pymc¶