Adding new problems to problem.zoo (#484)

* adding problems
* add tests
* update doc + formatting

---------

Co-authored-by: Dario Coscia <dariocos99@gmail.com>

Author: GiovanniCanali

pina/problem/zoo/__init__.py

@@ -1,15 +1,19 @@
 """TODO"""
 
 __all__ = [
-    "Poisson2DSquareProblem",
     "SupervisedProblem",
-    "InversePoisson2DSquareProblem",
+    "HelmholtzProblem",
+    "AllenCahnProblem",
+    "AdvectionProblem",
+    "Poisson2DSquareProblem",
     "DiffusionReactionProblem",
-    "InverseDiffusionReactionProblem",
+    "InversePoisson2DSquareProblem",
 ]
 
-from .poisson_2d_square import Poisson2DSquareProblem
 from .supervised_problem import SupervisedProblem
-from .inverse_poisson_2d_square import InversePoisson2DSquareProblem
+from .helmholtz import HelmholtzProblem
+from .allen_cahn import AllenCahnProblem
+from .advection import AdvectionProblem
+from .poisson_2d_square import Poisson2DSquareProblem
 from .diffusion_reaction import DiffusionReactionProblem
-from .inverse_diffusion_reaction import InverseDiffusionReactionProblem
+from .inverse_poisson_2d_square import InversePoisson2DSquareProblem

pina/problem/zoo/advection.py

@@ -0,0 +1,107 @@
+"""Formulation of the advection problem."""
+
+import torch
+from ... import Condition
+from ...operator import grad
+from ...equation import Equation
+from ...domain import CartesianDomain
+from ...utils import check_consistency
+from ...problem import SpatialProblem, TimeDependentProblem
+
+
+class AdvectionEquation(Equation):
+    """
+    Implementation of the advection equation.
+    """
+
+    def __init__(self, c):
+        """
+        Initialize the advection equation.
+
+        :param c: The advection velocity parameter.
+        :type c: float | int
+        """
+        self.c = c
+        check_consistency(self.c, (float, int))
+
+        def equation(input_, output_):
+            """
+            Implementation of the advection equation.
+
+            :param LabelTensor input_: Input data of the problem.
+            :param LabelTensor output_: Output data of the problem.
+            :return: The residual of the advection equation.
+            :rtype: LabelTensor
+            """
+            u_x = grad(output_, input_, components=["u"], d=["x"])
+            u_t = grad(output_, input_, components=["u"], d=["t"])
+            return u_t + self.c * u_x
+
+        super().__init__(equation)
+
+
+def initial_condition(input_, output_):
+    """
+    Implementation of the initial condition.
+
+    :param LabelTensor input_: Input data of the problem.
+    :param LabelTensor output_: Output data of the problem.
+    :return: The residual of the initial condition.
+    :rtype: LabelTensor
+    """
+    return output_ - torch.sin(input_.extract("x"))
+
+
+class AdvectionProblem(SpatialProblem, TimeDependentProblem):
+    r"""
+    Implementation of the advection problem in the spatial interval
+    :math:`[0, 2 \pi]` and temporal interval :math:`[0, 1]`.
+
+    .. seealso::
+
+        **Original reference**: Wang, Sifan, et al. *An expert's guide to
+        training physics-informed neural networks*.
+        arXiv preprint arXiv:2308.08468 (2023).
+        DOI: `arXiv:2308.08468  <https://arxiv.org/abs/2308.08468>`_.
+    """
+
+    output_variables = ["u"]
+    spatial_domain = CartesianDomain({"x": [0, 2 * torch.pi]})
+    temporal_domain = CartesianDomain({"t": [0, 1]})
+
+    domains = {
+        "D": CartesianDomain({"x": [0, 2 * torch.pi], "t": [0, 1]}),
+        "t0": CartesianDomain({"x": [0, 2 * torch.pi], "t": 0.0}),
+    }
+
+    conditions = {
+        "t0": Condition(domain="t0", equation=Equation(initial_condition)),
+    }
+
+    def __init__(self, c=1.0):
+        """
+        Initialize the advection problem.
+
+        :param c: The advection velocity parameter.
+        :type c: float | int
+        """
+        super().__init__()
+
+        self.c = c
+        check_consistency(self.c, (float, int))
+
+        self.conditions["D"] = Condition(
+            domain="D", equation=AdvectionEquation(self.c)
+        )
+
+    def solution(self, pts):
+        """
+        Implementation of the analytical solution of the advection problem.
+
+        :param LabelTensor pts: Points where the solution is evaluated.
+        :return: The analytical solution of the advection problem.
+        :rtype: LabelTensor
+        """
+        sol = torch.sin(pts.extract("x") - self.c * pts.extract("t"))
+        sol.labels = self.output_variables
+        return sol

pina/problem/zoo/allen_cahn.py

@@ -0,0 +1,66 @@
+"""Formulation of the Allen Cahn problem."""
+
+import torch
+from ... import Condition
+from ...equation import Equation
+from ...domain import CartesianDomain
+from ...operator import grad, laplacian
+from ...problem import SpatialProblem, TimeDependentProblem
+
+
+def allen_cahn_equation(input_, output_):
+    """
+    Implementation of the Allen Cahn equation.
+
+    :param LabelTensor input_: Input data of the problem.
+    :param LabelTensor output_: Output data of the problem.
+    :return: The residual of the Allen Cahn equation.
+    :rtype: LabelTensor
+    """
+    u_t = grad(output_, input_, components=["u"], d=["t"])
+    u_xx = laplacian(output_, input_, components=["u"], d=["x"])
+    return u_t - 0.0001 * u_xx + 5 * output_**3 - 5 * output_
+
+
+def initial_condition(input_, output_):
+    """
+    Definition of the initial condition of the Allen Cahn problem.
+
+    :param LabelTensor input_: Input data of the problem.
+    :param LabelTensor output_: Output data of the problem.
+    :return: The residual of the initial condition.
+    :rtype: LabelTensor
+    """
+    x = input_.extract("x")
+    u_0 = x**2 * torch.cos(torch.pi * x)
+    return output_ - u_0
+
+
+class AllenCahnProblem(TimeDependentProblem, SpatialProblem):
+    r"""
+    Implementation of the Allen Cahn problem in the spatial interval
+    :math:`[-1, 1]` and temporal interval :math:`[0, 1]`.
+
+    .. seealso::
+        **Original reference**: Sokratis J. Anagnostopoulos, Juan D. Toscano,
+        Nikolaos Stergiopulos, and George E. Karniadakis.
+        *Residual-based attention and connection to information
+        bottleneck theory in PINNs*.
+        Computer Methods in Applied Mechanics and Engineering 421 (2024): 116805
+        DOI: `10.1016/
+        j.cma.2024.116805 <https://doi.org/10.1016/j.cma.2024.116805>`_.
+    """
+
+    output_variables = ["u"]
+    spatial_domain = CartesianDomain({"x": [-1, 1]})
+    temporal_domain = CartesianDomain({"t": [0, 1]})
+
+    domains = {
+        "D": CartesianDomain({"x": [-1, 1], "t": [0, 1]}),
+        "t0": CartesianDomain({"x": [-1, 1], "t": 0.0}),
+    }
+
+    conditions = {
+        "D": Condition(domain="D", equation=Equation(allen_cahn_equation)),
+        "t0": Condition(domain="t0", equation=Equation(initial_condition)),
+    }

pina/problem/zoo/diffusion_reaction.py

@@ -1,22 +1,26 @@
-"""Definition of the diffusion-reaction problem."""
+"""Formulation of the diffusion-reaction problem."""
 
 import torch
-from pina import Condition
-from pina.problem import SpatialProblem, TimeDependentProblem
-from pina.equation.equation import Equation
-from pina.domain import CartesianDomain
-from pina.operator import grad
+from ... import Condition
+from ...domain import CartesianDomain
+from ...operator import grad, laplacian
+from ...equation import Equation, FixedValue
+from ...problem import SpatialProblem, TimeDependentProblem
 
 
 def diffusion_reaction(input_, output_):
     """
     Implementation of the diffusion-reaction equation.
+
+    :param LabelTensor input_: Input data of the problem.
+    :param LabelTensor output_: Output data of the problem.
+    :return: The residual of the diffusion-reaction equation.
+    :rtype: LabelTensor
     """
     x = input_.extract("x")
     t = input_.extract("t")
-    u_t = grad(output_, input_, d="t")
-    u_x = grad(output_, input_, d="x")
-    u_xx = grad(u_x, input_, d="x")
+    u_t = grad(output_, input_, components=["u"], d=["t"])
+    u_xx = laplacian(output_, input_, components=["u"], d=["x"])
     r = torch.exp(-t) * (
         1.5 * torch.sin(2 * x)
         + (8 / 3) * torch.sin(3 * x)
@@ -26,30 +30,72 @@ def diffusion_reaction(input_, output_):
     return u_t - u_xx - r
 
 
-class DiffusionReactionProblem(TimeDependentProblem, SpatialProblem):
+def initial_condition(input_, output_):
+    """
+    Definition of the initial condition of the diffusion-reaction problem.
+
+    :param LabelTensor input_: Input data of the problem.
+    :param LabelTensor output_: Output data of the problem.
+    :return: The residual of the initial condition.
+    :rtype: LabelTensor
     """
-    Implementation of the diffusion-reaction problem on the spatial interval
-    [-pi, pi] and temporal interval [0,1].
+    x = input_.extract("x")
+    u_0 = (
+        torch.sin(x)
+        + (1 / 2) * torch.sin(2 * x)
+        + (1 / 3) * torch.sin(3 * x)
+        + (1 / 4) * torch.sin(4 * x)
+        + (1 / 8) * torch.sin(8 * x)
+    )
+    return output_ - u_0
+
+
+class DiffusionReactionProblem(TimeDependentProblem, SpatialProblem):
+    r"""
+    Implementation of the diffusion-reaction problem in the spatial interval
+    :math:`[-\pi, \pi]` and temporal interval :math:`[0, 1]`.
+
+    .. seealso::
+        **Original reference**: Si, Chenhao, et al. *Complex Physics-Informed
+        Neural Network.* arXiv preprint arXiv:2502.04917 (2025).
+        DOI: `arXiv:2502.04917 <https://arxiv.org/abs/2502.04917>`_.
     """
 
     output_variables = ["u"]
     spatial_domain = CartesianDomain({"x": [-torch.pi, torch.pi]})
     temporal_domain = CartesianDomain({"t": [0, 1]})
 
+    domains = {
+        "D": CartesianDomain({"x": [-torch.pi, torch.pi], "t": [0, 1]}),
+        "g1": CartesianDomain({"x": -torch.pi, "t": [0, 1]}),
+        "g2": CartesianDomain({"x": torch.pi, "t": [0, 1]}),
+        "t0": CartesianDomain({"x": [-torch.pi, torch.pi], "t": 0.0}),
+    }
+
     conditions = {
-        "D": Condition(
-            domain=CartesianDomain({"x": [-torch.pi, torch.pi], "t": [0, 1]}),
-            equation=Equation(diffusion_reaction),
-        )
+        "D": Condition(domain="D", equation=Equation(diffusion_reaction)),
+        "g1": Condition(domain="g1", equation=FixedValue(0.0)),
+        "g2": Condition(domain="g2", equation=FixedValue(0.0)),
+        "t0": Condition(domain="t0", equation=Equation(initial_condition)),
     }
 
-    def _solution(self, pts):
+    def solution(self, pts):
+        """
+        Implementation of the analytical solution of the diffusion-reaction
+        problem.
+
+        :param LabelTensor pts: Points where the solution is evaluated.
+        :return: The analytical solution of the diffusion-reaction problem.
+        :rtype: LabelTensor
+        """
         t = pts.extract("t")
         x = pts.extract("x")
-        return torch.exp(-t) * (
+        sol = torch.exp(-t) * (
             torch.sin(x)
             + (1 / 2) * torch.sin(2 * x)
             + (1 / 3) * torch.sin(3 * x)
             + (1 / 4) * torch.sin(4 * x)
             + (1 / 8) * torch.sin(8 * x)
         )
+        sol.labels = self.output_variables
+        return sol