Recently, a new 🆕 PyTorch 🔦 optimizer is published [1]. This post will be about testing this new optimizer to solve linear elasticity problems using Data-less Physics Informed Neural Network. The details of optimizer are mentioned in [1]. To run the optimizer as provided with the code in this post, install the optimizer by placing the soap.py from [1] to the envs\torch\Lib\site-packages folder, where torch is the PyTorch environment. The soap.py can also be copied to the working directory. The changes in code to implement the optimizer are very straight forward.

The first case solved is that of a flat plate in tension. Symmetry boundary conditions are applied in both x and y axes to take advantage of symmetric loadings and constraints. The new optimizer offer very quicker ⚡ convergence as compared to ADAM, it must be said. Within Fig. 1, the results are shown including von Mises stress, resultant displacement and equivalent strain. These results quite same as the results produced by a commercial FEM software which shall not be named 😁, under same boundary conditions. The plate is in pure tension. Fig. 2 shows the same specimen under compression. The def linear_elasticity_equation(lee): function is made available as well along with its loss function.

The function to implement Brazilian test on concrete cubes is made available. Remaining code is same as before. The Brazilian test puts the specimen in tension under compressive loading. The 1st principal stress is shown in Fig. 3.

The function to implement 3-point bending test on concrete slabs is made available. Remaining code is same as before. The deformed beam is shown in Fig. 4. The beam dimensions are taken from [2].

Fig. 1, The results from the new optimizer

Fig. 2, The results of compressive loading

Fig. 3, The results from Brazilian test

Fig. 4, The results from 3-Point bending test

Code

The code to reproduce Fig. 1 is made available.

#%% import necessary libraries

import os

import time

from datetime import datetime

import numpy as np

import random as rn

import matplotlib.pyplot as plt

import torch

import torch.nn as nn

from soap import SOAP # copy soap.py to working directory

import glob

device = torch.device("cpu") # cuda for CUDA gpu, cpu for cpu, mps for Apple GPU, torch_directml.device() for any other gpu

#%% 2D linear elasticity equations and boundary conditions

def linear_elasticity_equation(lee):

lee.requires_grad_(True) # watch lee for gradient computation

out = model(lee) # output from model

u = out[:, 0] # x displacement

v = out[:, 1] # y displacement

sxx = out[:, 2] # normal stress x

syy = out[:, 3] # normal stress y

sxy = out[:, 4] # shear stress xy

exx = out[:, 5] # normal strain x

eyy = out[:, 6] # normal strain y

exy = out[:, 7] # shear strain xy

du = torch.autograd.grad(u, lee, torch.ones_like(u), create_graph = True)[0] # du/dlee

dv = torch.autograd.grad(v, lee, torch.ones_like(v), create_graph = True)[0] # dv/dlee

dsxx = torch.autograd.grad(sxx, lee, torch.ones_like(sxx), create_graph = True)[0] # dsxx/dlee

dsyy = torch.autograd.grad(syy, lee, torch.ones_like(syy), create_graph = True)[0] # dsyy/dlee

dsxy = torch.autograd.grad(sxy, lee, torch.ones_like(sxy), create_graph = True)[0] # dsxy/dlee

dexx = torch.autograd.grad(exx, lee, torch.ones_like(exx), create_graph = True)[0] # dexx/dlee

deyy = torch.autograd.grad(eyy, lee, torch.ones_like(eyy), create_graph = True)[0] # deyy/dlee

dexy = torch.autograd.grad(exy, lee, torch.ones_like(exy), create_graph = True)[0] # dexy/dlee

u_x = du[:, 0] # du/dx

u_y = du[:, 1] # du/dy

v_x = dv[:, 0] # dv/dx

v_y = dv[:, 1] # dv/dy

sxx_x = dsxx[:, 0] # dsxx/dx

syy_y = dsyy[:, 1] # dsyy/dy

sxy_x = dsxy[:, 0] # dtxy/dx

sxy_y = dsxy[:, 1] # dtxy/dy

exx_y = dexx[:, 1] # dexx/dy

eyy_x = deyy[:, 0] # deyy/dx

exy_x = dexy[:, 0] # dexy/dx

x_mom = (D * sxx_x) + (L * sxy_y) # x momentum

y_mom = (D * sxy_x) + (L * syy_y) # y momentum

x_kin = exx - u_x # kinematic x

y_kin = eyy - v_y # kinematic y

xy_kin = exy - (0.5 * (u_y + v_x)) # kinematic xy

ps_xx = sxx - (alpha * (((1 - v_ps) * exx) + (v_ps * eyy))) # plane strain xx

ps_yy = syy - (alpha * (((1 - v_ps) * eyy) + (v_ps * exx))) # plane strain yy

ps_xy = sxy - (beta * exy) # plane strain xy

# ps_xx = sxx - (alpha_plane_stress * (exx + (v_ps * eyy))) # plane stress xx

