aimz 0.10.0

Scalable probabilistic impact modeling

aimz: Scalable probabilistic impact modeling

Installation | Tutorial | User Guide | FAQs | Changelog

Overview

aimz is a Python library for scalable probabilistic impact modeling—estimating how interventions affect outcomes while quantifying uncertainty. It provides an intuitive interface for fitting Bayesian models, drawing posterior samples, generating large-scale posterior predictive simulations, and estimating intervention effects with minimal boilerplate.

Key capabilities

• Flexible model specification: Built on NumPyro and JAX, bring a NumPyro model as a "kernel"—aimz does not enforce a fixed architecture.
• Scalable predictive sampling: JIT-compiled, sharded sampling streams results to chunked Zarr stores, enabling large-scale posterior predictive simulations.
• Structured outputs: Predictions, samples, and effect estimates are materialized as Xarray objects backed by Zarr, integrating cleanly with the scientific Python ecosystem.
• Intervention handling: Specify interventions declaratively and estimate effects from posterior predictive distributions.
• Experiment tracking: MLflow integration for logging runs, parameters, metrics, and model artifacts with full lineage.

Installation

Install aimz using either pip or conda:

pip install -U aimz

conda install -c conda-forge aimz

For additional details, see the full installation guide.

Quick start

from aimz import ImpactModel

# Define a probabilistic model (kernel) using NumPyro primitives
def model(X, y=None):
    ...

# Load or prepare data
X, y = ...

# Initialize ImpactModel with SVI or MCMC inference
im = ImpactModel(
    model,
    rng_key=...,      # e.g., jax.random.key(0)
    inference=...,    # e.g., SVI (or MCMC)
)

# Fit model and draw posterior samples
im.fit(X, y)

# Generate posterior predictive samples
dt = im.predict(X)

# Estimate intervention effects
dt_baseline = im.predict(X)
dt_intervention = im.predict(X, intervention={"treatment": 1.0})
effect = im.estimate_effect(dt_baseline, dt_intervention)

Contributing

See the Contributing Guide to get started.