Merge pull request #615 from mathLab/dev

Dev Updates

Author: dario-coscia

README.md

@@ -90,20 +90,21 @@ Do you want to learn more about it? Look at our [Tutorials](https://github.com/m
 ### Solve Data Driven Problems
 Data driven modelling aims to learn a function that given some input data gives an output (e.g. regression, classification, ...). In PINA you can easily do this by:
 ```python
+import torch
 from pina import Trainer
 from pina.model import FeedForward
 from pina.solver import SupervisedSolver
 from pina.problem.zoo import SupervisedProblem
 
 input_tensor  = torch.rand((10, 1))
-output_tensor = input_tensor.pow(3)
+target_tensor = input_tensor.pow(3)
 
 # Step 1. Define problem
 problem = SupervisedProblem(input_tensor, target_tensor)
 # Step 2. Design model (you can use your favourite torch.nn.Module in here)
 model   = FeedForward(input_dimensions=1, output_dimensions=1, layers=[64, 64])
 # Step 3. Define Solver
-solver  = SupervisedSolver(problem, model)
+solver  = SupervisedSolver(problem, model, use_lt=False)
 # Step 4. Train
 trainer = Trainer(solver, max_epochs=1000, accelerator='gpu')
 trainer.train()
@@ -149,6 +150,7 @@ class SimpleODE(SpatialProblem):
 
 # Step 1. Define problem
 problem = SimpleODE()
+problem.discretise_domain(n=100, mode="grid", domains=["D", "x0"])
 # Step 2. Design model (you can use your favourite torch.nn.Module in here)
 model   = FeedForward(input_dimensions=1, output_dimensions=1, layers=[64, 64])
 # Step 3. Define Solver

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/callback/optimizer_callback.py

@@ -21,45 +21,52 @@ def __init__(self, new_optimizers, epoch_switch):
             single :class:`torch.optim.Optimizer` instance or a list of them
             for multiple model solver.
         :type new_optimizers: pina.optim.TorchOptimizer | list
-        :param epoch_switch: The epoch at which the optimizer switch occurs.
-        :type epoch_switch: int
+        :param int epoch_switch: The epoch at which the optimizer switch occurs.
 
         Example:
-            >>> switch_callback = SwitchOptimizer(new_optimizers=optimizer,
-            >>>                                  epoch_switch=10)
+            >>> optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
+            >>> switch_callback = SwitchOptimizer(
+            >>>     new_optimizers=optimizer, epoch_switch=10
+            >>> )
         """
         super().__init__()
 
+        # Check if epoch_switch is greater than 1
         if epoch_switch < 1:
             raise ValueError("epoch_switch must be greater than one.")
 
+        # If new_optimizers is not a list, convert it to a list
         if not isinstance(new_optimizers, list):
             new_optimizers = [new_optimizers]
 
-        # check type consistency
+        # Check consistency
+        check_consistency(epoch_switch, int)
         for optimizer in new_optimizers:
             check_consistency(optimizer, TorchOptimizer)
-        check_consistency(epoch_switch, int)
-        # save new optimizers
+
+        # Store the new optimizers and epoch switch
         self._new_optimizers = new_optimizers
         self._epoch_switch = epoch_switch
 
     def on_train_epoch_start(self, trainer, __):
         """
         Switch the optimizer at the start of the specified training epoch.
 
-        :param trainer: The trainer object managing the training process.
-        :type trainer: pytorch_lightning.Trainer
+        :param lightning.pytorch.Trainer trainer: The trainer object managing
+            the training process.
         :param _: Placeholder argument (not used).
-
-        :return: None
-        :rtype: None
         """
+        # Check if the current epoch matches the switch epoch
         if trainer.current_epoch == self._epoch_switch:
             optims = []
 
+            # Hook the new optimizers to the model parameters
             for idx, optim in enumerate(self._new_optimizers):
                 optim.hook(trainer.solver._pina_models[idx].parameters())
                 optims.append(optim)
 
+            # Update the solver's optimizers
             trainer.solver._pina_optimizers = optims
+
+            # Update the trainer's strategy optimizers
+            trainer.strategy.optimizers = [o.instance for o in optims]

pina/equation/system_equation.py

@@ -8,18 +8,51 @@
 
 class SystemEquation(EquationInterface):
     """
-    Implementation of the System of Equations. Every ``equation`` passed to a
-    :class:`~pina.condition.condition.Condition` object must be either a
-    :class:`~pina.equation.equation.Equation` or a
-    :class:`~pina.equation.system_equation.SystemEquation` instance.
+    Implementation of the System of Equations, to be passed to a
+    :class:`~pina.condition.condition.Condition` object.
+
+    Unlike the :class:`~pina.equation.equation.Equation` class, which represents
+    a single equation, the :class:`SystemEquation` class allows multiple
+    equations to be grouped together into a system. This is particularly useful
+    when dealing with multi-component outputs or coupled physical models, where
+    the residual must be computed collectively across several constraints.
+
+    Each equation in the system must be either:
+    - An instance of :class:`~pina.equation.equation.Equation`;
+    - A callable function.
+
+    The residuals from each equation are computed independently and then
+    aggregated using an optional reduction strategy (e.g., ``mean``, ``sum``).
+    The resulting residual is returned as a single :class:`~pina.LabelTensor`.
+
+    :Example:
+
+    >>> from pina.equation import SystemEquation, FixedValue, FixedGradient
+    >>> from pina import LabelTensor
+    >>> import torch
+    >>> pts = LabelTensor(torch.rand(10, 2), labels=["x", "y"])
+    >>> pts.requires_grad = True
+    >>> output_ = torch.pow(pts, 2)
+    >>> output_.labels = ["u", "v"]
+    >>> system_equation = SystemEquation(
+    ...     [
+    ...         FixedValue(value=1.0, components=["u"]),
+    ...         FixedGradient(value=0.0, components=["v"],d=["y"]),
+    ...     ],
+    ...     reduction="mean",
+    ... )
+    >>> residual = system_equation.residual(pts, output_)
+
     """
 
     def __init__(self, list_equation, reduction=None):
         """
         Initialization of the :class:`SystemEquation` class.
 
-        :param Callable equation: A ``torch`` callable function used to compute
-            the residual of a mathematical equation.
+        :param list_equation: A list containing either callable functions or
+            instances of :class:`~pina.equation.equation.Equation`, used to
+            compute the residuals of mathematical equations.
+        :type list_equation: list[Callable] | list[Equation]
         :param str reduction: The reduction method to aggregate the residuals of
             each equation. Available options are: ``None``, ``mean``, ``sum``,
             ``callable``.
@@ -32,9 +65,10 @@ def __init__(self, list_equation, reduction=None):
         check_consistency([list_equation], list)
 
         # equations definition
-        self.equations = []
-        for _, equation in enumerate(list_equation):
-            self.equations.append(Equation(equation))
+        self.equations = [
+            equation if isinstance(equation, Equation) else Equation(equation)
+            for equation in list_equation
+        ]
 
         # possible reduction
         if reduction == "mean":
@@ -45,7 +79,7 @@ def __init__(self, list_equation, reduction=None):
             self.reduction = reduction
         else:
             raise NotImplementedError(
-                "Only mean and sum reductions implemented."
+                "Only mean and sum reductions are currenly supported."
             )
 
     def residual(self, input_, output_, params_=None):

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,89 @@
+"""Module for the PirateNet block class."""
+
+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. It must have the
+            same shape as ``x``.
+        :param torch.Tensor V: The second shared gating tensor. It must have the
+            same shape as ``x``.
+        :return: The output tensor of the block.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        # Compute the output of the first dense layer with scaling
+        f = self.activation(self.linear1(x) * torch.exp(self.scale1))
+        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
+
+    @property
+    def alpha(self):
+        """
+        Return the alpha parameter.
+
+        :return: The alpha parameter controlling the residual connection.
+        :rtype: torch.nn.Parameter
+        """
+        return self._alpha