tutorial.py

#!/usr/bin/env python
# coding: utf-8

# # Tutorial: Introduction to PINA `Equation` class
# 
# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PINA-org/PINA-tutorials/blob/main/tutorials/tutorial12/tutorial.ipynb)
# 
# 
# 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.

# In[1]:


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

    IN_COLAB = True
except:
    IN_COLAB = False
if IN_COLAB:
    get_ipython().system('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


# 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.

# In[2]:


# 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](https://PINA-org.github.io/PINA-Tutorials/tutorial16/tutorial.html) to learn how to build a Problem from scratch.**

# In[3]:


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$:
# 
# ```python
# 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:

# In[4]:


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}$:

# In[5]:


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](https://PINA-org.github.io/PINA/).