aimz: Scalable probabilistic impact modeling

Installation | User Guide | Examples | FAQs | Changelog

Overview

aimz is a Python library for scalable probabilistic impact modeling: estimating how interventions affect outcomes while quantifying uncertainty.

It provides a high-level, object-oriented interface on top of NumPyro and JAX for building, fitting, and scaling Bayesian models: a user-defined NumPyro model is wrapped as a "kernel" inside a single class, augmented with capabilities for scalable predictive sampling, structured outputs, and experiment tracking.

Key capabilities

• Object-oriented interface for NumPyro models: Bring any NumPyro model as a "kernel" and access fit, predict, sample, and related methods through a single class; aimz does not enforce a fixed architecture.
• Scalable predictive sampling: JIT-compiled, sharded sampling streams results to chunked Zarr stores (or accumulates them in memory), enabling large-scale posterior predictive simulations that do not need to fit in memory.
• Structured outputs: Predictions, samples, and effect estimates are materialized as Xarray objects backed by Zarr or host memory, integrating cleanly with the scientific Python ecosystem.
• Intervention handling and impact modeling: 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 with minibatch SVI and draw posterior samples
im.fit(X, y, batch_size=32, epochs=10)

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

# Estimate the effect of fixing the sample site "treatment" to 1.0
dt_intervention = im.predict(X, intervention={"treatment": 1.0})
effect = im.estimate_effect(dt, dt_intervention)

Contributing

See the Contributing Guide to get started.

Citation

If you use aimz in your work, please cite the accompanying paper in the Journal of Open Source Software:

@article{Kim2026,
  title        = {aimz: Scalable probabilistic impact modeling},
  author       = {Kim, Eunseop},
  year         = 2026,
  journal      = {Journal of Open Source Software},
  publisher    = {The Open Journal},
  volume       = 11,
  number       = 120,
  pages        = 9738,
  doi          = {10.21105/joss.09738},
  url          = {https://doi.org/10.21105/joss.09738}
}