Delete lecture_2

lecture_2/README.md

@@ -1,99 +0,0 @@
-# Lecture 2: Flow Matching (Chapter 2)
-
-A velocity field assigns one vector to every state and every time. Following it from a starting
-point traces a trajectory, the trajectories from every starting point form the flow map, and the
-continuity equation ties the motion of the density to the same field. Training such a field by
-maximum likelihood costs an ODE solve and a divergence trace per gradient step, so flow matching
-replaces it with a regression: pair a Gaussian sample with a data sample, join the pair by a
-straight line, and regress a network onto the derivative of that line. Least squares recovers the
-average velocity of all the paths through a state, which is the field that transports the prior
-onto the data, and generation is Euler integration of the learned field.
-
-`flow_matching_mnist.py` is the lecture's main script, the one Section 2.20 of the notes walks
-through line by line. The other three files are the mathematics underneath it: the fields and the
-solver before any training, the full method on a target whose density is known exactly, and the
-family of interpolants the straight line belongs to.
-
-| File | Command | What it does |
-| --- | --- | --- |
-| `flow_matching_mnist.py` | `python lecture_2/flow_matching_mnist.py` | Trains the small U-Net on MNIST and writes a 4-by-4 grid of digits and a trajectory strip. `--dataset toy` runs offline. |
-| `velocity_fields.py` | `python lecture_2/velocity_fields.py` | No training. Trajectories, the flow map, the Euler refinement ladder, the pushforward of a Gaussian, the flux balance and the continuity equation. |
-| `flow_matching_2d.py` | `python lecture_2/flow_matching_2d.py` | Trains a velocity MLP on the four-Gaussian target and samples by Euler. Measures the conditional flow-matching equivalence against the exact marginal velocity. |
-| `interpolants.py` | `python lecture_2/interpolants.py` | Trains the same MLP three times, on the linear, the t-squared and the trigonometric schedule, and compares the samples at a fine and a coarse step budget. |
-
-Nothing is trained in `velocity_fields.py`, so its `--steps` sets the largest rung of the Euler
-refinement ladder. In the other three files `--steps` counts optimizer steps, and in
-`flow_matching_mnist.py` it is the budget for the whole run, which `--epochs` would otherwise set.
-
-## The numbers a default run produces
-
-`velocity_fields.py`, in about two seconds, reproduces the chapter's arithmetic exactly.
-
-```
-euler steps   1  landing 2.000000  error 0.718282  error x steps 0.7183
-euler steps   2  landing 2.250000  error 0.468282  error x steps 0.9366
-euler steps   4  landing 2.441406  error 0.276876  error x steps 1.1075
-flux_balance  -0.3354      enclosed_mass_rate  -0.3354
-continuity_time_derivative  0.1123      continuity_minus_flux_divergence  0.1123
-density_ratio_at_origin  2.0000         log_density_change  -0.6931
-```
-
-`flow_matching_2d.py`, in about six seconds on one core, with 83,202 parameters.
-
-```
-crossing_conditional_loss  [10.0, 9.0, 10.0]     crossing_marginal_loss  [1.0, 0.0, 1.0]
-crossing_constant_gap      [9.0, 9.0, 9.0]       crossing_target_variance  9.0000
-loss_cfm  10.3243   loss_fm  5.1196   loss_gap  5.2506   loss_cross_term  -0.0459
-gradient_cosine  0.9989                          final_loss  5.3169
-modes_found  4      mode_entropy  1.3657         nearest_mode_distance  0.4505
-mean_log_density  -2.2648                        data_mean_log_density  -2.1402
-```
-
-`interpolants.py`, in about thirteen seconds, trains three fields from the same initialization.
-
-```
-                 endpoint_error  covariance_deviation  mean_log_density  coarse_log_density
-linear                   0.0000                0.4999           -2.5276             -3.0309
-quadratic                0.0000                0.4939           -2.5748             -3.6419
-trigonometric            0.0000                0.0179           -2.6621             -2.9566
-```
-
-All three schedules reach the four modes. The trigonometric one is the only schedule whose
-interpolation keeps identity covariance at every time, and it is the one that loses least when the
-sampler drops from 64 Euler steps to 8, while the t-squared schedule, whose velocity grows linearly
-in t, loses most.
-
-`flow_matching_mnist.py` downloads MNIST on its first run. The offline path, `--dataset toy`,
-trains the same 471,265-parameter U-Net on 512 bars and blobs in about two minutes on one core.
-
-```
-epoch=20 squared_norm_loss=977.133
-final loss 977.133  sample range [-5.15, 3.50]
-```
-
-Twenty epochs over 512 images is 80 optimizer steps, and the generated range is still wider than
-the data's [-1, 1]. With `--epochs 100` the loss reaches 387.136 and the range [-3.25, 3.04], in
-about nine minutes. The MNIST default is 20 epochs over 60,000 images, and it is the one script
-here that does not finish while you watch it. `--steps 20` trains for twenty optimizer steps and
-then samples, which is the quick check that the file runs.
-
-## What is simplified
-
-- The marginal velocity `v_t(x) = E[U_t | X_t = x]` is written in closed form in
-  `flow_matching_2d.py`. That is available because the target is a Gaussian mixture and the prior is
-  Gaussian; Section 2.14 of the notes is about the case where it is not, which is why the equivalence
-  it proves is worth proving.
-- The equivalence is measured on one large batch, so the two losses agree up to that batch's own
-  sampling error. The file prints the cross term that accounts for the residual exactly.
-- The identity-covariance check uses two independent standard Gaussian endpoints, which is the
-  hypothesis the chapter states. The four-Gaussian target is centered but not white.
-- `velocity_fields.py` differentiates the densities with autograd where the chapter differentiates
-  by hand, so no derivative in that file is transcribed from the notes. It also runs in double
-  precision, because the Euler ladder it prints is arithmetic quoted to six decimal places.
-- The Gaussian translation of Section 2.17 is left out. Its coupling is deterministic, so the
-  conditional velocity is the constant 4 at every state and time and already equals the marginal
-  velocity.
-- The cubic Hermite interpolant of Section 2.18 is left out of `interpolants.py`. It needs endpoint
-  slopes, which the four-Gaussian data do not carry.
-- The three two-dimensional scripts run on the CPU. Each state there has two coordinates, so a GPU
-  would be slower. `flow_matching_mnist.py` picks CUDA, then Apple MPS, then the CPU.

lecture_2/flow_matching_2d.py

