# -*- coding: utf-8 -*- """ Created on Fri Mar 14 21:22:39 2025 @author: malij """ import torch import torch.nn as nn import torch.optim as optim import numpy as np import matplotlib.pyplot as plt from torch.utils.data import Dataset, DataLoader, random_split from sklearn.metrics import r2_score import seaborn as sns import pickle # Load the data with open('sensorDataTotal.pkl', 'rb') as f: sensorDataTotal = pickle.load(f) with open('elementStatesTotal.pkl', 'rb') as f: elementStatesTotal = pickle.load(f) with open('stiffnessMatrixToatal.pkl', 'rb') as f: stiffnessMatrixToatal = pickle.load(f) stiffnessMatrixToatal = np.array(stiffnessMatrixToatal) #font size plt.rcParams.update({'font.size': 21}) # Increase base font size # Data preprocessing functions (your existing functions) def normalize_forces(forces): forces = np.array(forces) normalized_forces = np.zeros_like(forces) normalization_factors = [] for i in range(9): component = forces[:, :, i] max_val = np.max(np.abs(component)) normalized_forces[:, :, i] = component / (max_val + 1e-8) normalization_factors.append(max_val) return normalized_forces, tuple(normalization_factors) def normalize_accelerations(acc_data): """ Normalize acceleration data across all samples. Parameters: ----------- acc_data : numpy.ndarray Acceleration data with shape (num_scenarios, num_timesteps, num_nodes, num_dofs) Returns: -------- normalized_data : numpy.ndarray Normalized acceleration data with the same shape normalization_factors : dict Dictionary containing 'mean' and 'std' for each DOF of each node """ # Original shape: (num_scenarios, num_timesteps, num_nodes, num_dofs) num_scenarios, num_timesteps, num_nodes, num_dofs = acc_data.shape # Reshape to combine scenarios and timesteps for easier normalization # New shape: (num_scenarios * num_timesteps, num_nodes, num_dofs) reshaped_data = acc_data.reshape(-1, num_nodes, num_dofs) normalized_data = np.zeros_like(reshaped_data) normalization_factors = {'mean': np.zeros((num_nodes, num_dofs)), 'std': np.zeros((num_nodes, num_dofs))} # Normalize each node and DOF separately for node in range(num_nodes): for dof in range(num_dofs): values = reshaped_data[:, node, dof] mean = np.mean(values) std = np.std(values) # Avoid division by zero if std < 1e-10: std = 1.0 normalized_data[:, node, dof] = (values - mean) / std normalization_factors['mean'][node, dof] = mean normalization_factors['std'][node, dof] = std # Reshape back to original shape normalized_data = normalized_data.reshape(num_scenarios, num_timesteps, num_nodes, num_dofs) return normalized_data, normalization_factors def extract_sensor_arrays(sensorDataTotal): first_scenario = sensorDataTotal[0] num_scenarios = len(sensorDataTotal) first_acc = list(first_scenario['accelerations'].values())[0] num_timesteps = first_acc.shape[0] num_nodes = len(first_scenario['accelerations']) num_dofs = first_acc.shape[1] print(first_acc.shape) all_accelerations = np.zeros((num_scenarios, num_timesteps, num_nodes, num_dofs)) for scenario_idx, scenario in enumerate(sensorDataTotal): for node_idx, (node_id, acc_data) in enumerate(scenario['accelerations'].items()): all_accelerations[scenario_idx, :, node_idx, :] = acc_data print(all_accelerations.shape) # Remove the time steps from degrees of freedom (time steps, x, y, teta) all_accelerations = all_accelerations[:,:,:,1:] print(all_accelerations.shape) # Normalize the acceleration data all_accelerations, normalization_factors = normalize_accelerations(all_accelerations) # Reshape to flatten nodes and DOFs # From (scenarios, timesteps, nodes, dofs) to (scenarios, timesteps, nodes*dofs) flattened_accelerations = all_accelerations.reshape(num_scenarios, num_timesteps, -1) return flattened_accelerations, normalization_factors # NEW FUNCTIONS FOR NOISE AUGMENTATION # def add_gaussian_noise_for_trainingAndVal(data, noise_level): # """ # Add Gaussian noise to sensor data # Parameters: # ----------- # data : numpy.ndarray # Input sensor data # noise_level : float # Standard deviation of the Gaussian noise as a fraction of data standard deviation # Returns: # -------- # noisy_data : numpy.ndarray # Data with added Gaussian noise # """ # # Calculate standard deviation of the data # data_std = np.std(data) # # Generate Gaussian noise # noise = np.random.normal(0, noise_level * data_std, data.shape) # # Add noise to data # noisy_data = data + noise # return noisy_data # def add_gaussian_noise_degradation_eval(data, noise_level, normalization_factors=None): # """ # Add Gaussian noise to sensor data, properly calibrated for normalized data # Parameters: # ----------- # data : numpy.ndarray # Input sensor data (normalized) # noise_level : float # Standard deviation of the Gaussian noise as a fraction of original data std # normalization_factors : dict, optional # Dictionary containing normalization factors from original data # Returns: # -------- # noisy_data : numpy.ndarray # Data with added Gaussian noise # """ # if normalization_factors is not None: # # Use the average std from original data for proper scaling # avg_std = np.mean(normalization_factors['std']) # effective_noise_level = noise_level * avg_std # print(f"Nominal noise level: {noise_level}, Effective noise std: {effective_noise_level}") # else: # # If no normalization factors provided, use the normalized data std # data_std = np.std(data) # effective_noise_level = noise_level * data_std # # Generate Gaussian noise with properly scaled magnitude # noise = np.random.normal(0, effective_noise_level, data.shape) # # Add noise to data # noisy_data = data + noise # return noisy_data def add_gaussian_noise(data, noise_level, normalization_factors=None): """ Add Gaussian noise to sensor data, properly calibrated for normalized data Parameters: ----------- data : numpy.ndarray Input sensor data (normalized) noise_level : float Standard deviation of the Gaussian noise to apply directly normalization_factors : dict, optional Not used in this version Returns: -------- noisy_data : numpy.ndarray Data with added Gaussian noise """ # Since data is already normalized, apply noise directly noise = np.random.normal(0, noise_level, data.shape) # Add noise to data noisy_data = data + noise # print(f"Applied noise with std: {noise_level}") return noisy_data def add_impulse_noise(data, noise_level, prob=0.05): """ Add impulse noise (random spikes) to sensor data Parameters: ----------- data : numpy.ndarray Input sensor data noise_level : float Magnitude of impulse noise as a fraction of data range prob : float Probability of impulse noise at each point Returns: -------- noisy_data : numpy.ndarray Data with added impulse noise """ # Calculate data range data_range = np.max(data) - np.min(data) # Generate random mask for impulse locations mask = np.random.random(data.shape) < prob # Generate impulse values impulses = np.random.uniform(-noise_level * data_range, noise_level * data_range, data.shape) # Apply impulses only at mask locations noisy_data = data.copy() noisy_data[mask] += impulses[mask] return noisy_data def add_missing_data(data, missing_prob): """ Simulate missing data by zeroing out random values Parameters: ----------- data : numpy.ndarray Input sensor data missing_prob : float Probability of a data point being missing Returns: -------- data_with_missing : numpy.ndarray Data with simulated missing values """ # Create a copy of the data data_with_missing = data.copy() # Create random mask for missing data mask = np.random.random(data.shape) < missing_prob # Set masked values to zero (simulating missing data) data_with_missing[mask] = 0 return data_with_missing def add_sensor_drift(data, drift_factor=0.0005): """ Add gradual drift to sensor readings over time Parameters: ----------- data : numpy.ndarray Input sensor data with shape (batch_size, time_steps, features) drift_factor : float Factor controlling the magnitude of drift Returns: -------- data_with_drift : numpy.ndarray Data with added sensor drift """ data_with_drift = data.copy() batch_size, time_steps, features = data.shape for b in range(batch_size): for f in range(features): # Generate random drift direction and magnitude drift_direction = np.random.choice([-1, 1]) drift_magnitude = np.random.uniform(0, drift_factor) # Create linear drift over time drift = drift_direction * drift_magnitude * np.arange(time_steps) # Apply drift to this feature data_with_drift[b, :, f] += drift return data_with_drift # Preprocess data stiffnessMatrixToatal, stiffness_norm_factors = normalize_forces(stiffnessMatrixToatal) acc_Sensors_ALL, acc_norm_factors = extract_sensor_arrays(sensorDataTotal) print(acc_Sensors_ALL[0]) # NEW DIRECT MODEL: Sensor data → Element states prediction # class DirectSensorToElementModel(nn.Module): # def __init__(self, input_features, output_size=21): # output_size is the number of element states # super().__init__() # # Encoder part - similar to the original SensorToStiffnessModel # self.encoder = nn.Sequential( # nn.Conv1d(input_features, 512, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(512), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(512, 256, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(256), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(256, 128, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(128), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(128, 64, kernel_size=3, stride=1, padding=1), # nn.BatchNorm1d(64), # nn.ReLU(), # nn.Dropout(0.2), # ) # # Calculate the flattened size after convolutions # with torch.no_grad(): # sample_input = torch.randn(1, input_features, 50) # sample_output = self.encoder(sample_input) # self.flatten_size = sample_output.numel() # # Direct element state projector # self.element_projector = nn.Sequential( # nn.Linear(self.flatten_size, 512), # nn.BatchNorm1d(512), # nn.ReLU(), # nn.Dropout(0.3), # nn.Linear(512, 256), # nn.BatchNorm1d(256), # nn.ReLU(), # nn.Dropout(0.2), # nn.Linear(256, 128), # nn.BatchNorm1d(128), # nn.ReLU(), # nn.Dropout(0.2), # nn.Linear(128, output_size) # ) # def forward(self, x): # x = x.transpose(1, 2) # Change from [batch, time, features] to [batch, features, time] # x = self.encoder(x) # x = x.view(x.size(0), -1) # Flatten # element_states = self.element_projector(x) # return element_states class DirectSensorToElementModel(nn.Module): def __init__(self, input_features, output_size=21, timesteps=50): super().__init__() # 1. Improved encoder with residual connections (from model_1) self.encoder_block1 = nn.Sequential( nn.Conv1d(input_features, 512, kernel_size=3, stride=2, padding=1), nn.BatchNorm1d(512), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) self.encoder_block2 = nn.Sequential( nn.Conv1d(512, 512, kernel_size=3, stride=1, padding=1), nn.BatchNorm1d(512), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) self.encoder_block3 = nn.Sequential( nn.Conv1d(512, 256, kernel_size=3, stride=2, padding=1), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) self.encoder_block4 = nn.Sequential( nn.Conv1d(256, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm1d(256), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) self.encoder_block5 = nn.Sequential( nn.Conv1d(256, 128, kernel_size=3, stride=2, padding=1), nn.BatchNorm1d(128), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) self.encoder_block6 = nn.Sequential( nn.Conv1d(128, 128, kernel_size=3, stride=1, padding=1), nn.BatchNorm1d(128), nn.LeakyReLU(0.2), nn.Dropout(0.2) ) # 2. Attention mechanism (from model_1) self.attention = nn.Sequential( nn.Conv1d(128, 1, kernel_size=1), nn.Sigmoid() ) # Calculate the flattened size after convolutions with torch.no_grad(): sample_input = torch.randn(1, input_features, timesteps) # Forward through encoder blocks x = self.encoder_block1(sample_input) x = self.encoder_block2(x) + x # Residual connection x = self.encoder_block3(x) x = self.encoder_block4(x) + x # Residual connection x = self.encoder_block5(x) x = self.encoder_block6(x) + x # Residual connection # Apply attention att = self.attention(x) x = x * att self.flatten_size = x.numel() print(f"Flatten size for element model: {self.flatten_size}") # 3. Improved element state projector with skip connections hidden_size1 = 512 hidden_size2 = 256 hidden_size3 = 128 # First layer self.fc1 = nn.Linear(self.flatten_size, hidden_size1) self.bn1 = nn.BatchNorm1d(hidden_size1) self.act1 = nn.LeakyReLU(0.2) self.drop1 = nn.Dropout(0.3) # Second layer with skip connection self.fc2 = nn.Linear(hidden_size1, hidden_size2) self.bn2 = nn.BatchNorm1d(hidden_size2) self.act2 = nn.LeakyReLU(0.2) self.drop2 = nn.Dropout(0.3) # Projection for skip connection self.skip_proj = nn.Linear(hidden_size1, hidden_size2) # Third layer self.fc3 = nn.Linear(hidden_size2, hidden_size3) self.bn3 = nn.BatchNorm1d(hidden_size3) self.act3 = nn.LeakyReLU(0.2) self.drop3 = nn.Dropout(0.2) # Output layer self.output_layer = nn.Linear(hidden_size3, output_size) def forward(self, x): batch_size = x.size(0) # Transpose for CNN x = x.transpose(1, 2) # [batch, time, features] -> [batch, features, time] # Encode with residual connections x1 = self.encoder_block1(x) x2 = self.encoder_block2(x1) + x1 # Residual connection x3 = self.encoder_block3(x2) x4 = self.encoder_block4(x3) + x3 # Residual connection x5 = self.encoder_block5(x4) x6 = self.encoder_block6(x5) + x5 # Residual connection # Apply attention mechanism attention_weights = self.attention(x6) x_attended = x6 * attention_weights # Flatten x_flat = x_attended.view(batch_size, -1) # Apply fully connected layers with skip connection x = self.fc1(x_flat) x = self.bn1(x) x = self.act1(x) x = self.drop1(x) # Skip connection with projection (since dimensions might differ) skip = self.skip_proj(x) x = self.fc2(x) x = self.bn2(x) x = self.act2(x) x = self.drop2(x) # Add skip connection x = x + skip x = self.fc3(x) x = self.bn3(x) x = self.act3(x) x = self.drop3(x) # Final output element_states = self.output_layer(x) return element_states # Add a method to apply physics-based regularization if needed def add_physics_regularization(self, element_states, alpha=0.01): """ Add regularization to ensure the element states follow physical constraints (e.g., states should be between 0 and 1 for damage indicators) """ batch_size = element_states.size(0) # Example constraint: penalize values outside [0,1] range below_zero_penalty = torch.relu(-element_states).sum() above_one_penalty = torch.relu(element_states - 1.0).sum() # Total penalty penalty = below_zero_penalty + above_one_penalty return alpha * penalty / batch_size # # Keep the original models for comparison # class ElementStateDecoder(nn.Module): # def __init__(self, input_size=21*9, output_size=21): # super().__init__() # self.decoder = nn.Sequential( # nn.Linear(input_size, 128), # nn.LeakyReLU(0.1), # nn.BatchNorm1d(128), # nn.Dropout(0.2), # nn.Linear(128, 64), # nn.LeakyReLU(0.1), # nn.BatchNorm1d(64), # nn.Dropout(0.2), # nn.Linear(64, 32), # nn.LeakyReLU(0.1), # nn.BatchNorm1d(32), # nn.Dropout(0.1), # nn.Linear(32, output_size) # ) # def forward(self, x): # return self.decoder(x) # class SensorToStiffnessModel(nn.Module): # def __init__(self, input_features, output_size=21*9): # 21*9 for stiffness matrix size # super().__init__() # # Keep the encoder part similar # self.encoder = nn.Sequential( # nn.Conv1d(input_features, 512, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(512), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(512, 256, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(256), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(256, 128, kernel_size=3, stride=2, padding=1), # nn.BatchNorm1d(128), # nn.ReLU(), # nn.Dropout(0.2), # nn.Conv1d(128, 64, kernel_size=3, stride=1, padding=1), # nn.BatchNorm1d(64), # nn.ReLU(), # nn.Dropout(0.2), # ) # # Calculate the flattened size after convolutions # with torch.no_grad(): # sample_input = torch.randn(1, input_features, 50) # sample_output = self.encoder(sample_input) # self.flatten_size = sample_output.numel() # # Replace bottleneck projector with stiffness matrix projector # self.stiffness_projector = nn.Sequential( # nn.Linear(self.flatten_size, 768), # nn.BatchNorm1d(768), # nn.ReLU(), # nn.Dropout(0.3), # nn.Linear(768, 384), # nn.BatchNorm1d(384), # nn.ReLU(), # nn.Dropout(0.2), # nn.Linear(384, output_size) # ) # def forward(self, x): # x = x.transpose(1, 2) # x = self.encoder(x) # x = x.view(x.size(0), -1) # stiffness_matrix = self.stiffness_projector(x) # return stiffness_matrix class DirectModel(nn.Module): def __init__(self, sensor_model, element_decoder): super().__init__() self.sensor_model = sensor_model self.element_decoder = element_decoder def forward(self, x): stiffness_matrix = self.sensor_model(x) states = self.element_decoder(stiffness_matrix) return states, stiffness_matrix # Dataset classes class DirectSensorDataset(Dataset): def __init__(self, sensor_data, element_states): self.sensor_data = torch.FloatTensor(sensor_data) self.element_states = torch.FloatTensor(element_states) def __len__(self): return len(self.sensor_data) def __getitem__(self, idx): return self.sensor_data[idx], self.element_states[idx] class StiffnessDataset(Dataset): def __init__(self, stiffness_matrices, element_states): # Flatten the stiffness matrices from (numSamples, 21, 9) to (numSamples, 21*9) self.stiffness_matrices = torch.FloatTensor(stiffness_matrices.reshape(len(stiffness_matrices), -1)) self.element_states = torch.FloatTensor(element_states) def __len__(self): return len(self.stiffness_matrices) def __getitem__(self, idx): return self.stiffness_matrices[idx], self.element_states[idx] class CombinedDataset(Dataset): def __init__(self, sensor_data, stiffness_matrices, element_states): # sensor_data is already flattened in the preprocessing step self.sensor_data = torch.FloatTensor(sensor_data) self.stiffness_matrices = torch.FloatTensor(stiffness_matrices.reshape(len(stiffness_matrices), -1)) self.element_states = torch.FloatTensor(element_states) def __len__(self): return len(self.sensor_data) def __getitem__(self, idx): return self.sensor_data[idx], self.stiffness_matrices[idx], self.element_states[idx] class TrainingHistory: def __init__(self): self.history = { 'train_loss': [], 'val_loss': [], 'train_r2': [], 'val_r2': [], # # Add these missing metrics for the two-stage model # 'train_stiffness_loss': [], 'val_stiffness_loss': [], # 'train_states_loss': [], 'val_states_loss': [], # 'train_r2_states': [], 'val_r2_states': [], # 'train_r2_stiffness': [], 'val_r2_stiffness': [] } def update(self, metrics, phase='train'): for key, value in metrics.items(): self.history[f'{phase}_{key}'].append(value) def get_latest(self, metric, phase='train'): return self.history[f'{phase}_{metric}'][-1] # NEW TRAINING FUNCTION FOR DIRECT MODEL def train_direct_sensor_model(model, train_loader, val_loader, num_epochs=100, device='cuda'): model = model.to(device) criterion = nn.MSELoss() # optimizer = optim.Adam(model.parameters(), lr=0.001) optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=1e-4) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=10) history = TrainingHistory() best_val_loss = float('inf') best_model = None # Noise augmentation parameters # initial_noise_level = 0.000 # max_noise_level = 0.000 # noise_increase_rate = 0.000 # Noise augmentation parameters remain the same pre_train_no_noise_epoch = 0 initial_noise_level = 0.0 max_noise_level = 0.33 noise_increase_rate = 0.01 for epoch in range(num_epochs): if epoch > pre_train_no_noise_epoch: # noise_level = min(initial_noise_level + noise_increase_rate * (epoch-pre_train_no_noise_epoch), max_noise_level) noise_level = min(initial_noise_level + noise_increase_rate * epoch, max_noise_level) else: noise_level = 0.00 for phase in ['train', 'val']: if phase == 'train': model.train() loader = train_loader else: model.eval() loader = val_loader running_loss = 0.0 all_pred_states = [] all_true_states = [] for sensor_data, true_states in loader: # Apply noise augmentation only during training if phase in ['train', 'val']: sensor_data_np = sensor_data.numpy() noise_type = epoch % 4 if noise_type == 0: # sensor_data_np = add_gaussian_noise_for_trainingAndVal(sensor_data_np, noise_level) # sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level, acc_norm_factors) sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level, acc_norm_factors) elif noise_type == 1: sensor_data_np = add_impulse_noise(sensor_data_np, noise_level) elif noise_type == 2: sensor_data_np = add_missing_data(sensor_data_np, noise_level * 0.5) else: sensor_data_np = add_sensor_drift(sensor_data_np, noise_level * 0.01) sensor_data = torch.FloatTensor(sensor_data_np) sensor_data = sensor_data.to(device) true_states = true_states.to(device) if phase == 'train': optimizer.zero_grad() with torch.set_grad_enabled(phase == 'train'): pred_states = model(sensor_data) loss = criterion(pred_states, true_states) # Add physics regularization if needed physics_reg = model.add_physics_regularization(pred_states, alpha=0.01) loss = loss + physics_reg if phase == 'train': loss.backward() optimizer.step() running_loss += loss.item() * sensor_data.size(0) all_pred_states.append(pred_states.detach().cpu().numpy()) all_true_states.append(true_states.cpu().numpy()) epoch_loss = running_loss / len(loader.dataset) all_pred_states = np.concatenate(all_pred_states) all_true_states = np.concatenate(all_true_states) r2 = r2_score(all_true_states.flatten(), all_pred_states.flatten()) metrics = { 'loss': epoch_loss, 'r2': r2 } history.update(metrics, phase) if phase == 'val' and epoch_loss < best_val_loss: best_val_loss = epoch_loss best_model = model.state_dict() print(f'Epoch {epoch+1}/{num_epochs} - {phase.capitalize()}:') print(f'Loss: {epoch_loss:.4f}, R²: {r2:.4f}') print(f'Current noise level: {noise_level:.4f}') scheduler.step(history.get_latest('loss', 'val')) return best_model, history, max_noise_level # NEW EVALUATION FUNCTION FOR DIRECT MODEL def evaluate_direct_sensor_model_robustness(model, test_loader, device='cuda', noise_levels=[0, 0.05, 0.1, 0.2]): model.eval() results = {} for noise_level in noise_levels: all_pred_states = [] all_true_states = [] running_loss = 0.0 criterion = nn.MSELoss() with torch.no_grad(): for sensor_data, true_states in test_loader: if noise_level > 0: sensor_data_np = sensor_data.numpy() # sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level) # sensor_data_np = add_gaussian_noise_degradation_eval(sensor_data_np, noise_level, acc_norm_factors) sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level, acc_norm_factors) sensor_data = torch.FloatTensor(sensor_data_np) sensor_data = sensor_data.to(device) true_states = true_states.to(device) pred_states = model(sensor_data) loss = criterion(pred_states, true_states) running_loss += loss.item() * sensor_data.size(0) all_pred_states.append(pred_states.cpu().numpy()) all_true_states.append(true_states.cpu().numpy()) all_pred_states = np.concatenate(all_pred_states) all_true_states = np.concatenate(all_true_states) test_loss = running_loss / len(test_loader.dataset) r2_states = r2_score(all_true_states.flatten(), all_pred_states.flatten()) results[noise_level] = { 'loss': test_loss, 'r2': r2_states } print(f"Noise Level {noise_level}:") print(f" Loss = {test_loss:.4f}") print(f" R² = {r2_states:.4f}") return results # Keep the original evaluation function for comparison def evaluate_model_robustness(model, test_loader, device='cuda', noise_levels=[0, 0.05, 0.1, 0.2]): model.eval() results = {} for noise_level in noise_levels: all_pred_states = [] all_true_states = [] all_pred_stiffness = [] all_true_stiffness = [] running_loss = 0.0 running_stiffness_loss = 0.0 running_states_loss = 0.0 criterion = nn.MSELoss() with torch.no_grad(): for sensor_data, true_stiffness, true_states in test_loader: if noise_level > 0: sensor_data_np = sensor_data.numpy() # sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level) # sensor_data_np = add_gaussian_noise_degradation_eval(sensor_data_np, noise_level, acc_norm_factors) sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level, acc_norm_factors) sensor_data = torch.FloatTensor(sensor_data_np) sensor_data = sensor_data.to(device) true_states = true_states.to(device) true_stiffness = true_stiffness.to(device) pred_states, pred_stiffness = model(sensor_data) stiffness_loss = criterion(pred_stiffness, true_stiffness) states_loss = criterion(pred_states, true_states) loss = 0.7 * states_loss + 0.3 * stiffness_loss running_loss += loss.item() * sensor_data.size(0) running_stiffness_loss += stiffness_loss.item() * sensor_data.size(0) running_states_loss += states_loss.item() * sensor_data.size(0) all_pred_states.append(pred_states.cpu().numpy()) all_true_states.append(true_states.cpu().numpy()) all_pred_stiffness.append(pred_stiffness.cpu().numpy()) all_true_stiffness.append(true_stiffness.cpu().numpy()) all_pred_states = np.concatenate(all_pred_states) all_true_states = np.concatenate(all_true_states) all_pred_stiffness = np.concatenate(all_pred_stiffness) all_true_stiffness = np.concatenate(all_true_stiffness) test_loss = running_loss / len(test_loader.dataset) test_stiffness_loss = running_stiffness_loss / len(test_loader.dataset) test_states_loss = running_states_loss / len(test_loader.dataset) r2_states = r2_score(all_true_states.flatten(), all_pred_states.flatten()) r2_stiffness = r2_score(all_true_stiffness.flatten(), all_pred_stiffness.flatten()) results[noise_level] = { 'loss': test_loss, 'stiffness_loss': test_stiffness_loss, 'states_loss': test_states_loss, 'r2_states': r2_states, 'r2_stiffness': r2_stiffness } print(f"Noise Level {noise_level}:") print(f" Total Loss = {test_loss:.4f}") print(f" Stiffness Loss = {test_stiffness_loss:.4f}, R² = {r2_stiffness:.4f}") print(f" States Loss = {test_states_loss:.4f}, R² = {r2_states:.4f}") return results # NEW PLOTTING FUNCTIONS FOR DIRECT MODEL def plot_direct_training_history(history, title_prefix=''): plt.figure(figsize=(20, 10)) plt.subplot(1, 2, 1) plt.plot(history.history['train_loss'], label='Train') plt.plot(history.history['val_loss'], label='Val') plt.title(f'{title_prefix} Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.ylim(0, 0.12) # Set y-axis limits from 0 to 0.12 plt.legend() plt.subplot(1, 2, 2) plt.plot(history.history['train_r2'], label='Train') plt.plot(history.history['val_r2'], label='Val') plt.title(f'{title_prefix} R²') plt.xlabel('Epoch') plt.ylabel('R²') plt.ylim(0, 1) # Set y-axis limits from 0 to 1 plt.legend() plt.tight_layout() plt.show() # NEW FUNCTION TO PLOT DIRECT MODEL ROBUSTNESS RESULTS def plot_direct_robustness_results(results): noise_levels = list(results.keys()) r2_values = [results[nl]['r2'] for nl in noise_levels] loss_values = [results[nl]['loss'] for nl in noise_levels] plt.figure(figsize=(12, 5)) plt.subplot(1, 2, 1) plt.plot(noise_levels, r2_values, 'o-', linewidth=2) plt.xlabel('Noise Level') plt.ylabel('R² Score') plt.title('Element States R² vs Noise Level') plt.grid(True) plt.subplot(1, 2, 2) plt.plot(noise_levels, loss_values, 'o-', linewidth=2, color='red') plt.xlabel('Noise Level') plt.ylabel('Loss') plt.title('Element States Loss vs Noise Level') plt.grid(True) plt.tight_layout() plt.show() # Keep the original robustness plotting function for comparison def plot_robustness_results(results): noise_levels = list(results.keys()) r2_states_values = [results[nl]['r2_states'] for nl in noise_levels] r2_stiffness_values = [results[nl]['r2_stiffness'] for nl in noise_levels] states_loss_values = [results[nl]['states_loss'] for nl in noise_levels] stiffness_loss_values = [results[nl]['stiffness_loss'] for nl in noise_levels] plt.figure(figsize=(15, 10)) plt.subplot(2, 2, 1) plt.plot(noise_levels, r2_states_values, 'o-', linewidth=2) plt.xlabel('Noise Level') plt.ylabel('R² Score') plt.title('Element States R² vs Noise Level') plt.grid(True) plt.subplot(2, 2, 2) plt.plot(noise_levels, r2_stiffness_values, 'o-', linewidth=2, color='green') plt.xlabel('Noise Level') plt.ylabel('R² Score') plt.title('Stiffness Matrix R² vs Noise Level') plt.grid(True) plt.subplot(2, 2, 3) plt.plot(noise_levels, states_loss_values, 'o-', linewidth=2, color='red') plt.xlabel('Noise Level') plt.ylabel('Loss') plt.title('Element States Loss vs Noise Level') plt.grid(True) plt.subplot(2, 2, 4) plt.plot(noise_levels, stiffness_loss_values, 'o-', linewidth=2, color='orange') plt.xlabel('Noise Level') plt.ylabel('Loss') plt.title('Stiffness Matrix Loss vs Noise Level') plt.grid(True) plt.tight_layout() plt.show() # NEW VISUALIZATION FUNCTION FOR DIRECT MODEL def visualize_direct_prediction(model, test_loader, device='cuda', num_samples=4): model.eval() with torch.no_grad(): all_pred_states = [] all_true_states = [] for sensor_data, true_states in test_loader: sensor_data = sensor_data.to(device) pred_states = model(sensor_data) all_pred_states.append(pred_states.cpu().numpy()) all_true_states.append(true_states.numpy()) if len(all_pred_states) * sensor_data.size(0) >= num_samples: break all_pred_states = np.concatenate(all_pred_states) all_true_states = np.concatenate(all_true_states) # Visualization for element states plt.figure(figsize=(20, 15)) for i in range(min(num_samples, len(all_pred_states))): # Scatter plot plt.subplot(3, num_samples, i + 1) plt.scatter(all_true_states[i], all_pred_states[i], alpha=0.6) min_val = min(all_true_states[i].min(), all_pred_states[i].min()) max_val = max(all_true_states[i].max(), all_pred_states[i].max()) plt.plot([min_val, max_val], [min_val, max_val], 'r--', alpha=0.5) plt.xlabel('True States') plt.ylabel('Predicted States') plt.title(f'Sample {i+1} States') r2 = r2_score(all_true_states[i], all_pred_states[i]) plt.text(0.05, 0.95, f'R² = {r2:.3f}', transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8)) # Bar plot comparison plt.subplot(3, num_samples, i + num_samples + 1) x = np.arange(len(all_true_states[i])) width = 0.35 plt.bar(x - width/2, all_true_states[i], width, label='True', alpha=0.6) plt.bar(x + width/2, all_pred_states[i], width, label='Predicted', alpha=0.6) plt.xlabel('Element State Index') plt.ylabel('State Value') plt.title(f'Sample {i+1} Comparison') plt.legend() if len(x) > 20: plt.xticks(x[::5], rotation=0) # Show every 5th tick elif len(x) > 10: plt.xticks(x[::2], rotation=0) # Show every 2nd tick else: plt.xticks(x, rotation=0) # Error plot plt.subplot(3, num_samples, i + 2*num_samples + 1) errors = all_pred_states[i] - all_true_states[i] plt.bar(x, errors, alpha=0.6, color='red' if np.any(errors < 0) else 'green') plt.axhline(y=0, color='k', linestyle='-', alpha=0.3) plt.xlabel('Element State Index') plt.ylabel('Prediction Error') plt.title(f'Sample {i+1} Error') if len(x) > 20: plt.xticks(x[::5], rotation=0) # Show every 5th tick elif len(x) > 10: plt.xticks(x[::2], rotation=0) # Show every 2nd tick else: plt.xticks(x, rotation=0) plt.tight_layout() plt.suptitle('Element States Prediction Analysis', y=1.02) plt.show() def extract_important_sensors(flattened_accelerations, num_nodes=12, num_dofs=3): """ Extract only the top 7 most important sensors from flattened acceleration data. Parameters: ----------- flattened_accelerations : numpy.ndarray Flattened acceleration data with shape (scenarios, timesteps, nodes*dofs) num_nodes : int Number of nodes in the model num_dofs : int Number of degrees of freedom per node Returns: -------- important_sensors : numpy.ndarray for example 5: Acceleration data with only the important sensors, shape (scenarios, timesteps, 5) """ # Define the important sensors (node, dof) # X is index 0 in the DOF dimension (after removing time). its (X,Y,θ) important_sensors_indices = [ # ## Node 1 sensors (1-1, 0), # Node 1 - X (0-indexed) #48% @ solo-run @ 110k samples # (1-1, 1), # Node 1 - Y #38% # (1-1, 2), # Node 1 - θ #36% # # Node 2 sensors (2-1, 0), # Node 2 - X #46% # (2-1, 1), # Node 2 - Y #36% # (2-1, 2), # Node 2 - θ #36% # # Node 3 sensors (3-1, 0), # Node 3 - X #46% # (3-1, 1), # Node 3 - Y #35% # (3-1, 2), # Node 3 - θ #36% # # Node 4 sensors (4-1, 0), # Node 4 - X #48% # (4-1, 1), # Node 4 - Y # (4-1, 2), # Node 4 - θ # # Node 5 sensors (5-1, 0), # Node 5 - X #51% # (5-1, 1), # Node 5 - Y # (5-1, 2), # Node 5 - θ # # Node 6 sensors (6-1, 0), # Node 6 - X #49% # (6-1, 1), # Node 6 - Y # (6-1, 2), # Node 6 - θ # # Node 7 sensors (7-1, 0), # Node 7 - X #48% # (7-1, 1), # Node 7 - Y # (7-1, 2), # Node 7 - θ # # Node 8 sensors (8-1, 0), # Node 8 - X #51% # (8-1, 1), # Node 8 - Y # (8-1, 2), # Node 8 - θ # # Node 9 sensors # (9-1, 0), # Node 9 - X #45% # (9-1, 1), # Node 9 - Y # (9-1, 2), # Node 9 - θ # # Node 10 sensors # (10-1, 0), # Node 10 - X #45% # (10-1, 1), # Node 10 - Y # (10-1, 2), # Node 10 - θ # # Node 11 sensors # (11-1, 0), # Node 11 - X #45% # (11-1, 1), # Node 11 - Y # (11-1, 2), # Node 11 - θ # # Node 12 sensors # (12-1, 0), # Node 12 - X #45% # (12-1, 1), # Node 12 - Y #39% # (12-1, 2), # Node 12 - θ #29% #@ 9 sensors ~90% #@ 8 sensors ~85% #@ 2 sensors with 0.05 noise >> ~55% #@ 6 sensors with 0.14 noise >> ~65-68% ] # Convert to flattened indices flat_indices = [node * num_dofs + dof for node, dof in important_sensors_indices] # Extract only the important sensors scenarios, timesteps, _ = flattened_accelerations.shape important_sensors = np.zeros((scenarios, timesteps, len(flat_indices))) for i, idx in enumerate(flat_indices): important_sensors[:, :, i] = flattened_accelerations[:, :, idx] return important_sensors def plot_noise_degradation_curve(model, val_loader, device='cuda', noise_levels=[0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3]): """ Plot model performance degradation across different noise levels. """ model.eval() results = {level: {'r2': [], 'loss': []} for level in noise_levels} criterion = nn.MSELoss() with torch.no_grad(): for noise_level in noise_levels: all_pred_states = [] all_true_states = [] running_loss = 0.0 # Evaluate on the validation set with this noise level for sensor_data, true_states in val_loader: # Apply noise at the current level sensor_data_np = sensor_data.numpy() # Apply different types of noise for a more comprehensive test # Or just stick with Gaussian noise for simplicity # sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level) # sensor_data_np = add_gaussian_noise_degradation_eval(sensor_data_np, noise_level, acc_norm_factors) sensor_data_np = add_gaussian_noise(sensor_data_np, noise_level, acc_norm_factors) sensor_data = torch.FloatTensor(sensor_data_np) sensor_data = sensor_data.to(device) true_states = true_states.to(device) pred_states = model(sensor_data) loss = criterion(pred_states, true_states) running_loss += loss.item() * sensor_data.size(0) all_pred_states.append(pred_states.cpu().numpy()) all_true_states.append(true_states.cpu().numpy()) # Calculate metrics all_pred_states = np.concatenate(all_pred_states) all_true_states = np.concatenate(all_true_states) test_loss = running_loss / len(val_loader.dataset) r2 = r2_score(all_true_states.flatten(), all_pred_states.flatten()) results[noise_level]['r2'] = r2 results[noise_level]['loss'] = test_loss print(f"Noise Level {noise_level}: R² = {r2:.4f}, Loss = {test_loss:.4f}") # Plot the degradation curve plt.figure(figsize=(12, 8)) # Plot R² degradation plt.subplot(2, 1, 1) plt.plot(noise_levels, [results[nl]['r2'] for nl in noise_levels], 'o-', linewidth=2) plt.xlabel('Noise Level') plt.ylabel('R² Score') plt.title('Model Performance Degradation with Increasing Noise') plt.grid(True) # Plot Loss increase plt.subplot(2, 1, 2) plt.plot(noise_levels, [results[nl]['loss'] for nl in noise_levels], 'o-', linewidth=2, color='red') plt.xlabel('Noise Level') plt.ylabel('Loss') plt.title('Loss Increase with Increasing Noise') plt.grid(True) plt.tight_layout() plt.show() return results # MODIFIED MAIN FUNCTION TO INCLUDE DIRECT MODEL def main(): # Set device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # Extract sensor data and preprocess acc_Sensors_ALL, acc_norm_factors = extract_sensor_arrays(sensorDataTotal) # Extract only important sensors print("Original sensor shape:", acc_Sensors_ALL.shape) acc_Sensors_IMPORTANT = extract_important_sensors(acc_Sensors_ALL, num_nodes=12, num_dofs=3) print("Reduced sensor shape:", acc_Sensors_IMPORTANT.shape) # Create datasets for both approaches direct_dataset = DirectSensorDataset(acc_Sensors_IMPORTANT, elementStatesTotal) combined_dataset = CombinedDataset(acc_Sensors_IMPORTANT, stiffnessMatrixToatal, elementStatesTotal) # Split data for training train_size = int(0.7 * len(direct_dataset)) val_size = int(0.15 * len(direct_dataset)) test_size = len(direct_dataset) - train_size - val_size # Split the direct dataset train_direct, val_direct, test_direct = random_split( direct_dataset, [train_size, val_size, test_size], generator=torch.Generator().manual_seed(42) ) # Split the combined dataset (for comparison) train_combined, val_combined, test_combined = random_split( combined_dataset, [train_size, val_size, test_size], generator=torch.Generator().manual_seed(42) ) # Create data loaders batch_size = 64 train_direct_loader = DataLoader(train_direct, batch_size=batch_size, shuffle=True, drop_last=True) val_direct_loader = DataLoader(val_direct, batch_size=batch_size, drop_last=True) test_direct_loader = DataLoader(test_direct, batch_size=batch_size, drop_last=True) train_combined_loader = DataLoader(train_combined, batch_size=batch_size, shuffle=True, drop_last=True) val_combined_loader = DataLoader(val_combined, batch_size=batch_size, drop_last=True) test_combined_loader = DataLoader(test_combined, batch_size=batch_size, drop_last=True) # Create and train the direct sensor to element model print("Training Direct Sensor to Element Model...") # direct_sensor_model = DirectSensorToElementModel( # input_features=acc_Sensors_IMPORTANT.shape[2], # output_size=21 # Number of element states # ) direct_sensor_model = DirectSensorToElementModel( input_features=acc_Sensors_IMPORTANT.shape[2], output_size=21, # Number of element states timesteps=50 # Number of time steps in your data ) best_direct_model_state, direct_history, max_noise_level = train_direct_sensor_model( direct_sensor_model, train_direct_loader, val_direct_loader, num_epochs=4, device=device ) # Load best model direct_sensor_model.load_state_dict(best_direct_model_state) degradation_results = plot_noise_degradation_curve(direct_sensor_model, val_direct_loader) # Evaluate direct model robustness print("Evaluating direct model robustness across noise levels...") direct_robustness_results = evaluate_direct_sensor_model_robustness( direct_sensor_model, test_direct_loader, device=device, noise_levels=[0, 0.05, 0.1, 0.15, 0.2, 0.25] ) # Plot direct model results plot_direct_training_history(direct_history, title_prefix='Direct Sensor to Element Model') plot_direct_robustness_results(direct_robustness_results) visualize_direct_prediction(direct_sensor_model, test_direct_loader, device=device) # Save the direct model torch.save(direct_sensor_model.state_dict(), 'direct_sensor_to_element_model.pth') print("Direct model saved as 'direct_sensor_to_element_model.pth'") # *************************************************************** # *** ADD THIS SECTION TO SAVE VALIDATION R² SCORES *** # *************************************************************** print("\nSaving Direct Model Validation R² Scores...") try: # Extract the validation R2 scores from the history object # Use .get() with an empty list default for safety in case 'val_r2' wasn't populated val_r2_scores = direct_history.history.get('val_r2', []) if val_r2_scores: # Check if the list is not empty # Construct the filename using the PREVIOUSLY DEFINED max noise level variable r2_filename = f"Direct_under_{max_noise_level}_noise_val_r2_scores.txt" # Save the scores using numpy.savetxt np.savetxt( r2_filename, np.array(val_r2_scores), fmt='%.6f', # Format scores to 6 decimal places header=f'Direct Model Validation R2 Scores (trained with max_noise_level={max_noise_level})', comments='' # Remove the default '#' comment prefix from numpy output ) print(f"--> Successfully saved validation R² scores to '{r2_filename}'") else: print("--> Warning: No validation R² scores found in history object to save.") except AttributeError: print("--> Error: The 'direct_history' object does not seem to have a 'history' attribute or 'val_r2' key.") except Exception as e: print(f"--> Error saving Direct model validation R² scores: {e}") # *************************************************************** # *** END OF ADDED SECTION *** # *************************************************************** # Compare the performance of both models print("\nModel Performance Comparison:") print("Direct Sensor to Element Model:") print(f" R² at 0% noise: {direct_robustness_results[0]['r2']:.4f}") print(f" R² at 10% noise: {direct_robustness_results[0.1]['r2']:.4f}") print(f" R² at 20% noise: {direct_robustness_results[0.2]['r2']:.4f}") return direct_sensor_model,direct_robustness_results, if __name__ == "__main__": direct_model,direct_results = main()