Blame for lecture_1/calculus.py

"""Jacobians, change of variables, divergence and the continuity equation.

A map moves points, and the Jacobian is the linear map it looks like near one of them. Its
determinant is the local volume factor, so the density at an image point is the density at its
source divided by that factor. The log-determinants add across a composition, and for a
many-to-one map the density is a sum of one corrected term per inverse branch. In continuous time
that factor becomes the divergence of a velocity field, and the continuity equation holds it in
balance with the change in density at every point.

Chapter 1, Sections 1.22 to 1.30 of the notes. Runs on the maps of Sections 1.22.2 and 1.23, the
squaring map of Section 1.26, and the expanding flow v(x) = x of Section 1.30.2. Every Jacobian
is taken by automatic differentiation and checked against the chapter's closed form.

    python lecture_1/calculus.py               # every worked example
    python lecture_1/calculus.py --grid 4000   # finer grids in the numerical checks
    python lecture_1/calculus.py --steps 2000  # a finer Euler walk along the worked flow
"""

# %% 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.


A_SCALE = torch.tensor([[2.0, 0.0], [0.0, 3.0]])  # The map of Section 1.23, a scaling of area by 6.


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 phi_mixed(x: torch.Tensor) -> torch.Tensor:
    """The map (x_1^2 + x_2, sin x_1 + 3 x_2) of Section 1.22.2, nonlinear in its first component."""
    return torch.stack([x[0] ** 2 + x[1], torch.sin(x[0]) + 3.0 * x[1]])


def phi_square(x: torch.Tensor) -> torch.Tensor:
    """The map (x_1^2, x_2) of Section 1.23, one-to-one once the domain is restricted to x_1 > 0."""
    return torch.stack([x[0] ** 2, x[1]])


def psi_scale(y: torch.Tensor) -> torch.Tensor:
    """The linear map y -> A y with A = diag(2, 3), composed after phi_square in Section 1.25."""
    return A_SCALE @ y


def f_bowl(x: torch.Tensor) -> torch.Tensor:
    """The scalar field f(x_1, x_2) = x_1^2 + 2 x_2^2 of Section 1.27.1."""
    return x[0] ** 2 + 2.0 * x[1] ** 2


def divergence(field, x: torch.Tensor) -> float:
    """Section 1.29: the divergence is the sum of d v_i / d x_i, the trace of J_v."""
    return float(torch.autograd.functional.jacobian(field, x).trace())


# %% 02. The Jacobian matrix and the local linear approximation
def jacobian(quiet: bool = False) -> dict:
    """Sections 1.22.2 and 1.22.3: J_phi at (1, 0), and the displacement it predicts."""
    point = torch.tensor([1.0, 0.0])  # The point at which the chapter evaluates the Jacobian.
    computed = torch.autograd.functional.jacobian(phi_mixed, point)  # [2, 2], by backpropagation.
    closed_form = torch.tensor([[2.0 * point[0], 1.0],  # The chapter's entries, 2 x_1 and 1.
                                [torch.cos(point[0]), 3.0]])  # cos x_1 and 3; [2, 2].
    delta = torch.tensor([0.01, -0.02])  # The small input displacement of Section 1.22.3.
    predicted = computed @ delta  # J_phi(x) Delta x, the first-order change in the output; [2].
    exact = phi_mixed(point + delta) - phi_mixed(point)  # The change the full map produces; [2].
    show("J_phi(1, 0)_11", 2.0, float(computed[0, 0]), quiet)
    show("J_phi(1, 0)_21 = cos(1)", 0.5403, float(computed[1, 0]), quiet)
    show("largest |autograd - closed form|", 0.0, float((computed - closed_form).abs().max()), quiet)
    show("predicted change, first output", 0.0, float(predicted[0]), quiet)
    show("predicted change, second output", -0.0546, float(predicted[1]), quiet)
    show("exact change, first output", 0.0001, float(exact[0]), quiet)
    show("exact change, second output", -0.05464, float(exact[1]), quiet)
    return {"jacobian": computed, "predicted": predicted, "exact": exact}


