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.

Particle transform restrictions

The flow’s particle_transform must not contain a PeriodicBoundaryTransform, including inside nested CompositeTransform objects. Both adapters reject these transforms on construction, assignment, fitting, and checkpoint restoration. Periodic wrapping is many-to-one: a correct density would require summing over periodic images, which these adapters do not implement. A Jacobian correction alone is insufficient.

Existing checkpoints containing periodic particle transforms are also rejected. Nonperiodic transforms, such as AffineTransform, remain supported.

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.