Normalizing flows

Install the optional dependencies for your chosen backend:

python -m pip install "sequince-density[zuko]"
# Alternatively, for the JAX backend:
python -m pip install "sequince-density[flowjax]"

Fit a Zuko flow to a two-dimensional sample population, then draw samples and evaluate their log densities:

import numpy as np

from sequince_density import ZukoFlow

rng = np.random.default_rng(42)
particles = rng.normal(size=(2000, 2))

flow = ZukoFlow(dims=2, flow_class="MAF", hidden_features=[32, 32])
flow.fit(particles, n_epochs=100, batch_size=500, rng=42)
samples, log_q = flow.sample_and_log_prob(256, rng=43)

ZukoFlow returns PyTorch tensors; FlowJax returns JAX arrays.

Preconditioned sampling

For preconditioned SMC, pass an unfitted flow to FlowTransform and use the transform as the particle_transform argument to make_standard_bayesian_smc. The mutation step fits it to the current particle population:

from sequince.transforms import FlowTransform

transform = FlowTransform(
    ZukoFlow(dims=2, flow_class="MAF", hidden_features=[32, 32]),
    fit_kwargs={"n_epochs": 100, "batch_size": 500},
)

See Proposal distributions to use a flow as a fitted starting distribution instead.