models/FNO.py

"""Convolutional Neural Operator used by the CNO Navier--Stokes experiment.

The filename is fixed by the reproduction request.  This module implements a
CNO, not a Fourier Neural Operator.  The implementation is written from the
architecture and filter description in Sections 2 and C.1.4 of arXiv:2302.01178.
No source code from the authors' repository is included here.
"""

from __future__ import annotations

from typing import Literal

import torch
from torch import Tensor, nn
import torch.nn.functional as F


def _design_windowed_sinc(
    num_taps: int,
    resample_factor: int,
    cutoff_denominator: float,
    half_width: float,
) -> Tensor:
    """Construct a finite, symmetric low-pass windowed-sinc filter.

    The paper fixes ``N_tap=12``, ``c_h=0.8`` and a cutoff arbitrarily close
    to the target Nyquist frequency, ``s/2.0001``.  At a high-rate grid used
    for factor-r resampling, this corresponds to a normalized cutoff of
    ``1 / (r * 2.0001)`` cycles per sample.  The exact finite window is not
    specified in the paper; a Kaiser window is used as an explicit,
    configurable approximation to ``scipy.signal.firwin``.
    """
    if num_taps < 2:
        raise ValueError(f"num_taps must be at least 2, got {num_taps}")
    if resample_factor < 1:
        raise ValueError("resample_factor must be positive")
    if cutoff_denominator <= 2.0:
        raise ValueError("cutoff_denominator must be greater than 2")
    if half_width <= 0:
        raise ValueError("half_width must be positive")

    dtype = torch.float64
    positions = torch.arange(num_taps, dtype=dtype) - (num_taps - 1) / 2
    cutoff = 1.0 / (resample_factor * cutoff_denominator)
    ideal = 2.0 * cutoff * torch.sinc(2.0 * cutoff * positions)

    # c_h=0.8 maps to a conventional beta=8.6 window.  Keeping the relation
    # explicit makes the paper-unspecified window choice auditable.
    beta = 8.6 * half_width / 0.8
    window = torch.kaiser_window(num_taps, periodic=False, beta=beta, dtype=dtype)
    kernel = ideal * window
    kernel = kernel / kernel.sum()
    return kernel.to(torch.float32)


class FixedSincResample2d(nn.Module):
    """Separable periodic 2-D windowed-sinc up/downsampling."""

    def __init__(
        self,
        factor: int = 2,
        num_taps: int = 12,
        cutoff_denominator: float = 2.0001,
        half_width: float = 0.8,
    ) -> None:
        super().__init__()
        if factor < 1:
            raise ValueError("factor must be positive")
        self.factor = int(factor)
        kernel = _design_windowed_sinc(
            num_taps=num_taps,
            resample_factor=factor,
            cutoff_denominator=cutoff_denominator,
            half_width=half_width,
        )
        self.register_buffer("kernel", kernel, persistent=True)

    def _filter(self, x: Tensor, gain: float = 1.0) -> Tensor:
        if x.ndim != 4:
            raise ValueError(f"expected BCHW input, got shape {tuple(x.shape)}")
        channels = x.shape[1]
        kernel = self.kernel.to(device=x.device, dtype=x.dtype)
        taps = int(kernel.numel())
        pad_left = (taps - 1) // 2
        pad_right = taps - 1 - pad_left

        weight_x = (kernel * gain).view(1, 1, 1, taps).repeat(channels, 1, 1, 1)
        x = F.pad(x, (pad_left, pad_right, 0, 0), mode="circular")
        x = F.conv2d(x, weight_x, groups=channels)

        weight_y = kernel.view(1, 1, taps, 1).repeat(channels, 1, 1, 1)
        x = F.pad(x, (0, 0, pad_left, pad_right), mode="circular")
        return F.conv2d(x, weight_y, groups=channels)

    def upsample(self, x: Tensor) -> Tensor:
        if self.factor == 1:
            return x
        batch, channels, height, width = x.shape
        up = x.new_zeros(batch, channels, height * self.factor, width * self.factor)
        up[..., :: self.factor, :: self.factor] = x
        return self._filter(up, gain=float(self.factor * self.factor))

    def downsample(self, x: Tensor) -> Tensor:
        if self.factor == 1:
            return x
        if x.shape[-2] % self.factor or x.shape[-1] % self.factor:
            raise ValueError(
                f"spatial shape {tuple(x.shape[-2:])} is not divisible by {self.factor}"
            )
        return self._filter(x)[..., :: self.factor, :: self.factor]


