models/DeepONet.py

"""Independent PyTorch implementation of the DeepONet in arXiv:1910.03193.

The paper is the architectural authority.  No implementation from the official
repository is imported or copied.  ReLU and Xavier-normal initialization are
configurable details used only because the paper leaves them unspecified.
"""

from __future__ import annotations

from typing import Any, Dict, Mapping

import torch
from torch import Tensor, nn


def _activation(name: str) -> nn.Module:
    choices = {
        "relu": nn.ReLU,
        "tanh": nn.Tanh,
        "gelu": nn.GELU,
        "silu": nn.SiLU,
    }
    try:
        return choices[name.lower()]()
    except KeyError as exc:
        raise ValueError(f"Unsupported activation {name!r}; choose {sorted(choices)}") from exc


class DenseNetwork(nn.Module):
    """A dense network where ``depth`` counts all Linear layers."""

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        depth: int,
        width: int,
        activation: str,
        *,
        activate_output: bool,
        dense_bias: bool = True,
        output_bias: bool = True,
    ) -> None:
        super().__init__()
        if depth < 1:
            raise ValueError("depth must be at least one")
        if min(input_dim, output_dim, width) < 1:
            raise ValueError("input_dim, output_dim and width must be positive")

        layers = []
        current_dim = input_dim
        for layer_index in range(depth):
            is_output = layer_index == depth - 1
            next_dim = output_dim if is_output else width
            layers.append(
                nn.Linear(
                    current_dim,
                    next_dim,
                    bias=output_bias if is_output else dense_bias,
                )
            )
            if not is_output or activate_output:
                layers.append(_activation(activation))
            current_dim = next_dim
        self.layers = nn.Sequential(*layers)

    def forward(self, inputs: Tensor) -> Tensor:
        return self.layers(inputs)


class DeepONet(nn.Module):
    """Stacked or unstacked DeepONet with the paper's branch/trunk fusion."""

    def __init__(
        self,
        branch_input_dim: int,
        trunk_input_dim: int,
        latent_dim: int,
        *,
        branch_depth: int = 2,
        trunk_depth: int = 3,
        width: int = 40,
        activation: str = "relu",
        stacked: bool = False,
        dense_bias: bool = True,
        branch_output_bias: bool = True,
        global_bias: bool = True,
        initializer: str = "xavier_normal",
    ) -> None:
        super().__init__()
        self.branch_input_dim = int(branch_input_dim)
        self.trunk_input_dim = int(trunk_input_dim)
        self.latent_dim = int(latent_dim)
        self.stacked = bool(stacked)

        branch_kwargs = dict(
            input_dim=self.branch_input_dim,
            output_dim=1 if self.stacked else self.latent_dim,
            depth=branch_depth,
            width=width,
            activation=activation,
            activate_output=False,
            dense_bias=dense_bias,
            output_bias=branch_output_bias,
        )
        if self.stacked:
            self.branch = nn.ModuleList(
                DenseNetwork(**branch_kwargs) for _ in range(self.latent_dim)
            )
        else:
            self.branch = DenseNetwork(**branch_kwargs)

        self.trunk = DenseNetwork(
            input_dim=self.trunk_input_dim,
            output_dim=self.latent_dim,
            depth=trunk_depth,
            width=width,
            activation=activation,
            activate_output=True,
            dense_bias=dense_bias,
            output_bias=dense_bias,
        )
        if global_bias:
            self.output_bias = nn.Parameter(torch.zeros(1))
        else:
            self.register_parameter("output_bias", None)
        self.reset_parameters(initializer)

    def reset_parameters(self, initializer: str = "xavier_normal") -> None:
        for module in self.modules():
            if not isinstance(module, nn.Linear):
                continue
            if initializer == "xavier_normal":
                nn.init.xavier_normal_(module.weight)
            elif initializer == "xavier_uniform":
                nn.init.xavier_uniform_(module.weight)
            else:
                raise ValueError(f"Unsupported initializer {initializer!r}")
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def encode_branch(self, branch_inputs: Tensor) -> Tensor:
        if branch_inputs.ndim != 2 or branch_inputs.shape[1] != self.branch_input_dim:
            raise ValueError(
                f"branch input must have shape [N,{self.branch_input_dim}], "
                f"got {tuple(branch_inputs.shape)}"
            )
        if self.stacked:
            return torch.cat([head(branch_inputs) for head in self.branch], dim=-1)
        return self.branch(branch_inputs)

    def forward(self, branch_inputs: Tensor, trunk_inputs: Tensor) -> Tensor:
        if trunk_inputs.ndim != 2 or trunk_inputs.shape[1] != self.trunk_input_dim:
            raise ValueError(
                f"trunk input must have shape [N,{self.trunk_input_dim}], "
                f"got {tuple(trunk_inputs.shape)}"
            )
        if branch_inputs.shape[0] != trunk_inputs.shape[0]:
            raise ValueError("branch and trunk batches must contain the same number of rows")
        branch_features = self.encode_branch(branch_inputs)
        trunk_features = self.trunk(trunk_inputs)
        prediction = torch.sum(branch_features * trunk_features, dim=-1, keepdim=True)
        if self.output_bias is not None:
            prediction = prediction + self.output_bias
        return prediction


