lecture_1/fields.py

"""Score functions, the Laplacian and the Fokker-Planck equation.

The score of a density is the gradient of its logarithm, a vector field of the same shape as a
velocity field, and the normalizing constant differentiates away. The Laplacian is the divergence
of the gradient and measures curvature, so the heat equation sets the rate of change of a density
proportional to it. The Fokker-Planck equation holds both mechanisms at once, and it returns the
continuity equation when the noise magnitude is zero and the heat equation when the drift is zero.

Chapter 1, Sections 1.31 and 1.32 of the notes. Runs on the Gaussians of Section 1.31, the bowl and
the saddle of Section 1.32, the widening Gaussian p_t = N(0, t) of Section 1.32.1 and the expanding
flow v(x) = x of Section 1.30.2. Prints the chapter's number beside the computed one.

    python lecture_1/fields.py               # every worked example
    python lecture_1/fields.py --steps 2000  # a finer Euler-Maruyama walk
    python lecture_1/fields.py --quiet       # the closing report only
"""

# %% 01. Imports and configuration
from __future__ import annotations

import argparse
import math

import torch

from cis6270.runner import print_report, save_run, set_seed

torch.set_default_dtype(torch.float64)  # Double precision, because the chapter prints five decimal places.


SCORE_POINTS = torch.tensor([3.0, -2.0, 0.0])  # The three states scored in Section 1.31.
WIDE_MEAN, WIDE_VARIANCE = 1.0, 4.0  # The general Gaussian N(mu, sigma^2) of Section 1.31.
DRIFT = 0.5  # The constant drift b_t(x) = 0.5 of the combined process in Section 1.32.2.
NOISE = 1.0  # The noise magnitude sigma_t, so that D = sigma_t^2 / 2 = 0.5 as in Section 1.32.1.


WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.


def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
    if not quiet:  # Under --quiet only the closing report is printed.
        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
              f"   gap {abs(by_hand - computed):.1e}")


def log_gaussian(x: torch.Tensor, mu: float = 0.0, variance: float = 1.0) -> torch.Tensor:
    """log N(x; mu, sigma^2) = -(x - mu)^2 / (2 sigma^2) - (1/2) log(2 pi sigma^2)."""
    normalizer = 0.5 * math.log(2.0 * math.pi * variance)  # The term with no x in it.
    return -((x - mu) ** 2) / (2.0 * variance) - normalizer


def gaussian(t: float, x: torch.Tensor) -> torch.Tensor:
    """p_t = N(0, t) in one dimension, the widening Gaussian of Section 1.32.1."""
    return (-(x ** 2) / (2.0 * t)).exp() / math.sqrt(2.0 * math.pi * t)


def score(log_density, x: torch.Tensor) -> torch.Tensor:
    """Section 1.31: s(x) = grad_x log p(x), taken by automatic differentiation."""
    point = x.detach().clone().requires_grad_(True)  # A leaf to differentiate against; [N].
    total = log_density(point).sum()  # The points do not interact, so one pass computes every score.
    gradient, = torch.autograd.grad(total, point)  # d/dx log p at each point; [N].
    return gradient


def laplacian(field, x: torch.Tensor) -> float:
    """Section 1.32: the Laplacian is the divergence of the gradient, the trace of the Hessian."""
    return float(torch.autograd.functional.hessian(field, x).trace())


def fokker_planck_residual(p, b, sigma: float, x: torch.Tensor, t: float) -> float:
    """Section 1.32.2: d p_t / d t + div(b_t p_t) - (1/2) sigma_t^2 Laplacian p_t, by differences.

    We recover the continuity equation of Section 1.30 by setting sigma to zero, and the heat
    equation of Section 1.32.1 by setting b to zero, so one residual covers all three statements.
    """
    h = 1e-4  # The central-difference step, small enough for five digits in double precision.
    time_rate = (p(t + h, x) - p(t - h, x)) / (2.0 * h)  # d p_t(x) / d t at fixed x; [grid].
    flux = (b(x + h) * p(t, x + h) - b(x - h) * p(t, x - h)) / (2.0 * h)  # div(b_t p_t); [grid].
    curvature = (p(t, x + h) - 2.0 * p(t, x) + p(t, x - h)) / h ** 2  # Laplacian p_t; [grid].
    return float((time_rate + flux - 0.5 * sigma ** 2 * curvature).abs().max())


