Untitled

import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

# Define the PINN architecture
class PINN(nn.Module):
    def __init__(self, input_dim=2, output_dim=2, hidden_layers=[50]*6):
        super(PINN, self).__init__()
        layers = [nn.Linear(input_dim, hidden_layers[0]), nn.Tanh()]
        for i in range(len(hidden_layers)-1):
            layers += [nn.Linear(hidden_layers[i], hidden_layers[i+1]), nn.Tanh()]
        layers.append(nn.Linear(hidden_layers[-1], output_dim))
        self.network = nn.Sequential(*layers)
        
    def forward(self, x):
        return self.network(x)

# Instantiate the PINN model
model = PINN(input_dim=2, output_dim=2, hidden_layers=[50]*6)

# Define the PDE residual function (linear elasticity equation in 2D)
def pde_residual(u_xx, u_yy, u_xy, f):
    return u_xx + u_yy + u_xy - f

# Compute the loss based on the PDE
def compute_pde_loss(model, x_colloc, E, nu, L):
    u = model(x_colloc)
    
    # Compute strain tensor ε
    u_x = torch.autograd.grad(u[:, 0], x_colloc, grad_outputs=torch.ones_like(u[:, 0]), create_graph=True)[0]
    u_y = torch.autograd.grad(u[:, 1], x_colloc, grad_outputs=torch.ones_like(u[:, 1]), create_graph=True)[0]
    
    u_xx = torch.autograd.grad(u_x[:, 0], x_colloc, grad_outputs=torch.ones_like(u_x[:, 0]), create_graph=True)[0][:, 0]
    u_yy = torch.autograd.grad(u_y[:, 1], x_colloc, grad_outputs=torch.ones_like(u_y[:, 1]), create_graph=True)[0][:, 1]
    u_xy = (torch.autograd.grad(u_x[:, 1], x_colloc, grad_outputs=torch.ones_like(u_x[:, 1]), create_graph=True)[0][:, 0] + 
            torch.autograd.grad(u_y[:, 0], x_colloc, grad_outputs=torch.ones_like(u_y[:, 0]), create_graph=True)[0][:, 1]) / 2
    
    # Lamé parameters
    lam = E * nu / ((1+nu) * (1-2*nu))
    mu = E / (2 * (1 + nu))
    
    f = torch.zeros_like(u_x[:, 0])
    
    pde_residuals = pde_residual(u_xx, u_yy, u_xy, f)
    
    return torch.mean(pde_residuals**2)

# Compute the loss based on boundary conditions
def compute_data_loss(model, x_fix, y_fix):
    u_fix = model(x_fix)
    
    loss_data = (u_fix[:, 0] - y_fix[:, 0]).pow(2).mean() + (u_fix[:, 1] - y_fix[:, 1]).pow(2).mean()
    
    return loss_data

# Generate collocation points
def generate_collocation_points(L, b, h, N=100):
    x = np.linspace(0, L, N)
    y = np.linspace(-h/2, h/2, N)
    
    X, Y = np.meshgrid(x, y)
    collocation_points = np.vstack([X.flatten(), Y.flatten()]).T
    return torch.tensor(collocation_points, dtype=torch.float32)

# Generate fixed boundary points
def generate_fixed_boundary_points(L, b, h):
    x_fix = np.zeros((1, 2))
    x_fix[0, 0] = 0.0
    x_fix[0, 1] = -h/2
    
    y_fix_x = np.zeros(1)
    y_fix_y = np.zeros(1)
    
    return torch.tensor(x_fix, dtype=torch.float32), torch.tensor(y_fix_x, dtype=torch.float32), torch.tensor(y_fix_y, dtype=torch.float32)

# Hyperparameters
learning_rate = 0.001
lambda_D = 1.0
epochs = 10000

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

L = 182.88e-3  # Length in meters
b = 25.4e-3    # Width in meters
h = 12.7e-3    # Height in meters

collocation_points = generate_collocation_points(L, b, h)
x_fix, y_fix_x, y_fix_y = generate_fixed_boundary_points(L, b, h)

E = 15e9       # Young's modulus in Pa
nu = 0.3       # Poisson's ratio

for epoch in range(epochs):
    optimizer.zero_grad()
    
    loss_pde = compute_pde_loss(model, collocation_points, E, nu, L)
    loss_data = compute_data_loss(model, x_fix, torch.hstack([y_fix_x[:, None], y_fix_y[:, None]]))
    
    total_loss = loss_pde + lambda_D * loss_data
    
    total_loss.backward()
    optimizer.step()
    
    if epoch % 1000 == 0:
        print(f'Epoch [{epoch}/{epochs}], Loss: {total_loss.item():.4f}, PDE Loss: {loss_pde.item():.4f}, Data Loss: {loss_data.item():.4f}')