ML for Structural Beam Design: PyTorch example comparing linear and non-linear neural networks to traditional bending moment calculations for steel beam depth. Demonstrates augmenting engineering formulas with data-driven learning. (Flocode Newsletter #062)

061-pytorch_beam_example.py

"""
Flocode Newsletter #062 - Machine Learning for Engineers: When and How Should We Use it?

This script demonstrates a simple application of machine learning (ML) to a
structural engineering problem: predicting the required depth of a steel beam
based on its span and applied uniformly distributed load (UDL). It compares
three ML models (linear, non-linear, and deep non-linear neural networks)
with a traditional calculation based on the bending moment formula.

The key concept is that the ML models are trained on data *derived* from the
traditional formula, but with added random noise to simulate real-world
imperfections and variations. This demonstrates the potential of ML to
augment traditional engineering methods by learning from data and adapting
to complexities not captured by idealized equations.

This code accompanies the article and is intended for educational purposes.
It is not intended for use in real-world structural design without further
validation and refinement.

Key Features:
- Generates synthetic beam data (span, load, depth) based on the traditional
  bending moment formula.
- Adds random noise to the calculated depths to simulate real-world variations.
- Defines three PyTorch neural network models:
    - LinearModel: A simple linear regression model.
    - NonLinearModel: A moderately deep neural network.
    - DeepNonLinearModel: A deeper neural network.
- Trains the models using the generated data.
- Evaluates model performance using R-squared and RMSE.
- Creates visualizations:
    - Training loss comparison.
    - Predictions vs. traditional method.
    - Error distribution.
    - Depth vs. span (fixed load).
    - Depth vs. load (fixed span).
    - Parameter Importance Scatter Plot (3D).
- Prints example predictions and performance metrics.
- Demonstrates simple use case scenarios.

Usage:
1.  Install the necessary libraries: `pip install numpy torch scikit-learn matplotlib`
2.  Run the script: `061-pytorch_beam_example.py`

Author: James O'Reilly, Flocode.
"""
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score, mean_squared_error
from sklearn.preprocessing import StandardScaler

# Set random seeds for reproducibility
np.random.seed(42)
torch.manual_seed(42)

# Steel beam design parameters
fy_nominal = 275  # Steel yield strength (MPa)
E = 210000  # Young's modulus (MPa), unused in this calc but kept for context
safety_factor = 1.15
b_nominal = 100  # Nominal width (mm)

def calculate_traditional_beam_depth(span, udl, fy=fy_nominal, b=b_nominal):
    """Traditional steel beam depth calculation with fixed width."""
    M = (udl * span**2) / 8  # kNm
    required_Zx = (M * 1e6) / (fy / safety_factor)  # mm³
    depth = np.sqrt(6 * required_Zx / b) / 10  # cm
    return depth

def generate_beam_data(n_samples, fy_nominal=275, b_nominal=100):
    """Generate synthetic beam data with realistic variations.

    Simulates variations in load, yield strength, and beam width
    to create a dataset reflecting real-world uncertainties.

    Args:
        n_samples: The number of data points to generate.
        fy_nominal: Nominal yield strength of steel (MPa).
        b_nominal: Nominal beam width (mm).

    Returns:
        A tuple containing:
            - A NumPy array with span and load values.
            - A NumPy array of "true" depths (with variations).
            - A NumPy array of traditional depths (no variations).
    """
    spans = np.random.uniform(4, 12, n_samples)  # m
    loads = np.random.uniform(10, 50, n_samples)  # kN/m

    # Baseline traditional depth (nominal conditions)
    M = (loads * spans**2) / 8  # kNm
    Zx_nominal = (M * 1e6) / (fy_nominal / safety_factor)  # mm³
    traditional_depths = np.sqrt(6 * Zx_nominal / b_nominal) / 10  # cm

    # Realistic variations
    load_adjustment = np.random.normal(1, 0.05, n_samples)  # ±5% load variation
    M_adjusted = M * load_adjustment

    fy_variation = np.random.normal(1, 0.03, n_samples)  # ±3% fy variation
    fy_adjusted = fy_nominal * fy_variation

    b_variation = np.random.normal(1, 0.05, n_samples)  # ±5% width variation
    b_adjusted = b_nominal * b_variation

    # "True" depth with adjusted parameters
    Zx_adjusted = (M_adjusted * 1e6) / (fy_adjusted / safety_factor)  # mm³
    true_depths = np.sqrt(6 * Zx_adjusted / b_adjusted) / 10  # cm

    return np.column_stack((spans, loads)), true_depths, traditional_depths

# Parameters
n_samples = 1000
n_epochs = 2000

# Model definitions
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 1)
        )
        for layer in self.network:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight)
                nn.init.constant_(layer.bias, 0.1)
    
    def forward(self, x):
        return self.network(x)

class NonLinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(2, 32),
            nn.ReLU(),
            nn.Linear(32, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.network(x)

class DeepNonLinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(2, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, 1)
        )
    
    def forward(self, x):
        return self.network(x)

def train_model(model, X_tensor, y_tensor, epochs):
    """Train the model with custom parameters."""
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
    loss_fn = nn.MSELoss()
    losses = []
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        predictions = model(X_tensor)
        loss = loss_fn(predictions, y_tensor)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        if epoch % 500 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.6f}")
    
    return losses

# --- Main script ---
print("Generating synthetic beam data...")
X, y, traditional_depths = generate_beam_data(n_samples)

# Normalize inputs and outputs
scaler_X = StandardScaler()
X_normalized = scaler_X.fit_transform(X)
X_tensor = torch.tensor(X_normalized, dtype=torch.float32)

scaler_y = StandardScaler()
y_normalized = scaler_y.fit_transform(y.reshape(-1, 1))
y_tensor = torch.tensor(y_normalized, dtype=torch.float32)

# Create and train models
models = {
    'Linear': LinearModel(),
    'NonLinear': NonLinearModel(),
    'DeepNonLinear': DeepNonLinearModel()
}

losses_dict = {}
predictions_dict = {}
metrics_dict = {}

print("\nTraining models:")
print("-" * 50)
for name, model in models.items():
    print(f"\nTraining {name} model:")
    losses_dict[name] = train_model(model, X_tensor, y_tensor, n_epochs)
    predictions_normalized = model(X_tensor).detach().numpy()
    predictions_dict[name] = scaler_y.inverse_transform(predictions_normalized).flatten()
    metrics_dict[name] = {
        'R2': r2_score(y, predictions_dict[name]),
        'RMSE': np.sqrt(mean_squared_error(y, predictions_dict[name]))
    }

# --- Create visualizations ---
plt.figure(figsize=(14, 15))
plt.suptitle('Machine Learning for Structural Beam Design', fontsize=16)

colors = {
    'Linear': 'skyblue',
    'NonLinear': 'lightgreen',
    'DeepNonLinear': 'lightcoral',
    'Traditional': 'gray'
}

# Plot 1: Training Loss Comparison
plt.subplot(3, 2, 1)
for name, losses in losses_dict.items():
    plt.plot(losses, label=name, color=colors[name])
plt.title('Training Loss Comparison')
plt.xlabel('Epoch')
plt.ylabel('Loss (MSE)')
plt.yscale('log')
plt.legend()
plt.grid(True)

# Plot 2: Predictions vs Actual
plt.subplot(3, 2, 2)
for name, predictions in predictions_dict.items():
    plt.scatter(y, predictions, alpha=0.5, label=name, color=colors[name])
plt.plot([min(y), max(y)], [min(y), max(y)], '--', color='gray', label='Perfect Match')
plt.title('Predictions vs Actual Depths')
plt.xlabel('Actual Depth (cm)')
plt.ylabel('Predicted Depth (cm)')
plt.legend()
plt.grid(True)

