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()
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.
mse_f
$\endgroup$