aimz: Scalable probabilistic impact modeling

Run Pytest

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, 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, 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 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.

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}
}

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

aimz-0.14.0.tar.gz (87.5 kB view details)

Uploaded Jul 25, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

aimz-0.14.0-py3-none-any.whl (74.6 kB view details)

Uploaded Jul 25, 2026 Python 3

File details

Details for the file aimz-0.14.0.tar.gz.

File metadata

Download URL: aimz-0.14.0.tar.gz
Upload date: Jul 25, 2026
Size: 87.5 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/6.1.0 CPython/3.13.14

File hashes

Hashes for aimz-0.14.0.tar.gz
Algorithm	Hash digest
SHA256	`ad6b8dce8a41c04cfe4ce989b34ddebc65ddff14a6a62ef860038d2b4c4a5e2d`
MD5	`de1a9da4b4169f76ebd80df39747e4df`
BLAKE2b-256	`db87e2516d4f45dd5bbccf3d2c8dc725bf1296f332bb61c0d91a69864291c810`

See more details on using hashes here.

Provenance

The following attestation bundles were made for aimz-0.14.0.tar.gz:

Publisher: publish.yaml on markean/aimz

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: aimz-0.14.0.tar.gz
- Subject digest: ad6b8dce8a41c04cfe4ce989b34ddebc65ddff14a6a62ef860038d2b4c4a5e2d
- Sigstore transparency entry: 2241207884
- Sigstore integration time: Jul 25, 2026
Source repository:
- Permalink: markean/aimz@2b3c82c831e26b62c906a94776d8ed856d7b5c3f
- Branch / Tag: refs/heads/main
- Owner: https://github.com/markean
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yaml@2b3c82c831e26b62c906a94776d8ed856d7b5c3f
- Trigger Event: workflow_dispatch