Skip to content

mapi_sdk

This package is the xarray runtime. It runs in notebooks and in runner images. It never runs in the control plane.

Fit and score

The fit/score harness: the platform side of every run (plan §4.2).

:func:fit_spec implements fit.mode:

  • bind the data, validate it, fit the stateless prefix on all data up to the split stage;
  • split; fit the remaining step states and the model on train (the evaluation fit);
  • evaluate the holdout through the score path (context from train) - never a separate code path - and compute the spec's metrics;
  • evaluate_then_refit: refit states and model on train ∪ holdout so the context tails and terminal state come from the full data; the production bundle links the evaluation bundle and its metrics. train_only keeps the train fit as the result.
  • derived quantities the spec requests (outputs.derived) whose manifest declaration says cache: fit are computed on the production fit by the plugin's derive hook, over the first outputs.ui_draws draws of the thinning index, with ScoreContext.to_original mapping them to the output space.

:func:score_bundle implements scoring: CoordPlan checks, binding, interventions, entry at inputs_stage, context replay, canonical coords, model.score, context removal, selection, inverse transforms to the output space, aggregation, summaries, thinned draws and reverse binding. It returns an xr.DataTree.

Tree datasets (M4): both accept an xr.DataTree wherever they take the dataset. The nodes the spec's data.nodes selects are merged into one flat dataset first (:func:mapi_sdk.tree.select_nodes, aligned per data.join) and everything else runs on it exactly as for a plain dataset. A model whose capabilities declare input="tree" also gets the whole tree, bound per node, as FitContext.tree / ScoreContext.tree.

Image collections (v0.15): both accept a :class:~mapi_sdk.collection.Collection as the dataset. A model that reads datasets gets a seeded sample as the v0.14 image dataset (:func:~mapi_sdk.collection.collection_read, per data.collection). A model whose capabilities declare input="collection" runs on the manifest dataset (one row per image, no pixels; binding, schema checks, row splits, metrics and bundles as usual) and only the model call changes: fit_batches/score_batches receive a loader of the images of the rows that reached the model.

FitHooks

FitHooks(
    log: Callable[[str], None] | None = None,
    on_progress: Callable[[float, str], None] | None = None,
    cancel: CancelToken | None = None,
    checkpoint_sink: Any | None = None,
    tracker: Tracker | None = None,
)

Optional callbacks for progress, logging and cooperative cancellation.

checkpoint_sink is handed to the model's :class:~mapi_sdk.models.FitContext so engines with segmented sampling can checkpoint (and resume) long fits.

tracker (v0.11) receives the metrics, params, tags and log lines the model logs (FitContext.tracker / ScoreContext.tracker) plus the spec params of a fit; the runner attaches one that streams them as run events.

window

window(
    lo: float, hi: float
) -> Callable[[float, str], None] | None

A progress callback mapping a sub-task's [0, 1] onto [lo, hi].

fit_spec

fit_spec(
    spec: ModelSpec | str,
    data: Dataset | DataTree | Collection,
    side_tables: Mapping[str, Dataset] | None = None,
    registry: Registry | None = None,
    *,
    raw_schema: DataSchema | None = None,
    hooks: FitHooks | None = None,
    keep_evaluation_bundle: bool = True,
    hash_data: bool = True,
    tree_schema: TreeSchema | None = None,
    tree: DataTree | None = None,
    record_metrics: bool = True,
) -> Bundle

Fit spec on data and return the production :class:~mapi_sdk.bundle.Bundle.

