KPZ latent evolution block / kpz_experiment.py
Beats tuned baseline
1#!/usr/bin/env python3
2"""Numerical checks for a small periodic 3-D KPZ evolution cell.
3Run with: python3 kpz_experiment.py
4"""
5import json
6from pathlib import Path
7import numpy as np
8
9
10def derivatives(h, dx=1.0):
11 lap = np.zeros_like(h)
12 grad2 = np.zeros_like(h)
13 for ax in range(3):
14 hp = np.roll(h, -1, axis=ax)
15 hm = np.roll(h, 1, axis=ax)
16 lap += (hp - 2.0 * h + hm) / dx**2
17 grad = (hp - hm) / (2.0 * dx)
18 grad2 += grad * grad
19 return lap, grad2
20
21
22def kpz_step(h, nu, lam, D, dt, dx, rng, noise=True):
23 lap, grad2 = derivatives(h, dx)
24 z = rng.standard_normal(h.shape) if noise and D > 0 else 0.0
25 return h + dt * (nu * lap + 0.5 * lam * grad2) + np.sqrt(2 * D * dt / dx**3) * z
26
27
28def noise_variance_check():
29 # Prediction: Var(noise increment) = 2 D dt / dx^3.
30 rng = np.random.default_rng(11)
31 shape = (8, 8, 8)
32 D, dt = 0.37, 0.021
33 rows = []
34 for dx in (0.5, 1.0, 2.0):
35 vals = []
36 for _ in range(1000):
37 vals.append(kpz_step(np.zeros(shape), 0, 0, D, dt, dx, rng)[0, 0, 0])
38 observed = float(np.var(vals, ddof=1))
39 predicted = 2 * D * dt / dx**3
40 rows.append({"dx": dx, "predicted": predicted, "observed": observed,
41 "relative_error": abs(observed-predicted)/predicted})
42 return rows
43
44
45def nonlinear_drift_check():
46 # Prediction: for a fixed shape, mean height velocity is linear in lambda,
47 # with slope 0.5 * spatial mean(|grad h|^2); lambda=0 has zero nonlinear drift.
48 rng = np.random.default_rng(22)
49 shape = (16, 16, 16)
50 h = rng.normal(size=shape)
51 _, grad2 = derivatives(h)
52 q = float(grad2.mean())
53 rows = []
54 for lam in (-2.0, -1.0, 0.0, 1.0, 2.0):
55 out = kpz_step(h, 0, lam, 0, 0.01, 1, rng, noise=False)
56 observed_velocity = float((out.mean() - h.mean()) / 0.01)
57 predicted_velocity = 0.5 * lam * q
58 rows.append({"lambda": lam, "predicted_velocity": predicted_velocity,
59 "observed_velocity": observed_velocity,
60 "relative_error": abs(observed_velocity-predicted_velocity)/(abs(predicted_velocity)+1e-12)})
61 return {"mean_grad2": q, "rows": rows}
62
63
64def stability_check():
65 # For the 7-point Laplacian, explicit diffusion is stable when
66 # dt*nu*2*d/dx^2 <= 1 (d=3). A highest-frequency checkerboard mode
67 # should decay below the boundary and grow above it.
68 n = 12
69 x = np.indices((n, n, n)).sum(axis=0)
70 h0 = (-1.0) ** x
71 nu, dx = 1.0, 1.0
72 boundary = dx**2 / (2 * 3 * nu)
73 rows = []
74 for factor in (0.8, 1.0, 1.2):
75 dt = factor * boundary
76 h = h0.copy()
77 amplitudes = []
78 rng = np.random.default_rng(33)
79 for _ in range(8):
80 h = kpz_step(h, nu, 0, 0, dt, dx, rng, noise=False)
81 amplitudes.append(float(np.max(np.abs(h))))
82 # exact checkerboard amplification is |1-12 nu dt/dx^2|.
83 predicted_amp = abs(1 - 12 * nu * dt / dx**2)
84 observed_amp = amplitudes[0]
85 rows.append({"factor_of_bound": factor, "dt": dt,
86 "predicted_one_step_amplification": predicted_amp,
87 "observed_one_step_amplification": observed_amp,
88 "final_amplitude": amplitudes[-1],
89 "stable_observed": amplitudes[-1] <= 1.0 + 1e-9})
90 return {"predicted_dt_bound": boundary, "rows": rows}
91
92
93def roughening_control():
94 # Small same-setup rollout: KPZ has nonlinear mean growth whereas
95 # diffusion-only preserves the mean exactly (periodic boundaries).
96 rng_k, rng_d = np.random.default_rng(44), np.random.default_rng(44)
97 shape = (16, 16, 16)
98 h_k = np.zeros(shape); h_d = np.zeros(shape)
99 nu, lam, D, dt, dx = 0.15, 2.0, 0.03, 0.01, 1.0
100 kpz_means, diff_means, kpz_rms, diff_rms = [], [], [], []
101 for _ in range(100):
102 h_k = kpz_step(h_k, nu, lam, D, dt, dx, rng_k)
103 h_d = kpz_step(h_d, nu, 0, D, dt, dx, rng_d)
104 kpz_means.append(float(h_k.mean())); diff_means.append(float(h_d.mean()))
105 kpz_rms.append(float(h_k.std())); diff_rms.append(float(h_d.std()))
106 return {"kpz_final_mean": kpz_means[-1], "diffusion_final_mean": diff_means[-1],
107 "kpz_final_rms": kpz_rms[-1], "diffusion_final_rms": diff_rms[-1],
108 "kpz_mean_abs_change": abs(kpz_means[-1]),
109 "diffusion_mean_abs_change": abs(diff_means[-1])}
110
111
112def main():
113 result = {"noise_variance": noise_variance_check(),
114 "nonlinear_drift": nonlinear_drift_check(),
115 "stability": stability_check(),
116 "same_setup_control": roughening_control()}
117 Path("results.json").write_text(json.dumps(result, indent=2))
118 print(json.dumps(result, indent=2))
119
120
121if __name__ == "__main__":
122 main()
123
124try:
125 import torch
126 from torch import nn
127
128 class KPZCell(nn.Module):
129 """Differentiable periodic KPZ transition for [B,C,X,Y,Z] latent fields."""
130 def __init__(self, nu=0.1, lam=0.0, D=0.0, dt=0.01, dx=1.0,
131 stochastic=True):
132 super().__init__()
133 self.nu_raw = nn.Parameter(torch.tensor(float(np.log(np.expm1(nu)))))
134 self.D_raw = nn.Parameter(torch.tensor(float(np.log(np.expm1(D))))) if D > 0 else nn.Parameter(torch.tensor(-20.0))
135 self.lambda_ = nn.Parameter(torch.tensor(float(lam)))
136 self.dt, self.dx, self.stochastic = float(dt), float(dx), stochastic
137
138 def forward(self, h, noise=None):
139 lap = torch.zeros_like(h)
140 grad2 = torch.zeros_like(h)
141 for axis in (2, 3, 4):
142 hp = torch.roll(h, -1, dims=axis)
143 hm = torch.roll(h, 1, dims=axis)
144 lap = lap + (hp - 2 * h + hm) / self.dx ** 2
145 grad2 = grad2 + ((hp - hm) / (2 * self.dx)) ** 2
146 nu = torch.nn.functional.softplus(self.nu_raw)
147 D = torch.nn.functional.softplus(self.D_raw)
148 if noise is None:
149 noise = torch.randn_like(h) if self.training and self.stochastic else torch.zeros_like(h)
150 return h + self.dt * (nu * lap + 0.5 * self.lambda_ * grad2) + torch.sqrt(2 * D * self.dt / self.dx ** 3) * noise
151except ImportError:
152 KPZCell = None