Untitled
unknown
plain_text
10 months ago
34 kB
18
Indexable
# --- Core Imports ---
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, TensorDataset
import numpy as np
import pandas as pd
import random
import os
import json
import time
import joblib
from scipy import interpolate
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.linear_model import Ridge
from sklearn.kernel_ridge import KernelRidge
from sklearn.preprocessing import StandardScaler, PolynomialFeatures
from sklearn.pipeline import Pipeline
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
from tqdm import tqdm
# --- General Setup ---
RANDOM_SEED = 42 # Use a fixed seed for random numbers to make experiments repeatable.
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Use GPU if available, otherwise use CPU.
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) if '__file__' in locals() else os.getcwd() # Get the directory where the script is located.
# Set random seeds for all relevant libraries for reproducibility.
torch.manual_seed(RANDOM_SEED)
np.random.seed(RANDOM_SEED)
random.seed(RANDOM_SEED)
if torch.cuda.is_available():
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# --- Path Configuration ---
# Define output folders for each training pipeline.
SENSOR_DATA_DIR = SCRIPT_DIR
OUTPUT_DIR_SHIFTING = os.path.join(SCRIPT_DIR, "output_shifting_peak_model")
OUTPUT_DIR_FWHM = os.path.join(SCRIPT_DIR, "output_varying_fwhm_model")
OUTPUT_DIR_MERGING = os.path.join(SCRIPT_DIR, "output_merging_peaks_model")
# Define paths to the sensor response curve CSV files.
SENSOR_FILES = {
name: os.path.join(SENSOR_DATA_DIR, filename)
for name, filename in {
'Blue': 'BlueSensor.csv', 'Clear': 'ClearSensor.csv', 'Green': 'GreenSensor.csv',
'IR': 'IrSensor.csv', 'Red': 'RedSensor.csv', 'Visible': 'VisibleSensor.csv'
}.items()
}
# --- Shared Model & Training Parameters ---
# Neural Network architecture and training settings.
NN_HIDDEN_LAYERS = [256, 512, 256] # Number of neurons in each hidden layer.
NN_LEARNING_RATE = 5e-4
NN_DROPOUT_RATE = 0.2
NN_WEIGHT_DECAY = 1e-5
NN_EPOCHS = 200 # Maximum number of training cycles.
NN_BATCH_SIZE = 64
NN_SCHEDULER_PATIENCE = 5 # How many epochs to wait for improvement before reducing learning rate.
NN_SCHEDULER_FACTOR = 0.5
NN_EARLY_STOPPING_PATIENCE = 10 # How many epochs to wait for improvement before stopping training.
# --- Hyperparameter Tuning Configuration for Baseline Models ---
CV_FOLDS = 3 # Number of cross-validation folds for GridSearchCV.
# Define parameter ranges to test for the Ridge models.
LINEAR_RIDGE_PARAM_GRID = {
'ridge__alpha': [1e-3, 1e-2, 1e-1, 1, 10, 100],
}
RIDGE_PARAM_GRID = {
'poly__degree': [2, 3],
'ridge__alpha': [1e-2, 1e-1, 1, 10, 100],
}
KRR_PARAM_GRID = {
'kernel_ridge__kernel': ['rbf'],
'kernel_ridge__alpha': [1e-2, 1e-1, 1.0, 10.0],
'kernel_ridge__gamma': np.logspace(-2, 2, 5),
}
# --- Data Augmentation Parameters ---
ADD_NOISE_TO_TRAIN_DATA = False # Set to True to add random noise to sensor data during training.
NOISE_LEVEL = 0.03 # The amount of noise to add, as a fraction of the max sensor reading.
# --- Shared Data & Evaluation Parameters ---
COMMON_WL_RANGE = (400, 700) # Wavelength range in nanometers.
N_POINTS_RESAMPLE = int(COMMON_WL_RANGE[1] - COMMON_WL_RANGE[0]) + 1 # Number of points for a 1nm grid.
VAL_RATIO = 0.2 # 20% of the data will be used for validation.
NUM_EXAMPLES_PLOT = 5 # Number of random validation examples to plot.
# --- Parameter for incremental testing ---
INCREMENTAL_STEP = 1.0 # The step size for testing how models perform as parameters change.
# --- Pipeline 1: Shifting Peak Configuration ---
SHIFTING_NUM_SAMPLES = 50000
SHIFTING_CENTER_WL_RANGE = (400, 700)
SHIFTING_FWHM = 26
SHIFTING_AMP = 30000.0
# --- Pipeline 2: Varying FWHM Configuration ---
FWHM_NUM_SAMPLES = 50000
FWHM_CENTER_WL = 525
FWHM_RANGE = (1, 100)
FWHM_AMP_RANGE = (30000.0, 30000.0)
# --- Pipeline 3: Merging Peaks Configuration ---
MERGING_NUM_SAMPLES = 50000
MERGING_PEAK1_WL = 400
MERGING_PEAK2_WL = 700
MERGING_FWHM = 18
MERGING_AMP = 30000.0
# Converts Full-Width at Half-Maximum (FWHM) to standard deviation (sigma).
def fwhm_to_sigma(fwhm):
return fwhm / (2 * np.sqrt(2 * np.log(2)))
# Creates a directory if it doesn't already exist.
def ensure_output_dir(path):
if not os.path.exists(path):
os.makedirs(path)
print(f"Created directory: {path}")
# Creates a small example CSV (1 sample) from a larger one for easy viewing.
def save_baby_csv(full_csv_path):
try:
df = pd.read_csv(full_csv_path)
if 'sample_index' in df.columns:
baby_df = df[df['sample_index'] == 0].copy()
else:
baby_df = df.head(1).copy()
if not baby_df.empty:
baby_csv_path = full_csv_path.replace('.csv', '_baby_example.csv')
baby_df.to_csv(baby_csv_path, index=False, float_format='%.6f')
except Exception as e:
print(f" -> Could not generate baby CSV for {os.path.basename(full_csv_path)}. Reason: {e}")
# Defines the Neural Network architecture.
class SpectrumNet(nn.Module):
def __init__(self, input_dim, output_dim, hidden_dims, dropout_rate):
super(SpectrumNet, self).__init__()
layers = []
prev_dim = input_dim
# Create the network layer by layer.
for dim in hidden_dims:
layers.append(nn.Linear(prev_dim, dim)) # Fully connected layer
layers.append(nn.ReLU()) # Activation function
layers.append(nn.BatchNorm1d(dim)) # Normalize layer outputs
layers.append(nn.Dropout(dropout_rate)) # Prevent overfitting
prev_dim = dim
layers.append(nn.Linear(prev_dim, output_dim)) # Final output layer
layers.append(nn.ReLU()) # Ensure output is non-negative
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)
# Loads the sensor response curve data from CSV files.
def load_sensor_data(sensor_files_dict):
sensors = {}
print("Loading sensor response curves...")
for name, file_path in sensor_files_dict.items():
try:
data = pd.read_csv(file_path, header=None)
sensors[name] = {'wavelength': data.iloc[:, 0].values, 'response': data.iloc[:, 1].values}
print(f" -> Loaded '{name}' sensor from {os.path.basename(file_path)}")
except FileNotFoundError:
print(f" -> ERROR: Sensor file not found at {file_path}")
return None
except Exception as e:
print(f" -> ERROR: Could not read {file_path}. Reason: {e}")
return None
return sensors
# Simulates the output of each sensor for a given light spectrum.
def calculate_sensor_outputs(wavelengths, emission, sensors_dict):
sensor_outputs = []
# Ensure consistent order of sensors.
for name in sorted(sensors_dict.keys()):
sensor_data = sensors_dict[name]
# Interpolate the sensor's response to match the spectrum's wavelengths.
f = interpolate.interp1d(
sensor_data['wavelength'],
sensor_data['response'],
bounds_error=False,
fill_value=0.0
)
sensor_response_at_spectrum_wl = f(wavelengths)
# Integrate the product of the spectrum and the sensor response to get the sensor's reading.
output = np.trapz(emission * sensor_response_at_spectrum_wl, wavelengths)
sensor_outputs.append(output)
return np.array(sensor_outputs)
# Trains the neural network model.
def train_nn_model(model, criterion, optimizer, scheduler, train_loader, val_loader, output_dir):
print("Training neural network model...")
start_time = time.time()
best_val_loss = float('inf')
epochs_no_improve = 0
history = {'train_loss': [], 'val_loss': []}
best_model_path = os.path.join(output_dir, 'best_nn_model.pth')
# Loop through all epochs.
for epoch in range(NN_EPOCHS):
model.train() # Set the model to training mode.
train_loss = 0.0
# Process data in batches.
for inputs, targets in tqdm(train_loader, desc=f"Epoch {epoch+1}/{NN_EPOCHS} [Train]", leave=False):
optimizer.zero_grad() # Reset gradients.
outputs = model(inputs) # Forward pass: get predictions.
loss = criterion(outputs, targets) # Calculate the loss.
loss.backward() # Backward pass: calculate gradients.
optimizer.step() # Update model weights.
train_loss += loss.item() * inputs.size(0)
history['train_loss'].append(train_loss / len(train_loader.dataset))
# Evaluate the model on the validation set after each epoch.
model.eval() # Set the model to evaluation mode.
val_loss = 0.0
with torch.no_grad(): # No need to calculate gradients during validation.
for inputs, targets in tqdm(val_loader, desc=f"Epoch {epoch+1}/{NN_EPOCHS} [Val]", leave=False):
outputs = model(inputs)
loss = criterion(outputs, targets)
val_loss += loss.item() * inputs.size(0)
val_loss /= len(val_loader.dataset)
history['val_loss'].append(val_loss)
scheduler.step(val_loss) # Adjust learning rate based on validation loss.
if (epoch + 1) % 5 == 0:
print(f"Epoch {epoch+1}/{NN_EPOCHS} - Train Loss: {history['train_loss'][-1]:.6e}, Val Loss: {val_loss:.6e}")
# Early stopping logic.
if val_loss < best_val_loss:
best_val_loss = val_loss
epochs_no_improve = 0
torch.save(model.state_dict(), best_model_path) # Save the best model.
else:
epochs_no_improve += 1
if epochs_no_improve >= NN_EARLY_STOPPING_PATIENCE:
print(f"Early stopping triggered after {epoch+1} epochs.")
break
model.load_state_dict(torch.load(best_model_path)) # Load the best model weights.
print(f"NN training complete. Best model saved to {best_model_path}")
print(f"Training time: {time.time() - start_time:.2f} seconds")
return model, history
# Trains the scikit-learn baseline models.
def train_baseline_models(X_train, y_train, output_dir):
print("🚀 Optimizing and training baseline models with GridSearchCV...")
# --- Linear Ridge Regression Optimization ---
print("\n -> Searching for optimal Linear Ridge model...")
linear_ridge_pipeline = Pipeline([
('scaler', StandardScaler()),
('ridge', Ridge())
])
# Use GridSearchCV to find the best 'alpha' regularization parameter.
linear_ridge_grid_search = GridSearchCV(
linear_ridge_pipeline,
param_grid=LINEAR_RIDGE_PARAM_GRID,
cv=CV_FOLDS,
scoring='neg_mean_squared_error',
n_jobs=1,
verbose=1
)
linear_ridge_grid_search.fit(X_train, y_train)
best_linear_ridge_model = linear_ridge_grid_search.best_estimator_
print(f" Optimal Linear Ridge params: {linear_ridge_grid_search.best_params_}")
joblib.dump(best_linear_ridge_model, os.path.join(output_dir, 'linear_ridge_model.joblib'))
# --- Polynomial Ridge Regression Optimization ---
print("\n -> Searching for optimal Polynomial Ridge model...")
ridge_pipeline = Pipeline([
('scaler', StandardScaler()),
('poly', PolynomialFeatures(include_bias=False)), # Create polynomial features.
('ridge', Ridge())
])
# Use GridSearchCV to find the best polynomial degree and alpha.
ridge_grid_search = GridSearchCV(
ridge_pipeline,
param_grid=RIDGE_PARAM_GRID,
cv=CV_FOLDS,
scoring='neg_mean_squared_error',
n_jobs=1,
verbose=1
)
ridge_grid_search.fit(X_train, y_train)
best_ridge_model = ridge_grid_search.best_estimator_
print(f" Optimal Polynomial Ridge params: {ridge_grid_search.best_params_}")
joblib.dump(best_ridge_model, os.path.join(output_dir, 'ridge_model.joblib'))
# --- Kernel Ridge Regression Optimization with Sub-sampling ---
print("\n -> Searching for optimal Kernel Ridge model...")
# Use a smaller subset of data for tuning KRR because it's very slow.
n_samples_for_tuning = 2000
if len(X_train) > n_samples_for_tuning:
print(f" -> Using a random subset of {n_samples_for_tuning} samples for KRR tuning to conserve memory.")
subset_indices = np.random.choice(X_train.shape[0], n_samples_for_tuning, replace=False)
X_train_subset = X_train[subset_indices]
y_train_subset = y_train[subset_indices]
else:
X_train_subset, y_train_subset = X_train, y_train
krr_pipeline = Pipeline([
('scaler', StandardScaler()),
('kernel_ridge', KernelRidge())
])
# Run GridSearchCV on the smaller data subset to find best kernel parameters.
krr_grid_search = GridSearchCV(
krr_pipeline,
param_grid=KRR_PARAM_GRID,
cv=CV_FOLDS,
scoring='neg_mean_squared_error',
n_jobs=1,
verbose=1
)
krr_grid_search.fit(X_train_subset, y_train_subset)
print(f" Optimal Kernel Ridge params found: {krr_grid_search.best_params_}")
# Save the best model found on the subset.
best_krr_model = krr_grid_search.best_estimator_
joblib.dump(best_krr_model, os.path.join(output_dir, 'kernel_ridge_model.joblib'))
print("\nOptimized baseline models trained and saved.")
return best_linear_ridge_model, best_ridge_model, best_krr_model
# Saves a detailed CSV report for a single example.
def save_single_example_report(idx, wl_grid, sensor_names, sensor_inputs, y_true, predictions_dict, report_dir):
df_data = {
'wavelength_nm': wl_grid,
'ground_truth': y_true
}
for model_name, pred in predictions_dict.items():
col_name = model_name.lower().replace(" ", "_").replace("-", "_")
df_data[f'{col_name}_prediction'] = pred
report_df = pd.DataFrame(df_data)
# Calculate and add all performance metrics to the report.
all_metrics = {}
for model_name, pred in predictions_dict.items():
model_metrics = calculate_and_compile_metrics(y_true, pred, model_name.lower().replace(" ", "_"))
all_metrics.update(model_metrics)
# Add the normalized sensor inputs to the report.
sensor_data = {f'norm_sensor_{name}': sensor_inputs[i] for i, name in enumerate(sensor_names)}
summary_data = {**all_metrics, **sensor_data}
for key, value in summary_data.items():
report_df[key] = value
file_path = os.path.join(report_dir, f"report_sample_{idx}.csv")
report_df.to_csv(file_path, index=False, float_format='%.8f')
# Evaluates all models on the validation set and creates plots.
def evaluate_and_plot_results(models, data, wavelength_grid, output_dir, sensor_names):
print("\n--- Evaluating Models on Validation Set ---")
nn_model, linear_ridge_model, ridge_model, kernel_ridge_model = models
X_val, y_val, val_loader = data
# Get predictions from the Neural Network.
nn_model.eval()
nn_preds_list = []
with torch.no_grad():
for inputs, _ in tqdm(val_loader, desc="Evaluating NN", leave=False):
outputs = nn_model(inputs)
nn_preds_list.append(outputs.cpu().numpy())
nn_preds = np.vstack(nn_preds_list)
# Get predictions from the baseline models.
print("Predicting with optimized baseline models...")
linear_ridge_preds = linear_ridge_model.predict(X_val)
ridge_preds = ridge_model.predict(X_val)
kr_preds = kernel_ridge_model.predict(X_val)
# Calculate overall performance metrics.
metrics = {
"nn_mse": mean_squared_error(y_val, nn_preds),
"linear_ridge_mse": mean_squared_error(y_val, linear_ridge_preds),
"ridge_mse": mean_squared_error(y_val, ridge_preds),
"kr_mse": mean_squared_error(y_val, kr_preds),
}
print("Overall Model Performance (MSE):")
print(f" NN Model: {metrics['nn_mse']:.6f}")
print(f" Linear Ridge Model: {metrics['linear_ridge_mse']:.6f}")
print(f" Polynomial Ridge Model: {metrics['ridge_mse']:.6f}")
print(f" Kernel Ridge Model: {metrics['kr_mse']:.6f}")
# Save metrics to a JSON file.
with open(os.path.join(output_dir, 'evaluation_metrics.json'), 'w') as f:
json.dump(metrics, f, indent=4)
print("Generating random example plots and detailed CSV reports...")
indices_to_plot = random.sample(range(len(X_val)), min(NUM_EXAMPLES_PLOT, len(X_val)))
report_dir = os.path.join(output_dir, "random_example_reports")
ensure_output_dir(report_dir)
# Create plots and reports for a few random examples.
for idx in tqdm(indices_to_plot, desc="Generating Reports"):
predictions_dict = {
"NN": nn_preds[idx],
"Linear Ridge": linear_ridge_preds[idx],
"Polynomial Ridge": ridge_preds[idx],
"Kernel Ridge": kr_preds[idx]
}
plot_example(idx, wavelength_grid, y_val[idx], predictions_dict, output_dir, f"Random Validation Sample {idx}")
save_single_example_report(idx, wavelength_grid, sensor_names, X_val[idx], y_val[idx], predictions_dict, report_dir)
# Plots a single reconstruction example.
def plot_example(idx, wl, y_true, predictions_dict, output_dir, title):
plt.figure(figsize=(12, 8))
colors = ['#9f86c0', '#00b4d8', '#ffbf00', '#57cc99', '#720026']
styles = ['--', '-.', ':', (0, (3, 1, 1, 1))]
# Plot the ground truth spectrum.
plt.plot(wl, y_true, color=colors[0], linestyle='-', linewidth=2.5, label='True Spectrum')
# Plot each model's prediction.
for i, (name, pred) in enumerate(predictions_dict.items()):
mse = mean_squared_error(y_true, pred)
label_text = f"{name} (MSE: {mse:.6f})"
plt.plot(wl, pred, color=colors[i+1], linestyle=styles[i], linewidth=2, label=label_text)
plt.title(f'Spectrum Reconstruction: {title}')
plt.xlabel('Wavelength (nm)')
plt.ylabel('Normalized Intensity')
plt.ylim(0, 1.1)
plt.grid(True, alpha=0.4)
plt.legend(loc='upper right', fontsize='medium')
plt.tight_layout()
plt.savefig(os.path.join(output_dir, f"example_{idx}.png"))
plt.close()
# Saves the validation dataset (inputs and outputs) to a CSV for later analysis.
def save_sensor_and_truth_data_to_csv(X_raw, X_normalized, y_true, params, sensor_names, wl_grid, output_path):
df_params = pd.DataFrame(params)
raw_sensor_cols = [f"raw_sensor_{name}" for name in sensor_names]
norm_sensor_cols = [f"norm_sensor_{name}" for name in sensor_names]
wl_cols = [f"wl_{wl:.2f}" for wl in wl_grid]
df_raw_sensors = pd.DataFrame(X_raw, columns=raw_sensor_cols)
df_norm_sensors = pd.DataFrame(X_normalized, columns=norm_sensor_cols)
df_truth = pd.DataFrame(y_true, columns=wl_cols)
combined_df = pd.concat([df_params, df_raw_sensors, df_norm_sensors, df_truth], axis=1)
combined_df.insert(0, 'sample_index', combined_df.index)
combined_df.to_csv(output_path, index=False)
print(f"-> Saved validation data to {os.path.basename(output_path)}")
save_baby_csv(output_path) # Also save a small version.
# Generates a single Gaussian peak spectrum.
def generate_gaussian_spectrum(wavelengths, mu, sigma, amp):
return amp * np.exp(-((wavelengths - mu) ** 2) / (2 * sigma ** 2))
# Computes sensor outputs and the normalized spectrum shape from a raw spectrum.
def compute_raw_sensor_outputs_and_spectrum(raw_spectrum, common_wl_grid, sensors):
# Simulate sensor readings.
raw_sensor_outputs = calculate_sensor_outputs(common_wl_grid, raw_spectrum, sensors)
# Normalize the spectrum shape to have a max value of 1.
max_spec_val = np.max(raw_spectrum)
y_spectrum_normalized = raw_spectrum / max_spec_val if max_spec_val > 1e-9 else raw_spectrum
return raw_sensor_outputs, y_spectrum_normalized
# Generates training data where a single Gaussian peak shifts its position.
def prepare_shifting_gaussian_data(sensors, common_wl_grid):
print(f"Generating {SHIFTING_NUM_SAMPLES} shifting Gaussian spectra...")
X_list, y_list, params_list = [], [], []
sigma = fwhm_to_sigma(SHIFTING_FWHM)
for _ in tqdm(range(SHIFTING_NUM_SAMPLES), desc="Generating Shifting Peak Data"):
mu = random.uniform(*SHIFTING_CENTER_WL_RANGE)
raw_spectrum = generate_gaussian_spectrum(common_wl_grid, mu, sigma, SHIFTING_AMP)
X_raw, y = compute_raw_sensor_outputs_and_spectrum(raw_spectrum, common_wl_grid, sensors)
X_list.append(X_raw)
y_list.append(y)
params_list.append({'parameter_name': 'center_wl', 'parameter_value': mu})
return np.array(X_list), np.array(y_list), np.array(params_list)
# Generates training data where a single peak changes its width (FWHM).
def prepare_varying_fwhm_data(sensors, common_wl_grid):
print(f"Generating {FWHM_NUM_SAMPLES} varying FWHM Gaussian spectra...")
X_list, y_list, params_list = [], [], []
for _ in tqdm(range(FWHM_NUM_SAMPLES), desc="Generating Varying FWHM Data"):
fwhm = random.uniform(*FWHM_RANGE)
sigma = fwhm_to_sigma(fwhm)
amp = random.uniform(*FWHM_AMP_RANGE)
raw_spectrum = generate_gaussian_spectrum(common_wl_grid, FWHM_CENTER_WL, sigma, amp)
X_raw, y = compute_raw_sensor_outputs_and_spectrum(raw_spectrum, common_wl_grid, sensors)
X_list.append(X_raw)
y_list.append(y)
params_list.append({'parameter_name': 'fwhm', 'parameter_value': fwhm})
return np.array(X_list), np.array(y_list), np.array(params_list)
# Generates training data where two peaks start far apart and move closer together.
def prepare_merging_gaussians_data(sensors, common_wl_grid):
print(f"Generating {MERGING_NUM_SAMPLES} merging Gaussian spectra...")
X_list, y_list, params_list = [], [], []
sigma = fwhm_to_sigma(MERGING_FWHM)
center_point = (MERGING_PEAK1_WL + MERGING_PEAK2_WL) / 2
for i in range(MERGING_NUM_SAMPLES):
separation = random.uniform(1.0, MERGING_PEAK2_WL - MERGING_PEAK1_WL)
current_half_dist = separation / 2.0
mu1, mu2 = center_point - current_half_dist, center_point + current_half_dist
peak1 = generate_gaussian_spectrum(common_wl_grid, mu1, sigma, MERGING_AMP)
peak2 = generate_gaussian_spectrum(common_wl_grid, mu2, sigma, MERGING_AMP)
# The final spectrum is the maximum envelope of the two peaks.
raw_spectrum = np.maximum(peak1, peak2)
X_raw, y = compute_raw_sensor_outputs_and_spectrum(raw_spectrum, common_wl_grid, sensors)
X_list.append(X_raw)
y_list.append(y)
params_list.append({'parameter_name': 'separation', 'parameter_value': separation})
return np.array(X_list), np.array(y_list), np.array(params_list)
# Adds Gaussian noise to sensor data to simulate real-world conditions.
def add_noise(sensor_data, noise_level):
noisy_data = np.copy(sensor_data).astype(np.float32)
for i in range(noisy_data.shape[0]):
max_val = np.max(noisy_data[i])
if max_val > 0:
noise_sigma = max_val * noise_level
noise = np.random.normal(0, noise_sigma, noisy_data[i].shape)
noisy_data[i] += noise
return noisy_data
# Normalizes each sensor reading vector by its own max value.
def normalize_sensor_data(sensor_data):
normalized_data = np.copy(sensor_data).astype(np.float32)
for i in range(normalized_data.shape[0]):
max_val = np.max(normalized_data[i])
# Avoid division by zero.
if max_val > 1e-9:
normalized_data[i] /= max_val
return normalized_data
# Calculates a dictionary of performance metrics for a single prediction.
def calculate_and_compile_metrics(y_true, y_pred, model_name):
y_true, y_pred = np.asarray(y_true).flatten(), np.asarray(y_pred).flatten()
metrics = {}
metrics[f'{model_name}_mse'] = mean_squared_error(y_true, y_pred)
metrics[f'{model_name}_mae'] = mean_absolute_error(y_true, y_pred)
metrics[f'{model_name}_r2'] = r2_score(y_true, y_pred)
metrics[f'{model_name}_cos_sim'] = cosine_similarity(y_true.reshape(1, -1), y_pred.reshape(1, -1))[0, 0]
return metrics
# Saves the sensor inputs used during the incremental test to a CSV file.
def save_incremental_sensor_data_csv(model_name, param_info_list, X_raw, X_norm, sensor_names, output_dir):
df_params = pd.DataFrame(param_info_list)
raw_sensor_cols = [f"raw_sensor_{name}" for name in sensor_names]
norm_sensor_cols = [f"norm_sensor_{name}" for name in sensor_names]
df_raw_sensors = pd.DataFrame(X_raw, columns=raw_sensor_cols)
df_norm_sensors = pd.DataFrame(X_norm, columns=norm_sensor_cols)
combined_df = pd.concat([df_params, df_raw_sensors, df_norm_sensors], axis=1)
output_path = os.path.join(output_dir, f'incremental_sensor_outputs_{model_name}.csv')
combined_df.to_csv(output_path, index=False, float_format='%.6f')
save_baby_csv(output_path)
# Saves the reconstruction results from an incremental test to a CSV file.
def save_model_reconstruction_csv(y_true_all, y_pred_model, model_name, wl_grid, param_info_list, output_dir):
wl_columns = [f"wl_{wl:.2f}" for wl in wl_grid]
df_true = pd.DataFrame(y_true_all, columns=wl_columns)
df_pred = pd.DataFrame(y_pred_model, columns=wl_columns)
df_params = pd.DataFrame(param_info_list)
# Add metadata columns.
for df in [df_true, df_pred]:
df['parameter_name'] = df_params['parameter_name']
df['parameter_value'] = df_params['parameter_value']
df_true['type'] = 'ground_truth'
df_pred['type'] = 'reconstruction'
combined_df = pd.concat([df_true, df_pred], ignore_index=True)
output_path = os.path.join(output_dir, f'incremental_reconstruction_{model_name}.csv')
combined_df.to_csv(output_path, index=False, float_format='%.6f')
save_baby_csv(output_path)
# Runs a deterministic test where one parameter is varied incrementally.
def evaluate_and_save_incremental_tests(models, pipeline_id, sensors, wavelength_grid, output_dir, sensor_names):
print(f"\n--- Evaluating Models on Incremental Test Cases (Step: {INCREMENTAL_STEP} nm) ---")
nn_model, linear_ridge_model, ridge_model, kernel_ridge_model = models
nn_model.eval()
test_case_dir = os.path.join(output_dir, "incremental_tests")
ensure_output_dir(test_case_dir)
y_true_list, nn_preds_list, linear_ridge_preds_list, ridge_preds_list, kr_preds_list = [], [], [], [], []
x_raw_list, x_norm_list = [], []
param_info_list = []
# Logic for SHIFTING PEAK incremental test.
if pipeline_id == 'shifting':
param_range = np.arange(SHIFTING_CENTER_WL_RANGE[0], SHIFTING_CENTER_WL_RANGE[1] + INCREMENTAL_STEP, INCREMENTAL_STEP)
# Loop through every center wavelength value.
for param_val in tqdm(param_range, desc="Center WL"):
sigma = fwhm_to_sigma(SHIFTING_FWHM)
raw_spectrum = generate_gaussian_spectrum(wavelength_grid, param_val, sigma, SHIFTING_AMP)
x_raw, y_true = compute_raw_sensor_outputs_and_spectrum(raw_spectrum, wavelength_grid, sensors)
x_norm = normalize_sensor_data(x_raw.reshape(1, -1))
# Get predictions from all models.
nn_pred = nn_model(torch.FloatTensor(x_norm).to(DEVICE)).detach().cpu().numpy().flatten()
linear_ridge_pred = linear_ridge_model.predict(x_norm).flatten()
ridge_pred = ridge_model.predict(x_norm).flatten()
kr_pred = kernel_ridge_model.predict(x_norm).flatten()
# Store results.
param_info_list.append({'parameter_name': 'center_wl', 'parameter_value': param_val})
y_true_list.append(y_true)
nn_preds_list.append(nn_pred)
linear_ridge_preds_list.append(linear_ridge_pred)
ridge_preds_list.append(ridge_pred)
kr_preds_list.append(kr_pred)
x_raw_list.append(x_raw)
x_norm_list.append(x_norm.flatten())
# ... (Similar logic for 'fwhm' and 'merging' pipelines) ...
# Save all the collected results to CSV files.
save_model_reconstruction_csv(y_true_list, nn_preds_list, 'nn', wavelength_grid, param_info_list, test_case_dir)
save_model_reconstruction_csv(y_true_list, linear_ridge_preds_list, 'linear_ridge', wavelength_grid, param_info_list, test_case_dir)
save_model_reconstruction_csv(y_true_list, ridge_preds_list, 'ridge', wavelength_grid, param_info_list, test_case_dir)
save_model_reconstruction_csv(y_true_list, kr_preds_list, 'kernel_ridge', wavelength_grid, param_info_list, test_case_dir)
# Save the raw and normalized sensor data used for the tests.
save_incremental_sensor_data_csv(pipeline_id, param_info_list, x_raw_list, x_norm_list, sensor_names, test_case_dir)
# Main function to run a full training and evaluation pipeline.
def run_pipeline(pipeline_id, data_prep_func, output_dir):
print("\n" + "="*80)
print(f"--- STARTING PIPELINE: {pipeline_id.upper()} ---")
print("="*80 + "\n")
ensure_output_dir(output_dir)
sensors = load_sensor_data(SENSOR_FILES)
if not sensors:
return
sensor_names = sorted(sensors.keys())
# --- Data Preparation ---
common_wl_grid = np.linspace(COMMON_WL_RANGE[0], COMMON_WL_RANGE[1], N_POINTS_RESAMPLE)
# Generate the synthetic data.
X_raw, y_spectra, params = data_prep_func(sensors, common_wl_grid)
# Optionally add noise.
if ADD_NOISE_TO_TRAIN_DATA:
print(f"Adding Gaussian noise with level {NOISE_LEVEL} to training data...")
X_raw = add_noise(X_raw, NOISE_LEVEL)
# Normalize the sensor data.
X_normalized = normalize_sensor_data(X_raw)
# Split data into training and validation sets.
X_train, X_val, y_train, y_val, params_train, params_val = train_test_split(
X_normalized, y_spectra, params, test_size=VAL_RATIO, random_state=RANDOM_SEED
)
X_raw_train, X_raw_val = train_test_split(
X_raw, test_size=VAL_RATIO, random_state=RANDOM_SEED
)
print(f"\nData split: {len(X_train)} training samples, {len(X_val)} validation samples.")
# --- Model Training ---
# Prepare data for PyTorch.
train_dataset = TensorDataset(torch.FloatTensor(X_train).to(DEVICE), torch.FloatTensor(y_train).to(DEVICE))
val_dataset = TensorDataset(torch.FloatTensor(X_val).to(DEVICE), torch.FloatTensor(y_val).to(DEVICE))
train_loader = DataLoader(train_dataset, batch_size=NN_BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=NN_BATCH_SIZE, shuffle=False)
# 1. Train the Neural Network.
input_dim = X_train.shape[1]
output_dim = y_train.shape[1]
nn_model = SpectrumNet(input_dim, output_dim, NN_HIDDEN_LAYERS, NN_DROPOUT_RATE).to(DEVICE)
criterion = nn.MSELoss()
optimizer = optim.Adam(nn_model.parameters(), lr=NN_LEARNING_RATE, weight_decay=NN_WEIGHT_DECAY)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=NN_SCHEDULER_PATIENCE, factor=NN_SCHEDULER_FACTOR)
nn_model, _ = train_nn_model(nn_model, criterion, optimizer, scheduler, train_loader, val_loader, output_dir)
# 2. Train the Baseline Models.
linear_ridge_model, ridge_model, kernel_ridge_model = train_baseline_models(X_train, y_train, output_dir)
# --- Evaluation ---
models = (nn_model, linear_ridge_model, ridge_model, kernel_ridge_model)
validation_data = (X_val, y_val, val_loader)
evaluate_and_plot_results(models, validation_data, common_wl_grid, output_dir, sensor_names)
# --- Incremental Testing ---
evaluate_and_save_incremental_tests(models, pipeline_id, sensors, common_wl_grid, output_dir, sensor_names)
print("\n" + "="*80)
print(f"--- PIPELINE {pipeline_id.upper()} COMPLETE ---")
print("="*80 + "\n")
# This is the main entry point when the script is run.
if __name__ == '__main__':
start_total_time = time.time()
# Run all three pipelines sequentially.
run_pipeline('shifting', prepare_shifting_gaussian_data, OUTPUT_DIR_SHIFTING)
run_pipeline('fwhm', prepare_varying_fwhm_data, OUTPUT_DIR_FWHM)
run_pipeline('merging', prepare_merging_gaussians_data, OUTPUT_DIR_MERGING)
end_total_time = time.time()
print(f"All pipelines finished in {(end_total_time - start_total_time) / 60:.2f} minutes.")Editor is loading...
Leave a Comment