Tutorial: Introduction to PINA Equation class

In this tutorial, we will explore how to use the Equation class in PINA. We will focus on how to leverage this class, along with its inherited subclasses, to enforce residual minimization in Physics-Informed Solvers.

By the end of this guide, you'll understand how to integrate physical laws and constraints directly into your model training, ensuring that the solution adheres to the underlying differential equations.

Example: The Burgers 1D equation

We will start implementing the viscous Burgers 1D problem Class, described as follows:

$$ \begin{equation} \begin{cases} \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} &= \nu \frac{\partial^2 u}{ \partial x^2}, \quad x\in(0,1), \quad t>0\\ u(x,0) &= -\sin (\pi x), \quad x\in(0,1)\\ u(x,t) &= 0, \quad x = \pm 1, \quad t>0\\ \end{cases} \end{equation} $$

where we set $ \nu = \frac{0.01}{\pi}$.

In the class that models this problem we will see in action the Equation class and one of its inherited classes, the FixedValue class.

## routine needed to run the notebook on Google Colab
try:
    import google.colab

    IN_COLAB = True
except:
    IN_COLAB = False
if IN_COLAB:
    !pip install "pina-mathlab[tutorial]"

import torch

# useful imports
from pina import Condition
from pina.problem import SpatialProblem, TimeDependentProblem
from pina.equation import Equation
from pina.equation.zoo import FixedValue
from pina.domain import CartesianDomain
from pina.operator import grad, laplacian

/opt/hostedtoolcache/Python/3.10.22/x64/lib/python3.10/site-packages/torch/jit/_script.py:1491: FutureWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(

Let's begin by defining the Burgers equation and its initial condition as Python functions. These functions will take the model's input_ (spatial and temporal coordinates) and output_ (predicted solution) as arguments. The goal is to compute the residuals for the Burgers equation, which we will minimize during training.

# define the burgers equation
def burgers_equation(input_, output_):
    du = grad(output_, input_)
    ddu = laplacian(output_, input_, components="x")
    return du["dudt"] + output_["u"] * du["dudx"] - (0.01 / torch.pi) * ddu


# define initial condition
def initial_condition(input_, output_):
    u_expected = -torch.sin(torch.pi * input_["x"])
    return output_["u"] - u_expected

Above we use the grad operator from pina.operator to compute the gradient. In PINA each differential operator takes the following inputs:

• output_: A tensor on which the operator is applied.
• input_: A tensor with respect to which the operator is computed.
• components: The names of the output variables for which the operator is evaluated.
• d: The names of the variables with respect to which the operator is computed.

Each differential operator has its fast version, which performs no internal checks on input and output tensors. For these methods, the user is always required to specify both components and d as lists of strings.

Let's define now the problem!

👉 Do you want to learn more on Problems? Check the dedicated tutorial to learn how to build a Problem from scratch.

class Burgers1D(TimeDependentProblem, SpatialProblem):

    # assign output/ spatial and temporal variables
    output_variables = ["u"]
    spatial_domain = CartesianDomain({"x": [-1, 1]})
    temporal_domain = CartesianDomain({"t": [0, 1]})

    domains = {
        "bound_cond": spatial_domain.partial().update(temporal_domain),
        "time_cond": spatial_domain.update(CartesianDomain({"t": 0.0})),
        "phys_cond": spatial_domain.update(temporal_domain),
    }
    # problem condition statement
    conditions = {
        "bound_cond": Condition(domain="bound_cond", equation=FixedValue(0.0)),
        "time_cond": Condition(
            domain="time_cond", equation=Equation(initial_condition)
        ),
        "phys_cond": Condition(
            domain="phys_cond", equation=Equation(burgers_equation)
        ),
    }

The Equation class takes as input a function that computes the residual of an equation, such as a PDE. In the example above, this is done twice, using initial_condition and burgers_equation. Once defined, each Equation object is passed to a specific Condition inside the problem class. When multiple residual equations need to be enforced on the same subdomain, they can be grouped using SystemEquation, which takes a list of equations defined in the problem.

For common use cases, PINA also provides several predefined equation classes. The FixedValue, FixedGradient, FixedFlux, and FixedLaplacian classes can be used to enforce fixed constraints on the solution or on its derivatives. In particular, FixedValue enforces a constant value with the same dimensions as the output function, making it suitable for conditions such as Dirichlet boundary conditions. Similarly, FixedGradient and FixedFlux enforce fixed values on the gradient or flux of the solution, respectively, while FixedLaplacian can be used to impose a fixed value on the Laplacian of the solution. Many ready-to-use equations are also available in the pina.equation.zoo module.

Once the equations are assigned to the problem conditions, the physics-informed solver aims to minimize the residuals defined by each equation during the training phase.

Defining a new Equation class

Equation classes can also be inherited to define a new class. For example, we can define a new class Burgers1D to represent the Burgers equation. During the class call, we can pass the viscosity parameter $\nu$:

class Burgers1D(Equation):
    def __init__(self, nu):
        self.nu = nu

    def equation(self, input_, output_):
        ...

In this case, the Burgers1D class will inherit from the Equation class and compute the residual of the Burgers equation. The viscosity parameter $\nu$ is passed when instantiating the class and used in the residual calculation. Let's see it in more details:

class Burgers1DEquation(Equation):

    def __init__(self, nu=0.0):
        """
        Burgers1D class. This class can be
        used to enforce the solution u to solve the viscous Burgers 1D Equation.

        :param torch.float32 nu: the viscosity coefficient. Default value is set to 0.
        """
        self.nu = nu

        def equation(input_, output_):
            return (
                grad(output_, input_, d="t")
                + output_ * grad(output_, input_, d="x")
                - self.nu * laplacian(output_, input_, d="x")
            )

        super().__init__(equation)

Now we can just pass the above class as input for the last condition, setting $\nu= \frac{0.01}{\pi}$:

class Burgers1D(TimeDependentProblem, SpatialProblem):

    # define initial condition
    def initial_condition(input_, output_):
        u_expected = -torch.sin(torch.pi * input_.extract(["x"]))
        return output_.extract(["u"]) - u_expected

    # assign output/ spatial and temporal variables
    output_variables = ["u"]
    spatial_domain = CartesianDomain({"x": [-1, 1]})
    temporal_domain = CartesianDomain({"t": [0, 1]})

    domains = {
        "bound_cond": spatial_domain.partial().update(temporal_domain),
        "time_cond": spatial_domain.update(CartesianDomain({"t": 0.0})),
        "phys_cond": spatial_domain.update(temporal_domain),
    }
    # problem condition statement
    conditions = {
        "bound_cond": Condition(domain="bound_cond", equation=FixedValue(0.0)),
        "time_cond": Condition(
            domain="time_cond", equation=Equation(initial_condition)
        ),
        "phys_cond": Condition(
            domain="phys_cond", equation=Burgers1DEquation(nu=0.01 / torch.pi)
        ),
    }

What's Next?

Congratulations on completing the Equation class tutorial of PINA! As we've seen, you can build new classes that inherit from Equation to store more complex equations, such as the 1D Burgers equation, by simply passing the characteristic coefficients of the problem.

From here, you can:

• Define Additional Complex Equation Classes: Create your own equation classes, such as SchrodingerEquation, NavierStokesEquation, etc.
• Define More FixedOperator Classes: Implement operators like FixedCurl for more advanced simulations.
• Integrate Custom Equations and Operators: Combine your custom equations and operators into larger systems for more complex simulations.
• and many more!: Explore for example different residual minimization techniques to improve the performance and accuracy of your models.

For more resources and tutorials, check out the PINA Documentation.