class BandlimitedActivation(nn.Module):
    """Paper Eq. (2.6): upsample, activate, then low-pass/downsample."""

    def __init__(
        self,
        upsampling_factor: int = 2,
        num_taps: int = 12,
        cutoff_denominator: float = 2.0001,
        half_width: float = 0.8,
        negative_slope: float = 0.2,
    ) -> None:
        super().__init__()
        self.negative_slope = float(negative_slope)
        self.resampler = FixedSincResample2d(
            factor=upsampling_factor,
            num_taps=num_taps,
            cutoff_denominator=cutoff_denominator,
            half_width=half_width,
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.resampler.upsample(x)
        x = F.leaky_relu(x, negative_slope=self.negative_slope)
        return self.resampler.downsample(x)


def _periodic_conv(in_channels: int, out_channels: int, kernel_size: int) -> nn.Conv2d:
    if kernel_size % 2 != 1:
        raise ValueError("CNO convolution kernel_size must be odd")
    return nn.Conv2d(
        in_channels,
        out_channels,
        kernel_size=kernel_size,
        padding=kernel_size // 2,
        padding_mode="circular",
    )


class CNOBlock(nn.Module):
    """Physical-space convolution followed by bandlimited activation/resampling."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        resample: Literal["same", "up", "down"],
        activation_kwargs: dict,
        batch_norm: bool = True,
    ) -> None:
        super().__init__()
        self.resample = resample
        self.conv = _periodic_conv(in_channels, out_channels, kernel_size)
        self.norm = nn.BatchNorm2d(out_channels) if batch_norm else nn.Identity()
        self.activation = BandlimitedActivation(**activation_kwargs)
        self.resampler = FixedSincResample2d(
            factor=activation_kwargs["upsampling_factor"],
            num_taps=activation_kwargs["num_taps"],
            cutoff_denominator=activation_kwargs["cutoff_denominator"],
            half_width=activation_kwargs["half_width"],
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.activation(self.norm(self.conv(x)))
        if self.resample == "down":
            return self.resampler.downsample(x)
        if self.resample == "up":
            return self.resampler.upsample(x)
        return x


class ResidualBlock(nn.Module):
    """Paper Eq. (2.7): identity plus K o Sigma o K."""

    def __init__(self, channels: int, kernel_size: int, activation_kwargs: dict) -> None:
        super().__init__()
        self.conv1 = _periodic_conv(channels, channels, kernel_size)
        self.norm1 = nn.BatchNorm2d(channels)
        self.activation = BandlimitedActivation(**activation_kwargs)
        self.conv2 = _periodic_conv(channels, channels, kernel_size)
        self.norm2 = nn.BatchNorm2d(channels)

    def forward(self, x: Tensor) -> Tensor:
        residual = self.norm1(self.conv1(x))
        residual = self.activation(residual)
        residual = self.norm2(self.conv2(residual))
        return x + residual


class InvariantBlock(nn.Module):
    """Paper Eq. (2.8): Sigma o K at an unchanged representation."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        activation_kwargs: dict,
    ) -> None:
        super().__init__()
        self.conv = _periodic_conv(in_channels, out_channels, kernel_size)
        self.norm = nn.BatchNorm2d(out_channels)
        self.activation = BandlimitedActivation(**activation_kwargs)

    def forward(self, x: Tensor) -> Tensor:
        return self.activation(self.norm(self.conv(x)))


