Notes on Physics-informed Neural Network (PINN)

Examples of Physics-informed Neural Network (PINN)¶

I have been working on a problem that involve solving multi-variate differential equations. As it turns out, recent developments in deep learning techniques revealed a fascinating connection between (deep) neural networks and differential equations.

It was one of the few occasions when my problems could be solved by recently developed numerical techniques. I have learned a lot during the endeavour, and this notebook shows two examples of making use of the physics-informed neural network (PINN) to solve partial differential equations (PDEs). I believe these two examples show the basic principles of implementing PINNs, and I have devised these examples by simplifying the problems I was working on down to very basic blocks.

The main problem I encountered with PINNs was that imposing boundary conditions (and initial conditions, for the same reason) adversely affect the non-linear composition of functions, as the network fails to properly minimize the loss function. This can happen for a number of reasons, but unfortunately no simple solution exists at this moment.

The methods used to derive the solutions for the PDEs shown in this notebook are largely based on the following works (sorted by publication date).

Raissi, M., and G. E. Karniadakis, 2018: Hidden physics models: Machine learning of nonlinear partial differential equations. Journal of Computational Physics, 357, 125–141. Raissi & Karniadakis (2018)
Raissi, M., P. Perdikaris, and G.E. Karniadakis, 2019: Physics-Informed Neural Networks: A Deep Learning Framework for Solving Forward and Inverse Problems Involving Nonlinear Partial Differential Equations. Journal of Computational Physics, 378, 686–707. Raissi et al. (2019)
Cai, S., Z. Wang, S. Wang, P. Perdikaris, and G. E. Karniadakis, 2021: Physics-Informed Neural Networks for Heat Transfer Problems. Journal of Heat Transfer, 143(6). Cai et al. (2021)
Sukumar, N., and A. Srivastava, 2022: Exact imposition of boundary conditions with distance functions in physics-informed deep neural networks. Computer Methods in Applied Mechanics and Engineering, 389, 114333. Sukumar & Srivastava (2022)

1D Heat Conduction with Neumann Boundary Conditions¶

We will start with the classical one-dimensional heat transfer problem, governed by the heat equation, which is a simple parabolic PDE.

\frac{\partial u}{\partial t} = \alpha \nabla u

(1)

where $u$ is temperature in C, $\alpha$ is thermal diffusivity (set to 0.1 for this example). The laplacian $\nabla$ is only defined along the x-axis. The values are defined between the range $[0, 1]$ .

Think of a very thin pipe with a unit length. At the beginning of the simulation, half of the pipe has been heated up to 1 C, while the rest of the pipe is at 0 C. The two edges of the pipe is well insulated, and therefore is subject to Neumann boundary conditions.

\frac{\partial u_{x = 0}}{\partial x} = \frac{\partial u_{x = 1}}{\partial x} = 0

(2)

If the simulation holds, we will see the right half of the pipe slowly becoming warmer over time. Because of the Neumann boundary conditions, there should be no heat loss.

import jax.numpy as jnp

import torch
import torch.nn as nn

import numpy as np

import plotly
import plotly.express as px
import plotly.graph_objects as go

from IPython.display import display, HTML

plotly.offline.init_notebook_mode()
display(HTML(
    '<script type="text/javascript" async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.1/MathJax.js?config=TeX-MML-AM_SVG"></script>'
))

Let’s first define the network. We will use Pytorch here. The following model is more or less standard Pytorch deep neural network that uses $\tanh$ as the activation function and takes in a list of (linear) layers.

class PINN(nn.Module):
    def __init__(self, layers):
        super(PINN, self).__init__()
        self.activation = nn.Tanh()

        layer_list = []
        for i in range(len(layers) - 1):
            layer_list.append(nn.Linear(layers[i], layers[i+1]))
        self.layers = nn.ModuleList(layer_list)

    def forward(self, x):
        for _, layer in enumerate(self.layers[:-1]):
            x = self.activation(layer(x))
        return self.layers[-1](x)

The heart of the PINN is in its loss function. This is largely made possible thanks to Pytorch’s autograd function. Here, we define the physics loss function based on the heat equation shown above, and add the losses from the initial and boundary conditions.