# %% 03. The Jacobian determinant as local volume change
def determinant(quiet: bool = False) -> dict:
    """Section 1.23: the absolute determinant is the factor by which a small region changes area."""
    corners = torch.tensor([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])  # The unit square.
    mapped = corners @ A_SCALE.T  # The image rectangle [0, 2] by [0, 3]; [4, 2].
    area = float((mapped[1, 0] - mapped[0, 0]) * (mapped[3, 1] - mapped[0, 1]))  # Width times height.
    collapsing = torch.tensor([[1.0, 0.0], [0.0, 0.0]])  # A map that flattens the plane onto a line.
    stretch = torch.autograd.functional.jacobian(phi_square, torch.tensor([2.0, 1.0]))  # At x_1 = 2.
    small = torch.tensor([0.25, 1.0])  # A point where the same map contracts, since 2 x_1 is below 1.
    squeeze = torch.autograd.functional.jacobian(phi_square, small)  # J_phi at x_1 = 0.25; [2, 2].
    show("det A for diag(2, 3)", 6.0, float(torch.linalg.det(A_SCALE)), quiet)
    show("area of the image rectangle", 6.0, area, quiet)
    show("det of the collapsing map", 0.0, float(torch.linalg.det(collapsing)), quiet)
    show("det J_phi at x_1 = 2", 4.0, float(torch.linalg.det(stretch)), quiet)
    show("det J_phi at x_1 = 0.25", 0.5, float(torch.linalg.det(squeeze)), quiet)
    return {"area": area, "stretch": float(torch.linalg.det(stretch))}


# %% 04. Change of variables in one and in two dimensions
def change_of_variables(grid: int, quiet: bool = False) -> dict:
    """Sections 1.24.1 and 1.24.2: p_Y(phi(x)) = p_X(x) / |det J_phi(x)| on the two worked maps."""
    inverse_slope = 0.5  # d phi^{-1}/dy for phi^{-1}(y) = (y - 1)/2, the reciprocal of phi' = 2.
    p_y = 1.0 * abs(inverse_slope)  # p_Y(y) = p_X(phi^{-1}(y)) |d phi^{-1}/dy|, so p_Y = 1/2.
    preimage = ((1.4 - 1.0) / 2.0, (2.0 - 1.0) / 2.0)  # The inverse image of B = [1.4, 2.0].
    mass = (preimage[1] - preimage[0]) * 1.0  # P(X in [0.2, 0.5]) under the uniform density.
    line = 1.0 + (torch.arange(grid) + 0.5) * (2.0 / grid)  # Midpoints of [1, 3]; [grid].
    normalization = float(torch.full_like(line, p_y).sum() * (2.0 / grid))  # Integral of p_Y there.
    flat = 1.0 / float(torch.linalg.det(A_SCALE))  # p_Y = 1/6 inside the image of the unit square.
    step = 1.0 / grid  # Midpoint spacing on the unit interval of the first input coordinate.
    x_one = (torch.arange(grid) + 0.5) * step  # Midpoints of (0, 1) in that coordinate; [grid].
    volume_factor = 2.0 * x_one  # |det J_phi| = 2 x_1 for phi = (x_1^2, x_2), at each point; [grid].
    output_density = 1.0 / (2.0 * (x_one ** 2).sqrt())  # p_Y(y) = 1/(2 sqrt(y_1)) at y_1 = x_1^2.
    residual = float((output_density * volume_factor - 1.0).abs().max())  # Should recover p_X = 1.
    transported = float((output_density * volume_factor).sum() * step)  # Total mass by substitution.
    show("preimage of [1.4, 2.0], lower end", 0.2, preimage[0], quiet)
    show("preimage of [1.4, 2.0], upper end", 0.5, preimage[1], quiet)
    show("P(X in the preimage)", 0.3, mass, quiet)
    show("p_Y on [1, 3] after Y = 2X + 1", 0.5, p_y, quiet)
    show("integral of p_Y over [1, 3]", 1.0, normalization, quiet)
    show("p_Y inside the image of A", 1.0 / 6.0, flat, quiet)
    show("density times area for A", 1.0, flat * 6.0, quiet)
    show("p_Y(y_1 = 1) for the x_1^2 map", 0.5, 1.0 / (2.0 * math.sqrt(1.0)), quiet)
    show("largest |p_Y |det J| - p_X|", 0.0, residual, quiet)
    show("total mass after the x_1^2 map", 1.0, transported, quiet)
    return {"p_y": p_y, "flat": flat, "residual": residual, "transported": transported}


