Problems solving 2D heat equation using physics-informed neural networks

Question

I am trying to solve 2D heat equation using the physics-informed neural networks approach. The training loss is decreasing, but my final network outputs make no sense. I am using Python/Pytorch.

2D heat equation: $$\frac{\partial u}{\partial t} = \alpha \left(\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}\right) $$

Here is what I am doing:

Defining initial and boundary conditions

I start by defining the initial and boundary conditions (BC: all sides are at 0C except top side which is at 100C / IC: same as boundary conditions with t=0)

import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn

num_points = 10000 # Number of randomly sampled points for IC and BC
alpha = 15 # thermal diffusivity

t_min, t_max, x_min, x_max, y_min, y_max = 0, 1, 0, 1, 0, 1
x_range, y_range, t_range = np.random.rand(num_points), np.random.rand(num_points), np.random.rand(num_points)
zeros = np.zeros(num_points)
ones = np.ones(num_points)
# [x, y, t, T]
bc1 = np.vstack([x_range, ones, t_range, ones*100]).T # Boundary top: u(x, y=1, t) = 100
bc2 = np.vstack([ones, y_range, t_range, zeros]).T # Boundary right: u(x=1, y, t) = 0
bc3 = np.vstack([x_range, zeros, t_range, zeros]).T # Boundary bottom: u(x, y=0, t) = 0
bc4 = np.vstack([zeros, y_range, t_range, zeros]).T # Boundary left: u(x=0, y0, t) = 0
ic = np.vstack([x_range, y_range, zeros, zeros]).T # IC u(x, y, t=0) = 0

initial_boundary_conditions = np.vstack([bc1, bc2, bc3, bc4, ic])

x_initial_bc = initial_boundary_conditions[:, 0]
y_initial_bc = initial_boundary_conditions[:, 1]
t_initial_bc = initial_boundary_conditions[:, 2]
u_initial_bc = initial_boundary_conditions[:, 3]


x_initial_bc = torch.autograd.Variable(torch.from_numpy(x_initial_bc).float(), requires_grad=True)
y_initial_bc = torch.autograd.Variable(torch.from_numpy(y_initial_bc).float(), requires_grad=True)
t_initial_bc = torch.autograd.Variable(torch.from_numpy(t_initial_bc).float(), requires_grad=True)
u_initial_bc = torch.autograd.Variable(torch.from_numpy(u_initial_bc).float(), requires_grad=True)

Network architecture

Next, I define my network architecture, 3 inputs (2 spatial, one temporal) layers, 7 hidden layers with ReLU() activation, each with 20 neurones, and one output layer.

collocation_points = 100000 # Number of randomly sampled collocation points to be evaluated for the PDE-based loss
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.neurons_per_layer = 20
        self.fc1 = nn.Linear(3, self.neurons_per_layer)
        self.fc2 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc3 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc4 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc5 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc6 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc7 = nn.Linear(self.neurons_per_layer, self.neurons_per_layer)
        self.fc8 = nn.Linear(self.neurons_per_layer, 1)
        self.relu = nn.ReLU()
        
    def forward(self, x, y, t):
        inputs = torch.cat([x.reshape(-1, 1), y.reshape(-1, 1), t.reshape(-1, 1)], axis=1)
        output = self.relu(self.fc1(inputs))
        output = self.relu(self.fc2(output))
        output = self.relu(self.fc3(output))
        output = self.relu(self.fc4(output))
        output = self.relu(self.fc5(output))
        output = self.relu(self.fc6(output))
        output = self.relu(self.fc7(output))
        output = self.fc8(output)
        return output
net = Net()
epochs = 1000
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
criterion = torch.nn.MSELoss()

PDE loss

I define the loss from both the PDE evaluated at the collocation points (using 100.000 points):

$$\frac{\partial u}{\partial t} - \alpha \left(\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}\right) = 0$$

