enrich equation factory

Author: GiovanniCanali

docs/source/_rst/equation/equation_factory.rst

@@ -15,5 +15,29 @@ Equation Factory
    :show-inheritance:
 
 .. autoclass:: FixedLaplacian
+   :members:
+   :show-inheritance:
+
+.. autoclass:: Laplace
+   :members:
+   :show-inheritance:
+
+.. autoclass:: AdvectionEquation
+   :members:
+   :show-inheritance:
+
+.. autoclass:: AllenCahnEquation
+   :members:
+   :show-inheritance:
+
+.. autoclass:: DiffusionReactionEquation
+   :members:
+   :show-inheritance:
+
+.. autoclass:: HelmholtzEquation
+   :members:
+   :show-inheritance:
+
+.. autoclass:: PoissonEquation
    :members:
    :show-inheritance:

pina/equation/__init__.py

@@ -7,6 +7,12 @@
     "FixedGradient",
     "FixedFlux",
     "FixedLaplacian",
+    "Laplace",
+    "AdvectionEquation",
+    "AllenCahnEquation",
+    "DiffusionReactionEquation",
+    "HelmholtzEquation",
+    "PoissonEquation",
 ]
 
 from .equation import Equation
@@ -15,5 +21,11 @@
     FixedGradient,
     FixedLaplacian,
     FixedValue,
+    Laplace,
+    AdvectionEquation,
+    AllenCahnEquation,
+    DiffusionReactionEquation,
+    HelmholtzEquation,
+    PoissonEquation,
 )
 from .system_equation import SystemEquation

pina/equation/equation_factory.py

@@ -1,7 +1,9 @@
 """Module for defining various general equations."""
-
+from typing import Callable
+import torch
 from .equation import Equation
 from ..operator import grad, div, laplacian
+from ..utils import check_consistency
 
 
 class FixedValue(Equation):
@@ -121,7 +123,7 @@ def __init__(self, value, components=None, d=None):
 
         :param float value: The fixed value to be enforced to the laplacian.
         :param list[str] components: The name of the output variables for which
-            the null laplace condition is applied. It should be a subset of the
+            the fixed laplace condition is applied. It should be a subset of the
             output labels. If ``None``, all output variables are considered.
             Default is ``None``.
         :param list[str] d: The name of the input variables on which the
@@ -146,3 +148,297 @@ def equation(input_, output_):
             )
 
         super().__init__(equation)