The loss function really determines if the PINN will be able to find the solution, as we are trying to satisfy both the PDE and the boundary conditions without any data. This portion of the PINN closely resembles supervised learning; if enough guidance can be given in the solution space, the network will eventually find the solution that fully satisfies the PDE and the boundary conditions. For more complex multi-variate PDEs, this becomes the main problem that needs to be resolved.

def phys_loss(model, t, x):
    def u(t, x):
        return model(torch.cat((t, x), dim=1))

    def _gr(_o, _i):
        return torch.autograd.grad(
            outputs = _o,
            inputs = _i,
            grad_outputs=torch.ones_like(_o),
            create_graph=True,
            retain_graph=True
        )[0]

    u_x = _gr(u(t, x), x)
    u_t = _gr(u(t, x), t)
    u_xx = _gr(u_x, x)

    loss_p = nn.MSELoss()(
        u_t,
        0.1 * u_xx
    )

    x_0 = torch.zeros(100, requires_grad=True).reshape(-1, 1)
    x_1 = torch.ones(100, requires_grad=True).reshape(-1, 1)
    t_0 = torch.zeros(100, requires_grad=True).reshape(-1, 1)

    u_0 = torch.zeros(100)
    u_0[:50] = 5
    u_0 = u_0.reshape(-1, 1)

    # Initial condition
    loss_ic = nn.MSELoss()(
        u(t_0, x),
        u_0
    )

    # Boundary condition
    loss_bd = nn.MSELoss()(
        _gr(u(t, x_0), x_0),
        torch.zeros(100).reshape(-1, 1)
    )

    loss_bd += nn.MSELoss()(
        _gr(u(t, x_1), x_1),
        torch.zeros(100).reshape(-1, 1)
    )

    return loss_p + loss_ic + loss_bd

Here we define our PINN model, which has one input layer (that takes in a $(t, x)$ pair), three hidden layers with 50 neurons, and an output layer that outputs $u(t, x)$ .

pinn = PINN([2, 50, 50, 50, 50, 1])

# Co-location points
t = torch.linspace(0, 1, 100, requires_grad=True).reshape(-1, 1)
x = torch.linspace(0, 1, 100, requires_grad=True).reshape(-1, 1)

def u(t, x):
    u = pinn(torch.stack((t, x), dim=1))
    return u

print(pinn)

PINN(
  (activation): Tanh()
  (layers): ModuleList(
    (0): Linear(in_features=2, out_features=50, bias=True)
    (1-3): 3 x Linear(in_features=50, out_features=50, bias=True)
    (4): Linear(in_features=50, out_features=1, bias=True)
  )
)

optimizer = torch.optim.Adam(pinn.parameters(), lr=1e-4)
n_epoch = int(6e3)

for epoch in range(n_epoch):
    optimizer.zero_grad()

    loss = phys_loss(pinn, t, x)

    loss.backward()
    optimizer.step()

    if epoch % 500 == 0:
        print(f"Epoch {epoch}: Total Loss = {loss.item():.4f}")

print("Training complete!")

Epoch 0: Total Loss = 13.0334
Epoch 500: Total Loss = 5.9018
Epoch 1000: Total Loss = 4.5405
Epoch 1500: Total Loss = 0.8725
Epoch 2000: Total Loss = 0.6965
Epoch 2500: Total Loss = 0.4473
Epoch 3000: Total Loss = 0.2364
Epoch 3500: Total Loss = 0.1446
Epoch 4000: Total Loss = 0.1077
Epoch 4500: Total Loss = 0.0880
Epoch 5000: Total Loss = 0.0749
Epoch 5500: Total Loss = 0.0654
Training complete!

t_test = torch.linspace(0, 1, 100)
x_test = torch.linspace(0, 1, 100)
T, X = torch.meshgrid(t_test, x_test, indexing='ij')

t_flat = T.flatten()
x_flat = X.flatten()

with torch.no_grad():
    u_pred = pinn(torch.stack((t_flat, x_flat), dim=1)).detach()

U = u_pred.view(100, 100).numpy()

fig = px.imshow(
    U,
    x=x_test,
    y=t_test,
    color_continuous_scale="viridis",
    origin="lower",
    aspect="auto"
)

fig.update_layout(
        title_text="1D Heat Conduction",
        xaxis_title_text="x",
        yaxis_title_text="t",
        coloraxis_colorbar_title_text = "u(t, x)",
        autosize=False,
        height=540,
        width=720
    )

