CIS6270 / lecture_3 /sde.py
pranamanam's picture
Upload 87 files
0ba9d09 verified
Raw History Blame Contribute Delete
13.3 kB
"""The forward corruption as a stochastic differential equation.
A deterministic flow is fixed by its velocity at every point, and a random path needs one more
piece, the perturbation applied over each small interval. A Brownian increment of an interval of
length h has size sqrt(h), and one Euler-Maruyama step adds that increment to the Euler step of
Chapter 2. A population started at one point spreads at a rate the noise scale sets, and a
restoring drift keeps that spread bounded. Ito's lemma and the Fokker-Planck equation then turn a
single drift and noise scale into an equation for the density. Every quantity here is a simulation
or a closed form.
In this file, as in Chapter 3, t = 0 holds a clean sample and corruption grows with t.
Chapter 3, Sections 3.2 through 3.10 of the notes. Prints the chapter's worked step from x = 2, the
Ornstein-Uhlenbeck mean and variance at t = log 2, the Ito check on E[W_t^2], and the
Fokker-Planck estimate of the Gaussian peak at t = 0.1.
python lecture_3/sde.py # 400 Euler-Maruyama steps, 16384 paths
python lecture_3/sde.py --steps 20 # a coarse grid, for a first look
"""
# %% 01. Imports and configuration
import math
import torch
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
# `--steps` is the number of Euler-Maruyama steps on the grid and `--samples` the number of paths,
# and together they set the whole cost of a run.
TERMINAL_TIME = 1.0 # T, the end of the simulated interval.
# %% 02. Drift and Brownian increments
#
# dX_t = f(X_t, t) dt + g(t) dW_t. The drift f fixes the average displacement over a short interval
# and the Brownian increment dW_t spreads the possible next states around that average. If we set
# g = 0 we recover the ordinary differential equation of Chapter 2, with f as the velocity field.
def brownian_increment(h: float, shape: tuple[int, ...], generator=None) -> torch.Tensor:
"""One Brownian increment over an interval of length h, drawn as sqrt(h) * epsilon."""
epsilon = torch.randn(shape, generator=generator) # [*shape] standard Gaussian draw.
return math.sqrt(h) * epsilon # Covariance h I, so each coordinate has variance h.
# %% 03. The size of a Brownian increment
#
# When we multiply a standard Gaussian by sqrt(h) its covariance scales by h, so a coordinate of
# the increment has standard deviation sqrt(h). If we halve the interval, the typical displacement
# drops by sqrt(2) = 1.414. The ratio sqrt(h)/h = 1/sqrt(h) grows without bound as the grid
# refines. That is why dW_t appears in the equation and never a velocity dW_t/dt.
def increment_scale(h: float, f_value: float = 1.0, g_value: float = 1.0) -> dict:
"""The standard deviation of one coordinate of the increment, and two ratios built from it."""
std = math.sqrt(h) # Typical displacement of the noise over an interval of length h.
return {
"std": std,
"std_over_h": std / h, # Grows without bound as the grid refines.
"noise_over_drift": g_value * std / (abs(f_value) * h), # g / (|f| sqrt(h)).
}
# %% 04. The Euler-Maruyama step
#
# x_{k+1} = x_k + f(x_k, t_k) h + g(t_k) sqrt(h) epsilon_k. The drift term is the Euler step of
# Chapter 2 unchanged. The noise term is the only addition, carrying the exact law of a Brownian
# increment over an interval of length h.
def euler_maruyama_step(x_k, f_value, g_value, h: float, epsilon):
"""One step of the scheme, with the Gaussian draw epsilon_k supplied by the caller."""
drift_part = f_value * h # Displacement the drift alone would produce.
noise_part = g_value * math.sqrt(h) * epsilon # Displacement from this draw.
return x_k + drift_part + noise_part, drift_part, noise_part
# %% 05. Fresh randomness and the spreading cloud
#
# Each step draws a new epsilon_k, so two runs from the same state give two trajectories. If we
# start from a common point x, the next state is Gaussian with mean x + f h and covariance
# g^2 h I, so the drift moves the center and the noise sets the width. Without a drift the
# variances of independent increments add, so after N steps the variance is N g^2 h = g^2 T and is
# unbounded.
def driftless_cloud(num_paths: int, steps: int, g_value: float, terminal: float, generator=None):
"""Simulate f = 0 on a grid of ``steps`` intervals and return the terminal states. [N]."""
h = terminal / steps # Step size, a time.
x = torch.zeros(num_paths) # [N] every path starts at the origin.
for _ in range(steps): # Independent increments accumulate additively in variance.
x = x + g_value * brownian_increment(h, (num_paths,), generator)
return x # [N] terminal states, with variance g^2 T in the limit of many paths.
# %% 06. A restoring drift
#
# dX_t = -kappa X_t dt + g dW_t. We multiply by the integrating factor e^{kappa t} to clear the
# drift. Under the Ito isometry the variance of the remaining stochastic integral is an ordinary
# integral, so X_t given x_0 is Gaussian with mean e^{-kappa t} x_0 and variance
# g^2 (1 - e^{-2 kappa t})/(2 kappa). The stationary variance g^2/(2 kappa) is the level at
# which the inward pull of the drift offsets the outward spread of the noise.
def ornstein_uhlenbeck_law(x_0: float, kappa: float, g_value: float, t: float) -> dict:
"""The exact conditional mean and variance of the Ornstein-Uhlenbeck process at time t."""
decay = math.exp(-kappa * t) # Fraction of the initial signal that remains.
stationary = g_value ** 2 / (2.0 * kappa) # Variance the process settles on.
variance = stationary * (1.0 - decay ** 2) # Rises from zero toward the stationary value.
return {"mean": decay * x_0, "variance": variance, "std": math.sqrt(variance),
"stationary_variance": stationary}
def simulate_ornstein_uhlenbeck(num_paths, steps, x_0, kappa, g_value, t, generator=None):
"""Euler-Maruyama on the same process, for comparison with the closed form. Returns [N]."""
h = t / steps
x = torch.full((num_paths,), float(x_0)) # [N] all paths start at x_0.
for _ in range(steps):
x = x - kappa * x * h + g_value * brownian_increment(h, (num_paths,), generator)
return x # [N] states at time t.
# %% 07. Ito's lemma
#
# d phi(X_t) = (phi_t + f phi_x + (1/2) g^2 phi_xx) dt + g phi_x dW_t. The squared increment
# (Delta X)^2 tends to g^2 Delta t and does not vanish, so the second-order Taylor term is as
# large as the first-order drift term and belongs in the dt part. The chapter's test function is
# phi(x) = x^2 on X_t = W_t, where phi_x = 2x and phi_xx = 2, giving d(W_t^2) = dt + 2 W_t dW_t.
def ito_drift(f_value, g_value, phi_x, phi_xx):
"""The dt bracket of Ito's lemma for a probe with no explicit time dependence."""
return f_value * phi_x + 0.5 * g_value ** 2 * phi_xx # Drift term plus curvature correction.
def check_ito_on_square(num_paths: int, steps: int, terminal: float, generator=None) -> dict:
"""Integrate d(W_t^2) = dt + 2 W_t dW_t on a grid and compare with the exact E[W_t^2] = t."""
h = terminal / steps
w = torch.zeros(num_paths) # [N] Brownian position.
accumulated = torch.zeros(num_paths) # [N] right-hand side of Ito's lemma, integrated.
without_curvature = torch.zeros(num_paths) # [N] the same sum with the dt term dropped.
for _ in range(steps):
dw = brownian_increment(h, (num_paths,), generator) # [N] one increment.
accumulated = accumulated + ito_drift(0.0, 1.0, 2.0 * w, 2.0) * h + 2.0 * w * dw
without_curvature = without_curvature + 2.0 * w * dw # Ordinary chain rule.
w = w + dw # Advance the path itself.
return {
"realized_mean_square": float((w ** 2).mean()), # Direct average of W_T^2.
"ito_mean": float(accumulated.mean()), # Same quantity through Ito's lemma.
"no_curvature_mean": float(without_curvature.mean()), # Near zero, the wrong answer.
"exact_mean_square": terminal, # E[W_T^2] = T.
}
# %% 08. The Fokker-Planck equation
#
# d_t p_t = -div(f p_t) + (1/2) g^2 laplacian(p_t). With f = 0 and g = 1 and a start at N(0, 1),
# the solution stays Gaussian with variance 1 + t. Both sides are then in closed form. Autograd
# supplies the two derivatives, so we can check the identity directly.
def spreading_gaussian_log_density(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""log N(x; 0, 1 + t), the exact solution of d_t p = (1/2) d_x^2 p from N(0, 1)."""
variance = 1.0 + t # Variance grows at the rate the diffusion injects it.
return -0.5 * torch.log(2.0 * math.pi * variance) - x ** 2 / (2.0 * variance)
def fokker_planck_residual(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
"""d_t p - (1/2) d_x^2 p on the closed-form solution, zero at every (x, t)."""
x = x.clone().requires_grad_(True) # [M] evaluation points.
t = t.clone().requires_grad_(True) # [M] evaluation times.
p = spreading_gaussian_log_density(x, t).exp() # [M] the density itself.
ones = torch.ones_like(p)
dp_dt, dp_dx = torch.autograd.grad(p, (t, x), ones, create_graph=True) # First derivatives.
d2p_dx2, = torch.autograd.grad(dp_dx, x, torch.ones_like(dp_dx)) # Second derivative in x.
return (dp_dt - 0.5 * d2p_dx2).detach() # The equation with f = 0 and g = 1.
def peak_extrapolation(step: float) -> dict:
"""The chapter's first-order estimate of the peak density after a short interval."""
p_0 = 1.0 / math.sqrt(2.0 * math.pi) # p_0(0), the height of the standard Gaussian.
curvature = -p_0 # p_0''(x) = (x^2 - 1) p_0(x), so p_0''(0) = -p_0(0).
rate = 0.5 * 1.0 * curvature # d_t p_t(0) at t = 0, with g = 1.
linear = p_0 + step * rate # One-step extrapolation from the rate.
exact = 1.0 / math.sqrt(2.0 * math.pi * (1.0 + step)) # Peak of N(0, 1 + step).
return {"p_0": p_0, "rate": rate, "linear": linear, "exact": exact, "gap": exact - linear}
# %% 09. The report
def main() -> None:
parser = common_parser(__doc__.split("\n")[0])
parser.set_defaults(steps=400, samples=16384)
args = parser.parse_args()
set_seed(args.seed)
device = pick_device(args.device) # Printed for consistency; this file runs on the CPU.
generator = torch.Generator().manual_seed(args.seed) # Local generator, so the report repeats.
# The chapter's step: x_k = 2, f = -1, g = 0.5, h = 0.04 and the draw epsilon_k = 1.5.
h = 0.04
landed, drift_part, noise_part = euler_maruyama_step(2.0, -1.0, 0.5, h, 1.5)
scale = increment_scale(h, f_value=-1.0, g_value=0.5)
three_draws = [euler_maruyama_step(2.0, -1.0, 0.5, h, e)[0] for e in (1.5, -0.5, 0.2)]
# The cloud without a drift: variance g^2 T after the whole interval.
terminal_states = driftless_cloud(args.samples, args.steps, 0.5, TERMINAL_TIME, generator)
# The Ornstein-Uhlenbeck example: x_0 = 2, kappa = 1, g = sqrt(2), evaluated at t = log 2.
ou_time = math.log(2.0)
ou = ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), ou_time)
ou_paths = simulate_ornstein_uhlenbeck(
args.samples, args.steps, 2.0, 1.0, math.sqrt(2.0), ou_time, generator,
)
ito = check_ito_on_square(args.samples, args.steps, 0.25, generator) # t = 0.25.
grid = torch.linspace(-3.0, 3.0, 61) # [61] evaluation points for the density equation.
residual = fokker_planck_residual(grid, torch.full_like(grid, 0.5)) # At t = 0.5.
peak = peak_extrapolation(0.1)
report = {
"euler_maruyama x_{k+1}": landed, # 2.11 in the notes.
"drift displacement f h": drift_part, # -0.04.
"noise displacement g sqrt(h) eps": noise_part, # 0.15.
"typical noise over drift": scale["noise_over_drift"], # 2.5 at this step size.
"increment std at h=0.04": scale["std"], # 0.2.
"three endpoints, mean": sum(three_draws) / 3.0, # 2.00 against the drift's 1.96.
"driftless terminal variance": float(terminal_states.var(unbiased=False)), # g^2 T = 0.25.
"OU mean at t=log 2": ou["mean"], # 1, half the initial signal.
"OU variance at t=log 2": ou["variance"], # 0.75 of the stationary 1.
"OU simulated mean": float(ou_paths.mean()),
"OU simulated variance": float(ou_paths.var(unbiased=False)),
"E[W_t^2] at t=0.25, exact": ito["exact_mean_square"], # 0.25.
"E[W_t^2] simulated directly": ito["realized_mean_square"],
"E[W_t^2] through Ito's lemma": ito["ito_mean"],
"E[W_t^2] without the curvature": ito["no_curvature_mean"], # Near zero.
"Fokker-Planck residual, largest": float(residual.abs().max()), # Zero up to float error.
"peak p_0(0)": peak["p_0"], # 0.3989.
"peak rate at t=0": peak["rate"], # -0.1995.
"peak linear estimate at t=0.1": peak["linear"], # 0.3790.
"peak exact at t=0.1": peak["exact"], # 0.3804.
"second-order gap": peak["gap"], # 0.0014.
}
if not args.quiet:
print(f"device={device} steps={args.steps} paths={args.samples} seed={args.seed}")
print_report("Forward corruption without a network", report)
save_run(args.out, config=vars(args), report=report)
# %% 10. Run the script
if __name__ == "__main__":
main()