# %% 02. The score of a Gaussian
def scores(quiet: bool = False) -> dict:
    """Section 1.31: the score of N(0, 1) is -x, and of N(mu, sigma^2) it is -(x - mu)/sigma^2."""
    standard = score(log_gaussian, SCORE_POINTS)  # s(x) at x = 3, -2 and 0; [3].
    closed_form = -SCORE_POINTS  # The chapter's result, with the log(2 pi) term differentiated away.
    unnormalized = score(lambda z: -(z ** 2) / 2.0, SCORE_POINTS)  # The same density without its Z.
    wide = score(lambda z: log_gaussian(z, WIDE_MEAN, WIDE_VARIANCE), torch.tensor([5.0]))  # [1].
    show("s(x) at x = 3", -3.0, float(standard[0]), quiet)
    show("s(x) at x = -2", 2.0, float(standard[1]), quiet)
    show("s(x) at the mode x = 0", 0.0, float(standard[2]), quiet)
    show("largest |s(x) + x|", 0.0, float((standard - closed_form).abs().max()), quiet)
    show("largest change from dropping Z", 0.0, float((unnormalized - standard).abs().max()), quiet)
    show("s(5) under N(1, 4)", -1.0, float(wide[0]), quiet)
    return {"standard": standard, "wide": float(wide[0])}


# %% 03. The Laplacian of a scalar field
def curvatures(quiet: bool = False) -> dict:
    """Section 1.32: the bowl x_1^2 + 3 x_2^2 has Laplacian 8, and the saddle x_1^2 - x_2^2 has 0."""
    def bowl(x: torch.Tensor) -> torch.Tensor:
        return x[0] ** 2 + 3.0 * x[1] ** 2  # The field whose two curvatures are 2 and 6.

    def saddle(x: torch.Tensor) -> torch.Tensor:
        return x[0] ** 2 - x[1] ** 2  # The field whose two curvatures cancel.

    point = torch.tensor([1.0, 2.0])  # The Laplacians below are the same at every point.
    second = torch.autograd.functional.hessian(bowl, point)  # The two pure second derivatives; [2, 2].
    h = 1e-3  # The spacing of the four neighbors used for the average reading of the Laplacian.
    steps = torch.tensor([[h, 0.0], [-h, 0.0], [0.0, h], [0.0, -h]])  # The four offsets; [4, 2].
    neighbors = sum(bowl(point + step) for step in steps) - 4.0 * bowl(point)  # Above the center.
    show("d^2 f / d x_1^2", 2.0, float(second[0, 0]), quiet)
    show("d^2 f / d x_2^2", 6.0, float(second[1, 1]), quiet)
    show("Laplacian of x_1^2 + 3 x_2^2", 8.0, laplacian(bowl, point), quiet)
    show("Laplacian from the four neighbors", 8.0, float(neighbors) / h ** 2, quiet)
    show("Laplacian of the saddle", 0.0, laplacian(saddle, point), quiet)
    return {"bowl": laplacian(bowl, point), "saddle": laplacian(saddle, point)}


# %% 04. The heat equation
def heat(grid: int, quiet: bool = False) -> dict:
    """Section 1.32.1: at D = 1/2 the Gaussian N(0, t) solves d p_t / d t = D Laplacian p_t."""
    diffusion = 0.5 * NOISE ** 2  # D = sigma_t^2 / 2, the coefficient the chapter sets to 1/2.
    center = 1.0 / math.sqrt(2.0 * math.pi)  # p_1(0), the height at the middle at t = 1.
    time_rate = -0.5 / math.sqrt(2.0 * math.pi)  # d p_t(0)/dt = -(1/2)(2 pi)^{-1/2} t^{-3/2} at t = 1.
    curvature = laplacian(lambda z: gaussian(1.0, z), torch.tensor([0.0]))  # Laplacian p_1 at 0.
    x = torch.linspace(-3.0, 3.0, grid)  # Points covering the bulk of the mass at t = 1; [grid].
    residual = fokker_planck_residual(
        gaussian, torch.zeros_like, NOISE, x, 1.0)  # The heat equation, with no drift.
    show("p_1(0) = 1 / sqrt(2 pi)", 0.39894, center, quiet)
    show("d p_t(0) / d t at t = 1", -0.19947, time_rate, quiet)
    show("Laplacian of p_1 at 0", -0.39894, curvature, quiet)
    show("-p_1(0) / sigma^2 at sigma^2 = 1", -0.39894, -center / 1.0, quiet)
    show("D times the Laplacian of p_1 at 0", -0.19947, diffusion * curvature, quiet)
    show("largest heat-equation residual", 0.0, residual, quiet)
    return {"center": center, "time_rate": time_rate, "curvature": curvature,
            "flow": diffusion * curvature, "residual": residual}