fig.show()

The resulting plot on the $(x, t)$ domain looks good. Over time (as you go up on the y-axis), the heat from the left-hand side of the 1D pipe slowly dissipates into the right-hand side, as expected. We did not quite reach the equilibrium here, but I stopped early for the sake of visualization; here, you can clearly see the abrupt boundary between the warm side and the cold side of the pipe moving towards the right edge.

Solving Non-linear Boundary Value Problem with PINN¶

Let’s try a slightly different formulation. This time, consider a one-dimensional elliptical PDE, defined on $[0, 2 \pi]$ and subject to the boundary condition $u(0) = u(1) = 0$ .

\frac{\partial^2 u}{\partial x^2} = f(x)

(3)

This is only slightly different from the first example; now the pipe is temperature-controlled at its edges and is being heated up by the amount defined by $f(x)$ .

Also, to make this a non-linear problem, we will make $u$ a function of x.

\frac{\partial }{\partial x} \left( k(x) \frac{\partial u}{\partial x} \right) = f(x)

(4)

Still, we do not want to make this example too convoluted. Consider the case where $k(x) = x$ and $u(x) = \sin^2(x)$ . Then we can re-write the equation above.

\frac{\partial }{\partial x} \left( k(x) \frac{\partial u}{\partial x} \right) = \sin{(2 x)} + 2 \cos{(2 x)}

(5)

# Co-location points
xn = 100
x = torch.linspace(0, 2 * jnp.pi, xn, requires_grad=True)

layers = [1] + [xn for _ in range(5)] + [1]

pinn = PINN(layers)

print(pinn)

PINN(
  (activation): Tanh()
  (layers): ModuleList(
    (0): Linear(in_features=1, out_features=100, bias=True)
    (1-4): 4 x Linear(in_features=100, out_features=100, bias=True)
    (5): Linear(in_features=100, out_features=1, bias=True)
  )
)

def p_loss(model, x):
    def u(x):
        return model(x)
    
    def f_true(x):
        return torch.sin(2.0 * x) + 2.0 * x * torch.cos(2 * x)

    def _gr(m, i):
        return torch.autograd.grad(
            outputs = m,
            inputs = i,
            grad_outputs=torch.ones_like(i),
            create_graph=True
        )[0]

    u_x = _gr(u(x), x)
    u_xx = _gr(u_x * x, x)

    loss_p = nn.MSELoss()(
        u_xx,
        f_true(x)
    )

    x_0 = torch.zeros(100, requires_grad=True)

    # Boundary condition
    loss_bd = nn.MSELoss()(
        u(x_0[:, None]),
        f_true(x_0)
    )

    return loss_p + loss_bd

optimizer = torch.optim.Adam(pinn.parameters(), lr=1e-3)
n_epoch = int(4e3)

for epoch in range(n_epoch):
    optimizer.zero_grad()

    loss = p_loss(pinn, x[:, None])

    loss.backward()
    optimizer.step()

    if epoch % 499 == 0:
        print(f"Epoch {epoch}: Total Loss = {loss.item():.6f}")

print("Training complete!")

Epoch 0: Total Loss = 27.105961

/Users/lorenoh/Workspace/repos/latent-ice/.venv/lib/python3.13/site-packages/torch/nn/modules/loss.py:610: UserWarning:

