scripts/grid1d.py

"""Deterministic (grid) density propagation for 1-D NON-quadratic potentials.

Needed because a quadratic V makes every scheme a linear-Gaussian recursion, so
the *first-order* weak error that separates LMC from ULMC vanishes identically.
On a general smooth potential it does not, and the comparison becomes faithful.

Each 1-d coordinate potential f(x) is handled exactly:
  ULMC  : 2-d chain (x,p); one step = deterministic map + fixed 2-d Gaussian
          convolution (Eq. 3.2 with grad V(x) = f'(x))
  LMC   : 1-d chain; x' = x - h f'(x) + N(0, 2h)
  comp. : 1-d chain; exponential integrator that treats the alpha-strongly
          convex quadratic part exactly (Freund et al. 2022 composite scheme)

Propagation is by mass-conserving bilinear scatter of the deterministic image
followed by an FFT convolution with the (fixed) noise kernel.  No sampling, no
seeds.  Accuracy is validated against the closed-form Gaussian answer in
scripts/test_grid1d.py.
"""

import numpy as np

import ulmc_core as U


class Grid2D:
    def __init__(self, Lx, nx, Lp, np_, f, df, gamma):
        self.x = np.linspace(-Lx, Lx, nx, endpoint=False)
        self.p = np.linspace(-Lp, Lp, np_, endpoint=False)
        self.dx = self.x[1] - self.x[0]
        self.dp = self.p[1] - self.p[0]
        self.X, self.P = np.meshgrid(self.x, self.p, indexing="ij")
        self.f, self.df, self.g = f, df, gamma
        lp = -(f(self.X) + self.P**2 / 2)
        lp -= np.log(np.sum(np.exp(lp)) * self.dx * self.dp)
        self.logpi = lp
        self.pi = np.exp(lp)

    def _kernel(self, Q):
        """Periodic 2-d Gaussian kernel with covariance Q, centred at 0."""
        xs = np.fft.fftfreq(len(self.x), 1.0 / len(self.x)) * self.dx
        ps = np.fft.fftfreq(len(self.p), 1.0 / len(self.p)) * self.dp
        XX, PP = np.meshgrid(xs, ps, indexing="ij")
        Qi = np.linalg.inv(Q)
        k = np.exp(
            -0.5 * (Qi[0, 0] * XX**2 + 2 * Qi[0, 1] * XX * PP + Qi[1, 1] * PP**2)
        )
        return k / k.sum()

    def ulmc_step_setup(self, h):
        g = self.g
        e = np.exp(-g * h)
        c1 = (1 - e) / g
        G = (h - c1) / g
        Q = np.array(
            [
                [U.var_xi1(h, g), U.cov_xi1_xi2(h, h, g)],
                [U.cov_xi1_xi2(h, h, g), U.var_xi2(h, g)],
            ]
        )
        dv = self.df(self.X)
        xn = self.X + c1 * self.P - G * dv
        pn = e * self.P - c1 * dv
        return xn, pn, np.fft.rfft2(self._kernel(Q))

    def scatter(self, xn, pn, w):
        """Mass-conserving bilinear deposit of weights w at positions (xn,pn)."""
        nx, np_ = len(self.x), len(self.p)
        fx = (xn - self.x[0]) / self.dx
        fp = (pn - self.p[0]) / self.dp
        i0 = np.floor(fx).astype(np.int64)
        j0 = np.floor(fp).astype(np.int64)
        tx = fx - i0
        tp = fp - j0
        out = np.zeros((nx, np_))
        for di in (0, 1):
            for dj in (0, 1):
                ii = np.clip(i0 + di, 0, nx - 1)
                jj = np.clip(j0 + dj, 0, np_ - 1)
                wt = w * (tx if di else 1 - tx) * (tp if dj else 1 - tp)
                np.add.at(out, (ii.ravel(), jj.ravel()), wt.ravel())
        return out

    def run_ulmc(self, h, rho0, nsteps, record_every=1):
        xn, pn, K = self.ulmc_step_setup(h)
        rho = rho0.copy()
        kls = []
        for n in range(nsteps):
            rho = self.scatter(xn, pn, rho)
            rho = np.fft.irfft2(np.fft.rfft2(rho) * K, s=rho.shape)
            rho = np.maximum(rho, 0.0)
            rho /= rho.sum() * self.dx * self.dp
            if (n + 1) % record_every == 0:
                kls.append(self.kl(rho))
        return rho, np.array(kls)

    def kl(self, rho):
        m = rho > 1e-300
        return float(
            np.sum(rho[m] * (np.log(rho[m]) - self.logpi[m])) * self.dx * self.dp
        )


