add burgers equation

Author: GiovanniCanali

docs/source/_rst/_code.rst

@@ -224,6 +224,7 @@ Equation Zoo
     Acoustic Wave Equation <equation/zoo/acoustic_wave_equation.rst>
     Advection Equation <equation/zoo/advection_equation.rst>
     Allen-Cahn Equation <equation/zoo/allen_cahn_equation.rst>
+    Burgers Equation <equation/zoo/burgers_equation.rst>
     Diffusion-Reaction Equation <equation/zoo/diffusion_reaction_equation.rst>
     Fixed Flux <equation/zoo/fixed_flux.rst>
     Fixed Gradient <equation/zoo/fixed_gradient.rst>

docs/source/_rst/equation/zoo/burgers_equation.rst

@@ -0,0 +1,7 @@
+Burgers Equation
+====================
+.. currentmodule:: pina.equation.zoo.burgers_equation
+
+.. automodule:: pina._src.equation.zoo.burgers_equation
+    :members:
+    :show-inheritance:

pina/_src/equation/zoo/burgers_equation.py

@@ -0,0 +1,84 @@
+"""Module for defining the Burgers equation."""
+
+from pina._src.core.operator import laplacian, grad
+from pina._src.core.utils import check_consistency
+from pina._src.equation.equation import Equation
+import torch
+
+
+class BurgersEquation(Equation):
+    r"""
+    Implementation of the N-dimensional Burgers equation, defined as follows:
+
+    .. math::
+
+        \frac{\partial u}{\partial t} + u \cdot \nabla u = \nu \Delta u
+
+    Here, :math:`\nu` is the viscosity coefficient.
+    """
+
+    def __init__(self, nu):
+        """
+        Initialization of the :class:`BurgersEquation` class.
+
+        :param nu: The viscosity coefficient.
+        :type nu: float | int
+        :raises ValueError: If ``nu`` is not a float or an int.
+        :raises ValueError: If ``nu`` is negative.
+        """
+        # Check consistency
+        check_consistency(nu, (float, int))
+        if nu < 0:
+            raise ValueError(
+                "The viscosity ``nu`` must be a positive float or int."
+            )
+
+        # Store viscosity coefficient
+        self.nu = nu
+
+        def equation(input_, output_):
+            """
+            Implementation of the Burgers equation.
+
+            :param LabelTensor input_: The input data of the problem.
+            :param LabelTensor output_: The output data of the problem.
+            :raises ValueError: If the number of output components does not
+                match the number of spatial dimensions.
+            :raises ValueError: If the ``input_`` labels do not contain the time
+                variable 't'.
+            :return: The residual of the Burgers equation.
+            :rtype: LabelTensor
+            """
+            # Store labels
+            spatial_d = [di for di in input_.labels if di != "t"]
+
+            # Ensure consistency between output and spatial dimensions
+            if len(output_.labels) != len(spatial_d):
+                raise ValueError(
+                    f"The number of output components must match the number of "
+                    f"spatial dimensions. Got {len(output_.labels)} and "
+                    f"{len(spatial_d)}."
+                )
+
+            # 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 differential terms
+            u_t = grad(output_, input_, d=["t"])
+            u_x = grad(output_, input_, d=spatial_d)
+            u_xx = laplacian(output_, input_, d=spatial_d)
+
+            # Compute the convective term componentwise
+            convection = torch.zeros_like(output_)
+            for i, c in enumerate(output_.labels):
+                convection[:, i] = sum(
+                    output_[output_.labels[j]] * u_x[f"d{c}d{spatial_d[j]}"]
+                    for j in range(len(spatial_d))
+                ).reshape(-1)
+
+            return u_t + convection - self.nu * u_xx
+
+        super().__init__(equation)

pina/equation/__init__.py

@@ -36,6 +36,7 @@
     "AdvectionEquation": ".zoo",
     "AllenCahnEquation": ".zoo",
     "DiffusionReactionEquation": ".zoo",
+    "BurgersEquation": ".zoo",
 }
 
 

pina/equation/zoo.py

@@ -12,6 +12,7 @@
     "Laplace",
     "PoissonEquation",
     "AcousticWaveEquation",
+    "BurgersEquation",
 ]
 
 from pina._src.equation.zoo.acoustic_wave_equation import AcousticWaveEquation
@@ -22,6 +23,7 @@
 )
 from pina._src.equation.zoo.helmholtz_equation import HelmholtzEquation
 from pina._src.equation.zoo.poisson_equation import PoissonEquation
+from pina._src.equation.zoo.burgers_equation import BurgersEquation
 from pina._src.equation.zoo.fixed_value import FixedValue
 from pina._src.equation.zoo.fixed_gradient import FixedGradient
 from pina._src.equation.zoo.fixed_flux import FixedFlux

tests/test_equation_zoo/test_burgers_equation.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+from pina import LabelTensor
+from pina.equation.zoo import BurgersEquation
+
+
+# Define input and output values
+pts = LabelTensor(torch.rand(10, 3, requires_grad=True), labels=["x", "y", "t"])
+u = torch.sin(pts["x", "y"]) * torch.cos(pts["y", "t"])
+u.labels = ["u", "v"]
+
+
+@pytest.mark.parametrize("nu", [0, 1, 2.5])
+def test_burgers_equation(nu):
+
+    # Constructor
+    equation = BurgersEquation(nu=nu)
+
+    # Should fail if nu is not a float or int
+    with pytest.raises(ValueError):
+        BurgersEquation(nu="invalid")
+
+    # Should fail if nu is negative
+    with pytest.raises(ValueError):
+        BurgersEquation(nu=-1)
+
+    # Residual
+    residual = equation.residual(pts, u)
+    assert residual.shape == u.shape
+
+    # Should fail if the input has no 't' label
+    with pytest.raises(ValueError):
+        residual = equation.residual(pts["x", "y"], u)
+
+    # Should fail if output and spatial dimensions do not match
+    with pytest.raises(ValueError):
+        residual = equation.residual(pts, u["u"])