aimz: Scalable probabilistic impact modeling#

Version: 0.12.0

Useful links: Home | Code Repository | Issues | Releases

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.

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