Score-Based Models & Flow Matching

"""Score matching, the probability-flow ODE, and flow matching side by side.

Three formulations that look nothing alike and are the same object:

  DENOISING SCORE MATCHING   regress the score of the noised distribution
  PROBABILITY FLOW ODE       a deterministic path with the same marginals
  FLOW MATCHING              regress a velocity along a chosen path

The point of putting them in one file is the conversion functions at the
bottom: score, noise and velocity are AFFINE TRANSFORMS of each other. Mixing
conventions between frameworks produces plausible nonsense with no error
anywhere, which is the most common failure when exchanging checkpoints.

Plain loops, no libraries.
"""

import math
import random

SEED = 31


def network(x, t, params):
    """Predicts a direction. Which direction depends on the convention.

    That ambiguity is not sloppiness - it is the actual situation. The same
    architecture is trained to predict a score, a noise vector or a velocity
    depending on the framework, and the caller has to know which.
    """
    embedding = math.sin(t * math.pi), math.cos(t * math.pi)

    hidden = []
    for i, row in enumerate(params['w1']):
        total = params['b1'][i] + params['w_time'][i] * embedding[i % 2]
        for weight, value in zip(row, x):
            total += weight * value
        hidden.append(math.tanh(total))

    return [
        bias + sum(w * h for w, h in zip(row, hidden))
        for row, bias in zip(params['w2'], params['b2'])
    ]


# --------------------------------------------------------------------------
# Denoising score matching
# --------------------------------------------------------------------------

def denoising_score_matching_loss(x0, sigma, params, rng):
    """Regress the score of the NOISED distribution, not of the data.

    The trick that makes this tractable: the score of the data distribution is
    unknown, but the score of the data-plus-known-noise distribution has a
    closed form - it is just the scaled noise. So an intractable objective
    becomes a regression against a target the loop constructed.

        score of q(x_t | x_0) = -(x_t - x_0) / sigma^2 = -eps / sigma
    """
    epsilon = [rng.gauss(0.0, 1.0) for _ in x0]
    x_t = [value + sigma * e for value, e in zip(x0, epsilon)]

    target = [-e / sigma for e in epsilon]
    predicted = network(x_t, sigma, params)

    return sum((t - p) ** 2 for t, p in zip(target, predicted)) / len(target)


# --------------------------------------------------------------------------
# Flow matching
# --------------------------------------------------------------------------

def flow_matching_loss(x0, x1, params, rng):
    """Regress a VELOCITY along a chosen path.

    With a straight-line interpolation the target velocity is simply the
    difference between the endpoints - a CONSTANT along the whole path. That
    is the entire reason flow matching samples in few steps: a constant
    velocity means one large step is exact, and in practice a handful suffice.

        x_t = (1 - t) * x0 + t * x1
        u_t = x1 - x0                 <- constant in t
    """
    t = rng.random()
    x_t = [(1.0 - t) * a + t * b for a, b in zip(x0, x1)]

    target = [b - a for a, b in zip(x0, x1)]
    predicted = network(x_t, t, params)

    return sum((u - p) ** 2 for u, p in zip(target, predicted)) / len(target)


# --------------------------------------------------------------------------
# Sampling: the same trained field, two different equations
# --------------------------------------------------------------------------

def sde_sample_step(x, t, step, params, rng, g):
    """Stochastic: inject noise at every step.

    Gives diversity and does NOT give a likelihood, because the trajectory is
    not a deterministic map.
    """
    score = network(x, t, params)
    drift = [-g * g * s for s in score]
    noise_scale = math.sqrt(abs(step)) * g

    return [
        value + step * d + noise_scale * rng.gauss(0.0, 1.0)
        for value, d in zip(x, drift)
    ]


