PyTorch PINN: Inaccurate sigma_xx (stress) predictions while other parameters match FEM results

[AP1S1T]

Select Topic Area

Question

Body

I'm implementing a Physics-Informed Neural Network (PINN) to solve a 2D elasticity problem, comparing results with FEM analysis. While most parameters (ux, uy, sigma_yy, sigma_xy) closely match the FEM results, the sigma_xx predictions are significantly off.

Problem Description

• Domain: Irregular polygon with fixed bottom edge and traction-free top surface
• Material properties: E = 5.0, ν = 0.3
• Loading: Body force in y-direction (γ = -1.8)

Network Architecture

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.path import Path
import pandas as pd
from torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingLR
import copy
# Set random seed for reproducibility
torch.manual_seed(0)
np.random.seed(0)

# Check if GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Set default tensor type to float32
torch.set_default_tensor_type(torch.FloatTensor)

# Define the neural network architecture
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.hidden1 = nn.Linear(2, 64)
        self.hidden2 = nn.Linear(64, 128)
        self.hidden3 = nn.Linear(128, 256)
        self.hidden4 = nn.Linear(256, 128)  
        self.hidden5 = nn.Linear(128, 64)
        self.output = nn.Linear(64, 2)
    def forward(self, x):
        x = torch.tanh(self.hidden1(x))
        x = torch.tanh(self.hidden2(x))
        x = torch.tanh(self.hidden3(x))
        x = torch.tanh(self.hidden4(x))
        x = torch.tanh(self.hidden5(x))
        return self.output(x)  # return [u, v]
    # Define the problem domain using the given vertices
vertices = np.array([
    [0, 0],
    [0, 3],
    [3, 3],
    [5, 2],
    [7, 2],
    [7,0],
    [0,0]  # Closing the polygon
],dtype=np.float32)
path = Path(vertices)
def in_domain(x, y):
    points = np.column_stack((x.cpu().numpy(), y.cpu().numpy()))
    return torch.tensor(path.contains_points(points), dtype=torch.bool, device=device)

# Define boundary conditions
def BC_bottom(x, y):
    return ((y == 0) & (x >= 0) & (x <= 7)).squeeze()

def BC_left(x, y):
    return ((x == 0) & (y >= 0) & (y <= 3)).squeeze()

def BC_top(x, y):
    return (((y == 3) & (x >= 0) & (x <= 3)) | 
            ((y - 4.5) * -2 == (x - 0)) & (x > 3) & (x < 5) |
            (y == 2) & (x >= 5) & (x <= 7)).squeeze()

def BC_right(x, y):
    return ((x == 7) & (y >= 0) & (y <= 2)).squeeze()

def BC_Slope(x, y):
    return (((y - 4.5) * -2 == (x - 0)) & (x > 3) & (x < 5)).squeeze()


class HuberLoss:
    def __init__(self, delta=1.0):
        self.delta = delta
        
    def __call__(self, y_true, y_pred):
        error = torch.abs(y_true - y_pred)
        quadratic = torch.min(error, torch.tensor(self.delta))
        linear = error - quadratic
        return torch.mean(0.5 * quadratic ** 2 + self.delta * linear)

class MultiViewPINN:
    def __init__(self):
        # Keep track of learned knowledge from different views
        self.view_knowledge = {}
        self.huber = HuberLoss(delta=1.0)
        
    def transfer_knowledge(self, source_net, target_net, layer_mapping):
        """Transfer weights from source network to target network"""
        for source_layer, target_layer in layer_mapping:
            target_net.state_dict()[target_layer].data.copy_(
                source_net.state_dict()[source_layer].data
            )

