Tutorial 33: Physics Informed Neural Networks using JaxModel & PINN_Model ¶

• Vignesh Venkataraman

Contents ¶

• Physics Informed Neural Networks
• Setup
• Brief about Jax and Autodiff
• Burger's Equation
• Data Visualisation
• Explanation of the Solution using Jax
• Usage of PINN Model
• Visualize the final results

Physics Informed Neural Networks ¶

PINNs was introduced by Maziar Raissi et. al in their paper Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equations which are used for solving supervised learning tasks and also follow an underlying differential equation derived from understanding the Physics. In more simple terms, we try solving a differential equation with a neural network and using the differential equation as the regulariser in the loss function.

---

Here is an illustration of PINNs using a simple differential equation- ¶

$\quad \quad \quad \frac{df}{dt} = f(u, t), \quad $ where initial condition is $\ \ u(t=0) = u_0$

We approximate function $u(t)$ using a Neural Network as $NN(t)$ and apply the following loss function -

Initial Loss: $\quad L_0 = (NN(t=0) - u_0)^2$

Regulariser Loss:= $\quad L_r = | \frac{d NN (t) }{d t} - f(NN(t), t) |$

And we minimise the $Total Loss$ using Backpropagation-

Total Loss = Initial Loss + Regulariser Loss

Here is a technical definition of PINNs taken from the author's official blog- ¶

$\quad \quad u_t + \mathcal{N}[u] = 0,\ x \in \Omega, \ t\in[0,T],$

where $\ u(t,x) \ $ denotes the latent (hidden) solution, $\ N[⋅] \ $ is a nonlinear differential operator, and $Ω$ is a subset of $\ \mathbb{R}^D \ $ , and proceed by approximating $u(t,x)$ by a deep neural network. We define $\ f(t,x) \ $ to be given by

$\quad \quad f := u_t + \mathcal{N}[u],$

This assumption results in a physics informed neural network $ \ f(t,x) \ $. This network can be derived by the calculus on computational graphs: Backpropagation.

---

Useful Resources to learn more about PINNs ¶

• Maziar Raissi et. al. official blog on PINNs
• Chris Rackauckas's lecture on PINNs lecture 3 : these lectures are in Julia programming language but still are a great source of learning.

Setup ¶

Burgers Equation ¶

Burgers’ equation is a partial differential equation that was originally proposed as a simplified model of turbulence as exhibited by the full-fledged Navier-Stokes equations. It is a nonlinear equation for which exact solutions are known and is therefore important as a benchmark problem for numerical methods. More Refrence

Here is the differential Equation we are trying to solve

$\begin{array}{l} \ \ \ u_t + u u_x - (0.01/\pi) u_{xx} = 0,\ \ \ x \in [-1,1],\ \ \ t \in [0,1] \end{array}$

Here are the initial conditions

$\ \ \ u(x, 0) = -\sin(\pi x),$

$\ \ \ u(-1, t) = u(1, t) = 0.0$

Now let us define:

$ \ \ \ f := u_t + u u_x - (0.01/\pi) u_{xx}, $

and we approximate $u(x, t)$ using Neural Network as $NN(\theta, x, t)$ where $\theta$ are the weights of neural networks

---

Now here are the three main data points that will be used for training our Neural Network to approximate $u(x,t)$

1. We will train points lying between $x \in [-1,1]$ and $t=0$ to follow as part of the L2 Loss

$min\ \ _\theta \ \ (NN(\theta, x, t) + \sin(\pi x))^2$

2. We will train points lying between $t \in [0,1]$ and $x= \pm1 $ as part of the L2 Loss

$min\ \ _\theta \ \ (NN(\theta, x, t) + \sin(\pi x))^2$

3. We will train points lying between $x \in [-1,1],\ \ \ t \in [0,1]$ as part of the regulariser loss

$f(\theta, x, t):= \ \ \frac{\partial NN(\theta, x, t)}{\partial t} + NN(\theta, x, t)\frac{\partial NN(\theta, x, t)}{\partial x} - (0.01/\pi)\frac{\partial^2 NN(\theta, x, t)}{\partial^2 x} $

