Introduction

This blog covers Physics Informed Neural Networks, which are basically neural networks that are constrained to obey physical constraints. The blog broadly discusses two directions:

In particular, it deals with two distinct types of algorithms, namely

  • A new family of data-efficient spatio-temporal function approximators
  • Arbitrary accurate RK time steppers with potentially unlimited number of stages

We know that deep neural networks are powerful function approximators. They can approximate a nonlinear map from a few – potentially very high dimensional – input and output pairs of data. However, the approximation does not take into account any underlying physical constraints imposed by fundamental principles (such as conservation laws) which are
often given in terms of PDEs. PINNs precisely solve this problem by regularizing the loss functions with these constraints. The resulting approximation enjoys both the power of deep neural nets as function approximators and obeying the underlying physical conservation laws.

PDEs

Let us consider a parametric nonlinear PDE

ut+N[u;λ]=0,x∈Ω⊂RD, t∈[0,T];(⋅)t:=∂(⋅)∂t \color{BrickRed} u_{t}+\mathcal{N}[u ; \lambda]=0, x \in \Omega \subset \mathbb{R}^D,\, t \in[0, T];\quad ({\cdot})_t := \frac{\partial ({\cdot})}{\partial t}
where u(t,x)u(t, x) denotes the latent (hidden) solution and N[⋅;λ]\mathcal{N}[\cdot ; \lambda] is a nonlinear operator parametrized by λ\lambda. This encapsulation covers a wide range of PDEs in math, physics, atmospheric sciences, biology, finance, including conservation laws, diffusion, reac-diff-advec. PDEs, kinetics etc. The burger’s equation in 22d is an apt starting step to investigate if indeed PINNs can efficiently solve PDEs, and therefore the authors begin by applying PINNs to it.

N[u;λ]=λ1uux−λ2uxx and λ=(λ1,λ2);(⋅)x:=∂(⋅)∂x(⋅)xx:=∂2(⋅)∂x2 \mathcal{N}[u ; \lambda]=\lambda_{1} u u_{x}-\lambda_{2} u_{x x} \text { and } \lambda=\left(\lambda_{1}, \lambda_{2}\right);\quad ({\cdot})_x := \frac{\partial ({\cdot})}{\partial x} \quad ({\cdot})_{xx} := \frac{\partial^2 ({\cdot})}{\partial x^2}

The two directions highlighted above can now be restated in context of Burger’s equation as follows:

  • Given λ\lambda what is u(t,x)u(t, x) (data-driven solutions of PDEs)
  • Find λ\lambda that best describes observations u(ti,xj)u({t_i, x_j}) (data-driven discovery of PDEs)

PDEs + PINNs

First, let us rewrite the PDE as f(u;t,x)=0f(u; t,x) = 0
f(u;t,x)≐ut+N[u],along withu=uθ(t,x) f(u; t,x) \doteq u_{t}+\mathcal{N}[u], \quad \text{along with}\quad u = u_\theta ({t, x})

The next step consists of appending the loss function used for the training process (finding the parameters θ\theta) as follows:

L=Lu+Lf \mathcal{L}=\mathcal{L}_{u}+\mathcal{L}_{f}

where

Lu=1Nu∑i=1Nu∣u(tui,xui)−ui∣2 ;Lf=1Nf∑i=1Nf∣f(tfi,xfi)∣2. \mathcal{L}_{u}=\frac{1}{N_{u}} \sum_{i=1}^{N_{u}}\left|u\left(t_{u}^{i}, x_{u}^{i}\right)-u^{i}\right|^{2}\, ;\quad \mathcal{L}_{f}=\frac{1}{N_{f}} \sum_{i=1}^{N_{f}}\left|f\left(t_{f}^{i}, x_{f}^{i}\right)\right|^{2}.

Here Lu\mathcal{L}_{u} consists of the approximation error at the boundary of the domain (initial and boundary conditions), and Lf\mathcal{L}_{f} denotes the error incurred inside the domain. Also, {tfi,xfi}i=1Nf\left\{t_{f}^{i}, x_{f}^{i}\right\}_{i=1}^{N_{f}} specify the collocation points — the discrete points at which the physical constraints are imposed through a regularized (penalty) formulation. Another way to look at this is that Lu\mathcal{L}_u helps to enforce initial and boundary data accurately, while Lf\mathcal{L}_f imposes the structure of the PDE into the training process. An important thing to consider is the relative weighting of the two losses, namely,
L=λ Lu+(2−λ) Lf \mathcal{L}=\lambda\,\mathcal{L}_{u}+(2-\lambda)\,\mathcal{L}_{f}
where λ=1\lambda = 1 corresponds to the original loss above. Increasing λ\lambda towards 22 leads to a an approximation that is accurate inside the domain but performs poorly on the boundaries. Similarly, λ→0+\lambda\rightarrow 0^+ leads to an approximation that is accurate on the boundary but poor inside the domain.

