Metadata-Version: 2.4
Name: metabeta
Version: 0.4.2
Summary: Amortized inference for Bayesian GLMMs
License-Expression: CC-BY-NC-SA-4.0
Project-URL: Homepage, https://github.com/adkipnis/metabeta
Project-URL: Repository, https://github.com/adkipnis/metabeta
Project-URL: Checkpoints, https://huggingface.co/adkipnis/metabeta
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: <3.13,>=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: arviz>=0.23.4
Requires-Dist: huggingface-hub>=1.14.0
Requires-Dist: joblib>=1.5.2
Requires-Dist: matplotlib>=3.10.7
Requires-Dist: numpy>=2.0.2
Requires-Dist: pandas>=3.0.3
Requires-Dist: pydantic>=2.13.4
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: scikit-learn>=1.8.0
Requires-Dist: scipy>=1.16.3
Requires-Dist: seaborn>=0.13.2
Requires-Dist: tabulate>=0.9.0
Requires-Dist: torch>=2.9.1
Requires-Dist: tqdm>=4.67.3
Provides-Extra: research
Requires-Dist: bambi>=0.17.2; extra == "research"
Requires-Dist: datasets>=4.4.1; extra == "research"
Requires-Dist: pmlb>=1.0.1.post3; extra == "research"
Requires-Dist: pymc>=5.28.5; extra == "research"
Requires-Dist: rpy2>=3.6.7; extra == "research"
Requires-Dist: schedulefree>=1.4.1; extra == "research"
Requires-Dist: statsmodels>=0.14.6; extra == "research"
Requires-Dist: ucimlrepo>=0.0.7; extra == "research"
Requires-Dist: wandb>=0.25.0; extra == "research"
Dynamic: license-file

# metabeta

[![Python 3.12](https://img.shields.io/badge/python-3.12-blue)](https://www.python.org/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.x-ee4c2c)](https://pytorch.org/)
[![uv](https://img.shields.io/badge/env-uv-6b47ff)](https://docs.astral.sh/uv/)
[![Release checks](https://github.com/adkipnis/metabeta/actions/workflows/release-checks.yml/badge.svg)](https://github.com/adkipnis/metabeta/actions/workflows/release-checks.yml)
[![Hugging Face](https://img.shields.io/badge/🤗-checkpoints-yellow)](https://huggingface.co/adkipnis/metabeta)
[![License: CC BY-NC-SA 4.0](https://img.shields.io/badge/license-CC%20BY--NC--SA%204.0-lightgrey)](LICENSE)

`metabeta` is an amortized model for Bayesian generalized linear mixed-effects models (GLMMs).
Given a grouped dataset, a model formula, and an optional prior specification, it returns
posterior samples for fixed effects, random effects, variance components, and
correlations with the speed and batching capabilities of a PyTorch forward pass.

- Supports Normal, Bernoulli, and Poisson outcomes.
- Accepts convenient lme4-style formulas.
- Conditions on prior family and hyperparameters at inference time.
- Supports batched prior-sensitivity analysis in one `sample()` call.
- Includes diagnostics for posterior contraction and prediction accuracy.
- Ships the pretrained inference API, plotting helpers, diagnostics, and demo notebooks.
- Keeps simulation, training, reference-method evaluation, and experiment scripts in the
  source repository for research workflows.

Pretrained checkpoints are hosted on [Hugging Face](https://huggingface.co/adkipnis/metabeta)
and are downloaded automatically on first use.

## Quick start

For local development, install from source with [`uv`](https://docs.astral.sh/uv/):

```bash
git clone https://github.com/adkipnis/metabeta.git
cd metabeta
uv sync
uv pip install -e .
```

Run posterior inference on a grouped dataframe:

```python
import statsmodels.api as sm
from metabeta.models.api import Api

mb = Api.from_pretrained("normal")
df = sm.datasets.get_rdataset("sleepstudy", "lme4").data
result = mb.sample(df, formula="Reaction ~ Days + (Days | Subject)", n_samples=1000)
print(mb.posteriorSummary(result))
```

See [demos/intro.ipynb](demos/intro.ipynb) for the full `sleepstudy`
walkthrough and [demos/priors.ipynb](demos/priors.ipynb) for an exemplary prior-sensitivity
analysis.

## Development and research workflows

The default PyPI install is for pretrained-model inference. To reproduce training runs,
generate synthetic datasets, benchmark against reference methods, or extend the package,
clone the repository and install the optional research dependencies:

```bash
git clone https://github.com/adkipnis/metabeta.git
cd metabeta
uv sync --extra research --group simulation --dev
uv pip install -e ".[research]"
```

This enables the repository-only paths for synthetic data generation, model training,
benchmarking against reference methods, simulation-based calibration studies, and experiment
scripts. The `simulation` dependency group installs the GitHub-only `scamd` dependency used
by SCM dataset generation. The `metabeta.posthoc` package is included for experimental
posterior refinement, but it is not part of the production pretrained API and may change
without a deprecation window.

## From simulation to deployment

Each pretrained model is built through the same pipeline, from a dataset simulator to a checkpoint that can be loaded by the public API.

1. **Define the model family.** Choose the likelihood, GLMM dimensions, group structure,
   covariate styles, and prior families covered by a checkpoint.
2. **Generate training data.** Simulate hierarchical datasets and posterior reference targets
   across the configured design space.
3. **Format inputs.** Convert grouped tabular data into padded tensors, masks, and prior encodings shared by training and inference.
4. **Train amortized posteriors.** Use set-transformers (to learn low-dimensional permutation-invariant summaries of the datasets) and conditional coupling
   flows to approximate posteriors for fixed effects, variance parameters, correlations, and group-wise random effects.
5. **Validate behavior.** Check parameter recovery, credible interval coverage, simulation-based calibration, posterior
   predictive accuracy and run comparisons against reference methods.
6. **Package checkpoints.** Bundle weights, configs, routing metadata, and
   preprocessing expectations into a joint checkpoint.
7. **Deploy through the API.** Load joint checkpoints with
   `Api.from_pretrained(...)` for convenient posterior estimation and diagnostics.

## Repository map

| Path | Contents |
|---|---|
| [metabeta/analytical/](metabeta/analytical/) | GLMM analytical fits and helpers |
| [metabeta/configs/](metabeta/configs/) | model and preset configuration files |
| [metabeta/datasets/](metabeta/datasets/) | preprocessing and source-specific dataset fetchers |
| [metabeta/evaluation/](metabeta/evaluation/) | parameter recovery, coverage, SBC, posterior predictive checks and summary metrics |
| [metabeta/models/](metabeta/models/) | model API, set transformers, normalizing flows |
| [metabeta/plotting/](metabeta/plotting/) | plot functions for posterior samples, recovery, calibration, and runtime |
| [metabeta/posthoc/](metabeta/posthoc/) | experimental post-hoc posterior refinement helpers |
| [metabeta/simulation/](metabeta/simulation/) | synthetic hierarchical data generation and reference fitting with PyMC |
| [metabeta/training/](metabeta/training/) | training entry point and checkpoint loop |
| [metabeta/utils/](metabeta/utils/) | config, dataloading, routing, IO, and shared helper code |
| [experiments/](experiments/) | reproducible experiment scripts grouped by package area |
| [tests/](tests/) | pytest suite for models, simulation, evaluation, datasets, and utils |
| [demos/](demos/) | demo notebooks |