# Plot 3: Error Distribution
plt.subplot(3, 2, 3)
for name, predictions in predictions_dict.items():
    error = (predictions - y) / y * 100
    plt.hist(error, bins=30, alpha=0.6, label=name, color=colors[name])
plt.title('Error Distribution')
plt.xlabel('Percentage Error (%)')
plt.ylabel('Frequency')
plt.legend()
plt.grid(True)

# Plot 4: Depth vs Span (fixed load)
plt.subplot(3, 2, 4)
load_fixed = 30  # kN/m
spans_test = np.linspace(4, 12, 100)
test_data = np.column_stack((spans_test, np.ones_like(spans_test) * load_fixed))
test_normalized = scaler_X.transform(test_data)
test_tensor = torch.tensor(test_normalized, dtype=torch.float32)
traditional_test = np.array([calculate_traditional_beam_depth(s, load_fixed) for s in spans_test])

for name, model in models.items():
    predictions_normalized = model(test_tensor).detach().numpy()
    predictions = scaler_y.inverse_transform(predictions_normalized).flatten()
    plt.plot(spans_test, predictions, label=name, color=colors[name])

plt.plot(spans_test, traditional_test, '--', color='gray', label='Traditional')
plt.title(f'Beam Depth vs Span (Load = {load_fixed} kN/m)')
plt.xlabel('Span (m)')
plt.ylabel('Depth (cm)')
plt.legend()
plt.grid(True)

# Plot 5: Depth vs Load (fixed span)
plt.subplot(3, 2, 5)
span_fixed = 8  # m
loads_test = np.linspace(10, 50, 100)
test_data = np.column_stack((np.ones_like(loads_test) * span_fixed, loads_test))
test_normalized = scaler_X.transform(test_data)
test_tensor = torch.tensor(test_normalized, dtype=torch.float32)
traditional_test = np.array([calculate_traditional_beam_depth(span_fixed, l) for l in loads_test])

for name, model in models.items():
    predictions_normalized = model(test_tensor).detach().numpy()
    predictions = scaler_y.inverse_transform(predictions_normalized).flatten()
    plt.plot(loads_test, predictions, label=name, color=colors[name])

plt.plot(loads_test, traditional_test, '--', color='gray', label='Traditional')
plt.title(f'Beam Depth vs Load (Span = {span_fixed} m)')
plt.xlabel('Load (kN/m)')
plt.ylabel('Depth (cm)')
plt.legend()
plt.grid(True)

# Plot 6: Parameter Scatter Plot
plt.subplot(3, 2, 6)
plt.scatter(X[:, 0], X[:, 1], c=y, cmap='viridis', alpha=0.6)
plt.colorbar(label='Beam Depth (cm)')
plt.title('Span vs Load (Color: Actual Depth)')
plt.xlabel('Span (m)')
plt.ylabel('Load (kN/m)')
plt.grid(True)

plt.tight_layout(rect=[0, 0, 1, 0.95])
plt.show()

# Print metrics
print("\nModel Performance Metrics:")
print("-" * 60)
for name, metrics in metrics_dict.items():
    print(f"{name} Model:")
    print(f"  R² Score: {metrics['R2']:.4f}")
    print(f"  RMSE: {metrics['RMSE']:.4f} cm")
    print("-" * 40)

# Print example predictions
print("\nExample Predictions:")
header = "Span(m)  Load(kN/m)  Actual(cm)  Trad(cm)"
for name in models.keys():
    header += f"  {name:12s}"
print(header)
print("-" * (47 + len(models) * 14))

for i in range(10):
    row = f"{X[i,0]:6.1f}   {X[i,1]:9.1f}    {y[i]:10.2f}  {traditional_depths[i]:8.2f}"
    for name, predictions in predictions_dict.items():
        row += f"    {predictions[i]:8.2f}"
    print(row)