Implementation details and Code

Since most of the examples considered here involve a small number of training data points, the authors choose to exploit quasi-second order optimization methods such as L-BFGS. They do not consider mini-batching for the same reason and consider the full batch on the update. There are no theoretical guarantees on the existence of the minimizer but the authors observe empirically that as long as the PDE is well-posed the optimization algorithm converges to the correct solution.

The original implementation can be accessed at https://github.com/maziarraissi/PINNs which builds on top of Tensorflow. Corresponding PyTorch and Julia (Flux) implementations can be accessed at https://github.com/idrl-lab/idrlnet and https://neuralpde.sciml.ai/dev/

Examples: Schrodinger Equation

As a first example, consider a complex-valued differential equation (h(t,x)=u(t,x)+i v(t,x)h(t, x) = u(t, x) + \mathrm{i}\, v(t,x)), namely the Schrodinger equation,

f≐ iht+0.5hxx+∣h∣2h=0,x∈[−5,5],t∈[0,π/2]h(0,x)=2sech⁡(x)h(t,−5)=h(t,5)hx(t,−5)=hx(t,5) \begin{aligned} f\doteq \, &i h_{t}+0.5 h_{x x}+|h|^{2} h=0, \quad x \in[-5,5], \quad t \in[0, \pi / 2] \\ &h(0, x)=2 \operatorname{sech}(x) \\ &h(t,-5)=h(t, 5) \\ &h_{x}(t,-5)=h_{x}(t, 5) \end{aligned}

In order to generate the training data, the authors employ a classical pseudo-spectral solver in space (xx) coupled with a high-order RK time stepper.
The results are shown in the figure below (Fig. 1 in the paper). PINN provides an accurate solution to Schrodinger equation and can handle periodic boundary conditions, nonlinearities and complex valued nature of the PDE efficiently.

Logo

Although PINN provides a good approximation to the solution above, training the network requires a large number of data points. This is where the adaptive time-stepping using RK methods comes into the picture. The authors propose an adaptive time-stepping with a neural net at each time-step. This significantly improves the approximation quality, and allows one to take much larger time-steps compared to traditional solvers.

Examples: Allen-Cahn Equation

In order to test the adaptive time-stepping scheme, the authors next take a look at the Allen-Cahn equation.

ut−0.0001uxx+5u3−5u=0,x∈[−1,1],t∈[0,1],u(0,x)=x2cos⁡(πx),u(t,−1)=u(t,1),ux(t,−1)=ux(t,1). \begin{aligned} &u_{t}-0.0001 u_{x x}+5 u^{3}-5 u=0, \quad x \in[-1,1], \quad t \in[0,1], \\ &u(0, x)=x^{2} \cos (\pi x), \\ &u(t,-1)=u(t, 1), \\ &u_{x}(t,-1)=u_{x}(t, 1) . \end{aligned}

The results show excellent agreement between the predicted and exact (numerical) solutions. The only difference here is that the authors consider the sum-of-squared errors as the loss function for training, instead of the MSE used before, i.e.

SSEn=∑j=1q+1∑i=1Nn∣ujn(xn,i)−un,i∣2 S S E_{n}=\sum_{j=1}^{q+1} \sum_{i=1}^{N_{n}}\left|u_{j}^{n}\left(x^{n, i}\right)-u^{n, i}\right|^{2}
SSEb=∑i=1q∣un+ci(−1)−un+ci(1)∣2+∣un+1(−1)−un+1(1)∣2+∑i=1q∣uxn+ci(−1)−uxn+ci(1)∣2+∣uxn+1(−1)−uxn+1(1)∣2 \begin{aligned} S S E_{b} &= \sum_{i=1}^{q}\left|u^{n+c_{i}}(-1)-u^{n+c_{i}}(1)\right|^{2}+\left|u^{n+1}(-1)-u^{n+1}(1)\right|^{2} \\ &+ \sum_{i=1}^{q}\left|u_{x}^{n+c_{i}}(-1)-u_{x}^{n+c_{i}}(1)\right|^{2}+\left|u_{x}^{n+1}(-1)-u_{x}^{n+1}(1)\right|^{2} \end{aligned}

Logo

Examples: Navier-Stokes Equation

In this section, we take a look at how data-driven discovery of PDEs can be carried out using PINNs. Consider the NS equations in 2-dimensions,

