Tuesday, August 13, 2024

Teaching Linear Elasticity to Sand: Data-Less Physics Informed Neural Network (Includes Free PyTorch Code)

This post is about solving linear ⏐ elasticity 〰️ problems using PINNs. As you can guess from the title, after successfully teaching neural 🧠 networks to solve the unsolvable i.e. the Navier-Stokes 🌬 equations, yours truly has now taught 🧑‍🏫 linear elasticity to the neural networks 🖧 as well. The results of a simple case under plane-strain assumption as defined in Fig. 1 are presented in Fig. 2. The results are in excellent agreement 🤝 with those obtained from a "shall not be named" commercial 💸 software.

The equations of motion are in the form of Young's modulus and Poisson's ratio. The material properties are that of mild-steel. Of course the equations are scaled and non-dimensionalized as the magnitude difference between Young's modulus and Poisson's ratio for exemple is way too large and this becomes a problem while trying to minimize 📉 the custom loss 🗠 function. The equations of motion are in strong form and do not require a mesh.

Cite as: Fahad Butt (2024). LE-PINN (https://whenrobotslove.blogspot.com/2024/08/teaching-linear-elasticity-to-sand-data.html), Blogger. Retrieved Month Date, Year

Fig. 1, The boundary conditions

Fig. 2, PINN VS FEM

The coloring scale for each photo is same and are mentioned within the code 01!

The code 🖥 has now the capability to solve problems of symmetric tension or compression. For example a flat plate being compressed or put in tension without no constraints. Code 02 is mentioned below. The results from Code 02 are shown in Fig. 4 while the boundary conditions are shown in Fig. 3. The colocation points are shown in Fig. 5. Within Fig.4, blue means low and red means high values. Scales are mention in the caption. Fig. 6 shows unloaded and loaded (red) plates. It can be clearly seen that the plate gets compressed in horizontal axis and expands in the vertical axis, as expected in this set of loading.

Fig. 3, The boundary conditions

Fig. 4, Scales: u = -0.002 - 0.002 mm; v = -0.0002 - 0.0002 mm; σxx = -1.00015 - -1.00005 MPa; σyy = -4e-5 - 1e-5 MPa [essentially zero]; σxy = -4e-5 - 4e-5 MPa [essentially zero]; εxx = -4.553e-6 - -4.546e-6; εyy = 1.948e-6 - 1.951e-6; εxy = -2e-9 - 2e-9

Fig. 5, A comparison

The code (Code 03) is now capable to learn about stresses, strains and displacements in a plate with a hole with considerable accuracy! The results of such a case are shown in Fig. 7. The boundary conditions are shown in Fig. 6.

Fig. 6, The boundary conditions

Fig. 7, The results

The code is now capable of predicting stresses in plates with one end fixed and one end loaded in tension or compression. Of course! There is a hole 🕳 in the center 😆. The code is similar to Code 03, and will not be shared here. The results and boundary conditions are shown in Fig. 8.

Fig. 8, The results and boundary conditions

Code 01

#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

# import torch_directml # enable for AMD and Intel gpu

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

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_bc = u[lee[:, 0] == 0] # x = 0

v_left_bc = v[lee[:, 0] == 0] # x = 0

sxy_left_bc = sxy[lee[:, 0] == 0] # x = 0

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

sxy_bottom_bc = sxy[lee[:, 1] == 0] # y = 0

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

syy_bottom_bc = syy[lee[:, 1] == 0] # y = 0

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, \

u_left_bc, v_left_bc, \

sxy_left_bc, sxy_right_bc, sxy_bottom_bc, sxy_top_bc, \

syy_bottom_bc, syy_top_bc, \

sxx_right_bc

#%% define loss function

def loss_fn(lee):

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, \

u_left_bc, v_left_bc, \

sxy_left_bc, sxy_right_bc, sxy_bottom_bc, sxy_top_bc, \

syy_bottom_bc, syy_top_bc, \

sxx_right_bc = 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_bc, torch.zeros_like(u_left_bc)) # u boundary condition loss

v_bc_loss = nn.MSELoss()(v_left_bc, torch.zeros_like(v_left_bc)) # v boundary condition 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_left_bc, torch.zeros_like(sxy_left_bc)) + nn.MSELoss()(sxy_right_bc, torch.zeros_like(sxy_right_bc)) \

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

syy_bc_loss = nn.MSELoss()(syy_bottom_bc, torch.zeros_like(syy_bottom_bc)) + 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) / 30 # 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) / 100 # 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 = 1 # 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

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

num_training_points = 10 # per length

# Generate points for boundaries

x_l = np.random.uniform(0, 0, round(num_training_points * D)) # left boundary

y_l = np.random.uniform(0, D, round(num_training_points * D))

x_r = np.random.uniform(L, L, round(num_training_points * D)) # right boundary

y_r = np.random.uniform(0, D, round(num_training_points * D))

x_b = np.random.uniform(0, L, round(num_training_points * L)) # bottom boundary

y_b = np.random.uniform(0, 0, round(num_training_points * L))

x_t = np.random.uniform(0, L, round(num_training_points * L)) # top boundary

y_t = np.random.uniform(D, D, round(num_training_points * L))

x_m = np.random.uniform(0, L, 200)

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

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 mesh

plt.figure(dpi = 300)

plt.scatter(x, y, s = 0.1, c = 'k', 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 coordinates

nn.Tanh(), # hidden layer

nn.Linear(h, h),

nn.Tanh(), # hidden layer

nn.Linear(h, 8)).to(device) # 8 outputs for 3 stresses, 3 strains, 2 displacements

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, 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

#%% train using Adam

resume_training = False # True if resuming training

start_epoch = 1

optimizer = optim.Adam(model.parameters(), lr = 0.00001) # 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(200001):

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

#%% prediction

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_training_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, 10000)

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

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]

v = predict[:, 1]

sxx = predict[:, 2]

syy = predict[:, 3]

sxy = predict[:, 4]

exx = predict[:, 5]

eyy = predict[:, 6]

exy = predict[:, 7]

# plotting

plt.figure(dpi = 300)

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

alpha = 1, vmin = 0, vmax = 0.0045)

plt.colorbar(orientation = 'vertical')

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, vmin = -0.0004, vmax = 0.0004)

plt.colorbar(orientation = 'vertical')

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, vmin = 0.9, vmax = 1.1)

plt.colorbar(orientation = 'vertical')

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, vmin = 0, vmax = 0.4)

plt.colorbar(orientation = 'vertical')

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, vmin = -0.4, vmax = 0.4)

plt.colorbar(orientation = 'vertical')

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, vmin = 3.5e-6, vmax = 5e-6)

plt.colorbar(orientation = 'vertical')

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, vmin = -2.7e-6, vmax = 0)

plt.colorbar(orientation = 'vertical')

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, vmin = -1e-6, vmax = 1e-6)

plt.colorbar(orientation = 'vertical')

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

plt.axis('off')

plt.title("gxy")

plt.show()

Code 02

# 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.