├── environment.yml ├── README.md ├── OpenPyStruct_FrameOpt_Discrete_Beta.py ├── OpenPyStruct_BeamOpt_training_SingleCore.py ├── OpenPyStruct_BeamOpt_training_MultiCore.py ├── OpenPyStruct_BeamOpt.py ├── OpenPyStruct_BeamOpt_training_GPU.py ├── OpenPyStruct_GNN_MultiCase_Beta.py ├── OpenPyStruct_FNN_MultiCase.py ├── OpenPyStruct_FNO_MultiCase_Beta.py └── OpenPyStruct_Bayesian_TFDModule_MultiCase_Beta.py /environment.yml: -------------------------------------------------------------------------------- 1 | name: OpenPyStruct 2 | channels: 3 | - defaults 4 | - conda-forge 5 | dependencies: 6 | - python=3.8 7 | - numpy 8 | - pytorch 9 | - matplotlib 10 | - seaborn 11 | - scikit-learn 12 | - pip 13 | - pip: 14 | - openseespy -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # **OpenPyStruct** 2 | 3 | **OpenPyStruct** is an open-source toolkit designed for machine learning-based structural optimization. Leveraging **Physics-Informed Neural Networks (PINNs)**, **Transformer-Diffusion Modules**, and other state-of-the-art techniques, the framework provides powerful tools for tackling **single and multi-load case optimization** problems with diverse boundary and loading conditions. 4 | 5 | ## **Table of Contents** 6 | 7 | - [**OpenPyStruct**](#openpystruct) 8 | - [**Table of Contents**](#table-of-contents) 9 | - [**Features**](#features) 10 | - [**Requirements**](#requirements) 11 | - [Option 1, Manual Install](#option-1-manual-install) 12 | - [Option 2, Conda Environment Install](#option-2-conda-environment-install) 13 | 23 | - [**License**](#license) 24 | 25 | --- 26 | 27 | ## **Features** 28 | 29 | - **Classical Model-Based Structural Optimizers:** Used for generating training data 30 | - **OpenSeesPy Integration:** Facilitates physics-based finite element simulations. 31 | - **Multi-Core and GPU-Accelerated Optimization:** Enables large-scale data generation and rapid computations. 32 | - **Physics-Informed Neural Networks (PINNs) - ML Structural Optimization:** Embeds structural mechanics into the learning process for highly accurate predictions. 33 | - **Transformer-Diffusion Modules - ML Structural Optimization:** Incorporates advanced attention mechanisms and diffusion-based techniques for modeling complex structural behavior. 34 | - **Feedforward Neural Networks (FNNs) - ML Structural Optimization:** Provides scalable solutions for simpler structural optimization tasks. 35 | - **Flexible Loss Functions and Parameter Design:** Supports user-defined constraints, objectives, and optimization goals. 36 | 37 | --- 38 | 39 | ## **Requirements** 40 | 41 | 42 | ### Option 1, Manual Install 43 | 44 | Create a new conda environment with python 3.8+ installed 45 | 46 | ```zsh 47 | conda install python=3.8 48 | ``` 49 | 50 | First, install OpenSeesPy: 51 | 52 | ```zsh 53 | pip install openseespy 54 | ``` 55 | 56 | Install rest of packages 57 | 58 | ```zsh 59 | conda install numpy, torch, matplotlib, seaborn, scikit-learn 60 | ``` 61 | 62 | 63 | 64 | ### Option 2, Conda Environment Install 65 | 66 | To create a Conda environment with all dependencies run: 67 | 68 | ```zsh 69 | conda env create -f environment.yml 70 | conda activate OpenPyStruct 71 | ``` 72 | 73 | --- 74 | 141 | 142 | --- 143 | 144 | ## **License** 145 | 146 | **MIT License** 147 | 148 | ``` 149 | MIT License 150 | 151 | Copyright (c) 2025 Danny Smyl 152 | 153 | 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: 154 | 155 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 156 | 157 | 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. 158 | ``` 159 | -------------------------------------------------------------------------------- /OpenPyStruct_FrameOpt_Discrete_Beta.py: -------------------------------------------------------------------------------- 1 | ################################################################ 2 | #### OpenPyStruct Single Frame Optimizer #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ################################################################ 5 | 6 | import openseespy.opensees as ops 7 | import random 8 | import numpy as np 9 | import torch 10 | import matplotlib.pyplot as plt 11 | 12 | ############################## 13 | # User-Defined and Material Parameters 14 | ############################## 15 | 16 | # Maximum numbers for random generation: 17 | max_bays = 10 # Maximum number of bays (span divisions) 18 | max_stories = 10 # Maximum number of stories 19 | 20 | # Geometric dimensions (in meters) 21 | bay_width = 6.0 # Width of each bay 22 | story_height = 3.0 # Height of each story 23 | 24 | # Material and cross-sectional properties: 25 | E = 200e9 # Young's Modulus (Pa) 26 | nu = 0.3 # Poisson's ratio 27 | G = E / (2*(1+nu)) # Shear Modulus (Pa) 28 | A = 0.02 # Cross-sectional area (m^2) used in the analysis (assumed constant) 29 | I0 = 5e-4 # Initial guess for moment of inertia (m^4) 30 | 31 | # Loss function coefficients 32 | alpha_moment = 1e-2 33 | alpha_shear = 1e-2 34 | k = 0.03 # Coefficient to define local area: A_local = k * sqrt(I) 35 | 36 | # Applied loads: 37 | lateral_load = 1e4 # Lateral nodal load (N) applied at left-hand side nodes (x = 0) 38 | vertical_load = -1e4 # Uniform vertical load (N) applied to beam elements (negative for downward) 39 | 40 | # Optimization settings 41 | num_epochs = 5000 42 | lr = 0.005 43 | tolerance = 1e-3 44 | patience = 10 45 | 46 | ############################## 47 | # Random Generation of Frame Geometry 48 | ############################## 49 | 50 | num_bays = random.randint(1, max_bays) 51 | num_stories = random.randint(1, max_stories) 52 | print(f"Generated frame with {num_bays} bay(s) and {num_stories} story(ies).") 53 | 54 | # Create node grid and store original coordinates (nodes numbered from 1) 55 | node_coords = {} 56 | num_nodes = (num_stories + 1) * (num_bays + 1) 57 | for i in range(num_stories + 1): 58 | for j in range(num_bays + 1): 59 | tag = i * (num_bays + 1) + j + 1 60 | x = j * bay_width 61 | y = i * story_height 62 | node_coords[tag] = (x, y) 63 | 64 | # Count elements: 65 | # Columns: vertical members (each column between successive stories) 66 | num_columns = num_stories * (num_bays + 1) 67 | # Beams: horizontal members (each beam on each story except the ground) 68 | num_beams = num_stories * num_bays 69 | total_elems = num_columns + num_beams 70 | 71 | ############################## 72 | # OpenSees Model Setup Functions 73 | ############################## 74 | 75 | def setup_frame_model(I_tensor): 76 | """ 77 | Build the 2D frame model in OpenSees using the moment-of-inertia vector I_tensor. 78 | The first num_columns entries of I_tensor are used for columns; the remainder for beams. 79 | 80 | Loads are applied as follows: 81 | - Nodal lateral loads on all left-hand side nodes (x=0) except ground (y=0). 82 | - Uniform vertical loads (using eleLoad) on all beam elements. 83 | """ 84 | ops.wipe() 85 | ops.model('basic', '-ndm', 2, '-ndf', 3) 86 | ops.geomTransf('Linear', 1) 87 | 88 | # Create nodes using stored coordinates 89 | for tag, (x, y) in node_coords.items(): 90 | ops.node(tag, x, y) 91 | 92 | # Fix nodes at the ground (nodes with y=0) 93 | for tag, (x, y) in node_coords.items(): 94 | if y == 0.0: 95 | ops.fix(tag, 1, 1, 1) 96 | 97 | # Create column elements (vertical members) 98 | elem_tag = 1 # element numbering starts at 1 99 | for i in range(num_stories): 100 | for j in range(num_bays + 1): 101 | node_i = i * (num_bays + 1) + j + 1 102 | node_j = (i + 1) * (num_bays + 1) + j + 1 103 | I_val = I_tensor[elem_tag - 1].item() # columns occupy indices 0 .. num_columns-1 104 | ops.element('elasticBeamColumn', elem_tag, node_i, node_j, A, E, I_val, 1) 105 | elem_tag += 1 106 | 107 | # Create beam elements (horizontal members) 108 | for i in range(1, num_stories + 1): 109 | for j in range(num_bays): 110 | node_i = i * (num_bays + 1) + j + 1 111 | node_j = i * (num_bays + 1) + j + 2 112 | I_val = I_tensor[elem_tag - 1].item() # beams occupy remaining indices 113 | ops.element('elasticBeamColumn', elem_tag, node_i, node_j, A, E, I_val, 1) 114 | elem_tag += 1 115 | 116 | # Define load pattern: 117 | ops.timeSeries('Linear', 1) 118 | ops.pattern('Plain', 1, 1) 119 | 120 | # --- Apply lateral loads on left-hand side nodes (x = 0) except the ground --- 121 | # Left-hand side nodes are those with j == 0 and y > 0. 122 | for tag, (x, y) in node_coords.items(): 123 | if x == 0.0 and y != 0.0: 124 | ops.load(tag, lateral_load, 0.0, 0.0) 125 | 126 | # --- Apply uniform vertical load to all beam elements --- 127 | # Beam elements have tags from num_columns+1 to total_elems. 128 | for ele in range(num_columns + 1, total_elems + 1): 129 | # For a uniform beam load (applied in the local vertical direction), 130 | # the command below applies vertical load (downward) at both ends of the element. 131 | ops.eleLoad('-ele', ele, '-type', '-beamUniform', vertical_load, vertical_load) 132 | 133 | # Analysis settings 134 | ops.system('BandGeneral') 135 | ops.numberer('RCM') 136 | ops.constraints('Plain') 137 | ops.integrator('LoadControl', 1.0) 138 | ops.algorithm('Newton') 139 | ops.analysis('Static') 140 | 141 | def compute_combined_loss(I_tensor): 142 | """ 143 | Run the analysis and compute the combined loss: 144 | total_loss = sum(I) + α_moment*(∑[bending_moment²/(2·E·I)]) + α_shear*(∑[shear_force²/(G·A_local)]) 145 | where A_local = k * sqrt(I). 146 | """ 147 | bending_energy = 0.0 148 | shear_energy = 0.0 149 | for elem_id in range(1, total_elems + 1): 150 | # Retrieve element forces (assumed ordering: [axial, shear, bending moment, ...]) 151 | response = ops.eleResponse(elem_id, 'forces') 152 | shear_force = response[1] 153 | bending_moment = response[2] 154 | I_val = I_tensor[elem_id - 1] 155 | bending_energy += (bending_moment**2) / (2 * E * I_val + 1e-8) 156 | A_local = k * (I_val**0.5) 157 | shear_energy += (shear_force**2) / (G * A_local) 158 | primary_loss = torch.sum(I_tensor) 159 | total_loss = primary_loss + alpha_moment * bending_energy + alpha_shear * shear_energy 160 | return total_loss, primary_loss, alpha_moment * bending_energy, alpha_shear * shear_energy 161 | 162 | ############################## 163 | # Optimization Initialization 164 | ############################## 165 | 166 | # Initialize moment-of-inertia values for all elements (columns + beams) 167 | I_values = [I0 for _ in range(total_elems)] 168 | I_tensor = torch.tensor(I_values, dtype=torch.float32, requires_grad=True) 169 | 170 | optimizer = torch.optim.Adam([I_tensor], lr=lr) 171 | loss_history = [] 172 | best_loss = float('inf') 173 | no_improve = 0 174 | 175 | ############################## 176 | # Optimization Loop 177 | ############################## 178 | 179 | for epoch in range(num_epochs): 180 | optimizer.zero_grad() 181 | ops.wipe() # Clear previous model 182 | setup_frame_model(I_tensor) # Rebuild the model with current I values 183 | ops.analyze(1) 184 | total_loss, primary_loss, bending_loss, shear_loss = compute_combined_loss(I_tensor) 185 | total_loss.backward() 186 | optimizer.step() 187 | 188 | # Prevent I from going to zero or negative 189 | with torch.no_grad(): 190 | I_tensor.clamp_(min=1e-8) 191 | 192 | current_loss = total_loss.item() 193 | loss_history.append(current_loss) 194 | 195 | if current_loss < best_loss - tolerance: 196 | best_loss = current_loss 197 | no_improve = 0 198 | else: 199 | no_improve += 1 200 | 201 | if (epoch + 1) % 10 == 0: 202 | print(f"Epoch {epoch+1:3d}: Total Loss = {current_loss:.6e}, Primary = {primary_loss.item():.6e}") 203 | 204 | if no_improve >= patience: 205 | print(f"Stopping early at epoch {epoch+1} (no improvement for {patience} epochs).") 206 | break 207 | 208 | print(f"\nOptimization complete. Best Loss: {best_loss:.6e}") 209 | 210 | # Extract the optimized I values as a NumPy array. 211 | opt_I = I_tensor.detach().numpy() 212 | 213 | ############################## 214 | # Re-run the Analysis with Optimized I 215 | ############################## 216 | 217 | ops.wipe() 218 | setup_frame_model(I_tensor) 219 | ops.analyze(1) 220 | 221 | ############################## 222 | # Plot Loss History 223 | ############################## 224 | 225 | plt.figure(figsize=(8, 5)) 226 | plt.plot(loss_history, 'b-', linewidth=2) 227 | plt.xlabel("Epoch") 228 | plt.ylabel("Total Loss") 229 | plt.title("Optimization Loss History") 230 | plt.grid(True) 231 | plt.show() 232 | 233 | ############################## 234 | # Visualization: Plot the Frame with Optimized I Distribution 235 | ############################## 236 | 237 | plt.figure(figsize=(12, 8)) 238 | 239 | # Plot the original (undeformed) frame in light gray for reference 240 | for i in range(num_stories): 241 | for j in range(num_bays + 1): 242 | node_i = i * (num_bays + 1) + j + 1 243 | node_j = (i + 1) * (num_bays + 1) + j + 1 244 | x_i, y_i = node_coords[node_i] 245 | x_j, y_j = node_coords[node_j] 246 | plt.plot([x_i, x_j], [y_i, y_j], '--', color='lightgray') 247 | for i in range(1, num_stories + 1): 248 | for j in range(num_bays): 249 | node_i = i * (num_bays + 1) + j + 1 250 | node_j = i * (num_bays + 1) + j + 2 251 | x_i, y_i = node_coords[node_i] 252 | x_j, y_j = node_coords[node_j] 253 | plt.plot([x_i, x_j], [y_i, y_j], '--', color='lightgray') 254 | 255 | # Plot the optimized frame: scale line width based on optimized I (cube-root scaling) 256 | max_I_val = max(opt_I) 257 | elem_index = 0 258 | 259 | # Plot columns (in blue) 260 | for i in range(num_stories): 261 | for j in range(num_bays + 1): 262 | node_i = i * (num_bays + 1) + j + 1 263 | node_j = (i + 1) * (num_bays + 1) + j + 1 264 | x_i, y_i = node_coords[node_i] 265 | x_j, y_j = node_coords[node_j] 266 | I_elem = opt_I[elem_index] 267 | lw = 15 * (I_elem / max_I_val)**(1/3) 268 | plt.plot([x_i, x_j], [y_i, y_j], 'b-', linewidth=lw, 269 | label="Column" if elem_index == 0 else "") 270 | elem_index += 1 271 | 272 | # Plot beams (in red) 273 | for i in range(1, num_stories + 1): 274 | for j in range(num_bays): 275 | node_i = i * (num_bays + 1) + j + 1 276 | node_j = i * (num_bays + 1) + j + 2 277 | x_i, y_i = node_coords[node_i] 278 | x_j, y_j = node_coords[node_j] 279 | I_elem = opt_I[elem_index] 280 | lw = 15 * (I_elem / max_I_val)**(1/3) 281 | plt.plot([x_i, x_j], [y_i, y_j], 'r-', linewidth=lw, 282 | label="Beam" if elem_index == num_columns else "") 283 | elem_index += 1 284 | 285 | plt.xlabel("X Coordinate (m)") 286 | plt.ylabel("Y Coordinate (m)") 287 | plt.title("Frame Structure with Optimized Moment of Inertia Distribution") 288 | plt.legend() 289 | plt.grid(True) 290 | plt.axis('equal') 291 | plt.show() 292 | 293 | -------------------------------------------------------------------------------- /OpenPyStruct_BeamOpt_training_SingleCore.py: -------------------------------------------------------------------------------- 1 | ################################################################ 2 | #### OpenPyStruct Single-Core Optimizer / Data Generator #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ################################################################ 5 | 6 | import os 7 | import openseespy.opensees as ops 8 | import numpy as np 9 | import random 10 | import torch 11 | import json 12 | from torch.optim.lr_scheduler import ExponentialLR 13 | import time 14 | from tqdm import tqdm # Progress bar library 15 | 16 | ############## 17 | # Parameters # 18 | ############## 19 | 20 | E = 200e9 # Young's Modulus (Pa) 21 | nu = 0.3 # Poisson ratio 22 | G = E / (2 * (1 + nu)) # Shear modulus (Pa) 23 | A = 0.01 # Cross-sectional area (m^2) 24 | L_max = 200.0 # Length of the beam (m) 25 | num_nodes = 101 # Number of nodes 26 | num_elements = num_nodes - 1 # Number of elements 27 | N_rollers_max = 4 # Number of additional roller supports 28 | M_forces_max = 4 # Number of point forces 29 | L_min = 15 # Min distance between rollers (m) 30 | max_force = -355857 # Max point load (N) - 80,0000 lb semi 31 | min_force = max_force / 10 32 | uniform_udl = -1000 # Uniformly distributed load (N) 33 | I_0 = 0.5 # Initial guess for I 34 | 35 | ## Optimization parameters ## 36 | max_e = 600 37 | lr = 0.01 # Initial optimization learning rate 38 | gamma = 0.98 # Learning rate decay rate 39 | alpha_moment = 1e-2 # Coefficient for bending energy loss term 40 | alpha_shear = 1e-2 # Coefficient for shear energy loss term 41 | 42 | # Parameters for stopping criterion # 43 | tolerance = 5e-3 # Minimum improvement in loss 44 | patience = 5 # Number of epochs to wait for improvement 45 | 46 | num_samples = 100000 # Number of training data samples to generate 47 | 48 | random_bridge = 0 # Input 1 if you want the bridge length and roller locations randomized 49 | flag = random_bridge # Flag for generating fixed or random bridge 50 | 51 | # Single-core execution (no parallel processing) 52 | # num_workers = 22 # Removed as we're using single-core 53 | 54 | ################################################################################################# 55 | ## Generate a random geometry and support condition (only used if not randomizing a bridge) ## 56 | ################################################################################################# 57 | 58 | L = L_max 59 | node_positions = np.linspace(0, L, num_nodes) 60 | 61 | # Fixed roller node locations 62 | roller_nodes = [10, 30, 70, 85, num_nodes - 1] # Assuming node numbering starts at 1 63 | available_nodes = list(range(2, num_nodes)) # Exclude node 1 64 | for node in roller_nodes: 65 | if node in available_nodes: 66 | available_nodes.remove(node) 67 | 68 | ################################################################################################# 69 | ################################################################################################# 70 | ################################################################################################# 71 | 72 | # Training data structure 73 | training_data = { 74 | "roller_x_locations": [], 75 | "force_x_locations": [], 76 | "force_values": [], 77 | "I_values": [], 78 | "shear_forces": [], 79 | "bending_moments": [], 80 | "node_positions": [], 81 | "roller_nodes": [], 82 | "force_nodes": [], 83 | "num_nodes": [], 84 | "L": [], 85 | "rotations": [], 86 | "deflections": [], 87 | } 88 | 89 | def setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl): 90 | """ 91 | Set up the OpenSees model using updated moments of inertia and apply uniform UDL to all elements. 92 | """ 93 | ops.model('basic', '-ndm', 2, '-ndf', 3) 94 | 95 | # Define nodes 96 | for i, x in enumerate(node_positions): 97 | ops.node(i + 1, x, 0.0) 98 | 99 | # Apply supports 100 | ops.fix(1, 1, 1, 0) # First node (constrained in X-Y, rotationally free) 101 | for node in roller_nodes: 102 | ops.fix(node, 0, 1, 0) # All other nodes (rotationally free, free in X) 103 | 104 | # Define elements with updated moments of inertia 105 | ops.geomTransf('Linear', 1) 106 | for i in range(len(node_positions) - 1): 107 | ops.element('elasticBeamColumn', i + 1, i + 1, i + 2, A, E, I_tensor[i].item(), 1) 108 | 109 | # Apply point loads 110 | ops.timeSeries('Linear', 1) 111 | ops.pattern('Plain', 1, 1) 112 | for node, force in zip(force_nodes, force_values): 113 | ops.load(node, 0.0, force, 0.0) 114 | 115 | # Apply uniform UDL to all elements 116 | for elem_id in range(1, len(node_positions)): 117 | ops.eleLoad('-ele', elem_id, '-type', '-beamUniform', uniform_udl, uniform_udl) 118 | 119 | # Static analysis setup 120 | ops.system('BandSPD') 121 | ops.numberer('RCM') 122 | ops.constraints('Plain') 123 | ops.integrator('LoadControl', 1.0) 124 | ops.algorithm('Linear') 125 | 126 | def generate_sample(sample_idx, num_nodes, flag, L, node_positions, roller_nodes, available_nodes, patience=10): 127 | """ 128 | Generate a single sample for training data with optimized OpenSeesPy modeling and early stopping. 129 | """ 130 | import openseespy.opensees as ops # Import inside function 131 | 132 | # Random bridge length / rollers/ nodes 133 | if flag == 1: 134 | L = L_min + random.uniform(0, L_max) 135 | node_positions = np.linspace(0, L, num_nodes) 136 | 137 | # Randomize roller node locations 138 | roller_nodes_sample, available_nodes_sample = [], list(range(2, num_nodes)) # Exclude node 1 139 | num_rollers = random.randint(1, N_rollers_max) 140 | 141 | # Add first roller node 142 | first_roller_node = random.choice(available_nodes_sample) 143 | roller_nodes_sample.append(first_roller_node) 144 | available_nodes_sample.remove(first_roller_node) 145 | 146 | # Add additional rollers without minimum distance constraint 147 | for _ in range(num_rollers - 1): 148 | if available_nodes_sample: # Ensure there are available nodes 149 | new_roller_node = random.choice(available_nodes_sample) 150 | roller_nodes_sample.append(new_roller_node) 151 | available_nodes_sample.remove(new_roller_node) 152 | else: 153 | roller_nodes_sample = roller_nodes.copy() 154 | available_nodes_sample = available_nodes.copy() 155 | 156 | # Randomize point forces and their values 157 | num_forces = random.randint(1, M_forces_max) 158 | num_forces = min(num_forces, len(available_nodes_sample)) 159 | force_nodes = random.sample(available_nodes_sample, num_forces) 160 | force_values = [random.uniform(min_force, max_force) for _ in force_nodes] 161 | 162 | # Initialize moments of inertia tensor 163 | I_tensor = torch.tensor([I_0] * num_elements, dtype=torch.float32, requires_grad=True) 164 | 165 | # Setup optimizer and scheduler 166 | optimizer = torch.optim.Adam([I_tensor], lr=lr) 167 | scheduler = ExponentialLR(optimizer, gamma=gamma) 168 | 169 | # Optimization loop with early stopping 170 | best_loss = float('inf') 171 | patience_counter = 0 172 | 173 | # Progress bar for epochs within each sample 174 | for epoch in tqdm(range(max_e), desc=f"Sample {sample_idx+1}/{num_samples}", leave=False): 175 | optimizer.zero_grad() 176 | ops.wipe() # Reset OpenSees model 177 | 178 | setup_model(I_tensor, node_positions, roller_nodes_sample, force_nodes, force_values, A, E, uniform_udl) 179 | 180 | ops.analysis('Static') 181 | try: 182 | ops.analyze(1) 183 | except Exception as e: 184 | print(f"OpenSees analysis failed for sample {sample_idx+1} at epoch {epoch+1}: {e}") 185 | break # Skip to the next sample 186 | 187 | # Compute losses 188 | try: 189 | bending_moments = torch.tensor([ops.eleResponse(i, 'forces')[2] for i in range(1, len(I_tensor) + 1)], dtype=torch.float32) 190 | shear_forces = torch.tensor([ops.eleResponse(i, 'forces')[1] for i in range(1, len(I_tensor) + 1)], dtype=torch.float32) 191 | except Exception as e: 192 | print(f"Failed to retrieve responses for sample {sample_idx+1} at epoch {epoch+1}: {e}") 193 | break # Skip to the next sample 194 | 195 | bending_energy = torch.sum((bending_moments ** 2) / (2 * E * I_tensor + 1e-6)) 196 | A_approx = 0.03 * I_tensor ** 0.5 197 | shear_energy = torch.sum(shear_forces ** 2 / (G * A_approx)) 198 | primary_loss = torch.sum(I_tensor) 199 | total_loss = primary_loss + alpha_moment * bending_energy + alpha_shear * shear_energy 200 | 201 | # Backpropagate and update 202 | total_loss.backward() 203 | optimizer.step() 204 | scheduler.step() 205 | 206 | # Prevent negative inertia values 207 | with torch.no_grad(): 208 | I_tensor.clamp_(min=1e-8) 209 | 210 | # Early stopping logic 211 | if total_loss.item() < best_loss - tolerance: 212 | best_loss = total_loss.item() 213 | patience_counter = 0 214 | else: 215 | patience_counter += 1 216 | 217 | if patience_counter >= patience: 218 | #print(f"Stopping early at epoch {epoch} due to no improvement in {patience} epochs.") 219 | break 220 | 221 | # Extract rotations and deflections 222 | rotations = [] 223 | deflections = [] 224 | for i in range(1, len(node_positions) + 1): 225 | try: 226 | rotations.append(ops.nodeDisp(i, 3)) 227 | except: 228 | rotations.append(0.0) 229 | try: 230 | deflections.append(ops.nodeDisp(i, 2)) 231 | except: 232 | deflections.append(0.0) 233 | 234 | # Return the results 235 | return { 236 | "roller_x_locations": [node_positions[node - 1] for node in roller_nodes_sample], 237 | "force_x_locations": [node_positions[node - 1] for node in force_nodes], 238 | "force_values": force_values, 239 | "I_values": I_tensor.detach().numpy().tolist(), 240 | "shear_forces": shear_forces.detach().tolist(), 241 | "bending_moments": bending_moments.detach().tolist(), 242 | "node_positions": node_positions.tolist(), 243 | "roller_nodes": roller_nodes_sample, 244 | "force_nodes": force_nodes, 245 | "num_nodes": num_nodes, 246 | "L": L, 247 | "rotations": rotations, 248 | "deflections": deflections, 249 | } 250 | 251 | def main(): 252 | start_time = time.time() # Start the timer 253 | 254 | # Initialize progress bar for sample generation 255 | with tqdm(total=num_samples, desc="Generating Samples") as pbar: 256 | for i in range(num_samples): 257 | result = generate_sample(i, num_nodes, flag, L, node_positions, roller_nodes, available_nodes, patience=patience) 258 | for key, value in result.items(): 259 | training_data[key].append(value) 260 | pbar.update(1) # Update the progress bar 261 | 262 | # Save the training data to a JSON file 263 | with open("training_data_PINN_mini.json", "w") as f: 264 | json.dump(training_data, f) 265 | 266 | end_time = time.time() # End the timer 267 | 268 | print("Data generation complete.") 269 | print(f"Total execution time: {end_time - start_time:.2f} seconds") 270 | 271 | if __name__ == "__main__": 272 | main() 273 | 274 | # For sanity: Load the training data from the JSON file # 275 | with open("training_data_PINN_mini.json", "r") as file: 276 | training_data = json.load(file) 277 | 278 | # Print a summary of the dataset 279 | print("Data loaded successfully!") 280 | print(f"Number of samples: {len(training_data['roller_x_locations'])}") 281 | print("Keys available in the dataset:") 282 | for key in training_data.keys(): 283 | print(f"- {key} (Number of entries: {len(training_data[key])})") 284 | 285 | # Access all data into variables 286 | roller_x_locations = training_data["roller_x_locations"] 287 | force_x_locations = training_data["force_x_locations"] 288 | force_values = training_data["force_values"] 289 | I_values = training_data["I_values"] 290 | shear_forces = training_data["shear_forces"] 291 | bending_moments = training_data["bending_moments"] 292 | node_positions = training_data["node_positions"] 293 | roller_nodes = training_data["roller_nodes"] 294 | force_nodes = training_data["force_nodes"] 295 | num_nodes = training_data["num_nodes"] 296 | beam_lengths = training_data["L"] 297 | rotations = training_data["rotations"] 298 | deflections = training_data["deflections"] 299 | -------------------------------------------------------------------------------- /OpenPyStruct_BeamOpt_training_MultiCore.py: -------------------------------------------------------------------------------- 1 | ################################################################ 2 | #### OpenPyStruct Multicore Optimizer / Data Generator #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ################################################################ 5 | 6 | from joblib import Parallel, delayed 7 | import os 8 | import openseespy.opensees as ops 9 | import numpy as np 10 | import random 11 | import torch 12 | import json 13 | from torch.optim.lr_scheduler import ExponentialLR 14 | import time 15 | 16 | ############## 17 | # Parameters # 18 | ############## 19 | 20 | E = 200e9 # Young's Modulus (Pa) 21 | nu = 0.3 # Poisson ratio 22 | G = E / (2 * (1 + nu)) # shear modulus (Pa) 23 | A = 0.01 # Cross-sectional area (m^2) 24 | L_max = 200.0 # Length of the beam (m) 25 | num_nodes = 101 # Number of nodes 26 | num_elements = num_nodes - 1 # Number of elements 27 | N_rollers_max = 4 # Number of additional roller supports 28 | M_forces_max = 4 # Number of point forces 29 | L_min = 15 # min distance between rollers (m) 30 | max_force = -355857 # Max point load (N) - 80,0000 lb semi 31 | min_force = max_force / 10 32 | uniform_udl = -1000 # uniformly distributed load (N) 33 | I_0 = 0.5 # initial guess for I 34 | 35 | ## opt parameters ## 36 | max_e = 600 37 | lr = 0.01 # initial opt (momentum) rate 38 | gamma = 0.98 # lr decay rate 39 | alpha_moment = 1e-2 # coefficient for bending energy loss term 40 | alpha_shear = 1e-2 # coefficient for shear energy loss term 41 | 42 | # Parameters for stopping criterion # 43 | tolerance = 5e-3 # Minimum improvement in loss 44 | patience = 5 # Number of epochs to wait for improvement 45 | 46 | num_samples = 100000 # Number of training data samples to generate 47 | 48 | random_bridge = 0 # input 1 if you want the bridge length and roller locations randomized 49 | flag = random_bridge # flag for generating fixed or random bridge 50 | 51 | # User-specified number of CPU cores # 52 | num_workers = 22 53 | 54 | ################################################################################################# 55 | ## generate a random geometry and support condition (only used if not randomizing a bridge) ## 56 | ################################################################################################# 57 | 58 | L = L_max 59 | node_positions = np.linspace(0, L, num_nodes) 60 | 61 | # Randomize roller node locations - or fix, if you prefer # 62 | roller_nodes, available_nodes = [], list(range(2, num_nodes)) # Exclude node 1 63 | num_rollers = 4 64 | 65 | # assign rollers to nodes # 66 | roller_nodes = list([10, 30, 70, 85, num_nodes-1]) 67 | 68 | # Remove the fixed roller nodes from available_nodes # 69 | for node in roller_nodes: 70 | available_nodes.remove(node) 71 | 72 | ################################################################################################# 73 | ################################################################################################# 74 | ################################################################################################# 75 | 76 | # Training data structure 77 | training_data = { 78 | "roller_x_locations": [], 79 | "force_x_locations": [], 80 | "force_values": [], 81 | "I_values": [], 82 | "shear_forces": [], 83 | "bending_moments": [], 84 | "node_positions": [], 85 | "roller_nodes": [], 86 | "force_nodes": [], 87 | "num_nodes": [], 88 | "L": [], 89 | "rotations": [], 90 | "deflections": [], 91 | } 92 | 93 | def setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl): 94 | """ 95 | Set up the OpenSees model using updated moments of inertia and apply uniform UDL to all elements. 96 | """ 97 | ops.model('basic', '-ndm', 2, '-ndf', 3) 98 | 99 | # Define nodes 100 | for i, x in enumerate(node_positions): 101 | ops.node(i + 1, x, 0.0) 102 | 103 | # Apply supports 104 | ops.fix(1, 1, 1, 0) # First node (constrained in X-Y, rotationally free) 105 | for node in roller_nodes: 106 | ops.fix(node, 0, 1, 0) # All other nodes (rotationally free, free in X) 107 | 108 | # Define elements with updated moments of inertia 109 | ops.geomTransf('Linear', 1) 110 | for i in range(len(node_positions) - 1): 111 | ops.element('elasticBeamColumn', i + 1, i + 1, i + 2, A, E, I_tensor[i].item(), 1) 112 | 113 | # Apply point loads 114 | ops.timeSeries('Linear', 1) 115 | ops.pattern('Plain', 1, 1) 116 | for node, force in zip(force_nodes, force_values): 117 | ops.load(node, 0.0, force, 0.0) 118 | 119 | # Apply uniform UDL to all elements 120 | for elem_id in range(1, len(node_positions)): 121 | ops.eleLoad('-ele', elem_id, '-type', '-beamUniform', uniform_udl, uniform_udl) 122 | 123 | # Static analysis setup 124 | ops.system('BandSPD') 125 | ops.numberer('RCM') 126 | ops.constraints('Plain') 127 | ops.integrator('LoadControl', 1.0) 128 | ops.algorithm('Linear') 129 | 130 | def generate_sample(sample_idx, num_nodes, flag, L, node_positions, roller_nodes, available_nodes, patience=10): 131 | """ 132 | Generate a single sample for training data with optimized OpenSeesPy modeling and early stopping. 133 | """ 134 | import openseespy.opensees as ops # import in function (may or may not work without this) 135 | 136 | # Random bridge length / rollers/ nodes 137 | if flag == 1: 138 | L = L_min + random.uniform(0, L_max) 139 | node_positions = np.linspace(0, L, num_nodes) 140 | 141 | # Randomize roller node locations 142 | roller_nodes, available_nodes = [], list(range(2, num_nodes)) # Exclude node 1 143 | roller_nodes.sort(reverse=True) 144 | num_rollers = random.randint(1, N_rollers_max) 145 | 146 | # Add first roller node 147 | first_roller_node = random.choice(available_nodes) 148 | roller_nodes.append(first_roller_node) 149 | available_nodes.remove(first_roller_node) 150 | 151 | # Add additional rollers without minimum distance constraint 152 | for _ in range(num_rollers - 1): 153 | if available_nodes: # Ensure there are available nodes 154 | new_roller_node = random.choice(available_nodes) 155 | roller_nodes.append(new_roller_node) 156 | available_nodes.remove(new_roller_node) 157 | 158 | 159 | # Randomize point forces and their values 160 | num_forces = random.randint(1, M_forces_max) 161 | force_nodes = random.sample(available_nodes, min(num_forces, len(available_nodes))) 162 | force_values = [random.uniform(min_force, max_force) for _ in force_nodes] 163 | 164 | # Initialize moments of inertia tensor 165 | I_tensor = torch.tensor([I_0] * num_elements, dtype=torch.float32, requires_grad=True) 166 | 167 | # Setup optimizer and scheduler 168 | optimizer = torch.optim.Adam([I_tensor], lr=0.01) 169 | scheduler = ExponentialLR(optimizer, gamma=0.98) 170 | 171 | # Optimization loop with early stopping 172 | best_loss = float('inf') 173 | patience_counter = 0 174 | 175 | for epoch in range(max_e): 176 | optimizer.zero_grad() 177 | ops.wipe() # Reset OpenSees model 178 | 179 | setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl) 180 | 181 | ops.analysis('Static') 182 | success = ops.analyze(1) 183 | 184 | if success != 0: 185 | # If analysis fails, skip this sample 186 | return None 187 | 188 | # Compute losses 189 | try: 190 | bending_moments = torch.tensor([ops.eleResponse(i, 'forces')[2] for i in range(1, len(I_tensor) + 1)], dtype=torch.float32) 191 | shear_forces = torch.tensor([ops.eleResponse(i, 'forces')[1] for i in range(1, len(I_tensor) + 1)], dtype=torch.float32) 192 | except: 193 | # If response retrieval fails, skip this sample 194 | return None 195 | 196 | bending_energy = torch.sum((bending_moments ** 2) / (2 * E * I_tensor + 1e-6)) 197 | A_approx = 0.03 * I_tensor ** 0.5 198 | shear_energy = torch.sum(shear_forces ** 2 / (G * A_approx)) 199 | primary_loss = torch.sum(I_tensor) 200 | total_loss = primary_loss + alpha_moment * bending_energy + alpha_shear * shear_energy 201 | 202 | # Backpropagate and update 203 | total_loss.backward() 204 | optimizer.step() 205 | scheduler.step() 206 | 207 | # Prevent negative inertia values 208 | with torch.no_grad(): 209 | I_tensor.clamp_(min=1e-8) 210 | 211 | # Early stopping logic 212 | if total_loss.item() < best_loss - tolerance: 213 | best_loss = total_loss.item() 214 | patience_counter = 0 215 | else: 216 | patience_counter += 1 217 | 218 | if patience_counter >= patience: 219 | break 220 | 221 | # Extract rotations and deflections 222 | rotations = [ops.nodeDisp(i, 3) if i < len(node_positions) else 0.0 for i in range(1, len(node_positions) + 1)] 223 | deflections = [ops.nodeDisp(i, 2) if i < len(node_positions) else 0.0 for i in range(1, len(node_positions) + 1)] 224 | 225 | # Return the results 226 | return { 227 | "roller_x_locations": [node_positions[node - 1] for node in roller_nodes], 228 | "force_x_locations": [node_positions[node - 1] for node in force_nodes], 229 | "force_values": force_values, 230 | "I_values": I_tensor.detach().numpy().tolist(), 231 | "shear_forces": shear_forces.detach().tolist(), 232 | "bending_moments": bending_moments.detach().tolist(), 233 | "node_positions": node_positions.tolist(), 234 | "roller_nodes": roller_nodes, 235 | "force_nodes": force_nodes, 236 | "num_nodes": num_nodes, 237 | "L": L, 238 | "rotations": rotations, 239 | "deflections": deflections, 240 | } 241 | 242 | def main(): 243 | start_time = time.time() # Start the timer 244 | 245 | # Define batch size 246 | batch_size = 500 247 | total_batches = (num_samples + batch_size - 1) // batch_size # Ceiling division 248 | 249 | for batch_num in range(total_batches): 250 | start_idx = batch_num * batch_size 251 | end_idx = min(start_idx + batch_size, num_samples) 252 | current_batch_size = end_idx - start_idx 253 | 254 | # Generate current batch indices 255 | current_indices = range(start_idx, end_idx) 256 | 257 | # Process the current batch in parallel 258 | results = Parallel(n_jobs=num_workers, backend="loky")( 259 | delayed(generate_sample)( 260 | i, num_nodes, flag, L, node_positions, roller_nodes.copy(), available_nodes.copy() 261 | ) for i in current_indices 262 | ) 263 | 264 | # Filter out any None results due to failures 265 | filtered_results = [res for res in results if res is not None] 266 | 267 | # Append the results to the training data 268 | for result in filtered_results: 269 | for key, value in result.items(): 270 | training_data[key].append(value) 271 | 272 | # Update and print the counter every 500 samples 273 | processed = end_idx 274 | print(f"{processed} samples processed.") 275 | 276 | # Save the training data to a JSON file 277 | with open("training_data_PINN_mini.json", "w") as f: 278 | json.dump(training_data, f) 279 | 280 | end_time = time.time() # End the timer 281 | 282 | print("Data generation complete.") 283 | print(f"Total execution time: {end_time - start_time:.2f} seconds") 284 | 285 | if __name__ == "__main__": 286 | main() 287 | 288 | 289 | # For sanity: Load the training data from the JSON file # 290 | with open("training_data_PINN_mini.json", "r") as file: 291 | training_data = json.load(file) 292 | 293 | # Print a summary of the dataset 294 | print("Data loaded successfully!") 295 | print(f"Number of samples: {len(training_data['roller_x_locations'])}") 296 | print("Keys available in the dataset:") 297 | for key in training_data.keys(): 298 | print(f"- {key} (Number of entries: {len(training_data[key])})") 299 | 300 | # Access all data into variables 301 | roller_x_locations = training_data["roller_x_locations"] 302 | force_x_locations = training_data["force_x_locations"] 303 | force_values = training_data["force_values"] 304 | I_values = training_data["I_values"] 305 | shear_forces = training_data["shear_forces"] 306 | bending_moments = training_data["bending_moments"] 307 | node_positions = training_data["node_positions"] 308 | roller_nodes = training_data["roller_nodes"] 309 | force_nodes = training_data["force_nodes"] 310 | num_nodes = training_data["num_nodes"] 311 | beam_lengths = training_data["L"] 312 | rotations = training_data["rotations"] 313 | deflections = training_data["deflections"] 314 | -------------------------------------------------------------------------------- /OpenPyStruct_BeamOpt.py: -------------------------------------------------------------------------------- 1 | 2 | ################################################################ 3 | #### OpenPyStruct Single Load Optimizer #### 4 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 5 | ################################################################ 6 | 7 | import os 8 | import openseespy.opensees as ops 9 | import numpy as np 10 | import random 11 | import torch 12 | import matplotlib.pyplot as plt 13 | from torch.optim.lr_scheduler import ExponentialLR 14 | 15 | # Disable torch._dynamo optimizations 16 | os.environ["TORCH_DISTRIBUTED_DEBUG"] = "NONE" 17 | 18 | 19 | 20 | ############## 21 | # Parameters # 22 | ############## 23 | 24 | E = 200e9 # Young's Modulus (Pa) 25 | nu = 0.3 # Poisson ratio 26 | G = E / (2 * (1 + nu)) # shear modulus (Pa) 27 | A = 0.01 # Cross-sectional area (m^2) 28 | L = 200.0 # Length of the beam (m) 29 | num_nodes = 101 # Number of nodes 30 | num_elements = num_nodes - 1 # Number of elements 31 | N_rollers = 5 # Number of additional roller supports 32 | M_forces = 5 # Number of point forces 33 | L_min = 15 # min distance between rollers (m) 34 | max_force = -355857 # Max point load (N) - 80,0000 lb semi 35 | uniform_udl = -5000 # unniformly distributed load (N) 36 | I_0 = 0.5 # initial guess for I 37 | 38 | 39 | ## opt parameters ## 40 | num_epochs = 1000 41 | lr = 0.01 #initial opt (momentum) rate 42 | gamma = 0.98 # lr decay rate 43 | alpha_moment = 1e-2 # coefficient for bending energy loss term 44 | alpha_shear = 1e-2 # coefficient for shear energy loss term 45 | 46 | # Parameters for stopping criterion # 47 | tolerance = 1e-2 # Minimum improvement in loss 48 | patience = 10 # Number of epochs to wait for improvement 49 | 50 | # Initialize the moments of inertia 51 | I_values = [I_0 for _ in range(num_elements)] 52 | 53 | # Define nodal positions 54 | node_positions = np.linspace(0, L, num_nodes) 55 | 56 | # Predefine roller nodes, ensuring a minimum distance between them 57 | roller_nodes = [] 58 | available_nodes = [n for n in range(2, num_nodes) if n != num_nodes] # List of nodes to choose from 59 | 60 | # Start by selecting the first roller node 61 | first_roller_node = random.choice(available_nodes) 62 | roller_nodes.append(first_roller_node) 63 | available_nodes.remove(first_roller_node) 64 | 65 | # Now iteratively select the remaining roller nodes 66 | for _ in range(1, N_rollers): 67 | valid_choice = False 68 | while not valid_choice: 69 | # Randomly pick a new roller node 70 | new_roller_node = random.choice(available_nodes) 71 | 72 | # Check that the distance to all existing roller nodes is >= L_min 73 | if all(abs(new_roller_node - existing_node) >= L_min for existing_node in roller_nodes): 74 | roller_nodes.append(new_roller_node) 75 | available_nodes.remove(new_roller_node) 76 | valid_choice = True 77 | 78 | force_nodes = [n for n in range(2, num_nodes) if n not in roller_nodes and n != 1] 79 | force_nodes = random.sample(force_nodes, min(M_forces, len(force_nodes))) 80 | force_values = [random.uniform(0.5 * max_force, max_force) for _ in force_nodes] 81 | 82 | 83 | 84 | 85 | 86 | ###################################### 87 | ## Configure Model Helper Functions ## 88 | ###################################### 89 | 90 | 91 | def setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl): 92 | """ 93 | Set up the OpenSees model using updated moments of inertia and apply uniform UDL to all elements. 94 | """ 95 | ops.model('basic', '-ndm', 2, '-ndf', 3) 96 | 97 | # Define nodes 98 | for i, x in enumerate(node_positions): 99 | ops.node(i + 1, x, 0.0) 100 | 101 | # Apply supports 102 | ops.fix(1, 1, 1, 0) # First node (constrained in X, rotationally free) 103 | for node in roller_nodes: 104 | ops.fix(node, 0, 1, 0) # All other nodes (rotationally free, free in X) 105 | 106 | # Define elements with updated moments of inertia 107 | ops.geomTransf('Linear', 1) 108 | for i in range(len(node_positions) - 1): 109 | ops.element('elasticBeamColumn', i + 1, i + 1, i + 2, A, E, I_tensor[i].item(), 1) 110 | 111 | # Apply point loads 112 | ops.timeSeries('Linear', 1) 113 | ops.pattern('Plain', 1, 1) 114 | for node, force in zip(force_nodes, force_values): 115 | ops.load(node, 0.0, force, 0.0) 116 | 117 | # Apply uniform UDL to all elements 118 | for elem_id in range(1, len(node_positions)): 119 | ops.eleLoad('-ele', elem_id, '-type', '-beamUniform', uniform_udl, uniform_udl) 120 | 121 | # Static analysis setup 122 | ops.system('BandSPD') 123 | ops.numberer('RCM') 124 | ops.constraints('Plain') 125 | ops.integrator('LoadControl', 1.0) 126 | ops.algorithm('Linear') 127 | 128 | def compute_combined_loss(I_tensor, E , G, alpha_moment, alpha_shear, node_positions): 129 | shear_forces = [] 130 | bending_moments = [] 131 | rotations = [] 132 | deflections = [] 133 | 134 | # Retrieve element responses 135 | for elem_id in range(1, len(I_tensor) + 1): 136 | response = ops.eleResponse(elem_id, 'forces') 137 | shear_forces.append(response[1]) # Shear force 138 | bending_moments.append(response[2]) # Moment at the start of the element 139 | 140 | for node_id in range(1, len(node_positions) + 1): # Node IDs start from 1 141 | uy = ops.nodeDisp(node_id, 2) # Translation in Y 142 | theta = ops.nodeDisp(node_id, 3) # Rotation 143 | deflections.append(uy) 144 | rotations.append(theta) 145 | 146 | 147 | deflections = torch.tensor(np.array(deflections)) 148 | 149 | # Convert bending moments and shear forces to PyTorch tensors 150 | bending_moments = torch.tensor(bending_moments, dtype=torch.float32, requires_grad=True) 151 | shear_forces = torch.tensor(shear_forces, dtype=torch.float32, requires_grad=True) 152 | 153 | # Compute bending energy 154 | bending_energy = torch.sum((bending_moments**2) / (2 * E * I_tensor + 1e-6)) # Avoid division by zero 155 | 156 | # make a proportional A 157 | k = 0.03 # assuming 0.01 is suitable for a built up section 158 | A = k* I_tensor** 0.5 159 | # Compute shear energy 160 | shear_energy = torch.sum(shear_forces**2 / (G*A)) # Relative to A 161 | 162 | # Compute primary loss: minimize the sum of moments of inertia 163 | primary_loss = torch.sum(I_tensor) 164 | 165 | # Combine losses 166 | total_loss = primary_loss + alpha_moment * bending_energy + alpha_shear * shear_energy 167 | 168 | return total_loss, primary_loss, alpha_moment * bending_energy, alpha_shear * shear_energy 169 | 170 | 171 | 172 | 173 | 174 | ############################################## 175 | #### Optimization Initialization and loop #### 176 | ############################################## 177 | 178 | 179 | # Convert I_values to a PyTorch tensor for optimization 180 | I_tensor = torch.tensor(I_values, dtype=torch.float32, requires_grad=True) 181 | 182 | # Define the optimizer 183 | optimizer = torch.optim.Adam([I_tensor], lr=lr) 184 | scheduler = ExponentialLR(optimizer, gamma=gamma) 185 | 186 | # Store loss history for debugging 187 | loss_history = { 188 | "total": [], 189 | "primary": [], 190 | "bending_energy": [], 191 | "shear_energy": [] 192 | } 193 | 194 | # Initialize variables for stopping criterion 195 | best_loss = float('inf') # Initialize best loss as infinity 196 | no_improvement_epochs = 0 # Counter for epochs with no improvement 197 | 198 | # Optimization loop 199 | for epoch in range(num_epochs): 200 | optimizer.zero_grad() 201 | 202 | # Clear and rebuild the model 203 | ops.wipe() 204 | setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl) 205 | 206 | ops.analysis('Static') 207 | ops.analyze(1) 208 | 209 | # Compute the combined loss 210 | total_loss, primary_loss, bending_energy, shear_energy = compute_combined_loss( 211 | I_tensor, E, G, alpha_moment, alpha_shear, node_positions) 212 | 213 | # Backpropagate and optimize 214 | total_loss.backward() 215 | optimizer.step() 216 | scheduler.step() 217 | 218 | # Clamp values to prevent negative inertia 219 | with torch.no_grad(): 220 | I_tensor.clamp_(min=1e-8) 221 | 222 | # Store loss components for debugging 223 | loss_history["total"].append(total_loss.item()) 224 | loss_history["primary"].append(primary_loss.item()) 225 | loss_history["bending_energy"].append(bending_energy.item()) 226 | loss_history["shear_energy"].append(shear_energy.item()) 227 | 228 | # Check for stopping criterion 229 | if total_loss.item() < best_loss - tolerance: 230 | best_loss = total_loss.item() 231 | no_improvement_epochs = 0 # Reset counter if improvement 232 | else: 233 | no_improvement_epochs += 1 234 | 235 | if no_improvement_epochs >= patience: 236 | print(f"Stopping early at epoch {epoch + 1}: No improvement in loss for {patience} epochs.") 237 | break 238 | 239 | # Print progress 240 | if (epoch + 1) % 100 == 0: 241 | print(f"Epoch {epoch + 1}/{num_epochs}") 242 | print(f"Total Loss: {total_loss.item():.6f}") 243 | print(f"Primary Loss: {primary_loss.item():.6f}") 244 | print(f"Bending Energy: {bending_energy.item():.6f}, Shear Energy: {shear_energy.item():.6f}") 245 | 246 | # Plot loss history 247 | plt.figure(figsize=(10, 6)) 248 | plt.plot(loss_history["total"], label="Total Loss") 249 | plt.plot(loss_history["primary"], label="Primary Loss (I Sum)") 250 | plt.plot(loss_history["bending_energy"], label="Bending Energy Loss") 251 | plt.plot(loss_history["shear_energy"], label="Shear Energy Loss") 252 | plt.xlabel("Epochs") 253 | plt.ylabel("Loss") 254 | plt.legend() 255 | plt.title("Loss Components During Optimization") 256 | plt.show() 257 | 258 | 259 | 260 | 261 | 262 | ######################################### 263 | # Update I_values with optimized values # 264 | ######################################### 265 | 266 | I_values = I_tensor.detach().numpy() 267 | 268 | # Retrieve Shear and Moment Data # 269 | shear_forces = [] 270 | bending_moments = [] 271 | for elem_id in range(1, num_elements + 1): 272 | response = ops.eleResponse(elem_id, 'forces') 273 | shear_forces.append(response[1]) # Shear force 274 | bending_moments.append(response[2]) # Moment at the start of the element 275 | 276 | # Convert Shear Forces to kN and Bending Moments to kN·m 277 | shear_forces_kn = [sf / 1e3 for sf in shear_forces] # Convert to kN 278 | bending_moments_knm = [bm / 1e3 for bm in bending_moments] # Convert to kN·m 279 | 280 | 281 | 282 | 283 | ################# 284 | # Visualization # 285 | ################# 286 | 287 | 288 | fig, axs = plt.subplots(3, 1, figsize=(20, 10), sharex=True) 289 | 290 | # Plot Moments of Inertia as Scaled Beam Elements 291 | for i in range(len(I_values)): 292 | thickness = 15 * (I_values[i] / max(I_values))**(1/3) # Scale thickness 293 | axs[0].plot([node_positions[i], node_positions[i + 1]], [0, 0], linewidth=thickness, color='blue', alpha=0.3, label='Approx section height' if i == 0 else "") 294 | 295 | # Highlight Pin and Roller Supports 296 | axs[0].scatter(0, 0, color='green', s=200, marker='^', label='Pin Support') # Pin at x = 0 297 | for node in roller_nodes: 298 | x = node_positions[node - 1] 299 | axs[0].scatter(x, 0, color='red', s=200, marker='o', label='Roller Support' if node == roller_nodes[0] else "") 300 | if num_nodes not in roller_nodes: 301 | axs[0].scatter(L, 0, color='red', s=200, marker='o') # Explicitly add roller at x = L 302 | 303 | # Plot Point Loads 304 | for node, force in zip(force_nodes, force_values): 305 | x = node_positions[node - 1] 306 | arrow_length = -0.0125 # Arrow length downward 307 | axs[0].arrow(x, -arrow_length + 0.0125, 0, arrow_length, head_width=3.5, head_length=0.0125, 308 | fc='red', ec='red', label='Point Load' if node == force_nodes[0] else "") 309 | 310 | # Set titles, labels, and legend for Moment of Inertia plot 311 | #axs[0].set_title('Beam Elements Highlighted by Moment of Inertia', fontsize=20) 312 | axs[0].set_ylabel('(Normalized I)$^{1/3}$', fontsize=20) 313 | axs[0].grid(True) 314 | axs[0].legend(fontsize=22) 315 | axs[0].set_xlim([0, max(node_positions)]) # Ensure x-axis spans the entire beam 316 | 317 | # Plot Shear Force Diagram (in kN) 318 | axs[1].step(node_positions[:-1], shear_forces_kn, where='post', color='red', label='Shear Force') 319 | axs[1].axhline(0, color='gray', linestyle='--', linewidth=0.8) 320 | axs[1].set_title('Shear Force Diagram', fontsize=20) 321 | axs[1].set_ylabel('Shear Force (kN)', fontsize=20) 322 | axs[1].grid(True) 323 | axs[1].set_xlim([0, L]) 324 | 325 | # Plot Bending Moment Diagram (in kN·m) 326 | moment_positions = (node_positions[:-1] + node_positions[1:]) / 2 # Midpoints 327 | axs[2].plot(moment_positions, bending_moments_knm, color='blue', marker='o', label='Bending Moment') 328 | axs[2].axhline(0, color='gray', linestyle='--', linewidth=0.8) 329 | axs[2].set_title('Bending Moment Diagram', fontsize=20) 330 | axs[2].set_ylabel('Bending Moment (kN·m)', fontsize=20) 331 | axs[2].set_xlabel('Beam Span (m)', fontsize=20) # Label x-axis with units 332 | axs[2].grid(True) 333 | axs[2].set_xlim([0, L]) 334 | 335 | # Adjust layout and display 336 | plt.tight_layout() 337 | plt.show() 338 | -------------------------------------------------------------------------------- /OpenPyStruct_BeamOpt_training_GPU.py: -------------------------------------------------------------------------------- 1 | ################################################################ 2 | #### OpenPyStruct GPU-Accelerated Single-Core Optimizer / Data Generator #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ################################################################ 5 | 6 | import os 7 | import openseespy.opensees as ops 8 | import numpy as np 9 | import random 10 | import torch 11 | import json 12 | from torch.optim.lr_scheduler import ExponentialLR 13 | import time 14 | from tqdm import tqdm # Progress bar library 15 | 16 | ############## 17 | # Parameters # 18 | ############## 19 | 20 | # Material Properties 21 | E = 200e9 # Young's Modulus (Pa) 22 | nu = 0.3 # Poisson ratio 23 | G = E / (2 * (1 + nu)) # Shear modulus (Pa) 24 | A = 0.01 # Cross-sectional area (m^2) 25 | 26 | # Geometry Parameters 27 | L_max = 200.0 # Maximum length of the beam (m) 28 | num_nodes = 101 # Number of nodes 29 | num_elements = num_nodes - 1 # Number of elements 30 | N_rollers_max = 4 # Maximum number of additional roller supports 31 | M_forces_max = 4 # Maximum number of point forces 32 | L_min = 15 # Minimum distance between rollers (m) 33 | 34 | # Load Parameters 35 | max_force = -355857 # Maximum point load (N) - approximately 80,000 lb 36 | min_force = max_force / 10 37 | uniform_udl = -1000 # Uniformly distributed load (N) 38 | 39 | # Initial Guess 40 | I_0 = 0.5 # Initial guess for Moments of Inertia (I) 41 | 42 | ## Optimization Parameters ## 43 | max_e = 600 # Maximum number of epochs 44 | lr = 0.01 # Initial optimization learning rate 45 | gamma = 0.98 # Learning rate decay rate 46 | alpha_moment = 1e-2 # Coefficient for bending energy loss term 47 | alpha_shear = 1e-2 # Coefficient for shear energy loss term 48 | 49 | # Stopping Criterion Parameters # 50 | tolerance = 1e-2 # Minimum improvement in loss 51 | patience = 100 # Number of epochs to wait for improvement before stopping 52 | 53 | # Data Generation Parameters # 54 | num_samples = 100000 # Number of training data samples to generate 55 | random_bridge = 0 # Set to 1 to randomize bridge length and roller locations 56 | flag = random_bridge # Flag for generating fixed or random bridge 57 | 58 | ################################################################################################# 59 | ## Generate a random geometry and support condition (only used if not randomizing a bridge) ## 60 | ################################################################################################# 61 | 62 | # Fixed Bridge Configuration 63 | L = L_max 64 | node_positions = np.linspace(0, L, num_nodes) 65 | 66 | # Fixed roller node locations (assuming node numbering starts at 1) 67 | roller_nodes = [10, 30, 70, 85, num_nodes - 1] # Node indices 68 | available_nodes = list(range(2, num_nodes)) # Exclude node 1 69 | for node in roller_nodes: 70 | if node in available_nodes: 71 | available_nodes.remove(node) 72 | 73 | ################################################################################################# 74 | ################################################################################################# 75 | ################################################################################################# 76 | 77 | # Training data structure 78 | training_data = { 79 | "roller_x_locations": [], 80 | "force_x_locations": [], 81 | "force_values": [], 82 | "I_values": [], 83 | "shear_forces": [], 84 | "bending_moments": [], 85 | "node_positions": [], 86 | "roller_nodes": [], 87 | "force_nodes": [], 88 | "num_nodes": [], 89 | "L": [], 90 | "rotations": [], 91 | "deflections": [], 92 | } 93 | 94 | def setup_model(I_tensor, node_positions, roller_nodes, force_nodes, force_values, A, E, uniform_udl): 95 | """ 96 | Set up the OpenSees model using updated moments of inertia and apply uniform UDL to all elements. 97 | """ 98 | ops.model('basic', '-ndm', 2, '-ndf', 3) 99 | 100 | # Define nodes 101 | for i, x in enumerate(node_positions): 102 | ops.node(i + 1, x, 0.0) 103 | 104 | # Apply supports 105 | ops.fix(1, 1, 1, 0) # First node (constrained in X-Y, rotationally free) 106 | for node in roller_nodes: 107 | ops.fix(node, 0, 1, 0) # All other nodes (rotationally free, free in X) 108 | 109 | # Define elements with updated moments of inertia 110 | ops.geomTransf('Linear', 1) 111 | for i in range(len(node_positions) - 1): 112 | ops.element('elasticBeamColumn', i + 1, i + 1, i + 2, A, E, I_tensor[i].item(), 1) 113 | 114 | # Apply point loads 115 | ops.timeSeries('Linear', 1) 116 | ops.pattern('Plain', 1, 1) 117 | for node, force in zip(force_nodes, force_values): 118 | ops.load(node, 0.0, force, 0.0) 119 | 120 | # Apply uniform UDL to all elements 121 | for elem_id in range(1, len(node_positions)): 122 | ops.eleLoad('-ele', elem_id, '-type', '-beamUniform', uniform_udl, uniform_udl) 123 | 124 | # Static analysis setup 125 | ops.system('BandSPD') 126 | ops.numberer('RCM') 127 | ops.constraints('Plain') 128 | ops.integrator('LoadControl', 1.0) 129 | ops.algorithm('Linear') 130 | 131 | def generate_sample(sample_idx, num_nodes, flag, L, node_positions, roller_nodes, available_nodes, patience=100, device='cpu'): 132 | """ 133 | Generate a single sample for training data with optimized OpenSeesPy modeling and early stopping. 134 | """ 135 | import openseespy.opensees as ops # Import inside function 136 | 137 | # Initialize OpenSees model 138 | ops.wipe() 139 | 140 | # Random bridge length / rollers/ nodes 141 | if flag == 1: 142 | L_sample = L_min + random.uniform(0, L_max) 143 | node_positions_sample = np.linspace(0, L_sample, num_nodes) 144 | 145 | # Randomize roller node locations 146 | roller_nodes_sample, available_nodes_sample = [], list(range(2, num_nodes)) # Exclude node 1 147 | num_rollers = random.randint(1, N_rollers_max) 148 | 149 | # Add first roller node 150 | first_roller_node = random.choice(available_nodes_sample) 151 | roller_nodes_sample.append(first_roller_node) 152 | available_nodes_sample.remove(first_roller_node) 153 | 154 | # Add additional rollers without minimum distance constraint 155 | for _ in range(num_rollers - 1): 156 | if available_nodes_sample: # Ensure there are available nodes 157 | new_roller_node = random.choice(available_nodes_sample) 158 | roller_nodes_sample.append(new_roller_node) 159 | available_nodes_sample.remove(new_roller_node) 160 | else: 161 | roller_nodes_sample = roller_nodes.copy() 162 | available_nodes_sample = available_nodes.copy() 163 | node_positions_sample = node_positions.copy() 164 | L_sample = L 165 | 166 | # Randomize point forces and their values 167 | num_forces = random.randint(1, M_forces_max) 168 | num_forces = min(num_forces, len(available_nodes_sample)) 169 | force_nodes = random.sample(available_nodes_sample, num_forces) 170 | force_values = [random.uniform(min_force, max_force) for _ in force_nodes] 171 | 172 | # Initialize moments of inertia tensor on the specified device 173 | I_tensor = torch.tensor([I_0] * num_elements, dtype=torch.float32, requires_grad=True, device=device) 174 | 175 | # Setup optimizer and scheduler 176 | optimizer = torch.optim.Adam([I_tensor], lr=lr) 177 | scheduler = ExponentialLR(optimizer, gamma=gamma) 178 | 179 | # Optimization loop with early stopping 180 | best_loss = float('inf') 181 | patience_counter = 0 182 | 183 | for epoch in range(1, max_e + 1): 184 | optimizer.zero_grad() 185 | ops.wipe() # Reset OpenSees model 186 | 187 | setup_model(I_tensor, node_positions_sample, roller_nodes_sample, force_nodes, force_values, A, E, uniform_udl) 188 | 189 | ops.analysis('Static') 190 | try: 191 | ops.analyze(1) 192 | except Exception as e: 193 | print(f"OpenSees analysis failed for sample {sample_idx+1} at epoch {epoch}: {e}") 194 | break # Skip to the next sample 195 | 196 | # Compute losses 197 | try: 198 | bending_moments = torch.tensor( 199 | [ops.eleResponse(i, 'forces')[2] for i in range(1, len(I_tensor) + 1)], 200 | dtype=torch.float32, 201 | device=device 202 | ) 203 | shear_forces = torch.tensor( 204 | [ops.eleResponse(i, 'forces')[1] for i in range(1, len(I_tensor) + 1)], 205 | dtype=torch.float32, 206 | device=device 207 | ) 208 | except Exception as e: 209 | print(f"Failed to retrieve responses for sample {sample_idx+1} at epoch {epoch}: {e}") 210 | break # Skip to the next sample 211 | 212 | bending_energy = torch.sum((bending_moments ** 2) / (2 * E * I_tensor + 1e-6)) 213 | A_approx = 0.03 * I_tensor ** 0.5 214 | shear_energy = torch.sum(shear_forces ** 2 / (G * A_approx)) 215 | primary_loss = torch.sum(I_tensor) 216 | total_loss = primary_loss + alpha_moment * bending_energy + alpha_shear * shear_energy 217 | 218 | # Backpropagate and update 219 | total_loss.backward() 220 | optimizer.step() 221 | scheduler.step() 222 | 223 | # Prevent negative inertia values 224 | with torch.no_grad(): 225 | I_tensor.clamp_(min=1e-8) 226 | 227 | # Early stopping logic 228 | if total_loss.item() < best_loss - tolerance: 229 | best_loss = total_loss.item() 230 | patience_counter = 0 231 | else: 232 | patience_counter += 1 233 | 234 | if patience_counter >= patience: 235 | break 236 | 237 | # Print progress every 100 epochs 238 | if epoch % 100 == 0: 239 | print(f"Sample {sample_idx+1}/{num_samples} - Epoch {epoch}/{max_e} - Loss: {total_loss.item():.4f}") 240 | 241 | # Extract rotations and deflections 242 | rotations = [] 243 | deflections = [] 244 | for i in range(1, len(node_positions_sample) + 1): 245 | try: 246 | rotations.append(ops.nodeDisp(i, 3)) 247 | except: 248 | rotations.append(0.0) 249 | try: 250 | deflections.append(ops.nodeDisp(i, 2)) 251 | except: 252 | deflections.append(0.0) 253 | 254 | # Detach tensors and move to CPU for storage 255 | I_values_cpu = I_tensor.detach().cpu().numpy().tolist() 256 | shear_forces_cpu = shear_forces.detach().cpu().tolist() 257 | bending_moments_cpu = bending_moments.detach().cpu().tolist() 258 | 259 | # Return the results 260 | return { 261 | "roller_x_locations": [node_positions_sample[node - 1] for node in roller_nodes_sample], 262 | "force_x_locations": [node_positions_sample[node - 1] for node in force_nodes], 263 | "force_values": force_values, 264 | "I_values": I_values_cpu, 265 | "shear_forces": shear_forces_cpu, 266 | "bending_moments": bending_moments_cpu, 267 | "node_positions": node_positions_sample.tolist(), 268 | "roller_nodes": roller_nodes_sample, 269 | "force_nodes": force_nodes, 270 | "num_nodes": num_nodes, 271 | "L": L_sample, 272 | "rotations": rotations, 273 | "deflections": deflections, 274 | } 275 | 276 | def main(): 277 | start_time = time.time() # Start the timer 278 | 279 | # Determine the device to use (GPU if available, else CPU) 280 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 281 | print(f"Using device: {device}") 282 | 283 | # Initialize progress bar for sample generation 284 | with tqdm(total=num_samples, desc="Generating Samples") as pbar: 285 | for i in range(num_samples): 286 | result = generate_sample( 287 | i, 288 | num_nodes, 289 | flag, 290 | L, 291 | node_positions, 292 | roller_nodes, 293 | available_nodes, 294 | patience=patience, 295 | device=device 296 | ) 297 | for key, value in result.items(): 298 | training_data[key].append(value) 299 | pbar.update(1) # Update the progress bar 300 | 301 | # Save the training data to a JSON file 302 | with open("training_data_PINN_mini.json", "w") as f: 303 | json.dump(training_data, f) 304 | 305 | end_time = time.time() # End the timer 306 | 307 | print("Data generation complete.") 308 | print(f"Total execution time: {end_time - start_time:.2f} seconds") 309 | 310 | if __name__ == "__main__": 311 | main() 312 | 313 | # For sanity: Load the training data from the JSON file # 314 | with open("training_data_PINN_mini.json", "r") as file: 315 | training_data = json.load(file) 316 | 317 | # Print a summary of the dataset 318 | print("Data loaded successfully!") 319 | print(f"Number of samples: {len(training_data['roller_x_locations'])}") 320 | print("Keys available in the dataset:") 321 | for key in training_data.keys(): 322 | print(f"- {key} (Number of entries: {len(training_data[key])})") 323 | 324 | # Access all data into variables 325 | roller_x_locations = training_data["roller_x_locations"] 326 | force_x_locations = training_data["force_x_locations"] 327 | force_values = training_data["force_values"] 328 | I_values = training_data["I_values"] 329 | shear_forces = training_data["shear_forces"] 330 | bending_moments = training_data["bending_moments"] 331 | node_positions = training_data["node_positions"] 332 | roller_nodes = training_data["roller_nodes"] 333 | force_nodes = training_data["force_nodes"] 334 | num_nodes = training_data["num_nodes"] 335 | beam_lengths = training_data["L"] 336 | rotations = training_data["rotations"] 337 | deflections = training_data["deflections"] 338 | -------------------------------------------------------------------------------- /OpenPyStruct_GNN_MultiCase_Beta.py: -------------------------------------------------------------------------------- 1 | ############################################################################ 2 | #### OpenPyStruct Chain GNN With Speedups #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ############################################################################ 5 | 6 | 7 | #### EXPERIMENTAL CODE (FUNCTIONAL BUT NOT OPTIMIZED) #### 8 | 9 | import os 10 | import json 11 | import time 12 | import math 13 | import random 14 | import numpy as np 15 | import matplotlib.pyplot as plt 16 | import matplotlib.cm as cm 17 | from matplotlib.lines import Line2D 18 | from matplotlib.patches import FancyArrowPatch, Rectangle 19 | import seaborn as sns 20 | 21 | import torch 22 | import torch.nn as nn 23 | import torch.optim as optim 24 | from torch.utils.data import DataLoader, TensorDataset 25 | from sklearn.preprocessing import StandardScaler 26 | from sklearn.metrics import r2_score 27 | from scipy.stats import mode 28 | 29 | # For faster adjacency multiplication 30 | from einops import rearrange 31 | # For Mixed Precision 32 | from torch.cuda.amp import autocast, GradScaler 33 | 34 | ####################################### 35 | # 1) CONFIGURATION & HYPERPARAMETERS 36 | ####################################### 37 | 38 | n_cases = 6 39 | nelem = 100 40 | box_constraint_coeff = 5e-1 41 | encoder_hidden_dim = 128 42 | gnn_hidden_dim = 128 43 | num_gnn_layers = 2 44 | dropout_rate = 0.5 45 | num_epochs = 500 46 | batch_size = 512 47 | patience = 10 48 | learning_rate = 3e-3 49 | weight_decay = 1e-2 50 | train_split = 0.8 51 | sigma_0 = 0.01 52 | gamma_noise = 0.99 53 | gamma = 0.975 54 | initial_alpha = 0.5 55 | c = 0.5 56 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 57 | 58 | print("Using device:", device) 59 | 60 | ####################################### 61 | # 2) DATA LOADING & PREPROCESSING 62 | ####################################### 63 | 64 | def pad_sequences(data_list, max_length, pad_val=0.0): 65 | out = np.full((len(data_list), max_length), pad_val, dtype=np.float32) 66 | for i, arr in enumerate(data_list): 67 | arr_np = np.array(arr, dtype=np.float32) 68 | length = min(len(arr_np), max_length) 69 | out[i, :length] = arr_np[:length] 70 | return out 71 | 72 | def unify_label_with_c(I_3d, cval): 73 | I_mean = I_3d.mean(axis=1) 74 | I_std = I_3d.std(axis=1) 75 | return I_mean + cval * I_std 76 | 77 | def fit_transform_3d(arr_3d, scaler): 78 | B, NC, M = arr_3d.shape 79 | flat = arr_3d.reshape(B*NC, M) 80 | scaled = scaler.fit_transform(flat) 81 | return scaled.reshape(B, NC, M) 82 | 83 | def merge_sub_features(*arrays): 84 | return np.concatenate(arrays, axis=2) 85 | 86 | def scale_user_inputs(user_roller, user_force_x, user_force_vals, user_node_pos, 87 | scalers, n_cases, max_lengths): 88 | """ 89 | Scales a single user input sample, returns shape => (1, n_cases, feat_dim) 90 | """ 91 | def pad_to_length(seq, req_len): 92 | arr = np.zeros((req_len,), dtype=np.float32) 93 | length = min(len(seq), req_len) 94 | arr[:length] = seq[:length] 95 | return arr 96 | 97 | feat_arrays = [] 98 | for i in range(n_cases): 99 | r_pad = pad_to_length(user_roller[i], max_lengths['roller_x']) 100 | fx_pad = pad_to_length(user_force_x[i], max_lengths['force_x']) 101 | fv_pad = pad_to_length(user_force_vals[i], max_lengths['force_values']) 102 | nd_pad = pad_to_length(user_node_pos[i], max_lengths['node_positions']) 103 | 104 | r_scaled = scalers["roller_x"].transform(r_pad.reshape(1, -1)).flatten() 105 | fx_scaled = scalers["force_x"].transform(fx_pad.reshape(1, -1)).flatten() 106 | fv_scaled = scalers["force_values"].transform(fv_pad.reshape(1, -1)).flatten() 107 | nd_scaled = scalers["node_positions"].transform(nd_pad.reshape(1, -1)).flatten() 108 | 109 | sub_feat = np.concatenate([r_scaled, fx_scaled, fv_scaled, nd_scaled]) 110 | feat_arrays.append(sub_feat) 111 | 112 | feat_2d = np.stack(feat_arrays, axis=0) # (n_cases, total_feat_dim) 113 | feat_3d = feat_2d[np.newaxis, ...] # (1, n_cases, total_feat_dim) 114 | return feat_3d 115 | 116 | # Load data 117 | try: 118 | with open("training_data_PINN_Case_two.json", "r") as f: 119 | data = json.load(f) 120 | except FileNotFoundError: 121 | raise FileNotFoundError("The file 'training_data_PINN_Case_two.json' was not found.") 122 | 123 | roller_x = data.get("roller_x_locations", []) 124 | force_x = data.get("force_x_locations", []) 125 | force_values = data.get("force_values", []) 126 | node_positions = data.get("node_positions", []) 127 | I_values = data.get("I_values", []) 128 | 129 | num_samples = len(I_values) 130 | req_keys = ["roller_x_locations","force_x_locations","force_values","node_positions"] 131 | if not all(len(data.get(k, [])) == num_samples for k in req_keys): 132 | raise ValueError("Mismatch in sample counts among roller_x, force_x, force_values, node_positions.") 133 | 134 | max_lengths = { 135 | "roller_x": max(len(r) for r in roller_x) if roller_x else 0, 136 | "force_x": max(len(r) for r in force_x) if force_x else 0, 137 | "force_values": max(len(r) for r in force_values) if force_values else 0, 138 | "node_positions": max(len(r) for r in node_positions) if node_positions else 0, 139 | "I_values": max(len(r) for r in I_values) if I_values else 0 140 | } 141 | 142 | # Pad 143 | roller_x_pad = pad_sequences(roller_x, max_lengths["roller_x"]) 144 | force_x_pad = pad_sequences(force_x, max_lengths["force_x"]) 145 | force_val_pad = pad_sequences(force_values,max_lengths["force_values"]) 146 | node_pos_pad = pad_sequences(node_positions, max_lengths["node_positions"]) 147 | I_values_pad = pad_sequences(I_values, max_lengths["I_values"]) 148 | 149 | # Group by n_cases 150 | total_grouped = num_samples // n_cases 151 | if total_grouped == 0: 152 | raise ValueError(f"n_cases={n_cases} > total samples={num_samples}.") 153 | 154 | trim_len = total_grouped * n_cases 155 | roller_x_pad = roller_x_pad[:trim_len] 156 | force_x_pad = force_x_pad[:trim_len] 157 | force_val_pad = force_val_pad[:trim_len] 158 | node_pos_pad = node_pos_pad[:trim_len] 159 | I_values_pad = I_values_pad[:trim_len] 160 | 161 | roller_grouped = roller_x_pad.reshape(total_grouped, n_cases, -1) 162 | force_x_grouped = force_x_pad.reshape(total_grouped, n_cases, -1) 163 | force_val_grouped = force_val_pad.reshape(total_grouped, n_cases, -1) 164 | node_grouped = node_pos_pad.reshape(total_grouped, n_cases, -1) 165 | I_grouped = I_values_pad.reshape(total_grouped, n_cases, -1) 166 | 167 | indices = np.random.permutation(total_grouped) 168 | train_sz = int(train_split * total_grouped) 169 | train_idx = indices[:train_sz] 170 | val_idx = indices[train_sz:] 171 | 172 | roller_train = roller_grouped[train_idx] 173 | roller_val = roller_grouped[val_idx] 174 | force_x_train = force_x_grouped[train_idx] 175 | force_x_val = force_x_grouped[val_idx] 176 | force_val_train = force_val_grouped[train_idx] 177 | force_val_val = force_val_grouped[val_idx] 178 | node_train = node_grouped[train_idx] 179 | node_val = node_grouped[val_idx] 180 | I_train = I_grouped[train_idx] 181 | I_val = I_grouped[val_idx] 182 | 183 | scalers_inputs = { 184 | "roller_x": StandardScaler(), 185 | "force_x": StandardScaler(), 186 | "force_values": StandardScaler(), 187 | "node_positions": StandardScaler() 188 | } 189 | scaler_Y = StandardScaler() 190 | 191 | roller_train_std = fit_transform_3d(roller_train, scalers_inputs["roller_x"]) 192 | force_x_train_std = fit_transform_3d(force_x_train, scalers_inputs["force_x"]) 193 | force_val_train_std = fit_transform_3d(force_val_train, scalers_inputs["force_values"]) 194 | node_train_std = fit_transform_3d(node_train, scalers_inputs["node_positions"]) 195 | 196 | roller_val_std = fit_transform_3d(roller_val, scalers_inputs["roller_x"]) 197 | force_x_val_std = fit_transform_3d(force_x_val, scalers_inputs["force_x"]) 198 | force_val_val_std = fit_transform_3d(force_val_val, scalers_inputs["force_values"]) 199 | node_val_std = fit_transform_3d(node_val, scalers_inputs["node_positions"]) 200 | 201 | # Merge features => (B, n_cases, feat_in) 202 | X_train_3d = merge_sub_features( 203 | roller_train_std, 204 | force_x_train_std, 205 | force_val_train_std, 206 | node_train_std 207 | ) 208 | X_val_3d = merge_sub_features( 209 | roller_val_std, 210 | force_x_val_std, 211 | force_val_val_std, 212 | node_val_std 213 | ) 214 | 215 | B_train, nc, fin = X_train_3d.shape 216 | B_val, _, _ = X_val_3d.shape 217 | enc_in_dim = nc * fin # flatten sub-cases 218 | 219 | # Flatten => (B, enc_in_dim) 220 | X_train_2d = X_train_3d.reshape(B_train, -1) 221 | X_val_2d = X_val_3d.reshape(B_val, -1) 222 | 223 | Y_train_2d = unify_label_with_c(I_train, c) 224 | Y_val_2d = unify_label_with_c(I_val, c) 225 | 226 | scaler_Y.fit(Y_train_2d) 227 | Y_train_std = scaler_Y.transform(Y_train_2d) 228 | Y_val_std = scaler_Y.transform(Y_val_2d) 229 | 230 | X_train_tensor = torch.tensor(X_train_2d, dtype=torch.float32) 231 | Y_train_tensor = torch.tensor(Y_train_std, dtype=torch.float32) 232 | X_val_tensor = torch.tensor(X_val_2d, dtype=torch.float32) 233 | Y_val_tensor = torch.tensor(Y_val_std, dtype=torch.float32) 234 | 235 | min_constraint = torch.min(Y_train_tensor) 236 | max_constraint = torch.max(Y_train_tensor) 237 | 238 | train_dataset = TensorDataset(X_train_tensor, Y_train_tensor) 239 | val_dataset = TensorDataset(X_val_tensor, Y_val_tensor) 240 | 241 | print(f"Train set shape => X: {X_train_2d.shape}, Y: {Y_train_2d.shape}") 242 | print(f"Val set shape => X: {X_val_2d.shape}, Y: {Y_val_2d.shape}") 243 | 244 | 245 | ####################################### 246 | # 3) DEFINE A CHAIN GNN WITH SPEEDUPS 247 | ####################################### 248 | 249 | def precompute_normalized_adjacency(n): 250 | """ 251 | Returns a precomputed (n, n) adjacency for a chain, normalized by D^{-1/2}AD^{-1/2}. 252 | """ 253 | A = torch.zeros((n, n), dtype=torch.float32) 254 | for i in range(n - 1): 255 | A[i, i+1] = 1.0 256 | A[i+1, i] = 1.0 257 | degrees = A.sum(dim=1) 258 | D_inv_sqrt = torch.pow(degrees + 1e-8, -0.5) 259 | # elementwise multiply for normalization => A_hat = D^-1/2 * A * D^-1/2 260 | A_hat = A * D_inv_sqrt.unsqueeze(0) 261 | A_hat = A_hat * D_inv_sqrt.unsqueeze(1) 262 | return A_hat 263 | 264 | class GCNLayer(nn.Module): 265 | """ 266 | Single GCN layer: out = A_hat @ (X W). 267 | We'll use a single adjacency across the entire batch. 268 | """ 269 | def __init__(self, in_dim, out_dim): 270 | super().__init__() 271 | self.linear = nn.Linear(in_dim, out_dim, bias=False) 272 | 273 | def forward(self, x, A_hat): 274 | """ 275 | x shape => (B, n, in_dim) 276 | A_hat => (n, n), normalized adjacency 277 | returns => (B, n, out_dim) 278 | """ 279 | # Step 1: Wx => (B, n, out_dim) 280 | Wx = self.linear(x) 281 | # Step 2: multiply adjacency => out = A_hat @ Wx per sample in batch 282 | # We'll do => out[b] = A_hat (n,n) @ Wx[b] (n, out_dim) 283 | # Use einops: A_hat => (n,n), Wx => (B, n, out_dim) 284 | # out => (B,n,out_dim) => (b,i,d) = sum_j A_hat[i,j]*Wx[b,j,d] 285 | out = torch.einsum('ij,bjd->bid', A_hat, Wx) 286 | return out 287 | 288 | class ChainGNN(nn.Module): 289 | """ 290 | 1) Encoder MLP => (B, n_elem, gnn_hidden_dim) 291 | 2) Several GCN Layers 292 | 3) Output => (B, n_elem) 293 | """ 294 | def __init__(self, enc_in_dim, n_elem, enc_hidden_dim, gnn_hidden_dim, 295 | num_gnn_layers, dropout): 296 | super().__init__() 297 | self.n_elem = n_elem 298 | self.enc_in_dim = enc_in_dim 299 | self.enc_hidden_dim = enc_hidden_dim 300 | self.gnn_hidden_dim = gnn_hidden_dim 301 | self.num_gnn_layers = num_gnn_layers 302 | 303 | # Precompute chain adjacency 304 | A_hat = precompute_normalized_adjacency(n_elem) 305 | self.register_buffer("A_hat", A_hat) # stays on same device as model 306 | 307 | # Encoder: flatten sub-cases => node embeddings 308 | self.encoder = nn.Sequential( 309 | nn.Linear(enc_in_dim, enc_hidden_dim), 310 | nn.ReLU(), 311 | nn.Linear(enc_hidden_dim, n_elem * gnn_hidden_dim) 312 | ) 313 | 314 | # GCN stack 315 | layers = [] 316 | norms = [] 317 | drops = [] 318 | in_dim = gnn_hidden_dim 319 | for _ in range(num_gnn_layers): 320 | layers.append(GCNLayer(in_dim, in_dim)) # keep dimension stable 321 | norms.append(nn.LayerNorm(in_dim)) 322 | drops.append(nn.Dropout(dropout)) 323 | self.gcn_layers = nn.ModuleList(layers) 324 | self.norms = nn.ModuleList(norms) 325 | self.drops = nn.ModuleList(drops) 326 | 327 | # Output readout => (n_elem, 1) 328 | self.out_layer = nn.Linear(in_dim, 1) 329 | 330 | def forward(self, x): 331 | """ 332 | x => (B, enc_in_dim) 333 | returns => (B, n_elem) 334 | """ 335 | B = x.size(0) 336 | # 1) Encode => (B, n_elem * gnn_hidden_dim) 337 | enc = self.encoder(x) 338 | node_feats = enc.view(B, self.n_elem, self.gnn_hidden_dim) 339 | 340 | # 2) GCN layers 341 | out = node_feats 342 | for i, gcn in enumerate(self.gcn_layers): 343 | out_in = self.norms[i](out) 344 | out_g = gcn(out_in, self.A_hat) # (B, n, gnn_hidden_dim) 345 | out = out + self.drops[i](out_g) # residual 346 | 347 | # 3) Output => (B, n_elem, 1) => (B, n_elem) 348 | out = self.out_layer(out).squeeze(-1) 349 | return out 350 | 351 | ####################################### 352 | # 4) CUSTOM LOSS WITH CONSTRAINT PENALTY 353 | ####################################### 354 | 355 | class TrainableL1L2Loss(nn.Module): 356 | def __init__(self, initial_alpha, min_constraint, max_constraint, penalty_weight): 357 | super().__init__() 358 | self.alpha = nn.Parameter(torch.tensor(initial_alpha, dtype=torch.float32, requires_grad=True)) 359 | self.l1 = nn.L1Loss() 360 | self.l2 = nn.MSELoss() 361 | self.min_constraint = min_constraint 362 | self.max_constraint = max_constraint 363 | self.penalty_weight = penalty_weight 364 | 365 | def forward(self, preds, targets): 366 | alpha = torch.clamp(self.alpha, 1e-6, 1.0) 367 | l1_loss = self.l1(preds, targets) 368 | l2_loss = self.l2(preds, targets) 369 | 370 | penalty = 0.0 371 | if self.min_constraint is not None: 372 | below_min = torch.sum(torch.relu(self.min_constraint - preds)) 373 | penalty += below_min 374 | if self.max_constraint is not None: 375 | above_max = torch.sum(torch.relu(preds - self.max_constraint)) 376 | penalty += above_max 377 | 378 | total_loss = alpha * l1_loss + (1 - alpha) * l2_loss + self.penalty_weight * penalty 379 | return total_loss 380 | 381 | 382 | ####################################### 383 | # 5) INIT MODEL & TRAIN (SPEEDUPS) 384 | ####################################### 385 | 386 | model = ChainGNN( 387 | enc_in_dim = enc_in_dim, 388 | n_elem = nelem, 389 | enc_hidden_dim = encoder_hidden_dim, 390 | gnn_hidden_dim = gnn_hidden_dim, 391 | num_gnn_layers = num_gnn_layers, 392 | dropout = dropout_rate 393 | ).to(device) 394 | 395 | optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay) 396 | scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma) 397 | criterion = TrainableL1L2Loss( 398 | initial_alpha = initial_alpha, 399 | min_constraint = min_constraint.to(device), 400 | max_constraint = max_constraint.to(device), 401 | penalty_weight = box_constraint_coeff 402 | ).to(device) 403 | 404 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 405 | val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 406 | 407 | scaler_amp = GradScaler() 408 | 409 | # We'll do live plotting every 5 epochs to reduce overhead 410 | plot_frequency = 5 411 | plt.ion() 412 | fig, ax = plt.subplots(figsize=(10,6)) 413 | 414 | def live_plot(epoch, train_losses, val_losses): 415 | ax.clear() 416 | ax.plot(range(1, epoch + 1), train_losses, label="Train Loss", marker='o', color='blue') 417 | ax.plot(range(1, epoch + 1), val_losses, label="Val Loss", marker='x', color='red') 418 | ax.set_xlabel("Epochs") 419 | ax.set_ylabel("Loss") 420 | ax.set_title("Training and Validation Loss") 421 | ax.legend() 422 | ax.grid(True, linestyle='--', alpha=0.7) 423 | plt.pause(0.01) 424 | 425 | train_losses, val_losses = [], [] 426 | best_val_loss = float('inf') 427 | epochs_no_improve = 0 428 | 429 | for epoch in range(1, num_epochs+1): 430 | model.train() 431 | noise_level = sigma_0 * (gamma_noise ** epoch) 432 | total_train_loss = 0.0 433 | t0 = time.time() 434 | 435 | for Xb, Yb in train_loader: 436 | Xb, Yb = Xb.to(device), Yb.to(device) 437 | # Add optional noise to inputs 438 | Xb_noisy = Xb + torch.randn_like(Xb) * noise_level 439 | 440 | optimizer.zero_grad() 441 | with autocast(): 442 | preds = model(Xb_noisy) 443 | loss = criterion(preds, Yb) 444 | 445 | scaler_amp.scale(loss).backward() 446 | scaler_amp.unscale_(optimizer) 447 | torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) 448 | scaler_amp.step(optimizer) 449 | scaler_amp.update() 450 | 451 | total_train_loss += loss.item() 452 | 453 | avg_train_loss = total_train_loss / len(train_loader) 454 | train_losses.append(avg_train_loss) 455 | 456 | model.eval() 457 | total_val_loss = 0.0 458 | with torch.no_grad(), autocast(): 459 | for Xb, Yb in val_loader: 460 | Xb, Yb = Xb.to(device), Yb.to(device) 461 | preds = model(Xb) 462 | val_loss = criterion(preds, Yb) 463 | total_val_loss += val_loss.item() 464 | 465 | avg_val_loss = total_val_loss / len(val_loader) 466 | val_losses.append(avg_val_loss) 467 | scheduler.step() 468 | 469 | if avg_val_loss < best_val_loss: 470 | best_val_loss = avg_val_loss 471 | epochs_no_improve = 0 472 | torch.save(model.state_dict(), "best_gnn_model.pth") 473 | else: 474 | epochs_no_improve += 1 475 | if epochs_no_improve >= patience: 476 | print(f"Early stopping at epoch {epoch}") 477 | break 478 | 479 | dt = time.time() - t0 480 | print(f"Epoch {epoch}/{num_epochs} | " 481 | f"Train Loss={avg_train_loss:.6f}, " 482 | f"Val Loss={avg_val_loss:.6f}, " 483 | f"Time={dt:.2f}s") 484 | 485 | live_plot(epoch, train_losses, val_losses) 486 | 487 | 488 | ####################################### 489 | # 6) EVALUATION 490 | ####################################### 491 | 492 | model.load_state_dict(torch.load("best_gnn_model.pth", map_location=device)) 493 | model.eval() 494 | 495 | val_loader_eval = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 496 | all_preds, all_labels = [], [] 497 | 498 | with torch.no_grad(): 499 | for Xb, Yb in val_loader_eval: 500 | Xb = Xb.to(device) 501 | preds = model(Xb) 502 | all_preds.append(preds.cpu()) 503 | all_labels.append(Yb) 504 | 505 | all_preds = torch.cat(all_preds, dim=0).numpy() 506 | all_labels = torch.cat(all_labels, dim=0).numpy() 507 | 508 | all_preds_unstd = scaler_Y.inverse_transform(all_preds) 509 | all_labels_unstd = scaler_Y.inverse_transform(all_labels) 510 | 511 | all_preds_unstd = np.clip(all_preds_unstd, 0.0, 1e10) 512 | all_labels_unstd = np.clip(all_labels_unstd, 0.0, 1e10) 513 | 514 | r2_val = r2_score(all_labels_unstd.ravel(), all_preds_unstd.ravel()) 515 | print(f"R² on Validation: {r2_val:.4f}") 516 | 517 | 518 | ####################################### 519 | # 7) EXAMPLE INFERENCE & PLOT 520 | ####################################### 521 | 522 | # Example user input 523 | L_beam = 200 524 | Fmin_user = -355857 525 | Fmax_user = Fmin_user / 10 526 | user_rollers = [2*9, 2*29, 2*69, 2*85, 2*100] 527 | user_roller = [user_rollers for _ in range(n_cases)] 528 | 529 | user_force_x = [] 530 | user_force_vals = [] 531 | for _ in range(n_cases): 532 | num_forces = random.randint(1,3) 533 | fx = sorted([random.uniform(0, L_beam) for _ in range(num_forces)]) 534 | fv = [random.uniform(Fmin_user, Fmax_user) for _ in range(num_forces)] 535 | user_force_x.append(fx) 536 | user_force_vals.append(fv) 537 | 538 | user_node_pos = [np.linspace(0, L_beam, nelem+1).tolist() for _ in range(n_cases)] 539 | 540 | def build_user_input_no_agg(user_roller, user_force_x, user_force_vals, user_node_pos, 541 | scalers, n_cases, max_lengths): 542 | feat_3d = scale_user_inputs( 543 | user_roller, user_force_x, user_force_vals, user_node_pos, 544 | scalers, n_cases, max_lengths 545 | ) 546 | return feat_3d 547 | 548 | X_user_3d = build_user_input_no_agg( 549 | user_roller, user_force_x, user_force_vals, user_node_pos, 550 | scalers_inputs, n_cases, max_lengths 551 | ) 552 | 553 | B_usr, Nc_usr, Fin_usr = X_user_3d.shape 554 | X_user_2d = X_user_3d.reshape(B_usr, -1) 555 | X_user_t = torch.tensor(X_user_2d, dtype=torch.float32).to(device) 556 | 557 | model.eval() 558 | with torch.no_grad(): 559 | pred_1x = model(X_user_t) # => (1, n_elem) 560 | pred_1x_np = pred_1x.cpu().numpy().squeeze() 561 | pred_1x_unstd = scaler_Y.inverse_transform(pred_1x_np.reshape(1, -1)).squeeze() 562 | 563 | # Visualization 564 | unique_rollers = sorted(set([x for sublist in user_roller for x in sublist] + [L_beam])) 565 | case_colors = sns.color_palette("Set1", n_colors=n_cases) 566 | case_labels = [f'Force Case {i+1} (N)' for i in range(n_cases)] 567 | 568 | beam_y = 0 569 | beam_x = [0, L_beam] 570 | beam_y_vals = [beam_y, beam_y] 571 | 572 | force_positions = [] 573 | force_vals_plot = [] 574 | for fx, fv in zip(user_force_x, user_force_vals): 575 | for xx, val in zip(fx, fv): 576 | force_positions.append(xx) 577 | force_vals_plot.append(val) 578 | 579 | all_force_vals = force_vals_plot 580 | max_force = max(abs(val) for val in all_force_vals) if all_force_vals else 1.0 581 | desired_max_arrow_length = 2.0 582 | arrow_scale = desired_max_arrow_length / max_force if max_force != 0 else 1.0 583 | 584 | beam_positions = user_node_pos[0][:nelem] 585 | I_normalized = (pred_1x_unstd - pred_1x_unstd.min()) / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8) 586 | cmap = cm.winter 587 | norm = plt.Normalize(pred_1x_unstd.min(), pred_1x_unstd.max()) 588 | 589 | block_width = L_beam / nelem * 0.8 590 | block_height = 1 591 | 592 | fig, ax = plt.subplots(figsize=(18, 7)) 593 | ax.plot(beam_x, beam_y_vals, color='black', linewidth=3, label='Beam') 594 | ax.scatter(beam_x[0], beam_y - 0.15, marker='^', color='red', s=300, zorder=6) 595 | 596 | # Plot Rollers 597 | ax.scatter(unique_rollers, [beam_y]*len(unique_rollers), 598 | marker='o', color='seagreen', s=200, label='Rollers', 599 | zorder=5, edgecolors='k') 600 | 601 | # Plot Forces 602 | for case_idx in range(n_cases): 603 | fx_list = user_force_x[case_idx] 604 | fv_list = user_force_vals[case_idx] 605 | color = case_colors[case_idx] 606 | label = case_labels[case_idx] 607 | 608 | for idx, (fx, fv) in enumerate(zip(fx_list, fv_list)): 609 | arrow_length = abs(fv)*arrow_scale 610 | start_point = (fx, beam_y + arrow_length) 611 | end_point = (fx, beam_y) 612 | 613 | arrow = FancyArrowPatch( 614 | posA=start_point, posB=end_point, 615 | arrowstyle='-|>', 616 | mutation_scale=20, 617 | color=color, linewidth=2, alpha=0.8, 618 | label=label if idx == 0 else "" 619 | ) 620 | ax.add_patch(arrow) 621 | ax.text(fx, beam_y + arrow_length + desired_max_arrow_length*0.02, 622 | f"{fv:.0f}", ha='center', va='bottom', 623 | fontsize=10, color=color, fontweight='bold') 624 | 625 | # Plot predicted I 626 | for idx, (x_pos, I_val) in enumerate(zip(beam_positions, pred_1x_unstd)): 627 | color = cmap(norm(I_val)) 628 | rect_x = x_pos - block_width/2 629 | rect_y = beam_y - (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8))*block_height/2 630 | rect = Rectangle( 631 | (rect_x, rect_y), 632 | block_width, 633 | (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8))*block_height, 634 | linewidth=0, edgecolor=None, facecolor=color, alpha=0.6 635 | ) 636 | ax.add_patch(rect) 637 | 638 | sm = cm.ScalarMappable(cmap=cmap, norm=norm) 639 | sm.set_array([]) 640 | cbar = fig.colorbar(sm, ax=ax, orientation='vertical', fraction=0.046, pad=0.04) 641 | cbar.set_label('Predicted I (m$^4$)', fontsize=16) 642 | cbar.ax.tick_params(labelsize=10) 643 | 644 | ax.set_title("Beam Setup with Applied Forces & GNN (Speedups) Predicted I", 645 | fontsize=22, fontweight='bold', pad=20) 646 | ax.set_xlabel("Beam Length (m)", fontsize=16, fontweight='semibold') 647 | ax.set_xlim(-5, L_beam+5) 648 | ax.set_ylim(-2.5, 2.5) 649 | ax.set_xticks(np.arange(0, L_beam+5, 5)) 650 | ax.grid(True, linestyle='--', linewidth=0.5, alpha=0.7) 651 | 652 | legend_elements = [ 653 | Line2D([0], [0], color='black', lw=3, label='Beam'), 654 | Line2D([0], [0], marker=(3,0,-90), color='red', label='Pin', 655 | markerfacecolor='red', markersize=15), 656 | Line2D([0], [0], marker='o', color='seagreen', label='Rollers', 657 | markerfacecolor='seagreen', markeredgecolor='k', markersize=15), 658 | ] 659 | for color, label in zip(case_colors, case_labels): 660 | legend_elements.append(Line2D([0], [0], color=color, lw=2, label=label)) 661 | ax.legend(handles=legend_elements, loc='lower right', fontsize=12) 662 | 663 | plt.tight_layout() 664 | plt.show() 665 | -------------------------------------------------------------------------------- /OpenPyStruct_FNN_MultiCase.py: -------------------------------------------------------------------------------- 1 | ############################################################################ 2 | #### OpenPyStruct FNN with Residual Blocks Based Multi Load Case Optimizer # 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ############################################################################ 5 | 6 | import os 7 | import json 8 | import time 9 | import math 10 | import random 11 | import numpy as np 12 | import matplotlib.pyplot as plt 13 | import matplotlib.cm as cm 14 | from matplotlib.lines import Line2D 15 | from matplotlib.patches import FancyArrowPatch, RegularPolygon, Rectangle 16 | import seaborn as sns 17 | 18 | import torch 19 | import torch.nn as nn 20 | import torch.optim as optim 21 | from torch.optim.lr_scheduler import ExponentialLR 22 | from torch.utils.data import DataLoader, TensorDataset 23 | from sklearn.preprocessing import StandardScaler 24 | from sklearn.metrics import r2_score 25 | from scipy.stats import mode 26 | 27 | # Use new AMP API to avoid deprecation warnings 28 | from torch.cuda.amp import autocast, GradScaler 29 | 30 | ####################################### 31 | # 1) CONFIGURATION & HYPERPARAMETERS 32 | ####################################### 33 | 34 | # Model and training configuration # 35 | n_cases = 6 # Number of sub-cases per sample 36 | nelem = 100 # Final output dimension per sample: (B, n_elem) 37 | box_constraint_coeff = 5e-1 # Coefficient for box constraint penalty 38 | hidden_units = 128 # Number of hidden units in MLP 39 | dropout_rate = 0.5 # Dropout rate for regularization 40 | num_blocks = 3 # Number of blocks (unused in current model) 41 | num_epochs = 500 # Maximum number of training epochs 42 | batch_size = 128 # Batch size for training 43 | patience = 10 # Early stopping patience 44 | learning_rate = 2e-4 # Learning rate for optimizer 45 | weight_decay = 1e-2 # Weight decay (L2 regularization) for optimizer 46 | train_split = 0.8 # Fraction of data used for training 47 | sigma_0 = 0.03 # Initial Gaussian noise for input 48 | gamma_noise = 0.97 # Decay rate for noise during training 49 | gamma = 0.99 # Learning rate scheduler decay rate 50 | initial_alpha = 0.5 # Initial alpha value for loss weighting 51 | c = 1.0 # Parameter to adjust label aggregation (higher c = more more conservative I estimate) 52 | 53 | # Device configuration 54 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 55 | print("Using device:", device) 56 | 57 | ####################################### 58 | # 2) DATA LOADING & PREPROCESSING 59 | ####################################### 60 | 61 | def pad_sequences(data_list, max_length, pad_val=0.0): 62 | """ 63 | Pad each 1D array in data_list to 'max_length'. 64 | Returns a NumPy array of shape (num_samples, max_length). 65 | """ 66 | out = np.full((len(data_list), max_length), pad_val, dtype=np.float32) 67 | for i, arr in enumerate(data_list): 68 | arr_np = np.array(arr, dtype=np.float32) 69 | length = min(len(arr_np), max_length) 70 | out[i, :length] = arr_np[:length] 71 | return out 72 | 73 | ## mean ## 74 | def unify_label_with_c(I_3d, c=c): 75 | """ 76 | Aggregate labels by computing the mean across cases and adding c times the standard deviation. 77 | 78 | Parameters: 79 | - I_3d: NumPy array of shape (B, n_cases, n_elem) 80 | - c: Scalar multiplier for standard deviation 81 | 82 | Returns: 83 | - Y: NumPy array of shape (B, n_elem) 84 | """ 85 | I_mean = I_3d.mean(axis=1) # Mean across cases 86 | I_std = I_3d.std(axis=1) # Standard deviation across cases 87 | return I_mean + c * I_std 88 | 89 | def fit_transform_3d(arr_3d, scaler): 90 | """ 91 | Fit and transform a 3D array using the provided scaler over axis=0. 92 | 93 | Parameters: 94 | - arr_3d: NumPy array of shape (B, NC, M) 95 | - scaler: Scaler instance (e.g., StandardScaler) 96 | 97 | Returns: 98 | - scaled_arr: NumPy array of shape (B, NC, M) 99 | """ 100 | B, NC, M = arr_3d.shape 101 | flat = arr_3d.reshape(B * NC, M) # Combine B and NC for fitting 102 | scaled = scaler.fit_transform(flat) 103 | return scaled.reshape(B, NC, M) 104 | 105 | def merge_sub_features(*arrays): 106 | """ 107 | Concatenate multiple feature arrays along the feature dimension. 108 | 109 | Parameters: 110 | - arrays: Variable number of NumPy arrays to concatenate 111 | 112 | Returns: 113 | - merged_array: Concatenated NumPy array 114 | """ 115 | return np.concatenate(arrays, axis=2) 116 | 117 | def pad_feat_dim_to_multiple_of_nheads(X_3d, nheads): 118 | """ 119 | Pad the feature dimension to be a multiple of nheads. 120 | 121 | Parameters: 122 | - X_3d: NumPy array of shape (B, Nc, original_dim) 123 | - nheads: Integer, number of attention heads 124 | 125 | Returns: 126 | - X_3d_padded: Padded NumPy array 127 | - new_dim: New feature dimension after padding 128 | """ 129 | B, Nc, original_dim = X_3d.shape 130 | remainder = original_dim % nheads 131 | if remainder == 0: 132 | return X_3d, original_dim 133 | new_dim = ((original_dim // nheads) + 1) * nheads 134 | diff = new_dim - original_dim 135 | X_3d_padded = np.pad(X_3d, ((0,0), (0,0), (0,diff)), mode='constant') 136 | return X_3d_padded, new_dim 137 | 138 | def scale_user_inputs( 139 | user_roller, user_force_x, user_force_vals, user_node_pos, 140 | scalers, n_cases, max_lengths 141 | ): 142 | """ 143 | Scales the user inputs using the fitted scalers. 144 | 145 | Parameters: 146 | - user_roller: List of roller locations per case 147 | - user_force_x: List of force positions per case 148 | - user_force_vals: List of force values per case 149 | - user_node_pos: List of node positions per case 150 | - scalers: Dictionary of fitted scalers for each feature 151 | - n_cases: Number of cases 152 | - max_lengths: Dictionary of maximum lengths for padding 153 | 154 | Returns: 155 | - feat_3d: NumPy array of shape (1, n_cases, feat_dim) 156 | """ 157 | def pad_to_length(seq, req_len): 158 | arr = np.zeros((req_len,), dtype=np.float32) 159 | length = min(len(seq), req_len) 160 | arr[:length] = seq[:length] 161 | return arr 162 | 163 | feat_arrays = [] 164 | for i in range(n_cases): 165 | # Pad each input type to its respective max length 166 | r_pad = pad_to_length(user_roller[i], max_lengths['roller_x']) 167 | fx_pad = pad_to_length(user_force_x[i], max_lengths['force_x']) 168 | fv_pad = pad_to_length(user_force_vals[i], max_lengths['force_values']) 169 | nd_pad = pad_to_length(user_node_pos[i], max_lengths['node_positions']) 170 | 171 | # Scale using the fitted scalers 172 | r_scaled = scalers["roller_x"].transform(r_pad.reshape(1, -1)).flatten() 173 | fx_scaled = scalers["force_x"].transform(fx_pad.reshape(1, -1)).flatten() 174 | fv_scaled = scalers["force_values"].transform(fv_pad.reshape(1, -1)).flatten() 175 | nd_scaled = scalers["node_positions"].transform(nd_pad.reshape(1, -1)).flatten() 176 | 177 | # Concatenate scaled features 178 | sub_feat = np.concatenate([r_scaled, fx_scaled, fv_scaled, nd_scaled]) 179 | feat_arrays.append(sub_feat) 180 | 181 | feat_2d = np.stack(feat_arrays, axis=0) # (n_cases, total_feat_dim) 182 | feat_3d = feat_2d[np.newaxis, ...] # (1, n_cases, total_feat_dim) 183 | return feat_3d 184 | 185 | # Load data 186 | try: 187 | with open("StructDataLite.json", "r") as f: 188 | data = json.load(f) 189 | except FileNotFoundError: 190 | raise FileNotFoundError("The file 'StructDataLite.json' was not found.") 191 | 192 | # Extract data 193 | roller_x = data.get("roller_x_locations", []) 194 | force_x = data.get("force_x_locations", []) 195 | force_values = data.get("force_values", []) 196 | node_positions = data.get("node_positions", []) 197 | I_values = data.get("I_values", []) 198 | 199 | num_samples = len(I_values) 200 | req_keys = ["roller_x_locations","force_x_locations","force_values","node_positions"] 201 | if not all(len(data.get(k, [])) == num_samples for k in req_keys): 202 | raise ValueError("Mismatch in sample counts among roller_x, force_x, force_values, node_positions.") 203 | 204 | # Determine maximum lengths for padding 205 | max_lengths = { 206 | "roller_x": max(len(r) for r in roller_x) if roller_x else 0, 207 | "force_x": max(len(r) for r in force_x) if force_x else 0, 208 | "force_values": max(len(r) for r in force_values) if force_values else 0, 209 | "node_positions": max(len(r) for r in node_positions) if node_positions else 0, 210 | "I_values": max(len(r) for r in I_values) if I_values else 0 211 | } 212 | 213 | # Pad sequences 214 | roller_x_pad = pad_sequences(roller_x, max_lengths["roller_x"]) 215 | force_x_pad = pad_sequences(force_x, max_lengths["force_x"]) 216 | force_val_pad = pad_sequences(force_values, max_lengths["force_values"]) 217 | node_pos_pad = pad_sequences(node_positions, max_lengths["node_positions"]) 218 | I_values_pad = pad_sequences(I_values, max_lengths["I_values"]) 219 | 220 | # Group data by n_cases 221 | total_grouped = num_samples // n_cases 222 | if total_grouped == 0: 223 | raise ValueError(f"n_cases={n_cases} > total samples={num_samples}.") 224 | 225 | trim_len = total_grouped * n_cases 226 | roller_x_pad = roller_x_pad[:trim_len] 227 | force_x_pad = force_x_pad[:trim_len] 228 | force_val_pad = force_val_pad[:trim_len] 229 | node_pos_pad = node_pos_pad[:trim_len] 230 | I_values_pad = I_values_pad[:trim_len] 231 | 232 | roller_grouped = roller_x_pad.reshape(total_grouped, n_cases, -1) 233 | force_x_grouped = force_x_pad.reshape(total_grouped, n_cases, -1) 234 | force_val_grouped = force_val_pad.reshape(total_grouped, n_cases, -1) 235 | node_grouped = node_pos_pad.reshape(total_grouped, n_cases, -1) 236 | I_grouped = I_values_pad.reshape(total_grouped, n_cases, -1) 237 | 238 | # Train/Validation Split 239 | indices = np.random.permutation(total_grouped) 240 | train_sz = np.int32(train_split * total_grouped) 241 | train_idx = indices[:train_sz] 242 | val_idx = indices[train_sz:] 243 | 244 | roller_train = roller_grouped[train_idx] 245 | roller_val = roller_grouped[val_idx] 246 | force_x_train = force_x_grouped[train_idx] 247 | force_x_val = force_x_grouped[val_idx] 248 | force_val_train = force_val_grouped[train_idx] 249 | force_val_val = force_val_grouped[val_idx] 250 | node_train = node_grouped[train_idx] 251 | node_val = node_grouped[val_idx] 252 | I_train = I_grouped[train_idx] 253 | I_val = I_grouped[val_idx] 254 | 255 | # Initialize Scalers 256 | scalers_inputs = { 257 | "roller_x": StandardScaler(), 258 | "force_x": StandardScaler(), 259 | "force_values": StandardScaler(), 260 | "node_positions": StandardScaler() 261 | } 262 | 263 | scaler_Y = StandardScaler() 264 | 265 | # Fit and transform training data 266 | roller_train_std = fit_transform_3d(roller_train, scalers_inputs["roller_x"]) 267 | force_x_train_std = fit_transform_3d(force_x_train, scalers_inputs["force_x"]) 268 | force_val_train_std = fit_transform_3d(force_val_train, scalers_inputs["force_values"]) 269 | node_train_std = fit_transform_3d(node_train, scalers_inputs["node_positions"]) 270 | 271 | # **Corrected**: Transform validation data using the already fitted scalers 272 | roller_val_std = scalers_inputs["roller_x"].transform(roller_val.reshape(-1, roller_val.shape[-1])).reshape(roller_val.shape) 273 | force_x_val_std = scalers_inputs["force_x"].transform(force_x_val.reshape(-1, force_x_val.shape[-1])).reshape(force_x_val.shape) 274 | force_val_val_std = scalers_inputs["force_values"].transform(force_val_val.reshape(-1, force_val_val.shape[-1])).reshape(force_val_val.shape) 275 | node_val_std = scalers_inputs["node_positions"].transform(node_val.reshape(-1, node_val.shape[-1])).reshape(node_val.shape) 276 | 277 | # Prepare final input features by merging sub-features 278 | X_train_3d = merge_sub_features( 279 | roller_train_std, 280 | force_x_train_std, 281 | force_val_train_std, 282 | node_train_std 283 | ) 284 | 285 | X_val_3d = merge_sub_features( 286 | roller_val_std, 287 | force_x_val_std, 288 | force_val_val_std, 289 | node_val_std 290 | ) 291 | 292 | # Flatten the feature dimension for FNN 293 | X_train_flat = X_train_3d.reshape(X_train_3d.shape[0], -1) # Shape: (B, n_cases * feat_dim) 294 | X_val_flat = X_val_3d.reshape(X_val_3d.shape[0], -1) # Shape: (B, n_cases * feat_dim) 295 | 296 | # Unify the label by aggregating across cases 297 | Y_train_2d = unify_label_with_c(I_train, c=c) # Shape: (B, n_elem) 298 | Y_val_2d = unify_label_with_c(I_val, c=c) # Shape: (B, n_elem) 299 | 300 | # Fit scaler_Y on training targets 301 | scaler_Y.fit(Y_train_2d) 302 | 303 | # Transform targets 304 | Y_train_std = scaler_Y.transform(Y_train_2d) # Shape: (B, n_elem) 305 | Y_val_std = scaler_Y.transform(Y_val_2d) # Shape: (B, n_elem) 306 | 307 | # Convert to PyTorch Tensors 308 | X_train_tensor = torch.tensor(X_train_flat, dtype=torch.float32) 309 | Y_train_tensor = torch.tensor(Y_train_std, dtype=torch.float32) 310 | X_val_tensor = torch.tensor(X_val_flat, dtype=torch.float32) 311 | Y_val_tensor = torch.tensor(Y_val_std, dtype=torch.float32) 312 | 313 | min_constraint = torch.min(Y_train_tensor) 314 | max_constraint = torch.max(Y_train_tensor) 315 | 316 | # Create TensorDatasets 317 | train_dataset = TensorDataset(X_train_tensor, Y_train_tensor) 318 | val_dataset = TensorDataset(X_val_tensor, Y_val_tensor) 319 | 320 | print("X_train_flat shape:", X_train_flat.shape, # e.g. (10000, 10*125) 321 | "Y_train_2d shape:", Y_train_2d.shape, # => (10000, 100) 322 | "\nBut after padding, X_train_flat shape:", 323 | X_train_flat.shape, # e.g. (10000, 1250) 324 | ) 325 | 326 | ####################################### 327 | # 3) DEFINE THE MODEL COMPONENTS 328 | ####################################### 329 | 330 | class ResidualBlock(nn.Module): 331 | """ 332 | A standard residual block with two linear layers and a skip connection. 333 | """ 334 | def __init__(self, input_dim, hidden_dim, dropout_rate): 335 | super(ResidualBlock, self).__init__() 336 | self.fc1 = nn.Linear(input_dim, input_dim) 337 | self.LeakyReLU = nn.LeakyReLU(0.01) 338 | self.dropout = nn.Dropout(dropout_rate) 339 | self.norm = nn.LayerNorm(input_dim) 340 | 341 | def forward(self, x): 342 | residual = x 343 | out = self.fc1(x) 344 | out = self.LeakyReLU(out) 345 | out = self.dropout(out) 346 | out += residual 347 | out = self.norm(out) 348 | out = self.LeakyReLU(out) 349 | return out 350 | 351 | class FNNWithResidual(nn.Module): 352 | """ 353 | Feedforward Neural Network with Residual Blocks. 354 | """ 355 | def __init__( 356 | self, 357 | input_dim, 358 | hidden_dim, 359 | num_residual_blocks, 360 | output_dim, 361 | dropout_rate 362 | ): 363 | super(FNNWithResidual, self).__init__() 364 | self.input_fc = nn.Linear(input_dim, hidden_dim) 365 | self.LeakyReLU = nn.LeakyReLU(0.01) 366 | self.dropout = nn.Dropout(dropout_rate) 367 | self.residual_blocks = nn.ModuleList([ 368 | ResidualBlock(hidden_dim, hidden_dim * 2, dropout_rate) 369 | for _ in range(num_residual_blocks) 370 | ]) 371 | self.output_fc = nn.Linear(hidden_dim, output_dim) 372 | 373 | def forward(self, x): 374 | out = self.input_fc(x) 375 | out = self.LeakyReLU(out) 376 | out = self.dropout(out) 377 | for block in self.residual_blocks: 378 | out = block(out) 379 | out = self.output_fc(out) 380 | return out 381 | 382 | ####################################### 383 | # 4) DEFINE CUSTOM LOSS 384 | ####################################### 385 | 386 | class TrainableL1L2Loss(nn.Module): 387 | """ 388 | Combines L1 and L2 loss with a trainable alpha parameter and 389 | penalizes predictions outside [min_constraint, max_constraint]. 390 | """ 391 | def __init__( 392 | self, 393 | initial_alpha=initial_alpha, 394 | min_constraint=min_constraint, 395 | max_constraint=max_constraint, 396 | penalty_weight=box_constraint_coeff 397 | ): 398 | super().__init__() 399 | # Initialize alpha as a learnable parameter 400 | self.alpha = nn.Parameter(torch.tensor(initial_alpha, dtype=torch.float32, requires_grad=True)) 401 | self.l1 = nn.L1Loss() 402 | self.l2 = nn.MSELoss() 403 | self.min_constraint = min_constraint 404 | self.max_constraint = max_constraint 405 | self.penalty_weight = penalty_weight 406 | 407 | def forward(self, preds, targets): 408 | """ 409 | Compute the combined loss. 410 | 411 | Parameters: 412 | - preds: Predicted tensor 413 | - targets: Target tensor 414 | 415 | Returns: 416 | - total_loss: Scalar tensor representing the loss 417 | """ 418 | # Clamp alpha to avoid extreme weighting 419 | alpha = torch.clamp(self.alpha, 1e-6, 1.0) 420 | 421 | # Compute L1 and L2 losses 422 | l1_loss = self.l1(preds, targets) 423 | l2_loss = self.l2(preds, targets) 424 | 425 | # Initialize penalty 426 | penalty = 0.0 427 | if self.min_constraint is not None: 428 | # Penalize predictions below the minimum constraint 429 | below_min_penalty = torch.sum(torch.relu(self.min_constraint - preds)) 430 | penalty += below_min_penalty 431 | if self.max_constraint is not None: 432 | # Penalize predictions above the maximum constraint 433 | above_max_penalty = torch.sum(torch.relu(preds - self.max_constraint)) 434 | penalty += above_max_penalty 435 | 436 | # Combine losses with trainable alpha and add penalty 437 | total_loss = alpha * l1_loss + (1 - alpha) * l2_loss + self.penalty_weight * penalty 438 | return total_loss 439 | 440 | def permute_data(X, Y): 441 | """ 442 | Permutes the data indices for both X and Y consistently. 443 | 444 | Parameters: 445 | - X (torch.Tensor): The input data tensor. 446 | - Y (torch.Tensor): The target data tensor. 447 | 448 | Returns: 449 | - X_permuted (torch.Tensor): The permuted input data. 450 | - Y_permuted (torch.Tensor): The permuted target data. 451 | """ 452 | assert X.size(0) == Y.size(0), "X and Y must have the same number of samples to permute." 453 | 454 | # Generate a random permutation of indices 455 | perm = torch.randperm(X.size(0), device=X.device) 456 | 457 | # Apply the permutation 458 | X_permuted = X[perm] 459 | Y_permuted = Y[perm] 460 | 461 | return X_permuted, Y_permuted 462 | 463 | ####################################### 464 | # 5) INITIALIZE & TRAIN 465 | ####################################### 466 | 467 | # Determine input dimension for FNN 468 | input_dim = X_train_flat.shape[1] 469 | output_dim = nelem # As defined earlier 470 | 471 | # Initialize the model 472 | model = FNNWithResidual( 473 | input_dim=input_dim, 474 | hidden_dim=hidden_units, 475 | num_residual_blocks=4, # Number of residual blocks can be adjusted 476 | output_dim=output_dim, 477 | dropout_rate=dropout_rate 478 | ).to(device) 479 | 480 | # Initialize optimizer, scheduler, and loss criterion 481 | optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay) 482 | scheduler = ExponentialLR(optimizer, gamma=gamma) 483 | criterion = TrainableL1L2Loss() 484 | 485 | # Initialize DataLoaders 486 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 487 | val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 488 | 489 | # Initialize AMP scaler for mixed precision training 490 | scaler_amp = GradScaler() # 'cuda' device is inferred automatically 491 | 492 | # Setup live plotting for training progress 493 | plt.ion() 494 | fig, ax = plt.subplots(figsize=(8, 6), dpi=300) 495 | 496 | def live_plot(epoch, train_losses, val_losses): 497 | """ 498 | Updates the training and validation loss plot during training. 499 | 500 | Parameters: 501 | - epoch: Current epoch number 502 | - train_losses: List of training losses up to current epoch 503 | - val_losses: List of validation losses up to current epoch 504 | """ 505 | ax.clear() 506 | ax.plot(range(1, epoch + 1), train_losses, label="Train Loss", marker='o', color='blue') 507 | ax.plot(range(1, epoch + 1), val_losses, label="Validation Loss", marker='x', color='red') 508 | 509 | ax.set_xlabel("Epochs") 510 | ax.set_ylabel("Loss") 511 | ax.set_title("Training and Validation Loss") 512 | ax.legend() 513 | ax.grid(True, linestyle='--', alpha=0.7) 514 | 515 | plt.pause(0.01) 516 | 517 | # Initialize lists to store losses 518 | train_losses = [] 519 | val_losses = [] 520 | 521 | best_val_loss = float('inf') 522 | epochs_no_improve = 0 523 | 524 | # Training loop 525 | for epoch in range(1, num_epochs + 1): 526 | model.train() 527 | noise_level = sigma_0 * (gamma_noise ** epoch) # Decaying noise level 528 | 529 | total_train_loss = 0.0 530 | t0 = time.time() 531 | 532 | for Xb, Yb in train_loader: 533 | Xb = Xb.to(device) 534 | Yb = Yb.to(device) 535 | 536 | # Permute Xb and Yb consistently 537 | Xb, Yb = permute_data(Xb, Yb) 538 | 539 | # Add optional Gaussian noise 540 | Xb_noisy = Xb + torch.randn_like(Xb) * noise_level 541 | 542 | optimizer.zero_grad() 543 | with autocast(): 544 | preds = model(Xb_noisy) # Predictions: (B, n_elem) 545 | # Mild penalty on deviation of alpha from its initial value 546 | L_alpha = (initial_alpha - criterion.alpha) ** 2 547 | loss = criterion(preds, Yb) + L_alpha 548 | 549 | # Backpropagation with mixed precision 550 | scaler_amp.scale(loss).backward() 551 | scaler_amp.unscale_(optimizer) 552 | torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Gradient clipping 553 | scaler_amp.step(optimizer) 554 | scaler_amp.update() 555 | 556 | total_train_loss += loss.item() 557 | 558 | avg_train_loss = total_train_loss / len(train_loader) 559 | train_losses.append(avg_train_loss) # Store the training loss 560 | 561 | # Validation phase 562 | model.eval() 563 | total_val_loss = 0.0 564 | with torch.no_grad(), autocast(): 565 | for Xb, Yb in val_loader: 566 | Xb = Xb.to(device) 567 | Yb = Yb.to(device) 568 | preds = model(Xb) 569 | val_loss = criterion(preds, Yb) 570 | total_val_loss += val_loss.item() 571 | 572 | avg_val_loss = total_val_loss / len(val_loader) 573 | val_losses.append(avg_val_loss) # Store the validation loss 574 | scheduler.step() # Update learning rate 575 | 576 | # Early Stopping Check 577 | if avg_val_loss < best_val_loss: 578 | best_val_loss = avg_val_loss 579 | epochs_no_improve = 0 580 | torch.save(model.state_dict(), "best_model_fnn_residual.pth") # Save the best model 581 | else: 582 | epochs_no_improve += 1 583 | if epochs_no_improve >= patience: 584 | print(f"Early stopping at epoch {epoch}") 585 | break 586 | 587 | dt = time.time() - t0 588 | print(f"Epoch {epoch}/{num_epochs} | " 589 | f"Train Loss={avg_train_loss:.6f}, " 590 | f"Val Loss={avg_val_loss:.6f}, " 591 | f"Time={dt:.2f}s") 592 | 593 | # Update live plot 594 | live_plot(epoch, train_losses, val_losses) 595 | 596 | 597 | 598 | ####################################### 599 | # 6) EVALUATION 600 | ####################################### 601 | 602 | # Load the best model 603 | model.load_state_dict(torch.load("best_model_fnn_residual.pth", map_location=device)) 604 | model.eval() 605 | 606 | # Create a DataLoader for evaluation 607 | val_loader_eval = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 608 | all_preds, all_labels = [], [] 609 | 610 | # Collect all predictions and labels 611 | with torch.no_grad(): 612 | for Xb, Yb in val_loader_eval: 613 | Xb = Xb.to(device) 614 | preds = model(Xb) # Predictions: (B, n_elem) 615 | all_preds.append(preds.cpu()) 616 | all_labels.append(Yb) 617 | 618 | # Concatenate all predictions and labels 619 | all_preds = torch.cat(all_preds, dim=0).numpy() # Shape: (N, n_elem) 620 | all_labels = torch.cat(all_labels, dim=0).numpy() # Shape: (N, n_elem) 621 | 622 | # Un-standardize using scaler_Y 623 | all_preds_unstd = scaler_Y.inverse_transform(all_preds) 624 | all_labels_unstd = scaler_Y.inverse_transform(all_labels) 625 | 626 | # Clip predictions and labels to a reasonable range if necessary 627 | all_preds_unstd = np.clip(all_preds_unstd, 0.0, 1e10) 628 | all_labels_unstd = np.clip(all_labels_unstd, 0.0, 1e10) 629 | 630 | # Compute R² score to evaluate model performance 631 | r2_val = r2_score(all_labels_unstd.ravel(), all_preds_unstd.ravel()) 632 | print(f"R² on Validation: {r2_val:.4f}") 633 | 634 | 635 | 636 | 637 | ####################################### 638 | # 7) EXAMPLE INFERENCE & PLOT 639 | ####################################### 640 | 641 | # Example User Inputs 642 | L_beam = 200 # Beam length (m) 643 | Fmin_user = -355857 # Min point load to be randomized (N) 644 | Fmax_user = Fmin_user / 10 # Max point load to be randomized (N) 645 | user_rollers = [2*9, 2*29, 2*69, 2*85, 2*100] # Roller locations (m) ## if using data from multi-core generator, subtract 1 from the node number 646 | 647 | def build_user_input_no_agg( 648 | roller_list, force_x_list, force_val_list, node_pos_list, 649 | scalers, n_cases, 650 | max_lengths 651 | ): 652 | feat_3d = scale_user_inputs( 653 | roller_list, force_x_list, force_val_list, node_pos_list, 654 | scalers, n_cases, max_lengths 655 | ) 656 | feat_flat = feat_3d.reshape(1, -1) # Flatten for FNN 657 | return feat_flat 658 | 659 | # Assign the same rollers to each case 660 | user_roller = [user_rollers.copy() for _ in range(n_cases)] 661 | 662 | # Generate diverse force positions and values for each case 663 | user_force_x = [] 664 | user_force_vals = [] 665 | for _ in range(n_cases): 666 | num_forces = random.randint(1, 3) # Each case has 1 to 3 forces 667 | fx = sorted([random.uniform(0, L_beam) for _ in range(num_forces)]) 668 | fv = [random.uniform(Fmin_user, Fmax_user) for _ in range(num_forces)] 669 | user_force_x.append(fx) 670 | user_force_vals.append(fv) 671 | 672 | # Node positions remain consistent across cases 673 | user_node_pos = [ 674 | np.linspace(0, L_beam, nelem + 1).tolist() for _ in range(n_cases) 675 | ] 676 | 677 | # Scale inputs for prediction using the correct function name 678 | X_user_flat = build_user_input_no_agg( 679 | user_roller, user_force_x, user_force_vals, user_node_pos, 680 | scalers_inputs, n_cases, max_lengths 681 | ) 682 | 683 | X_user_t = torch.tensor(X_user_flat, dtype=torch.float32).to(device) 684 | 685 | # Perform inference 686 | model.eval() 687 | with torch.no_grad(): 688 | pred_1x = model(X_user_t) # Predictions: (1, n_elem) 689 | 690 | # Un-standardize predictions 691 | pred_1x_np = pred_1x.cpu().numpy().squeeze() 692 | pred_1x_unstd = scaler_Y.inverse_transform(pred_1x_np.reshape(1, -1)).squeeze() 693 | 694 | # Visualization Parameters 695 | unique_rollers = sorted(set([x for sublist in user_roller for x in sublist] + [L_beam])) 696 | case_colors = sns.color_palette("Set1", n_colors=n_cases) 697 | case_labels = [f'Force Case {i+1} (N)' for i in range(n_cases)] 698 | 699 | beam_y = 0 700 | beam_x = [0, L_beam] 701 | beam_y_vals = [beam_y, beam_y] 702 | 703 | # Collect force positions and values for plotting 704 | force_positions = [] 705 | force_vals_plot = [] 706 | for fx, fv in zip(user_force_x, user_force_vals): 707 | for xx, val in zip(fx, fv): 708 | force_positions.append(xx) 709 | force_vals_plot.append(val) 710 | 711 | all_force_vals = force_vals_plot 712 | max_force = max(abs(val) for val in all_force_vals) if all_force_vals else 1.0 713 | desired_max_arrow_length = 2.0 714 | arrow_scale = desired_max_arrow_length / max_force if max_force != 0 else 1.0 715 | 716 | beam_positions = user_node_pos[0][:nelem] 717 | 718 | # Normalize I for visualization purposes using a colormap 719 | I_normalized = (pred_1x_unstd - pred_1x_unstd.min()) / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8) 720 | cmap = cm.winter # Changed from cm.viridis to cm.winter 721 | norm = plt.Normalize(pred_1x_unstd.min(), pred_1x_unstd.max()) 722 | 723 | # Define block dimensions 724 | block_width = L_beam / nelem * 0.8 # Slightly narrower blocks for better visibility 725 | block_height = 1 # Maximum height for visualization 726 | 727 | # Create a single plot with extra space for the colorbar 728 | fig, ax = plt.subplots(figsize=(18, 7), dpi=300) 729 | 730 | # Plot Beam 731 | ax.plot(beam_x, beam_y_vals, color='black', linewidth=3, label='Beam') 732 | ax.scatter(beam_x[0], beam_y - 0.15, marker='^', color='red', s=300, zorder=6) 733 | 734 | # Plot Rollers 735 | ax.scatter(unique_rollers, [beam_y]*len(unique_rollers), 736 | marker='o', color='seagreen', s=200, 737 | label='Rollers', zorder=5, edgecolors='k') 738 | 739 | # Plot Forces 740 | for case_idx in range(n_cases): 741 | fx_list = user_force_x[case_idx] 742 | fv_list = user_force_vals[case_idx] 743 | color = case_colors[case_idx] 744 | label = case_labels[case_idx] 745 | 746 | for idx, (fx, fv) in enumerate(zip(fx_list, fv_list)): 747 | arrow_length = abs(fv) * arrow_scale 748 | start_point = (fx, beam_y + arrow_length) 749 | end_point = (fx, beam_y) 750 | 751 | arrow = FancyArrowPatch( 752 | posA=start_point, posB=end_point, 753 | arrowstyle='-|>', 754 | mutation_scale=20, 755 | color=color, 756 | linewidth=2, 757 | alpha=0.8, 758 | label=label if idx == 0 else "" 759 | ) 760 | ax.add_patch(arrow) 761 | ax.text(fx, beam_y + arrow_length + desired_max_arrow_length * 0.02, 762 | f"{fv:.0f}", ha='center', va='bottom', 763 | fontsize=10, color=color, fontweight='bold') 764 | 765 | # Plot Moments of Inertia (I) as semi-translucent blocks centered on the beam 766 | for idx, (x_pos, I_val) in enumerate(zip(beam_positions, pred_1x_unstd)): 767 | # Normalize I for color mapping 768 | color = cmap(norm(I_val)) 769 | 770 | # Define the rectangle parameters 771 | rect_x = x_pos - block_width / 2 # Center the rectangle horizontally 772 | rect_y = beam_y - (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8)) * block_height / 2 # Center vertically 773 | 774 | rect = Rectangle((rect_x, rect_y), block_width, (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8)) * block_height, 775 | linewidth=0, edgecolor=None, 776 | facecolor=color, alpha=0.6) 777 | ax.add_patch(rect) 778 | 779 | # Create a ScalarMappable for the colorbar 780 | sm = cm.ScalarMappable(cmap=cmap, norm=norm) 781 | sm.set_array([]) # Only needed for older versions of matplotlib 782 | 783 | # Add the colorbar to the figure (vertically to the right) 784 | cbar = fig.colorbar(sm, ax=ax, orientation='vertical', fraction=0.046, pad=0.04) 785 | cbar.set_label('Predicted I (m$^4$)', fontsize=16) 786 | cbar.ax.tick_params(labelsize=10) 787 | 788 | 789 | # Set plot titles and labels 790 | ax.set_title("Beam Setup with Applied Forces and Predicted I", 791 | fontsize=22, fontweight='bold', pad=20) 792 | ax.set_xlabel("Beam Length (m)", fontsize=16, fontweight='semibold') 793 | #ax.set_ylabel("Force Representation and Moments of Inertia, I (m$^4$)", fontsize=16, fontweight='semibold') 794 | ax.set_xlim(-5, L_beam + 5) 795 | ax.set_ylim(-2.5, 2.5) # Adjusted to accommodate I blocks 796 | ax.set_xticks(np.arange(0, L_beam + 25, 25)) 797 | ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7) 798 | 799 | # Create custom legend 800 | legend_elements = [ 801 | Line2D([0], [0], color='black', lw=3, label='Beam'), 802 | Line2D([0], [0], marker=(3, 0, -90), color='red', label='Pin', 803 | markerfacecolor='red', markersize=15), 804 | Line2D([0], [0], marker='o', color='seagreen', label='Rollers', 805 | markerfacecolor='seagreen', markeredgecolor='k', markersize=15), 806 | ] 807 | 808 | for color, label in zip(case_colors, case_labels): 809 | legend_elements.append(Line2D([0], [0], color=color, lw=2, label=label)) 810 | 811 | # Remove the 'Predicted I' legend entry since the colorbar now represents it 812 | # legend_elements.append(Rectangle((0,0),1,1, facecolor='green', alpha=0.6, label='Predicted I')) 813 | 814 | ax.legend(handles=legend_elements, loc='lower right', fontsize=12) 815 | 816 | plt.tight_layout() 817 | plt.show() 818 | -------------------------------------------------------------------------------- /OpenPyStruct_FNO_MultiCase_Beta.py: -------------------------------------------------------------------------------- 1 | ############################################################################ 2 | #### OpenPyStruct FNO-Based Multi Load Case Optimizer #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ############################################################################ 5 | 6 | import os 7 | import json 8 | import time 9 | import math 10 | import random 11 | import numpy as np 12 | import matplotlib.pyplot as plt 13 | import matplotlib.cm as cm 14 | from matplotlib.lines import Line2D 15 | from matplotlib.patches import FancyArrowPatch, Rectangle 16 | import seaborn as sns 17 | 18 | import torch 19 | import torch.nn as nn 20 | import torch.optim as optim 21 | from torch.optim.lr_scheduler import ExponentialLR 22 | from torch.utils.data import DataLoader, TensorDataset 23 | from sklearn.preprocessing import StandardScaler 24 | from sklearn.metrics import r2_score 25 | from scipy.stats import mode 26 | 27 | # Use new AMP API to avoid deprecation warnings 28 | from torch.amp import autocast, GradScaler 29 | 30 | ####################################### 31 | # 1) CONFIGURATION & HYPERPARAMETERS 32 | ####################################### 33 | 34 | # Model and training configuration # 35 | n_cases = 6 # Number of sub-cases per sample 36 | nelem = 100 # Final output dimension per sample: (B, n_elem) 37 | box_constraint_coeff = 5e-1 # Coefficient for box constraint penalty 38 | hidden_units = 512 # Number of hidden units in final MLP 39 | dropout_rate = 0.1 # Dropout rate for regularization 40 | num_fno_layers = 4 # Number of FNO layers (analogous to "blocks") 41 | num_epochs = 500 # Maximum number of training epochs 42 | batch_size = 512 # Batch size for training 43 | patience = 10 # Early stopping patience 44 | learning_rate = 3e-3 # Learning rate for optimizer 45 | weight_decay = 1e-6 # Weight decay (L2 regularization) for optimizer 46 | train_split = 0.8 # Fraction of data used for training 47 | sigma_0 = 0.01 # Initial Gaussian noise for input 48 | gamma_noise = 0.95 # Decay rate for noise during training 49 | gamma = 0.975 # Learning rate scheduler decay rate 50 | initial_alpha = 0.5 # Initial alpha value for loss weighting 51 | c = 0.5 # Parameter to adjust label aggregation 52 | 53 | # FNO hyperparameters 54 | fno_modes = 4 # Adjusted to satisfy modes <= n_cases//2 + 1 (which is 4 for n_cases=6) 55 | fno_width = 128 # Channel width (hidden size) in FNO layers 56 | 57 | # For optional feature-dimension padding 58 | # (Typically for Transformers, but we leave it at 1 for "no-op" padding in FNO.) 59 | nheads_for_padding = 1 60 | 61 | # Device configuration 62 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 63 | print("Using device:", device) 64 | 65 | ####################################### 66 | # 2) DATA LOADING & PREPROCESSING 67 | ####################################### 68 | 69 | def pad_sequences(data_list, max_length, pad_val=0.0): 70 | """ 71 | Pad each 1D array in data_list to 'max_length'. 72 | Returns a NumPy array of shape (num_samples, max_length). 73 | """ 74 | out = np.full((len(data_list), max_length), pad_val, dtype=np.float32) 75 | for i, arr in enumerate(data_list): 76 | arr_np = np.array(arr, dtype=np.float32) 77 | length = min(len(arr_np), max_length) 78 | out[i, :length] = arr_np[:length] 79 | return out 80 | 81 | def unify_label_with_c(I_3d, c=0.5): 82 | """ 83 | Aggregate labels by computing the mean across cases and adding 84 | c times the standard deviation. 85 | 86 | Parameters: 87 | - I_3d: NumPy array of shape (B, n_cases, n_elem) 88 | - c: Scalar multiplier for standard deviation 89 | 90 | Returns: 91 | - Y: NumPy array of shape (B, n_elem) 92 | """ 93 | I_mean = I_3d.mean(axis=1) # Mean across cases 94 | I_std = I_3d.std(axis=1) # Standard deviation across cases 95 | return I_mean + c * I_std 96 | 97 | def fit_transform_3d(arr_3d, scaler): 98 | """ 99 | Fit and transform a 3D array using the provided scaler over axis=0. 100 | 101 | Parameters: 102 | - arr_3d: NumPy array of shape (B, NC, M) 103 | - scaler: Scaler instance (e.g., StandardScaler) 104 | 105 | Returns: 106 | - scaled_arr: NumPy array of shape (B, NC, M) 107 | """ 108 | B, NC, M = arr_3d.shape 109 | flat = arr_3d.reshape(B * NC, M) 110 | scaled = scaler.fit_transform(flat) 111 | return scaled.reshape(B, NC, M) 112 | 113 | def transform_3d(arr_3d, scaler): 114 | """ 115 | Transform (but do not fit) a 3D array using the provided scaler, 116 | maintaining (B, NC, M) shape. 117 | 118 | Parameters: 119 | - arr_3d: NumPy array of shape (B, NC, M) 120 | - scaler: Fitted scaler instance 121 | 122 | Returns: 123 | - scaled_arr: NumPy array of shape (B, NC, M) 124 | """ 125 | B, NC, M = arr_3d.shape 126 | flat = arr_3d.reshape(B * NC, M) 127 | scaled = scaler.transform(flat) 128 | return scaled.reshape(B, NC, M) 129 | 130 | def merge_sub_features(*arrays): 131 | """ 132 | Concatenate multiple feature arrays along the feature dimension. 133 | e.g., if each array is (B, NC, F?), we get (B, NC, sum_of_F). 134 | """ 135 | return np.concatenate(arrays, axis=2) 136 | 137 | def pad_feat_dim_to_multiple_of_nheads(X_3d, nheads=1): 138 | """ 139 | Pad the feature dimension to be a multiple of 'nheads'. 140 | For FNO, this is not strictly necessary, but left here as a no-op 141 | if nheads=1. 142 | """ 143 | B, Nc, original_dim = X_3d.shape 144 | remainder = original_dim % nheads 145 | if remainder == 0: 146 | return X_3d, original_dim 147 | new_dim = ((original_dim // nheads) + 1) * nheads 148 | diff = new_dim - original_dim 149 | X_3d_padded = np.pad(X_3d, ((0,0), (0,0), (0,diff)), mode='constant') 150 | return X_3d_padded, new_dim 151 | 152 | def scale_user_inputs( 153 | user_roller, user_force_x, user_force_vals, user_node_pos, 154 | scalers, n_cases, max_lengths 155 | ): 156 | """ 157 | Scales the user inputs using the fitted scalers. 158 | 159 | Parameters: 160 | - user_roller: List of roller locations per case 161 | - user_force_x: List of force positions per case 162 | - user_force_vals: List of force values per case 163 | - user_node_pos: List of node positions per case 164 | - scalers: Dictionary of fitted scalers for each feature 165 | - n_cases: Number of cases 166 | - max_lengths: Dictionary of maximum lengths for padding 167 | 168 | Returns: 169 | - feat_3d: NumPy array of shape (1, n_cases, feat_dim) 170 | """ 171 | def pad_to_length(seq, req_len): 172 | arr = np.zeros((req_len,), dtype=np.float32) 173 | length = min(len(seq), req_len) 174 | arr[:length] = seq[:length] 175 | return arr 176 | 177 | feat_arrays = [] 178 | for i in range(n_cases): 179 | r_pad = pad_to_length(user_roller[i], max_lengths['roller_x']) 180 | fx_pad = pad_to_length(user_force_x[i], max_lengths['force_x']) 181 | fv_pad = pad_to_length(user_force_vals[i], max_lengths['force_values']) 182 | nd_pad = pad_to_length(user_node_pos[i], max_lengths['node_positions']) 183 | 184 | r_scaled = scalers["roller_x"].transform(r_pad.reshape(1, -1)).flatten() 185 | fx_scaled = scalers["force_x"].transform(fx_pad.reshape(1, -1)).flatten() 186 | fv_scaled = scalers["force_values"].transform(fv_pad.reshape(1, -1)).flatten() 187 | nd_scaled = scalers["node_positions"].transform(nd_pad.reshape(1, -1)).flatten() 188 | 189 | # Concatenate scaled features 190 | sub_feat = np.concatenate([r_scaled, fx_scaled, fv_scaled, nd_scaled]) 191 | feat_arrays.append(sub_feat) 192 | 193 | feat_2d = np.stack(feat_arrays, axis=0) # (n_cases, total_feat_dim) 194 | feat_3d = feat_2d[np.newaxis, ...] # (1, n_cases, total_feat_dim) 195 | return feat_3d 196 | 197 | # Load data 198 | try: 199 | with open("StructDataMedium.json", "r") as f: 200 | data = json.load(f) 201 | except FileNotFoundError: 202 | raise FileNotFoundError("The file 'StructDataMedium.json' was not found.") 203 | 204 | # Extract data 205 | roller_x = data.get("roller_x_locations", []) 206 | force_x = data.get("force_x_locations", []) 207 | force_values = data.get("force_values", []) 208 | node_positions = data.get("node_positions", []) 209 | I_values = data.get("I_values", []) 210 | 211 | num_samples = len(I_values) 212 | req_keys = ["roller_x_locations","force_x_locations","force_values","node_positions"] 213 | if not all(len(data.get(k, [])) == num_samples for k in req_keys): 214 | raise ValueError("Mismatch in sample counts among roller_x, force_x, force_values, node_positions.") 215 | 216 | # Determine maximum lengths for padding 217 | max_lengths = { 218 | "roller_x": max(len(r) for r in roller_x) if roller_x else 0, 219 | "force_x": max(len(r) for r in force_x) if force_x else 0, 220 | "force_values": max(len(r) for r in force_values) if force_values else 0, 221 | "node_positions": max(len(r) for r in node_positions) if node_positions else 0, 222 | "I_values": max(len(r) for r in I_values) if I_values else 0 223 | } 224 | 225 | # Pad sequences 226 | roller_x_pad = pad_sequences(roller_x, max_lengths["roller_x"]) 227 | force_x_pad = pad_sequences(force_x, max_lengths["force_x"]) 228 | force_val_pad = pad_sequences(force_values, max_lengths["force_values"]) 229 | node_pos_pad = pad_sequences(node_positions, max_lengths["node_positions"]) 230 | I_values_pad = pad_sequences(I_values, max_lengths["I_values"]) 231 | 232 | # Group data by n_cases 233 | total_grouped = num_samples // n_cases 234 | if total_grouped == 0: 235 | raise ValueError(f"n_cases={n_cases} > total samples={num_samples}.") 236 | 237 | trim_len = total_grouped * n_cases 238 | roller_x_pad = roller_x_pad[:trim_len] 239 | force_x_pad = force_x_pad[:trim_len] 240 | force_val_pad = force_val_pad[:trim_len] 241 | node_pos_pad = node_pos_pad[:trim_len] 242 | I_values_pad = I_values_pad[:trim_len] 243 | 244 | roller_grouped = roller_x_pad.reshape(total_grouped, n_cases, -1) 245 | force_x_grouped = force_x_pad.reshape(total_grouped, n_cases, -1) 246 | force_val_grouped = force_val_pad.reshape(total_grouped, n_cases, -1) 247 | node_grouped = node_pos_pad.reshape(total_grouped, n_cases, -1) 248 | I_grouped = I_values_pad.reshape(total_grouped, n_cases, -1) 249 | 250 | # Train/Validation Split 251 | indices = np.random.permutation(total_grouped) 252 | train_sz = int(train_split * total_grouped) 253 | train_idx = indices[:train_sz] 254 | val_idx = indices[train_sz:] 255 | 256 | roller_train = roller_grouped[train_idx] 257 | roller_val = roller_grouped[val_idx] 258 | force_x_train = force_x_grouped[train_idx] 259 | force_x_val = force_x_grouped[val_idx] 260 | force_val_train = force_val_grouped[train_idx] 261 | force_val_val = force_val_grouped[val_idx] 262 | node_train = node_grouped[train_idx] 263 | node_val = node_grouped[val_idx] 264 | I_train = I_grouped[train_idx] 265 | I_val = I_grouped[val_idx] 266 | 267 | # Initialize Scalers 268 | scalers_inputs = { 269 | "roller_x": StandardScaler(), 270 | "force_x": StandardScaler(), 271 | "force_values": StandardScaler(), 272 | "node_positions": StandardScaler() 273 | } 274 | scaler_Y = StandardScaler() 275 | 276 | # Fit and transform training data 277 | roller_train_std = fit_transform_3d(roller_train, scalers_inputs["roller_x"]) 278 | force_x_train_std = fit_transform_3d(force_x_train, scalers_inputs["force_x"]) 279 | force_val_train_std = fit_transform_3d(force_val_train, scalers_inputs["force_values"]) 280 | node_train_std = fit_transform_3d(node_train, scalers_inputs["node_positions"]) 281 | 282 | # Transform validation data (without refit) 283 | roller_val_std = transform_3d(roller_val, scalers_inputs["roller_x"]) 284 | force_x_val_std = transform_3d(force_x_val, scalers_inputs["force_x"]) 285 | force_val_val_std = transform_3d(force_val_val, scalers_inputs["force_values"]) 286 | node_val_std = transform_3d(node_val, scalers_inputs["node_positions"]) 287 | 288 | # Prepare final input features by merging sub-features 289 | X_train_3d = merge_sub_features( 290 | roller_train_std, 291 | force_x_train_std, 292 | force_val_train_std, 293 | node_train_std 294 | ) 295 | X_val_3d = merge_sub_features( 296 | roller_val_std, 297 | force_x_val_std, 298 | force_val_val_std, 299 | node_val_std 300 | ) 301 | 302 | # (Optional) Pad feature dimensions 303 | X_train_3d_padded, feat_dim_padded = pad_feat_dim_to_multiple_of_nheads(X_train_3d, nheads=nheads_for_padding) 304 | X_val_3d_padded, _ = pad_feat_dim_to_multiple_of_nheads(X_val_3d, nheads=nheads_for_padding) 305 | 306 | # Unify the label by aggregating across cases 307 | Y_train_2d = unify_label_with_c(I_train, c=c) # Shape: (B, n_elem) 308 | Y_val_2d = unify_label_with_c(I_val, c=c) # Shape: (B, n_elem) 309 | 310 | # Fit scaler_Y on training targets 311 | scaler_Y.fit(Y_train_2d) 312 | 313 | # Transform targets 314 | Y_train_std = scaler_Y.transform(Y_train_2d) # (B, n_elem) 315 | Y_val_std = scaler_Y.transform(Y_val_2d) # (B, n_elem) 316 | 317 | # Convert to PyTorch Tensors 318 | X_train_tensor = torch.tensor(X_train_3d_padded, dtype=torch.float32) 319 | Y_train_tensor = torch.tensor(Y_train_std, dtype=torch.float32) 320 | X_val_tensor = torch.tensor(X_val_3d_padded, dtype=torch.float32) 321 | Y_val_tensor = torch.tensor(Y_val_std, dtype=torch.float32) 322 | 323 | min_constraint = torch.min(Y_train_tensor) 324 | max_constraint = torch.max(Y_train_tensor) 325 | 326 | # Create TensorDatasets 327 | train_dataset = TensorDataset(X_train_tensor, Y_train_tensor) 328 | val_dataset = TensorDataset(X_val_tensor, Y_val_tensor) 329 | 330 | print("X_train_3d shape:", X_train_3d.shape, 331 | "Y_train_2d shape:", Y_train_2d.shape, 332 | "\nAfter optional padding => X_train_3d_padded shape:", 333 | X_train_3d_padded.shape) 334 | 335 | 336 | ####################################### 337 | # 3) DEFINE THE FNO MODEL 338 | ####################################### 339 | 340 | class SpectralConv1d(nn.Module): 341 | """ 342 | 1D Fourier layer. Performs Fourier transform along the last dimension. 343 | """ 344 | def __init__(self, in_channels, out_channels, modes): 345 | super().__init__() 346 | self.in_channels = in_channels 347 | self.out_channels = out_channels 348 | self.modes = modes 349 | 350 | # Real & Imag parts of the Fourier coefficients 351 | self.scale = 1.0 / (in_channels * out_channels) 352 | self.weights_real = nn.Parameter( 353 | self.scale * torch.rand(in_channels, out_channels, self.modes) 354 | ) 355 | self.weights_imag = nn.Parameter( 356 | self.scale * torch.rand(in_channels, out_channels, self.modes) 357 | ) 358 | 359 | def forward(self, x): 360 | """ 361 | x: (B, in_channels, n) 362 | """ 363 | B, inC, n = x.shape 364 | x_ft = torch.fft.rfft(x, n=n) # => (B, inC, n//2+1) 365 | 366 | # Ensure that modes do not exceed the available frequency components 367 | actual_modes = min(self.modes, x_ft.shape[-1]) 368 | if actual_modes < self.modes: 369 | print(f"Adjusted modes from {self.modes} to {actual_modes} based on signal length {n}.") 370 | 371 | x_ft = x_ft[:, :, :actual_modes] 372 | 373 | # Adjust weights if actual_modes < self.modes 374 | if actual_modes < self.modes: 375 | w_r = self.weights_real[:, :, :actual_modes].unsqueeze(0) # => (1, inC, outC, actual_modes) 376 | w_i = self.weights_imag[:, :, :actual_modes].unsqueeze(0) 377 | else: 378 | w_r = self.weights_real.unsqueeze(0) # => (1, inC, outC, modes) 379 | w_i = self.weights_imag.unsqueeze(0) 380 | 381 | x_ft_real = x_ft.real 382 | x_ft_imag = x_ft.imag 383 | 384 | # Perform complex multiplication 385 | out_ft_real = torch.einsum("bim, iojm -> bojm", x_ft_real, w_r) - \ 386 | torch.einsum("bim, iojm -> bojm", x_ft_imag, w_i) 387 | out_ft_imag = torch.einsum("bim, iojm -> bojm", x_ft_real, w_i) + \ 388 | torch.einsum("bim, iojm -> bojm", x_ft_imag, w_r) 389 | 390 | # Sum over input channels 391 | out_ft_real = out_ft_real.sum(dim=2) # => (B, outC, modes) 392 | out_ft_imag = out_ft_imag.sum(dim=2) # => (B, outC, modes) 393 | 394 | out_ft = torch.complex(out_ft_real, out_ft_imag) 395 | 396 | # Pad back to original size in frequency domain 397 | pad_size = (0, (n//2 + 1) - actual_modes) 398 | if pad_size[1] > 0: 399 | out_ft = nn.functional.pad(out_ft, pad_size) 400 | 401 | # Inverse FFT to get back to spatial domain 402 | x_out = torch.fft.irfft(out_ft, n=n) # => (B, outC, n) 403 | return x_out 404 | 405 | class FNOBlock1d(nn.Module): 406 | """ 407 | One block of a 1D FNO: 408 | - SpectralConv1d 409 | - Pointwise Conv 410 | - BatchNorm + Activation 411 | """ 412 | def __init__(self, width, modes): 413 | super().__init__() 414 | self.conv = SpectralConv1d(width, width, modes) 415 | self.w = nn.Conv1d(width, width, 1) 416 | self.bn = nn.BatchNorm1d(width) 417 | 418 | def forward(self, x): 419 | """ 420 | x shape: (B, width, n) 421 | """ 422 | x1 = self.conv(x) 423 | x2 = self.w(x) 424 | x_out = x1 + x2 425 | x_out = self.bn(x_out) 426 | return nn.functional.gelu(x_out) 427 | 428 | class FNO1dModel(nn.Module): 429 | """ 430 | An FNO that operates on dimension = n_cases as the 1D dimension. 431 | Then flattens and passes an MLP to get final predictions of size (n_elem). 432 | """ 433 | def __init__( 434 | self, 435 | n_cases, 436 | feat_dim, 437 | n_elem, 438 | fno_modes, 439 | fno_width, 440 | num_fno_layers=4, 441 | hidden_units=512, 442 | dropout=0.1 443 | ): 444 | super().__init__() 445 | self.n_cases = n_cases 446 | self.feat_dim = feat_dim 447 | self.n_elem = n_elem 448 | self.modes = fno_modes 449 | self.width = fno_width 450 | self.num_fno_layers = num_fno_layers 451 | 452 | # Map input feat_dim -> fno_width 453 | self.fc0 = nn.Linear(feat_dim, fno_width) 454 | 455 | # Multiple FNO blocks 456 | self.fno_blocks = nn.ModuleList( 457 | [FNOBlock1d(fno_width, fno_modes) for _ in range(num_fno_layers)] 458 | ) 459 | 460 | # Output MLP: (fno_width*n_cases) -> hidden_units -> n_elem 461 | self.dropout = nn.Dropout(dropout) 462 | self.fc_out = nn.Sequential( 463 | nn.Linear(fno_width * n_cases, hidden_units), 464 | nn.LeakyReLU(0.1), 465 | nn.Dropout(dropout_rate), 466 | nn.Linear(hidden_units, n_elem), 467 | ) 468 | 469 | def forward(self, x): 470 | """ 471 | x shape: (B, n_cases, feat_dim) -> (B, n_elem) 472 | """ 473 | B, Nc, Fdim = x.shape 474 | assert Nc == self.n_cases and Fdim == self.feat_dim, \ 475 | f"Input shape {x.shape} does not match (B, {self.n_cases}, {self.feat_dim})." 476 | 477 | # Reshape to (B, feat_dim, n_cases) 478 | x = x.permute(0, 2, 1) # => (B, feat_dim, n_cases) 479 | 480 | # Map feat_dim -> fno_width 481 | # fc0 expects shape (B*n_cases, feat_dim), so we transpose last two dims 482 | x = self.fc0(x.transpose(-1, -2)) # => (B, n_cases, fno_width) 483 | x = x.transpose(-1, -2) # => (B, fno_width, n_cases) 484 | 485 | # FNO blocks 486 | for block in self.fno_blocks: 487 | x = block(x) 488 | 489 | # Flatten to (B, fno_width*n_cases) 490 | x = x.reshape(B, -1) 491 | x = self.dropout(x) 492 | 493 | # Final MLP 494 | out = self.fc_out(x) # => (B, n_elem) 495 | return out 496 | 497 | ####################################### 498 | # 4) DEFINE CUSTOM LOSS 499 | ####################################### 500 | 501 | class TrainableL1L2Loss(nn.Module): 502 | """ 503 | Combines L1 and L2 loss with a trainable alpha parameter and 504 | a penalty for predictions outside [min_constraint, max_constraint]. 505 | """ 506 | def __init__( 507 | self, 508 | initial_alpha=0.5, 509 | min_constraint=None, 510 | max_constraint=None, 511 | penalty_weight=0.5 512 | ): 513 | super().__init__() 514 | self.alpha = nn.Parameter(torch.tensor(initial_alpha, dtype=torch.float32, requires_grad=True)) 515 | self.l1 = nn.L1Loss() 516 | self.l2 = nn.MSELoss() 517 | self.min_constraint = min_constraint 518 | self.max_constraint = max_constraint 519 | self.penalty_weight = penalty_weight 520 | 521 | def forward(self, preds, targets): 522 | alpha_clamped = torch.clamp(self.alpha, 1e-6, 1.0) 523 | 524 | l1_loss = self.l1(preds, targets) 525 | l2_loss = self.l2(preds, targets) 526 | 527 | penalty = 0.0 528 | if self.min_constraint is not None: 529 | below_min_penalty = torch.sum(torch.relu(self.min_constraint - preds)) 530 | penalty += below_min_penalty 531 | if self.max_constraint is not None: 532 | above_max_penalty = torch.sum(torch.relu(preds - self.max_constraint)) 533 | penalty += above_max_penalty 534 | 535 | total_loss = alpha_clamped * l1_loss + (1 - alpha_clamped) * l2_loss 536 | total_loss += self.penalty_weight * penalty 537 | return total_loss 538 | 539 | def permute_data(X, Y): 540 | """ 541 | Permutes data indices for both X and Y consistently. 542 | """ 543 | assert X.size(0) == Y.size(0), "X and Y must have same batch size." 544 | perm = torch.randperm(X.size(0), device=X.device) 545 | return X[perm], Y[perm] 546 | 547 | ####################################### 548 | # 5) INITIALIZE & TRAIN 549 | ####################################### 550 | 551 | # Initialize the FNO model 552 | model = FNO1dModel( 553 | n_cases=n_cases, 554 | feat_dim=feat_dim_padded, # after optional padding 555 | n_elem=nelem, 556 | fno_modes=fno_modes, 557 | fno_width=fno_width, 558 | num_fno_layers=num_fno_layers, 559 | hidden_units=hidden_units, 560 | dropout=dropout_rate 561 | ).to(device) 562 | 563 | # Initialize optimizer, scheduler, and loss criterion 564 | optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay) 565 | scheduler = ExponentialLR(optimizer, gamma=gamma) 566 | criterion = TrainableL1L2Loss( 567 | min_constraint=min_constraint, 568 | max_constraint=max_constraint, 569 | penalty_weight=box_constraint_coeff 570 | ) 571 | 572 | # DataLoaders 573 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 574 | val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 575 | 576 | # Initialize GradScaler for AMP (Not using AMP due to FFT constraints) 577 | # If you decide to use AMP in future, ensure FFT dimensions are powers of two 578 | scaler_amp = GradScaler() 579 | 580 | # Live Plotting Setup 581 | plt.ion() 582 | fig, ax = plt.subplots(figsize=(10, 6)) 583 | train_losses = [] 584 | val_losses = [] 585 | best_val_loss = float('inf') 586 | epochs_no_improve = 0 587 | 588 | def live_plot(epoch, train_losses, val_losses): 589 | ax.clear() 590 | ax.plot(range(1, epoch + 1), train_losses, label="Train Loss", marker='o', color='blue') 591 | ax.plot(range(1, epoch + 1), val_losses, label="Validation Loss", marker='x', color='red') 592 | ax.set_xlabel("Epochs") 593 | ax.set_ylabel("Loss") 594 | ax.set_title("Training and Validation Loss") 595 | ax.legend() 596 | ax.grid(True, linestyle='--', alpha=0.7) 597 | plt.pause(0.01) 598 | 599 | # Training Loop 600 | for epoch in range(1, num_epochs + 1): 601 | model.train() 602 | noise_level = sigma_0 * (gamma_noise ** epoch) 603 | 604 | total_train_loss = 0.0 605 | t0 = time.time() 606 | 607 | for Xb, Yb in train_loader: 608 | Xb, Yb = Xb.to(device), Yb.to(device) 609 | 610 | # Optional permutation 611 | Xb, Yb = permute_data(Xb, Yb) 612 | 613 | # Add Gaussian noise 614 | Xb_noisy = Xb + torch.randn_like(Xb) * noise_level 615 | 616 | optimizer.zero_grad() 617 | # Disable AMP to avoid FFT precision issues 618 | with torch.cuda.amp.autocast(enabled=False): 619 | preds = model(Xb_noisy) 620 | L_alpha = (initial_alpha - criterion.alpha) ** 2 621 | loss = criterion(preds, Yb) + L_alpha 622 | 623 | loss.backward() 624 | torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) 625 | optimizer.step() 626 | 627 | total_train_loss += loss.item() 628 | 629 | avg_train_loss = total_train_loss / len(train_loader) 630 | train_losses.append(avg_train_loss) 631 | 632 | # Validation 633 | model.eval() 634 | total_val_loss = 0.0 635 | with torch.no_grad(): 636 | for Xb, Yb in val_loader: 637 | Xb, Yb = Xb.to(device), Yb.to(device) 638 | with torch.cuda.amp.autocast(enabled=False): 639 | preds = model(Xb) 640 | val_loss = criterion(preds, Yb) 641 | total_val_loss += val_loss.item() 642 | 643 | avg_val_loss = total_val_loss / len(val_loader) 644 | val_losses.append(avg_val_loss) 645 | scheduler.step() 646 | 647 | # Early Stopping 648 | if avg_val_loss < best_val_loss: 649 | best_val_loss = avg_val_loss 650 | epochs_no_improve = 0 651 | torch.save(model.state_dict(), "best_model_fno.pth") 652 | else: 653 | epochs_no_improve += 1 654 | if epochs_no_improve >= patience: 655 | print(f"Early stopping at epoch {epoch}") 656 | break 657 | 658 | dt = time.time() - t0 659 | print(f"Epoch {epoch}/{num_epochs} | " 660 | f"Train Loss={avg_train_loss:.6f}, " 661 | f"Val Loss={avg_val_loss:.6f}, " 662 | f"Time={dt:.2f}s") 663 | 664 | live_plot(epoch, train_losses, val_losses) 665 | 666 | ####################################### 667 | # 6) EVALUATION 668 | ####################################### 669 | 670 | # Load the best model 671 | model.load_state_dict(torch.load("best_model_fno.pth", map_location=device)) 672 | model.eval() 673 | 674 | # Evaluation on Validation Set 675 | val_loader_eval = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 676 | all_preds, all_labels = [], [] 677 | 678 | with torch.no_grad(): 679 | for Xb, Yb in val_loader_eval: 680 | Xb = Xb.to(device) 681 | preds = model(Xb) 682 | all_preds.append(preds.cpu()) 683 | all_labels.append(Yb) 684 | 685 | all_preds = torch.cat(all_preds, dim=0).numpy() 686 | all_labels = torch.cat(all_labels, dim=0).numpy() 687 | 688 | # Inverse transform to original scale 689 | all_preds_unstd = scaler_Y.inverse_transform(all_preds) 690 | all_labels_unstd = scaler_Y.inverse_transform(all_labels) 691 | 692 | # Example clipping (adjust based on domain knowledge) 693 | all_preds_unstd = np.clip(all_preds_unstd, 0.0, 1e10) 694 | all_labels_unstd = np.clip(all_labels_unstd, 0.0, 1e10) 695 | 696 | # Calculate R² score 697 | r2_val = r2_score(all_labels_unstd.ravel(), all_preds_unstd.ravel()) 698 | print(f"R² on Validation: {r2_val:.4f}") 699 | 700 | ####################################### 701 | # 7) EXAMPLE INFERENCE & PLOT 702 | ####################################### 703 | 704 | # Example user inputs 705 | L_beam = 200 706 | Fmin_user = -355857 707 | Fmax_user = Fmin_user / 10 708 | user_rollers = [2*9, 2*29, 2*69, 2*85, 2*100] 709 | 710 | def build_user_input_no_agg( 711 | roller_list, force_x_list, force_val_list, node_pos_list, 712 | scalers, n_cases, max_lengths 713 | ): 714 | """ 715 | Build user input without aggregation by scaling and padding. 716 | """ 717 | feat_3d = scale_user_inputs( 718 | roller_list, force_x_list, force_val_list, node_pos_list, 719 | scalers, n_cases, max_lengths 720 | ) 721 | return feat_3d 722 | 723 | # Create random multi-case loads 724 | user_roller = [user_rollers.copy() for _ in range(n_cases)] 725 | user_force_x = [] 726 | user_force_vals = [] 727 | for _ in range(n_cases): 728 | num_forces = random.randint(1, 3) 729 | fx = sorted([random.uniform(0, L_beam) for _ in range(num_forces)]) 730 | fv = [random.uniform(Fmin_user, Fmax_user) for _ in range(num_forces)] 731 | user_force_x.append(fx) 732 | user_force_vals.append(fv) 733 | 734 | user_node_pos = [ 735 | np.linspace(0, L_beam, nelem + 1).tolist() for _ in range(n_cases) 736 | ] 737 | 738 | # Build user input 739 | X_user_3d = build_user_input_no_agg( 740 | user_roller, user_force_x, user_force_vals, user_node_pos, 741 | scalers_inputs, n_cases, max_lengths 742 | ) 743 | 744 | # Optional padding 745 | X_user_3d_padded, _ = pad_feat_dim_to_multiple_of_nheads(X_user_3d, nheads=nheads_for_padding) 746 | X_user_t = torch.tensor(X_user_3d_padded, dtype=torch.float32).to(device) 747 | 748 | # Predict 749 | model.eval() 750 | with torch.no_grad(): 751 | pred_1x = model(X_user_t) # => (1, n_elem) 752 | 753 | pred_1x_np = pred_1x.cpu().numpy().squeeze() 754 | pred_1x_unstd = scaler_Y.inverse_transform(pred_1x_np.reshape(1, -1)).squeeze() 755 | 756 | # ------------------ PLOT ------------------ 757 | unique_rollers = sorted(set([x for sublist in user_roller for x in sublist] + [L_beam])) 758 | case_colors = sns.color_palette("Set1", n_colors=n_cases) 759 | case_labels = [f'Force Case {i+1} (N)' for i in range(n_cases)] 760 | 761 | beam_y = 0 762 | beam_x = [0, L_beam] 763 | beam_y_vals = [beam_y, beam_y] 764 | 765 | force_positions = [] 766 | force_vals_plot = [] 767 | for fx, fv in zip(user_force_x, user_force_vals): 768 | for xx, val in zip(fx, fv): 769 | force_positions.append(xx) 770 | force_vals_plot.append(val) 771 | 772 | max_force = max(abs(val) for val in force_vals_plot) if force_vals_plot else 1.0 773 | desired_max_arrow_length = 2.0 774 | arrow_scale = desired_max_arrow_length / max_force if max_force != 0 else 1.0 775 | 776 | beam_positions = user_node_pos[0][:nelem] 777 | 778 | I_normalized = (pred_1x_unstd - pred_1x_unstd.min()) / ( 779 | pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8 780 | ) 781 | cmap = cm.winter 782 | norm = plt.Normalize(pred_1x_unstd.min(), pred_1x_unstd.max()) 783 | 784 | block_width = L_beam / nelem * 0.8 785 | block_height = 1 786 | 787 | fig, ax = plt.subplots(figsize=(18, 7)) 788 | 789 | # Plot beam 790 | ax.plot(beam_x, beam_y_vals, color='black', linewidth=3, label='Beam') 791 | ax.scatter(beam_x[0], beam_y - 0.15, marker='^', color='red', s=300, zorder=6) 792 | 793 | # Plot rollers 794 | ax.scatter(unique_rollers, [beam_y]*len(unique_rollers), 795 | marker='o', color='seagreen', s=200, 796 | label='Rollers', zorder=5, edgecolors='k') 797 | 798 | # Plot forces 799 | for case_idx in range(n_cases): 800 | fx_list = user_force_x[case_idx] 801 | fv_list = user_force_vals[case_idx] 802 | color = case_colors[case_idx] 803 | label = case_labels[case_idx] 804 | 805 | for idx, (fx, fv) in enumerate(zip(fx_list, fv_list)): 806 | arrow_length = abs(fv) * arrow_scale 807 | start_point = (fx, beam_y + arrow_length) 808 | end_point = (fx, beam_y) 809 | 810 | arrow = FancyArrowPatch( 811 | posA=start_point, posB=end_point, 812 | arrowstyle='-|>', 813 | mutation_scale=20, 814 | color=color, 815 | linewidth=2, 816 | alpha=0.8, 817 | label=label if idx == 0 else "" 818 | ) 819 | ax.add_patch(arrow) 820 | ax.text(fx, beam_y + arrow_length + desired_max_arrow_length * 0.02, 821 | f"{fv:.0f}", ha='center', va='bottom', 822 | fontsize=10, color=color, fontweight='bold') 823 | 824 | # Plot predicted I as rectangles 825 | for idx, (x_pos, I_val) in enumerate(zip(beam_positions, pred_1x_unstd)): 826 | color = cmap(norm(I_val)) 827 | rect_x = x_pos - block_width / 2 828 | rect_y = beam_y - ( 829 | I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8) 830 | ) * block_height / 2 831 | 832 | rect = Rectangle( 833 | (rect_x, rect_y), 834 | block_width, 835 | (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8)) * block_height, 836 | linewidth=0, 837 | facecolor=color, 838 | alpha=0.6 839 | ) 840 | ax.add_patch(rect) 841 | 842 | # Colorbar 843 | sm = cm.ScalarMappable(cmap=cmap, norm=norm) 844 | sm.set_array([]) 845 | 846 | cbar = fig.colorbar(sm, ax=ax, orientation='vertical', fraction=0.046, pad=0.04) 847 | cbar.set_label('Predicted I (m$^4$)', fontsize=16) 848 | cbar.ax.tick_params(labelsize=10) 849 | 850 | ax.set_title("Beam Setup with Applied Forces (FNO-Predicted I)", 851 | fontsize=22, fontweight='bold', pad=20) 852 | ax.set_xlabel("Beam Length (m)", fontsize=16, fontweight='semibold') 853 | ax.set_xlim(-5, L_beam + 5) 854 | ax.set_ylim(-2.5, 2.5) 855 | ax.set_xticks(np.arange(0, L_beam + 5, 5)) 856 | ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7) 857 | 858 | # Create custom legend 859 | legend_elements = [ 860 | Line2D([0], [0], color='black', lw=3, label='Beam'), 861 | Line2D([0], [0], marker=(3, 0, -90), color='red', label='Pin', 862 | markerfacecolor='red', markersize=15), 863 | Line2D([0], [0], marker='o', color='seagreen', label='Rollers', 864 | markerfacecolor='seagreen', markeredgecolor='k', markersize=15), 865 | ] 866 | for color, label in zip(case_colors, case_labels): 867 | legend_elements.append(Line2D([0], [0], color=color, lw=2, label=label)) 868 | 869 | ax.legend(handles=legend_elements, loc='lower right', fontsize=12) 870 | 871 | plt.tight_layout() 872 | plt.show() 873 | -------------------------------------------------------------------------------- /OpenPyStruct_Bayesian_TFDModule_MultiCase_Beta.py: -------------------------------------------------------------------------------- 1 | ############################################################################ 2 | #### OpenPyStruct Transformer-Diffusion Based Multi-Load-Case + BNN #### 3 | #### Coder: Danny Smyl, PhD, PE, Georgia Tech, 2025 #### 4 | ############################################################################ 5 | 6 | #### torchbnn needed, example install: pip install torchbnn #### 7 | 8 | import os 9 | import json 10 | import time 11 | import math 12 | import random 13 | import numpy as np 14 | import matplotlib.pyplot as plt 15 | import matplotlib.cm as cm 16 | from matplotlib.lines import Line2D 17 | from matplotlib.patches import FancyArrowPatch, Rectangle 18 | import seaborn as sns 19 | 20 | import torch 21 | import torch.nn as nn 22 | import torch.optim as optim 23 | from torch.optim.lr_scheduler import ExponentialLR 24 | from torch.utils.data import DataLoader, TensorDataset 25 | from sklearn.preprocessing import StandardScaler 26 | from sklearn.metrics import r2_score 27 | from scipy.stats import mode 28 | 29 | # Use new AMP API to avoid deprecation warnings 30 | from torch.amp import autocast, GradScaler 31 | 32 | # -------------- NEW: torchbnn for Bayesian Layers -------------- 33 | import torchbnn as bnn 34 | 35 | ####################################### 36 | # 1) CONFIGURATION & HYPERPARAMETERS 37 | ####################################### 38 | 39 | # Model and training configuration # 40 | n_cases = 6 # Number of sub-cases per sample 41 | nelem = 100 # Final output dimension per sample: (B, n_elem) 42 | box_constraint_coeff = 5e-1 # Coefficient for box constraint penalty 43 | hidden_units = 512 # Number of hidden units in MLP 44 | dropout_rate = 0.1 # Dropout rate for regularization 45 | num_blocks = 2 # Number of blocks (unused in current model) 46 | num_epochs = 500 # Maximum number of training epochs 47 | batch_size = 512 # Batch size for training 48 | patience = 10 # Early stopping patience 49 | learning_rate = 3e-4 # Learning rate for optimizer 50 | weight_decay = 1e-6 # Weight decay (L2 regularization) for optimizer 51 | train_split = 0.8 # Fraction of data used for training 52 | sigma_0 = 0.01 # Initial Gaussian noise for input 53 | gamma_noise = 0.95 # Decay rate for noise during training 54 | gamma = 0.99 # Learning rate scheduler decay rate 55 | initial_alpha = 0.5 # Initial alpha value for loss weighting 56 | c = 0.5 # Parameter to adjust label aggregation 57 | bnn_kl_scale = 1e-6 # Scaling factor for KL-Divergence in Bayesian layers 58 | 59 | # Additional diffusion & Transformer hyperparameters # 60 | num_transformer_layers = 4 # Number of Transformer encoder layers 61 | dim_feedforward = 512 # Dimension of feedforward network in Transformer 62 | num_heads = 24 # Number of attention heads in Transformer 63 | max_len = 512 # Maximum sequence length for positional encoding 64 | diffusion_hidden_dim = 512 # Hidden dimension in diffusion MLP 65 | diffusion_T = 512 # Total number of diffusion steps 66 | 67 | # Device configuration 68 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 69 | print("Using device:", device) 70 | 71 | ####################################### 72 | # 2) DATA LOADING & PREPROCESSING 73 | ####################################### 74 | 75 | def pad_sequences(data_list, max_length, pad_val=0.0): 76 | """ 77 | Pad each 1D array in data_list to 'max_length'. 78 | Returns a NumPy array of shape (num_samples, max_length). 79 | """ 80 | out = np.full((len(data_list), max_length), pad_val, dtype=np.float32) 81 | for i, arr in enumerate(data_list): 82 | arr_np = np.array(arr, dtype=np.float32) 83 | length = min(len(arr_np), max_length) 84 | out[i, :length] = arr_np[:length] 85 | return out 86 | 87 | def unify_label_with_c(I_3d, c=c): 88 | """ 89 | Aggregate labels by computing the mean across cases and adding c times the standard deviation. 90 | 91 | Parameters: 92 | - I_3d: NumPy array of shape (B, n_cases, n_elem) 93 | - c: Scalar multiplier for standard deviation 94 | 95 | Returns: 96 | - Y: NumPy array of shape (B, n_elem) 97 | """ 98 | I_mean = I_3d.mean(axis=1) # Mean across cases 99 | I_std = I_3d.std(axis=1) # Standard deviation across cases 100 | return I_mean + c * I_std 101 | 102 | def fit_transform_3d(arr_3d, scaler): 103 | """ 104 | Fit and transform a 3D array using the provided scaler over axis=0 105 | on the flattened shape (B*NC, M). 106 | """ 107 | B, NC, M = arr_3d.shape 108 | flat = arr_3d.reshape(B * NC, M) # Combine B and NC for fitting 109 | scaled = scaler.fit_transform(flat) 110 | return scaled.reshape(B, NC, M) 111 | 112 | def transform_3d(arr_3d, scaler): 113 | """ 114 | Transform (but do not fit) a 3D array using the provided scaler, 115 | maintaining (B, NC, M) shape. 116 | """ 117 | B, NC, M = arr_3d.shape 118 | flat = arr_3d.reshape(B * NC, M) 119 | scaled = scaler.transform(flat) 120 | return scaled.reshape(B, NC, M) 121 | 122 | def merge_sub_features(*arrays): 123 | """ 124 | Concatenate multiple feature arrays along the feature dimension. 125 | 126 | Parameters: 127 | - arrays: Variable number of NumPy arrays to concatenate 128 | 129 | Returns: 130 | - merged_array: Concatenated NumPy array 131 | """ 132 | return np.concatenate(arrays, axis=2) 133 | 134 | def pad_feat_dim_to_multiple_of_nheads(X_3d, nheads): 135 | """ 136 | Pad the feature dimension to be a multiple of nheads. 137 | 138 | Parameters: 139 | - X_3d: NumPy array of shape (B, Nc, original_dim) 140 | - nheads: Integer, number of attention heads 141 | 142 | Returns: 143 | - X_3d_padded: Padded NumPy array 144 | - new_dim: New feature dimension after padding 145 | """ 146 | B, Nc, original_dim = X_3d.shape 147 | remainder = original_dim % nheads 148 | if remainder == 0: 149 | return X_3d, original_dim 150 | new_dim = ((original_dim // nheads) + 1) * nheads 151 | diff = new_dim - original_dim 152 | X_3d_padded = np.pad(X_3d, ((0,0), (0,0), (0,diff)), mode='constant') 153 | return X_3d_padded, new_dim 154 | 155 | def scale_user_inputs( 156 | user_roller, user_force_x, user_force_vals, user_node_pos, 157 | scalers, n_cases, max_lengths 158 | ): 159 | """ 160 | Scales the user inputs using the fitted scalers. 161 | 162 | Parameters: 163 | - user_roller: List of roller locations per case 164 | - user_force_x: List of force positions per case 165 | - user_force_vals: List of force values per case 166 | - user_node_pos: List of node positions per case 167 | - scalers: Dictionary of fitted scalers for each feature 168 | - n_cases: Number of cases 169 | - max_lengths: Dictionary of maximum lengths for padding 170 | 171 | Returns: 172 | - feat_3d: NumPy array of shape (1, n_cases, feat_dim) 173 | """ 174 | def pad_to_length(seq, req_len): 175 | arr = np.zeros((req_len,), dtype=np.float32) 176 | length = min(len(seq), req_len) 177 | arr[:length] = seq[:length] 178 | return arr 179 | 180 | feat_arrays = [] 181 | for i in range(n_cases): 182 | # Pad each input type to its respective max length 183 | r_pad = pad_to_length(user_roller[i], max_lengths['roller_x']) 184 | fx_pad = pad_to_length(user_force_x[i], max_lengths['force_x']) 185 | fv_pad = pad_to_length(user_force_vals[i], max_lengths['force_values']) 186 | nd_pad = pad_to_length(user_node_pos[i], max_lengths['node_positions']) 187 | 188 | # Scale using the fitted scalers 189 | r_scaled = scalers["roller_x"].transform(r_pad.reshape(1, -1)).flatten() 190 | fx_scaled = scalers["force_x"].transform(fx_pad.reshape(1, -1)).flatten() 191 | fv_scaled = scalers["force_values"].transform(fv_pad.reshape(1, -1)).flatten() 192 | nd_scaled = scalers["node_positions"].transform(nd_pad.reshape(1, -1)).flatten() 193 | 194 | # Concatenate scaled features 195 | sub_feat = np.concatenate([r_scaled, fx_scaled, fv_scaled, nd_scaled]) 196 | feat_arrays.append(sub_feat) 197 | 198 | feat_2d = np.stack(feat_arrays, axis=0) # (n_cases, total_feat_dim) 199 | feat_3d = feat_2d[np.newaxis, ...] # (1, n_cases, total_feat_dim) 200 | return feat_3d 201 | 202 | # Load data 203 | try: 204 | with open("StructDataMedium.json", "r") as f: 205 | data = json.load(f) 206 | except FileNotFoundError: 207 | raise FileNotFoundError("The file 'StructDataLite.json' was not found.") 208 | 209 | # Extract data 210 | roller_x = data.get("roller_x_locations", []) 211 | force_x = data.get("force_x_locations", []) 212 | force_values = data.get("force_values", []) 213 | node_positions = data.get("node_positions", []) 214 | I_values = data.get("I_values", []) 215 | 216 | num_samples = len(I_values) 217 | req_keys = ["roller_x_locations","force_x_locations","force_values","node_positions"] 218 | if not all(len(data.get(k, [])) == num_samples for k in req_keys): 219 | raise ValueError("Mismatch in sample counts among roller_x, force_x, force_values, node_positions.") 220 | 221 | # Determine maximum lengths for padding 222 | max_lengths = { 223 | "roller_x": max(len(r) for r in roller_x) if roller_x else 0, 224 | "force_x": max(len(r) for r in force_x) if force_x else 0, 225 | "force_values": max(len(r) for r in force_values) if force_values else 0, 226 | "node_positions": max(len(r) for r in node_positions) if node_positions else 0, 227 | "I_values": max(len(r) for r in I_values) if I_values else 0 228 | } 229 | 230 | # Pad sequences 231 | roller_x_pad = pad_sequences(roller_x, max_lengths["roller_x"]) 232 | force_x_pad = pad_sequences(force_x, max_lengths["force_x"]) 233 | force_val_pad = pad_sequences(force_values, max_lengths["force_values"]) 234 | node_pos_pad = pad_sequences(node_positions, max_lengths["node_positions"]) 235 | I_values_pad = pad_sequences(I_values, max_lengths["I_values"]) 236 | 237 | # Group data by n_cases 238 | total_grouped = num_samples // n_cases 239 | if total_grouped == 0: 240 | raise ValueError(f"n_cases={n_cases} > total samples={num_samples}.") 241 | 242 | trim_len = total_grouped * n_cases 243 | roller_x_pad = roller_x_pad[:trim_len] 244 | force_x_pad = force_x_pad[:trim_len] 245 | force_val_pad = force_val_pad[:trim_len] 246 | node_pos_pad = node_pos_pad[:trim_len] 247 | I_values_pad = I_values_pad[:trim_len] 248 | 249 | roller_grouped = roller_x_pad.reshape(total_grouped, n_cases, -1) 250 | force_x_grouped = force_x_pad.reshape(total_grouped, n_cases, -1) 251 | force_val_grouped = force_val_pad.reshape(total_grouped, n_cases, -1) 252 | node_grouped = node_pos_pad.reshape(total_grouped, n_cases, -1) 253 | I_grouped = I_values_pad.reshape(total_grouped, n_cases, -1) 254 | 255 | # Train/Validation Split 256 | indices = np.random.permutation(total_grouped) 257 | train_sz = np.int32(train_split * total_grouped) 258 | train_idx = indices[:train_sz] 259 | val_idx = indices[train_sz:] 260 | 261 | roller_train = roller_grouped[train_idx] 262 | roller_val = roller_grouped[val_idx] 263 | force_x_train = force_x_grouped[train_idx] 264 | force_x_val = force_x_grouped[val_idx] 265 | force_val_train = force_val_grouped[train_idx] 266 | force_val_val = force_val_grouped[val_idx] 267 | node_train = node_grouped[train_idx] 268 | node_val = node_grouped[val_idx] 269 | I_train = I_grouped[train_idx] 270 | I_val = I_grouped[val_idx] 271 | 272 | # Initialize Scalers 273 | scalers_inputs = { 274 | "roller_x": StandardScaler(), 275 | "force_x": StandardScaler(), 276 | "force_values": StandardScaler(), 277 | "node_positions": StandardScaler() 278 | } 279 | scaler_Y = StandardScaler() 280 | 281 | # Fit/transform training data 282 | roller_train_std = fit_transform_3d(roller_train, scalers_inputs["roller_x"]) 283 | force_x_train_std = fit_transform_3d(force_x_train, scalers_inputs["force_x"]) 284 | force_val_train_std = fit_transform_3d(force_val_train, scalers_inputs["force_values"]) 285 | node_train_std = fit_transform_3d(node_train, scalers_inputs["node_positions"]) 286 | 287 | # Transform validation data (fit is done on training) 288 | roller_val_std = transform_3d(roller_val, scalers_inputs["roller_x"]) 289 | force_x_val_std = transform_3d(force_x_val, scalers_inputs["force_x"]) 290 | force_val_val_std = transform_3d(force_val_val, scalers_inputs["force_values"]) 291 | node_val_std = transform_3d(node_val, scalers_inputs["node_positions"]) 292 | 293 | # Prepare final input features by merging sub-features 294 | X_train_3d = merge_sub_features( 295 | roller_train_std, 296 | force_x_train_std, 297 | force_val_train_std, 298 | node_train_std 299 | ) 300 | X_val_3d = merge_sub_features( 301 | roller_val_std, 302 | force_x_val_std, 303 | force_val_val_std, 304 | node_val_std 305 | ) 306 | 307 | # Pad feature dimensions to be multiples of num_heads 308 | X_train_3d_padded, feat_dim_padded = pad_feat_dim_to_multiple_of_nheads(X_train_3d, num_heads) 309 | X_val_3d_padded, _ = pad_feat_dim_to_multiple_of_nheads(X_val_3d, num_heads) 310 | 311 | # Unify the label by aggregating across cases 312 | Y_train_2d = unify_label_with_c(I_train, c=c) # (B, n_elem) 313 | Y_val_2d = unify_label_with_c(I_val, c=c) # (B, n_elem) 314 | 315 | # Fit scaler_Y on training targets 316 | scaler_Y.fit(Y_train_2d) 317 | 318 | # Transform targets 319 | Y_train_std = scaler_Y.transform(Y_train_2d) # (B, n_elem) 320 | Y_val_std = scaler_Y.transform(Y_val_2d) # (B, n_elem) 321 | 322 | # Convert to PyTorch Tensors 323 | X_train_tensor = torch.tensor(X_train_3d_padded, dtype=torch.float32) 324 | Y_train_tensor = torch.tensor(Y_train_std, dtype=torch.float32) 325 | X_val_tensor = torch.tensor(X_val_3d_padded, dtype=torch.float32) 326 | Y_val_tensor = torch.tensor(Y_val_std, dtype=torch.float32) 327 | 328 | min_constraint = torch.min(Y_train_tensor) 329 | max_constraint = torch.max(Y_train_tensor) 330 | 331 | # Create TensorDatasets 332 | train_dataset = TensorDataset(X_train_tensor, Y_train_tensor) 333 | val_dataset = TensorDataset(X_val_tensor, Y_val_tensor) 334 | 335 | print("X_train_3d shape:", X_train_3d.shape, 336 | "Y_train_2d shape:", Y_train_2d.shape, 337 | "\nAfter padding, X_train_3d_padded shape:", 338 | X_train_3d_padded.shape) 339 | 340 | 341 | ####################################### 342 | # 3) DEFINE THE MODEL COMPONENTS 343 | ####################################### 344 | 345 | class PositionalEncoding(nn.Module): 346 | """ 347 | Sine/cosine positional encoding for even or odd d_model. 348 | Adds positional information to the input embeddings. 349 | """ 350 | def __init__(self, d_model, max_len=max_len): 351 | super().__init__() 352 | pe = torch.zeros(max_len, d_model) 353 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 354 | 355 | # Number of sin/cos pairs 356 | n_pairs = d_model // 2 357 | div_term = torch.exp(-math.log(10000.0) * torch.arange(0, n_pairs, dtype=torch.float) / d_model) 358 | 359 | # Fill in sin for even indices 360 | pe[:, 0 : 2*n_pairs : 2] = torch.sin(position * div_term) 361 | # Fill in cos for odd indices 362 | pe[:, 1 : 2*n_pairs : 2] = torch.cos(position * div_term) 363 | # If d_model is odd, the last column remains zeros 364 | 365 | pe = pe.unsqueeze(0) # (1, max_len, d_model) 366 | self.register_buffer('pe', pe) 367 | 368 | def forward(self, x): 369 | """ 370 | Adds positional encoding to input tensor. 371 | x shape: (B, seq_len, d_model) 372 | """ 373 | seq_len = x.size(1) 374 | return x + self.pe[:, :seq_len, :] 375 | 376 | 377 | ####################################### 378 | # 4) DIFFUSION & BAYESIAN MODULES 379 | ####################################### 380 | 381 | class DiffusionSchedule: 382 | """ 383 | Defines the diffusion schedule for adding noise. 384 | """ 385 | def __init__(self, T, beta_start=1e-12, beta_end=1e-5): 386 | self.T = T 387 | self.beta = torch.linspace(beta_start, beta_end, T) 388 | self.alpha = 1.0 - self.beta 389 | self.alpha_cumprod = torch.cumprod(self.alpha, dim=0) 390 | 391 | 392 | class BayesianDiffusionMLP(nn.Module): 393 | """ 394 | A small MLP for noise prediction & denoising in the diffusion process, 395 | using BayesianLinear layers. 396 | """ 397 | def __init__(self, in_features, hidden_features, prior_mu=0.0, prior_sigma=0.01): 398 | super().__init__() 399 | self.lin1 = bnn.BayesLinear( 400 | in_features=in_features, 401 | out_features=hidden_features, 402 | prior_mu=prior_mu, 403 | prior_sigma=prior_sigma 404 | ) 405 | self.lin2 = bnn.BayesLinear( 406 | in_features=hidden_features, 407 | out_features=in_features, 408 | prior_mu=prior_mu, 409 | prior_sigma=prior_sigma 410 | ) 411 | self.relu = nn.LeakyReLU(0.1) 412 | self.dropout = nn.Dropout(dropout_rate) 413 | self.norm = nn.LayerNorm(hidden_features) 414 | 415 | def forward(self, x): 416 | x = self.lin1(x) 417 | x = self.norm(x) 418 | x = self.relu(x) 419 | x = self.dropout(x) 420 | x = self.lin2(x) 421 | return x 422 | 423 | 424 | class DiffusionModule(nn.Module): 425 | """ 426 | Diffusion module that adds and removes noise using a learned Bayesian MLP. 427 | """ 428 | def __init__(self, feat_dim, hidden_dim=diffusion_hidden_dim, T=diffusion_T): 429 | super().__init__() 430 | self.T = T 431 | self.schedule = DiffusionSchedule(T) 432 | self.mlp = BayesianDiffusionMLP( 433 | in_features=feat_dim, 434 | hidden_features=hidden_dim 435 | ) 436 | 437 | def forward(self, x): 438 | """ 439 | Applies diffusion process to input tensor x. 440 | x shape: (B, n_cases, feat_dim) 441 | """ 442 | B, Nc, Fdim = x.shape 443 | device = x.device 444 | 445 | # For each sub-case, pick a random time step 446 | t = torch.randint(0, self.T, (B, Nc), device=device).long() # (B, n_cases) 447 | alpha_cumprod = self.schedule.alpha_cumprod.to(device) 448 | 449 | sqrt_alpha_cumprod = torch.sqrt(alpha_cumprod[t]) 450 | sqrt_one_minus_alpha_cumprod = torch.sqrt(1 - alpha_cumprod[t]) 451 | 452 | # Expand for broadcasting 453 | sqrt_alpha_cumprod = sqrt_alpha_cumprod.unsqueeze(-1) # (B, n_cases, 1) 454 | sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod.unsqueeze(-1) 455 | 456 | # Sample noise 457 | eps = torch.randn_like(x) 458 | 459 | # Noisy input 460 | x_noisy = sqrt_alpha_cumprod * x + sqrt_one_minus_alpha_cumprod * eps 461 | 462 | # Predict noise via Bayesian MLP 463 | eps_pred = self.mlp(x_noisy) 464 | 465 | # Denoise 466 | x_denoised = (x_noisy - sqrt_one_minus_alpha_cumprod * eps_pred) / sqrt_alpha_cumprod 467 | return x_denoised 468 | 469 | 470 | ####################################### 471 | # 5) BAYESIAN TRANSFORMER MODEL 472 | ####################################### 473 | class BayesianOutputMLP(nn.Module): 474 | """ 475 | Simple 2-layer Bayesian MLP: [feat_dim -> hidden_units -> n_elem]. 476 | """ 477 | def __init__(self, in_features, hidden_features, out_features, prior_mu=0.0, prior_sigma=0.01): 478 | super().__init__() 479 | self.lin1 = bnn.BayesLinear( 480 | in_features=in_features, 481 | out_features=hidden_features, 482 | prior_mu=prior_mu, 483 | prior_sigma=prior_sigma 484 | ) 485 | self.lin2 = bnn.BayesLinear( 486 | in_features=hidden_features, 487 | out_features=out_features, 488 | prior_mu=prior_mu, 489 | prior_sigma=prior_sigma 490 | ) 491 | self.norm = nn.LayerNorm(hidden_features) 492 | self.relu = nn.LeakyReLU(0.1) 493 | self.dropout = nn.Dropout(dropout_rate) 494 | 495 | def forward(self, x): 496 | x = self.lin1(x) 497 | x = self.norm(x) 498 | x = self.relu(x) 499 | x = self.dropout(x) 500 | x = self.lin2(x) 501 | return x 502 | 503 | class ModelOnePassTransformerWithDiffusion(nn.Module): 504 | def __init__( 505 | self, 506 | n_cases, 507 | feat_dim, 508 | n_elem, 509 | hidden_units, 510 | num_transformer_layers, 511 | num_heads, 512 | dim_feedforward, 513 | dropout, 514 | max_len, 515 | diffusion_hidden_dim, 516 | diffusion_T 517 | ): 518 | super().__init__() 519 | self.n_cases = n_cases 520 | self.feat_dim = feat_dim 521 | self.diffusion = DiffusionModule( 522 | feat_dim=feat_dim, 523 | hidden_dim=diffusion_hidden_dim, 524 | T=diffusion_T 525 | ) 526 | 527 | # Transformer encoder 528 | self.cls_token = nn.Parameter(torch.zeros(1, 1, feat_dim)) 529 | self.pos_encoder = PositionalEncoding(feat_dim, max_len=max_len) 530 | encoder_layer = nn.TransformerEncoderLayer( 531 | d_model=feat_dim, 532 | nhead=num_heads, 533 | dim_feedforward=dim_feedforward, 534 | dropout=dropout, 535 | batch_first=True 536 | ) 537 | self.transformer_encoder = nn.TransformerEncoder( 538 | encoder_layer, 539 | num_layers=num_transformer_layers 540 | ) 541 | 542 | # Bayesian MLP for final output 543 | self.bnn_output = BayesianOutputMLP( 544 | in_features=feat_dim, 545 | hidden_features=hidden_units, 546 | out_features=n_elem, 547 | prior_mu=0.0, 548 | prior_sigma=0.01 549 | ) 550 | 551 | def forward(self, x): 552 | """ 553 | Forward pass of the model. 554 | 555 | x shape: (B, n_cases, feat_dim) 556 | """ 557 | B, Nc, Fdim = x.shape 558 | msg = (f"Input dims {tuple(x.shape)} do not match " 559 | f"(B, {self.n_cases}, {self.feat_dim}).") 560 | assert Nc == self.n_cases and Fdim == self.feat_dim, msg 561 | 562 | # Diffusion: add & remove noise with a Bayesian MLP 563 | x = self.diffusion(x) # (B, n_cases, feat_dim) 564 | 565 | # Insert [CLS] token 566 | cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, feat_dim) 567 | x = torch.cat((cls_tokens, x), dim=1) # (B, 1 + n_cases, feat_dim) 568 | 569 | # Positional encoding 570 | x = self.pos_encoder(x) # (B, seq_len, feat_dim) 571 | 572 | # Transformer Encoder 573 | x = self.transformer_encoder(x) # (B, 1 + n_cases, feat_dim) 574 | 575 | # Extract [CLS] token representation 576 | cls_rep = x[:, 0, :] # (B, feat_dim) 577 | 578 | # Bayesian MLP for final output 579 | out = self.bnn_output(cls_rep) # (B, n_elem) 580 | return out 581 | 582 | 583 | ####################################### 584 | # 6) DEFINE CUSTOM LOSS 585 | ####################################### 586 | 587 | class TrainableL1L2Loss(nn.Module): 588 | """ 589 | Combines L1 and L2 loss with a trainable alpha parameter and 590 | penalizes predictions outside [min_constraint, max_constraint]. 591 | """ 592 | def __init__( 593 | self, 594 | initial_alpha=initial_alpha, 595 | min_constraint=min_constraint, 596 | max_constraint=max_constraint, 597 | penalty_weight=box_constraint_coeff 598 | ): 599 | super().__init__() 600 | self.alpha = nn.Parameter(torch.tensor(initial_alpha, dtype=torch.float32, requires_grad=True)) 601 | self.l1 = nn.L1Loss() 602 | self.l2 = nn.MSELoss() 603 | self.min_constraint = min_constraint 604 | self.max_constraint = max_constraint 605 | self.penalty_weight = penalty_weight 606 | 607 | def forward(self, preds, targets): 608 | # Clamp alpha 609 | alpha_clamped = torch.clamp(self.alpha, 1e-6, 1.0) 610 | 611 | # Compute L1 and L2 612 | l1_loss = self.l1(preds, targets) 613 | l2_loss = self.l2(preds, targets) 614 | 615 | # Box constraint penalty 616 | penalty = 0.0 617 | if self.min_constraint is not None: 618 | below_min_penalty = torch.sum(torch.relu(self.min_constraint - preds)) 619 | penalty += below_min_penalty 620 | if self.max_constraint is not None: 621 | above_max_penalty = torch.sum(torch.relu(preds - self.max_constraint)) 622 | penalty += above_max_penalty 623 | 624 | total_loss = alpha_clamped * l1_loss + (1 - alpha_clamped) * l2_loss 625 | total_loss = total_loss + self.penalty_weight * penalty 626 | return total_loss 627 | 628 | def permute_data(X, Y): 629 | """ 630 | Permutes the data indices for both X and Y consistently. 631 | """ 632 | assert X.size(0) == Y.size(0), "X and Y must have the same number of samples." 633 | perm = torch.randperm(X.size(0), device=X.device) 634 | return X[perm], Y[perm] 635 | 636 | 637 | ####################################### 638 | # 7) INITIALIZE & TRAIN 639 | ####################################### 640 | 641 | model = ModelOnePassTransformerWithDiffusion( 642 | n_cases=n_cases, 643 | feat_dim=feat_dim_padded, # Padded feature dimension (multiple of num_heads) 644 | n_elem=nelem, 645 | hidden_units=hidden_units, 646 | num_transformer_layers=num_transformer_layers, 647 | num_heads=num_heads, 648 | dim_feedforward=dim_feedforward, 649 | dropout=dropout_rate, 650 | max_len=max_len, 651 | diffusion_hidden_dim=diffusion_hidden_dim, 652 | diffusion_T=diffusion_T 653 | ).to(device) 654 | 655 | optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay) 656 | scheduler = ExponentialLR(optimizer, gamma=gamma) 657 | criterion = TrainableL1L2Loss() 658 | 659 | # DataLoaders 660 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 661 | val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 662 | 663 | # AMP scaler 664 | scaler_amp = GradScaler() 665 | 666 | # For live plotting 667 | plt.ion() 668 | fig, ax = plt.subplots(figsize=(10, 6)) 669 | 670 | def live_plot(epoch, train_losses, val_losses): 671 | ax.clear() 672 | ax.plot(range(1, epoch + 1), train_losses, label="Train Loss", marker='o', color='blue') 673 | ax.plot(range(1, epoch + 1), val_losses, label="Validation Loss", marker='x', color='red') 674 | ax.set_xlabel("Epochs") 675 | ax.set_ylabel("Loss") 676 | ax.set_title("Training and Validation Loss") 677 | ax.legend() 678 | ax.grid(True, linestyle='--', alpha=0.7) 679 | plt.pause(0.01) 680 | 681 | train_losses = [] 682 | val_losses = [] 683 | best_val_loss = float('inf') 684 | epochs_no_improve = 0 685 | 686 | for epoch in range(1, num_epochs + 1): 687 | model.train() 688 | noise_level = sigma_0 * (gamma_noise ** epoch) # Decaying noise 689 | 690 | total_train_loss = 0.0 691 | t0 = time.time() 692 | 693 | for Xb, Yb in train_loader: 694 | Xb, Yb = Xb.to(device), Yb.to(device) 695 | # Permute 696 | Xb, Yb = permute_data(Xb, Yb) 697 | # Add optional Gaussian noise 698 | Xb_noisy = Xb + torch.randn_like(Xb) * noise_level 699 | 700 | optimizer.zero_grad() 701 | with autocast(device_type='cuda', enabled=(device.type == 'cuda')): 702 | preds = model(Xb_noisy) 703 | # Mild penalty on alpha deviation 704 | L_alpha = (initial_alpha - criterion.alpha)**2 705 | 706 | # -------------- Compute KL loss for all Bayesian layers -------------- 707 | kl_loss = sum(m.kl_loss() for m in model.modules() if hasattr(m, 'kl_loss')) 708 | 709 | loss = criterion(preds, Yb) + L_alpha + bnn_kl_scale * kl_loss 710 | 711 | scaler_amp.scale(loss).backward() 712 | scaler_amp.unscale_(optimizer) 713 | torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) 714 | scaler_amp.step(optimizer) 715 | scaler_amp.update() 716 | 717 | total_train_loss += loss.item() 718 | 719 | avg_train_loss = total_train_loss / len(train_loader) 720 | train_losses.append(avg_train_loss) 721 | 722 | # Validation 723 | model.eval() 724 | total_val_loss = 0.0 725 | with torch.no_grad(), autocast(device_type='cuda', enabled=(device.type == 'cuda')): 726 | for Xb, Yb in val_loader: 727 | Xb, Yb = Xb.to(device), Yb.to(device) 728 | preds = model(Xb) 729 | kl_loss = sum(m.kl_loss() for m in model.modules() if hasattr(m, 'kl_loss')) 730 | val_loss = criterion(preds, Yb) + bnn_kl_scale * kl_loss 731 | total_val_loss += val_loss.item() 732 | 733 | avg_val_loss = total_val_loss / len(val_loader) 734 | val_losses.append(avg_val_loss) 735 | scheduler.step() 736 | 737 | # Early Stopping 738 | if avg_val_loss < best_val_loss: 739 | best_val_loss = avg_val_loss 740 | epochs_no_improve = 0 741 | torch.save(model.state_dict(), "best_model_onepass_bnn.pth") 742 | else: 743 | epochs_no_improve += 1 744 | if epochs_no_improve >= patience: 745 | print(f"Early stopping at epoch {epoch}") 746 | break 747 | 748 | dt = time.time() - t0 749 | print(f"Epoch {epoch}/{num_epochs} | " 750 | f"Train Loss={avg_train_loss:.6f}, " 751 | f"Val Loss={avg_val_loss:.6f}, " 752 | f"Time={dt:.2f}s") 753 | 754 | live_plot(epoch, train_losses, val_losses) 755 | 756 | 757 | ####################################### 758 | # 8) EVALUATION 759 | ####################################### 760 | 761 | model.load_state_dict(torch.load("best_model_onepass_bnn.pth", map_location=device)) 762 | model.eval() 763 | 764 | val_loader_eval = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 765 | all_preds, all_labels = [], [] 766 | 767 | with torch.no_grad(): 768 | for Xb, Yb in val_loader_eval: 769 | Xb = Xb.to(device) 770 | preds = model(Xb) 771 | all_preds.append(preds.cpu()) 772 | all_labels.append(Yb) 773 | 774 | all_preds = torch.cat(all_preds, dim=0).numpy() 775 | all_labels = torch.cat(all_labels, dim=0).numpy() 776 | 777 | # Un-standardize 778 | all_preds_unstd = scaler_Y.inverse_transform(all_preds) 779 | all_labels_unstd = scaler_Y.inverse_transform(all_labels) 780 | 781 | # Clip if desired (example: clip to [0, 1e10]) 782 | all_preds_unstd = np.clip(all_preds_unstd, 0.0, 1e10) 783 | all_labels_unstd = np.clip(all_labels_unstd, 0.0, 1e10) 784 | 785 | r2_val = r2_score(all_labels_unstd.ravel(), all_preds_unstd.ravel()) 786 | print(f"R² on Validation: {r2_val:.4f}") 787 | 788 | 789 | ####################################### 790 | # 9) EXAMPLE INFERENCE & PLOT 791 | ####################################### 792 | 793 | # Example user inputs 794 | L_beam = 200 795 | Fmin_user = -355857 796 | Fmax_user = Fmin_user / 10 797 | user_rollers = [2*9, 2*29, 2*69, 2*85, 2*100] 798 | 799 | def build_user_input_no_agg( 800 | roller_list, force_x_list, force_val_list, node_pos_list, 801 | scalers, n_cases, 802 | max_lengths 803 | ): 804 | feat_3d = scale_user_inputs( 805 | roller_list, force_x_list, force_val_list, node_pos_list, 806 | scalers, n_cases, max_lengths 807 | ) 808 | return feat_3d 809 | 810 | # Construct example multi-case loads 811 | user_roller = [user_rollers.copy() for _ in range(n_cases)] 812 | user_force_x = [] 813 | user_force_vals = [] 814 | for _ in range(n_cases): 815 | num_forces = random.randint(1, 3) 816 | fx = sorted([random.uniform(0, L_beam) for _ in range(num_forces)]) 817 | fv = [random.uniform(Fmin_user, Fmax_user) for _ in range(num_forces)] 818 | user_force_x.append(fx) 819 | user_force_vals.append(fv) 820 | 821 | user_node_pos = [np.linspace(0, L_beam, nelem + 1).tolist() for _ in range(n_cases)] 822 | 823 | X_user_3d = build_user_input_no_agg( 824 | user_roller, user_force_x, user_force_vals, user_node_pos, 825 | scalers_inputs, n_cases, max_lengths 826 | ) 827 | X_user_3d_padded, _ = pad_feat_dim_to_multiple_of_nheads(X_user_3d, num_heads) 828 | X_user_t = torch.tensor(X_user_3d_padded, dtype=torch.float32).to(device) 829 | 830 | model.eval() 831 | with torch.no_grad(): 832 | pred_1x = model(X_user_t) # (1, n_elem) 833 | 834 | pred_1x_np = pred_1x.cpu().numpy().squeeze() 835 | pred_1x_unstd = scaler_Y.inverse_transform(pred_1x_np.reshape(1, -1)).squeeze() 836 | 837 | # ---- PLOT RESULTS ---- 838 | unique_rollers = sorted(set([x for sublist in user_roller for x in sublist] + [L_beam])) 839 | case_colors = sns.color_palette("Set1", n_colors=n_cases) 840 | case_labels = [f'Force Case {i+1} (N)' for i in range(n_cases)] 841 | 842 | beam_y = 0 843 | beam_x = [0, L_beam] 844 | beam_y_vals = [beam_y, beam_y] 845 | 846 | # Collect forces 847 | force_positions = [] 848 | force_vals_plot = [] 849 | for fx, fv in zip(user_force_x, user_force_vals): 850 | for xx, val in zip(fx, fv): 851 | force_positions.append(xx) 852 | force_vals_plot.append(val) 853 | 854 | max_force = max(abs(val) for val in force_vals_plot) if force_vals_plot else 1.0 855 | desired_max_arrow_length = 2.0 856 | arrow_scale = desired_max_arrow_length / max_force if max_force != 0 else 1.0 857 | 858 | beam_positions = user_node_pos[0][:nelem] 859 | 860 | # Normalize predicted I for color 861 | I_normalized = (pred_1x_unstd - pred_1x_unstd.min()) / ( 862 | pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8 863 | ) 864 | cmap = cm.winter 865 | norm = plt.Normalize(pred_1x_unstd.min(), pred_1x_unstd.max()) 866 | 867 | block_width = L_beam / nelem * 0.8 868 | block_height = 1 869 | 870 | fig, ax = plt.subplots(figsize=(18, 7)) 871 | 872 | # Plot Beam 873 | ax.plot(beam_x, beam_y_vals, color='black', linewidth=3, label='Beam') 874 | ax.scatter(beam_x[0], beam_y - 0.15, marker='^', color='red', s=300, zorder=6) 875 | 876 | # Plot Rollers 877 | ax.scatter(unique_rollers, [beam_y]*len(unique_rollers), 878 | marker='o', color='seagreen', s=200, 879 | label='Rollers', zorder=5, edgecolors='k') 880 | 881 | # Plot Forces 882 | for case_idx in range(n_cases): 883 | fx_list = user_force_x[case_idx] 884 | fv_list = user_force_vals[case_idx] 885 | color = case_colors[case_idx] 886 | label = case_labels[case_idx] 887 | 888 | for idx, (fx, fv) in enumerate(zip(fx_list, fv_list)): 889 | arrow_length = abs(fv) * arrow_scale 890 | start_point = (fx, beam_y + arrow_length) 891 | end_point = (fx, beam_y) 892 | 893 | arrow = FancyArrowPatch( 894 | posA=start_point, posB=end_point, 895 | arrowstyle='-|>', 896 | mutation_scale=20, 897 | color=color, 898 | linewidth=2, 899 | alpha=0.8, 900 | label=label if idx == 0 else "" 901 | ) 902 | ax.add_patch(arrow) 903 | ax.text(fx, beam_y + arrow_length + desired_max_arrow_length * 0.02, 904 | f"{fv:.0f}", ha='center', va='bottom', 905 | fontsize=10, color=color, fontweight='bold') 906 | 907 | # Plot predicted I as rectangles 908 | for idx, (x_pos, I_val) in enumerate(zip(beam_positions, pred_1x_unstd)): 909 | color = cmap(norm(I_val)) 910 | rect_x = x_pos - block_width / 2 911 | rect_y = beam_y - ( 912 | I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8) 913 | ) * block_height / 2 914 | 915 | rect = Rectangle((rect_x, rect_y), 916 | block_width, 917 | (I_val / (pred_1x_unstd.max() - pred_1x_unstd.min() + 1e-8)) * block_height, 918 | linewidth=0, 919 | edgecolor=None, 920 | facecolor=color, 921 | alpha=0.6) 922 | ax.add_patch(rect) 923 | 924 | # Colorbar 925 | sm = cm.ScalarMappable(cmap=cmap, norm=norm) 926 | sm.set_array([]) 927 | cbar = fig.colorbar(sm, ax=ax, orientation='vertical', fraction=0.046, pad=0.04) 928 | cbar.set_label('Predicted I (m$^4$)', fontsize=16) 929 | cbar.ax.tick_params(labelsize=10) 930 | 931 | ax.set_title("Beam Setup with Applied Forces and Bayesian-Predicted I", 932 | fontsize=22, fontweight='bold', pad=20) 933 | ax.set_xlabel("Beam Length (m)", fontsize=16, fontweight='semibold') 934 | ax.set_xlim(-5, L_beam + 5) 935 | ax.set_ylim(-2.5, 2.5) 936 | ax.set_xticks(np.arange(0, L_beam + 5, 5)) 937 | ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7) 938 | 939 | legend_elements = [ 940 | Line2D([0], [0], color='black', lw=3, label='Beam'), 941 | Line2D([0], [0], marker=(3, 0, -90), color='red', label='Pin', 942 | markerfacecolor='red', markersize=15), 943 | Line2D([0], [0], marker='o', color='seagreen', label='Rollers', 944 | markerfacecolor='seagreen', markeredgecolor='k', markersize=15), 945 | ] 946 | for color, label in zip(case_colors, case_labels): 947 | legend_elements.append(Line2D([0], [0], color=color, lw=2, label=label)) 948 | 949 | ax.legend(handles=legend_elements, loc='lower right', fontsize=12) 950 | plt.tight_layout() 951 | plt.show() 952 | --------------------------------------------------------------------------------