nbi: neural bayesian inference

Documentation

Do you have challenging inference problems that are difficult to solve with standard optimization and/or MCMC methods? Are you looking to fit the same forward model to thousands or millions of observed targets? nbi may be your solution.

nbi is an engine for Neural Posterior Estimation (NPE) focused on out-of-the-box functionality for astronomical data, particularly light curves and spectra. nbi provides effective embedding/featurizer networks for spectra and light-curve data, along with importance-sampling integration that enables asymptotically exact inference so that the inference results are interpretable and trustworthy.

Installation

You may either install nbi with pip install nbi or directly from source. As nbi is currently under active development, installing from source may be preferable at this stage.

git clone https://github.com/kmzzhang/nbi.git
cd nbi
pip install .

If you are using Mac ARM CPU (i.e. M1/M2/M3), you might want to install PyTorch from source and disable NNPACK, which is known to reduce performance (see issue). Note that currently the MPS Also support for weight_norm on Mac M1-M3 GPUs is recently implemented but has not been included in a stable release yet. Installing the nightly version from source also enables weight_norm for the MPS device.

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
USE_NNPACK=0 python setup.py install

Quick Start

The examples/ directory contains complete examples that demonstrates the functionality of nbi. A bare-bone example below illustrates the basic API, which follows the scikit-learn style. The default featurizer network for sequential data is resnet-gru, which is a hybrid CNN-RNN architecture.

Here are a rule of thumb for resnet-gru hyperparameters:

• dim_in: this is your number of input data channels
• depth: number of ResNet blocks. Start near log2(L)-5, where L is length of your sequential data.
• max_hidden: Maximum hidden dimensions for ResNet. Hidden dimensions double (from hidden_conv=32 by default) every depth. At least a few times D^2, where D is the dimension of the physical parameter space.

import nbi

# hyperparameters
featurizer = {
    "type": "resnet-gru",
    "dim_in": 1,
    "max_hidden": 64
}

flow = {
    "n_dims": 1,        # parameter space dimension
    "flow_hidden": 32,  # generally no larger than max_hidden
    "num_blocks": 10    # depends on complexity of posterior shape
}

engine = nbi.NBI(
    flow,
    featurizer,
    simulator,
    noise,
    priors,
    device='cpu'        # 'cuda', 'cuda:0', 'mps' for M1/M2 Mac GPU
)
engine.fit(
    n_sims=1000,
    n_rounds=1,
    n_epochs=100
)
y_pred, weights = engine.predict(x_obs, x_err, n_samples=2000)

References

nbi: the Astronomer's Package for Neural Posterior Estimation (Zhang et al. 2023). Accepted to NeurIPS 2023 Workshop on Deep Learning and Inverse Problems.

Masked Autoregressive Flow for Density Estimation (Papamakarios et al. 2017)
https://arxiv.org/abs/1705.07057

Featurizers: ResNet (He et al. 2015; https://arxiv.org/abs/1512.03385), Gated Recurrent Units (GRU; Cho et al. 2014; https://arxiv.org/abs/1406.1078), ResNet-GRU (Zhang et al. 2021; https://iopscience.iop.org/article/10.3847/1538-3881/abf42e)

Acknowledgments

The nbi package is expanded from code originally written for ''Real-time Likelihood-free Inference of Roman Binary Microlensing Events with Amortized Neural Posterior Estimation''' (Zhang et al. 2021). The Masked Autoregressive Flow in this package is partly adapted from the implementation in https://github.com/kamenbliznashki/normalizing_flows. Work on this project was supported by the National Science Foundation award #2206744 ("CDS&E: Accelerating Astrophysical Insight at Scale with Likelihood-Free Inference").