raw_schema describes data under its dataset names; when omitted it is read from data.attrs["mapi_schema"] or inferred (variables bound to the model's targets are marked as targets). side_tables maps data.side_tables names to datasets.

data may be a tree dataset (M4): the nodes spec.data.nodes selects are merged into the flat dataset the pipeline reads (:func:mapi_sdk.tree.select_nodes); tree_schema describes the tree (inferred when omitted) and raw_schema, if given, the merged flat dataset. A tree-input model also gets the bound tree (ctx.tree); tree (dataset names) stands in for it when data was already selected from a tree (a refit on modified flat data). With a data.split the evaluation fit and the holdout score see only their part of the tree.

Run tracking (v0.11): the spec params (:func:mapi_sdk.tracking.spec_params) go to hooks.tracker once; what the model logs goes to hooks.tracker (an evaluation fit's metrics with the label fit=evaluation) and, with record_metrics, into each bundle's metrics_history model group.

data may be an image collection (v0.15, see the module docstring): the pipeline must then be empty for an input="collection" model, and only row splits are taken.

score_bundle

score_bundle(
    bundle: Bundle,
    request: ScoreRequest | str,
    data: Dataset | DataTree | Collection | None = None,
    *,
    scenario: Dataset | DataTree | None = None,
    side_tables: Mapping[str, Dataset] | None = None,
    registry: Registry | None = None,
    max_draws: int | None = None,
    return_full: bool = False,
    metrics: Sequence[str] | None = None,
    metric_mask: DataArray | None = None,
    hooks: FitHooks | None = None,
    observed: Dataset | DataTree | None = None,
    history: Dataset | DataTree | None = None,
    tree: DataTree | None = None,
    metric_history: Dataset | None = None,
    collection: Collection | None = None,
) -> DataTree

Score a fitted bundle; returns a DataTree.

data is the request's data payload and scenario its scenario payload. Outside counterfactuals a scenario simply is the payload. A counterfactual scores the scenario (or the data, when no scenario is given) with the interventions applied against the baseline data; a scenario without baseline data is rejected.

Groups: summary (stats along a stat dim), draws (when requested; joint, from the bundle thinning index), predictions (full representation, when return_full), and for counterfactuals baseline/... and delta/.... Root attrs carry the representation, output kind, units/output space, retransformation labels, non-error issues and - in evaluate mode - the metrics. observed (evaluation of cell-level splits) gives the observed targets for metrics when data has them masked. history (M3) is the fitted dataset (dataset names): forecasts whose future_covariates policy is carry_last or seasonal_naive extend the covariates the payload lacks from it (:mod:mapi_sdk.future).

Tree datasets (M4): data, scenario, observed and history may be xr.DataTree versions entering at the raw stage; the nodes the bundle's data.nodes selects are merged into the flat payload first. A model whose capabilities declare input="tree" gets the payload's tree, bound, as ScoreContext.tree; tree (dataset names) stands in for it when the payload is already flat (the evaluation of a fit from a tree dataset).

metric_history (v0.10) is the training data at req.inputs_stage (any names) of an evaluation; mase takes its scale from it (:func:mapi_sdk.metrics.mase_scale). Its targets go through the pipeline to the model stage and back to the output space exactly like the observed values (or stay at the model stage for a model that owns inverse_units), so the naive error is in the units of the error it divides. Without it, mase is skipped.

data may be an image collection (v0.15): a model that reads datasets scores a seeded sample (data.collection of the fit spec); an input="collection" model scores the manifest dataset with score_batches on the requested images. A manifest payload (a subset of the collection's manifest dataset, such as the evaluation holdout) names its collection instead.

Steps

The preprocessing step contract: Step, StepState and StepContext.

A step is a class with static metadata (id, version, kind, a Pydantic Settings model, an optional State model) and four methods::

fit(ds, settings, ctx) -> State           # learn on train data only
transform(ds, settings, state, ctx) -> ds # never learns; must tolerate an absent target
inverse(ds, settings, state, ctx) -> ds   # invertible VALUE steps only

Rules the contract kit (:mod:mapi_sdk.testing) enforces:

  • transform(train, fit(train)) equals the fit-time output exactly, and transform is deterministic and does not mutate its input;
  • reductions name the dims they reduce over; the state stays labelled by the rest;
  • group-local operations take along/within and declare a :class:ContextNeed;
  • unseen levels go through :meth:StepContext.align_state;
  • the static :class:~mapi_spec.effects.SchemaEffect (effect_template) matches the real output schema.

The static metadata is exported as a :class:~mapi_spec.catalog.StepCatalogEntry, so the control plane validates specs without importing step code.

Step

Base class for preprocessing steps. See the module docstring for the contract.

linear

linear: bool = True

Whether the inverse is affine (point retransformation is exact for means).

additive

additive: bool = False

Whether the inverse is linear without an offset (inv(a + b) = inv(a) + inv(b), e.g. multiplying by a scale or a rate). Only then do additive decompositions (contributions(component, ...)) invert component by component.

provides

provides: tuple[Literal['scaling', 'carryover'], ...] = ()

Preprocessing a model can own (Capabilities.owns): scaling or carry-over.

about

about: About | None = None

M3.5: the plain-language About card (mapi_spec.about.About). Built-in core.* steps take theirs from :mod:mapi_spec.abouts when this is unset.

about_card

about_card() -> About | None

The step's About card: about, else the built-in card of its id.

catalog_entry

catalog_entry() -> StepCatalogEntry

Static JSON describing this step (no data, no code) for the control plane.

effect

effect(s: SettingsT) -> SchemaEffect

The schema effect for concrete settings (default: render effect_template).

context

context(s: SettingsT) -> ContextNeed | None

Per-group tail needed at score time (default: render context_template).

touched_vars

touched_vars(s: SettingsT) -> set[str]

Variables whose values this step changes (for the inverse chain).

fit

fit(ds: Dataset, s: SettingsT, ctx: StepContext) -> StateT

Learn the state on training data. Stateless steps keep the default.

StepSettings

Bases: BaseModel

Base for step settings: unknown keys are rejected and instances are immutable.

StepState

Bases: BaseModel

Learned state of one step: JSON (learned and subclass fields) plus arrays.

arrays holds labelled arrays (e.g. scale(market, channel)) and context the per-group tail of the step's input stage recorded at fit time. Subclass fields must be JSON-serialisable; arrays must use persistable dtypes. Nothing is ever pickled.

json_payload

json_payload() -> dict[str, Any]

Every JSON field (including subclass fields), without the arrays.

array

array(name: str) -> DataArray

A named array from arrays (StepError if absent).

StepContext

StepContext(
    step_id: str,
    phase: Literal[
        "fit", "transform", "score", "inverse"
    ] = "fit",
    schema: DataSchema | None = None,
    side_tables: Mapping[str, Dataset] = dict(),
    seed: int = 0,
    log: Callable[[str], None] | None = None,
    cancel: CancelToken | None = None,
)

What a step may use besides its data, settings and state.

side_table

side_table(name: str) -> Dataset

A side table by its name in data.side_tables (SIDE_TABLE_COVERAGE if absent).

align_state

align_state(
    state: DataArray,
    ds: Dataset | DataArray,
    policy: UnseenPolicy = "error",
    *,
    pooled: DataArray | float | None = None,
    fill_value: float = nan,
    dims: Sequence[str] | None = None,
) -> DataArray

Reindex a labelled state array to the coords of ds, handling unseen levels.

error raises UNSEEN_LEVELS; pooled uses pooled (default: the state's mean over the dim) for new levels; fill uses fill_value. Levels the state knows but ds lacks are dropped, so the result aligns with ds by label.

CancelToken

CancelToken()

Cooperative cancellation flag shared between an executor and a running job.

Models

The model plugin contract (plan §5).

A plugin turns a model-stage xr.Dataset into a :class:FitOutput and, later, scores a canonicalised dataset into an xr.DataTree whose predictions group holds the declared outputs in the plugin's :class:~mapi_spec.capabilities.Representation::

draws      (chain, draw, *dims)
quantiles  (quantile, *dims)       quantile coords in (0, 1)
interval   (bound, *dims)          bound coords include "lower", "upper" (+ "point")
point      (*dims)
moments    (moment, *dims)         moment coords ["mean", "sd"]

The platform owns everything around that: binding, CoordPlan, context, canonical coords, inverse transforms, reverse binding and summaries. :class:CustomModel is the base class for hand-written plugins; mapi_pymc.PyMCModel and mapi_sklearn.SklearnModel build on the same protocol.

PREDICTIONS

PREDICTIONS = 'predictions'

Name of the DataTree group a plugin's score must fill.

CustomModel

Base class for hand-written plugins: declare metadata, implement fit/score.

Optional hooks: :meth:diagnostics, :meth:postprocess, :meth:derive, :meth:objective, :meth:recommend, :meth:translate_priors, :meth:optimize, :meth:forecast_with_actuals (with capabilities.conditions_on_actuals), :meth:fit_batches/:meth:score_batches (with capabilities.input="collection"). Set pad_extend = False if the model does not index time positionally (scoring then skips padding extend dims to the fit range). List in structural_missing the canonical variables whose NaN are part of the design of the data (not gaps), so that the scorecard check data.missing does not count them.

structural_missing

structural_missing: Sequence[str] = ()

Canonical variables whose NaN are expected by design (for example cells without a sensor). The scorecard check data.missing does not count their missing values.

about

about: About | None = None

M3.5: the plain-language About card. Built-in models take theirs from :data:mapi_spec.abouts.MODEL_ABOUTS when this is unset.

family

family: str | None = None

M3.5: the family governance policies match (default: the id's namespace).

side_table_params

side_table_params: tuple[str, ...] = ()

v0.15: settings keys whose value names a data.side_tables entry the model reads (from FitContext.side_tables), such as pretrained weights.

fit_batches

fit_batches(
    loader: ImageLoader,
    s: SettingsT,
    cfg: FitConfig,
    ctx: FitContext,
) -> FitOutput

Fit from a lazy image loader (v0.15; capabilities.input="collection").

The harness calls it instead of :meth:fit when the data is an image collection: loader holds the images of the rows that reached the model (the training rows of the split), in the order of the model-stage manifest dataset. The output is the same :class:FitOutput :meth:fit returns.

score_batches

score_batches(
    fitted: Fitted,
    loader: ImageLoader,
    req: ScoreRequest,
    ctx: ScoreContext,
) -> DataTree

Score a lazy image loader (v0.15): predictions along the loader's dim in the loader's order (the requested images only). The harness puts them back in the canonical order (NaN for padded rows) before its usual post-processing.

forecast_with_actuals

forecast_with_actuals(
    fitted: Fitted,
    ds: Dataset,
    actuals: ForecastActuals,
    req: ScoreRequest,
    ctx: ScoreContext,
) -> DataTree

A forecast conditioned on actuals (:mod:mapi_spec.actuals): plugins that declare capabilities.conditions_on_actuals override it; the harness calls it instead of :meth:score whenever a forecast payload carries actuals.

ds is the canonicalised model-stage payload (actuals, periods to forecast and their drivers); actuals says which periods are which per series (:class:~mapi_sdk.actuals.ForecastActuals). Update the model's state with the actuals using the fitted parameters - never re-estimate - and return the predictions group over every period of ds, as :meth:score does: the conditioned multi-step forecast on actuals.future, one-step-ahead predictions elsewhere. Supplied actuals inside the fitted range replace the fitted values.

The default refuses (ACTUALS_NOT_SUPPORTED), so a model never ignores actuals.

optimize

optimize(
    fitted: Fitted,
    request: OptimizeRequest,
    ctx: ScoreContext,
) -> BudgetPlan

Budget optimisation (M3.5, POST …:optimize; plugins that declare capabilities.decision): split request's budget across channels from the fitted model's posterior and return the plan with its uncertainty. The runner fills in the KPI name and units.

translate_priors

translate_priors(
    ds: Dataset,
    s: SettingsT,
    job: PriorTranslationJob,
    ctx: FitContext,
) -> PriorTranslation

Business priors to model-space priors over model-stage ds (M3; plugins that declare business_priors in Capabilities.calibration). ctx.to_original carries the data-derived preprocessing scales.

ModelPlugin

Bases: Protocol

Structural protocol every model plugin satisfies.

ModelSettings

Bases: BaseModel

Base for plugin settings: unknown keys are rejected, instances are immutable.

FitContext

FitContext(
    seed: int = 0,
    schema: DataSchema | None = None,
    policies: Mapping[str, DimPolicy] = dict(),
    coords: Mapping[str, ndarray] = dict(),
    workdir: Path | None = None,
    log: Callable[[str], None] | None = None,
    on_progress: ProgressFn | None = None,
    cancel: CancelToken | None = None,
    checkpoint_sink: Any | None = None,
    stage_schema: DataSchema | None = None,
    calibration: Calibration | None = None,
    to_original: Callable[[DataArray, str], DataArray]
    | None = None,
    tree: DataTree | None = None,
    tracker: Tracker | None = None,
    side_tables: Mapping[str, Dataset] = dict(),
    fit_key: str | None = None,
)

Bases: TrackingMixin

Everything a plugin's fit may use besides data, settings and fit config.

schema is the model's resolved input schema (its template); stage_schema is the schema of the pipeline stage the data comes from, with every variable and its role.

checkpoint_sink is an engine-specific sink for segmented fits (for example mapi_pymc.CheckpointSink); the runner supplies one that writes to the run prefix. None disables checkpointing.

calibration (M3) is the spec's calibration block - business priors and lift tests - for plugins whose capabilities declare them (Capabilities.calibration). to_original (M3) maps a model-stage array in the units of input variable src back to the output space with the fitted preprocessing state (as :attr:ScoreContext.to_original does for derived quantities), so a plugin can turn business-unit calibration data into model units. Both are None outside the harness.

tree (M4) is the whole tree dataset, binding renames applied per node (:meth:~mapi_sdk.binding.BindingMap.apply_tree), for models whose capabilities declare input="tree" and were fitted from a tree dataset; None otherwise. ds is still the flat pipeline output (the nodes data.nodes selects, preprocessed).

tracker (v0.11) receives what the plugin logs with the :class:~mapi_sdk.tracking.TrackingMixin methods (log_metric, log_params, set_tag, log_line, chains); the harness attaches one that records the bundle's metrics_history group and forwards to the runner. None: validate only.

side_tables (v0.15) holds the model's side tables (the settings named by ModelManifest.side_table_params), by data.side_tables name, raw as given. fit_key (v0.15) is 32 hex characters, stable for (spec hash, data hash, role): the base of checkpoint keys, so a restarted fit finds its own checkpoints. None outside the harness.

FitOutput

FitOutput(
    groups: dict[str, Dataset] = dict(),
    learned: dict[str, JsonValue] = dict(),
    native: dict[str, bytes] = dict(),
    terminal_state: Dataset | None = None,
    derived: dict[str, Dataset] = dict(),
    diagnostics: list[Diagnostic] = list(),
)

What a plugin's fit returns. Everything here is persisted in the bundle.

groups uses ArviZ group names (posterior, sample_stats, log_likelihood, observed_data, prior, ...). native holds safe native formats only (skops, xgboost JSON, safetensors) - never pickle. learned is small JSON state.

Fitted

Fitted(
    manifest: ModelManifest,
    settings: dict[str, Any],
    fit_config: FitConfig,
    groups: dict[str, Dataset] = dict(),
    learned: dict[str, JsonValue] = dict(),
    native: dict[str, bytes] = dict(),
    terminal_state: Dataset | None = None,
    derived: dict[str, Dataset] = dict(),
    schema: DataSchema | None = None,
    coords: dict[str, ndarray] = dict(),
)

A fitted model as score sees it (rebuilt from the bundle).

ScoreContext

ScoreContext(
    mode: ScoreMode = "predict",
    kind: Literal["expected", "predictive"] = "expected",
    seed: int = 0,
    schema: DataSchema | None = None,
    policies: Mapping[str, DimPolicy] = dict(),
    horizon: Horizon | None = None,
    draw_index: ndarray | None = None,
    log: Callable[[str], None] | None = None,
    cancel: CancelToken | None = None,
    to_original: Callable[[DataArray, str], DataArray]
    | None = None,
    tree: DataTree | None = None,
    tracker: Tracker | None = None,
)

Bases: TrackingMixin

Everything a plugin's score may use besides the fitted model and data.

The dataset handed to score is canonicalised: fixed and new-level dims are in fit order (padded to the full fit set), new levels are appended and flagged by the bool coordinate <dim>__is_new, and <dim>__requested marks real (non-padding) labels.

to_original is set when the harness computes fit-cached derived quantities (:meth:CustomModel.derive): to_original(da, src) maps a model-stage array in the units of input variable src (e.g. a contribution in y units, a spend level in spend units) to the output space - the pipeline inverse back to its last reshape step - and labels it with units and output_space attributes.

tree (M4) is the bound tree dataset the payload came from, for models whose capabilities declare input="tree" (as :attr:FitContext.tree); None when the payload is a flat dataset or the model reads datasets only.

tracker (v0.11): as :attr:FitContext.tracker (score runs log through it too).

new_level_mask

new_level_mask(ds: Dataset, dim: str) -> DataArray

Bool mask along dim: True for levels unseen at fit (from canonicalisation).

build_manifest

build_manifest(
    *,
    id: str,
    version: str,
    capabilities: Capabilities,
    schema_template: DataSchemaTemplate
    | DataSchema
    | Mapping[str, Any],
    settings_model: type[BaseModel] | None = None,
    fit_schema: Mapping[str, Any] | None = None,
    outputs: Mapping[str, OutputDecl | Mapping[str, Any]]
    | None = None,
    derived: Sequence[DerivedDecl | Mapping[str, Any]] = (),
    plots: Sequence[str] = (),
    resources: ResourceHint | None = None,
    artifact_format: str = "zarr.zip/v1",
    description: str | None = None,
    about: About | None = None,
    family: str | None = None,
    side_table_params: Sequence[str] = (),
) -> ModelManifest

Assemble and validate a :class:~mapi_spec.capabilities.ModelManifest.

Pipelines

The pipeline engine: stages, fitting, score-time replay and the inverse chain (plan §4.2).

A stage is the dataset after a named step (raw is the bound input before any step). :meth:Pipeline.fit fits and transforms a range of stages, recording each step's state and its per-group context tail at the step's input stage. :meth:Pipeline.transform replays stored states over a range of stages, prepending each step's context before it runs (:mod:mapi_sdk.context). :meth:Pipeline.inverse_outputs maps model outputs back through invertible VALUE steps up to the last RESHAPE step and labels the output space.

Pipeline

Pipeline(
    steps: Sequence[PipelineStep],
    *,
    schemas: Mapping[str, DataSchema | None] | None = None,
)

An ordered list of :class:PipelineStep with optional per-stage schemas.

build

build(
    items: Sequence[
        tuple[
            str, type[Step[Any, Any]] | Step[Any, Any], Any
        ]
    ],
    *,
    schemas: Mapping[str, DataSchema | None] | None = None,
) -> Pipeline

Build from (step_id, step class or instance, settings) triples.

from_spec

from_spec(
    spec: ModelSpec,
    registry: Registry,
    resolved: ResolvedSpec | None = None,
) -> Pipeline

Instantiate a spec's pipeline from a registry (settings with defaults applied).

index

index(stage: str) -> int

Position of stage (0 = raw, i = after the i-th step).

between

between(
    start: str = RAW, end: str | None = None
) -> list[PipelineStep]

Steps executed to go from stage start to stage end (default: last).

fit

fit(
    ds: Dataset,
    *,
    start: str = RAW,
    end: str | None = None,
    side_tables: Mapping[str, Dataset] | None = None,
    seed: int = 0,
    cancel: CancelToken | None = None,
    log: Callable[[str], None] | None = None,
    keep_stages: bool = False,
) -> FitTrace

Fit and transform the steps from start to end.

Each step's state is learned on its input, its context tail is recorded at that input stage, and the transformed data feeds the next step.

transform

transform(
    ds: Dataset,
    states: Mapping[str, StepState],
    *,
    start: str = RAW,
    end: str | None = None,
    replay: bool = True,
    side_tables: Mapping[str, Dataset] | None = None,
    seed: int = 0,
    cancel: CancelToken | None = None,
    log: Callable[[str], None] | None = None,
    keep_stages: bool = False,
    fit_coords: Mapping[str, Dataset] | None = None,
) -> ReplayTrace

Replay stored states from start to end; never learns.

With replay=True each step's stored context is prepended (per group, checked for contiguity) before it runs; the returned marks say which rows to drop later. fit_coords (per-stage fitted coordinates) let a replay tell an in-sample window that starts with the fitted series from one that loses look-back history.

inverse_chain

inverse_chain(
    src: str,
    *,
    to_stage: str,
    from_stage: str | None = None,
) -> tuple[list[PipelineStep], str]

Invertible steps to apply to src (newest first) and the stage reached.

Walks back from from_stage and stops at to_stage, at any RESHAPE step, or at the first non-invertible step that changes src.

inverse_outputs

inverse_outputs(
    ds: Dataset,
    states: Mapping[str, StepState],
    var_sources: Mapping[str, str],
    *,
    to_stage: str | None = None,
    from_stage: str | None = None,
    representation: Representation = DRAWS,
    retransform: Retransform = "median",
    residuals: Mapping[str, DataArray] | None = None,
    side_tables: Mapping[str, Dataset] | None = None,
    decompositions: Collection[str] = (),
) -> InverseResult

Map output variables back to their output space.

var_sources maps each output variable to the input variable whose transforms it inherits (y_hat -> y). Inversion stops at to_stage (default: the last RESHAPE stage) or at the first non-invertible step that touched the variable. Point outputs through a non-linear inverse are labelled median (plain inverse) or use Duan's smearing over stored model-space residuals; moments use the first-order delta method (g(mu), |g'(mu)| sd with a central difference).

decompositions name additive decompositions (contributions(component, ...)). They are inverted component by component only through additive steps (Step.additive); through any other inverse (expm1, x * sd + mean) the components would no longer sum to the inverted prediction, so inversion stops right after that step and :attr:InverseResult.kept says why.

PipelineStep

PipelineStep(
    id: str,
    step: Step[Any, Any],
    settings: StepSettings,
    use: str = "",
)

A step instance bound to a pipeline position and concrete settings.

Bundles

Fitted-model bundles: one bundle.zarr.zip DataTree plus JSON sidecars (plan §8.6).

Layout of bundle.zarr.zip::

/                       attrs: format, bundle_id, role
/spec                   attrs: spec (canonical JSON), resolved (ResolvedSpec JSON)
/schema/raw             attrs: json (dataset-named raw DataSchema)
/schema/model           attrs: json (model input schema, effective policies)
/binding                attrs: json
/preprocess/<step_id>   attrs: state (JSON); children: arrays, context
/model                  attrs: manifest, settings, fit_config, learned (JSON)
/model/<group>          ArviZ-named groups (posterior, sample_stats, log_likelihood ...)
/terminal_state         plugin-owned state for sequential models
/native                 uint8 arrays of safe native artifacts (skops, xgb json, ...)
/derived/<name>         derived quantities cached at fit
/fit_coords/<stage>     coordinates of the fitted data at every stage
/side_tables/<name>     pinned side tables used at fit
/residuals              model-space residuals (point models, smearing)
/thin                   the bundle-wide draw thinning index

Sidecars: manifest.json (content hashes of every file: size, CRC32C and MD5 in the GCS base64 form, SHA-256), diagnostics.json, plots.json and metrics.json. An evaluate_then_refit bundle keeps its train-only evaluation bundle under evaluation/. Pickle is never written and is rejected on read.

Bundle

Bundle(
    spec: ModelSpec,
    resolved: ResolvedSpec,
    raw_schema: DataSchema,
    states: dict[str, StepState],
    fit_coords: dict[str, Dataset],
    model: Fitted,
    manifest: BundleManifest,
    thin: ndarray | None = None,
    side_tables: dict[str, Dataset] = dict(),
    residuals: dict[str, DataArray] = dict(),
    diagnostics: list[Diagnostic] = list(),
    metrics: MetricsReport = MetricsReport(),
    plots: list[dict[str, Any]] = list(),
    evaluation: Bundle | None = None,
)

A fitted model: spec, schemas, preprocessing states, model groups and sidecars.

to_tree

to_tree() -> DataTree

The bundle.zarr.zip DataTree (see the module docstring for the layout).

write

write(path: str | Path) -> BundleManifest

Write the bundle directory; returns the manifest with file hashes filled in.

read

read(
    path: str | Path,
    registry: Registry,
    *,
    verify: bool = True,
    limits: WireLimits = DEFAULT_LIMITS,
) -> Bundle

Read a bundle directory (verifying file hashes against manifest.json).

from_tree

from_tree(
    tree: DataTree,
    registry: Registry,
    *,
    manifest: BundleManifest,
) -> Bundle

Rebuild a :class:Bundle from its DataTree (typed states need the registry).

BundleManifest

Bases: BaseModel

Bundle identity, provenance and the content hashes of its files.

MetricsReport

Bases: BaseModel

Holdout metrics computed through the score path, plus skipped metrics and why.

FileDigest

Bases: BaseModel

Content hashes of one file, in the forms GCS reports (CRC32C/MD5 base64).

Coordinates

CoordPlan checks, coordinate canonicalisation and horizon coordinates (plan §4.2-4.3).

  • :class:CoordPlan runs the static checks of :func:mapi_spec.validate.validate_score_request against the bundle's fit coordinates and the actual score payload, raising typed errors (UNSEEN_LEVELS, COORD_NOT_ALLOWED, HORIZON_EXCEEDS_FIT, MISSING_FUTURE_COVARIATE, NON_CONTIGUOUS_EXTENSION, ...).
  • :func:canonicalise_coords reorders model-stage score data to the fit order, pads requested subsets to the full fit set, appends unseen levels (flagged by <dim>__is_new) and returns a :class:CanonicalPlan that selects the requested coords back afterwards.
  • :func:horizon_coords / :func:horizon_frame generate future coordinates from a dim's frequency (or integer step).

CoordPlan

CoordPlan(
    request: ScoreRequest,
    resolved: ResolvedSpec,
    fit_coords: Mapping[str, Dataset],
    *,
    max_draws: int | None = None,
    live_capabilities: Capabilities | None = None,
    data_space: DataSpace | None = None,
)

Typed early checks for a :class:~mapi_spec.spec.ScoreRequest.

fit_coords maps each stage name to a coordinates-only dataset of the fitted data at that stage (stored in the bundle). Raises the most specific error for the first problem and returns the remaining warnings/infos.

check

check(inputs: Dataset | None = None) -> list[Issue]

Run the checks; inputs is the payload at inputs_stage (canonical names).

CanonicalPlan

CanonicalPlan(dims: dict[str, DimCanon] = dict())

Records canonicalisation so outputs can be mapped back to the requested coords.

restore

restore(obj: Dataset) -> Dataset

Select the requested labels (request order) and drop helper coordinates.

canonicalise_coords

canonicalise_coords(
    ds: Dataset,
    fit_coords: Dataset,
    schema: DataSchema | None,
    *,
    pad_extend: bool = True,
    marks: Collection[ContextMark] = (),
    pad_fill: float = 0.0,
) -> tuple[Dataset, CanonicalPlan]

Reindex model-stage score data to the fit coordinate order.

  • fixed dims: labels must be fit labels (UNSEEN_LEVELS); data is reindexed to the full fit set (padding unrequested labels);
  • new_levels dims: fit labels in fit order, then unseen labels appended (the <dim>__is_new coordinate flags them; strategy error raises);
  • extend dims: labels may not start before the fit range (COORD_NOT_ALLOWED), must continue it without gaps (NON_CONTIGUOUS_EXTENSION) and may not exceed max_extension (HORIZON_EXCEEDS_FIT); with pad_extend the full fit range is included so positional models see a contiguous axis.

Padded cells of numeric non-target variables are filled with pad_fill and target variables with NaN; <dim>__requested marks real labels.

horizon_coords

horizon_coords(
    fit: ndarray | Sequence[Any],
    periods: int,
    *,
    freq: str | None = None,
) -> ndarray

periods labels continuing fit (dates by freq, integers by their step).

Binding

Dataset ↔ canonical names (geo: market), applied to inputs and reversed on outputs.

BindingMap

BindingMap(
    binding: Binding | Mapping[str, Any] | None = None,
)

Runtime view of a :class:~mapi_spec.spec.Binding.

apply renames dataset names to canonical ones (what steps, models and priors use); reverse renames canonical names back on outputs. Unbound names pass through.

canonical_keys

canonical_keys(
    mapping: Mapping[str, V],
) -> dict[str, V]

Re-key a {dataset name: value} mapping by canonical names.

apply

apply(
    ds: Dataset,
    *,
    missing: Literal["error", "ignore"] = "error",
) -> Dataset

Rename dataset names to canonical names.

missing="error" raises BINDING_INVALID when a bound dataset name is absent (fit time); "ignore" skips absent names (score payloads may omit the target). A canonical name that collides with an unbound dataset name is always an error.

reverse

reverse(ds: Dataset) -> Dataset

Rename canonical names (dims, coords, variables) back to dataset names.

reverse_tree

reverse_tree(tree: DataTree) -> DataTree

:meth:reverse applied to every node of a DataTree.

Data checks

Array-level schema checks, schema inference and coordinate summaries.

  • :func:validate_dataset is the sandbox-side DataSchema.validate(ds): it returns :class:SchemaIssue objects with plain-language messages ("3 markets are missing 2 weeks of revenue in March 2024") that the UI can show as-is.
  • :func:infer_schema drafts a :class:~mapi_spec.schema.DataSchema from a dataset.
  • :func:summarize_coords / :func:data_summary produce the metadata-only :class:~mapi_spec.coords.CoordSummary objects the control plane uses for CoordPlan checks.

SchemaIssue

Bases: BaseModel

One array-level problem, described for a non-expert reader.

to_issue

to_issue() -> Issue

As a :class:~mapi_spec.errors.Issue pointing into the schema.

infer_schema

infer_schema(
    ds: Dataset,
    *,
    targets: Collection[str] = (),
    roles: Mapping[str, Role | str] | None = None,
    policies: Mapping[str, DimPolicy | str] | None = None,
    units: Mapping[str, str] | None = None,
) -> DataSchema

Draft a :class:DataSchema from ds.

Time dims get policy extend and an inferred frequency; other dims are fixed unless policies says otherwise. Variables are covariates unless listed in targets or roles. Missingness is allow when the data has missing values.

validate_dataset

validate_dataset(
    schema: DataSchema,
    ds: Dataset,
    *,
    optional_roles: Collection[Role] = (),
    check_values: bool = True,
    check_coords: bool = True,
) -> list[SchemaIssue]

Check ds against schema and describe every problem in plain language.

optional_roles lists roles whose variables may be absent (e.g. the target when scoring). check_values controls missingness, sentinel and bounds checks.

check_dataset

check_dataset(
    schema: DataSchema, ds: Dataset, **kwargs: Any
) -> list[SchemaIssue]

:func:validate_dataset, raising :class:SchemaValidationError on any error.

data_summary

data_summary(
    ds: Dataset, schema: DataSchema | None = None
) -> DataSummary

:class:~mapi_spec.coords.DataSummary of ds (variables plus coord summaries).

With a schema, observed says where each target - and each driver a forecast needs (future_required) - holds values along every ordered dim (time, extend or ordered): the actuals of a forecast payload and the drivers on its rows.

Summaries

Representation-aware summaries, aggregation and draw thinning (plan §5).

Outputs come in one of five representations (see :mod:mapi_sdk.models). Which summaries and aggregations are valid depends on the representation; the rules live in :data:mapi_spec.capabilities.REPRESENTATION_RULES and are enforced here:

  • draws: every summary; any aggregation (applied per draw);
  • quantiles: median/quantiles/ETI (interpolated on the stored quantile grid, never extrapolated); no aggregation;
  • interval: the interval and its point; no aggregation;
  • point: mean/point; linear aggregation only, flagged;
  • moments: mean/sd and normal ETIs; no aggregation.

Summaries are returned as one variable per output with a leading stat dim whose labels read like mean, hdi_90_lower, hdi_90_upper, q5 ... HDIs are computed in numpy (narrowest window over sorted draws).

Circular outputs (OutputDecl.circular, periods of :func:summarise): angles that wrap at period (one turn). Draw summaries are then taken on the circle: mean is the circular mean, sd the circular standard deviation sqrt(-2 ln R), and the median, quantiles, ETIs and HDIs come from the draws' wrapped deviations from the circular mean, so an interval is centred on the mean and may cross the seam. Every value is wrapped into [0, period) ([-period/2, period/2) for a difference, signed); an interval with lower > upper crosses the seam. This matches the evaluate metrics of :mod:mapi_sdk.metrics.

aggregate

aggregate(
    ds: Dataset,
    rep: Representation,
    agg: Aggregation | None,
) -> tuple[Dataset, bool]

Reduce outputs over agg.over if the representation allows it.

Returns the reduced dataset and whether the result must be flagged (point outputs). Raises CAPABILITY_UNSUPPORTED for statistically invalid aggregations.

hdi

hdi(
    values: ndarray, prob: float, axis: int = -1
) -> tuple[ndarray, ndarray]

Highest-density interval of draws along axis (narrowest window). NaN-propagating.

thin_index

thin_index(
    chains: int, draws: int, *, seed: int = 0
) -> ndarray

A bundle-wide random permutation of (chain, draw) pairs, shape (n, 2).

Every group and every score request uses the first k pairs, so subsets of draws stay joint across variables and requests.

Registry

Registries of steps, models and plot recipes.

Discovery uses entry points (mapi.steps, mapi.models, mapi.plots) - only inside runner images and notebooks, never in the control plane - plus in-memory registration for tests and notebooks. :meth:Registry.catalog exports the static JSON (:class:~mapi_spec.catalog.Catalog) that the control plane validates specs against.

An entry point may load a step class, a model class, a plot recipe, an iterable of those, or a zero-argument callable returning one of those.

Registry

Registry()

Steps, models and plot recipes by id and semver version.

register_step

register_step(
    cls: type[Step[Any, Any]], *, replace: bool = False
) -> type[Step[Any, Any]]

Register a step class (usable as a decorator).

register_model

register_model(
    cls: type[Any], *, replace: bool = False
) -> type[Any]

Register a model plugin class (usable as a decorator).

register_plot

register_plot(recipe: Any, *, replace: bool = False) -> Any

Register a plot recipe (anything with catalog_entry() and render).

register

register(obj: Any, *, replace: bool = False) -> None

Register a step, model or plot recipe (or an iterable of them).

load_entry_points

load_entry_points(
    groups: Iterable[str] = ENTRY_POINT_GROUPS,
    *,
    strict: bool = False,
) -> list[str]

Load and register every entry point in groups; return the errors.

Nothing an entry point exports may replace what is already registered (the first registration wins, whatever the iteration order of installed distributions), and ids in the built-in namespaces (:data:BUILTIN_NAMESPACES) are only accepted from the built-in distributions. Collisions are reported in load_errors (raised with strict).

catalog

catalog() -> Catalog

Static catalog JSON for the control plane (no code, no data).

plot_examples

plot_examples(
    *,
    on_error: Callable[[str, Exception], None]
    | None = None,
) -> PlotExamples

Every plot recipe rendered on its synthetic example data, keyed by id@version (recipes without an example_spec() or without example data are left out).

A recipe that fails, or whose example is not self-contained, is left out too and reported to on_error: one broken example never hides the others.

default_registry

default_registry(
    *,
    reload: bool = False,
    loader: Callable[[Registry], Any] | None = None,
) -> Registry

The process-wide registry, populated from entry points on first use.

In-process executor

In-process executor (isolation NONE) for notebooks, tests and mapi fit on own code.

Implements the :class:mapi_spec.executor.Executor protocol: submit starts the run in a worker thread, status is the source of truth, cancel is cooperative (checked between pipeline phases) and event_source yields :class:~mapi_spec.executor.RunEvent objects. Payloads are JSON only; dataset references are resolved by a resolver.

Supported run kinds: fit (payload spec, data, optional side_tables and raw_schema) and score/evaluate (payload request, bundle - a bundle directory or a key registered with :meth:InprocExecutor.put_bundle - and optional data/scenario/side_tables; otherwise the request's data/scenario refs are used). A counterfactual resolves data as the baseline and scenario as the scenario.

What the model logs through the tracking API (v0.11) is emitted as metrics, params and tags events (I-2 payload dicts; metric points batched per second) and log events with stream user.

Never use this executor for untrusted code: it runs plugins in the caller's process.

InprocExecutor

InprocExecutor(
    *,
    registry: Registry | None = None,
    resolver: DataResolver
    | Mapping[str, Dataset]
    | None = None,
    output_dir: str | Path | None = None,
)

Runs fit/score work in a thread of the calling process.

submit

submit(rs: RunSubmit, nonce: str) -> Handle

Start rs; resubmitting the same nonce returns the existing handle.

event_source

event_source() -> AsyncIterator[RunEvent]

Events of every run submitted after the iterator starts (runs forever).

put_bundle

put_bundle(key: str, bundle: Bundle) -> None

Make an in-memory bundle available to score payloads as {"bundle": key}.

result

result(h: Handle) -> Any

The run's result (a :class:Bundle for fit, an xr.DataTree for score).

wait

wait(h: Handle, timeout: float | None = None) -> ExecStatus

Wait until the run is terminal (TimeoutError after timeout seconds).

mapping_resolver

mapping_resolver(
    datasets: Mapping[str, Dataset],
) -> DataResolver

Resolve dataset refs from a mapping; file:// refs are read as zarr zips.

Wire formats

Wire formats: XrJSON, single-object zarr v3 zips, DataTree zips, stacking and Arrow IPC.

Encoding rules (plan §3.2), applied by every writer here:

  • any pandas.MultiIndex is reset before persisting (and recorded so readers restore it);
  • string coords and variables are zarr v3 vlen-utf8 (data type string);
  • datetimes are int64 nanoseconds with explicit units and calendar attrs;
  • only the zstd and blosc compressors are written or accepted;
  • :func:stack_rows drops empty rows and records a row mask so :func:unstack_rows restores them.

Pickle (and joblib/cloudpickle) is never used for anything. Readers treat every payload as hostile: zip entries, zarr metadata, codecs, dtypes and sizes are checked before any array is decoded.

to_xrjson

to_xrjson(
    ds: Dataset, *, limits: WireLimits = DEFAULT_LIMITS
) -> dict[str, Any]

Encode ds as XrJSON (the xr.Dataset.to_dict() shape plus explicit dtypes).

NaN/NaT become null; datetimes are ISO-8601 strings; MultiIndexes are reset.

from_xrjson

from_xrjson(
    obj: Mapping[str, Any],
    *,
    limits: WireLimits = DEFAULT_LIMITS,
) -> Dataset

Decode an XrJSON object (as produced by :func:to_xrjson or to_dict()).

Every element is counted against limits.max_elements before any array is built; shapes must be rectangular and consistent with dims.

write_zarr_zip

write_zarr_zip(
    ds: Dataset,
    dest: str | Path | IO[bytes] | None = None,
    *,
    compressor: Literal["zstd", "blosc"] = "zstd",
) -> bytes | None

Write ds as a single-object zarr v3 zip (deterministic bytes, STORED entries).

Returns the bytes when dest is None.

read_zarr_zip

read_zarr_zip(
    src: str | Path | bytes | IO[bytes],
    *,
    limits: WireLimits = DEFAULT_LIMITS,
) -> Dataset

Read a single-dataset zarr v3 zip after validating entries, codecs and dtypes.

write_datatree_zip

write_datatree_zip(
    tree: DataTree,
    dest: str | Path | IO[bytes] | None = None,
    *,
    compressor: Literal["zstd", "blosc"] = "zstd",
) -> bytes | None

Write a DataTree (every node encoded as by :func:write_zarr_zip).

read_datatree_zip

read_datatree_zip(
    src: str | Path | bytes | IO[bytes],
    *,
    limits: WireLimits = DEFAULT_LIMITS,
    node_pattern: str | None = None,
) -> DataTree

Read a DataTree written by :func:write_datatree_zip (same checks).

Every node gets the checks of :func:read_zarr_zip; the store as a whole is capped at limits.max_tree_nodes groups no deeper than limits.max_tree_depth, and when node_pattern is given every group path (/, /a/b) must match it - all checked before any array is decoded.

stack_rows

stack_rows(
    ds: Dataset,
    dims: Sequence[str],
    new_dim: str = "row",
    *,
    drop: Literal["all", "any", "none"] = "all",
    mask_var: str | None = None,
) -> Dataset

Stack dims into one new_dim and drop empty rows, recording a row mask.

A row is dropped when all (drop="all") or any ("any") of the variables that use every stacked dim are missing there (lower-rank variables, which are broadcast, only decide when no variable is full-rank), or where mask_var is false. The result has no MultiIndex: the component coords are plain coords along new_dim and the boolean <new_dim>__rowmask variable (over the original dims, with their full coords) says which rows were kept, so :func:unstack_rows can restore the dropped rows as missing.

unstack_rows

unstack_rows(ds: Dataset, new_dim: str = 'row') -> Dataset

Inverse of :func:stack_rows: dropped rows come back as missing values.

to_arrow_table

to_arrow_table(
    obj: Dataset | DataArray, *, max_rows: int = 2000000
) -> Table

Tidy (long) Arrow table: one column per dim coordinate plus one per variable.

Variable attrs (units, ...) are attached as JSON schema metadata under mapi.attrs, and the names of the dim columns under mapi.dims (a reader cannot tell a numeric dim, for example a grid position in cm, from a variable by its type).

Contract kit

Contract-kit checks for step and model authors (plan §13).

:func:check_step and :func:check_model run every contract check that is possible without a sandbox and return a :class:ContractReport (raising :class:ContractViolation by default when any check fails). Each check is independent: one failure never hides the others.

M3.5: both also check the item's plain-language About card (about): an incomplete card always fails; a missing one fails with require_about=True (what the catalog requires) and is skipped otherwise. :func:check_about checks cards on their own.

ModelFixture

ModelFixture(
    spec: ModelSpec | str,
    data: Dataset,
    side_tables: Mapping[str, Dataset] = dict(),
    raw_schema: DataSchema | None = None,
    registry: Registry | None = None,
    window: int | None = None,
)

A tiny end-to-end fixture for :func:check_model.

spec must reference the plugin under test; keep fit settings tiny (few draws).

ContractReport

ContractReport(
    subject: str, results: list[CheckResult] = list()
)

Outcome of a contract-kit run.

ContractViolation

ContractViolation(report: ContractReport)

Bases: AssertionError

Raised when a contract check fails; the report is attached.

check_step

check_step(
    step: type[Step[Any, Any]] | Step[Any, Any],
    ds: Dataset,
    settings: Mapping[str, Any]
    | StepSettings
    | None = None,
    *,
    schema: DataSchema | None = None,
    side_tables: Mapping[str, Dataset] | None = None,
    holdout: int = 2,
    raise_on_failure: bool = True,
    require_about: bool = False,
) -> ContractReport

Run the step contract checks on ds (the step's input) with settings.

schema describes ds (inferred when omitted); targets in it are used for the "absent target" check. holdout is the number of trailing rows used for the context-window check.

check_model

check_model(
    plugin: type[Any],
    fixture: ModelFixture,
    *,
    raise_on_failure: bool = True,
    checks: Sequence[str] | None = None,
    require_about: bool = False,
    collection: Collection | None = None,
) -> ContractReport

Run the model contract checks: fit, bundle round trip, coordinate handling, typed errors for every dim policy (unseen levels; gaps and out-of-range coords on datetime and integer extend dims; horizons past max_extension; forecasts missing a future-required covariate; horizons longer than a target lag), forecasts with actuals (refused without conditions_on_actuals; with it, the fitted values as actuals reproduce the plain forecast and a new actual moves it), unsupported modes, context-window equivalence, evaluate path == score path, output space, aggregation legality, additive contributions, no pickle, seed determinism, fast scoring == exact scoring within MCSE (via the plugin's check_fast_score hook) and that the manifest's plots sanitize.

collection (v0.15): for a model that declares input="collection", the batch path runs once - fit_batches on the collection with the fixture spec, then score_batches on a few images - and must give the outputs the dataset path gives (same variables and dims), only for the requested images.