# ps_yy = syy - (alpha_plane_stress * (eyy + (v_ps * exx))) # plane stress yy

# ps_xy = sxy - (beta * exy) # plane stess xy

exx_yy = torch.autograd.grad(exx_y, lee, torch.ones_like(exx_y), create_graph=True)[0][:, 1] # d2exx/dy2

eyy_xx = torch.autograd.grad(eyy_x, lee, torch.ones_like(eyy_x), create_graph=True)[0][:, 0] # d2eyy/dx2

exy_xy = torch.autograd.grad(exy_x, lee, torch.ones_like(exy_x), create_graph=True)[0][:, 1] # d2exy/dxdy

comp = ((L / D) * exx_yy) - (2 * exy_xy) + ((D / L) * eyy_xx) # compatibility condition

# apply boundary conditions

u_left = u[(lee[:, 0] == 0)] # x = 0

v_bottom = v[(lee[:, 1] == 0)] # y = 0

sxy_right_bc = sxy[lee[:, 0] == L] # x = L

sxy_top_bc = sxy[lee[:, 1] == D] # y = D

syy_top_bc = syy[lee[:, 1] == D] # y = D

sxx_right_bc = sxx[lee[:, 0] == L] - 1 # x = L

return x_mom, y_mom, x_kin, y_kin, xy_kin, ps_xx, ps_yy, ps_xy, comp, \

sxy_right_bc, sxy_top_bc, \

syy_top_bc, \

sxx_right_bc, \

u_left, v_bottom

#%% define loss function

def loss_fn(lee):

global u_bc_loss, v_bc_loss, sxx_bc_loss, sxy_bc_loss, syy_bc_loss, sigma_r_hole_bc_loss, tau_rtheta_hole_bc_loss

lee = lee.clone().detach().requires_grad_(True).to(device) # move to "device" and use x32 data type

x_mom, y_mom, x_kin, y_kin, xy_kin, ps_xx, ps_yy, ps_xy, comp, \

sxy_right_bc, sxy_top_bc, \

syy_top_bc, \

sxx_right_bc, \

u_left, v_bottom = linear_elasticity_equation(lee) # compute residuals for equations with specified boundary conditions

x_mom_loss = nn.MSELoss()(x_mom, torch.zeros_like(x_mom)) # x momentum loss

y_mom_loss = nn.MSELoss()(y_mom, torch.zeros_like(y_mom)) # y momentum loss

x_kin_loss = nn.MSELoss()(x_kin, torch.zeros_like(x_kin)) # kinematic equation loss x

y_kin_loss = nn.MSELoss()(y_kin, torch.zeros_like(y_kin)) # kinematic equation loss y

xy_kin_loss = nn.MSELoss()(xy_kin, torch.zeros_like(xy_kin)) # kinematic equation loss xy

ps_xx_loss = nn.MSELoss()(ps_xx, torch.zeros_like(ps_xx)) # plane strain loss xx

ps_yy_loss = nn.MSELoss()(ps_yy, torch.zeros_like(ps_yy)) # plane strain loss yy

ps_xy_loss = nn.MSELoss()(ps_xy, torch.zeros_like(ps_xy)) # plane strain loss xy

comp_loss = nn.MSELoss()(comp, torch.zeros_like(comp)) # compatibility loss

u_bc_loss = nn.MSELoss()(u_left, torch.zeros_like(u_left)) # x symmetry loss

v_bc_loss = nn.MSELoss()(v_bottom, torch.zeros_like(v_bottom)) # y symmetry loss

sxx_bc_loss = nn.MSELoss()(sxx_right_bc, torch.zeros_like(sxx_right_bc)) # sxx boundary condition loss

sxy_bc_loss = nn.MSELoss()(sxy_right_bc, torch.zeros_like(sxy_right_bc)) \

+ nn.MSELoss()(sxy_top_bc, torch.zeros_like(sxy_top_bc)) # sxy boundary condition loss

syy_bc_loss = nn.MSELoss()(syy_top_bc, torch.zeros_like(syy_top_bc)) # syy boundary condition loss

bc_loss = (u_bc_loss + v_bc_loss + sxx_bc_loss + sxy_bc_loss + syy_bc_loss) / 25 # total boundary condition loss

mse_loss = (x_mom_loss + y_mom_loss + \

+ x_kin_loss + y_kin_loss + xy_kin_loss + \

ps_xx_loss + ps_yy_loss + ps_xy_loss + comp_loss) / 525 # total equation loss

total_loss = mse_loss + bc_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)

s_0 = 1e6 # applied load [Pa]

E = 2e11 # Young's modulus [Pa]

v_ps = 0.3 # Poisson's ratio

L = 0.5 # length of plate [m]

D = 0.5 # width of plate [m]

E_0 = E / s_0 # non dimensional Young's modulus

U = s_0 * L / E # scaled u displacement