+
+
+class Laplace(FixedLaplacian):
+    """
+    Equation to enforce a null laplacian for a specific condition.
+    """
+
+    def __init__(self, components=None, d=None):
+        """
+        Initialization of the :class:`Laplace` class.
+
+        :param list[str] components: The name of the output variables for which
+            the null laplace condition is applied. It should be a subset of the
+            output labels. If ``None``, all output variables are considered.
+            Default is ``None``.
+        :param list[str] d: The name of the input variables on which the
+            laplacian is computed. It should be a subset of the input labels.
+            If ``None``, all the input variables are considered.
+            Default is ``None``.
+        """
+        super().__init__(0.0, components=components, d=d)
+
+
+class AdvectionEquation(Equation):
+    r"""
+    Implementation of the N-dimensional advection equation with constant
+    velocity parameter. The equation is defined as follows:
+
+    .. math::
+
+        \frac{\partial u}{\partial t} + c \cdot \nabla u = 0
+
+    Here, :math:`c` is the advection velocity parameter.
+    """
+
+    def __init__(self, c):
+        """
+        Initialization of the :class:`AdvectionEquation`.
+
+        :param c: The advection velocity. If a scalar is provided, the same
+            velocity is applied to all spatial dimensions. If a list is
+            provided, it must contain one value per spatial dimension.
+        :type c: float | int | List[float] | List[int]
+        :raises ValueError: If ``c`` is an empty list.
+        """
+        # Check consistency
+        check_consistency(c, (float, int, list))
+        if isinstance(c, list):
+            all(check_consistency(ci, (float, int)) for ci in c)
+            if len(c) < 1:
+                raise ValueError("'c' cannot be an empty list.")
+        else:
+            c = [c]
+
+        # Store advection velocity parameter
+        self.c = torch.tensor(c).unsqueeze(0)
+
+        def equation(input_, output_):
+            """
+            Implementation of the advection equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :return: The residual of the advection equation.
+            :rtype: LabelTensor
+            :raises ValueError: If the ``input_`` labels do not contain the time
+                variable 't'.
+            :raises ValueError: If ``c`` is a list and its length is not
+                consistent with the number of spatial dimensions.
+            """
+            # Store labels
+            input_lbl = input_.labels
+            spatial_d = [di for di in input_lbl if di != "t"]
+
+            # Ensure time is passed as input
+            if "t" not in input_lbl:
+                raise ValueError(
+                    "The ``input_`` labels must contain the time 't' variable."
+                )
+
+            # Ensure consistency of c length
+            if len(self.c) != (len(input_lbl) - 1) and len(self.c) > 1:
+                raise ValueError(
+                    "If 'c' is passed as a list, its length must be equal to "
+                    "the number of spatial dimensions."
+                )
+
+            # Repeat c to ensure consistent shape for advection
+            self.c = self.c.repeat(output_.shape[0], 1)
+            if self.c.shape[1] != (len(input_lbl) - 1):
+                self.c = self.c.repeat(1, len(input_lbl) - 1)
+
+            # Add a dimension to c for the following operations
+            self.c = self.c.unsqueeze(-1)
+
+            # Compute the time derivative and the spatial gradient
+            time_der = grad(output_, input_, components=None, d="t")
+            grads = grad(output_=output_, input_=input_, d=spatial_d)
+
+            # Reshape and transpose
+            tmp = grads.reshape(*output_.shape, len(spatial_d))
+            tmp = tmp.transpose(-1, -2)
+
+            # Compute advection term
+            adv = (tmp * self.c).sum(dim=tmp.tensor.ndim - 2)
+
+            return time_der + adv
+
+        super().__init__(equation)
+
+
+class AllenCahnEquation(Equation):
+    r"""
+    Implementation of the N-dimensional Allen-Cahn equation, defined as follows:
+
+    .. math::
+
+        \frac{\partial u}{\partial t} - \alpha \Delta u + \beta(u^3 - u) = 0
+
+    Here, :math:`\alpha` and :math:`\beta` are parameters of the equation.
+    """
+
+    def __init__(self, alpha, beta):
+        """
+        Initialization of the :class:`AllenCahnEquation`.
+
+        :param alpha: The diffusion coefficient.
+        :type alpha: float | int
+        :param beta: The reaction coefficient.
+        :type beta: float | int
+        """
+        check_consistency(alpha, (float, int))
+        check_consistency(beta, (float, int))
+        self.alpha = alpha
+        self.beta = beta
+
+        def equation(input_, output_):
+            """
+            Implementation of the Allen-Cahn equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :return: The residual of the Allen-Cahn equation.
+            :rtype: LabelTensor
+            :raises ValueError: If the ``input_`` labels do not contain the time
+                variable 't'.
+            """
+            # Ensure time is passed as input
+            if "t" not in input_.labels:
+                raise ValueError(
+                    "The ``input_`` labels must contain the time 't' variable."
+                )
+
+            # Compute the time derivative and the spatial laplacian
+            u_t = grad(output_, input_, d=["t"])
+            u_xx = laplacian(
+                output_, input_, d=[di for di in input_.labels if di != "t"]
+            )
+
+            return u_t - self.alpha * u_xx + self.beta * (output_**3 - output_)
+
+        super().__init__(equation)
+
+
+class DiffusionReactionEquation(Equation):
+    r"""
+    Implementation of the N-dimensional Diffusion-Reaction equation,
+    defined as follows:
+
+    .. math::
+
+        \frac{\partial u}{\partial t} - \alpha \Delta u - f = 0
+
+    Here, :math:`\alpha` is a parameter of the equation, while :math:`f` is the
+    reaction term.
+    """
+
+    def __init__(self, alpha, forcing_term):
+        """
+        Initialization of the :class:`DiffusionReactionEquation`.
+
+        :param alpha: The diffusion coefficient.
+        :type alpha: float | int
+        :param Callable forcing_term: The forcing field function, taking as
+            input the points on which evaluation is required.
+        """
+        check_consistency(alpha, (float, int))
+        check_consistency(forcing_term, (Callable))
+        self.alpha = alpha
+        self.forcing_term = forcing_term
+
+        def equation(input_, output_):
+            """
+            Implementation of the Diffusion-Reaction equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :return: The residual of the Diffusion-Reaction equation.
+            :rtype: LabelTensor
+            :raises ValueError: If the ``input_`` labels do not contain the time
+                variable 't'.
+            """
+            # Ensure time is passed as input
+            if "t" not in input_.labels:
+                raise ValueError(
+                    "The ``input_`` labels must contain the time 't' variable."
+                )
+
+            # Compute the time derivative and the spatial laplacian
+            u_t = grad(output_, input_, d=["t"])
+            u_xx = laplacian(
+                output_, input_, d=[di for di in input_.labels if di != "t"]
+            )
+
+            return u_t - self.alpha * u_xx - self.forcing_term(input_)
+
+        super().__init__(equation)
+
+
+class HelmholtzEquation(Equation):
+    r"""
+    Implementation of the Helmholtz equation, defined as follows:
+
+    .. math::
+
+            \Delta u + k u - f = 0
+
+    Here, :math:`k` is a parameter of the equation, while :math:`f` is the
+    forcing term.
+    """
+
+    def __init__(self, k, forcing_term):
+        """
+        Initialization of the :class:`HelmholtzEquation` class.
+
+        :param k: The parameter of the equation.
+        :type k: float | int
+        :param Callable forcing_term: The forcing field function, taking as
+            input the points on which evaluation is required.
+        """
+        check_consistency(k, (int, float))
+        check_consistency(forcing_term, (Callable))
+        self.k = k
+        self.forcing_term = forcing_term
+
+        def equation(input_, output_):
+            """
+            Implementation of the Helmholtz equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :return: The residual of the Helmholtz equation.
+            :rtype: LabelTensor
+            """
+            lap = laplacian(output_, input_)
+            return lap + self.k * output_ - self.forcing_term(input_)
+
+        super().__init__(equation)
+
+
+class PoissonEquation(Equation):
+    r"""
+    Implementation of the Poisson equation, defined as follows:
+
+    .. math::
+
+            \Delta u - f = 0
+
+    Here, :math:`f` is the forcing term.
+    """
+
+    def __init__(self, forcing_term):
+        """
+        Initialization of the :class:`PoissonEquation` class.
+
+        :param Callable forcing_term: The forcing field function, taking as
+            input the points on which evaluation is required.
+        """
+        check_consistency(forcing_term, (Callable))
+        self.forcing_term = forcing_term
+
+        def equation(input_, output_):
+            """
+            Implementation of the Poisson equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :return: The residual of the Poisson equation.
+            :rtype: LabelTensor
+            """
+            lap = laplacian(output_, input_)
+            return lap - self.forcing_term(input_)
+
+        super().__init__(equation)