Download lecture_3/sde.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 13.3 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_3/sde.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_3/sde.py
-
curl -L -o sde.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_3/sde.py
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() | |