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

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]



# 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
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)
loss_f = u_t - alpha * (u_xx  + u_yy)
return loss_f


# Training the network

losses = []
for epoch in range(epochs):
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)

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)

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])
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.

• As a layman are you sure you are training this the correct way? It doesn't look like you have anything differentiate the separate timesteps from one another so it is training all timesteps with the same dirichlet conditions. So it makes sense that you get basically the same answer for each selected time step. Maybe it would be worthwhile to visit the literature on solving forward problems in PINNs and show what you are trying to reproduce.
– wwfe
May 15 at 17:07
• @wwfe don't Dirichlet boundary conditions by definition specify the values of the target function? I can assume the temperatures are fixed at my boundaries at every time step. The interior temperature is indeed different at different time steps, which is defined by the PDE loss term mse_f May 15 at 17:30
• I am not well-versed in ML, but have you tried using the variational formulation as a loss, and the PDE as the gradient instead? May 16 at 10:41
• It seems to be just reproducing the initial condition, but the loss is decreasing as training progresses. Have you tried giving it an initial condition closer to steady state to see if this pattern continues? May 16 at 16:48

• I guess you are correct! I did not think about it at all. I changed it to tanh(). The output is still not what it should be, but at least the model does not output the same thing at each time step. May 17 at 15:06