V = s_0 * D / E # scaled v displacement

e_0 = s_0 / E # scaled strain

alpha = 1 / ((1 + v_ps) * (1 - (2 * v_ps))) # non dimensional parameter plane strain

alpha_plane_stress = 1 / (1 - (v_ps * v_ps)) # non dimensional parameter plane stress

beta = 1 / (2 * (1 + v_ps)) # non dimensional parameter

num_training_points = 25 # per length

# generate the original points (as in your code)

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, 525) # mid points

y_m = np.random.uniform(0, D, 525)

x = np.concatenate([x_l, x_r, x_b, x_t, x_m]) # x co-ordinates

y = np.concatenate([y_l, y_r, y_b, y_t, y_m]) # y co-ordinates

#%% display points

plt.figure(dpi = 300)

plt.scatter(x, y, s = 2, c = 'r', alpha = 1) # scatter plot

plt.xlim(0, L)

plt.ylim(0, D)

plt.grid(True)

plt.axis('equal')

plt.grid(False)

plt.show()

#%% create neural network model

h = 50 # number of neurons per hidden layer

model = nn.Sequential(

nn.Linear(2, h), # 2 inputs for x and y

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, 8)).to(device) # 3 outputs for 3 stresses, 3 strains, 2 displacements

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, y)).T # stack x, y together

train_points = torch.tensor(train_points, dtype=torch.float32).to(device) # convert to tensor and move training points to GPU / CPU, use float64 if more precision is required

#%% training

start_time = time.time()

print_time = datetime.now()

print(f"Calculation started at: {print_time.strftime('%Y-%m-%d %H:%M:%S')}")

resume_training = False # True if resuming training

start_epoch = 1

lr = 1e-3 # learning rate

optimizer = SOAP(model.parameters(), lr = lr)

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 param_group in optimizer.param_groups:

param_group['lr'] = lr

for param_group in optimizer.param_groups:

current_lr = param_group['lr'] # get learning rate

print(f"Learning Rate: {current_lr}")

for epoch in range(10001):

optimizer.zero_grad() # clear gradients of all optimized variables

loss_value = loss_fn(train_points) # compute prediction by passing inputs to model

loss_value.backward() # compute gradient of loss with respect to model parameters

optimizer.step() # perform parameter up date

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(f"Learning Rate: {current_lr}")

end_time = time.time()

elapsed_time = end_time - start_time

print(f"Time taken for the run: {elapsed_time:.2f} seconds")

#%% prediction

num_testing_points = 200 # per length

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, 40000) # mid points

y_m_te = np.random.uniform(0, D, 40000)

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 = torch.tensor(test_points, dtype=torch.float32).to(device) # convert to tensor and move training points to GPU / CPU, use float64 if more precision is required

predict = model(test_points).detach().cpu().numpy() # predict and move back to CPU

u = predict[:, 0] * 2

v = predict[:, 1] * 2

sxx = predict[:, 2]

syy = predict[:, 3]

sxy = predict[:, 4]

exx = predict[:, 5]

eyy = predict[:, 6]

exy = predict[:, 7]

von_mises = np.sqrt(sxx**2 + syy**2 - sxx * syy + 3 * sxy**2)

resultant_displacement = np.sqrt((u * U * 1000)**2 + (v * V * 1000)**2)

equivalent_strain = np.sqrt((2 / 3) * (exx**2 + eyy**2 + 2 * exy**2))

average_stress = (sxx + syy) / 2

radius_stress = np.sqrt(((sxx - syy) / 2) ** 2 + sxy ** 2)

sigma_1 = average_stress + radius_stress # 1st (maximum) principal stress

sigma_3 = average_stress - radius_stress # 3rd (minimum) principal stress

# plotting

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = u * U * 1000, cmap = 'jet', s = 1, edgecolor='none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("u")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = v * V * 1000, cmap = 'jet', s = 1, edgecolor='none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("v")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = sxx * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("sxx")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = syy * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("syy")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = sxy * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("txy")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = exx * e_0, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("exx")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = eyy * e_0, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("eyy")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = exy * e_0, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("gxy")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = von_mises * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'vertical')

plt.gca().set_aspect('equal', adjustable = 'box')

# plt.axis('off')

plt.title("von Mises [MPa]")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = equivalent_strain * e_0, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'vertical')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("Equivalent Strain")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = resultant_displacement, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'vertical')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("Resultant Displacement [mm]")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = sigma_1 * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("sigma_1 ")

plt.show()

plt.figure(dpi = 300)

sc = plt.scatter(x_te, y_te, c = sigma_3 * s_0 * 1e-6, cmap = 'jet', s = 1, edgecolor = 'none', alpha = 1)

plt.colorbar(orientation = 'horizontal')

plt.gca().set_aspect('equal', adjustable = 'box')

plt.axis('off')

plt.title("sigma_3 ")

plt.show()

Compression