File size: 7,554 Bytes
b887635 | 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 167 168 | """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()
|