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)

enter image description here

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):
    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
    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_ylim([y_min, y_max])

And this is what I get: enter image description here

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.

  • $\begingroup$ 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. $\endgroup$
    – wwfe
    Commented May 15, 2023 at 17:07
  • $\begingroup$ @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 $\endgroup$ Commented May 15, 2023 at 17:30
  • $\begingroup$ 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? $\endgroup$
    – lightxbulb
    Commented May 16, 2023 at 10:41
  • $\begingroup$ 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? $\endgroup$
    – whpowell96
    Commented May 16, 2023 at 16:48

2 Answers 2


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.

  • $\begingroup$ 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. $\endgroup$ Commented May 17, 2023 at 15:06
  • 1
    $\begingroup$ I do not have a complete information on your problem, However these are my points i, Your training epochs seems to be very less.. for problems of this size in deepxde's github page, they are recommending to run epochs in the range close to 100k 2, Further increase the number of collocation points and reduce the timestep space to very few timesteps at start. Let the residual reach atleast a magnitude of 1e-2 and then check the solution. $\endgroup$ Commented May 20, 2023 at 11:45

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.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.