File size: 8,087 Bytes
e456924 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """Independent implementation of RONOM, arXiv:2507.12814v2, Sec. 2 and SM3.
The projection solves weighted regularized least squares in a Gaussian basis;
it is not a learned Fourier operator. All geometric matrices are path independent.
"""
import math
from typing import List
import numpy as np
from scipy.special import erf
import torch
from torch import nn
def gaussian(x, centers, sigma):
return np.exp(-0.5 * ((x[:, None] - centers[None, :]) / sigma) ** 2)
def laplacian_gram(centers, sigma):
"""Closed-form integral of phi_i'' phi_j'' on [0,1], in float64.
Complete the square at the midpoint and integrate Gaussian moments 0,2,4.
The recurrence follows integration by parts, without numerical quadrature.
"""
mid = (centers[:, None] + centers[None, :]) / 2
dist = centers[:, None] - centers[None, :]
lo, hi = -mid, 1 - mid
ea, eb = np.exp(-(lo / sigma)**2), np.exp(-(hi / sigma)**2)
j0 = sigma * math.sqrt(math.pi) / 2 * (erf(hi / sigma) - erf(lo / sigma))
j2 = sigma**2 / 2 * (j0 - hi * eb + lo * ea)
j4 = sigma**2 / 2 * (3 * j2 - hi**3 * eb + lo**3 * ea)
value = np.exp(-dist**2 / (4 * sigma**2)) / sigma**8 * (
j4 - (dist**2 / 2 + 2 * sigma**2) * j2
+ (dist**2 / 4 - sigma**2)**2 * j0)
return (value + value.T) / 2
def projection_matrix(x_in, x_measure, centers, sigma, alpha, sobolev_reg, bc_reg):
basis = gaussian(x_in, centers, sigma)
# Trapezoidal integral weights; all input-generalization points are uniform.
w = np.empty(len(x_in), dtype=np.float64)
w[0], w[-1] = (x_in[1] - x_in[0]) / 2, (x_in[-1] - x_in[-2]) / 2
w[1:-1] = (x_in[2:] - x_in[:-2]) / 2
edge = gaussian(np.array([0., 1.]), centers, sigma)
# Official residual convention scales the derivative residual by sigma.
deriv = -(np.array([0., 1.])[:, None] - centers) / sigma * edge
residual = np.stack([edge[0] - edge[1], deriv[0] - deriv[1]])
rhs = basis.T * w
normal = (rhs @ basis + sobolev_reg * laplacian_gram(centers, sigma)
+ bc_reg * (residual.T @ residual) + alpha * np.eye(len(centers)))
return gaussian(x_measure, centers, sigma) @ np.linalg.solve(normal, rhs)
class LatentFlow(nn.Module):
"""Autonomous neural ODE with the explicit fourth order 3/8 tableau."""
def __init__(self, dimension, widths, step_size):
super().__init__()
layers = []
prev = dimension
for width in widths:
layers += [nn.Linear(prev, width), nn.GELU()]
prev = width
layers.append(nn.Linear(prev, dimension))
self.velocity = nn.Sequential(*layers)
self.step_size = float(step_size)
def forward(self, initial: torch.Tensor, times: torch.Tensor) -> torch.Tensor:
state = initial
states = torch.jit.annotate(List[torch.Tensor], [])
states.append(state)
for i in range(times.numel() - 1):
# Sampling times are 0,.1,...,1. Substeps also allow a coarser query.
interval = float(times[i + 1] - times[i])
count = max(1, int(math.ceil(interval / self.step_size - 1e-6)))
h = interval / count
for _ in range(count):
a = self.velocity(state)
b = self.velocity(state + (h / 3.) * a)
c = self.velocity(state + h * (b - a / 3.))
d = self.velocity(state + h * (a - b + c))
state = state + (h / 8.) * (a + 3. * b + 3. * c + d)
states.append(state)
return torch.stack(states, dim=1)
class RONOM(nn.Module):
def __init__(self, grid_size=33, features=(64, 128, 128, 256),
latent_channels=4, fourier_features=32, velocity_widths=(64, 64, 64),
padding_size=2, alpha=1e-5, sobolev_reg=1e-8, bc_reg=10.,
step_size=0.1):
super().__init__()
self.grid_size = int(grid_size)
self.latent_channels = int(latent_channels)
self.alpha, self.sobolev_reg, self.bc_reg = alpha, sobolev_reg, bc_reg
self.sigma = 1 / ((grid_size - 1) * math.sqrt(-2 * math.log(0.5)))
self.register_buffer('frequencies', torch.randn(fourier_features, 1))
f = list(features)
self.lift = nn.Sequential(nn.Conv1d(1 + 2 * fourier_features, 64, 3, padding=1),
nn.ReLU(), nn.Conv1d(64, f[0] // 2, 3, padding=1))
self.output_project = nn.Sequential(
nn.ConvTranspose1d(f[0] // 2, 64, 3, padding=1), nn.ReLU(),
nn.ConvTranspose1d(64, 1, 3, padding=1))
enc = []
prev = f[0] // 2
for width in f:
enc += [nn.Conv1d(prev, width, 3, padding=1), nn.ReLU()]
prev = width
self.encoder = nn.Sequential(*enc)
self.to_latent = nn.Conv1d(f[-1], latent_channels, 3, padding=1)
dec = []
back = list(reversed(f)) + [f[0] // 2]
for left, right in zip(back[:-1], back[1:]):
dec += [nn.ConvTranspose1d(left, right, 3, padding=1), nn.ReLU()]
self.decoder = nn.Sequential(*dec)
self.from_latent = nn.ConvTranspose1d(latent_channels, f[-1], 3, padding=1)
x = np.linspace(0., 1., grid_size)
centers = np.arange(-padding_size, grid_size + padding_size) / (grid_size - 1)
proj = projection_matrix(x, x, centers, self.sigma, alpha, sobolev_reg, bc_reg)
kernel = gaussian(x, x, self.sigma)
recovery = np.linalg.solve(kernel + alpha * np.eye(grid_size), kernel.T).T
self.register_buffer('grid', torch.tensor(x, dtype=torch.float32))
self.register_buffer('centers', torch.tensor(centers, dtype=torch.float64))
self.register_buffer('projection', torch.tensor(proj, dtype=torch.float32))
self.register_buffer('recovery', torch.tensor(recovery, dtype=torch.float32))
self.latent_flow = LatentFlow(grid_size * latent_channels, velocity_widths, step_size)
def encode(self, initial, x_input=None):
if initial.ndim == 3:
initial = initial.squeeze(-1)
if x_input is None:
projected = initial @ self.projection.T
else:
matrix = projection_matrix(x_input.detach().cpu().double().numpy(),
self.grid.detach().cpu().double().numpy(), self.centers.cpu().numpy(),
self.sigma, self.alpha, self.sobolev_reg, self.bc_reg)
projected = initial @ torch.as_tensor(matrix.T, device=initial.device, dtype=initial.dtype)
phase = 2 * math.pi * self.grid[:, None] @ self.frequencies.T
positional = torch.cat([phase.cos(), phase.sin()], dim=-1).T
lifted_input = torch.cat([projected[:, None, :],
positional[None].expand(initial.shape[0], -1, -1)], dim=1)
return self.to_latent(self.encoder(self.lift(lifted_input)))
def decode(self, latent, x_output=None):
b, nt = latent.shape[:2]
features = latent.reshape(b * nt, self.latent_channels, self.grid_size)
values = self.output_project(self.decoder(self.from_latent(features)))
values = values.reshape(b, nt, self.grid_size)
if x_output is None:
return values @ self.recovery.T
train_x = self.grid.detach().cpu().double().numpy()
target_x = x_output.detach().cpu().double().numpy()
kernel = gaussian(train_x, train_x, self.sigma)
matrix = np.linalg.solve(kernel + self.alpha * np.eye(self.grid_size),
gaussian(target_x, train_x, self.sigma).T).T
return values @ torch.as_tensor(matrix.T, device=values.device, dtype=values.dtype)
def forward(self, initial, times, return_latent=False):
# CPU time coordinates avoid a host/device synchronization at every RK stage.
z0 = self.encode(initial)
zt = self.latent_flow(z0.flatten(1), times)
zt = zt.reshape(initial.shape[0], times.numel(), self.latent_channels, self.grid_size)
prediction = self.decode(zt)
if return_latent:
return prediction, zt
return prediction
|