Installation | Tutorial | User Guide | FAQs | Changelog

aimz is a Python library for scalable probabilistic impact modeling, enabling assessment of intervention effects on outcomes with a streamlined interface for fitting, sampling, prediction, and effect estimation—minimal boilerplate, accelerated execution, and powered by NumPyro, JAX, Xarray, and Zarr.

• Intuitive API combining the ease of use from ML frameworks with the flexibility of probabilistic modeling.
• Accelerated computation via parallelism and distributed data.
• Support for interventional causal inference for counterfactuals and causal effects.
• MLflow integration for experiment tracking and model management.

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.

from aimz import ImpactModel

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

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

# Initialize ImpactModel
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)

# Make predictions or posterior predictive samples
dt = im.predict(X)

See the Contributing Guide to get started.