def ode_sample_step(x, t, step, params, g):
    """Deterministic: the SAME marginals, with HALF the score coefficient.

    That factor of one half is the only difference between the two equations,
    and it is the whole result: the stochastic process and this deterministic
    ODE agree on their time-marginal distributions while sharing no
    trajectories at all.

    What the deterministic route buys: reproducibility, an exact likelihood,
    and a trajectory whose intermediate points can be SKIPPED - which is what
    turned sampling from a fixed procedure into a solver problem.
    """
    score = network(x, t, params)
    drift = [-0.5 * g * g * s for s in score]
    return [value + step * d for value, d in zip(x, drift)]


def euler_sample(dimension, params, steps, rng, g=1.0):
    """First-order integration. The simplest solver, and rarely the best.

    A higher-order solver buys quality at the SAME step count for free, which
    is worth trying before increasing the step count.
    """
    x = [rng.gauss(0.0, 1.0) for _ in range(dimension)]
    step = -1.0 / steps

    # Stop slightly short of zero: the true score DIVERGES there and the
    # network has seen little data that close, so integrating all the way
    # amplifies the worst-estimated part of the field.
    for index in range(steps):
        t = 1.0 - index / steps
        x = ode_sample_step(x, max(t, 1e-3), step, params, g)
    return x


def heun_sample(dimension, params, steps, rng, g=1.0):
    """Second-order: evaluate, step, re-evaluate, average the two slopes.

    Twice the evaluations per step and a much better step, so at a fixed
    evaluation budget it usually wins. That is the cheapest quality
    improvement available in this whole family.
    """
    x = [rng.gauss(0.0, 1.0) for _ in range(dimension)]
    step = -1.0 / steps

    for index in range(steps):
        t = max(1.0 - index / steps, 1e-3)
        first = network(x, t, params)
        predicted = [value + step * (-0.5 * g * g * s) for value, s in zip(x, first)]

        next_t = max(t + step, 1e-3)
        second = network(predicted, next_t, params)

        x = [
            value + step * (-0.5 * g * g) * 0.5 * (a + b)
            for value, a, b in zip(x, first, second)
        ]
    return x


# --------------------------------------------------------------------------
# The conversions. This is the part that causes real production failures.
# --------------------------------------------------------------------------

def noise_to_score(epsilon_prediction, sigma):
    """score = -eps / sigma."""
    return [-e / sigma for e in epsilon_prediction]


def score_to_noise(score_prediction, sigma):
    """eps = -sigma * score."""
    return [-sigma * s for s in score_prediction]


def velocity_to_score(velocity, x_t, t):
    """For a straight-line path, velocity and score are related through the
    interpolation coefficients.

    All three quantities are AFFINE TRANSFORMS of each other. A checkpoint
    predicting velocity, sampled by code expecting noise, produces output that
    looks like a badly trained model rather than like a bug - which is exactly
    why this is worth writing down rather than leaving implicit.
    """
    return [(v - value) / max(1.0 - t, 1e-6) for v, value in zip(velocity, x_t)]


def path_curvature(trajectory):
    """How far the sampled path deviates from a straight line.

    This is what determines how few steps are viable, and it is worth
    MEASURING rather than assuming: choosing a straight interpolant does not
    guarantee the learned field produces straight trajectories.
    """
    if len(trajectory) < 3:
        return 0.0

    start, end = trajectory[0], trajectory[-1]
    straight = [b - a for a, b in zip(start, end)]
    length = math.sqrt(sum(v * v for v in straight)) or 1.0

    total = 0.0
    for index in range(1, len(trajectory) - 1):
        fraction = index / (len(trajectory) - 1)
        expected = [a + fraction * d for a, d in zip(start, straight)]
        deviation = math.sqrt(
            sum((p - e) ** 2 for p, e in zip(trajectory[index], expected))
        )
        total += deviation

    return total / (len(trajectory) - 2) / length