def BC(xy, net):
    x, y = xy[:, 0].unsqueeze(1), xy[:, 1].unsqueeze(1)
    
    output = net(xy)
    u = output[:, 0:1]  # แยก output channel แรกเป็น u
    v = output[:, 1:2]  # แยก output channel ที่สองเป็น v
    
    bc_b = BC_bottom(x, y)
    bc_l = BC_left(x, y)
    bc_t = BC_top(x, y)
    bc_r = BC_right(x, y)
    bc_slope = BC_Slope(x, y)
    
    huber = HuberLoss(delta=1.0)
    
    loss = huber(u[bc_b], torch.zeros_like(u[bc_b])) + huber(v[bc_b], torch.zeros_like(v[bc_b]))
    loss += huber(u[bc_l], torch.zeros_like(u[bc_l]))
    loss += huber(u[bc_r], torch.zeros_like(u[bc_r]))
    
    xy_top = xy[bc_t].requires_grad_(True)
    output_top = net(xy_top)
    u_top = output_top[:, 0:1]
    v_top = output_top[:, 1:2]
    
    u_x_top = torch.autograd.grad(u_top.sum(), xy_top, create_graph=True)[0][:, 0]
    u_y_top = torch.autograd.grad(u_top.sum(), xy_top, create_graph=True)[0][:, 1]
    v_x_top = torch.autograd.grad(v_top.sum(), xy_top, create_graph=True)[0][:, 0]
    v_y_top = torch.autograd.grad(v_top.sum(), xy_top, create_graph=True)[0][:, 1]
    
    E = 5.0  
    nu = 0.3 
    
    sigma_xx_top = (E / ((1 + nu) * (1 - 2 * nu))) * (u_x_top + nu * v_y_top)
    sigma_yy_top = (E / ((1 + nu) * (1 - 2 * nu))) * (v_y_top + nu * u_x_top)
    sigma_xy_top = E / (2 * (1 + nu)) * (u_y_top + v_x_top)
    
    loss += huber(sigma_xy_top, torch.zeros_like(sigma_xy_top))
    loss += huber(sigma_xx_top, torch.zeros_like(sigma_xx_top))
    loss += huber(sigma_yy_top, torch.zeros_like(sigma_yy_top))
    
    xy_right = xy[bc_r].requires_grad_(True)
    output_right = net(xy_right)
    u_right = output_right[:, 0:1]
    v_right = output_right[:, 1:2]
    
    u_y_right = torch.autograd.grad(u_right.sum(), xy_right, create_graph=True)[0][:, 1]
    v_x_right = torch.autograd.grad(v_right.sum(), xy_right, create_graph=True)[0][:, 0]
    
    sigma_xy_right = E / (2 * (1 + nu)) * (u_y_right + v_x_right)
    loss += huber(sigma_xy_right, torch.zeros_like(sigma_xy_right))

    xy_left = xy[bc_l].requires_grad_(True)
    output_left = net(xy_left)
    u_left = output_left[:, 0:1]
    v_left = output_left[:, 1:2]
    
    u_y_left = torch.autograd.grad(u_left.sum(), xy_left, create_graph=True)[0][:, 1]
    v_x_left = torch.autograd.grad(v_left.sum(), xy_left, create_graph=True)[0][:, 0]
    
    sigma_xy_left = E / (2 * (1 + nu)) * (u_y_left + v_x_left)
    loss += huber(sigma_xy_left, torch.zeros_like(sigma_xy_left))

    return loss


def load_fem_data(filename='FEM2_data.csv'):
    df = pd.read_csv(filename)
    x = torch.tensor(df['X'].values, dtype=torch.float32).unsqueeze(1).to(device)
    y = torch.tensor(df['Y'].values, dtype=torch.float32).unsqueeze(1).to(device)
    return x, y


def generate_training_data(fem_data_file, n_boundary=2000):
    x, y = load_fem_data(fem_data_file)
    
    # Keep only points inside the domain
    mask = in_domain(x, y)
    x, y = x[mask], y[mask]
    
    # Generate boundary points
    t = torch.linspace(0, 1, n_boundary, device=device).unsqueeze(1)
    
    # Define the boundary segments
    segments = [
        ([0, 0, 0], [0, 3, 3]),  # Left boundary
        ([0, 3, 5, 7, 7], [3, 3, 2, 2, 0]),  # Top and right boundary
        ([7, 0], [0, 0])  # Bottom boundary
    ]
    
    x_b = []
    y_b = []
    
    for segment in segments:
        x_seg = torch.tensor(np.interp(t.cpu().numpy(), np.linspace(0, 1, len(segment[0])), segment[0]), dtype=torch.float32, device=device)
        y_seg = torch.tensor(np.interp(t.cpu().numpy(), np.linspace(0, 1, len(segment[1])), segment[1]), dtype=torch.float32, device=device)
        x_b.append(x_seg)
        y_b.append(y_seg)
    
    x_b = torch.cat(x_b)
    y_b = torch.cat(y_b)
    
    return x, y, x_b, y_b