class FNNBaseline(nn.Module):
    """Paper baseline that concatenates sensor values and the query coordinate."""

    def __init__(
        self,
        branch_input_dim: int,
        trunk_input_dim: int,
        *,
        depth: int = 3,
        width: int = 40,
        activation: str = "relu",
        output_bias: bool = True,
        initializer: str = "xavier_normal",
    ) -> None:
        super().__init__()
        self.branch_input_dim = int(branch_input_dim)
        self.trunk_input_dim = int(trunk_input_dim)
        self.network = DenseNetwork(
            input_dim=self.branch_input_dim + self.trunk_input_dim,
            output_dim=1,
            depth=depth,
            width=width,
            activation=activation,
            activate_output=False,
            output_bias=output_bias,
        )
        for module in self.modules():
            if isinstance(module, nn.Linear):
                if initializer == "xavier_normal":
                    nn.init.xavier_normal_(module.weight)
                elif initializer == "xavier_uniform":
                    nn.init.xavier_uniform_(module.weight)
                else:
                    raise ValueError(f"Unsupported initializer {initializer!r}")
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, branch_inputs: Tensor, trunk_inputs: Tensor) -> Tensor:
        if branch_inputs.ndim != 2 or branch_inputs.shape[1] != self.branch_input_dim:
            raise ValueError("invalid branch input shape")
        if trunk_inputs.ndim != 2 or trunk_inputs.shape[1] != self.trunk_input_dim:
            raise ValueError("invalid trunk input shape")
        return self.network(torch.cat((branch_inputs, trunk_inputs), dim=-1))


def _merged_model_config(config: Mapping[str, Any], experiment: str) -> Dict[str, Any]:
    if experiment not in config.get("experiments", {}):
        raise KeyError(f"Unknown experiment {experiment!r}")
    merged = dict(config.get("model_defaults", {}))
    for key in ("branch_depth", "trunk_depth", "width", "latent_dim"):
        if key in config["experiments"][experiment]:
            merged[key] = config["experiments"][experiment][key]
    return merged


def build_model(
    config: Mapping[str, Any],
    experiment: str,
    variant: str | None = None,
) -> nn.Module:
    """Build a model from the YAML-compatible configuration mapping."""

    experiment_config = config["experiments"][experiment]
    variant_name = variant or experiment_config["default_variant"]
    try:
        variant_config = config["variants"][variant_name]
    except KeyError as exc:
        raise KeyError(f"Unknown model variant {variant_name!r}") from exc
    model_config = _merged_model_config(config, experiment)
    common = dict(
        branch_input_dim=int(experiment_config["sensor_points"]),
        trunk_input_dim=int(experiment_config["trunk_dim"]),
        activation=str(model_config["activation"]),
        initializer=str(model_config["initializer"]),
    )
    if variant_config["architecture"] == "fnn":
        return FNNBaseline(
            **common,
            depth=int(variant_config["depth"]),
            width=int(variant_config["width"]),
            output_bias=bool(variant_config.get("output_bias", True)),
        )
    if variant_config["architecture"] != "deeponet":
        raise ValueError(f"Unsupported architecture {variant_config['architecture']!r}")
    return DeepONet(
        **common,
        latent_dim=int(model_config["latent_dim"]),
        branch_depth=int(model_config["branch_depth"]),
        trunk_depth=int(model_config["trunk_depth"]),
        width=int(model_config["width"]),
        stacked=bool(variant_config["stacked"]),
        dense_bias=bool(model_config.get("dense_bias", True)),
        branch_output_bias=bool(variant_config["branch_output_bias"]),
        global_bias=bool(variant_config["global_bias"]),
    )


def count_parameters(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)


__all__ = ["DeepONet", "FNNBaseline", "build_model", "count_parameters"]