funcflows 0.1.2__py3-none-any.whl

FuncFlows/base_measures/__init__.py

@@ -0,0 +1,3 @@
+from .bases import Basis, CosineBasis, FourierBasis
+from .abstract_reference_measure import ReferenceMeasure
+from .gaussian_reference_measure import GaussianReferenceMeasure

FuncFlows/base_measures/abstract_reference_measure.py

@@ -0,0 +1,28 @@
+import torch
+
+from .bases import Basis
+
+
+class ReferenceMeasure:
+    """Probability measure on coefficient vectors [..., num_functions] in the given basis."""
+
+    def __init__(self, basis: Basis, dtype=torch.float64):
+        self.basis = basis
+        self.num_functions = basis.num_functions
+        self.dtype = dtype
+
+    def sample(self, num_samples: int):
+        """-> coeffs [num_samples, num_functions]"""
+        raise NotImplementedError
+
+    def log_density_diff(self, coeffs_from, coeffs_to):
+        """-> log p(coeffs_from) - log p(coeffs_to), shape [...]. The thing transport needs."""
+        raise NotImplementedError
+
+    def in_support(self, coeffs):
+        # always true: Gaussian / Student-t have full support. Adjust for bounded domains.
+        return torch.ones(coeffs.shape[:-1], dtype=torch.bool, device=coeffs.device)
+
+    def evaluate(self, coeffs, points):
+        """coeffs [..., num_functions] -> function values [..., num_points]"""
+        return coeffs @ self.basis.evaluate(points).T.to(coeffs)

FuncFlows/base_measures/bases.py