# %% 05. Composed maps and the log-density form
def composition(quiet: bool = False) -> dict:
    """Section 1.25: log|det J| adds across a composition, so two factors become one sum."""
    point = torch.tensor([2.0, 1.0])  # The chapter evaluates the composition at x_1 = 2.
    inner = torch.autograd.functional.jacobian(phi_square, point)  # J_phi(x), det 2 x_1 = 4; [2, 2].
    outer = torch.autograd.functional.jacobian(psi_scale, phi_square(point))  # J_psi at phi(x); [2, 2].
    chained = outer @ inner  # The chain rule of Proposition 1.1, giving diag(8, 3); [2, 2].
    composed = torch.autograd.functional.jacobian(lambda x: psi_scale(phi_square(x)), point)
    log_inner = math.log(abs(float(torch.linalg.det(inner))))  # log 4 = 1.3863.
    log_outer = math.log(abs(float(torch.linalg.det(outer))))  # log 6 = 1.7918.
    log_density = math.log(1.0) - (log_inner + log_outer)  # log p_Y = log p_X - log|det J|.
    show("det J_psi", 6.0, float(torch.linalg.det(outer)), quiet)
    show("chained Jacobian, entry 11", 8.0, float(chained[0, 0]), quiet)
    show("det of the composition", 24.0, float(torch.linalg.det(composed)), quiet)
    show("log 4", 1.3863, log_inner, quiet)
    show("log 6", 1.7918, log_outer, quiet)
    show("log 4 + log 6 = log 24", 3.1781, log_inner + log_outer, quiet)
    show("log p_Y from a unit base density", -3.1781, log_density, quiet)
    return {"log_sum": log_inner + log_outer, "log_density": log_density}


# %% 06. Local invertibility and several inverse branches
def branches(grid: int, quiet: bool = False) -> dict:
    """Section 1.26: for Y = X^2 on a uniform X, the two preimages each contribute to p_Y(y)."""
    def branch_density(y: float) -> float:
        roots = (math.sqrt(y), -math.sqrt(y))  # The two solutions of x^2 = y inside the support.
        return sum(0.5 / abs(2.0 * root) for root in roots)  # Sum of p_X(x) / |phi'(x)| over roots.

    lower = 0.01  # We cut the integral at 0.01 because the integrand grows without bound at zero.
    step = (1.0 - lower) / grid  # Midpoint spacing on [0.01, 1].
    points = lower + (torch.arange(grid) + 0.5) * step  # Midpoints of the integration interval.
    partial = float((1.0 / (2.0 * points.sqrt())).sum() * step)  # Midpoint integral of 1/(2 sqrt y).
    show("p_Y(0.25), two branches at ±0.5", 1.0, branch_density(0.25), quiet)
    show("p_Y(0.04), two branches at ±0.2", 2.5, branch_density(0.04), quiet)
    show("integral of p_Y over [0.01, 1]", 1.0 - math.sqrt(lower), partial, quiet)
    return {"at_quarter": branch_density(0.25), "partial": partial}


# %% 07. Scalar fields, the gradient and the directional derivative
def gradient(quiet: bool = False) -> dict:
    """Sections 1.27.1 and 1.27.2: grad f at (1, 2), its steepest direction and a directional rate."""
    point = torch.tensor([1.0, 2.0])  # The point at which the chapter evaluates the gradient.
    grad = torch.autograd.functional.jacobian(f_bowl, point)  # (2 x_1, 4 x_2) = (2, 8); [2].
    steepest = grad / grad.norm()  # The unit vector of fastest increase; [2].
    along_x1 = float(grad @ torch.tensor([1.0, 0.0]))  # D_u f for u the first basis direction.
    along_x2 = float(grad @ torch.tensor([0.0, 1.0]))  # D_u f for u the second basis direction.
    show("grad f(1, 2), first entry", 2.0, float(grad[0]), quiet)
    show("grad f(1, 2), second entry", 8.0, float(grad[1]), quiet)
    show("|grad f(1, 2)| = sqrt(68)", math.sqrt(68.0), float(grad.norm()), quiet)
    show("steepest direction, first entry", 0.243, float(steepest[0]), quiet)
    show("steepest direction, second entry", 0.970, float(steepest[1]), quiet)
    show("D_u f along (1, 0)", 2.0, along_x1, quiet)
    show("D_u f along (0, 1)", 8.0, along_x2, quiet)
    return {"gradient": grad, "along_x1": along_x1}


