model/gpinn.py

from __future__ import annotations

from collections.abc import Sequence

import numpy as np
import torch
import torch.autograd as autograd
import torch.nn as nn


class SubNet(nn.Module):
    """Fully connected network with trainable adaptive activations."""

    def __init__(self, layers: Sequence[int], activation: str = "tanh") -> None:
        super().__init__()
        if len(layers) < 2:
            raise ValueError("layers must contain at least an input and an output size")
        if activation not in {"tanh", "sin", "cos"}:
            raise ValueError(f"unsupported activation: {activation}")

        self.activation = activation
        self.linears = nn.ModuleList(
            nn.Linear(layers[index], layers[index + 1], dtype=torch.float64)
            for index in range(len(layers) - 1)
        )
        # Keep one value per layer for compatibility with the published checkpoints.
        self.a = nn.ParameterList(
            nn.Parameter(torch.tensor(0.05, dtype=torch.float64))
            for _ in range(len(layers) - 1)
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        hidden = inputs
        for index, linear in enumerate(self.linears[:-1]):
            hidden = 20.0 * self.a[index] * linear(hidden)
            hidden = getattr(torch, self.activation)(hidden)
        return self.linears[-1](hidden)


class gPINN(nn.Module):
    """Gradient-enhanced physics-informed neural network backbone."""

    def __init__(self, layers: Sequence[int], activation: str = "tanh") -> None:
        super().__init__()
        self.net = SubNet(layers, activation)

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


class gPINNPoisson2D(nn.Module):
    """2D Poisson model whose output transform enforces zero boundaries."""

    def __init__(self, layers: Sequence[int]) -> None:
        super().__init__()
        self.net = gPINN(layers)

    def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
        x = coordinates[:, 0:1]
        y = coordinates[:, 1:2]
        return x * y * (1.0 - x) * (1.0 - y) * self.net(coordinates)


def exact_poisson1d(x: np.ndarray) -> np.ndarray:
    solution = x + np.sin(8.0 * x) / 8.0
    for frequency in range(1, 5):
        solution += np.sin(frequency * x) / frequency
    return solution


def gpinn_loss_poisson1d(
    model: nn.Module,
    interior: torch.Tensor,
    boundary: torch.Tensor,
    boundary_values: torch.Tensor,
    gradient_weight: float,
) -> tuple[torch.Tensor, dict[str, float]]:
    x = interior.detach().requires_grad_(True)
    prediction = model(x)
    first = autograd.grad(prediction, x, torch.ones_like(prediction), create_graph=True)[0]
    second = autograd.grad(first, x, torch.ones_like(first), create_graph=True)[0]
    third = autograd.grad(second, x, torch.ones_like(second), create_graph=True)[0]

    source = 8.0 * torch.sin(8.0 * x) + sum(
        frequency * torch.sin(frequency * x) for frequency in range(1, 5)
    )
    source_gradient = (
        torch.cos(x)
        + 4.0 * torch.cos(2.0 * x)
        + 9.0 * torch.cos(3.0 * x)
        + 16.0 * torch.cos(4.0 * x)
        + 64.0 * torch.cos(8.0 * x)
    )

    residual = -second - source
    residual_gradient = -third - source_gradient
    residual_loss = torch.mean(residual.square())
    boundary_loss = torch.mean((model(boundary) - boundary_values).square())
    gradient_loss = torch.mean(residual_gradient.square())
    loss = residual_loss + boundary_loss + gradient_weight * gradient_loss
    return loss, {
        "residual": residual_loss.item(),
        "boundary": boundary_loss.item(),
        "gradient": gradient_loss.item(),
    }


def exact_poisson2d(x: np.ndarray, y: np.ndarray, exponent: float) -> np.ndarray:
    return (16.0 * x * y * (1.0 - x) * (1.0 - y)) ** exponent


def _poisson2d_source(coordinates: torch.Tensor, exponent: float) -> torch.Tensor:
    x = coordinates[:, 0:1]
    y = coordinates[:, 1:2]
    exact = (16.0 * x * y * (1.0 - x) * (1.0 - y)) ** exponent
    exact_gradient = autograd.grad(
        exact, coordinates, torch.ones_like(exact), create_graph=True
    )[0]
    exact_xx = autograd.grad(
        exact_gradient[:, 0:1],
        coordinates,
        torch.ones_like(exact_gradient[:, 0:1]),
        create_graph=True,
    )[0][:, 0:1]
    exact_yy = autograd.grad(
        exact_gradient[:, 1:2],
        coordinates,
        torch.ones_like(exact_gradient[:, 1:2]),
        create_graph=True,
    )[0][:, 1:2]
    return -(exact_xx + exact_yy)


def gpinn_loss_poisson2d(
    model: nn.Module,
    interior: torch.Tensor,
    gradient_weight: float,
    exponent: float,
) -> tuple[torch.Tensor, dict[str, float]]:
    coordinates = interior.detach().requires_grad_(True)
    prediction = model(coordinates)
    prediction_gradient = autograd.grad(
        prediction, coordinates, torch.ones_like(prediction), create_graph=True
    )[0]
    prediction_xx = autograd.grad(
        prediction_gradient[:, 0:1],
        coordinates,
        torch.ones_like(prediction_gradient[:, 0:1]),
        create_graph=True,
    )[0][:, 0:1]
    prediction_yy = autograd.grad(
        prediction_gradient[:, 1:2],
        coordinates,
        torch.ones_like(prediction_gradient[:, 1:2]),
        create_graph=True,
    )[0][:, 1:2]

    residual = prediction_xx + prediction_yy + _poisson2d_source(coordinates, exponent)
    residual_gradient = autograd.grad(
        residual, coordinates, torch.ones_like(residual), create_graph=True
    )[0]
    residual_loss = torch.mean(residual.square())
    gradient_loss = torch.mean(residual_gradient[:, 0:1].square()) + torch.mean(
        residual_gradient[:, 1:2].square()
    )
    loss = residual_loss + gradient_weight * gradient_loss
    return loss, {"residual": residual_loss.item(), "gradient": gradient_loss.item()}


def output_transform_burgers(
    coordinates: torch.Tensor, raw_prediction: torch.Tensor
) -> torch.Tensor:
    x = coordinates[:, 0:1]
    t = coordinates[:, 1:2]
    return (
        (1.0 - x) * (1.0 + x) * (1.0 - torch.exp(-t)) * raw_prediction
        - torch.sin(torch.pi * x)
    )


class gPINNBurgers(nn.Module):
    """Burgers model with hard initial and boundary constraints."""

    def __init__(self, layers: Sequence[int]) -> None:
        super().__init__()
        self.net = gPINN(layers)

    def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
        return output_transform_burgers(coordinates, self.net(coordinates))


def burgers_residual(
    model: nn.Module, coordinates: torch.Tensor, create_graph: bool = True
) -> tuple[torch.Tensor, torch.Tensor]:
    inputs = coordinates.detach().requires_grad_(True)
    prediction = model(inputs)
    prediction_gradient = autograd.grad(
        prediction, inputs, torch.ones_like(prediction), create_graph=True
    )[0]
    prediction_x = prediction_gradient[:, 0:1]
    prediction_t = prediction_gradient[:, 1:2]
    prediction_xx = autograd.grad(
        prediction_x,
        inputs,
        torch.ones_like(prediction_x),
        create_graph=create_graph,
    )[0][:, 0:1]
    viscosity = 0.01 / torch.pi
    residual = prediction_t + prediction * prediction_x - viscosity * prediction_xx
    return residual, inputs


def burgers_gpinn_terms(
    model: nn.Module, coordinates: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    residual, inputs = burgers_residual(model, coordinates, create_graph=True)
    residual_gradient = autograd.grad(
        residual, inputs, torch.ones_like(residual), create_graph=True
    )[0]
    return residual, residual_gradient[:, 0:1], residual_gradient[:, 1:2]