model/FASTDeepONet.py

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