funcyflows 0.2.0__py3-none-any.whl

FuncyFlows/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

FuncyFlows/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)

FuncyFlows/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

FuncyFlows/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)

FuncyFlows/diagnostics/__init__.py

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

FuncyFlows/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()

FuncyFlows/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]

FuncyFlows/examples/__init__.py

@@ -0,0 +1 @@
+"""Runnable examples. `python -m FuncyFlows.examples` copies them into the current directory."""

FuncyFlows/examples/__main__.py

@@ -0,0 +1,36 @@
+"""Copy the example scripts into the current directory:
+
+    python -m FuncyFlows.examples            # copies *.py (skips files that already exist)
+    python -m FuncyFlows.examples --force    # overwrite
+    python -m FuncyFlows.examples --list
+
+Then, e.g.:  python bimodal_posterior.py
+"""
+import pathlib
+import shutil
+import sys
+
+HERE = pathlib.Path(__file__).parent
+EXAMPLES = sorted(p for p in HERE.glob("*.py") if not p.name.startswith("__"))
+
+
+def main(argv):
+    if "--list" in argv:
+        for path in EXAMPLES:
+            doc = (path.read_text().split('"""')[1].strip().splitlines() or [""])[0]
+            print(f"  {path.name:24s} {doc}")
+        return
+    force = "--force" in argv
+    target = pathlib.Path.cwd()
+    for path in EXAMPLES:
+        dest = target / path.name
+        if dest.exists() and not force:
+            print(f"  exists, skipped: {dest.name}   (--force to overwrite)")
+            continue
+        shutil.copy(path, dest)
+        print(f"  wrote {dest.name}")
+    print("\nrun any of them with  python <name>.py  (they import _common.py from the same directory)")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])