After successfully solving the steady-state version here and one of the transient equation with second derivative of the quantity of interest here 🌊, this blog presents the code to solve the other transient version, the heat 🔥 / diffusion 💉 equation using a physics informed neural network 🖧, ⚛ PINN 🧠.
This 👽 technology works without any data 📚 requirements just like discrete methods like finite element or finite difference methods etc. The current code contains the mixed (Neumann and Dirichlet) boundary conditions which can be easily modified. The example also has a heat source.
I stopped the training early when by eye balling the results looked okay. Again, you are reading a blog, not a Q1 journal 😁. Furthermore, this is the first code I share with batch gradient descent.
Fig. 1 shows the results from the code at current boundary conditions. As always, the PINNs use hard form of the equation and do not require a mesh to be solved 😯.
Fig. 1, The animation
Cite as: Fahad Butt (2024). HEAT-PINN (), Blogger. Retrieved Month Date, Year
Code
#Copyright <2024> <FAHAD BUTT>
#Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#%% import necessary libraries
import os
import numpy as np
import random as rn
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import glob
# import torch_directml # enable for AMD and Intel gpu
device = torch.device("cpu") # cuda for CUDA gpu, cpu for cpu, mps for Apple GPU, torch_directml.device() for any other gpu
#%% 2D heat equation and boundary conditions
def heat_equation(he):
he.requires_grad_(True) # watch he for gradient computation
xx = he[:, 0] # x coordinates
yy = he[:, 1] # y coordinates
tt = he[:, 2] # time coordinates
out = model(he) # output from model
T = out[:, 0] # temperature
dT = torch.autograd.grad(T, he, torch.ones_like(T), create_graph=True)[0] # dT/dhe
T_x = dT[:, 0] # dT/dx
T_y = dT[:, 1] # dT/dy
T_t = dT[:, 2] # dT/dt
T_xx = torch.autograd.grad(T_x, he, torch.ones_like(T_x), create_graph=True)[0][:, 0] # d2T/dx2
T_yy = torch.autograd.grad(T_y, he, torch.ones_like(T_y), create_graph=True)[0][:, 1] # d2T/dy2
heat = T_t - (alpha * (T_xx + T_yy)) # heat equation
# apply boundary conditions
T_left_bc = T[xx == 0] # x = 0
T_right_bc = T[xx == L] # x = L
T_top_bc = T[yy == D] - 1 # y = D
T_bottom_bc = T_y[yy == 0] # y = 0
radius = 0.05 # radius of heat source
center_x = L / 2 # center points of heat source
center_y = D / 2
distance_from_center = torch.sqrt((xx - center_x) ** 2 + (yy - center_y) ** 2) # included points
inside_circle = distance_from_center <= radius
T_center = T[inside_circle] - torch.sin(he[inside_circle, 2]) # time-dependent heat source
# apply initial conditions
T_ic_I = T[tt == 0] # t = 0
return heat, \
T_left_bc, T_right_bc, T_top_bc, T_bottom_bc, \
T_ic_I, \
T_center
#%% define loss function
def loss_fn(he):
global heat_loss, T_bc_loss, T_ic_loss, mse_loss, bc_loss
he = he.clone().detach().requires_grad_(True).to(device) # move to "device" and use x32 data type
heat, \
T_left_bc, T_right_bc, T_top_bc, T_bottom_bc, \
T_ic_I, \
T_center = heat_equation(he) # compute residuals for heat equations with specified boundary and initial conditions
heat_loss = nn.MSELoss()(heat, torch.zeros_like(heat)) # heat equation loss
T_bc_loss = nn.MSELoss()(T_left_bc, torch.zeros_like(T_left_bc)) + nn.MSELoss()(T_right_bc, torch.zeros_like(T_right_bc)) \
+ nn.MSELoss()(T_top_bc, torch.zeros_like(T_top_bc)) + nn.MSELoss()(T_bottom_bc, torch.zeros_like(T_bottom_bc)) # boundary condition loss
T_ic_loss = nn.MSELoss()(T_ic_I, torch.zeros_like(T_ic_I)) # initial condition loss
T_center_loss = nn.MSELoss()(T_center, torch.zeros_like(T_center)) # heat source loss
bc_loss = T_bc_loss / 480 # boundary condition loss
mse_loss = (heat_loss + T_center_loss) / 1143 # equation loss
ic_loss = T_ic_loss / 41 # initial condition loss
total_loss = mse_loss + bc_loss + ic_loss # total loss
return total_loss
#%% save and resume training
def save_checkpoint(epoch, model, optimizer): # save check point
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()}
checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint_epoch_{epoch}.pt")
torch.save(checkpoint, checkpoint_path)
print(f"Checkpoint saved at epoch {epoch}")
def load_checkpoint(model, optimizer): # load check point
latest_checkpoint = max(glob.glob(os.path.join(checkpoint_dir, 'checkpoint_epoch_*.pt')), key=os.path.getctime)
checkpoint = torch.load(latest_checkpoint)
epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
print(f"Checkpoint loaded from epoch {epoch}")
return model, optimizer, epoch
#%% show model summary
def count_parameters(model):
return sum(p1.numel() for p1 in model.parameters() if p1.requires_grad) # show trainable parameters
#%% set same seeds for initialization
np.random.seed(20) # same random initializations every time numpy
rn.seed(20) # same random initializations every time for python
torch.manual_seed(20) # set seed for reproducibility
#%% generate sample points (grid)
alpha = 1e-4 # thermal diffusivity [m2/s]
L = 1 # Length of plate
D = 1 # Width of plate
TT = 2.5 * np.pi # total time [s]
# TT = 1 # total time [s]
num_training_points = 120 # per length
# generate points for boundaries
x_l = np.random.uniform(0, 0, num_training_points) # left boundary
y_l = np.random.uniform(0, D, num_training_points)
x_r = np.random.uniform(L, L, num_training_points) # right boundary
y_r = np.random.uniform(0, D, num_training_points)
x_b = np.random.uniform(0, L, num_training_points) # bottom boundary
y_b = np.random.uniform(0, 0, num_training_points)
x_t = np.random.uniform(0, L, num_training_points) # top boundary
y_t = np.random.uniform(D, D, num_training_points)
x_m = np.random.uniform(0, L, 1138) #interior points
y_m = np.random.uniform(0, D, 1138)
x_m_box = np.random.uniform(0.45, 0.55, 5) # hear source points
y_m_box = np.random.uniform(0.45, 0.55, 5)
x = np.concatenate([x_l, x_r, x_b, x_t, x_m, x_m_box]) # x co-ordinates
y = np.concatenate([y_l, y_r, y_b, y_t, y_m, y_m_box]) # y co-ordinates
t = np.linspace(0, TT, 41) # time points
x_repeated = np.repeat(x, len(t)) # make x, y and t same lengths
y_repeated = np.repeat(y, len(t))
t_tiled = np.tile(t, len(x))
# define circle parameters (for plotting)
center_x = L / 2 # x coordinate of center
center_y = D / 2 # y coordinate of center
radius = 0.05 # radius of circle
# generate theta values (angle) from 0 to 2*pi
theta = np.linspace(0, 2 * np.pi, 100)
x_circle = center_x + radius * np.cos(theta) # circle x co ordinates
y_circle = center_y + radius * np.sin(theta) # circle y co ordinates
#%% display mesh
plt.figure(dpi = 300)
plt.scatter(x, y, s = 1, c = 'k', alpha = 1) # scatter plot
plt.plot(x_circle, y_circle, 'r-', label = 'Circle') # a red circle
plt.xlim(0, L)
plt.ylim(0, D)
plt.gca().set_aspect('equal', adjustable = 'box')
plt.axis('off')
plt.show()
#%% create neural network model
h = 50 # number of neurons per hidden layer
model = nn.Sequential(
nn.Linear(3, h), # 3 inputs for x, y and t
nn.Tanh(), # hidden layer
nn.Linear(h, h),
nn.Tanh(), # hidden layer
nn.Linear(h, h),
nn.Tanh(), # hidden layer
nn.Linear(h, h),
nn.Tanh(), # hidden layer
nn.Linear(h, h),
nn.Tanh(), # hidden layer
nn.Linear(h, 1)).to(device) # 1 output for T
for m in model.modules():
if isinstance(m, nn.Linear):
init.xavier_normal_(m.weight) # weights initialization
init.constant_(m.bias, 0.0) # bias initialization
print(model) # show model details
num_trainable_params = count_parameters(model) # show trainable parameters
print("Number of trainable parameters:", num_trainable_params)
#%% prepare input data
train_points = np.vstack((x_repeated, y_repeated, t_tiled)).T # Stack x, y, t together
train_points = torch.tensor(train_points, dtype = torch.float32).to(device) # convert to tensor and move training points to GPU / CPU
train_dataset = TensorDataset(train_points) # mark points to create batches
train_loader = DataLoader(train_dataset, batch_size = 128, shuffle = True, drop_last = True) # create batches
#%% train using Adam
resume_training = False # True if resuming training
start_epoch = 1
optimizer = optim.Adam(model.parameters(), lr = 1e-4) # select optimizer and parameters
checkpoint_dir = './pytorch_checkpoints' # directory to save checkpoints
os.makedirs(checkpoint_dir, exist_ok = True) # create a directory to save checkpoints
if resume_training:
model, optimizer, start_epoch = load_checkpoint(model, optimizer)
start_epoch += 1 # start from next epoch
for epoch in range(2001):
epoch_loss = 0.0
for batch in train_loader:
batch_points = batch[0] # get batch data
optimizer.zero_grad() # clear gradients of all optimized variables
loss_value = loss_fn(batch_points) # compute prediction by passing inputs to model
loss_value.backward() # compute gradient of loss with respect to model parameters
optimizer.step() # perform parameter update2
epoch_loss += loss_value.item()
if epoch % 100 == 0:
print(f"{epoch} {loss_value.item()}")
save_checkpoint(epoch, model, optimizer) # save model details
print(f"{epoch} {loss_value.item()}") # show loss values
print("Final Loss Values:")
print("heat_loss:", heat_loss.item())
print("T_bc_loss:", T_bc_loss.item())
print("T_ic_loss:", T_ic_loss.item())
print("mse_loss:", mse_loss.item())
print("bc_loss:", bc_loss.item())
#%% prediction
target_time_step = 0 * np.pi # in seconds
num_testing_points = 250 # per length
# Generate points for boundaries
x_l_te = np.random.uniform(0, 0, num_testing_points) # left boundary
y_l_te = np.random.uniform(0, D, num_testing_points)
x_r_te = np.random.uniform(L, L, num_testing_points) # right boundary
y_r_te = np.random.uniform(0, D, num_testing_points)
x_b_te = np.random.uniform(0, L, num_testing_points) # bottom boundary
y_b_te = np.random.uniform(0, 0, num_testing_points)
x_t_te = np.random.uniform(0, L, num_testing_points) # top boundary
y_t_te = np.random.uniform(D, D, num_testing_points)
x_m_te = np.random.uniform(0, L, 62500) # mid field
y_m_te = np.random.uniform(0, D, 62500)
x_te = np.concatenate([x_l_te, x_r_te, x_b_te, x_t_te, x_m_te]) # x co-ordinates
y_te = np.concatenate([y_l_te, y_r_te, y_b_te, y_t_te, y_m_te]) # y co-ordinates
test_points = np.vstack((x_te, y_te)).T # stack x, y together
test_points_t = np.column_stack([test_points, np.full_like(test_points[:, 0], target_time_step)]) # generate test points for target time step
test_points_t = torch.tensor(test_points_t, dtype=torch.float32).to(device) # convert to tensor and move training points to GPU / CPU
predict = model(test_points_t).detach().cpu().numpy() # predict and move back to CPU
T = predict[:, 0] # predict and bring back to dimensional form
# plotting
plt.figure(dpi = 300)
sc = plt.scatter(x_te, y_te, c = T, cmap = 'jet', s = 1, edgecolor = 'none',\
alpha = 1, vmin = -1, vmax = 1)
plt.colorbar(orientation = 'vertical')
plt.gca().set_aspect('equal', adjustable = 'box')
plt.xticks([0, L])
plt.yticks([0, D])
plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.axis('off')
plt.show()
Thank you for reading!
