airfRANS-model-exploration/src/airfrans_frontier/models/frontier.py

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