class LiftProjectBlock(nn.Module):
    """Two convolutions with no BatchNorm, as specified for lift/project."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        latent_channels: int,
        kernel_size: int,
        activation_kwargs: dict,
    ) -> None:
        super().__init__()
        self.conv1 = _periodic_conv(in_channels, latent_channels, kernel_size)
        self.activation = BandlimitedActivation(**activation_kwargs)
        self.conv2 = _periodic_conv(latent_channels, out_channels, kernel_size)

    def forward(self, x: Tensor) -> Tensor:
        return self.conv2(self.activation(self.conv1(x)))


def _residual_stack(
    channels: int,
    count: int,
    kernel_size: int,
    activation_kwargs: dict,
) -> nn.Module:
    if count == 0:
        return nn.Identity()
    return nn.Sequential(
        *[
            ResidualBlock(channels, kernel_size, activation_kwargs)
            for _ in range(count)
        ]
    )


class CNO2d(nn.Module):
    """Operator U-Net CNO for the paper's 2-D Navier--Stokes experiment."""

    def __init__(
        self,
        in_channels: int = 1,
        out_channels: int = 1,
        base_width: int = 32,
        levels: int = 3,
        bottleneck_residual_blocks: int = 8,
        intermediate_residual_blocks: int = 1,
        kernel_size: int = 3,
        latent_channels: int = 64,
        activation_upsampling_factor: int = 2,
        filter_taps: int = 12,
        filter_half_width: float = 0.8,
        cutoff_denominator: float = 2.0001,
        leaky_relu_slope: float = 0.2,
    ) -> None:
        super().__init__()
        if levels < 1:
            raise ValueError("levels must be positive")
        if base_width % 2:
            raise ValueError("base_width must be even because lift width is d_e/2")

        self.in_channels = int(in_channels)
        self.out_channels = int(out_channels)
        self.base_width = int(base_width)
        self.levels = int(levels)
        self.required_divisor = 2**levels
        lift_width = base_width // 2
        encoder_widths = [base_width * (2**index) for index in range(levels)]
        activation_kwargs = {
            "upsampling_factor": activation_upsampling_factor,
            "num_taps": filter_taps,
            "cutoff_denominator": cutoff_denominator,
            "half_width": filter_half_width,
            "negative_slope": leaky_relu_slope,
        }

        self.lift = LiftProjectBlock(
            in_channels,
            lift_width,
            latent_channels,
            kernel_size,
            activation_kwargs,
        )

        down_blocks: list[nn.Module] = []
        intermediate_blocks: list[nn.Module] = []
        current_width = lift_width
        for level, next_width in enumerate(encoder_widths):
            down_blocks.append(
                CNOBlock(
                    current_width,
                    next_width,
                    kernel_size,
                    "down",
                    activation_kwargs,
                )
            )
            # The bottleneck has its own N_res,b stack; N_res,i belongs to the
            # genuinely intermediate resolutions only.
            count = intermediate_residual_blocks if level < levels - 1 else 0
            intermediate_blocks.append(
                _residual_stack(next_width, count, kernel_size, activation_kwargs)
            )
            current_width = next_width
        self.encoder = nn.ModuleList(down_blocks)
        self.encoder_residuals = nn.ModuleList(intermediate_blocks)

        self.bottleneck = _residual_stack(
            encoder_widths[-1],
            bottleneck_residual_blocks,
            kernel_size,
            activation_kwargs,
        )

        pre_patch: list[nn.Module] = []
        post_patch: list[nn.Module] = []
        up_blocks: list[nn.Module] = []
        decoder_current = encoder_widths[-1]
        output_widths = list(reversed([lift_width] + encoder_widths[:-1]))
        for next_width in output_widths:
            pre_patch.append(
                InvariantBlock(
                    decoder_current,
                    decoder_current,
                    kernel_size,
                    activation_kwargs,
                )
            )
            post_patch.append(
                InvariantBlock(
                    decoder_current * 2,
                    decoder_current,
                    kernel_size,
                    activation_kwargs,
                )
            )
            up_blocks.append(
                CNOBlock(
                    decoder_current,
                    next_width,
                    kernel_size,
                    "up",
                    activation_kwargs,
                )
            )
            decoder_current = next_width
        self.decoder_pre_patch = nn.ModuleList(pre_patch)
        self.decoder_post_patch = nn.ModuleList(post_patch)
        self.decoder = nn.ModuleList(up_blocks)

        self.project = LiftProjectBlock(
            lift_width * 2,
            out_channels,
            latent_channels,
            kernel_size,
            activation_kwargs,
        )

    def forward(self, x: Tensor) -> Tensor:
        if x.ndim != 4:
            raise ValueError(f"CNO2d expects BCHW input, got {tuple(x.shape)}")
        if x.shape[1] != self.in_channels:
            raise ValueError(
                f"expected {self.in_channels} channels, got {x.shape[1]}"
            )
        height, width = x.shape[-2:]
        if height % self.required_divisor or width % self.required_divisor:
            raise ValueError(
                f"spatial shape {(height, width)} must be divisible by "
                f"2**levels={self.required_divisor}"
            )

        lifted = self.lift(x)
        encoded = lifted
        skips: list[Tensor] = []
        for down, residuals in zip(self.encoder, self.encoder_residuals):
            encoded = residuals(down(encoded))
            skips.append(encoded)

        decoded = self.bottleneck(encoded)
        for pre, post, up, skip in zip(
            self.decoder_pre_patch,
            self.decoder_post_patch,
            self.decoder,
            reversed(skips),
        ):
            decoded = pre(decoded)
            if decoded.shape[-2:] != skip.shape[-2:]:
                raise RuntimeError(
                    "decoder/skip spatial mismatch before patching: "
                    f"{tuple(decoded.shape)} versus {tuple(skip.shape)}"
                )
            decoded = post(torch.cat((decoded, skip), dim=1))
            decoded = up(decoded)

        if decoded.shape[-2:] != lifted.shape[-2:]:
            raise RuntimeError(
                f"final decoder/lift mismatch: {decoded.shape} versus {lifted.shape}"
            )
        output = self.project(torch.cat((decoded, lifted), dim=1))
        if output.shape[-2:] != (height, width):
            raise RuntimeError(
                f"CNO changed output grid from {(height, width)} to {output.shape[-2:]}"
            )
        return output


def build_model(model_config: dict) -> CNO2d:
    """Build a CNO2d from the ``model`` section of config.yaml."""
    return CNO2d(**model_config)


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


__all__ = ["CNO2d", "build_model", "count_trainable_parameters"]