Using a target size (torch.Size([100])) that is different to the input size (torch.Size([100, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.

Epoch 499: Total Loss = 0.019959
Epoch 998: Total Loss = 0.006125
Epoch 1497: Total Loss = 0.002427
Epoch 1996: Total Loss = 0.018449
Epoch 2495: Total Loss = 0.311358
Epoch 2994: Total Loss = 0.144546
Epoch 3493: Total Loss = 0.088150
Epoch 3992: Total Loss = 0.004968
Training complete!

with torch.no_grad():
    def u_true(x):
        u = np.sin(x)**2
        return u
    
    ix = torch.FloatTensor(x)
    u = pinn(ix.reshape(xn, 1))[:, 0]

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=ix,
            y=u_true(ix),
            name="True",
            line=dict(color='royalblue', width=6, dash='dot')
        )
    )

    fig.add_trace(
        go.Scatter(
            x=ix,
            y=u,
            name="PINN",
            line=dict(color='firebrick', width=2)
        )
    )

    fig.update_layout(
        title_text=r"$\text{Solution for } u(x) = \sin{(2 x)}$",
        xaxis_title_text="x",
        yaxis_title_text="u(x)",
        bargap=0.05,
        bargroupgap=0,
    )

    fig.show()

def _gr(m, i):
    return torch.autograd.grad(
        outputs = m,
        inputs = i,
        grad_outputs=torch.ones_like(i),
        create_graph=True,
        allow_unused=True
    )[0]

ix = torch.FloatTensor(x)
u_x = _gr(pinn(ix.reshape(xn, 1))[:, 0], ix)
u_xx = _gr(u_x * ix, ix)

with torch.no_grad():
    def f_true(x):
        return torch.sin(2.0 * x) + 2.0 * x * torch.cos(2 * x)

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=ix,
            y=f_true(ix),
            name="True",
            line=dict(color='royalblue', width=6, dash='dot')
        )
    )

    fig.add_trace(
        go.Scatter(
            x=ix,
            y=u_xx,
            name="PINN",
            line=dict(color='firebrick', width=2)
        )
    )

    fig.update_layout(
        title_text=r"$f(x) = \sin{(2 x)} + 2 \cos{(2 x)}$",
        xaxis_title_text="x",
        yaxis_title_text="f(x)",
        bargap=0.05,
        bargroupgap=0,
    )

    fig.show()

The PINN-computed distributions of $u(t, x)$ and $f(t, x)$ are indistinguishable from the corresponding analytical solutions. There is a more rigorous version of this PDE in the original paper, but I have also found that with a bit of careful training, that also converges nicely. This is mainly because the problem is well-defined, and the solution space is narrow enough for the PINN.

Remarks¶

If you have read this far, congratulations. You might get an idea that PINNs can replace traditional PDE solvers, or that PINNs might outperform them. That was what the initial papers seemed to promise, at least. However, One should always remember that the purpose of research papers and blog posts is to show the best-case scenario; you should look no further than the following article.

McGreivy, N., and A. Hakim, 2024: Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelligence, 6(10), 1256–1269. McGreivy & Hakim (2024)

Traditional PDE solvers have become quite computationally efficient and are mathematically sound; if one needs mathematical accuracy, PINNs are not a suitable choice. In my experience, they work well as a fast estimator of multi-variate PDEs involving many (well-behaved) parameters, which was, fortunately for me, the problem I was trying to solve. They also work well as a surrogate model, but the accuracy of their predictions depend largely on the quality of their loss functions. This is not surprising, though, because in PINNs, the loss function behaves much like a guidance in supervised learning.

PINNs seem to have an inherent, fundamental dilemma: if we know so much about the nature of the PDE to come up with a physical loss function that can guide the DINN to converge nicely, we know enough to construct traditional PDE solvers, which will be faster, more computationally efficient, and more accurate. If we do not know so much about the problem, PINNs often fail to converge, as the network is not given enough information to find proper parameters to minimize the loss function. This is the same old problem we have had with supervised learning.

Still, I found them to be very useful for certain problems, especially when fast estimators were needed over a very large parameter space. It has a lot of issues that needs to be resolved, but I am interested in keeping an eye on what comes next for PINNs.

References¶

Raissi, M., & Karniadakis, G. E. (2018). Hidden physics models: Machine learning of nonlinear partial differential equations. Journal of Computational Physics, 357, 125–141. 10.1016/j.jcp.2017.11.039
Raissi, M., Perdikaris, P., & Karniadakis, G. E. (2019). Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378, 686–707. 10.1016/j.jcp.2018.10.045
Cai, S., Wang, Z., Wang, S., Perdikaris, P., & Karniadakis, G. E. (2021). Physics-Informed Neural Networks for Heat Transfer Problems. Journal of Heat Transfer, 143(6). 10.1115/1.4050542
Sukumar, N., & Srivastava, A. (2022). Exact imposition of boundary conditions with distance functions in physics-informed deep neural networks. Computer Methods in Applied Mechanics and Engineering, 389, 114333. 10.1016/j.cma.2021.114333
McGreivy, N., & Hakim, A. (2024). Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelligence, 6(10), 1256–1269. 10.1038/s42256-024-00897-5

notebooks

Async HTTP Request Example