"""The same formulations, with the parameterization made a type.

One change dominates. Score, noise and velocity are affine transforms of each
other, and every framework picks a different one — so a checkpoint predicting
velocity, sampled by code expecting noise, produces output that looks like a
badly trained model rather than like a bug. Making the convention part of the
type means that mismatch becomes an error rather than a mystery, and it is the
single highest-value piece of typing in this entry.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from enum import Enum
from typing import NamedTuple, Protocol, Sequence


class Parameterization(Enum):
    """What the network predicts. NOT interchangeable between frameworks.

    All three are affine transforms of each other, which is exactly what makes
    a mismatch so dangerous: the output is plausible rather than obviously
    wrong, so it reads as a poorly trained model.
    """

    SCORE = 'score'
    NOISE = 'noise'
    VELOCITY = 'velocity'


class ParameterizationMismatch(ValueError):
    """Raised when a sampler's convention differs from a checkpoint's.

    Its own type because the failure mode is specific and misleading: the
    sampler runs, produces output, and that output looks like a training
    problem rather than a conversion error.
    """


class DegeneratePath(ValueError):
    """Raised when an interpolant does not connect noise to data."""


class ShapeMismatch(ValueError):
    """Raised on a dimension violation instead of computing nonsense."""


class PathKind(Enum):
    """How noise and data are connected. The design variable that matters.

    LINEAR gives a constant target velocity, so in principle one step suffices
    and in practice a handful do. VARIANCE_PRESERVING is what DDPM uses and is
    curved, which is precisely why it needs many steps.
    """

    LINEAR = 'linear'
    VARIANCE_PRESERVING = 'vp'
    VARIANCE_EXPLODING = 've'


@dataclass(frozen=True)
class Interpolant:
    """The path between noise and data, and its exact target velocity.

    Frozen because the path is a property of the trained model: a sampler
    integrating a different path than the network was trained on will not stay
    on any trajectory at all.
    """

    kind: PathKind

    def position(self, x0: Sequence[float], x1: Sequence[float], t: float) -> list[float]:
        """x_t along the path. Closed form, which is what makes training
        parallel across t — no trajectory is ever simulated."""
        if self.kind is PathKind.LINEAR:
            return [(1.0 - t) * a + t * b for a, b in zip(x0, x1)]

        # Variance-preserving: a trigonometric path keeping unit variance.
        angle = t * math.pi / 2.0
        return [
            math.cos(angle) * a + math.sin(angle) * b for a, b in zip(x0, x1)
        ]

    def target_velocity(
        self, x0: Sequence[float], x1: Sequence[float], t: float
    ) -> list[float]:
        """The exact velocity along the path — the regression target.

        For the linear path this is CONSTANT in t, which is the entire reason
        flow matching samples in few steps: a constant velocity means one
        large step is exact.
        """
        if self.kind is PathKind.LINEAR:
            return [b - a for a, b in zip(x0, x1)]

        angle = t * math.pi / 2.0
        scale = math.pi / 2.0
        return [
            scale * (-math.sin(angle) * a + math.cos(angle) * b)
            for a, b in zip(x0, x1)
        ]

    @property
    def is_straight(self) -> bool:
        return self.kind is PathKind.LINEAR


class Prediction(NamedTuple):
    """A network output AND what convention it is in.

    Carrying the convention with the value is what turns a silent conversion
    error into a type error. The alternative — a bare list and a comment — is
    how checkpoints get sampled wrongly.
    """

    values: list[float]
    parameterization: Parameterization
    sigma: float
    t: float

    def as_score(self) -> list[float]:
        """Convert to the score convention. All three are affine transforms."""
        if self.parameterization is Parameterization.SCORE:
            return list(self.values)
        if self.parameterization is Parameterization.NOISE:
            return [-value / max(self.sigma, 1e-8) for value in self.values]
        # Velocity to score, for a straight path.
        return [value / max(1.0 - self.t, 1e-6) for value in self.values]

    def as_noise(self) -> list[float]:
        if self.parameterization is Parameterization.NOISE:
            return list(self.values)
        return [-self.sigma * s for s in self.as_score()]


class VelocityField(Protocol):
    """A trained model, with its conventions attached.

    A protocol rather than a class because the samplers below should work
    against any implementation — but the conventions are part of the contract,
    not an implementation detail.
    """

    @property
    def parameterization(self) -> Parameterization: ...

    @property
    def interpolant(self) -> Interpolant: ...

    def __call__(self, x: Sequence[float], t: float) -> list[float]: ...


@dataclass(frozen=True)
class SamplerConfig:
    """Solver, step count, and the convention it expects.

    The convention is here rather than assumed, because mismatching it is the
    most common way a working checkpoint produces nonsense.
    """

    steps: int
    order: int = 2
    expects: Parameterization = Parameterization.VELOCITY
    minimum_t: float = 1e-3

    def __post_init__(self) -> None:
        if self.steps < 1:
            raise ShapeMismatch('at least one sampling step is required')
        if self.order not in (1, 2):
            raise ShapeMismatch('only first- and second-order solvers are supported')
        # Guard clause for the quality loss nobody notices: the true score
        # DIVERGES at zero noise and the network has seen little data that
        # close, so integrating all the way amplifies the worst-estimated part
        # of the field.
        if self.minimum_t <= 0.0:
            raise ShapeMismatch(
                'integration must stop short of zero; the score diverges there and '
                'the estimate is at its worst'
            )

    def assert_compatible(self, field: VelocityField) -> None:
        """Guard clause for the failure that looks like something else."""
        if self.expects is not field.parameterization:
            raise ParameterizationMismatch(
                f'sampler expects {self.expects.value} but the model predicts '
                f'{field.parameterization.value}; these are affine transforms of each '
                'other, so the output will look like a training problem'
            )

    def timesteps(self) -> list[float]:
        """Descending from one, stopping short of zero."""
        return [
            max(1.0 - index / self.steps, self.minimum_t)
            for index in range(self.steps)
        ]


def flow_matching_loss(
    x0: Sequence[float],
    x1: Sequence[float],
    field: VelocityField,
    t: float,
) -> float:
    """Regress the velocity against a target the interpolant defines exactly.

    Nothing here is a bound or an estimate: the target is constructed, so the
    supervision is exact. That is what this family inherits from diffusion and
    why training is so much better behaved than anything adversarial.
    """
    if len(x0) != len(x1):
        raise ShapeMismatch('endpoints must share a dimension')

    x_t = field.interpolant.position(x0, x1, t)
    target = field.interpolant.target_velocity(x0, x1, t)
    predicted = field(x_t, t)

    return math.fsum((u - p) ** 2 for u, p in zip(target, predicted)) / len(target)


class SampleTrace(NamedTuple):
    """The sample, the trajectory, and the settings that produced it.

    The trajectory comes back because path curvature is worth measuring rather
    than assuming: choosing a straight interpolant does not guarantee the
    learned field produces straight trajectories, and curvature is what
    determines how few steps remain viable.
    """

    values: list[float]
    trajectory: list[list[float]]
    steps: int
    order: int

    def curvature(self) -> float:
        """Mean deviation from the straight line, relative to its length."""
        if len(self.trajectory) < 3:
            return 0.0

        start, end = self.trajectory[0], self.trajectory[-1]
        straight = [b - a for a, b in zip(start, end)]
        length = math.sqrt(math.fsum(v * v for v in straight)) or 1.0

        total = 0.0
        for index in range(1, len(self.trajectory) - 1):
            fraction = index / (len(self.trajectory) - 1)
            expected = [a + fraction * d for a, d in zip(start, straight)]
            total += math.sqrt(
                math.fsum((p - e) ** 2 for p, e in zip(self.trajectory[index], expected))
            )

        return total / (len(self.trajectory) - 2) / length


def sample(
    field: VelocityField,
    config: SamplerConfig,
    initial: Sequence[float],
) -> SampleTrace:
    """Integrate the velocity field. Second order by default, and for a reason.

    Heun costs two evaluations per step and takes a much better step, so at a
    FIXED evaluation budget it usually wins — which makes solver order the
    cheapest quality improvement available, and the one to exhaust before
    raising the step count.
    """
    config.assert_compatible(field)

    x = list(initial)
    trajectory = [list(x)]
    times = config.timesteps()
    step = -1.0 / config.steps

    for index, t in enumerate(times):
        first = field(x, t)

        if config.order == 1:
            x = [value + step * v for value, v in zip(x, first)]
        else:
            predicted = [value + step * v for value, v in zip(x, first)]
            next_t = times[index + 1] if index + 1 < len(times) else config.minimum_t
            second = field(predicted, next_t)
            x = [
                value + step * 0.5 * (a + b)
                for value, a, b in zip(x, first, second)
            ]

        trajectory.append(list(x))

    return SampleTrace(
        values=x, trajectory=trajectory, steps=config.steps, order=config.order
    )

"""Batched, with the divergence estimated rather than computed.