def PDE(x, y, net):
    xy = torch.cat([x, y], dim=1)
    xy.requires_grad = True
    
    output = net(xy)
    u = output[:, 0:1]
    v = output[:, 1:2]
    
    u_x = torch.autograd.grad(u.sum(), xy, create_graph=True)[0][:, 0].unsqueeze(1)
    u_y = torch.autograd.grad(u.sum(), xy, create_graph=True)[0][:, 1].unsqueeze(1)
    v_x = torch.autograd.grad(v.sum(), xy, create_graph=True)[0][:, 0].unsqueeze(1)
    v_y = torch.autograd.grad(v.sum(), xy, create_graph=True)[0][:, 1].unsqueeze(1)
    
    E = 5
    nu = 0.3
    gamma = 1.8
    
    sigma_xx = E / (1 - nu**2) * (u_x + nu * v_y)
    sigma_yy = E / (1 - nu**2) * (v_y + nu * u_x)
    sigma_xy = E / (2 * (1 + nu)) * (u_y + v_x)
    
    f_x = torch.zeros_like(x)
    f_y = -gamma * torch.ones_like(y)
    
    R_x = torch.autograd.grad(sigma_xx.sum(), xy, create_graph=True)[0][:, 0].unsqueeze(1) + \
          torch.autograd.grad(sigma_xy.sum(), xy, create_graph=True)[0][:, 1].unsqueeze(1) + f_x
    R_y = torch.autograd.grad(sigma_xy.sum(), xy, create_graph=True)[0][:, 0].unsqueeze(1) + \
          torch.autograd.grad(sigma_yy.sum(), xy, create_graph=True)[0][:, 1].unsqueeze(1) + f_y
    
    huber = HuberLoss(delta=1.0)
    loss_x = huber(R_x, torch.zeros_like(R_x))
    loss_y = huber(R_y, torch.zeros_like(R_y))
    
    return loss_x, loss_y


def train(net, optimizer, n_epochs, fem_data_file):
    multi_view = MultiViewPINN()
    huber = HuberLoss(delta=1.0)
    
    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=1e-6)
    
    best_loss = float('inf')
    patience = 0
    max_patience = 3000
    
    prev_net = None
    
    for epoch in range(n_epochs):
        optimizer.zero_grad()
        
        x, y, x_b, y_b = generate_training_data(fem_data_file, 30000)
        xy = torch.cat([x, y], dim=1)
        xy_b = torch.cat([x_b, y_b], dim=1)
        
        if prev_net is not None:
            layer_mapping = [
                ('hidden1.weight', 'hidden1.weight'),
                ('hidden2.weight', 'hidden2.weight'),
                ('hidden3.weight', 'hidden3.weight')
            ]
            multi_view.transfer_knowledge(prev_net, net, layer_mapping)
        
        loss_pde_x, loss_pde_y = PDE(x, y, net)
        loss_bc = BC(xy_b, net)
        
        loss = loss_pde_x + 2 * loss_pde_y + loss_bc
        
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=0.5)
        
        optimizer.step()
        scheduler.step()
        
        if loss < best_loss:
            best_loss = loss
            patience = 0
            
            prev_net = copy.deepcopy(net)
            
            torch.save({
                'net_state_dict': net.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'loss': loss,
                'epoch': epoch,
            }, 'best_model.pth')
        else:
            patience += 1
            if patience > max_patience:
                print(f"Early stopping at epoch {epoch}")
                break
                
        if (epoch + 1) % 10 == 0:
            print(f'Epoch {epoch+1}/{n_epochs}, Loss: {loss.item():.6f}, '
                  f'PDE_x: {loss_pde_x.item():.6f}, PDE_y: {loss_pde_y.item():.6f}, '
                  f'BC: {loss_bc.item():.6f}, LR: {scheduler.get_last_lr()[0]:.6f}')