# PDE as a loss function, this function calculates the first and second derivatives, then the PDE-based loss 
def f(x, y, t, net):
    u = net(x, y, t)
    u_x = torch.autograd.grad(u, x, create_graph=True, retain_graph=True, grad_outputs=torch.ones_like(u))[0]
    u_xx = torch.autograd.grad(u_x, x, create_graph=True, retain_graph=True, grad_outputs=torch.ones_like(u_x))[0]
    u_y = torch.autograd.grad(u, y, create_graph=True, retain_graph=True, grad_outputs=torch.ones_like(u))[0]
    u_yy = torch.autograd.grad(u_y, y, create_graph=True, retain_graph=True, grad_outputs=torch.ones_like(u_y))[0]
    u_t = torch.autograd.grad(u, t, create_graph=True, retain_graph=True, grad_outputs=torch.ones_like(u))[0]
    loss_f = u_t - alpha * (u_xx  + u_yy)
    return loss_f

Training the network

losses = []
for epoch in range(epochs):
    optimizer.zero_grad()
    predictions_initial_bc = net(x_initial_bc, y_initial_bc, t_initial_bc)
    mse_u = criterion(predictions_initial_bc.reshape(-1,), u_initial_bc) # This is the loss from boundary and initial conditions
    
    
    t_collocation = torch.FloatTensor(collocation_points, ).uniform_(t_min, t_max)
    x_collocation = torch.FloatTensor(collocation_points, ).uniform_(x_min, x_max)
    y_collocation = torch.FloatTensor(collocation_points, ).uniform_(y_min, y_max)
    t_collocation = torch.autograd.Variable(t_collocation, requires_grad=True)
    x_collocation = torch.autograd.Variable(x_collocation, requires_grad=True)
    y_collocation = torch.autograd.Variable(y_collocation, requires_grad=True)
   
    f_out = f(x_collocation, y_collocation, t_collocation, net)
    mse_f = criterion(torch.zeros_like(f_out), f_out) # This is the PDE-based loss evaluated at the randomly sampled collocation points
    
    loss = mse_u + mse_f
    losses.append(loss.item())
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        
        
        print(f'Epoch {epoch}/{epochs}: Loss = {loss.item()}')

Here is my training loss:

Epoch 0/1000: Loss = 2009.57
Epoch 100/1000: Loss = 1455.76
Epoch 200/1000: Loss = 726.14
Epoch 300/1000: Loss = 134.33
Epoch 400/1000: Loss = 64.62
Epoch 500/1000: Loss = 46.90
Epoch 600/1000: Loss = 35.36
Epoch 700/1000: Loss = 46.16
Epoch 800/1000: Loss = 28.22
Epoch 900/1000: Loss = 26.59

Finally, I try to visualize the results at 10 different time steps for time from 0 to 1s:

x = np.arange(x_min, x_max, 0.01)
y = np.arange(y_min, y_max, 0.01)
mesh_x, mesh_y = np.meshgrid(x, y)
x = np.ravel(mesh_x).reshape(-1, 1)
y = np.ravel(mesh_y).reshape(-1, 1)

pt_x = torch.autograd.Variable(torch.from_numpy(x).float(), requires_grad=True)
pt_y = torch.autograd.Variable(torch.from_numpy(y).float(), requires_grad=True)

period = np.linspace(0, t_max, 10)

fig, axes = plt.subplots(2, 5, figsize=(15, 5))

for index, axis in enumerate(axes.ravel()):
  t = torch.full(x.shape, period[index])
  pt_t = torch.autograd.Variable(t, requires_grad=True)
  u = net(pt_x, pt_y, pt_t).data.numpy()
  mesh_u = u.reshape(mesh_x.shape)
  cm = axis.pcolormesh(mesh_x, mesh_y, mesh_u, cmap='jet')#, vmin=-1, vmax=1)
  fig.colorbar(cm, ax=axis)
  axis.set_xlim([x_min, x_max])
  axis.set_xticks([])
  axis.set_yticks([])
  axis.set_ylim([y_min, y_max])
fig.tight_layout()

And this is what I get:

Obviously there is no heat flow at all. I tried to play around with the number of collocation and boundary/initial points, epochs, network architecture, but it just won't change anything.

Answer 1

The one thing that i can notice immediately where it might be going wrong is the use of ReLU Activation function for PINNs. In this problem,we have a double derivative in our loss function (Laplacian). So its better to use more smooth activation functions like tanh or sigmoid or swish, So that these gradients are preserved.

Answer 2

Pay attention to the activation function used. The second derivative of RELU would be zero (except at 0). In addition, for numerically solving problems like heat (diffusion) equation using physics informed neural networks (PINNs) take a look at some papers like this.