Two optimizations of different kinds:

  TRAINING batches exactly as DDPM does — the interpolant has a closed form,
  so every example takes its own random t and there is no rollout. Nothing
  clever is needed.

  LIKELIHOOD is where the real algorithmic change is. The probability-flow
  ODE gives an exact log-density, but it needs the DIVERGENCE of the velocity
  field along the trajectory - and computing that exactly costs one
  Jacobian-vector product per dimension, which is O(d) network evaluations per
  step. Hutchinson's estimator gets it in ONE, by noting that the expected
  value of v^T J v over random v with unit covariance IS the trace.

That substitution is what makes exact-likelihood evaluation affordable at all,
and it is why the resulting number has variance.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


def linear_interpolant(
    x0: NDArray[np.float32], x1: NDArray[np.float32], t: NDArray[np.float32]
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
    """Positions and target velocities for a whole batch at its own times.

    The target velocity for a straight path is the endpoint difference, which
    is CONSTANT in t — so it is computed once per pair rather than per
    timestep, and that constancy is the entire reason few steps work.
    """
    coefficient = t[:, None]
    velocity = x1 - x0

    # Fused: the position is built from the velocity already computed rather
    # than from both endpoints again.
    x_t = x0 + coefficient * velocity
    return x_t, velocity


def flow_matching_loss(
    field,
    x0: NDArray[np.float32],
    x1: NDArray[np.float32],
    rng: np.random.Generator,
) -> np.float32:
    """One batched step. Every example at its OWN random time.

    No rollout anywhere: the interpolant has a closed form, so training is
    embarrassingly parallel across examples and times alike.
    """
    t = rng.random(len(x0), dtype=FLOAT)
    x_t, target = linear_interpolant(x0, x1, t)

    residual = field(x_t, t)
    residual -= target
    # einsum contracts the squared norm in one pass rather than squaring into
    # a temporary and summing it.
    return np.einsum('ij,ij->', residual, residual, optimize=True) / residual.size


def hutchinson_divergence(
    field,
    x: NDArray[np.float32],
    t: NDArray[np.float32],
    rng: np.random.Generator,
    epsilon: float = 1e-3,
    probes: int = 1,
) -> NDArray[np.float32]:
    """Trace of the Jacobian, in O(1) network evaluations rather than O(d).

    The identity: for a random vector v with unit covariance,

        E[ v^T J v ] = trace(J)

    So one directional derivative per probe estimates the divergence,
    regardless of dimension. Computing it exactly would need one
    Jacobian-vector product PER DIMENSION, which at image scale is thousands
    of network evaluations per integration step - entirely impractical.

    The cost of the substitution is variance: the returned likelihood is an
    ESTIMATE, and calling it exact without averaging over probes overstates it.
    """
    base = field(x, t)
    total = np.zeros(len(x), dtype=FLOAT)

    for _ in range(probes):
        # Rademacher probes rather than Gaussian: same unit covariance, lower
        # estimator variance, and cheaper to generate.
        probe = rng.integers(0, 2, x.shape).astype(FLOAT)
        probe *= 2.0
        probe -= 1.0

        nudged = field(x + FLOAT(epsilon) * probe, t)
        nudged -= base
        nudged /= FLOAT(epsilon)
        total += np.einsum('ij,ij->i', probe, nudged, optimize=True)

    return total / probes


def ode_sample(
    field,
    shape: tuple[int, int],
    steps: int,
    rng: np.random.Generator,
    order: int = 2,
    minimum_t: float = 1e-3,
) -> NDArray[np.float32]:
    """Deterministic integration with Heun by default, and for a reason.

    Heun costs two evaluations per step and takes a much better step, so at a
    FIXED evaluation budget it usually wins - which makes solver order the
    cheapest quality improvement available, and the one to exhaust before
    raising the step count.

    Note what is NOT here: no batching over steps. The dependency is real, and
    the only lever is step count and solver order.
    """
    x = rng.standard_normal(shape, dtype=FLOAT)
    times = np.maximum(
        1.0 - np.arange(steps, dtype=FLOAT) / steps, FLOAT(minimum_t)
    )
    step = FLOAT(-1.0 / steps)

    for index in range(steps):
        t = np.full(shape[0], times[index], dtype=FLOAT)
        first = field(x, t)

        if order == 1:
            # In place: the state is updated through the prediction buffer.
            first *= step
            x += first
            continue

        predicted = x + step * first
        next_t = np.full(
            shape[0],
            times[index + 1] if index + 1 < steps else FLOAT(minimum_t),
            dtype=FLOAT,
        )
        second = field(predicted, next_t)

        # Fused average-and-step: neither slope is materialized separately.
        second += first
        second *= step * FLOAT(0.5)
        x += second

    return x


def exact_log_likelihood(
    field,
    x: NDArray[np.float32],
    steps: int,
    rng: np.random.Generator,
    probes: int = 4,
) -> NDArray[np.float32]:
    """Integrate the ODE BACKWARDS, accumulating the divergence.

        log p(x) = log p_prior(z) + integral of divergence along the path

    This is a real density rather than a bound, which is what the deterministic
    sampler buys over the stochastic one. It costs an ODE solve plus a
    divergence estimate per step, so it is orders of magnitude more expensive
    than sampling - and the divergence is estimated, so the result has
    variance that must be averaged away before it means anything.
    """
    state = x.copy()
    log_determinant = np.zeros(len(x), dtype=FLOAT)
    step = FLOAT(1.0 / steps)

    for index in range(steps):
        t = np.full(len(x), FLOAT(index) / steps, dtype=FLOAT)
        velocity = field(state, t)
        divergence = hutchinson_divergence(field, state, t, rng, probes=probes)

        state += step * velocity
        log_determinant += step * divergence

    # Standard Gaussian prior at the far end of the path.
    prior = -0.5 * np.einsum('ij,ij->i', state, state, optimize=True)
    prior -= 0.5 * x.shape[1] * np.log(2.0 * np.pi).astype(FLOAT)
    return prior + log_determinant


def path_curvature(trajectory: NDArray[np.float32]) -> NDArray[np.float32]:
    """Deviation from the straight line, per sample, as one vectorized pass.

    This is what determines how few steps are viable, and it is worth
    MEASURING rather than assuming: choosing a straight interpolant does not
    guarantee the learned field produces straight trajectories.

    trajectory is (steps, batch, dim).
    """
    start, end = trajectory[0], trajectory[-1]
    straight = end - start

    fractions = np.linspace(0.0, 1.0, len(trajectory), dtype=FLOAT)[1:-1, None, None]
    expected = start + fractions * straight

    deviation = trajectory[1:-1] - expected
    distances = np.sqrt(np.einsum('sbd,sbd->sb', deviation, deviation, optimize=True))
    lengths = np.sqrt(np.einsum('bd,bd->b', straight, straight, optimize=True))

    return distances.mean(axis=0) / np.maximum(lengths, 1e-8)