@@ -1,245 +0,0 @@
-"""Flow matching on four Gaussians, the whole method end to end.
-
-A coupling pairs a Gaussian prior sample with a data sample, a straight line joins the pair, and the
-derivative of that line is the regression target. Least squares against those per-pair velocities
-recovers the marginal velocity field, and Euler integration of the learned field turns fresh noise
-into data. The four-Gaussian target is small enough that the marginal velocity has a closed form, so
-the conditional flow-matching equivalence can be measured here and not only argued.
-
-Chapter 2, Sections 2.11 to 2.16 and 2.19 of the notes. Trained on
-``cis6270.data.gaussian_mixture_2d``. Prints the worked endpoint pair, the average at the crossing,
-the two losses with their constant gap, the agreement of their gradients, and the mode coverage and
-mean exact log density of the generated points.
-
-    python lecture_2/flow_matching_2d.py                # train and generate, about ten seconds
-    python lecture_2/flow_matching_2d.py --steps 20     # a first look, with no time to learn
-"""
-
-# %% 01. Imports and configuration
-import argparse
-
-import torch
-
-from cis6270.data import GAUSSIAN_CENTERS, GAUSSIAN_STD, gaussian_mixture_2d
-from cis6270.metrics import mean_log_density, mode_report
-from cis6270.nets import VelocityMLP, count_parameters
-from cis6270.runner import print_report, save_run, set_seed, train
-
-DIM = 2  # d, the width of the state.
-TRAIN_STEPS = 1000  # Optimizer steps, each on one fresh batch of endpoint pairs.
-BATCH_SIZE = 256  # Endpoint pairs per step.
-LEARNING_RATE = 3e-3  # Optimizer step size.
-WIDTH = 128  # Hidden width of the velocity network.
-SAMPLE_STEPS = 64  # Euler steps from t = 0 to t = 1, so Delta t = 1/64.
-NUM_SAMPLES = 1024  # Points generated for the mode coverage and the log density.
-DATA_POINTS = 4096  # Size of the training set drawn from the four-Gaussian target.
-EQUIVALENCE_POINTS = 16384  # Pairs used for the Monte Carlo check of Theorem 2.1.
-GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
-
-
-# %% 02. The coupling
-#
-# Section 2.11. A coupling is any joint law over the endpoint pair whose marginals are the prior and
-# the data. Drawing the two ends with no reference to each other gives the independent coupling.
-def sample_coupling(data, batch_size):
-    """Draw (X_0, X_1) ~ pi with pi(x_0, x_1) = p_0(x_0) p_1(x_1). Returns two [B, d] tensors."""
-    index = torch.randint(0, len(data), (batch_size,), device=data.device)  # [B] which data point.
-    x_1 = data[index]  # [B, d] the data endpoint, drawn from p_1.
-    x_0 = torch.randn_like(x_1)  # [B, d] the prior endpoint, drawn from p_0 = N(0, I).
-    return x_0, x_1
-
-
-# %% 03. The conditional path and its velocity
-#
-# Section 2.12. The straight segment at constant speed, and its derivative with the endpoints held
-# fixed. The velocity is called conditional because only someone who knows both endpoints has it.
-def conditional_path(x_0, x_1, t):
-    """X_t = (1 - t) X_0 + t X_1 and U_t = X_1 - X_0. Returns two [B, d] tensors."""
-    x_t = (1.0 - t)[:, None] * x_0 + t[:, None] * x_1  # [B, d] convex combination of the endpoints.
-    u_t = x_1 - x_0  # [B, d] the displacement, the same vector at every time.
-    return x_t, u_t
-
-
-def worked_pair():
-    """The chapter's one-dimensional pair x_0 = -2, x_1 = 3, evaluated on the grid of times."""
-    x_0 = torch.full((len(GRID), 1), -2.0)  # [5, 1] the same prior endpoint at every grid time.
-    x_1 = torch.full((len(GRID), 1), 3.0)  # [5, 1] the same data endpoint at every grid time.
-    x_t, u_t = conditional_path(x_0, x_1, GRID)  # The line x_t = -2 + 5t and its slope.
-    return {
-        "worked_pair_path": [round(float(v), 4) for v in x_t[:, 0]],
-        "worked_pair_velocity": float(u_t[0, 0]),  # The slope 5, so 1.25 units in every quarter.
-    }
-
-
-# %% 04. The marginal velocity
-#
-# Section 2.13. Several conditional paths pass through one state at one time, each with its own
-# U_t, and a deterministic field assigns one vector there, the average v_t(x) = E[U_t | X_t = x].
-# For this target the average is available exactly, and Section 2.14 is about the case where it
-# is not. Conditioning on component k makes (X_1, X_t) jointly Gaussian, so the posterior over
-# components is a softmax of the component log densities at X_t and v_t follows in closed form.
-def marginal_velocity(x_t, t):
-    """v_t(x) = E[U_t | X_t = x] for a standard Gaussian prior and the four-Gaussian target."""
-    centers = GAUSSIAN_CENTERS.to(x_t)  # [M, d] the four means of the mixture.
-    variance = (1.0 - t) ** 2 + (t * GAUSSIAN_STD) ** 2  # [B] variance of X_t given component k.
-    offset = x_t[:, None, :] - t[:, None, None] * centers  # [B, M, d] distance to each shifted mean.
-    log_weight = -0.5 * (offset ** 2).sum(-1) / variance[:, None]  # [B, M] up to a shared constant.
-    weight = torch.softmax(log_weight, dim=1)  # [B, M] posterior over components given X_t.
-    mean_center = (weight[:, :, None] * centers).sum(1)  # [B, d] posterior mean of the center.
-    # Writing E[X_1 | X_t] in terms of that posterior mean and cancelling the factor (1 - t) that
-    # U_t = (X_1 - X_t)/(1 - t) would otherwise divide by leaves one that stays finite at t = 1.
-    return ((1.0 - t)[:, None] * mean_center
-            - (1.0 - t - t * GAUSSIAN_STD ** 2)[:, None] * x_t) / variance[:, None]  # [B, d].
-
-
-def crossing_average():
-    """Two paths of Section 2.13 that meet at x = 2 at t = 0.5, with velocities 4 and -2."""
-    x_0 = torch.tensor([[0.0], [3.0]])  # [2, 1] the two prior endpoints.
-    x_1 = torch.tensor([[4.0], [1.0]])  # [2, 1] the two data endpoints.
-    x_t, u_t = conditional_path(x_0, x_1, torch.full((2,), 0.5))  # Both states land on x = 2.
-    return {
-        "crossing_state": float(x_t.mean()),  # Both paths occupy this state at t = 0.5.
-        "crossing_velocities": [round(float(v), 4) for v in u_t[:, 0]],
-        "crossing_average_velocity": float(u_t.mean()),  # The equally weighted average, 1.
-    }
-
-
-# %% 05. The two objectives
-#
-# Sections 2.14 and 2.15. The flow-matching loss regresses on the marginal velocity, which is not
-# available for real data, and the conditional loss regresses on the sampled U_t, which is.
-def objectives(model, x_0, x_1, t):
-    """L_CFM, L_FM, the gap E||U_t - v_t||^2 and the cross term, on one batch of endpoint pairs."""
-    x_t, u_t = conditional_path(x_0, x_1, t)  # [B, d] the state and its conditional velocity.
-    v_pred = model(x_t, t)  # [B, d] the network sees the state and the time, never the endpoints.
-    v_t = marginal_velocity(x_t, t)  # [B, d] the target of the loss we cannot evaluate in general.
-    loss_cfm = (v_pred - u_t).square().sum(1).mean()  # E||v_theta(X_t, t) - U_t||^2.
-    loss_fm = (v_pred - v_t).square().sum(1).mean()  # E||v_theta(X_t, t) - v_t(X_t)||^2.
-    gap = (u_t - v_t).square().sum(1).mean()  # The conditional variance of the target.
-    # Expanding the square leaves L_CFM = L_FM + cross + gap on any batch. The proof of Theorem 2.1
-    # rests on the cross term having conditional mean zero, so on a large batch it is near zero.
-    cross = 2.0 * ((v_pred - v_t) * (v_t - u_t)).sum(1).mean()
-    return loss_cfm, loss_fm, gap, cross
-
-
-# %% 06. The equivalence, measured
-#
-# Section 2.15. At the crossing the two losses are parabolas in the prediction a, nine units apart
-# at every a, with one derivative between them. The same identity holds in expectation on the
-# two-dimensional problem, up to the Monte Carlo error of the batch that estimates it.
-def crossing_gap(predictions=(0.0, 1.0, 2.0)):
-    """The conditional and marginal losses at x = 2, t = 0.5, where the targets are 4 and -2."""
-    targets = torch.tensor([4.0, -2.0])  # [2] the two conditional velocities meeting there.
-    marginal = targets.mean()  # The conditional mean of the target, which is 1.
-    conditional_values, marginal_values = [], []
-    for a in predictions:  # One candidate prediction at a time.
-        conditional_values.append(round(float((a - targets).square().mean()), 4))
-        marginal_values.append(round(float((a - marginal) ** 2), 4))
-    return {
-        "crossing_conditional_loss": conditional_values,
-        "crossing_marginal_loss": marginal_values,
-        "crossing_constant_gap": [round(c - m, 4) for c, m in
-                                  zip(conditional_values, marginal_values)],
-        "crossing_target_variance": float(targets.var(unbiased=False)),  # The same constant, 9.
-    }
-
-
-def equivalence_check(model, data, num_points):
-    """Theorem 2.1 on the four-Gaussian problem: the gap in value and the agreement in gradient."""
-    x_0, x_1 = sample_coupling(data, num_points)  # [N, d] pairs from the independent coupling.
-    t = torch.rand(num_points, device=x_1.device)  # [N] one uniform time per pair.
-    loss_cfm, loss_fm, gap, cross = objectives(model, x_0, x_1, t)
-    parameters = list(model.parameters())
-    grad_cfm = torch.autograd.grad(loss_cfm, parameters, retain_graph=True)  # Gradient of L_CFM.
-    grad_fm = torch.autograd.grad(loss_fm, parameters)  # Gradient of L_FM, on the same batch.
-    flat_cfm = torch.cat([g.flatten() for g in grad_cfm])  # [P] one vector per objective.
-    flat_fm = torch.cat([g.flatten() for g in grad_fm])  # [P].
-    return {
-        "loss_cfm": float(loss_cfm.detach()),
-        "loss_fm": float(loss_fm.detach()),
-        "loss_gap": float(gap.detach()),  # The constant that separates the two objectives.
-        "loss_cross_term": float(cross.detach()),  # Zero in expectation, so small on this batch.
-        "loss_fm_plus_gap": float((loss_fm + gap).detach()),  # L_CFM once the cross term is added.
-        "gradient_cosine": float(torch.nn.functional.cosine_similarity(flat_cfm, flat_fm, dim=0)),
-        # What separates the two gradients is the gradient of that cross term, which falls like one
-        # over the square root of the batch size.
-        "gradient_relative_difference": float((flat_cfm - flat_fm).norm() / flat_fm.norm()),
-    }
-
-
-# %% 07. Training
-#
-# Section 2.19. One iteration draws a pair and a time, builds the state in closed form, subtracts
-# the endpoints, and takes a gradient step. No ODE is solved while a training example is built.
-def make_loss(model, data, batch_size):
-    """The closure runner.train calls, returning L_CFM on one fresh batch."""
-    def loss_fn(step):
-        x_0, x_1 = sample_coupling(data, batch_size)  # [B, d] a fresh pair per step.
-        t = torch.rand(batch_size, device=x_1.device)  # [B] t ~ Uniform(0, 1), one per pair.
-        x_t, u_t = conditional_path(x_0, x_1, t)  # [B, d] the state and the regression target.
-        return (model(x_t, t) - u_t).square().sum(1).mean()  # ||v_theta(X_t, t) - U_t||^2.
-    return loss_fn
-
-
-# %% 08. Euler sampling
-#
-# Section 2.19. Generation starts from a fresh prior sample and integrates the learned field, so
-# one sample costs as many forward passes as there are steps.
-@torch.no_grad()
-def sample(model, num_samples, steps):
-    """Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
-    device = next(model.parameters()).device  # Use the model's own device.
-    x_t = torch.randn(num_samples, DIM, device=device)  # [N, d] the draw X_0 ~ p_0.
-    dt = 1.0 / steps  # Delta t, the width of one interval.
-    for step in range(steps):  # One network evaluation per step.
-        t = torch.full((num_samples,), step * dt, device=device)  # [N] the shared grid time.
-        x_t = x_t + dt * model(x_t, t)  # Euler: state += time increment times velocity.
-    return x_t
-
-
-# %% 09. The report
-def main():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
-    parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps")
-    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
-    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
-    parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
-    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
-    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="Euler steps")
-    parser.add_argument("--seed", type=int, default=0, help="random seed")
-    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 same seed gives the same report.
-    data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed)  # [N, d] the target p_1 = p_data.
-    model = VelocityMLP(dim=DIM, width=args.width)  # v_theta(x, t), the learned velocity field.
-
-    report = {"parameters": count_parameters(model)}
-    report.update(worked_pair())  # Section 2.12.
-    report.update(crossing_average())  # Section 2.13.
-    report.update(crossing_gap())  # Section 2.15.
-    # The gradient comparison runs before training, where the two gradients are far from zero.
-    report.update(equivalence_check(model, data, EQUIVALENCE_POINTS))  # Section 2.15.
-
-    losses = train(model, make_loss(model, data, args.batch_size),
-                   steps=args.steps, lr=args.lr, quiet=args.quiet)
-    generated = sample(model, args.samples, args.sample_steps)  # [N, d] fresh points.
-    coverage = mode_report(generated)  # Nearest mode, how many modes, how evenly they are visited.
-    report.update({
-        "final_loss": sum(losses[-100:]) / len(losses[-100:]),  # Averaged over the last steps.
-        "modes_found": coverage["modes_found"],
-        "mode_entropy": coverage["mode_entropy"],  # log 4 = 1.3863 when all four are equally used.
-        "nearest_mode_distance": coverage["nearest_distance"],
-        "mode_fractions": coverage["mode_fractions"],
-        "mean_log_density": mean_log_density(generated),  # Exact log density of the target.
-        "data_mean_log_density": mean_log_density(data),  # The same measurement on real points.
-    })
-
-    print_report("flow matching on four Gaussians", report)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses)
-
-
-# %% 10. Run the script
-if __name__ == "__main__":
-    main()

lecture_2/flow_matching_mnist.py

@@ -1,230 +0,0 @@
-"""Flow matching on MNIST, from noise to a handwritten digit in one script.
-
-We draw a Gaussian sample at time zero and a real image at time one, join them by a straight line,
-and train a U-Net to predict the velocity of that line at a random time in between. Generation then
-starts from fresh noise and takes Euler steps along the learned field. Chapter 2, Section 2.20 of
-the notes walks through this file section by section.
-
-    python lecture_2/flow_matching_mnist.py                  # 20 epochs on MNIST, downloads on first use
-    python lecture_2/flow_matching_mnist.py --steps 20       # 20 optimizer steps, for a first look
-    python lecture_2/flow_matching_mnist.py --epochs 1       # a single pass, about two minutes on a CPU
-    python lecture_2/flow_matching_mnist.py --dataset toy    # no download, for a first look or a test
-
-The network class and the checkpoint keys match the published `flow_unet_mnist.pt`, so the saved
-five-epoch run still loads.
-"""
-
-# %% 01. Imports
-import argparse
-from pathlib import Path
-
-import torch  # Tensors, gradients, and optimization.
-import torch.nn as nn  # Layers and model classes.
-import torch.nn.functional as F  # Resize feature maps in the decoder.
-
-
-# %% 02. Configuration
-BATCH_SIZE = 128  # Images per training batch.
-EPOCHS = 20  # Passes through the training set.
-LEARNING_RATE = 2e-4  # Optimizer step size.
-BASE_CHANNELS = 32  # U-Net widths: 32, 64, 128.
-SAMPLE_STEPS = 100  # Euler steps from t=0 to t=1.
-NUM_SAMPLES = 16  # Generate a 4-by-4 image grid.
-SEED = 7  # Seed for initialization and training.
-DATA_DIR = Path("data")  # MNIST download location.
-OUTPUT_DIR = Path("outputs/flow_matching_mnist")  # Checkpoint and generated image folder.
-
-DEVICE = torch.device(
-    "cuda" if torch.cuda.is_available()  # NVIDIA GPU.
-    else "mps" if torch.backends.mps.is_available()  # Apple GPU.
-    else "cpu"  # CPU fallback.
-)
-
-
-# %% 03. Load real images
-def make_loader(dataset="mnist", batch_size=BATCH_SIZE):
-    """Batches of images normalized to [-1, 1], the range the velocity field is trained on."""
-    from torch.utils.data import DataLoader, TensorDataset
-
-    if dataset == "toy":  # Shapes like MNIST, drawn offline, so a test needs no download.
-        from cis6270.data import toy_images
-
-        images = toy_images(num_images=512, size=28, seed=SEED)  # [512, 1, 28, 28].
-        labels = torch.zeros(len(images), dtype=torch.long)  # Unused; generation is unconditional.
-        return DataLoader(TensorDataset(images, labels), batch_size=batch_size, shuffle=True)
-
-    from torchvision import datasets
-    from torchvision.transforms import v2
-
-    transform = v2.Compose([
-        v2.ToImage(),  # Channel-first image: [1, 28, 28].
-        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 -> floats 0-1.
-        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Real data pixels: [0, 1] -> [-1, 1].
-    ])
-    mnist = datasets.MNIST(root=DATA_DIR, train=True, download=True, transform=transform)
-    return DataLoader(
-        mnist, batch_size=batch_size, shuffle=True,  # Batch shape: [B, 1, 28, 28].
-        num_workers=0, pin_memory=(DEVICE.type == "cuda"),  # Load in the main process.
-    )
-
-
-# %% 04. Sample the conditional path
-def sample_conditional_path(x1):  # x1 contains normalized real images.
-    x0 = torch.randn_like(x1)  # Independent Gaussian prior; same shape as x1.
-    t = torch.rand(x1.shape[0], device=x1.device)  # One random time per image; [B].
-    t_image = t[:, None, None, None]  # [B, 1, 1, 1]; share time across pixels.
-    xt = (1.0 - t_image) * x0 + t_image * x1  # Noise at t=0; real image at t=1.
-    ut = x1 - x0  # Derivative of the straight interpolation.
-    return xt, t, ut  # Network input, time, and target velocity.
-
-
-# %% 05. Define a simple convolution block
-def conv_block(in_channels, out_channels):
-    return nn.Sequential(
-        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),  # Keep H and W.
-        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),  # Refine features.
-    )
-
-
-# %% 06. Define the small U-Net
-class FlowUNet(nn.Module):
-    def __init__(self, base_channels=BASE_CHANNELS):
-        super().__init__()
-        c = base_channels  # Base width; c=32 by default.
-        self.encoder1 = conv_block(2, c)  # Image + time -> c channels.
-        self.encoder2 = conv_block(c, 2 * c)  # Encoder: c -> 2c.
-        self.middle = conv_block(2 * c, 4 * c)  # Bottleneck: 2c -> 4c.
-        self.decoder2 = conv_block(6 * c, 2 * c)  # Upsampled 4c + skip 2c -> 2c.
-        self.decoder1 = conv_block(3 * c, c)  # Upsampled 2c + skip c -> c.
-        self.output = nn.Conv2d(c, 1, 1)  # One signed velocity per pixel.
-        self.pool = nn.MaxPool2d(2)  # Halve the image height and width.
-
-    # %% 07. Compute the velocity field
-    def forward(self, xt, t):
-        t_image = t[:, None, None, None].expand_as(xt)  # Constant time channel.
-        x = torch.cat([xt, t_image], dim=1)  # [B, 2, 28, 28].
-        skip1 = self.encoder1(x)  # Keep full-resolution features.
-        skip2 = self.encoder2(self.pool(skip1))  # 28x28 -> 14x14.
-        x = self.middle(self.pool(skip2))  # 14x14 -> 7x7.
-        x = F.interpolate(x, size=skip2.shape[-2:], mode="nearest")  # 7x7 -> 14x14.
-        x = self.decoder2(torch.cat([x, skip2], dim=1))  # Restore encoder detail.
-        x = F.interpolate(x, size=skip1.shape[-2:], mode="nearest")  # 14x14 -> 28x28.
-        x = self.decoder1(torch.cat([x, skip1], dim=1))  # Restore fine detail.
-        v_pred = self.output(x)  # Predict v_theta(xt, t).
-        return v_pred  # [B, 1, 28, 28], matching xt.
-
-
-# %% 08. Train for one epoch
-def train_one_epoch(model, loader, optimizer, max_steps=0):
-    model.train()  # Select training mode.
-    total_loss, num_images = 0.0, 0  # Accumulate an image-weighted epoch loss.
-    for step, (x1, _) in enumerate(loader):  # Ignore digit labels: unconditional generation.
-        if max_steps and step >= max_steps:
-            break  # The step budget the caller left for this epoch is spent.
-        x1 = x1.to(DEVICE)  # Move real images to the model device.
-        xt, t, ut = sample_conditional_path(x1)  # Fresh noise and times for this batch.
-        v_pred = model(xt, t)  # The model receives only xt and t.
-        squared_error = (v_pred - ut).square()  # Velocity error at every pixel.
-        loss = squared_error.flatten(1).sum(1).mean()  # Sum pixels, average images.
-        optimizer.zero_grad(set_to_none=True)  # Clear old gradients.
-        loss.backward()  # Compute gradients with respect to weights.
-        optimizer.step()  # Update the U-Net weights.
-        total_loss += loss.item() * x1.shape[0]  # Weight by the actual batch size.
-        num_images += x1.shape[0]  # Include the smaller final batch.
-    return total_loss / num_images  # Average training loss per image.
-
-
-# %% 09. Generate images with Euler integration
-# Snapshots land every quarter of the default 100-step run, and the ODE states stay unclipped.
-@torch.no_grad()  # Sampling needs no gradient tracking.
-def sample_images(model, num_samples=NUM_SAMPLES, steps=SAMPLE_STEPS):
-    if num_samples < 1 or steps < 1:
-        raise ValueError("num_samples and steps must be positive")
-    model.eval()  # Select evaluation mode.
-    device = next(model.parameters()).device  # Use the model device.
-    xt = torch.randn(num_samples, 1, 28, 28, device=device)  # Initial x0.
-    dt = 1.0 / steps  # Positive time increment.
-    snapshots = [xt.cpu().clone()]  # Keep the initial noise for plotting.
-    for step in range(steps):  # Integrate forward in time.
-        t = torch.full((num_samples,), step * dt, device=device)  # Shared time grid.
-        v_pred = model(xt, t)  # Reevaluate velocity at the current state.
-        xt = xt + dt * v_pred  # Euler: state += time increment * velocity.
-        if (step + 1) % max(steps // 4, 1) == 0 or step + 1 == steps:
-            snapshots.append(xt.cpu().clone())  # Save intermediate and final states.
-    return xt, torch.stack(snapshots, dim=1)  # Raw images; [B, snapshots, 1, 28, 28].
-
-
-# %% 10. Set up the complete run
-def main():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
-    parser.add_argument("--epochs", type=int, default=EPOCHS)
-    parser.add_argument(
-        "--steps", type=int, default=0,
-        help="optimizer steps in the whole run; 0 runs every requested epoch in full",
-    )
-    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
-    parser.add_argument("--dataset", choices=["mnist", "toy"], default="mnist")
-    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS)
-    parser.add_argument("--samples", type=int, default=NUM_SAMPLES)
-    parser.add_argument("--seed", type=int, default=SEED)
-    parser.add_argument("--out", type=str, default=str(OUTPUT_DIR))
-    parser.add_argument("--quiet", action="store_true", help="print only the final report")
-    args = parser.parse_args()
-
-    torch.manual_seed(args.seed)  # Seed training randomness.
-    out = Path(args.out)
-    out.mkdir(parents=True, exist_ok=True)  # Create the output folder.
-    loader = make_loader(args.dataset, args.batch_size)  # Download, normalize, and batch.
-    model = FlowUNet().to(DEVICE)  # Create the velocity network.
-    optimizer = torch.optim.AdamW(
-        model.parameters(), lr=LEARNING_RATE, weight_decay=0.0,  # Optimize the flow loss.
-    )
-    num_parameters = sum(p.numel() for p in model.parameters())  # Count weights and biases.
-    if not args.quiet:
-        print(f"device={DEVICE} dataset={args.dataset} parameters={num_parameters:,}")
-
-    # %% 11. Train and save the learned parameters
-    losses = []  # Record the training loss after each epoch.
-    budget = args.steps  # Optimizer steps left in the run, counted across all epochs.
-    for epoch in range(1, args.epochs + 1):  # Complete the requested passes through the data.
-        loss = train_one_epoch(model, loader, optimizer, budget)  # Run the training loop.
-        losses.append(loss)  # Keep the loss history.
-        if not args.quiet:
-            print(f"epoch={epoch:02d} squared_norm_loss={loss:.3f}")
-        if args.steps:  # An epoch runs the smaller of the budget and the number of batches.
-            budget -= min(budget, len(loader))  # Optimizer steps this epoch used.
-            if budget <= 0:
-                break  # The run's budget is spent, so the remaining epochs do not run.
-    torch.save({
-        "model": model.state_dict(), "losses": losses,  # Weights and training curve.
-        "base_channels": BASE_CHANNELS,  # Width needed to reconstruct the model.
-        "seed": args.seed,  # Record the training seed.
-        "epochs_completed": len(losses),  # How far this run got.
-    }, out / "flow_unet_mnist.pt")
-
-    # %% 12. Sample and save images
-    torch.manual_seed(args.seed + 1)  # Fix the starting noise for generation.
-    samples, trajectory = sample_images(model, args.samples, args.sample_steps)
-    try:
-        from torchvision.utils import save_image  # Only needed to write the PNGs.
-
-        display_samples = ((samples.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # Clip for display only.
-        save_image(display_samples, out / "samples.png", nrow=4)  # Final images.
-        display_path = ((trajectory[:4] + 1.0) / 2.0).clamp(0.0, 1.0)  # First four sample paths.
-        save_image(
-            display_path.flatten(0, 1), out / "trajectory.png",  # One row per sample.
-            nrow=trajectory.shape[1],  # Time increases across columns.
-        )
-        if not args.quiet:
-            print(f"wrote {out}/samples.png and {out}/trajectory.png")
-    except ImportError:  # Without torchvision the run still trains and samples.
-        if not args.quiet:
-            print("torchvision is not installed, so no PNG was written")
-
-    print(f"final loss {losses[-1]:.3f}  sample range "
-          f"[{float(samples.min()):.2f}, {float(samples.max()):.2f}]")
-
-
-# %% 13. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Load data, train, sample, and save.

lecture_2/interpolants.py

@@ -1,210 +0,0 @@
-"""Three interpolants between the prior and the data.
-
-Replacing the weights 1 - t and t by any pair of schedules alpha_t and beta_t with alpha_0 = 1,
-beta_0 = 0, alpha_1 = 0 and beta_1 = 1 leaves the conditional flow-matching argument intact, and
-the regression target becomes the derivative alpha_dot X_0 + beta_dot X_1. The straight line, the
-t-squared schedule and the trigonometric schedule differ in where the path goes, how fast it gets
-there, and what the intermediate densities look like, so the same network trained on each produces
-different samples.
-
-Chapter 2, Section 2.18 of the notes. Trained on ``cis6270.data.gaussian_mixture_2d``. Prints the
-chapter's table of positions and velocities at t = 0.25, the endpoint check for each schedule, the
-identity covariance of the trigonometric schedule, and the samples each trained field produces at a
-fine and at a coarse step budget.
-
-    python lecture_2/interpolants.py                 # train three fields, about fifteen seconds
-    python lecture_2/interpolants.py --steps 20      # a first look, with no time to learn
-"""
-
-# %% 01. Imports and configuration
-import argparse
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.metrics import mean_log_density, mode_report
-from cis6270.nets import VelocityMLP
-from cis6270.runner import print_report, save_run, set_seed, train
-
-DIM = 2  # d, the width of the state.
-HALF_PI = math.pi / 2.0  # The constant in the trigonometric schedule and in its derivative.
-TRAIN_STEPS = 1000  # Optimizer steps for each of the three fields.
-BATCH_SIZE = 256  # Endpoint pairs per step.
-LEARNING_RATE = 3e-3  # Optimizer step size.
-WIDTH = 128  # Hidden width of the velocity network.
-SAMPLE_STEPS = 64  # Euler steps in the fine sampling budget.
-COARSE_STEPS = 8  # Euler steps in the coarse budget, where a varying velocity costs accuracy.
-NUM_SAMPLES = 1024  # Points generated per schedule and per budget.
-DATA_POINTS = 4096  # Size of the training set drawn from the four-Gaussian target.
-COVARIANCE_POINTS = 8192  # Endpoint pairs used for the empirical covariance of X_t.
-GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
-
-
-# %% 02. The three schedules
-#
-# Section 2.18. Each returns alpha_t, beta_t and their time derivatives, evaluated on a batch of
-# times. The pair of schedules fixes the path and the pair of derivatives fixes the target.
-def linear(t):
-    """alpha_t = 1 - t and beta_t = t, the straight segment traversed at constant speed."""
-    return 1.0 - t, t, -torch.ones_like(t), torch.ones_like(t)  # Four tensors of shape [B].
-
-
-def quadratic(t):
-    """alpha_t = 1 - t^2 and beta_t = t^2, the same segment entered slowly and left quickly."""
-    return 1.0 - t ** 2, t ** 2, -2.0 * t, 2.0 * t  # U_t = 2t (X_1 - X_0), zero speed at t = 0.
-
-
-def trigonometric(t):
-    """alpha_t = cos(pi t / 2) and beta_t = sin(pi t / 2), a quarter circle in the coefficients."""
-    angle = HALF_PI * t  # [B] the angle the coefficient pair has turned through.
-    return (torch.cos(angle), torch.sin(angle),
-            -HALF_PI * torch.sin(angle), HALF_PI * torch.cos(angle))
-
-
-SCHEDULES = {"linear": linear, "quadratic": quadratic, "trigonometric": trigonometric}
-
-
-def interpolate(schedule, x_0, x_1, t):
-    """X_t = alpha_t X_0 + beta_t X_1 and U_t = alpha_dot X_0 + beta_dot X_1. Two [B, d] tensors."""
-    alpha, beta, alpha_dot, beta_dot = schedule(t)  # [B] each, the two weights and their rates.
-    x_t = alpha[:, None] * x_0 + beta[:, None] * x_1  # [B, d] the state on the conditional path.
-    u_t = alpha_dot[:, None] * x_0 + beta_dot[:, None] * x_1  # [B, d] the regression target.
-    return x_t, u_t
-
-
-# %% 03. The endpoints and the velocity at one time
-#
-# Section 2.18. The two ends are correct when alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1,
-# and Table 2.1 reads the three paths at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
-def endpoint_and_table(name):
-    """The endpoint error of one schedule and its position and velocity at the chapter's input."""
-    x_0 = torch.tensor([[2.0, 0.0]])  # [1, d] the prior endpoint of the table.
-    x_1 = torch.tensor([[0.0, 2.0]])  # [1, d] the data endpoint of the table.
-    schedule = SCHEDULES[name]
-    ends, _ = interpolate(schedule, x_0.expand(2, DIM), x_1.expand(2, DIM),
-                          torch.tensor([0.0, 1.0]))  # [2, d] the states at t = 0 and t = 1.
-    start_error = (ends[0] - x_0[0]).abs().max()  # X_0 is returned at t = 0.
-    finish_error = (ends[1] - x_1[0]).abs().max()  # X_1 is returned at t = 1.
-    x_t, u_t = interpolate(schedule, x_0, x_1, torch.tensor([0.25]))  # One quarter of the interval.
-    return {
-        f"{name}_endpoint_error": float(torch.maximum(start_error, finish_error)),
-        f"{name}_position_at_quarter": [round(float(v), 4) + 0.0 for v in x_t[0]],
-        f"{name}_velocity_at_quarter": [round(float(v), 4) + 0.0 for v in u_t[0]],
-    }
-
-
-# %% 04. Identity covariance
-#
-# Section 2.18. For independent, centered endpoints with identity covariance the interpolation has
-# covariance (alpha_t^2 + beta_t^2) I. That factor is cos^2 + sin^2 = 1 at every time for the
-# trigonometric schedule, and below one between the ends for the other two.
-def covariance_check(name, num_points, seed):
-    """The coefficient squares on the grid, and the largest deviation of Cov(X_t) from the identity."""
-    set_seed(seed)  # Same draw for every schedule, so the comparison is between schedules alone.
-    x_0 = torch.randn(num_points, DIM)  # [N, d] a centered, white prior endpoint.
-    x_1 = torch.randn(num_points, DIM)  # [N, d] a centered, white data endpoint, independent of it.
-    identity = torch.eye(DIM)  # [d, d] the covariance the chapter's hypothesis gives each endpoint.
-    squares, deviations = [], []
-    for t in GRID:  # One time at a time, since the covariance is measured over the batch.
-        times = t.expand(num_points)  # [N] the same time for every pair.
-        alpha, beta, _, _ = SCHEDULES[name](times)
-        squares.append(round(float(alpha[0] ** 2 + beta[0] ** 2), 4))  # The exact variance factor.
-        x_t, _ = interpolate(SCHEDULES[name], x_0, x_1, times)  # [N, d] the interpolated states.
-        covariance = (x_t.T @ x_t) / num_points  # [d, d] the empirical second moment, centered at 0.
-        deviations.append(float((covariance - identity).abs().max()))
-    return {
-        f"{name}_coefficient_squares": squares,
-        f"{name}_covariance_deviation": max(deviations),
-    }
-
-
-# %% 05. Training one field per interpolant
-#
-# Section 2.18 and Section 2.19. The loop is the one of the straight-line case with the target
-# replaced, so any difference in the samples comes from the schedule and not from the setup.
-def train_one(name, data, args):
-    """Train v_theta against the conditional velocity of one schedule and return the loss history."""
-    set_seed(args.seed)  # The same initialization and the same batches for every schedule.
-    model = VelocityMLP(dim=DIM, width=args.width)  # v_theta(x, t).
-    schedule = SCHEDULES[name]
-
-    def loss_fn(step):
-        index = torch.randint(0, len(data), (args.batch_size,))  # [B] which data point to pair.
-        x_1 = data[index]  # [B, d] the data endpoint.
-        x_0 = torch.randn_like(x_1)  # [B, d] the prior endpoint, independent of it.
-        t = torch.rand(args.batch_size)  # [B] t ~ Uniform(0, 1).
-        x_t, u_t = interpolate(schedule, x_0, x_1, t)  # [B, d] the state and the target.
-        return (model(x_t, t) - u_t).square().sum(1).mean()  # ||v_theta(X_t, t) - U_t||^2.
-
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    return model, losses
-
-
-# %% 06. Euler sampling
-#
-# Section 2.19. Every schedule is sampled by the same integrator, so the step budget is the only
-# thing a reader has to hold fixed when comparing them.
-@torch.no_grad()
-def sample(model, num_samples, steps):
-    """Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
-    x_t = torch.randn(num_samples, DIM)  # [N, d] the draw X_0 ~ p_0.
-    dt = 1.0 / steps  # Delta t, the width of one interval.
-    for step in range(steps):  # One network evaluation per step.
-        x_t = x_t + dt * model(x_t, torch.full((num_samples,), step * dt))  # One Euler step.
-    return x_t
-
-
-def measure(model, args):
-    """Mode coverage and mean exact log density at the fine and the coarse step budget."""
-    set_seed(args.seed + 1)  # The same starting noise for every schedule.
-    fine = sample(model, args.samples, args.sample_steps)  # [N, d] points from the fine grid.
-    coarse = sample(model, args.samples, args.coarse_steps)  # [N, d] points from the coarse grid.
-    report = mode_report(fine)
-    return {
-        "modes_found": report["modes_found"],
-        "mode_entropy": report["mode_entropy"],
-        "nearest_distance": report["nearest_distance"],
-        "mean_log_density": mean_log_density(fine),
-        "coarse_log_density": mean_log_density(coarse),
-    }
-
-
-# %% 07. The report
-def main():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
-    parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps per field")
-    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
-    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
-    parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
-    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
-    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="fine Euler steps")
-    parser.add_argument("--coarse-steps", type=int, default=COARSE_STEPS, help="coarse Euler steps")
-    parser.add_argument("--seed", type=int, default=0, help="random seed")
-    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()
-
-    data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed)  # [N, d] the target p_1 = p_data.
-    report = {}
-    for name in SCHEDULES:  # Linear, quadratic, trigonometric, in the order the chapter gives them.
-        report.update(endpoint_and_table(name))  # Section 2.18, before any training.
-        report.update(covariance_check(name, COVARIANCE_POINTS, args.seed))
-        if not args.quiet:
-            print(f"training the {name} schedule")
-        model, losses = train_one(name, data, args)
-        report[f"{name}_final_loss"] = sum(losses[-100:]) / len(losses[-100:])
-        # The three losses are not comparable with each other, since each regresses on its own
-        # target, and the three measurements below are.
-        for key, value in measure(model, args).items():
-            report[f"{name}_{key}"] = value
-        save_run(f"{args.out}/{name}" if args.out else "", model=model, losses=losses)
-
-    report["data_mean_log_density"] = mean_log_density(data)  # The same measurement on real points.
-    print_report("three interpolants on four Gaussians", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-# %% 08. Run the script
-if __name__ == "__main__":
-    main()

lecture_2/velocity_fields.py

@@ -1,263 +0,0 @@
-"""Velocity fields, flow maps and the continuity equation.
-
-A velocity field assigns one vector to every state and every time, and following those vectors from
-a starting point traces a trajectory. The trajectories from every starting point together form the
-flow map, and the continuity equation relates the motion of the density to the same field. Nothing
-here is learned, so every number printed has a closed form beside it.
-
-Chapter 2, Sections 2.2 to 2.8 of the notes. Runs on the chapter's analytic fields in one and two
-dimensions. Prints the trajectories, the Euler error at one, two and four steps and the gap closing
-like 1/N, the pushforward of a Gaussian prior, the flux balance on [-1, 1], both sides of the
-continuity equation at x = 2, and the log density along one trajectory.
-
-    python lecture_2/velocity_fields.py                 # the full report
-    python lecture_2/velocity_fields.py --steps 20      # a shorter Euler refinement ladder
-"""
-
-# %% 01. Imports and configuration
-import argparse
-import math
-
-import torch
-
-from cis6270.runner import print_report, save_run, set_seed
-
-torch.set_default_dtype(torch.float64)  # Six printed digits of the Euler ladder need it.
-
-LOG2 = math.log(2.0)  # The rate a of the expansion field v_t(x) = a x, so that e^{at} = 2^t.
-GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
-START_POINTS = torch.tensor([-1.0, 0.0, 1.0])  # [3] the initial states followed through each field.
-EULER_STEPS = 64  # Largest step count in the refinement ladder, halved down to one step.
-NUM_SAMPLES = 4096  # Draws from the prior p_0 used for the pushforward moments.
-
-
-# %% 02. Three velocity fields and the trajectories they produce
-#
-# Section 2.2 and Section 2.3. Each field is a rule at a single instant, written as a function of
-# the state x and the time t, and each has an exact solution to compare a solver against.
-def constant_field(x, t):
-    """v_t(x) = 2, the field that translates every state at the same speed."""
-    return torch.full_like(x, 2.0)  # Same shape as x; the value depends on neither x nor t.
-
-
-def linear_field(x, t, a=LOG2):
-    """v_t(x) = a x, which expands away from the origin for a > 0 and contracts toward it for a < 0."""
-    return a * x  # Speed grows with distance from the origin and vanishes at x = 0.
-
-
-def rotation_field(x, t):
-    """v(x_1, x_2) = (-x_2, x_1), everywhere perpendicular to the position vector."""
-    return torch.stack([-x[..., 1], x[..., 0]], dim=-1)  # [..., 2] same shape as x.
-
-
-def exact_linear(x_0, t, a=LOG2):
-    """x_t = x_0 e^{at}, the solution of dx_t/dt = a x_t obtained by separating variables."""
-    return x_0 * torch.exp(a * t)  # Broadcasts a grid of times against a set of starting points.
-
-
-def trajectories():
-    """The chapter's worked trajectories for translation, expansion, contraction and rotation."""
-    speed = constant_field(torch.tensor(1.0), torch.tensor(0.0))  # The value 2, at any state.
-    translated = 1.0 + speed * GRID  # [5] the integral form x_t = x_0 + 2t from x_0 = 1.
-    expanded = exact_linear(START_POINTS, torch.tensor(1.0))  # [3] states at t = 1, each doubled.
-    contracted = exact_linear(START_POINTS, torch.tensor(1.0), a=-LOG2)  # [3] states each halved.
-    turns = torch.tensor([math.pi / 2, math.pi])  # [2] the quarter turn and the half turn.
-    circle = torch.stack([torch.cos(turns), torch.sin(turns)], dim=-1)  # [2, 2] x_t = (cos t, sin t).
-    # The squared norm is conserved because the field is orthogonal to the position: 2 x . v_t(x) = 0.
-    off_axis = torch.tensor([0.3, 0.7])  # [2] a state away from the axes, to test the conservation.
-    norm_rate = 2.0 * (off_axis * rotation_field(off_axis, torch.tensor(0.0))).sum()  # d||x||^2/dt.
-    return {
-        "constant_field_trajectory": [round(float(v), 4) for v in translated],
-        "expansion_x_half_from_two": float(exact_linear(torch.tensor(2.0), torch.tensor(0.5))),
-        "expansion_endpoints": [round(float(v), 4) + 0.0 for v in expanded],
-        "expansion_outer_spacing": float(expanded[2] - expanded[0]),
-        "contraction_endpoints": [round(float(v), 4) + 0.0 for v in contracted],
-        "contraction_outer_spacing": float(contracted[2] - contracted[0]),
-        "rotation_quarter_turn": [round(float(v), 4) + 0.0 for v in circle[0]],
-        "rotation_half_turn": [round(float(v), 4) + 0.0 for v in circle[1]],
-        "rotation_radius": float(circle.norm(dim=-1).mean()),
-        "rotation_squared_norm_rate": float(norm_rate),
-    }
-
-
-# %% 03. The flow map
-#
-# Section 2.4. Solving the ODE from every starting point at once defines phi_t, which is the identity
-# at t = 0 and, for the expansion field, multiplies the whole line by 2^t.
-def flow_map(x_0, t, a=LOG2):
-    """phi_t(x_0) = x_0 e^{at}, the accumulated transformation after following v_t for a time t."""
-    return exact_linear(x_0, t, a)  # [N] the state each starting point has reached by time t.
-
-
-def flow_map_report():
-    """The identity at t = 0, the Jacobian determinant at t = 1, and noncrossing trajectories."""
-    probe = torch.linspace(-3.0, 3.0, 101)  # [101] states spread over the line.
-    identity_error = (flow_map(probe, torch.tensor(0.0)) - probe).abs().max()  # phi_0 = id.
-    x_0 = torch.tensor(1.3, requires_grad=True)  # A state at which to differentiate phi_1.
-    jacobian = torch.autograd.grad(flow_map(x_0, torch.tensor(1.0)), x_0)[0]  # d phi_1 / d x_0 = 2.
-    # Proposition 2.1: distinct trajectories of a Lipschitz field never share a state at a shared time.
-    gaps = (flow_map(START_POINTS[2], GRID) - flow_map(START_POINTS[0], GRID)).abs()  # [5].
-    return {
-        "flow_map_identity_error": float(identity_error),
-        "flow_map_jacobian_det": float(jacobian),
-        "trajectory_gap_minimum": float(gaps.min()),
-    }
-
-
-# %% 04. Euler's method
-#
-# Section 2.5. Within one step the velocity stays fixed at its value at the start of the step, so
-# on dx_t/dt = x_t, where the velocity grows with the state, every step undershoots.
-def euler(field, x_0, steps):
-    """hat x_{k+1} = hat x_k + Delta t v_{t_k}(hat x_k) on a uniform grid of ``steps`` intervals."""
-    x_hat = x_0.clone()  # The running approximation, started from the exact initial state.
-    dt = 1.0 / steps  # Delta t, the width of one interval of the grid.
-    for k in range(steps):  # One velocity evaluation per step.
-        x_hat = x_hat + dt * field(x_hat, torch.tensor(k * dt))  # The Euler update at grid time t_k.
-    return x_hat
-
-
-def euler_ladder(largest):
-    """Euler on dx_t/dt = x_t from x_0 = 1, at every power of two up to ``largest`` steps."""
-    exact = math.e  # The exact solution x_1 = e^1, the chapter's 2.7183.
-    counts, landings, errors = [], [], []
-    steps = 1
-    largest = max(largest, 4)  # The chapter's three rungs, at one, two and four steps, always run.
-    while steps <= largest:  # 1, 2, 4, 8, ... up to the requested largest step count.
-        landing = float(euler(lambda x, t: x, torch.tensor(1.0), steps))  # The field is v_t(x) = x.
-        counts.append(steps)
-        landings.append(round(landing, 6))
-        errors.append(round(exact - landing, 6))
-        steps *= 2
-    return {
-        "euler_step_counts": counts,
-        "euler_landings": landings,
-        "euler_errors": errors,
-        # A first-order method has error proportional to 1/N, so the error times N settles at e/2.
-        "euler_error_times_steps": [round(e * n, 4) for e, n in zip(errors, counts)],
-        "euler_error_halving_ratios": [round(errors[i] / errors[i + 1], 4)
-                                       for i in range(len(errors) - 1)],
-        "euler_relative_error_four_steps": float((exact - landings[2]) / exact),
-    }
-
-
-# %% 05. A random starting point
-#
-# Section 2.6. The initial state becomes a random variable, each draw follows the same field, and
-# the law at time t is the pushforward p_t = (phi_t)_# p_0.
-def pushforward_moments(num_samples, seed):
-    """Moments of phi_1(X_0) for X_0 ~ N(0, 1) under the expansion field, against N(0, 4)."""
-    set_seed(seed)  # Same seed, same draw, same report.
-    x_0 = torch.randn(num_samples)  # [N] draws from the prior p_0 = N(0, 1).
-    x_1 = flow_map(x_0, torch.tensor(1.0))  # [N] states at t = 1, every one of them doubled.
-    return {
-        "pushforward_mean": float(x_1.mean()),  # Exact value 0, since the flow fixes the origin.
-        "pushforward_variance": float(x_1.var(unbiased=False)),  # Scaling by 2 gives variance 4.
-        "pushforward_exact_variance": 4.0,  # Scaling a centered law by c multiplies variance by c^2.
-    }
-
-
-# %% 06. Probability flux and the continuity equation
-#
-# Section 2.7. The flux j_t(x) = p_t(x) v_t(x) measures probability transported per unit time, and
-# the balance over an interval becomes the continuity equation once the interval shrinks to a point.
-def spreading_density(x, t):
-    """p_t(x) for X_t ~ N(0, 4^t), the law the expansion field produces from a standard normal."""
-    std = torch.pow(torch.tensor(2.0), t)  # 2^t, the standard deviation at time t.
-    return torch.exp(-x ** 2 / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))
-
-
-def flux_balance():
-    """The two endpoint fluxes on [-1, 1] at t = 0, and the rate of change of the enclosed mass."""
-    t_0 = torch.tensor(0.0)
-    endpoints = torch.tensor([-1.0, 1.0])  # [2] the interval's endpoints a and b.
-    density = spreading_density(endpoints, t_0)  # [2] p_0(-1) and p_0(1), both 0.2420.
-    flux = density * linear_field(endpoints, t_0)  # [2] j_0(x) = p_0(x) v_0(x).
-    # The same rate read off the distribution: M_t = erf(2^{-t} / sqrt 2), differentiated in t.
-    t = torch.tensor(0.0, requires_grad=True)
-    mass = torch.erf(torch.pow(torch.tensor(2.0), -t) / math.sqrt(2.0))  # P(|X_t| <= 1).
-    mass_rate = torch.autograd.grad(mass, t)[0]  # d M_t / dt at t = 0.
-    return {
-        "density_at_one": float(density[1]),
-        "flux_at_minus_one": float(flux[0]),
-        "flux_at_one": float(flux[1]),
-        "flux_balance": float(flux[0] - flux[1]),  # j_0(a) - j_0(b), the net rate into [-1, 1].
-        "enclosed_mass": float(mass.detach()),
-        "enclosed_mass_rate": float(mass_rate),
-    }
-
-
-def continuity_check():
-    """Both sides of the continuity equation at x = 2 and t = 0, in the expanded product-rule form."""
-    x = torch.tensor(2.0, requires_grad=True)  # The position held fixed while the density changes.
-    t = torch.tensor(0.0, requires_grad=True)  # The time at which both sides are evaluated.
-    p = spreading_density(x, t)  # p_0(2), the chapter's 0.05399.
-    dp_dt = torch.autograd.grad(p, t, create_graph=True)[0]  # Local time derivative of the density.
-    dp_dx = torch.autograd.grad(p, x, create_graph=True)[0]  # Spatial gradient of the density.
-    divergence = torch.autograd.grad(linear_field(x, t), x)[0]  # div v = log 2 in one dimension.
-    transport = linear_field(x, t) * dp_dx  # v_t(x) . grad p_t(x), the transport term.
-    volume = p * divergence  # p_t(x) div v_t(x), the local volume-change term.
-    return {
-        "density_at_two": float(p.detach()),
-        "density_gradient_at_two": float(dp_dx.detach()),
-        "velocity_divergence": float(divergence),
-        "continuity_time_derivative": float(dp_dt.detach()),
-        "continuity_minus_flux_divergence": float(-(transport + volume).detach()),
-    }
-
-
-# %% 07. The density along a moving trajectory
-#
-# Section 2.8. An observer riding with the sample sees no transport term, so the log density changes
-# at the rate -div v_t(x_t), and over a unit of time that agrees with the Jacobian of the flow map.
-def log_density_along_trajectory(x_start=0.7):
-    """The instantaneous change-of-variables equation on the expansion field, both ways."""
-    x_0 = torch.tensor(x_start)  # One starting point, carried from t = 0 to t = 1.
-    x_1 = flow_map(x_0, torch.tensor(1.0))  # The state the trajectory reaches at t = 1.
-    log_p0 = torch.log(spreading_density(x_0, torch.tensor(0.0)))  # log p_0(x_0).
-    log_p1 = torch.log(spreading_density(x_1, torch.tensor(1.0)))  # log p_1(x_1).
-    origin = torch.tensor(0.0)  # The fixed point of the flow, where phi_1(0) = 0.
-    prior_at_origin = float(spreading_density(origin, torch.tensor(0.0)))  # p_0(0) = 0.3989.
-    final_at_origin = float(spreading_density(origin, torch.tensor(1.0)))  # p_1(0) = 0.1995.
-    return {
-        "log_density_rate": -LOG2,  # d log p_t(x_t) / dt = -div v_t(x_t), a constant here.
-        "log_density_change": float(log_p1 - log_p0),  # Measured from the two densities.
-        "log_density_accumulated": -LOG2 * 1.0,  # The integral of -div v_s over s from 0 to 1.
-        "prior_density_at_origin": prior_at_origin,
-        "final_density_at_origin": final_at_origin,
-        "density_ratio_at_origin": prior_at_origin / final_at_origin,  # |det J_{phi_1}| = 2.
-    }
-
-
-# %% 08. The report
-def main():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
-    parser.add_argument("--steps", type=int, default=EULER_STEPS, help="largest Euler step count")
-    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="draws from the prior")
-    parser.add_argument("--seed", type=int, default=0, help="random seed")
-    parser.add_argument("--out", type=str, default="", help="directory for the run's report")
-    parser.add_argument("--quiet", action="store_true", help="print only the final report")
-    args = parser.parse_args()
-
-    report = {}
-    report.update(trajectories())  # Sections 2.2 and 2.3.
-    report.update(flow_map_report())  # Section 2.4.
-    report.update(euler_ladder(args.steps))  # Section 2.5.
-    report.update(pushforward_moments(args.samples, args.seed))  # Section 2.6.
-    report.update(flux_balance())  # Section 2.7.
-    report.update(continuity_check())  # Section 2.7.
-    report.update(log_density_along_trajectory())  # Section 2.8.
-
-    if not args.quiet:  # The Euler ladder, rung by rung, as the chapter develops it.
-        for n, landing, error in zip(report["euler_step_counts"], report["euler_landings"],
-                                     report["euler_errors"]):
-            print(f"  euler steps {n:3d}  landing {landing:.6f}  error {error:.6f}"
-                  f"  error x steps {error * n:.4f}")
-
-    print_report("velocity fields, flow maps and the continuity equation", report)
-    save_run(args.out, report=report)  # Writes report.json when --out names a directory.
-
-
-# %% 09. Run the script
-if __name__ == "__main__":
-    main()