funcyflows 0.2.0__tar.gz → 0.2.1__tar.gz

{funcyflows-0.2.0 → funcyflows-0.2.1}/FuncyFlows/samplers/latent_pcn.py

@@ -1,86 +1,5 @@
-r"""Preconditioned Crank-Nicolson MCMC in the latent space of a trained flow.
-
-What problem this solves
-------------------------
-You have a posterior on coefficients,
-
-    pi(dv)  proportional to  exp(-Phi(v)) rho(dv),
-
-with Phi the misfit (minus log likelihood) and rho a prior. Sampling it directly with a random-walk
-Metropolis proposal fails as the number of modes M grows: to keep the acceptance rate away from
-zero the step size has to shrink like M^(-1/2), so the chain needs O(M) steps to move anywhere.
-
-pCN (Cotter, Roberts, Stuart & White, Statistical Science 2013, arXiv:1202.0709) fixes that for a
-GAUSSIAN reference measure mu0. Its proposal is
-
-    z' = sqrt(1 - beta^2) z + beta xi,        xi ~ mu0,
-
-which leaves mu0 exactly invariant -- if z ~ mu0 then (z, z') is jointly Gaussian and exchangeable,
-so the pair is reversible. Because the proposal already carries the reference measure, the
-Metropolis ratio keeps only the likelihood:
-
-    accept with probability  min(1, exp(Phi(z) - Phi(z'))).
-
-No prior density appears, nothing scales with M, and beta can stay O(1) at any truncation.
-
-What the "latent" part adds
----------------------------
-pCN needs the reference measure to be Gaussian, which is a real restriction: it means the geometry
-the sampler assumes is the geometry of mu0. If the posterior is curved, multimodal or far from the
-prior, a Gaussian-geometry sampler crawls.
-
-So run pCN not on v but on the flow's latent variable z, where v = T(z) and T is the trained
-transport. The reference measure there IS Gaussian by construction -- it is the flow's base measure
--- and the flow has already absorbed the awkward geometry. This is transport-map preconditioning
-(Parno & Marzouk, arXiv:1412.5492; the neural version, with HMC rather than pCN, is NeuTra,
-arXiv:1903.03704).
-
-The potential to hand it
-------------------------
-The chain targets  nu(dz) proportional to exp(-Psi(z)) mu0(dz), and v = T(z) is then distributed as
-
-    exp(-Psi(T^-1 v)) q(dv),        q = T_# mu0  (the flow's pushforward).
-
-Setting that equal to pi gives the rule, once and for all:
-
-    Psi(z) = Phi(T z) + log (dq/dmu0)(T z) - log (drho/dmu0)(T z).
-
-Two cases cover everything in this package, and they are NOT the same call:
-
-  * T is a learned PRIOR, so rho = q. The two log terms cancel:
-
-        Psi = Phi . T                      potential = misfit
-
-    The flow's Jacobian never enters, so ANY trained prior flow works and the trace need not even
-    be available. This is what the inpainting examples used.
-
-  * T approximates the POSTERIOR and the prior is the base measure, rho = mu0. Then:
-
-        Psi = Phi . T + log_rn_at . T      potential = lambda v: misfit(v) + flow.log_rn_at(v)
-
-    Here the trace does enter, exactly. An estimated (Hutchinson) trace makes the chain converge to
-    the wrong measure, not to a noisy version of the right one. See bimodal_posterior.py.
-
-`potential` is called on COEFFICIENTS v, never on z; this function does the transport for you.
-
-When it works and when it does not
-----------------------------------
-If the flow is good, q is close to pi, Psi is nearly constant, and almost every proposal is accepted
-even at beta close to 1 -- the chain returns near-independent samples. If the flow is poor, or the
-likelihood is much sharper than anything the flow has seen, the adapted beta collapses and the
-chain barely moves. That failure is quiet: the chains sit on their initialisation and the output
-looks like a posterior. `info["moved"]` is reported for exactly this reason -- see its description
-in the Returns section, and prefer an amortised conditional flow when it is small.
-"""
-import math
-
-import torch
-
-try:                                                     # progress bar is optional
-    from tqdm.auto import trange
-except ImportError:                                      # pragma: no cover
-    def trange(n, **kwargs):
-        return range(n)
+import math, torch
+from tqdm.auto import trange
 
 
 def _as_coefficients(result):
@@ -93,7 +12,10 @@ def _as_coefficients(result):
 @torch.no_grad()
 def latent_pcn(flow, potential, num_chains=64, num_steps=2000, beta=0.2, init=None, burn=None,
                thin=1, adapt_to=None, context=None, progress=True, generator=None):
-    r"""Sample a posterior by running pCN in the latent space of `flow`.
+    r"""Sample a posterior by running pCN in the latent space of `flow`. If unfamiliar with pCN,
+    it's an MCMC method that makes proposals by scaling the current MCMC sample and adding
+    an adjusted sample from the prior. Scales better with dimension than standard MCMC because
+    of the implicit geometry added along with the prior samples.
 
     Parameters
     ----------
