add pirate network

Author: GiovanniCanali

docs/source/_rst/_code.rst

@@ -104,6 +104,7 @@ Models
     LowRankNeuralOperator <model/low_rank_neural_operator.rst>
     GraphNeuralOperator <model/graph_neural_operator.rst>
     GraphNeuralKernel <model/graph_neural_operator_integral_kernel.rst>
+    PirateNet <model/pirate_network.rst>
 
 Blocks
 -------------
@@ -121,6 +122,7 @@ Blocks
     Continuous Convolution Interface <model/block/convolution_interface.rst>
     Continuous Convolution Block <model/block/convolution.rst>
     Orthogonal Block <model/block/orthogonal.rst>
+    PirateNet Block <model/block/pirate_network_block.rst>
 
 Message Passing
 -------------------

docs/source/_rst/model/block/pirate_network_block.rst

@@ -0,0 +1,8 @@
+PirateNet Block
+=======================================
+.. currentmodule:: pina.model.block.pirate_network_block
+    
+.. autoclass:: PirateNetBlock
+    :members:
+    :show-inheritance:
+

docs/source/_rst/model/pirate_network.rst

@@ -0,0 +1,7 @@
+PirateNet
+=======================
+.. currentmodule:: pina.model.pirate_network
+
+.. autoclass:: PirateNet
+   :members:
+   :show-inheritance:

pina/model/__init__.py

@@ -13,6 +13,7 @@
     "LowRankNeuralOperator",
     "Spline",
     "GraphNeuralOperator",
+    "PirateNet",
 ]
 
 from .feed_forward import FeedForward, ResidualFeedForward
@@ -24,3 +25,4 @@
 from .low_rank_neural_operator import LowRankNeuralOperator
 from .spline import Spline
 from .graph_neural_operator import GraphNeuralOperator
+from .pirate_network import PirateNet

pina/model/block/__init__.py

@@ -18,6 +18,7 @@
     "LowRankBlock",
     "RBFBlock",
     "GNOBlock",
+    "PirateNetBlock",
 ]
 
 from .convolution_2d import ContinuousConvBlock
@@ -35,3 +36,4 @@
 from .low_rank_block import LowRankBlock
 from .rbf_block import RBFBlock
 from .gno_block import GNOBlock
+from .pirate_network_block import PirateNetBlock

pina/model/block/pirate_network_block.py