ut+λ1(uux+vuy)=−px+λ2(uxx+uyy)vt+λ1(uvx+vvy)=−py+λ2(vxx+vyy);where(⋅)x=∂(⋅)∂x. \begin{aligned} &u_{t}+\lambda_{1}\left(u u_{x}+v u_{y}\right)=-p_{x}+\lambda_{2}\left(u_{x x}+u_{y y}\right) \\ &v_{t}+\lambda_{1}\left(u v_{x}+v v_{y}\right)=-p_{y}+\lambda_{2}\left(v_{x x}+v_{y y}\right) \end{aligned}; \quad \text{where}\quad ({\cdot})_x = \frac{\partial({\cdot})}{\partial x}.

  • Here u(t,x,y)u(t, x, y) denotes the xx-component of the velocity, v(t,x,y)v(t, x, y) denotes the yy component and p(t,x,y)p(t, x, y) the pressure field.
  • Conservation of mass requires ux+vy=0  ⟹  u_{x}+v_{y}=0 \implies u=ψy,v=−ψxu=\psi_{y}, \quad v=-\psi_{x}
  • Given a set of observations: {ti,xi,yi,ui,vi}i=1N\left\{t^{i}, x^{i}, y^{i}, u^{i}, v^{i}\right\}_{i=1}^{N}
    f≐ut+λ1(uux+vuy)+px−λ2(uxx+uyy)g≐vt+λ1(uvx+vvy)+py−λ2(vxx+vyy) \begin{aligned} &f \doteq u_{t}+\lambda_{1}\left(u u_{x}+v u_{y}\right)+p_{x}-\lambda_{2}\left(u_{x x}+u_{y y}\right) \\ &g \doteq v_{t}+\lambda_{1}\left(u v_{x}+v v_{y}\right)+p_{y}-\lambda_{2}\left(v_{x x}+v_{y y}\right) \end{aligned}
  • The goal then is to learn λ={λ1,λ2}\lambda = \{\lambda_1, \lambda_2\}, and pressure field p(t,x,y)p(t, x, y) by jointly approximating [ψ(t,x,y)p(t,x,y)][\psi(t, x, y) \quad p(t, x, y)] with a single NN with two outputs. The total loss function is given by

L≐1N∑i=1N(∣u(ti,xi,yi)−ui∣2+∣v(ti,xi,yi)−vi∣2)+1N∑i=1N(∣f(ti,xi,yi)∣2+∣g(ti,xi,yi)∣2) \begin{aligned} \mathcal{L} &\doteq \frac{1}{N} \sum_{i=1}^{N}\left(\left|u\left(t^{i}, x^{i}, y^{i}\right)-u^{i}\right|^{2}+\left|v\left(t^{i}, x^{i}, y^{i}\right)-v^{i}\right|^{2}\right) \\ &+ \frac{1}{N} \sum_{i=1}^{N}\left(\left|f\left(t^{i}, x^{i}, y^{i}\right)\right|^{2}+\left|g\left(t^{i}, x^{i}, y^{i}\right)\right|^{2}\right) \end{aligned}

The training is carried out using a spectral solver NekTar and then randomly sampling points out of the grid for collocation. The figure on the bottom shows the locations of the training data-points.

Logo

PINN is able to successfully predict the pressure field with just 1%1\% of the available data as collocation points as shown below (and also able to learn the parameters λj\lambda_j in the process)

Logo

Examples: KDv Equation

As a final example illustrating data-driven discovery of PDEs, the authors choose an equation with higher order derivatives, namely the Korteweg-de Vries equation, that is encountered in the modeling of shallow water waves,

ut+λ1uux+λ2uxxx=0 u_{t}+\lambda_{1} u u_{x}+\lambda_{2} u_{x x x}=0

The problem is to learn the parameters in a similar fashion as done for Navier-Stokes Equation

N[un+cj]=λ1un+cjuxn+cj−λ2uxxxn+cj \mathcal{N}\left[u^{n+c_{j}}\right]=\lambda_{1} u^{n+c_{j}} u_{x}^{n+c_{j}}-\lambda_{2} u_{x x x}^{n+c_{j}}

PINN again is able to learn the parameters to the desired accuracy and provide an accurate resolution of the dynamics of the system.

Logo

Conclusion

The major take-away of the paper can be summarized in the following concluding remarks.

  • The authors introduced PINNs, a new class of universal function approximators that are capable of encoding any underlying physical laws that govern a given data-set (described by PDEs)
  • They also prove a design for data-driven algorithms that can be used for inferring solutions to general nonlinear PDEs, and constructing computationally efficient physics-informed surrogate models.

They also rightly point to some lingering questions that still remain unanswered in the original paper:

  • How deep/wide should the neural network be ? How much data is really needed ?
  • Why does the algorithm converge to unique values for the parameters of the differential operators, i.e., why is the algorithm not suffering from local optima for the parameters of the differential operator?
  • Does the network suffer from vanishing gradients for deeper architectures and higher order differential operators? Could this be mitigated by using different activation functions?
  • Can we improve on initializing the network weights or normalizing the data? How about the choices for the loss function choices (MSE, SSE)? What about the robustness of these networks especially when applied to solve chaotic PDEs/ODEs?

References