from __future__ import annotations import math from collections.abc import Sequence import torch from torch import nn from torch.nn import functional as F from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec class NeRFCFDMultiRes(nn.Module): def __init__( self, *, feature_names: Sequence[str], output_dim: int, coordinate_features: Sequence[str], encoding_levels: int, hidden_width: int, depth: int, condition_width: int, condition_depth: int, activation: str, coordinate_encoding_spec: CoordinateEncodingSpec | None = None, ) -> None: super().__init__() coordinate_indices = _indices(feature_names, coordinate_features) condition_indices = _complement_indices(len(feature_names), coordinate_indices) self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) spec = coordinate_encoding_spec or CoordinateEncodingSpec( type="nerf_multires", features=tuple(coordinate_features), levels=int(encoding_levels), scales=(), ) self.coordinate_encoder = CoordinateEncoder(spec, coordinate_dim=len(coordinate_indices)) encoded_dim = self.coordinate_encoder.output_dim condition_input_dim = len(condition_indices) if condition_indices else 1 self.condition_encoder = _mlp( input_dim=condition_input_dim, hidden_width=condition_width, output_dim=condition_width, depth=condition_depth, activation=activation, ) self.decoder = _mlp( input_dim=encoded_dim + condition_width, hidden_width=hidden_width, output_dim=output_dim, depth=depth, activation=activation, activate_output=False, ) def forward(self, features: torch.Tensor) -> torch.Tensor: coordinates = features.index_select(dim=1, index=self.coordinate_indices) condition = _gather_or_zeros(features, self.condition_indices) encoded = self.coordinate_encoder(coordinates) condition_embedding = self.condition_encoder(condition) return self.decoder(torch.cat((encoded, condition_embedding), dim=1)) class DeepONetBranchTrunk(nn.Module): def __init__( self, *, feature_names: Sequence[str], output_dim: int, coordinate_features: Sequence[str], fourier_scales: Sequence[float], hidden_width: int, depth: int, condition_width: int, condition_depth: int, activation: str, coordinate_encoding_spec: CoordinateEncodingSpec | None = None, ) -> None: super().__init__() coordinate_indices = _indices(feature_names, coordinate_features) condition_indices = _complement_indices(len(feature_names), coordinate_indices) self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) spec = coordinate_encoding_spec or CoordinateEncodingSpec( type="fixed_fourier", features=tuple(coordinate_features), scales=tuple(float(scale) for scale in fourier_scales), levels=0, ) self.coordinate_encoder = CoordinateEncoder(spec, coordinate_dim=len(coordinate_indices)) condition_input_dim = len(condition_indices) if condition_indices else 1 self.branch = _mlp( input_dim=condition_input_dim, hidden_width=condition_width, output_dim=hidden_width, depth=condition_depth, activation=activation, ) self.trunk = _mlp( input_dim=self.coordinate_encoder.output_dim, hidden_width=hidden_width, output_dim=hidden_width, depth=depth, activation=activation, ) self.head = nn.Linear(hidden_width, output_dim) def forward(self, features: torch.Tensor) -> torch.Tensor: coordinates = features.index_select(dim=1, index=self.coordinate_indices) condition = _gather_or_zeros(features, self.condition_indices) trunk = self.trunk(self.coordinate_encoder(coordinates)) branch = self.branch(condition) return self.head(trunk * branch) class PointContextPerceiver(nn.Module): def __init__( self, *, input_dim: int, output_dim: int, hidden_width: int, latent_width: int, context_points: int, attention_depth: int, activation: str, ) -> None: super().__init__() self.context_tokens = nn.Parameter(torch.empty(context_points, latent_width)) nn.init.normal_(self.context_tokens, std=latent_width ** -0.5) self.input_projection = nn.Linear(input_dim, latent_width) self.blocks = nn.ModuleList( [_PerceiverPointBlock(latent_width=latent_width, hidden_width=hidden_width, activation=activation) for _ in range(attention_depth)] ) self.head = _mlp( input_dim=latent_width, hidden_width=hidden_width, output_dim=output_dim, depth=max(1, attention_depth), activation=activation, activate_output=False, ) def forward(self, features: torch.Tensor) -> torch.Tensor: hidden = self.input_projection(features) tokens = self.context_tokens.to(dtype=hidden.dtype, device=hidden.device) for block in self.blocks: hidden = block(hidden, tokens) return self.head(hidden) class LocalPointTransformer(nn.Module): def __init__( self, *, feature_names: Sequence[str], output_dim: int, coordinate_features: Sequence[str], hidden_width: int, depth: int, neighbors: int, activation: str, ) -> None: super().__init__() coordinate_indices = _indices(feature_names, coordinate_features) self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) self.neighbors = int(neighbors) self.input_projection = nn.Linear(len(feature_names), hidden_width) self.blocks = nn.ModuleList([_LocalPointBlock(hidden_width=hidden_width, activation=activation) for _ in range(depth)]) self.head = nn.Linear(hidden_width, output_dim) def forward(self, features: torch.Tensor) -> torch.Tensor: coordinates = features.index_select(dim=1, index=self.coordinate_indices) hidden = self.input_projection(features) for block in self.blocks: hidden = block(hidden, coordinates, self.neighbors) return self.head(hidden) class RasterFNOUNet(nn.Module): def __init__( self, *, feature_names: Sequence[str], output_dim: int, coordinate_features: Sequence[str], grid_resolution: int, hidden_width: int, depth: int, condition_width: int, condition_depth: int, activation: str, ) -> None: super().__init__() coordinate_indices = _indices(feature_names, coordinate_features[:2]) condition_indices = _complement_indices(len(feature_names), coordinate_indices) self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) self.grid_resolution = int(grid_resolution) self.grid = nn.Parameter(torch.empty(self.grid_resolution, self.grid_resolution, hidden_width)) nn.init.normal_(self.grid, std=hidden_width ** -0.5) condition_input_dim = len(condition_indices) if condition_indices else 1 self.condition_encoder = _mlp( input_dim=condition_input_dim, hidden_width=condition_width, output_dim=condition_width, depth=condition_depth, activation=activation, ) self.decoder = _mlp( input_dim=hidden_width + condition_width, hidden_width=hidden_width, output_dim=output_dim, depth=depth, activation=activation, activate_output=False, ) def forward(self, features: torch.Tensor) -> torch.Tensor: xy = features.index_select(dim=1, index=self.coordinate_indices) sampled = _sample_grid(self.grid, xy) condition = self.condition_encoder(_gather_or_zeros(features, self.condition_indices)) return self.decoder(torch.cat((sampled, condition), dim=1)) class SirenConditionedINR(nn.Module): def __init__( self, *, feature_names: Sequence[str], output_dim: int, coordinate_features: Sequence[str], hidden_width: int, depth: int, condition_width: int, condition_depth: int, omega0: float, activation: str, ) -> None: super().__init__() coordinate_indices = _indices(feature_names, coordinate_features) condition_indices = _complement_indices(len(feature_names), coordinate_indices) self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) condition_input_dim = len(condition_indices) if condition_indices else 1 self.condition_encoder = _mlp( input_dim=condition_input_dim, hidden_width=condition_width, output_dim=condition_width, depth=condition_depth, activation=activation, ) layers: list[nn.Module] = [] input_dim = len(coordinate_indices) + condition_width for layer_index in range(depth): layers.append(_SineLayer(input_dim if layer_index == 0 else hidden_width, hidden_width, omega0=omega0, first=layer_index == 0)) self.net = nn.Sequential(*layers) self.head = nn.Linear(hidden_width, output_dim) def forward(self, features: torch.Tensor) -> torch.Tensor: coordinates = features.index_select(dim=1, index=self.coordinate_indices) condition = self.condition_encoder(_gather_or_zeros(features, self.condition_indices)) hidden = self.net(torch.cat((coordinates, condition), dim=1)) return self.head(hidden) class _PerceiverPointBlock(nn.Module): def __init__(self, *, latent_width: int, hidden_width: int, activation: str) -> None: super().__init__() self.norm = nn.LayerNorm(latent_width) self.ffn = _mlp( input_dim=latent_width, hidden_width=hidden_width, output_dim=latent_width, depth=2, activation=activation, activate_output=False, ) def forward(self, hidden: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor: scale = hidden.shape[1] ** -0.5 attention = torch.softmax(hidden @ tokens.T * scale, dim=1) context = attention @ tokens return hidden + self.ffn(self.norm(hidden + context)) class _LocalPointBlock(nn.Module): def __init__(self, *, hidden_width: int, activation: str) -> None: super().__init__() self.norm = nn.LayerNorm(hidden_width) self.update = _mlp( input_dim=hidden_width * 2, hidden_width=hidden_width, output_dim=hidden_width, depth=2, activation=activation, activate_output=False, ) def forward(self, hidden: torch.Tensor, coordinates: torch.Tensor, neighbors: int) -> torch.Tensor: if hidden.shape[0] <= 1 or neighbors <= 0: neighborhood = hidden else: k = min(neighbors + 1, hidden.shape[0]) distances = torch.cdist(coordinates.float(), coordinates.float()) indices = distances.topk(k=k, largest=False).indices[:, 1:] if k > 1 else distances.topk(k=k, largest=False).indices gathered = hidden.index_select(dim=0, index=indices.reshape(-1)).reshape(hidden.shape[0], -1, hidden.shape[1]) neighborhood = gathered.mean(dim=1) return hidden + self.update(torch.cat((self.norm(hidden), neighborhood), dim=1)) class _SineLayer(nn.Module): def __init__(self, input_dim: int, output_dim: int, *, omega0: float, first: bool) -> None: super().__init__() self.linear = nn.Linear(input_dim, output_dim) self.omega0 = float(omega0) with torch.no_grad(): bound = 1.0 / input_dim if first else math.sqrt(6.0 / input_dim) / self.omega0 self.linear.weight.uniform_(-bound, bound) def forward(self, values: torch.Tensor) -> torch.Tensor: return torch.sin(self.omega0 * self.linear(values)) def _indices(feature_names: Sequence[str], selected_names: Sequence[str]) -> tuple[int, ...]: indices: list[int] = [] for name in selected_names: try: indices.append(tuple(feature_names).index(name)) except ValueError as exc: raise ValueError(f"Coordinate feature {name!r} is not present in dataset features") from exc if not indices: raise ValueError("At least one coordinate feature is required") return tuple(indices) def _complement_indices(size: int, excluded: Sequence[int]) -> tuple[int, ...]: excluded_set = set(excluded) return tuple(index for index in range(size) if index not in excluded_set) def _gather_or_zeros(features: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: if indices.numel() == 0: return features.new_zeros((features.shape[0], 1)) return features.index_select(dim=1, index=indices) def _multires_encode(coordinates: torch.Tensor, levels: int) -> torch.Tensor: if levels <= 0: return coordinates pieces = [coordinates] for level in range(levels): scale = float(2**level) * math.pi pieces.append(torch.sin(coordinates * scale)) pieces.append(torch.cos(coordinates * scale)) return torch.cat(pieces, dim=1) def _fourier_features(coordinates: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: if scales.numel() == 0: return coordinates phases = coordinates.unsqueeze(-1) * scales.to(device=coordinates.device, dtype=coordinates.dtype) * math.pi return torch.cat((coordinates, torch.sin(phases).flatten(1), torch.cos(phases).flatten(1)), dim=1) def _sample_grid(grid: torch.Tensor, xy: torch.Tensor) -> torch.Tensor: resolution = grid.shape[0] if xy.shape[1] < 2: raise ValueError("Raster model requires at least x and y coordinate features") scaled = ((xy[:, :2].clamp(-1.0, 1.0) + 1.0) * 0.5) * float(resolution - 1) x = scaled[:, 0] y = scaled[:, 1] x0 = torch.floor(x).long().clamp(0, resolution - 1) y0 = torch.floor(y).long().clamp(0, resolution - 1) x1 = (x0 + 1).clamp(0, resolution - 1) y1 = (y0 + 1).clamp(0, resolution - 1) wx = (x - x0.to(x.dtype)).unsqueeze(1) wy = (y - y0.to(y.dtype)).unsqueeze(1) g00 = grid[y0, x0] g10 = grid[y0, x1] g01 = grid[y1, x0] g11 = grid[y1, x1] return (1 - wx) * (1 - wy) * g00 + wx * (1 - wy) * g10 + (1 - wx) * wy * g01 + wx * wy * g11 def _mlp( *, input_dim: int, hidden_width: int, output_dim: int, depth: int, activation: str, activate_output: bool = True, ) -> nn.Sequential: layers: list[nn.Module] = [] current_dim = input_dim for _ in range(max(depth - 1, 0)): layers.append(nn.Linear(current_dim, hidden_width)) layers.append(_activation(activation)) current_dim = hidden_width layers.append(nn.Linear(current_dim, output_dim)) if activate_output: layers.append(_activation(activation)) return nn.Sequential(*layers) def _activation(name: str) -> nn.Module: normalized = name.lower() if normalized == "gelu": return nn.GELU() if normalized == "relu": return nn.ReLU() if normalized == "silu": return nn.SiLU() if normalized == "tanh": return nn.Tanh() raise ValueError(f"Unsupported activation: {name}")