FuncyFlows

Normalizing flows and flow matching on function spaces.

A function is represented by its coefficients on a Laplacian eigenbasis (CosineBasis or FourierBasis), the reference measure is a Gaussian on those coefficients, and a transport is a neural ODE in coefficient space. The vector fields (LinearField, MatrixField) have closed-form divergences, so the density of the transported measure relative to the reference measure is exact — no Hutchinson estimators — which makes reverse-KL training, likelihood training and importance reweighting usable at hundreds of modes.

pip install funcyflows            # torch + tqdm

Examples

python -m FuncyFlows.examples        # copies the example scripts into the current directory
python bimodal_posterior.py         # then run any of them

The 2D examples and the two 1D prior examples use the spatial PointwiseField; the bimodal one is coefficient-space only. Each takes 3–12 minutes on a CPU. Except for cloud inpainting and phase inpainting which can take up to 50mins.

The pieces

Every field takes mode_scale=measure.scale: the nonlinearity sees whitened coefficients and the divergence is unchanged (similarity transform). Always pass it.

1. Learn a prior from samples (flow matching)

import torch
from FuncyFlows.base_measures import CosineBasis, GaussianReferenceMeasure
from FuncyFlows.transports.continuous import (ContinuousTransformation, SumField, LinearField,
                                             MatrixField, TimeBasisConditioner)
from FuncyFlows.objectives import FlowMatching
from FuncyFlows.utils.train import train

M = 64
basis = CosineBasis(M)
coeffs = ...                                  # [N, M] training functions projected onto the basis
measure = GaussianReferenceMeasure(basis, variances=coeffs.var(0))   # diagonal Gaussian fit = the GP baseline

field = SumField(
    LinearField(M, num_time_modes=4),
    MatrixField(TimeBasisConditioner(M, 256, num_time_modes=4), mode_scale=measure.scale))
flow = ContinuousTransformation(measure, field, num_steps=16)

losses = train(FlowMatching(flow, coeffs, batch_size=256, weights=1 / measure.scale, coupling="optimal"),
               flow.parameters(), num_steps=10000, learning_rate=2e-3)

draws = flow.transport(measure.sample(1000))              # [1000, M] new functions
points = torch.linspace(0, 1, 400).reshape(-1, 1)
values = draws @ basis.evaluate(points).T                 # [1000, 400] on a grid
log_q = flow.log_rn_at(draws)                             # exact log density w.r.t. the reference measure

weights=1/measure.scale whitens the loss so high modes count; coupling="optimal" needs scipy and only makes sense for an unconditional flow.

2. Posterior from a likelihood (no samples needed)

from FuncyFlows.objectives import ReverseKL
from FuncyFlows.utils.gaussian_misfit import GaussianMisfit
from FuncyFlows.diagnostics import ImportanceCorrection

prior = GaussianReferenceMeasure(basis, alpha=0.05, power=2.0)     # Matérn-like spectrum
phi_obs = basis.evaluate(obs_points)                               # [n_obs, M]
forward = lambda v: (v @ phi_obs.T) ** 2                           # any differentiable map
potential = GaussianMisfit(forward, data, noise_std)               # -log likelihood, [batch] -> [batch]

flow = ContinuousTransformation(prior, field, num_steps=16)        # same field recipe as above
train(ReverseKL(flow, potential, num_samples=64, path_gradient=True), flow.parameters(), num_steps=3000)

draws = flow.transport(prior.sample(2000))
check = ImportanceCorrection(flow, draws, potential)               # exact because the traces are exact
print(check.efficiency, check.log_evidence)                        # efficiency near 1 = posterior matched

For an exact sampler on top of a learned prior use latent_pcn(flow, potential): pCN in the flow's latent space, whose acceptance needs no Jacobian.

3. Amortised posterior (condition on data)

from FuncyFlows.transports.continuous import DataConditioner
from FuncyFlows.objectives import ConditionalFlowMatching

C = ...                                                    # context dimension (e.g. whitened projected data)
field = SumField(
    LinearField(M, C, num_time_modes=4, mode_scale=measure.scale),          # data drift: the linear-Gaussian mean
    MatrixField(DataConditioner(M, 256, C, num_time_modes=4), mode_scale=measure.scale))
flow = ContinuousTransformation(measure, field, num_steps=16)

def simulate(batch):                                       # -> (target coeffs [batch, M], context [batch, C])
    ...

train(ConditionalFlowMatching(flow, simulate, measure.sample, batch_size=256, weights=1 / measure.scale),
      flow.parameters(), num_steps=20000, learning_rate=2e-3)

posterior = flow.transport(measure.sample(1000), context.expand(1000, -1))   # one ODE solve per posterior

Whiten the context by measure.scale when it is itself a projected function; otherwise the drift weights for high modes have to reach 1/σ and never do. When a conjugate (linear-Gaussian) posterior is available in closed form, use it as the base measure and let the flow learn the residual.

Diagnostics

from FuncyFlows.diagnostics import coverage_curve, coverage_error
levels, coverage = coverage_curve(samples, truths, weights=1 / measure.scale)   # TARP; samples [cases, draws, M]
coverage_error(levels, coverage)     # signed max gap from the diagonal: negative = overconfident

Layout

FuncyFlows/
  base_measures/     bases, GaussianReferenceMeasure
  transports/        continuous/: fields, conditioners, ContinuousTransformation
                     (grid_fields.py: PointwiseField, one FNO-style spatial layer with exact trace)
  examples/          python -m FuncyFlows.examples
  objectives/        flow_matching, reverse_kl, negative_logl
  samplers/          latent_pcn
  diagnostics/       importance, coverage
  utils/             train, gaussian_misfit