@@ -0,0 +1,112 @@
+import torch
+
+
+class Basis:
+    """Abstract basis class.
+
+    Args:
+        num_functions (int): Number of functions you want to play around with.
+        physical_dim (int, optional): Number of physical dimensions the functions will be over. Defaults to 1.
+        dtype (optional): Data type. If you're reading this and don't know what that is. God help you. Defaults to torch.float64.
+    """
+
+    def __init__(self, num_functions: int, physical_dim=1, dtype=torch.float64):
+        self.num_functions = num_functions
+        self.physical_dim = physical_dim
+        self.dtype = dtype
+
+    def evaluate(self, points):
+        """points [num_points, physical_dim] -> [num_points, num_functions], L2-orthonormal columns on [0,1]^d."""
+        raise NotImplementedError("You are calling an abstract eval method on a Basis class.")
+
+    def quadrature(self, num_points_per_axis):
+        quad_points = torch.linspace(0, 1, num_points_per_axis, dtype=self.dtype)
+        quad_weights = torch.full_like(quad_points, 1 / (num_points_per_axis - 1))
+        quad_weights[[0, -1]] /= 2
+        points = torch.cartesian_prod(*[quad_points] * self.physical_dim).reshape(-1, self.physical_dim)
+        weights = torch.cartesian_prod(*[quad_weights] * self.physical_dim).reshape(-1, self.physical_dim).prod(1)
+        return points, weights
+
+
+def _leading_modes(num_functions, physical_dim, per_axis, sort_key):
+    """The first num_functions multi-indices from a per_axis^d grid, ordered by sort_key.
+
+    per_axis only has to reach past the ball that holds num_functions modes. The n-th mode in d
+    dimensions sits at radius ~ (n / volume of the unit d-ball)^(1/d), so a bound of order
+    n^(1/d) is enough; arange(2n) per axis — the previous choice — built a (2n)^d grid, which at
+    2000 modes in 2-D is 16M entries sorted on every construction (seconds), and in 3-D cannot be
+    allocated at all. The selection is only right if the ball fits inside the grid, so the first
+    point outside the grid, (per_axis, 0, ...), must sort strictly after the last mode kept.
+    """
+    axes = [torch.arange(per_axis)] * physical_dim
+    modes = torch.cartesian_prod(*axes).reshape(-1, physical_dim) if physical_dim > 1 else axes[0][:, None]
+    order = torch.argsort(sort_key(modes), stable=True)[:num_functions]
+    outside = torch.zeros(1, physical_dim, dtype=modes.dtype)
+    outside[0, 0] = per_axis
+    if len(order) < num_functions or sort_key(outside) <= sort_key(modes[order[-1:]]):
+        raise ValueError(f"per-axis bound {per_axis} too small for {num_functions} modes in {physical_dim}-D")
+    return modes[order]
+
+
+class CosineBasis(Basis):
+    """Neumann Laplacian eigenfunctions on [0,1]^d: prod_axis norm * cos(pi m x).
+
+    Args:
+        num_functions (int): Number of functions you want to play around with.
+        physical_dim (int, optional): Number of physical dimensions the functions will be over. Defaults to 1.
+        dtype (optional): Data type. If you're reading this and don't know what that is. God help you. Defaults to torch.float64.
+    """
+
+    def __init__(self, num_functions: int, physical_dim: int = 1, dtype=torch.float64):
+        super().__init__(num_functions, physical_dim, dtype=dtype)
+        per_axis = num_functions + 1 if physical_dim == 1 else int(2 * num_functions ** (1 / physical_dim)) + 3
+        self.modes = _leading_modes(num_functions, physical_dim, per_axis,
+                                    lambda modes: (modes ** 2).sum(1)).to(self.dtype)
+        self.eval_norm = 1 + (2 ** 0.5 - 1) * (self.modes != 0).to(self.dtype)
+        self.laplacian_eigenvalues = ((torch.pi * self.modes) ** 2).sum(1)
+
+    def evaluate(self, points):
+        # one axis at a time: peak memory is num_points x num_functions, not x physical_dim as
+        # well. .to(points): the mode tables follow the points' device and dtype.
+        points = points.reshape(-1, self.physical_dim)
+        modes, norms = self.modes.to(points), self.eval_norm.to(points)
+        values = None
+        for axis in range(self.physical_dim):
+            factor = norms[:, axis] * torch.cos(torch.pi * points[:, axis, None] * modes[:, axis])
+            values = factor if values is None else values * factor
+        return values
+
+
+class FourierBasis(Basis):
+    """Periodic Laplacian eigenfunctions on [0,1]^d: 1, sqrt2 cos(2 pi k t), sqrt2 sin(2 pi k t), ...
+
+    Per axis, mode index m -> wavenumber ceil(m/2); odd m is a cosine, even m > 0 a sine.
+    Columns are ordered by Laplacian eigenvalue (2 pi)^2 * sum_axis wavenumber^2.
+
+    Args:
+        num_functions (int): Number of functions you want to play around with.
+        physical_dim (int, optional): Number of physical dimensions the functions will be over. Defaults to 1.
+        dtype (optional): Data type. If you're reading this and don't know what that is. God help you. Defaults to torch.float64.
+    """
+
+    def __init__(self, num_functions: int, physical_dim=1, dtype=torch.float64):
+        super().__init__(num_functions, physical_dim, dtype)
+        # two mode indices per wavenumber, so twice the cosine bound, plus slack
+        per_axis = num_functions + 1 if physical_dim == 1 else 2 * int(2 * num_functions ** (1 / physical_dim)) + 3
+        modes = _leading_modes(num_functions, physical_dim, per_axis,
+                               lambda modes: (((modes + 1) // 2) ** 2).sum(1))
+        self.wavenumbers = ((modes + 1) // 2).to(dtype)
+        self.phases = -torch.pi / 2 * ((modes % 2 == 0) & (modes > 0)).to(dtype)    # sine = shifted cosine
+        self.norm_factors = 1 + (2 ** 0.5 - 1) * (modes > 0).to(dtype)
+        self.laplacian_eigenvalues = ((2 * torch.pi * self.wavenumbers) ** 2).sum(1)
+
+    def evaluate(self, points):
+        points = points.reshape(-1, self.physical_dim)
+        wavenumbers, phases, norms = (self.wavenumbers.to(points), self.phases.to(points),
+                                      self.norm_factors.to(points))
+        values = None
+        for axis in range(self.physical_dim):                  # same axis-by-axis product
+            angles = 2 * torch.pi * points[:, axis, None] * wavenumbers[:, axis] + phases[:, axis]
+            factor = norms[:, axis] * torch.cos(angles)
+            values = factor if values is None else values * factor
+        return values

FuncFlows/base_measures/gaussian_reference_measure.py

@@ -0,0 +1,37 @@
+import torch
+
+from .abstract_reference_measure import ReferenceMeasure
+
+
+class GaussianReferenceMeasure(ReferenceMeasure):
+    """N(0, C0) with C0 = (I - alpha Laplacian)^(-power), diagonal in the basis:
+    variance_k = (1 + alpha * sigma_k)^(-power), sigma_k = basis.laplacian_eigenvalues[k].
+
+    The measure lives wherever its variances live: pass variances already on a device (or call
+    .move_to(device)) and sample() draws there. Previously sample() always drew on the CPU in the
+    default dtype, which every GPU/MPS run then had to work around with a subclass.
+    """
+
+    def __init__(self, basis, alpha=0.1, power=2.0, variances=None, dtype=torch.float64):
+        if variances is None:
+            variances = ((1 + alpha * basis.laplacian_eigenvalues) ** (-power)).to(dtype)
+        elif not torch.is_tensor(variances):
+            variances = torch.as_tensor(variances, dtype=dtype)
+        super().__init__(basis, variances.dtype)                    # dtype follows the variances
+        self.variances, self.scale = variances, variances.sqrt()
+
+    def move_to(self, device):
+        self.variances, self.scale = self.variances.to(device), self.scale.to(device)
+        return self
+
+    def sample(self, num_samples):
+        return self.scale * torch.randn(num_samples, self.num_functions,
+                                        dtype=self.scale.dtype, device=self.scale.device)
+
+    def log_density_diff(self, coeffs_from, coeffs_to):
+        # only this simple because the modes are independent
+        return 0.5 * ((coeffs_to ** 2 - coeffs_from ** 2) / self.variances).sum(-1)
+
+    def cameron_martin_norm2(self, coeffs):
+        """Squared Cameron-Martin norm: the L2 norm after whitening by the covariance."""
+        return (coeffs ** 2 / self.variances).sum(-1)

FuncFlows/diagnostics/__init__.py

@@ -0,0 +1,2 @@
+from .importance import ImportanceCorrection
+from .coverage import coverage_curve, coverage_error

FuncFlows/diagnostics/coverage.py

@@ -0,0 +1,46 @@
+import torch
+
+
+def coverage_curve(samples, truths, references=None, weights=None, levels=None):
+    """TARP: expected coverage from samples alone (Lemos et al. 2023).
+
+    For each test case, count the fraction of posterior draws that land closer to a random
+    reference point than the truth does. If the posterior is right, those fractions are uniform,
+    so plotting the empirical CDF against the credibility level gives the diagonal.
+
+    Unlike rank statistics it needs no density and no per-dimension marginalisation, and unlike
+    simulation-based calibration it is necessary AND sufficient: a model that ignores the data and
+    returns the prior is perfectly self-consistent but fails this.
+
+    Read the deviation like this: BELOW the diagonal is overconfident (posteriors too narrow),
+    ABOVE is conservative (too wide). Both are failures, and they look nothing alike, which is
+    why the curve is worth more than the scalar.
+
+    samples [cases, draws, modes]; truths [cases, modes]; references default to the truths rolled
+    by one, which is a draw from the same marginal without needing the prior. weights scale the
+    distance, so pass the Cameron-Martin weights to measure in the metric the prior lives in.
+    """
+    levels = torch.linspace(0.05, 0.95, 19) if levels is None else levels
+    levels = levels.to(device=samples.device, dtype=samples.dtype)
+    if references is None:
+        references = truths.roll(1, dims=0)
+    scaled_samples = samples if weights is None else samples * weights
+    scaled_truths = truths if weights is None else truths * weights
+    scaled_references = references if weights is None else references * weights
+    truth_distance = (scaled_truths - scaled_references).norm(dim=-1)
+    draw_distance = (scaled_samples - scaled_references.unsqueeze(1)).norm(dim=-1)
+    closer = (draw_distance < truth_distance.unsqueeze(1)).to(samples.dtype).mean(1)
+    coverage = (closer.unsqueeze(0) <= levels.unsqueeze(1)).to(samples.dtype).mean(1)
+    return levels, coverage
+
+
+def coverage_error(levels, coverage):
+    """One number for the curve: the largest gap from the diagonal, signed by direction.
+
+    Positive means conservative on balance, negative overconfident. On a Gaussian check the exact
+    posterior scored 0.03 with 800 cases, a posterior 40% too narrow scored -0.18, and one 60% too
+    wide +0.19 -- so roughly 0.05 is the noise floor at that many cases.
+    """
+    deviation = coverage - levels
+    largest = deviation.abs().argmax()
+    return deviation[largest].item()

FuncFlows/diagnostics/importance.py

@@ -0,0 +1,60 @@
+import math
+
+import torch
+
+
+class ImportanceCorrection:
+    """Turn an approximate posterior into an asymptotically exact one, and measure how good it was
+    (Dax et al. 2023).
+
+    The proposal is whatever produced the draws -- an amortised posterior, a reverse-KL fit, a
+    guided sampler. The weights are
+
+        log w = -Phi(x) - log d(q)/d(mu0)(x)
+
+    i.e. the true unnormalised posterior over the proposal, both taken against the same reference
+    measure so the Radon-Nikodym derivatives are the ones the flow already computes. Nothing is
+    trained and nothing is assumed about the proposal beyond being able to evaluate its density.
+
+    Three things come out:
+      * a corrected posterior -- reweight, or resample, and every expectation is consistent;
+      * `efficiency`, the effective sample size over the sample count, which is a HARD verdict on
+        the proposal (Dax et al. report a median of about 10% and treat low values as failures);
+      * `log_evidence`, unbiased up to the reference measure's own normalisation, so model
+        comparison between structures comes free.
+
+    The catch is honest: efficiency falls with dimension, and a proposal narrower than the
+    posterior can score well on every self-consistency check and still be wrong -- which is
+    exactly what the efficiency number exposes.
+    """
+
+    def __init__(self, flow, draws, potential, context=None):
+        self.draws = draws
+        with torch.no_grad():
+            self.log_ratio = -potential(draws) - flow.log_rn_at(draws, context)
+        self.weights = torch.softmax(self.log_ratio, 0)
+
+    @property
+    def efficiency(self):
+        """Effective sample size over the sample count, in [1/count, 1]."""
+        return (1 / (len(self.weights) * self.weights.pow(2).sum())).item()
+
+    @property
+    def log_evidence(self):
+        """log mean w: the evidence, up to the reference measure's normalisation."""
+        return (torch.logsumexp(self.log_ratio, 0) - math.log(len(self.log_ratio))).item()
+
+    def mean(self):
+        return self.weights @ self.draws
+
+    def standard_deviation(self):
+        centred = self.draws - self.mean()
+        return (self.weights @ centred.pow(2)).sqrt()
+
+    def resample(self, count=None):
+        """Equally weighted draws, for anything downstream that cannot carry weights (plots)."""
+        count = count or len(self.draws)
+        offsets = (torch.rand((), device=self.weights.device, dtype=self.weights.dtype)
+                   + torch.arange(count, device=self.weights.device, dtype=self.weights.dtype)) / count
+        picked = torch.searchsorted(self.weights.cumsum(0), offsets).clamp(max=len(self.draws) - 1)
+        return self.draws[picked]

FuncFlows/objectives/__init__.py

@@ -0,0 +1,5 @@
+from .reverse_kl import ReverseKL
+from .negative_logl import NegativeLogL
+from .alpha_divergence import AlphaDivergence
+from .flow_matching import (FlowMatching, ConditionalFlowMatching, PairedFlowMatching,
+                            reflow_pairs, optimal_pairing, straight_path_loss)

FuncFlows/objectives/alpha_divergence.py

@@ -0,0 +1,67 @@
+import math
+
+import torch
+
+
+class AlphaDivergence:
+    """Renyi / alpha divergence between the flow q and the posterior, as a mass-covering
+    alternative to reverse KL:
+
+        D_alpha = log E_q[w^alpha] / (alpha (alpha - 1)),    w = ptilde / q.
+
+    alpha = 2 is the chi^2 divergence of FAB (Midgley et al. 2023): the variance of the
+    importance weights, which punishes exactly the regions where q is too narrow. 0 < alpha < 1
+    sits between reverse KL (alpha -> 0) and the evidence (alpha -> 1 from below in the Li &
+    Turner 2016 convention). The 1/(alpha - 1) is essential: for 0 < alpha < 1 the expectation is
+    MAXIMISED at q = p, so dividing by alpha alone would push the flow away from the posterior.
+
+    Where the samples come from decides whether this is usable. Estimated with draws from q
+    alone, E_q[w^alpha] for alpha > 1 only sees where q already has mass, and the gradient can
+    make q ever narrower (log q -> infinity, w -> 0, loss -> -infinity) without the estimator
+    ever noticing the posterior mass it left behind -- the true quantity is bounded below by
+    Z^alpha, the estimate is not. FAB fixes this with AIS toward ptilde^alpha q^(1-alpha). The fix
+    here is simpler: a defensive mixture proposal r = (1 - f) q + f mu0, evaluated by importance
+    weights,
+
+        E_q[w^alpha] = E_r[ ptilde^alpha q^(1-alpha) / r ],
+
+    with everything written against the reference measure, so only log dq/dmu0 (the quantity the
+    flow already computes) and Phi appear. The prior half covers the whole posterior support, so
+    a q that collapses is caught by the prior draws it no longer covers. prior_fraction = 0 is
+    the plain q-sample estimator, kept for comparison; the default 0.5 is what to use.
+
+    log-sum-exp throughout, so the weights never leave log space.
+    """
+
+    def __init__(self, transformation, potential, num_samples=30, alpha=2.0, context=None,
+                 prior_fraction=0.5):
+        if alpha in (0.0, 1.0):
+            raise ValueError("alpha = 0 and alpha = 1 are the KL limits; use ReverseKL")
+        self.transformation, self.potential, self.num_samples = transformation, potential, num_samples
+        self.alpha, self.context, self.prior_fraction = alpha, context, prior_fraction
+
+    def __call__(self):
+        flow, alpha = self.transformation, self.alpha
+        num_prior = int(round(self.prior_fraction * self.num_samples))
+        num_flow = self.num_samples - num_prior
+        points, log_ratio = [], []                              # log_ratio = log dq/dmu0 at the point
+        if num_flow:
+            coeffs_out, log_rn = flow.push_forward(flow.base_measure.sample(num_flow), self.context)
+            points.append(coeffs_out)
+            log_ratio.append(log_rn)
+        if num_prior:
+            prior_points = flow.base_measure.sample(num_prior)
+            points.append(prior_points)
+            log_ratio.append(flow.log_rn_at(prior_points, self.context))
+        points, log_ratio = torch.cat(points), torch.cat(log_ratio)
+        log_potential = -self.potential(points)
+        if num_prior == 0 or num_flow == 0:
+            # single proposal (q or mu0): log of the integrand relative to that proposal
+            log_proposal = log_ratio if num_prior == 0 else torch.zeros_like(log_ratio)
+        else:
+            fraction = num_prior / self.num_samples
+            log_proposal = torch.logaddexp(math.log(1 - fraction) + log_ratio,
+                                           torch.full_like(log_ratio, math.log(fraction)))
+        # ptilde^alpha q^(1-alpha) / r, all relative to mu0: alpha (-Phi) + (1 - alpha) log_ratio - log_proposal
+        log_terms = alpha * log_potential + (1 - alpha) * log_ratio - log_proposal
+        return (torch.logsumexp(log_terms, 0) - math.log(len(log_terms))) / (alpha * (alpha - 1))

FuncFlows/objectives/flow_matching.py

@@ -0,0 +1,132 @@
+import torch
+
+
+class FlowMatching:
+    """Regress the field onto the straight-path velocity from the reference measure to the data.
+
+    coupling: 'independent' pairs each base draw with a random target (Lipman et al. 2023;
+    FFM). 'optimal' re-pairs the minibatch by exact optimal transport in the training metric
+    (Tong et al. 2023 OT-CFM; in function space FOT-CFM, Li et al. 2026): straighter flows,
+    lower-variance targets, and fewer integration steps at sampling time.
+    """
+
+    def __init__(self, transformation, dataset_coeffs, batch_size, weights=None, coupling="independent"):
+        self.field, self.measure = transformation.vector_field, transformation.base_measure
+        self.total_time = transformation.total_time
+        self.dataset_coeffs, self.batch_size, self.weights = dataset_coeffs, batch_size, weights
+        self.coupling = coupling
+
+    def __call__(self):
+        if callable(self.dataset_coeffs):
+            target_coeffs = self.dataset_coeffs(self.batch_size)
+        else:
+            index = torch.randint(len(self.dataset_coeffs), (self.batch_size,),
+                                  device=self.dataset_coeffs.device)
+            target_coeffs = self.dataset_coeffs[index]
+        start_coeffs = self.measure.sample(self.batch_size)
+        return straight_path_loss(self.field, start_coeffs, target_coeffs, self.total_time, None,
+                                  self.weights, self.coupling)
+
+
+class PairedFlowMatching:
+    """Reflow (Liu et al. 2022): regress on the model's OWN (base, endpoint) pairs.
+
+    A trained flow defines a coupling whose marginals are the base measure and the model, and
+    whose paths cross. Refitting a field to those exact pairs -- never redrawing the base --
+    removes the crossings without changing either marginal, so the flow gets straighter and the
+    same sample needs fewer integration steps. Repeatable; each round compounds, and each round
+    also compounds whatever error the previous model had, so two or three is the usual limit.
+    """
+
+    def __init__(self, transformation, base_coeffs, target_coeffs, batch_size, weights=None):
+        self.field, self.total_time = transformation.vector_field, transformation.total_time
+        self.base_coeffs, self.target_coeffs = base_coeffs, target_coeffs
+        self.batch_size, self.weights = batch_size, weights
+
+    def __call__(self):
+        index = torch.randint(len(self.base_coeffs), (self.batch_size,),
+                              device=self.base_coeffs.device)
+        return straight_path_loss(self.field, self.base_coeffs[index], self.target_coeffs[index],
+                                  self.total_time, None, self.weights, "independent")
+
+
+def reflow_pairs(flow, count, batch_size=512):
+    """Base draws and where the flow sends them: the dataset PairedFlowMatching refits."""
+    base, target = [], []
+    with torch.no_grad():
+        for done in range(0, count, batch_size):
+            draws = flow.base_measure.sample(min(batch_size, count - done))
+            base.append(draws)
+            target.append(flow.transport(draws))
+    return torch.cat(base), torch.cat(target)
+
+
+class ConditionalFlowMatching:
+    """Same, but base and target come from a simulator and the field sees the context.
+
+    coupling='optimal' pairs base draws with targets ACROSS contexts, which biases which base
+    draw reaches each context, so the conditional flow no longer starts from the base measure.
+    context_weight is the fix (Chemseddine et al. 2024; Kerrigan et al. 2024): the transport cost
+    carries the squared context distance too, so a pair is only swapped when the contexts nearly
+    agree. Large weight -> only near-identical contexts are re-paired, and the conditional is
+    preserved; zero -> the biased coupling. It only does anything when the simulator repeats or
+    nearly repeats contexts within a batch.
+    """
+
+    def __init__(self, transformation, simulate, draw_base, batch_size, weights=None,
+                 coupling="independent", context_weight=0.0):
+        self.field, self.total_time = transformation.vector_field, transformation.total_time
+        self.simulate, self.draw_base = simulate, draw_base
+        self.batch_size, self.weights, self.coupling = batch_size, weights, coupling
+        self.context_weight = context_weight
+
+    def __call__(self):
+        target_coeffs, context = self.simulate(self.batch_size)
+        start_coeffs = self.draw_base(self.batch_size)
+        return straight_path_loss(self.field, start_coeffs, target_coeffs, self.total_time, context,
+                                  self.weights, self.coupling, self.context_weight)
+
+
+def squared_distances(left, right):
+    """[n, m] of squared distances, by expansion rather than an [n, m, dim] difference."""
+    return ((left ** 2).sum(-1)[:, None] + (right ** 2).sum(-1)[None, :] - 2 * left @ right.T)
+
+
+def optimal_pairing(start_coeffs, target_coeffs, weights, context=None, context_weight=0.0):
+    """Minibatch optimal transport: the permutation of the targets that minimises the summed
+    squared distance to the base draws, in the metric the loss uses (weights = the same per-mode
+    weights as the loss, so 'cameron-martin' and 'l2' stay consistent).
+
+    For uniform equal-size batches the OT plan is a permutation, so an exact assignment is the
+    whole solver: O(batch^3) on the CPU, about 2 ms at batch 128. One device sync per step.
+    """
+    from scipy.optimize import linear_sum_assignment       # only needed for this coupling
+    scaled_start = start_coeffs if weights is None else start_coeffs * weights
+    scaled_target = target_coeffs if weights is None else target_coeffs * weights
+    cost = squared_distances(scaled_start, scaled_target)
+    if context is not None and context_weight:
+        cost = cost + context_weight * squared_distances(context, context)
+    _, columns = linear_sum_assignment(cost.detach().cpu().numpy())
+    return torch.as_tensor(columns, device=target_coeffs.device)
+
+
+def straight_path_loss(field, start_coeffs, target_coeffs, total_time, context, weights,
+                       coupling="independent", context_weight=0.0):
+    """One time per sample when the field allows it, otherwise one time for the whole batch."""
+    if coupling == "optimal":
+        order = optimal_pairing(start_coeffs, target_coeffs, weights, context, context_weight)
+        target_coeffs = target_coeffs[order]
+        context = None if context is None else context[order]     # the context stays with its target
+    elif coupling != "independent":
+        raise ValueError(f"coupling must be 'independent' or 'optimal', not {coupling!r}")
+    batch = len(start_coeffs)
+    if getattr(field, "supports_batched_time", False):
+        times = torch.rand(batch, dtype=start_coeffs.dtype, device=start_coeffs.device) * total_time
+    else:
+        times = torch.as_tensor(torch.rand(()).item() * total_time, dtype=start_coeffs.dtype,
+                                device=start_coeffs.device)
+    path_coeffs = start_coeffs + times.reshape(-1, 1) / total_time * (target_coeffs - start_coeffs)
+    residual = field.velocity(path_coeffs, times, context) - (target_coeffs - start_coeffs) / total_time
+    if weights is not None:
+        residual = residual * weights
+    return residual.pow(2).sum(-1).mean()

FuncFlows/objectives/negative_logl.py

@@ -0,0 +1,15 @@
+import torch
+
+
+class NegativeLogL:
+    """P2 Alg. 1 (direct data): -mean_i log d(f#mu0)/dmu0 (u_i) over a minibatch of dataset coefficients."""
+
+    def __init__(self, transformation, dataset_coeffs, batch_size=64, context=None):
+        self.transformation, self.dataset_coeffs, self.batch_size = transformation, dataset_coeffs, batch_size
+        self.context = context
+
+    def __call__(self):
+        batch_indices = torch.randint(len(self.dataset_coeffs), (self.batch_size,),
+                                      device=self.dataset_coeffs.device)      # no host round trip
+        context = None if self.context is None else self.context[batch_indices]
+        return -self.transformation.log_rn_at(self.dataset_coeffs[batch_indices], context).mean()

FuncFlows/objectives/reverse_kl.py

@@ -0,0 +1,48 @@
+import torch
+
+
+class ReverseKL:
+    """E_{coeffs ~ mu0}[ log_rn_weight + Phi(f(coeffs)) ]: KL(f#mu0 || posterior) up to the evidence.
+
+    context, if given, is handed to the transformation unchanged (one context for the whole
+    objective -- the single-observation posterior case).
+
+    path_gradient=True uses the path-gradient (sticking-the-landing) estimator of Vaitl et al.
+    2022. The reparameterised gradient splits into a path term and a term that differentiates the
+    density at a FIXED sample; the second has zero mean and pure variance, and it vanishes only in
+    expectation, so dropping it leaves an unbiased estimator with strictly less noise. The gain
+    grows as the fit improves: in the Gaussian check in this package's notes the variance ratio is
+    about 1.4x far from the optimum and over 1000x near it, which is exactly where the ordinary
+    estimator stalls. It also skips the trace on the forward pass, so a step is no more expensive.
+    """
+
+    def __init__(self, transformation, potential, num_samples=30, context=None, path_gradient=False):
+        self.transformation, self.potential, self.num_samples = transformation, potential, num_samples
+        self.context, self.path_gradient = context, path_gradient
+
+    def __call__(self):
+        coeffs = self.transformation.base_measure.sample(self.num_samples)
+        if not self.path_gradient:
+            coeffs_out, log_rn_weight = self.transformation.push_forward(coeffs, self.context)
+            return (log_rn_weight + self.potential(coeffs_out)).mean()
+        coeffs_out = self.transformation.transport(coeffs, self.context)     # no trace needed
+        value, slope = self.integrand_and_slope(coeffs_out.detach())
+        surrogate = (slope * coeffs_out).sum(-1).mean()
+        return value + surrogate - surrogate.detach()       # value reports, gradient is the path term
+
+    def integrand_and_slope(self, points):
+        """log (dq/dmu0)(x) + Phi(x) and its x-derivative, with the flow's parameters held fixed."""
+        moving = [parameter for parameter in self.transformation.parameters()
+                  if parameter.requires_grad]
+        for parameter in moving:
+            parameter.requires_grad_(False)
+        try:
+            with torch.enable_grad():
+                state = points.detach().requires_grad_(True)
+                integrand = (self.transformation.log_rn_at(state, self.context)
+                             + self.potential(state))
+                slope, = torch.autograd.grad(integrand.sum(), state)
+        finally:
+            for parameter in moving:
+                parameter.requires_grad_(True)
+        return integrand.mean().detach(), slope.detach()

FuncFlows/samplers/__init__.py

@@ -0,0 +1 @@
+from .latent_pcn import latent_pcn

FuncFlows/samplers/latent_pcn.py

@@ -0,0 +1,53 @@
+import torch
+
+
+def latent_pcn(flow, potential, num_chains=128, num_steps=2000, beta=0.2, init=None, context=None,
+               burn=None, thin=10, adapt_to=None):
+    """Exact posterior sampling under a trained prior flow, by pCN in the flow's latent space.
+
+    The flow prior is the push-forward of the base measure mu0 (Gaussian) through T, so a draw is
+    f = T(z), z ~ mu0, and the posterior over z is
+
+        pi(dz) propto exp(-Phi(T(z))) mu0(dz)
+
+    which is the textbook setting for preconditioned Crank-Nicolson (Cotter et al. 2013): the
+    proposal z' = sqrt(1 - beta^2) z + beta xi, xi ~ mu0, is mu0-reversible, so the acceptance
+    ratio is exp(Phi(T(z)) - Phi(T(z'))) and the prior never enters -- no density of the flow,
+    no trace, no Jacobian. Each step costs one transport of the chain batch.
+
+    Compared with GuidedTransport / TwistedSampler this is asymptotically exact and has no
+    endpoint-estimate heuristic in it; the price is num_steps transports instead of one. Start it
+    from good draws (init = guided draws, pulled back through the flow) and the burn-in is short.
+
+    init: coefficients in the flow's output space, one per chain; pulled back to z. Default: mu0.
+    adapt_to: target acceptance rate; if set, beta is rescaled every 50 steps during burn-in
+    (pCN acceptance falls with beta; ~0.25 is a common target).
+    Returns (draws [kept, num_functions], info dict with acceptance and final beta).
+    """
+    measure = flow.base_measure
+    burn = num_steps // 4 if burn is None else burn
+    with torch.no_grad():
+        if init is None:
+            latent = measure.sample(num_chains)
+        else:
+            latent = flow.pull_back(init[:num_chains], context)
+            num_chains = len(latent)
+        state = flow.transport(latent, context)
+        energy = potential(state)
+        keep, accepted, window = [], 0, 0
+        for step in range(num_steps):
+            proposal_latent = (1 - beta ** 2) ** 0.5 * latent + beta * measure.sample(num_chains)
+            proposal = flow.transport(proposal_latent, context)
+            proposal_energy = potential(proposal)
+            accept = torch.log(torch.rand_like(energy)) < energy - proposal_energy
+            latent = torch.where(accept[:, None], proposal_latent, latent)
+            state = torch.where(accept[:, None], proposal, state)
+            energy = torch.where(accept, proposal_energy, energy)
+            accepted += accept.double().mean().item()
+            window += accept.double().mean().item()
+            if adapt_to is not None and step < burn and (step + 1) % 50 == 0:
+                beta = float(min(0.999, max(1e-3, beta * (1.0 + 0.5 * (window / 50 - adapt_to)))))
+                window = 0
+            if step >= burn and (step - burn) % thin == 0:
+                keep.append(state.clone())
+    return torch.cat(keep), dict(acceptance=accepted / num_steps, beta=beta)

FuncFlows/transports/__init__.py

@@ -0,0 +1,3 @@
+from .abstract_transformation import Transformation
+from .layers import DiscreteTransformation
+from .continuous import ContinuousTransformation