| """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) |
| |
| 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) |
| |
| 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): |
| |
| 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): |
| |
| 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 |
|
|
