Normalizing flows

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.

The MAF architecture used here is the masked autoregressive flow of Papamakarios, Pavlakou, and Murray (2017), Masked Autoregressive Flow for Density Estimation. Credit the architecture you select as well as the backend implementation: Zuko or FlowJAX.

Torch training controls#

ZukoFlow.fit accepts these optional controls, all disabled by default:

  • clip_grad=None: set a finite positive total gradient-norm limit.

  • lr_annealing=False: enable cosine annealing over n_epochs, starting at lr and stepping after each epoch.

  • patience=None: set a positive number of consecutive epochs without a strictly lower validation loss before stopping. Ties count; improvement resets the counter.

  • show_progress=False: display epochs and training/validation losses.

flow.fit(
    particles, n_epochs=100, clip_grad=5.0, lr_annealing=True,
    patience=10, show_progress=True, rng=42,
)

Pass the same options through fit_kwargs in FlowTransform or FlowProposal. See fit() for all training defaults.

fit returns the adapter with losses for every completed epoch in fit_history and restores the best validation parameters. With validation_fraction=0, validation is skipped, validation_loss is empty, and final parameters are retained; setting patience raises ValueError. For nonzero fractions, an empty split after rounding uses all particles for both training and validation.

Raw and transformed particles must be finite with nonempty shape (n_samples, dims). Invalid data raises ValueError; non-finite errors identify the stage and affected dimensions (zero-based).

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#

Flow-preconditioned SMC is described by Karamanis et al. (2022), Accelerating astronomical and cosmological inference with Preconditioned Monte Carlo. This is the method paper; the separate pocoMC software paper describes that package.

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.