@@ -104,8 +26,8 @@ def latent_pcn(flow, potential, num_chains=64, num_steps=2000, beta=0.2, init=No
         training, and every chain step costs one solve.
     potential : callable
         ``potential(v) -> [batch]``, the NEGATIVE log of the target's density relative to the
-        flow's pushforward, as a function of COEFFICIENTS. See the module docstring for which of
-        the two forms you want; getting this wrong is silent.
+        flow's pushforward, as a function of COEFFICIENTS. See explanation further down
+         in this docstring for which of the two forms you want.
     num_chains : int
         Chains advanced in lockstep. They share one ``transport`` call per step, so more chains are
         nearly free up to the batch the ODE can hold.
@@ -148,6 +70,79 @@ def latent_pcn(flow, potential, num_chains=64, num_steps=2000, beta=0.2, init=No
                          norm. Below ~0.1 the chain never left its initialisation and the "posterior"
                          you are looking at is whatever you passed as ``init``. Reported only when
                          ``init`` is given.
+
+
+    Preconditioned Crank-Nicolson MCMC in the latent space of a trained flow.
+
+    What we're doing
+    ------------------------
+    You have a posterior on coefficients,
+
+        pi(dv)  proportional to  exp(-Phi(v)) rho(dv),
+
+    with Phi the misfit (minus log likelihood) and rho a prior. Sampling it directly with a random-walk
+    Metropolis proposal fails as the number of modes M grows: to keep the acceptance rate away from
+    zero the step size has to shrink like M^(-1/2), so the chain needs O(M) steps to move anywhere.
+
+    pCN (Cotter, Roberts, Stuart & White, Statistical Science 2013, arXiv:1202.0709) fixes that for a
+    GAUSSIAN reference measure mu0. Its proposal is
+
+        z' = sqrt(1 - beta^2) z + beta xi,        xi ~ mu0,
+
+    which leaves mu0 exactly invariant -- if z ~ mu0 then (z, z') is jointly Gaussian and exchangeable,
+    so the pair is reversible. Because the proposal already carries the reference measure, the
+    Metropolis ratio keeps only the likelihood:
+
+        accept with probability  min(1, exp(Phi(z) - Phi(z'))).
+
+    No prior density appears, nothing scales with M, and beta can stay O(1) at any truncation.
+
+    What the "latent" part adds
+    ---------------------------
+    pCN needs the reference measure to be Gaussian, because the Gaussian family is closed under this operation 
+    and almost nothing else is: with v and xi independent N(0,C), the combination has covariance 
+    
+        (1-beta^2)C + beta^2 C = C, 
+        
+    and a mean-zero Gaussian is determined by its covariance, so v' is N(0,C) again.
+
+    If the posterior is curved, multimodal or far from the prior, a Gaussian-geometry sampler does not do very well. 
+    
+    The constraint is on the REFERENCE measure, not the prior. An arbitrary prior rho enters through -log(drho/dmu0) 
+    in the potential (formula above) with no change to the algorithm, or gets absorbed into the transport, which is what a learned prior flow is.
+    What must be closed under the proposal is the reference: Gaussian is the only finite-variance choice, 
+    though alpha-stable references work with rebalanced coefficients.
+
+    But, we run pCN not on v but on the flow's latent variable z, where v = T(z) and T is the trained
+    transport. The reference measure there *IS* Gaussian by construction -- it is the flow's base measure
+    -- and the flow has already absorbed the awkward geometry. This is transport-map preconditioning
+    (Parno & Marzouk, arXiv:1412.5492; the neural version, with HMC rather than pCN, is NeuTra,arXiv:1903.03704).
+
+
+
+    Which potential to pass
+    -----------------------
+    The chain targets  exp(-Psi(z)) mu0(dz), so v = T(z) comes out distributed as
+
+        Psi(z) = Phi(Tz) + log (dq/dmu0)(Tz) - log (drho/dmu0)(Tz),     q = T_# mu0
+
+    Two cases cover the below, they are not the same call :grimace: :
+
+      * T is a learned PRIOR, so rho = q and the log terms cancel:
+
+            potential = misfit
+
+        The flow's Jacobian never enters, so any trained prior flow works and the trace
+        need not even be computable.
+
+      * T approximates the POSTERIOR and the prior is the base measure, rho = mu0:
+
+            potential = lambda v: misfit(v) + flow.log_rn_at(v)
+
+        Here the trace DOES enter, and it unfortunately has to be exact. 
+        A Hutchinson estimate doesn't give a noisy version of the 'correct' chain:
+            - the randomness lands in the accept test
+            - so the chain has a different invariant measure. 
     """
     if not 0 < beta <= 1:
         raise ValueError(f"beta must be in (0, 1]; got {beta}")

{funcyflows-0.2.0 → funcyflows-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: funcyflows
-Version: 0.2.0
+Version: 0.2.1
 Summary: Normalizing flows and flow matching on function spaces, with exact-trace vector fields
 Author: Liam Pinchbeck
 License: MIT

{funcyflows-0.2.0 → funcyflows-0.2.1}/funcyflows.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: funcyflows
-Version: 0.2.0
+Version: 0.2.1
 Summary: Normalizing flows and flow matching on function spaces, with exact-trace vector fields
 Author: Liam Pinchbeck
 License: MIT

{funcyflows-0.2.0 → funcyflows-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "funcyflows"
-version = "0.2.0"
+version = "0.2.1"
 description = "Normalizing flows and flow matching on function spaces, with exact-trace vector fields"
 readme = "README.md"
 license = { text = "MIT" }