@@ -0,0 +1,81 @@
+import torch
+from ...utils import check_consistency, check_positive_integer
+
+
+class PirateNetBlock(torch.nn.Module):
+    """
+    The inner block of Physics-Informed residual adaptive network (PirateNet).
+
+    The block consists of three dense layers with dual gating operations and an
+    adaptive residual connection. The trainable ``alpha`` parameter controls
+    the contribution of the residual connection.
+
+    .. seealso::
+
+        **Original reference**:
+        Wang, S., Sankaran, S., Stinis., P., Perdikaris, P. (2025).
+        *Simulating Three-dimensional Turbulence with Physics-informed Neural
+        Networks*.
+        DOI: `arXiv preprint arXiv:2507.08972.
+        <https://arxiv.org/abs/2507.08972>`_
+    """
+
+    def __init__(self, inner_size, activation):
+        """
+        Initialization of the :class:`PirateNetBlock` class.
+
+        :param int inner_size: The number of hidden units in the dense layers.
+        :param torch.nn.Module activation: The activation function.
+        """
+        super().__init__()
+
+        # Check consistency
+        check_consistency(activation, torch.nn.Module, subclass=True)
+        check_positive_integer(inner_size, strict=True)
+
+        # Initialize the linear transformations of the dense layers
+        self.linear1 = torch.nn.Linear(inner_size, inner_size)
+        self.linear2 = torch.nn.Linear(inner_size, inner_size)
+        self.linear3 = torch.nn.Linear(inner_size, inner_size)
+
+        # Initialize the scales of the dense layers
+        self.scale1 = torch.nn.Parameter(torch.zeros(inner_size))
+        self.scale2 = torch.nn.Parameter(torch.zeros(inner_size))
+        self.scale3 = torch.nn.Parameter(torch.zeros(inner_size))
+
+        # Initialize the adaptive residual connection parameter
+        self.alpha = torch.nn.Parameter(torch.zeros(1))
+
+        # Initialize the activation function
+        self.activation = activation()
+
+    def forward(self, x, U, V):
+        """
+        Forward pass of the PirateNet block. It computes the output of the block
+        by applying the dense layers with scaling, and combines the results with
+        the input using the adaptive residual connection.
+
+        :param x: The input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor U: The first shared gating tensor.
+        :param torch.Tensor V: The second shared gating tensor.
+        :return: The output tensor of the block.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        # Compute the output of the first dense layer with scaling
+        print(f"{x.shape=}")
+        print(f"{self.linear1(x).shape=}")
+        print(f"{self.scale1.shape=}")
+        print(f"{torch.exp(self.scale1).shape=}")
+        print(f"{(self.linear1(x) * torch.exp(self.scale1)).shape=}")
+        f = self.activation(self.linear1(x) * torch.exp(self.scale1))
+        print(f.shape, U.shape)
+        z1 = f * U + (1 - f) * V
+
+        # Compute the output of the second dense layer with scaling
+        g = self.activation(self.linear2(z1) * torch.exp(self.scale2))
+        z2 = g * U + (1 - g) * V
+
+        # Compute the output of the block
+        h = self.activation(self.linear3(z2) * torch.exp(self.scale3))
+        return self.alpha * h + (1 - self.alpha) * x

pina/model/pirate_network.py

@@ -0,0 +1,106 @@
+import torch
+from .block import FourierFeatureEmbedding, PirateNetBlock
+from ..utils import check_consistency, check_positive_integer
+
+
+class PirateNet(torch.nn.Module):
+    """
+    Implementation of Physics-Informed residual adaptive network (PirateNet).
+
+    The model consists of a Fourier feature embedding layer, multiple PirateNet
+    blocks, and a final output layer. Each PirateNet block consist of three
+    dense layers with dual gating mechanism and an adaptive residual connection,
+    whose contribution is controlled by a trainable parameter ``alpha``.
+
+    The PirateNet, augmented with random weight factorization, is designed to
+    mitigate spectral bias in deep networks.
+
+    .. seealso::
+
+        **Original reference**:
+        Wang, S., Sankaran, S., Stinis., P., Perdikaris, P. (2025).
+        *Simulating Three-dimensional Turbulence with Physics-informed Neural
+        Networks*.
+        DOI: `arXiv preprint arXiv:2507.08972.
+        <https://arxiv.org/abs/2507.08972>`_
+    """
+
+    def __init__(
+        self,
+        input_dimension,
+        inner_size,
+        output_dimension,
+        sigma,
+        n_layers=3,
+        activation=torch.nn.Tanh,
+    ):
+        """
+        Initialization of the :class:`PirateNet` class.
+
+        :param int input_dimension: The number of input features.
+        :param int inner_size: The number of hidden units in the dense layers.
+        :param int output_dimension: The number of output features.
+        :param float sigma: The scaling factor for the Fourier embedding. This
+            value must reflect the granularity of the scale in the solution of
+            the problem to be solved.
+        :param int n_layers: The number of PirateNet blocks in the model.
+            Default is 3.
+        :param torch.nn.Module activation: The activation function to be used in
+            the blocks. Default is :class:`torch.nn.Tanh`.
+        """
+        super().__init__()
+
+        # Check consistency
+        check_consistency(sigma, (float, int))
+        check_consistency(activation, torch.nn.Module, subclass=True)
+        check_positive_integer(input_dimension, strict=True)
+        check_positive_integer(inner_size, strict=True)
+        check_positive_integer(output_dimension, strict=True)
+        check_positive_integer(n_layers, strict=True)
+
+        # Initialize the activation function
+        self.activation = activation()
+
+        # Initialize the Fourier embedding
+        self.fourier_embedding = FourierFeatureEmbedding(
+            input_dimension=input_dimension,
+            output_dimension=inner_size,
+            sigma=sigma,
+        )
+
+        # Initialize the shared dense layers
+        self.linear1 = torch.nn.Linear(inner_size, inner_size)
+        self.linear2 = torch.nn.Linear(inner_size, inner_size)
+
+        # Initialize the PirateNet blocks
+        self.blocks = torch.nn.ModuleList(
+            [PirateNetBlock(inner_size, activation) for _ in range(n_layers)]
+        )
+
+        # Initialize the output layer
+        self.output_layer = torch.nn.Linear(inner_size, output_dimension)
+
+    def forward(self, input_):
+        """
+        Forward pass of the PirateNet model. It applies the Fourier feature
+        embedding, computes the shared gating tensors U and V, and passes the
+        input through each block in the network. Finally, it applies the output
+        layer to produce the final output.
+
+        :param input_: The input tensor for the model.
+        :type input_: torch.Tensor | LabelTensor
+        :return: The output tensor of the model.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        # Apply the Fourier feature embedding
+        x = self.fourier_embedding(input_)
+
+        # Compute U and V from the shared dense layers
+        U = self.activation(self.linear1(x))
+        V = self.activation(self.linear2(x))
+
+        # Pass through each block in the network
+        for block in self.blocks:
+            x = block(x, U, V)
+
+        return self.output_layer(x)

tests/test_blocks/test_pirate_network_block.py

@@ -0,0 +1,53 @@
+import torch
+import pytest
+from pina.model.block import PirateNetBlock
+
+data = torch.rand((20, 3))
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_constructor(inner_size):
+
+    PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+
+    # Should fail if inner_size is negative
+    with pytest.raises(AssertionError):
+        PirateNetBlock(inner_size=-1, activation=torch.nn.Tanh)
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_forward(inner_size):
+
+    model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+
+    # Create dummy embedding
+    dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
+    x = dummy_embedding(data)
+
+    # Create dummy U and V tensors
+    U = torch.rand((data.shape[0], inner_size))
+    V = torch.rand((data.shape[0], inner_size))
+
+    output_ = model(x, U, V)
+    assert output_.shape == (data.shape[0], inner_size)
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_backward(inner_size):
+
+    model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+    data.requires_grad_()
+
+    # Create dummy embedding
+    dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
+    x = dummy_embedding(data)
+
+    # Create dummy U and V tensors
+    U = torch.rand((data.shape[0], inner_size))
+    V = torch.rand((data.shape[0], inner_size))
+
+    output_ = model(x, U, V)
+
+    loss = torch.mean(output_)
+    loss.backward()
+    assert data.grad.shape == data.shape