$min\ \ _\theta \ \ f(\theta, x, t)$

In this tutorial, we will be combing data conditions 1 and 2 under the same L2Loss

Data Visualisation of the Burgers Equation ¶

Now lets load the Burger's Data provided from the author. pre_process_shock_data is used to load the data in a format suitable for Neural Networks. Understanding this function is not neccesary for working through the tutorials.

We have three Numpy arrays labeled_X , labeled_y and unlabeled_X which will be used for training our neural network,

1. labeled_X consists of $x \in [-1,1]$ & $t=0$ and $t \in [0,1]$ & $x= \pm1 $. labeled_y has the value of $u(x, t)$:

Let us verify that labeled_X & labeled_y also consists of data points satisfying the condition of

$\ \ \ u(x, 0) = -\sin(\pi x), \quad \quad x \in [-1,1]$ & $t=0$

Let us verify that at labeled_X & labeled_y also consists of datapoints satisfying the condition of

$\ \ \ u(-1, t) = u(1, t) = 0.0, \quad \quad t \in [0,1]$ & $x= \pm1$

Explanation of the solution ¶

We will be using Deepchem's PINNModel class to solve Burger's Equation which is based out of Jax library. We will approximate $u(x, t)$ using a Neural Network represented as $NN(\theta, x, t)$

For our purpose, we will be using the Haiku library for building neural networks. Due to the functional nature of Jax, we define neural network with two things

• Parameters - which act as the weight matrices, upon which Backpropagation is applied for optimisation.
• forward_fn - This defines how the weights are used for computing the outputs. Ex- Feedforward, Convolution, etc

As per the docstrings of PINNModel, we require two additional functions in the given format -

1. Create a gradient_fn which tells us about how to compute the gradients of the function-

  >>>
  >> def gradient_fn(forward_fn, loss_outputs, initial_data):
  >>  def model_loss(params, target, weights, rng, ...):
  >>    # write code using the arguments.
  >>    # ... indicates the variable number of positional arguments.
  >>    return
  >>  return model_loss

And to understand more about PINNModel, you can see that the same gradient_fn gets called in the code for computing the gradients.

For our purpose, we have two variables $(x, t)$ and we need to tell the PINN Model how to compute the final gradient. For carrying out this process we will be using these main features from jax library for calculating the loss -

1. vmap - This for parallelising computations in batches. We will process each row of the dataset, but it will get batched automatically using this feature.
2. jacrev - This is used to calculate the jacobian matrix. In our case, the output is a single dimension and hence it can be thought of as the gradient function. We could directly use jax's grad function but using jacrev simplifies the array shapes and hence is easier.

We need to compute two losses for solving our differential equation-

1. Initial Loss

u_pred = forward_fn(params, rng, x_b, t_b)
initial_loss = jnp.mean((u_pred - boundary_target) ** 2)

2. Regulariser Loss

This is slightly complicated as we need to compute

$f(\theta, x, t):= \ \ \frac{\partial NN(\theta, x, t)}{\partial t} + NN(\theta, x, t)\frac{\partial NN(\theta, x, t)}{\partial x} - (0.01/\pi)\frac{\partial^2 NN(\theta, x, t)}{\partial^2 x} $

The partial derivative operation in the first and second terms can be calculated using jacrev function-

u_x, u_t = jacrev(forward_fn, argnums=(2, 3))(params, rng, x, t)

The second partial derivative operation in the third term can be applying jacrev twice-

u_xx = jacrev(jacrev(forward_fn, argnums=2), argnums=2)(params, rng, x, t)

2. We also need to provide an eval_fn in the below-given format for computing the weights

  >>>
  >> def create_eval_fn(forward_fn, params):
  >>  def eval_model(..., rng=None):
  >>    # write code here using arguments
  >>
  >>    return
  >>  return eval_model

Like previously we have two arguments for our model $(x, t)$ which get passed in function

Usage of PINN Model ¶

We will be using optax library for performing the optimisations. PINNModel executes the codes for training the models.

Visualize the final results ¶

• Code taken from authors for visualisation
• show both graphs