# %% 05. Drift and diffusion together
def fokker_planck(grid: int, quiet: bool = False) -> dict:
    """Section 1.32.2: a drift that translates and a noise that widens, in one equation.

    With a constant drift b and sigma_t = 1 the density N(bt, t) solves the equation, since the
    drift moves the center and the noise grows the variance at the rate of the heat equation.
    """
    def translating(t: float, x: torch.Tensor) -> torch.Tensor:
        return gaussian(t, x - DRIFT * t)  # N(bt, t), with the center moved along by the drift.

    def expanding(t: float, x: torch.Tensor) -> torch.Tensor:
        return torch.full_like(x, math.exp(-t))  # p_t = e^{-t} on [0, e^t], Section 1.30.2.

    x = torch.linspace(-3.0, 3.0, grid)  # Points across the translating Gaussian; [grid].
    both = fokker_planck_residual(translating, lambda z: torch.full_like(z, DRIFT), NOISE, x, 1.0)
    inside = torch.linspace(0.1, 1.4, grid)  # Interior points of the support at t = 0.5; [grid].
    transport = fokker_planck_residual(expanding, lambda z: z, 0.0, inside, 0.5)  # sigma_t = 0.
    center = DRIFT * 1.0  # The mean of p_t at t = 1, set by the drift alone.
    show("center of p_t at t = 1", 0.5, center, quiet)
    show("largest residual, drift and noise", 0.0, both, quiet)
    show("largest residual at sigma_t = 0", 0.0, transport, quiet)
    return {"both": both, "transport": transport, "center": center}


# %% 06. The process behind the density
def euler_maruyama(samples: int, steps: int, quiet: bool = False) -> dict:
    """Section 1.32.2: dX_t = b_t(X_t) dt + sigma_t dW_t, simulated, against the density above.

    The variance of p_t = N(bt, t) equals the elapsed time, so the sample mean and the sample
    variance of the walk track bt and t. Both gaps shrink as --samples grows.
    """
    dt = 1.0 / max(steps, 1)  # The time increment of one Euler-Maruyama step.
    x_t = torch.zeros(max(samples, 1))  # Every path starts at the origin, so p_0 is a point mass.
    mean_gap, variance_gap = 0.0, 0.0  # The largest departure from bt and from t along the walk.
    for step in range(max(steps, 1)):
        noise = torch.randn_like(x_t)  # The Brownian increment dW_t, scaled by sqrt(dt); [samples].
        x_t = x_t + DRIFT * dt + NOISE * math.sqrt(dt) * noise  # One step of the chapter's SDE.
        elapsed = (step + 1) * dt  # The time reached after this step.
        mean_gap = max(mean_gap, abs(float(x_t.mean()) - DRIFT * elapsed))
        variance_gap = max(variance_gap, abs(float(x_t.var(unbiased=False)) - elapsed))
    show("sample mean at t = 1, exact 0.5", 0.5, float(x_t.mean()), quiet)
    show("sample variance at t = 1, exact 1", 1.0, float(x_t.var(unbiased=False)), quiet)
    show("largest |sample mean - b t|", 0.0, mean_gap, quiet)
    show("largest |sample variance - t|", 0.0, variance_gap, quiet)
    return {"mean": float(x_t.mean()), "variance": float(x_t.var(unbiased=False)),
            "mean_gap": mean_gap, "variance_gap": variance_gap, "steps": max(steps, 1)}


# %% 07. The report
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--grid", type=int, default=1000, help="points in the numerical checks")
    parser.add_argument("--steps", type=int, default=200, help="Euler-Maruyama steps over [0, 1]")
    parser.add_argument("--samples", type=int, default=20000, help="paths simulated along the SDE")
    parser.add_argument("--seed", type=int, default=0, help="random seed for the simulated paths")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    set_seed(args.seed)  # Seed first, so the walk repeats under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.31 and 1.32, worked example by worked example\n")
    uphill = scores(args.quiet)
    bowl = curvatures(args.quiet)
    spreading = heat(args.grid, args.quiet)
    combined = fokker_planck(args.grid, args.quiet)
    walk = euler_maruyama(args.samples, args.steps, args.quiet)

    report = {
        "s(x) at x = 3": float(uphill["standard"][0]),
        "s(5) under N(1, 4)": uphill["wide"],
        "Laplacian of x_1^2 + 3 x_2^2": bowl["bowl"],
        "Laplacian of the saddle": bowl["saddle"],
        "p_1(0)": spreading["center"],
        "d p_t(0) / d t at t = 1": spreading["time_rate"],
        "Laplacian of p_1 at 0": spreading["curvature"],
        "D times that Laplacian": spreading["flow"],
        "heat-equation residual": spreading["residual"],
        "Fokker-Planck residual": combined["both"],
        "residual at sigma_t = 0": combined["transport"],
        "Euler-Maruyama steps": walk["steps"],
        "sample variance at t = 1": walk["variance"],
        "largest |sample variance - t|": walk["variance_gap"],
    }
    print_report("Fields, Sections 1.31 and 1.32", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work every example of Sections 1.31 and 1.32 and print the report.