"""Paper-faithful FAST-DeepONet for the localized-inclusion Darcy experiment.""" from __future__ import annotations import math from typing import Any import numpy as np import torch from torch import Tensor, nn def compute_spectral_basis( train_inputs: np.ndarray, requested_rank: int, relative_tolerance: float = 1.0e-6, ) -> tuple[np.ndarray, np.ndarray, int, int]: """Compute the uncentered thin-SVD basis used by FAST-DeepONet. Returns the float32 basis, singular values, numerical rank and effective rank. The SVD itself is evaluated in float64 as in the experiment implementation. """ matrix = np.asarray(train_inputs) if matrix.ndim != 2: raise ValueError(f"train_inputs must be rank 2, got {matrix.shape}") if requested_rank < 1: raise ValueError("requested_rank must be positive") _, singular_values, right_vectors_t = np.linalg.svd( matrix.astype(np.float64, copy=False), full_matrices=False ) if singular_values.size == 0 or singular_values[0] <= 0: raise ValueError("input matrix has no positive leading singular value") numerical_rank = int(np.count_nonzero(singular_values > relative_tolerance * singular_values[0])) effective_rank = max(1, min(int(requested_rank), numerical_rank)) basis = right_vectors_t[:effective_rank].T.astype(np.float32, copy=True) return basis, singular_values, numerical_rank, effective_rank class _MLP(nn.Module): """Three-hidden-layer ReLU MLP with an explicit first layer.""" def __init__(self, input_dim: int, hidden_width: int, hidden_layers: int, output_dim: int) -> None: super().__init__() if hidden_layers < 1: raise ValueError("hidden_layers must be at least one") layers: list[nn.Linear] = [nn.Linear(input_dim, hidden_width)] layers.extend(nn.Linear(hidden_width, hidden_width) for _ in range(hidden_layers - 1)) self.hidden = nn.ModuleList(layers) self.output = nn.Linear(hidden_width, output_dim) @property def first_layer(self) -> nn.Linear: return self.hidden[0] def forward(self, inputs: Tensor) -> Tensor: hidden = inputs for layer in self.hidden: hidden = torch.relu(layer(hidden)) return self.output(hidden) class FASTDeepONet(nn.Module): """Factor-augmented DeepONet with a directional residual regularizer.""" def __init__( self, spectral_basis: Tensor | np.ndarray, branch_input_dim: int = 4096, requested_factor_rank: int = 16, residual_width: int = 16, factor_scale: float = 1.0, branch_width: int = 128, trunk_input_dim: int = 18, trunk_width: int = 128, latent_dim: int = 128, hidden_layers: int = 3, ) -> None: super().__init__() basis = torch.as_tensor(spectral_basis, dtype=torch.float32) if basis.ndim != 2 or basis.shape[0] != branch_input_dim: raise ValueError( f"spectral_basis must have shape ({branch_input_dim}, rank), got {tuple(basis.shape)}" ) effective_rank = int(basis.shape[1]) if effective_rank < 1 or effective_rank > requested_factor_rank: raise ValueError("effective factor rank must be in [1, requested_factor_rank]") self.branch_input_dim = int(branch_input_dim) self.requested_factor_rank = int(requested_factor_rank) self.effective_factor_rank = effective_rank self.residual_width = int(residual_width) self.factor_scale = float(factor_scale) self.branch_width = int(branch_width) self.trunk_input_dim = int(trunk_input_dim) self.trunk_width = int(trunk_width) self.latent_dim = int(latent_dim) self.hidden_layers = int(hidden_layers) self.register_buffer("spectral_basis", basis.contiguous()) self.residual_projection = nn.Linear(branch_input_dim, residual_width, bias=False) self.branch_net = _MLP( effective_rank + residual_width, branch_width, hidden_layers, latent_dim ) self.trunk_net = _MLP(trunk_input_dim, trunk_width, hidden_layers, latent_dim) self.output_bias = nn.Parameter(torch.zeros(())) def model_arguments(self) -> dict[str, Any]: return { "branch_input_dim": self.branch_input_dim, "requested_factor_rank": self.requested_factor_rank, "effective_factor_rank": self.effective_factor_rank, "residual_width": self.residual_width, "factor_scale": self.factor_scale, "branch_width": self.branch_width, "trunk_input_dim": self.trunk_input_dim, "trunk_width": self.trunk_width, "latent_dim": self.latent_dim, "hidden_layers": self.hidden_layers, } def branch_features(self, branch_inputs: Tensor) -> Tensor: if branch_inputs.shape[-1] != self.branch_input_dim: raise ValueError( f"branch input last dimension must be {self.branch_input_dim}, got {branch_inputs.shape[-1]}" ) coefficients = branch_inputs @ self.spectral_basis factors = coefficients * (self.factor_scale / math.sqrt(self.branch_input_dim)) residual = branch_inputs - coefficients @ self.spectral_basis.transpose(0, 1) residual_features = self.residual_projection(residual) return torch.cat((factors, residual_features), dim=-1) def encode_branch(self, branch_inputs: Tensor) -> Tensor: return self.branch_net(self.branch_features(branch_inputs)) def encode_trunk(self, trunk_inputs: Tensor) -> Tensor: if trunk_inputs.shape[-1] != self.trunk_input_dim: raise ValueError( f"trunk input last dimension must be {self.trunk_input_dim}, got {trunk_inputs.shape[-1]}" ) return self.trunk_net(trunk_inputs) def combine_latents(self, branch_latent: Tensor, trunk_latent: Tensor) -> Tensor: return torch.sum(branch_latent * trunk_latent, dim=-1) + self.output_bias def predict_from_branch_encoding(self, branch_latent: Tensor, trunk_inputs: Tensor) -> Tensor: return self.combine_latents(branch_latent, self.encode_trunk(trunk_inputs)) def forward(self, branch_inputs: Tensor, trunk_inputs: Tensor) -> Tensor: return self.combine_latents(self.encode_branch(branch_inputs), self.encode_trunk(trunk_inputs)) def effective_residual_map(self) -> Tensor: """Return H_z Theta P_res without materializing the p-by-p projector.""" q = self.spectral_basis theta = self.residual_projection.weight theta_projected = theta - (theta @ q) @ q.transpose(0, 1) first_weight = self.branch_net.first_layer.weight residual_columns = first_weight[:, self.effective_factor_rank :] return residual_columns @ theta_projected def directional_regularizer(self, tau: float | Tensor) -> Tensor: tau_tensor = torch.as_tensor(tau, dtype=self.output_bias.dtype, device=self.output_bias.device) if torch.any(tau_tensor <= 0): raise ValueError("tau must be positive") effective_map = self.effective_residual_map() norms = torch.linalg.vector_norm(effective_map, ord=2, dim=1, keepdim=True) denominators = torch.where(norms == 0, torch.ones_like(norms), norms) normalized = effective_map / denominators return torch.clamp(torch.abs(normalized) / tau_tensor, max=1.0).mean()