class Grid1D:
    """Overdamped chains (LMC / composite)."""

    def __init__(self, Lx, nx, f, df):
        self.x = np.linspace(-Lx, Lx, nx, endpoint=False)
        self.dx = self.x[1] - self.x[0]
        self.f, self.df = f, df
        lp = -f(self.x)
        lp -= np.log(np.sum(np.exp(lp)) * self.dx)
        self.logpi = lp

    def _kernel(self, var):
        xs = np.fft.fftfreq(len(self.x), 1.0 / len(self.x)) * self.dx
        k = np.exp(-0.5 * xs**2 / var)
        return k / k.sum()

    def scatter(self, xn, w):
        nx = len(self.x)
        fx = (xn - self.x[0]) / self.dx
        i0 = np.floor(fx).astype(np.int64)
        tx = fx - i0
        out = np.zeros(nx)
        for di in (0, 1):
            ii = np.clip(i0 + di, 0, nx - 1)
            np.add.at(out, ii, w * (tx if di else 1 - tx))
        return out

    def run(self, scheme, h, rho0, nsteps, alpha=None, record_every=1):
        if scheme == "lmc":
            xn = self.x - h * self.df(self.x)
            var = 2 * h
        elif scheme == "composite":
            ea = np.exp(-alpha * h)
            c = (1 - ea) / alpha
            xn = ea * self.x - c * (self.df(self.x) - alpha * self.x)
            var = (1 - np.exp(-2 * alpha * h)) / alpha
        else:
            raise ValueError(scheme)
        K = np.fft.rfft(self._kernel(var))
        rho = rho0.copy()
        kls = []
        for n in range(nsteps):
            rho = self.scatter(xn, rho)
            rho = np.fft.irfft(np.fft.rfft(rho) * K, n=len(rho))
            rho = np.maximum(rho, 0.0)
            rho /= rho.sum() * self.dx
            if (n + 1) % record_every == 0:
                kls.append(self.kl(rho))
        return rho, np.array(kls)

    def kl(self, rho):
        m = rho > 1e-300
        return float(np.sum(rho[m] * (np.log(rho[m]) - self.logpi[m])) * self.dx)


class Grid2DPullback(Grid2D):
    """Semi-Lagrangian (exact pullback) ULMC propagator -- far more accurate
    than mass scattering.

    One ULMC step is  T(x,p) = (x + c1 p - G f'(x),  e p - c1 f'(x)).
    Eliminating p:   e x' - c1 p' = e x + (c1^2 - e G) f'(x) =: phi(x),
    so x = phi^{-1}(e x' - c1 p'),  p = (p' + c1 f'(x)) / e, and
    |det DT(x)| = phi'(x) = e + (c1^2 - e G) f''(x).
    """

    def __init__(self, Lx, nx, Lp, np_, f, df, ddf, gamma):
        super().__init__(Lx, nx, Lp, np_, f, df, gamma)
        self.ddf = ddf

    def ulmc_pullback_setup(self, h, nfine=200001, pad=6.0):
        from scipy import ndimage  # noqa: F401  (imported for run_ulmc_pb)
        g = self.g
        e = np.exp(-g * h)
        c1 = (1 - e) / g
        G = (h - c1) / g
        kap = c1 * c1 - e * G
        xf = np.linspace(self.x[0] - pad, self.x[-1] + pad, nfine)
        phif = e * xf + kap * self.df(xf)
        assert np.all(np.diff(phif) > 0), "phi not monotone: h too large"
        u = e * self.X - c1 * self.P
        xs = np.interp(u, phif, xf)
        ps = (self.P + c1 * self.df(xs)) / e
        jac = e + kap * self.ddf(xs)
        Q = np.array([[U.var_xi1(h, g), U.cov_xi1_xi2(h, h, g)],
                      [U.cov_xi1_xi2(h, h, g), U.var_xi2(h, g)]])
        ci = np.stack([(xs - self.x[0]) / self.dx, (ps - self.p[0]) / self.dp])
        return ci, jac, np.fft.rfft2(self._kernel(Q))

    def run_ulmc_pb(self, h, rho0, nsteps, record_every=1, order=3):
        from scipy import ndimage
        ci, jac, K = self.ulmc_pullback_setup(h)
        rho = rho0.copy()
        kls = []
        for n in range(nsteps):
            rho = ndimage.map_coordinates(rho, ci, order=order, mode="constant",
                                          cval=0.0, prefilter=True) / jac
            rho = np.maximum(rho, 0.0)
            rho = np.fft.irfft2(np.fft.rfft2(rho) * K, s=rho.shape)
            rho = np.maximum(rho, 0.0)
            rho /= rho.sum() * self.dx * self.dp
            if (n + 1) % record_every == 0:
                kls.append(self.kl(rho))
        return rho, np.array(kls)