# %% 08. Divergence
def divergences(quiet: bool = False) -> dict:
    """Section 1.29: the divergence of the chapter's three fields is -3, then 1, then 0."""
    point = torch.tensor([1.0, 1.0])  # The divergences below are the same at every point.
    contract = divergence(lambda x: torch.stack([-x[0], -2.0 * x[1]]), point)  # -1 + (-2) = -3.
    expand = divergence(lambda x: torch.stack([2.0 * x[0], -x[1]]), point)  # 2 - 1 = 1.
    rotate = divergence(lambda x: torch.stack([-x[1], x[0]]), point)  # 0 + 0 = 0.
    show("div of (-x_1, -2 x_2)", -3.0, contract, quiet)
    show("div of (2 x_1, -x_2)", 1.0, expand, quiet)
    show("div of (-x_2, x_1)", 0.0, rotate, quiet)
    show("area factor e^{div t} at t = 0.5", 1.649, math.exp(expand * 0.5), quiet)
    return {"contract": contract, "expand": expand, "rotate": rotate}


# %% 09. The continuity equation on the worked flow
def continuity(grid: int, steps: int, quiet: bool = False) -> dict:
    """Section 1.30.2: v(x) = x stretches a uniform source, with p_t(x) = e^{-t} on [0, e^t].

    Central differences give the continuity residual at interior points, and one Euler walk
    accumulates the log-density rate -div v of Section 1.30.1 along a trajectory.
    """
    def flux(time: float, position):
        return math.exp(-time) * position  # j_t(x) = p_t(x) v_t(x), with p_t = e^{-t} and v(x) = x.

    t, h = 0.5, 1e-4  # The time of the residual check, and the central-difference step.
    x = torch.linspace(0.1, 1.4, grid)  # Interior points of the support [0, e^{0.5}]; [grid].
    time_rate = (math.exp(-(t + h)) - math.exp(-(t - h))) / (2.0 * h)  # d p_t / d t at fixed x.
    flux_rate = (flux(t, x + h) - flux(t, x - h)) / (2.0 * h)  # d (p_t v) / d x; [grid].
    residual = float((time_rate + flux_rate).abs().max())  # The continuity equation, point by point.
    x_t, log_p, dt = 0.5, 0.0, 1.0 / max(steps, 1)  # One trajectory started inside the source.
    for _ in range(max(steps, 1)):
        log_p -= 1.0 * dt  # d/dt log p_t(x_t) = -div v = -1 for the field v(x) = x.
        x_t = x_t + x_t * dt  # Euler step on dx_t/dt = v(x_t) = x_t.
    show("largest continuity residual", 0.0, residual, quiet)
    show("log p along the trajectory at t = 1", -1.0, log_p, quiet)
    show("p_t(x_t) at t = 1", 0.3679, math.exp(log_p), quiet)
    show("x_t at t = 1, exact 0.5e", 0.5 * math.e, x_t, quiet)
    show("density at t = log 2", 0.5, math.exp(-math.log(2.0)), quiet)
    show("mass at t = log 2", 1.0, 2.0 * math.exp(-math.log(2.0)), quiet)
    show("mass at t = 1", 1.000, math.e * math.exp(-1.0), quiet)
    return {"residual": residual, "log_p": log_p, "x_t": x_t}


# %% 10. 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 steps along the worked flow")
    parser.add_argument("--seed", type=int, default=0, help="random seed, for a repeatable run")
    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 run repeats under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.22 to 1.30, worked example by worked example\n")
    local = jacobian(args.quiet)
    volume = determinant(args.quiet)
    moved = change_of_variables(args.grid, args.quiet)
    stacked = composition(args.quiet)
    folded = branches(args.grid, args.quiet)
    slope = gradient(args.quiet)
    fields = divergences(args.quiet)
    flow = continuity(args.grid, args.steps, args.quiet)

    report = {
        "J_phi(1, 0)_21 = cos(1)": float(local["jacobian"][1, 0]),
        "det J_phi at x_1 = 2": volume["stretch"],
        "p_Y on [1, 3] after Y = 2X + 1": moved["p_y"],
        "p_Y inside the image of A": moved["flat"],
        "mass after the x_1^2 map": moved["transported"],
        "log|det J| of the composition": stacked["log_sum"],
        "p_Y(0.25) summed over branches": folded["at_quarter"],
        "D_u f along (1, 0)": slope["along_x1"],
        "div of (2 x_1, -x_2)": fields["expand"],
        "largest continuity residual": flow["residual"],
        "log p along the flow at t = 1": flow["log_p"],
        "x_t at t = 1 after Euler steps": flow["x_t"],
    }
    print_report("Calculus, Sections 1.22 to 1.30", report)
    save_run(args.out, config=vars(args), report=report)


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