diff --git a/examples/prometheus_torch/README.md b/examples/prometheus_torch/README.md new file mode 100644 index 0000000..998d1e6 --- /dev/null +++ b/examples/prometheus_torch/README.md @@ -0,0 +1,72 @@ +# Prometheus Torch + +Example harness for Prometheus model reproduced in Torch. +Prodives script to train TemoralPrediction model in Keras and Torch, +as well as evaluation and comparison between two models. + +The example bundle contains a reproduction script: +the script trains a random intialized Prometheus model over each datetime labeled CSV data file. +It's important to note that the final operation before prediction is un-normalizing the +model outputs using statistics gathered over training data. +The script saves checkpoints after 1 epoch on each dataset file, saving both torch +model weights and a json with the current un-normalization terms. + +Using the `npm1_pwr_model.keras` checkpoint of Prometheus as a baseline, +the reproduced torch model works comparably with the keras model. + +test03_2025-09-18 + +Select an early checkpoint of the reproduced model for tests with APEIRON. +The harness checkpoints both the weights and the stats for comparisons against the base Prometheus model. + +## Getting Started + +Place the data and model weights under `examples/prometheus_torch/data/` +``` +data/ +├── test +│   ├── 2025-03-20.csv +│   ├── 2025-05-12.csv +│   ├── 2025-07-23.csv +│   └── 2025-09-18.csv +├── train +│   ├── 2025-02-27.csv +│   ├── 2025-03-12.csv +│   ├── 2025-03-19.csv +│   ├── 2025-03-27.csv +│   ├── 2025-04-23.csv +│   ├── 2025-04-28.csv +│   ├── 2025-04-30.csv +│   ├── 2025-05-20.csv +│   ├── 2025-06-02.csv +│   ├── 2025-06-04.csv +│   ├── 2025-06-10.csv +│   ├── 2025-06-12.csv +│   ├── 2025-06-26.csv +│   ├── 2025-07-21.csv +│   ├── 2025-07-22.csv +│   ├── 2025-07-30.csv +│   ├── 2025-07-31.csv +│   ├── 2025-08-25.csv +│   ├── 2025-08-26.csv +│   ├── 2025-09-02.csv +│   ├── 2025-09-16.csv +│   ├── 2025-09-17.csv +│   └── 2025-09-25.csv +├── npm1_pwr_config.pkl +├── npm1_pwr_model.h5 +└── npm1_pwr_model.keras + +3 directories, 30 files +``` + +### Reproduce Prometheus model and save as .pt +``` +cd ./examples/prometheus_torch/ +python reproduce_prometheus.py train --save ./output/prometheus_torch/reproduced_prometheus.pt +``` + +### Compare model checkpoint against base model on test set +``` +python reproduce_prometheus.py compare --model ./data/npm1_pwr_model.keras --torch-model ./output/apeiron/drift_adaptation_5.pt +``` diff --git a/examples/prometheus_torch/TemporalPredictUpdate.py b/examples/prometheus_torch/TemporalPredictUpdate.py new file mode 100644 index 0000000..e440e3d --- /dev/null +++ b/examples/prometheus_torch/TemporalPredictUpdate.py @@ -0,0 +1,1061 @@ +import numpy as np +import pandas as pd +import tensorflow as tf +from tensorflow import keras +from tensorflow.keras import layers +from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler +import pickle +import os +import glob +from pathlib import Path + + +class ModelGeneration: + """ + LSTM model class for handling batched time series data from CSV files. + Designed for reactor control rod position tracking and similar applications. + Supports variable-length and rolling-window sequence modes, optional + log-space target transformation, and relative/log-space loss functions + for accurate prediction across wide dynamic ranges (e.g. 0 W – 250 kW). + """ + + def __init__(self, model_path, testing_path, training_path, text_file_name, + feature_variables, target_variable, epochs=100, + lstm_units=[64, 32], dropout=0.2, normalization_method='minmax', + batch_size=32, validation_split=0.2, + use_rolling_windows=False, window_size=100, window_stride=10, + forecast_horizon=0, + target_log_transform=False, + loss_fn='mse', + random_seed=None): + """ + Initialize the ModelGeneration class. + + Parameters + ---------- + model_path : str + Path where the trained model will be saved. + testing_path : str + Directory containing testing CSV files. + training_path : str + Directory containing training CSV files. + text_file_name : str + Path for saving model summary and training information. + feature_variables : list + Column names to use as input features (e.g. ['Shim_1', 'Shim_2', 'Reg']). + target_variable : list + Column names to use as targets (e.g. ['power']). + epochs : int + Number of training epochs. + lstm_units : list + Units for each LSTM layer. + dropout : float + Dropout rate for regularization. + normalization_method : str + 'minmax' | 'standard' | 'robust' | 'none' + batch_size : int + Batch size for training. + validation_split : float + Fraction of training data reserved for validation. + use_rolling_windows : bool + If True, creates overlapping windows from each sequence. + If False, uses entire variable-length sequences. + window_size : int + Rolling window length in timesteps. + window_stride : int + Stride between successive windows. + forecast_horizon : int + 0 → predict current y_t (aligned / nowcasting). + >0 → predict y_{t+h} (forecasting h steps ahead). + target_log_transform : bool + If True, applies np.log1p() to y before normalization and + np.expm1() after denormalization. Compresses wide dynamic + ranges (e.g. 50 W – 250 kW) so low-power values are not + washed out by the dominant high-power training signal. + Requires all target values to be non-negative. + loss_fn : str + 'mse' – standard mean squared error (default). + 'relative_mse' – MSE / |y_true|²; equal relative weight + across all power levels. + 'log_mse' – MSE in log-space; implicit relative weighting + without modifying the normalization pipeline. + Note: combining target_log_transform=True with loss_fn='mse' + already trains in log-space and is the recommended pairing for + NRAD power prediction. + random_seed : int + Ability to set random seed when splitting training/validation data set + """ + self.model_path = model_path + self.testing_path = testing_path + self.training_path = training_path + self.text_file_name = text_file_name + self.feature_variables = feature_variables + self.target_variable = target_variable + self.epochs = epochs + self.lstm_units = lstm_units + self.dropout = dropout + self.normalization_method = normalization_method + self.batch_size = batch_size + self.validation_split = validation_split + self.use_rolling_windows = use_rolling_windows + self.window_size = window_size + self.window_stride = window_stride + self.forecast_horizon = int(forecast_horizon) + self.target_log_transform = target_log_transform + self.loss_fn = loss_fn + self.random_seed = random_seed + + if self.forecast_horizon < 0: + raise ValueError("forecast_horizon must be >= 0") + if self.loss_fn not in ('mse', 'relative_mse', 'log_mse'): + raise ValueError(f"Unknown loss_fn '{self.loss_fn}'. " + "Choose 'mse', 'relative_mse', or 'log_mse'.") + + # Data storage + self.X_train = None + self.y_train = None + self.X_val = None + self.y_val = None + self.X_test = None + self.y_test = None + + # Model components + self.model = None + self.scaler_X = None + self.scaler_y = None + self.history = None + + # Metadata + self.train_files = [] + self.test_files = [] + self.max_timesteps = None + + Path(self.model_path).parent.mkdir(parents=True, exist_ok=True) + Path(self.text_file_name).parent.mkdir(parents=True, exist_ok=True) + + self._initialize_scalers() + + # ───────────────────────────────────────────────────────────────────────── + # Scalers + # ───────────────────────────────────────────────────────────────────────── + + def _initialize_scalers(self): + """Initialize scalers based on chosen normalization method.""" + if self.normalization_method == 'minmax': + self.scaler_X = MinMaxScaler() + self.scaler_y = MinMaxScaler() + elif self.normalization_method == 'standard': + self.scaler_X = StandardScaler() + self.scaler_y = StandardScaler() + elif self.normalization_method == 'robust': + self.scaler_X = RobustScaler() + self.scaler_y = RobustScaler() + elif self.normalization_method == 'none': + self.scaler_X = None + self.scaler_y = None + else: + raise ValueError(f"Unknown normalization method: {self.normalization_method}") + + # ───────────────────────────────────────────────────────────────────────── + # I/O helpers + # ───────────────────────────────────────────────────────────────────────── + + def _read_csv_files(self, directory_path): + """Read all CSV files from a directory.""" + csv_files = glob.glob(os.path.join(directory_path, "*.csv")) + if not csv_files: + raise ValueError(f"No CSV files found in {directory_path}") + + dataframes, filenames = [], [] + for csv_file in sorted(csv_files): + try: + df = pd.read_csv(csv_file) + dataframes.append(df) + filenames.append(os.path.basename(csv_file)) + except Exception as e: + print(f"Warning: Could not read {csv_file}: {e}") + + return dataframes, filenames + + # ───────────────────────────────────────────────────────────────────────── + # Sequence extraction & windowing + # ───────────────────────────────────────────────────────────────────────── + + def _extract_sequences(self, dataframes, feature_cols, target_cols): + """Extract feature and target sequences from DataFrames.""" + X, y = [], [] + for df in dataframes: + missing_f = set(feature_cols) - set(df.columns) + missing_t = set(target_cols) - set(df.columns) + if missing_f: + raise ValueError(f"Missing feature columns: {missing_f}") + if missing_t: + raise ValueError(f"Missing target columns: {missing_t}") + X.append(df[feature_cols].values) + y.append(df[target_cols].values) + return X, y + + def _create_rolling_windows(self, X_sequences, y_sequences): + """ + Create rolling windows from sequences. + forecast_horizon=0 → aligned window targets. + forecast_horizon>0 → forecast_horizon future timesteps as targets. + """ + X_windows, y_windows = [], [] + for X_seq, y_seq in zip(X_sequences, y_sequences): + seq_len = len(X_seq) + min_req = self.window_size + self.forecast_horizon + if seq_len < min_req: + print(f"Warning: Skipping sequence of length {seq_len} (< {min_req})") + continue + for start in range(0, seq_len - min_req + 1, self.window_stride): + end = start + self.window_size + X_windows.append(X_seq[start:end]) + if self.forecast_horizon == 0: + y_windows.append(y_seq[start:end]) + else: + y_windows.append(y_seq[end:end + self.forecast_horizon]) + return X_windows, y_windows + + # ───────────────────────────────────────────────────────────────────────── + # Preprocessing + # ───────────────────────────────────────────────────────────────────────── + + def data_preprocessing(self): + """ + Load and preprocess data from CSV files. + Reads training and testing data, extracts sequences, and prepares + for modelling. Applies rolling windows when enabled. + """ + print("=" * 60) + print("DATA PREPROCESSING") + print("=" * 60) + + print(f"\nReading training data from: {self.training_path}") + train_dfs, self.train_files = self._read_csv_files(self.training_path) + print(f" Found {len(train_dfs)} training files") + + print(f"\nReading testing data from: {self.testing_path}") + test_dfs, self.test_files = self._read_csv_files(self.testing_path) + print(f" Found {len(test_dfs)} testing files") + + print("\nExtracting sequences...") + X_train_raw, y_train_raw = self._extract_sequences( + train_dfs, self.feature_variables, self.target_variable) + X_test_raw, y_test_raw = self._extract_sequences( + test_dfs, self.feature_variables, self.target_variable) + + # Forecast-horizon shift for full-sequence (non-rolling) mode + if (not self.use_rolling_windows) and (self.forecast_horizon > 0): + X_train_raw = [s[:-self.forecast_horizon] for s in X_train_raw + if len(s) > self.forecast_horizon] + y_train_raw = [s[self.forecast_horizon:] for s in y_train_raw + if len(s) > self.forecast_horizon] + X_test_raw = [s[:-self.forecast_horizon] for s in X_test_raw + if len(s) > self.forecast_horizon] + y_test_raw = [s[self.forecast_horizon:] for s in y_test_raw + if len(s) > self.forecast_horizon] + + if self.use_rolling_windows: + print(f"\nApplying rolling windows (size={self.window_size}, " + f"stride={self.window_stride}, horizon={self.forecast_horizon})") + X_train_raw, y_train_raw = self._create_rolling_windows(X_train_raw, y_train_raw) + X_test_raw, y_test_raw = self._create_rolling_windows(X_test_raw, y_test_raw) + print(f" Training windows : {len(X_train_raw)}") + print(f" Testing windows : {len(X_test_raw)}") + self.max_timesteps = self.window_size + else: + all_lengths = [len(s) for s in X_train_raw + X_test_raw] + self.max_timesteps = max(all_lengths) + + # Train / validation split + n_train = len(X_train_raw) + n_val = int(n_train * self.validation_split) + + # Set seed if specified + if self.random_seed is not None: + np.random.seed(self.random_seed) + + indices = np.random.permutation(n_train) + val_idx = indices[:n_val] + train_idx = indices[n_val:] + + self.X_train = [X_train_raw[i] for i in train_idx] + self.y_train = [y_train_raw[i] for i in train_idx] + self.X_val = [X_train_raw[i] for i in val_idx] + self.y_val = [y_train_raw[i] for i in val_idx] + self.X_test = X_test_raw + self.y_test = y_test_raw + + print("\nData Statistics:") + print(f" Training : {len(self.X_train)} samples") + print(f" Validation : {len(self.X_val)} samples") + print(f" Testing : {len(self.X_test)} samples") + print(f" Features : {len(self.feature_variables)} {self.feature_variables}") + print(f" Targets : {len(self.target_variable)} {self.target_variable}") + if self.target_log_transform: + print(f" Target log-transform : ENABLED (log1p/expm1)") + + if not self.use_rolling_windows: + tl = [len(s) for s in self.X_train] + el = [len(s) for s in self.X_test] + print(f"\nSequence lengths — Train: min={min(tl)}, max={max(tl)}, " + f"mean={np.mean(tl):.1f} | Test: min={min(el)}, max={max(el)}") + print(f" Max timesteps for model: {self.max_timesteps}") + + print(f"\nTarget Variable Statistics (across all timesteps):") + for i, var in enumerate(self.target_variable): + tv = np.concatenate([s[:, i] for s in self.y_train]) + ev = np.concatenate([s[:, i] for s in self.y_test]) + print(f" {var}:") + print(f" Train — min={tv.min():.4f}, max={tv.max():.4f}, mean={tv.mean():.4f}") + print(f" Test — min={ev.min():.4f}, max={ev.max():.4f}, mean={ev.mean():.4f}") + + print("\nData preprocessing complete!") + print("=" * 60) + + # ───────────────────────────────────────────────────────────────────────── + # Normalization + # ───────────────────────────────────────────────────────────────────────── + + def _normalize_data(self, X, y=None, fit=True): + """ + Normalize batched variable-length sequences. + If target_log_transform=True, applies log1p to y before scaling + so the scaler fits in log-space, expanding the low-value regime. + """ + if self.normalization_method == 'none': + return X, y + if not X or len(X) == 0: + return X, y + + # ── X ────────────────────────────────────────────────────────────── + X_flat = np.vstack(X) + X_scaled = (self.scaler_X.fit_transform(X_flat) if fit + else self.scaler_X.transform(X_flat)) + X_norm, idx = [], 0 + for seq in X: + l = len(seq) + X_norm.append(X_scaled[idx:idx+l]) + idx += l + + # ── y ────────────────────────────────────────────────────────────── + y_norm = None + if y is not None and len(y) > 0: + y_flat = np.vstack(y) + + if self.target_log_transform: + if np.any(y_flat < 0): + raise ValueError( + "target_log_transform=True requires non-negative targets. " + "Found negative values in y.") + y_flat = np.log1p(y_flat) + + if self.scaler_y is not None: + y_scaled = (self.scaler_y.fit_transform(y_flat) if fit + else self.scaler_y.transform(y_flat)) + else: + y_scaled = y_flat + + y_norm, idx = [], 0 + for seq in y: + l = len(seq) + y_norm.append(y_scaled[idx:idx+l]) + idx += l + + return X_norm, y_norm + + def _denormalize_data(self, data, is_target=False): + """ + Denormalize data back to original scale. + For targets with target_log_transform=True, applies expm1 after + the scaler's inverse_transform to undo the log1p step. + """ + if self.normalization_method == 'none': + return data + + scaler = self.scaler_y if is_target else self.scaler_X + result = scaler.inverse_transform(data) if scaler is not None else data + + if is_target and self.target_log_transform: + result = np.expm1(result) + result = np.clip(result, 0, None) # guard against sub-zero float noise + + return result + + # ───────────────────────────────────────────────────────────────────────── + # Padding helpers + # ───────────────────────────────────────────────────────────────────────── + + def _pad_sequences(self, sequences, max_length=None): + """Pad feature sequences to uniform length for batching.""" + if max_length is None: + max_length = self.max_timesteps + if self.use_rolling_windows: + return np.array(sequences) + n, feats = len(sequences), sequences[0].shape[-1] + padded = np.zeros((n, max_length, feats)) + for i, seq in enumerate(sequences): + l = min(len(seq), max_length) + padded[i, :l, :] = seq[:l] + return padded + + def _pad_targets(self, targets, max_length=None): + """Pad target sequences to uniform length for batching.""" + if max_length is None: + max_length = self.max_timesteps + if self.use_rolling_windows: + return np.array(targets) + n, feats = len(targets), targets[0].shape[-1] + padded = np.zeros((n, max_length, feats)) + for i, seq in enumerate(targets): + l = min(len(seq), max_length) + padded[i, :l, :] = seq[:l] + return padded + + # ───────────────────────────────────────────────────────────────────────── + # Loss function + # ───────────────────────────────────────────────────────────────────────── + + def _get_loss(self): + """ + Return the Keras-compatible loss for model.compile(). + + 'mse' – standard mean squared error. + 'relative_mse' – MSE / |y_true|²; equal relative penalty at all + power levels. + 'log_mse' – MSE in log-space; operates on normalized model + outputs so it complements (but does not duplicate) + target_log_transform. + """ + if self.loss_fn == 'mse': + return 'mse' + + elif self.loss_fn == 'relative_mse': + eps = tf.constant(1e-6, dtype=tf.float32) + def relative_mse(y_true, y_pred): + denom = tf.square(tf.abs(y_true) + eps) + return tf.reduce_mean(tf.square(y_true - y_pred) / denom) + return relative_mse + + elif self.loss_fn == 'log_mse': + eps = tf.constant(1.0, dtype=tf.float32) + def log_mse(y_true, y_pred): + log_t = tf.math.log(tf.abs(y_true) + eps) + log_p = tf.math.log(tf.abs(y_pred) + eps) + return tf.reduce_mean(tf.square(log_t - log_p)) + return log_mse + + # ───────────────────────────────────────────────────────────────────────── + # Model construction + # ───────────────────────────────────────────────────────────────────────── + + def _build_model(self): + """Build the LSTM model architecture.""" + print("\nBuilding LSTM model...") + + inputs = layers.Input(shape=(self.max_timesteps, len(self.feature_variables))) + + x = inputs if self.use_rolling_windows else layers.Masking(mask_value=0.0)(inputs) + + if self.forecast_horizon == 0: + # Sequence-to-sequence (aligned prediction) + for i, units in enumerate(self.lstm_units): + x = layers.LSTM(units, return_sequences=True, dropout=self.dropout)(x) + if i < len(self.lstm_units) - 1: + x = layers.Dropout(self.dropout)(x) + outputs = layers.TimeDistributed(layers.Dense(len(self.target_variable)))(x) + out_shape = f"(batch, {self.max_timesteps}, {len(self.target_variable)})" + arch_type = "Sequence-to-Sequence (aligned prediction)" + else: + # Encoder-Decoder (forecasting) + for i, units in enumerate(self.lstm_units[:-1]): + x = layers.LSTM(units, return_sequences=True, dropout=self.dropout)(x) + x = layers.Dropout(self.dropout)(x) + x = layers.LSTM(self.lstm_units[-1], return_sequences=False, dropout=self.dropout)(x) + x = layers.RepeatVector(self.forecast_horizon)(x) + x = layers.LSTM(self.lstm_units[-1], return_sequences=True, dropout=self.dropout)(x) + outputs = layers.TimeDistributed(layers.Dense(len(self.target_variable)))(x) + out_shape = f"(batch, {self.forecast_horizon}, {len(self.target_variable)})" + arch_type = f"Encoder-Decoder (forecast {self.forecast_horizon} steps ahead)" + + self.model = keras.Model(inputs=inputs, outputs=outputs) + self.model.compile( + optimizer=keras.optimizers.Adam(learning_rate=0.001), + loss=self._get_loss(), + metrics=['mae', 'mse'] + ) + + print(f" Input shape : (batch, {self.max_timesteps}, {len(self.feature_variables)})") + print(f" Output shape : {out_shape}") + print(f" LSTM layers : {self.lstm_units}") + print(f" Parameters : {self.model.count_params():,}") + print(f" Architecture : {arch_type}") + print(f" Loss function: {self.loss_fn}") + if self.use_rolling_windows: + print(f" Mode: Rolling Windows (size={self.window_size}, stride={self.window_stride})") + else: + print(f" Mode: Variable-length sequences (with masking)") + + # ───────────────────────────────────────────────────────────────────────── + # Training + # ───────────────────────────────────────────────────────────────────────── + + def train(self): + """Train the LSTM model on preprocessed data.""" + if self.X_train is None: + raise ValueError("Data not preprocessed. Call data_preprocessing() first.") + + print("\n" + "=" * 60) + print("MODEL TRAINING") + print("=" * 60) + + print("\nNormalizing data...") + X_train_norm, y_train_norm = self._normalize_data(self.X_train, self.y_train, fit=True) + + has_val = bool(self.X_val and len(self.X_val) > 0) + if has_val: + X_val_norm, y_val_norm = self._normalize_data(self.X_val, self.y_val, fit=False) + else: + print("Warning: No validation data. Training without validation.") + X_val_norm = y_val_norm = None + + print("Padding sequences...") + X_tr = self._pad_sequences(X_train_norm) + y_tr = self._pad_targets(y_train_norm) + val_data = None + if has_val: + val_data = (self._pad_sequences(X_val_norm), self._pad_targets(y_val_norm)) + + self._build_model() + + # Shape verification + print("\n=== SHAPE VERIFICATION ===") + exp_y = (len(self.y_train), + self.forecast_horizon if self.forecast_horizon > 0 else self.max_timesteps, + len(self.target_variable)) + if y_tr.shape != exp_y: + print(f"WARNING: target shape mismatch — expected {exp_y}, got {y_tr.shape}") + else: + print(f"✓ Target shape verified: {y_tr.shape}") + print("=" * 26 + "\n") + + monitor = 'val_loss' if has_val else 'loss' + callbacks = [ + keras.callbacks.EarlyStopping( + monitor=monitor, patience=15, restore_best_weights=True, verbose=1), + keras.callbacks.ReduceLROnPlateau( + monitor=monitor, factor=0.5, patience=5, min_lr=1e-7, verbose=1), + keras.callbacks.ModelCheckpoint( + self.model_path, monitor=monitor, save_best_only=True, verbose=1), + ] + + print(f"Training for up to {self.epochs} epochs...") + self.history = self.model.fit( + X_tr, y_tr, + validation_data=val_data, + epochs=self.epochs, + batch_size=self.batch_size, + callbacks=callbacks, + verbose=1 + ) + + print("\nTraining complete!") + print("=" * 60) + self._save_model_info() + + # ───────────────────────────────────────────────────────────────────────── + # Evaluation + # ───────────────────────────────────────────────────────────────────────── + + def evaluate(self, verbose=1): + """Evaluate the model on test data and return per-variable metrics.""" + if self.X_test is None: + raise ValueError("Test data not loaded. Call data_preprocessing() first.") + if self.model is None: + raise ValueError("Model not trained. Call train() first.") + + print("\n" + "=" * 60) + print("MODEL EVALUATION") + print("=" * 60) + + X_test_norm, y_test_norm = self._normalize_data(self.X_test, self.y_test, fit=False) + X_tp = self._pad_sequences(X_test_norm) + y_tp = self._pad_targets(y_test_norm) + + print("\nEvaluating on test data...") + self.model.evaluate(X_tp, y_tp, verbose=verbose) + + y_pred_norm = self.model.predict(X_tp, verbose=0) + bs, ts, feats = y_pred_norm.shape + y_pred = self._denormalize_data(y_pred_norm.reshape(-1, feats), + is_target=True).reshape(bs, ts, feats) + + print("\nTest Metrics (Original Scale):") + metrics = {} + for i, var in enumerate(self.target_variable): + true_all, pred_all = [], [] + for j in range(len(self.y_test)): + sl = len(self.y_test[j]) + true_all.extend(self.y_test[j][:, i]) + pred_all.extend(y_pred[j, :sl, i]) + t, p = np.array(true_all), np.array(pred_all) + metrics[var] = self._compute_metrics(t, p) + self._print_metrics(var, metrics[var]) + + print("\n" + "=" * 60) + return metrics, y_pred + + def evaluate_full_sequence(self, verbose=1): + """ + Evaluate on full test sequences by aggregating rolling-window predictions. + Falls back to standard evaluate() when use_rolling_windows=False. + """ + if self.X_test is None: + raise ValueError("Test data not loaded. Call data_preprocessing() first.") + if self.model is None: + raise ValueError("Model not trained. Call train() first.") + + print("\n" + "=" * 60) + print("FULL SEQUENCE EVALUATION") + print("=" * 60) + + if not self.use_rolling_windows: + print("\nNot using rolling windows — falling back to standard evaluation.") + return self.evaluate(verbose=verbose) + + print(f"\nReconstructing full sequences from rolling windows...") + all_preds, all_true = [], [] + for tf_ in self.test_files: + df = pd.read_csv(os.path.join(self.testing_path, tf_)) + all_preds.append(self.predict_full_sequence(df)) + all_true.append(df[self.target_variable].values) + + print("\nTest Metrics (Full Sequences, Original Scale):") + metrics = {} + for i, var in enumerate(self.target_variable): + t = np.concatenate([s[:, i] for s in all_true]) + p = np.concatenate([s[:, i] for s in all_preds]) + valid = ~np.isnan(t) & ~np.isnan(p) + t, p = t[valid], p[valid] + if len(t) == 0: + print(f"\nWarning: No valid predictions for {var}.") + continue + metrics[var] = self._compute_metrics(t, p) + self._print_metrics(var, metrics[var]) + print(f" Timesteps evaluated: {len(t):,}") + + print("\n" + "=" * 60) + return metrics, all_preds + + def evaluate_power_bands(self, y_true_flat, y_pred_flat, bands=None): + """ + Break out evaluation metrics by power band. + + Useful for verifying low-power capture independently of full-power + accuracy. Default bands are tuned for the NRAD 0–250 kW range with + emphasis on the startup regime. + + Parameters + ---------- + y_true_flat : 1-D array (original-scale, same units as training data) + y_pred_flat : 1-D array + bands : list of (label, low, high), optional + Custom power bands. Default NRAD bands used if None. + + Returns + ------- + dict keyed by band label with keys: n, RMSE, MAPE, R2 + """ + if bands is None: + bands = [ + ("startup 0–50 W", 0, 50), + ("low 50 W–1 kW", 50, 1_000), + ("mid 1–10 kW", 1_000, 10_000), + ("high 10–250 kW", 10_000, 250_000), + ] + + y_true_flat = np.asarray(y_true_flat) + y_pred_flat = np.asarray(y_pred_flat) + + results = {} + print("\nPower-Band Evaluation") + print(f" {'Band':<24} {'N':>8} {'RMSE':>12} {'MAPE %':>8} {'R²':>7}") + print(" " + "-" * 66) + + for label, lo, hi in bands: + mask = (y_true_flat >= lo) & (y_true_flat < hi) + n = mask.sum() + if n == 0: + print(f" {label:<24} {'—':>8}") + continue + t, p = y_true_flat[mask], y_pred_flat[mask] + rmse = np.sqrt(np.mean((t - p) ** 2)) + nz = t != 0 + mape = np.mean(np.abs((t[nz] - p[nz]) / t[nz])) * 100 if nz.any() else np.nan + ss_r = np.sum((t - p) ** 2) + ss_t = np.sum((t - t.mean()) ** 2) + r2 = 1 - ss_r / ss_t if ss_t > 0 else np.nan + results[label] = {'n': n, 'RMSE': rmse, 'MAPE': mape, 'R2': r2} + mape_s = f"{mape:8.1f}" if not np.isnan(mape) else " N/A" + r2_s = f"{r2:7.4f}" if not np.isnan(r2) else " N/A" + print(f" {label:<24} {n:>8,} {rmse:>12.4f} {mape_s} {r2_s}") + + return results + + # ───────────────────────────────────────────────────────────────────────── + # Metrics helpers + # ───────────────────────────────────────────────────────────────────────── + + @staticmethod + def _compute_metrics(t, p): + """Compute standard regression metrics given flat true/pred arrays.""" + mse = np.mean((t - p) ** 2) + mae = np.mean(np.abs(t - p)) + rmse = np.sqrt(mse) + me = np.max(np.abs(t - p)) + nz = t != 0 + mape = np.mean(np.abs((t[nz] - p[nz]) / t[nz])) * 100 if nz.any() else np.nan + ss_r = np.sum((t - p) ** 2) + ss_t = np.sum((t - t.mean()) ** 2) + r2 = 1 - ss_r / ss_t + return {'MSE': mse, 'MAE': mae, 'RMSE': rmse, 'ME': me, 'MAPE': mape, 'R²': r2} + + @staticmethod + def _print_metrics(var, m): + print(f"\n {var}:") + print(f" MSE : {m['MSE']:.6f}") + print(f" MAE : {m['MAE']:.6f}") + print(f" RMSE : {m['RMSE']:.6f}") + print(f" ME : {m['ME']:.6f}") + if not np.isnan(m['MAPE']): + print(f" MAPE : {m['MAPE']:.2f}%") + else: + print(f" MAPE : N/A (zero values present)") + print(f" R² : {m['R²']:.6f}") + + # ───────────────────────────────────────────────────────────────────────── + # Prediction + # ───────────────────────────────────────────────────────────────────────── + + def predict_full_sequence(self, csv_file_or_dataframe): + """ + Predict the target variable for an entire sequence by aggregating + rolling-window predictions. Overlapping windows are averaged. + + Returns + ------- + np.ndarray, shape (timesteps, n_targets) + forecast_horizon=0: predictions align with all input timesteps. + forecast_horizon>0: first window_size timesteps are NaN. + """ + if self.model is None: + raise ValueError("Model not trained. Call train() first.") + + print("\n" + "=" * 70) + print("DEBUG: predict_full_sequence()") + print("=" * 70) + + if isinstance(csv_file_or_dataframe, str): + df = pd.read_csv(csv_file_or_dataframe) + print(f"Loaded CSV: {os.path.basename(csv_file_or_dataframe)}") + else: + df = csv_file_or_dataframe + print("Using DataFrame") + + X_full = df[self.feature_variables].values + seq_len = len(X_full) + print(f"Sequence length: {seq_len}") + + if not self.use_rolling_windows: + print("Mode: Full sequence (non-rolling)") + print("=" * 70 + "\n") + return self.predict(csv_file_or_dataframe) + + print(f"\nWindow parameters: size={self.window_size}, " + f"stride={self.window_stride}, horizon={self.forecast_horizon}") + + max_start = seq_len - self.window_size - self.forecast_horizon + print(f"Max start index : {max_start}") + + X_windows, window_starts = [], [] + for start in range(0, max_start + 1, self.window_stride): + X_windows.append(X_full[start:start + self.window_size]) + window_starts.append(start) + + print(f"Windows created : {len(X_windows)}") + if not X_windows: + raise ValueError(f"No windows created. Seq length: {seq_len}, " + f"need >= {self.window_size + self.forecast_horizon}") + + X_norm, _ = self._normalize_data(X_windows, fit=False) + X_pad = self._pad_sequences(X_norm) + pred_norm = self.model.predict(X_pad, verbose=0) + + print(f"\nModel I/O — Input: {X_pad.shape}, Output: {pred_norm.shape}") + + bs, out_ts, feats = pred_norm.shape + pred_denorm = self._denormalize_data( + pred_norm.reshape(-1, feats), is_target=True).reshape(bs, out_ts, feats) + + full_pred = np.full((seq_len, len(self.target_variable)), np.nan, dtype=float) + counts = np.zeros(seq_len, dtype=float) + + if self.forecast_horizon == 0: + for i, s in enumerate(window_starts): + for j in range(out_ts): + t = s + j + if t < seq_len: + full_pred[t] = (pred_denorm[i, j] if np.isnan(full_pred[t, 0]) + else full_pred[t] + pred_denorm[i, j]) + counts[t] += 1 + else: + for i, s in enumerate(window_starts): + p_start = s + self.window_size + for j in range(out_ts): + t = p_start + j + if t < seq_len: + full_pred[t] = (pred_denorm[i, j] if np.isnan(full_pred[t, 0]) + else full_pred[t] + pred_denorm[i, j]) + counts[t] += 1 + + for t in range(seq_len): + if counts[t] > 1: + full_pred[t] /= counts[t] + + coverage = int(np.sum(~np.isnan(full_pred[:, 0]))) + print(f"Coverage: {coverage}/{seq_len} timesteps") + print("=" * 70 + "\n") + return full_pred + + def predict(self, csv_file_or_dataframe, return_all_windows=False): + """ + Make predictions on new data. + + Parameters + ---------- + csv_file_or_dataframe : str or DataFrame + return_all_windows : bool + Rolling-window mode only. If True, returns all window predictions + with shape (n_windows, output_timesteps, n_targets). + If False, returns the final window's predictions. + + Returns + ------- + np.ndarray + """ + if self.model is None: + raise ValueError("Model not trained. Call train() first.") + + df = pd.read_csv(csv_file_or_dataframe) if isinstance( + csv_file_or_dataframe, str) else csv_file_or_dataframe + X_full = df[self.feature_variables].values + + if self.use_rolling_windows: + seq_len = len(X_full) + if seq_len < self.window_size: + raise ValueError(f"Input length ({seq_len}) < window size ({self.window_size})") + X_windows = [X_full[s:s + self.window_size] + for s in range(0, seq_len - self.window_size + 1, self.window_stride)] + X_norm, _ = self._normalize_data(X_windows, fit=False) + X_pad = self._pad_sequences(X_norm) + pn = self.model.predict(X_pad, verbose=0) + bs, ts, f = pn.shape + preds = self._denormalize_data(pn.reshape(-1, f), + is_target=True).reshape(bs, ts, f) + return preds if return_all_windows else preds[-1] + else: + X_norm, _ = self._normalize_data([X_full], fit=False) + X_pad = self._pad_sequences(X_norm) + pn = self.model.predict(X_pad, verbose=0) + bs, ts, f = pn.shape + preds = self._denormalize_data(pn.reshape(-1, f), + is_target=True).reshape(bs, ts, f) + actual = len(df) + if self.forecast_horizon > 0: + out = np.full((actual, f), np.nan, dtype=float) + out[self.forecast_horizon:, :] = preds[0, :actual - self.forecast_horizon, :] + return out + return preds[0, :actual, :] + + # ───────────────────────────────────────────────────────────────────────── + # Persistence + # ───────────────────────────────────────────────────────────────────────── + + def save(self, base_path=None): + """ + Save the model, scalers, and configuration. + + Creates: + {base_path}_model.keras + {base_path}_model.h5 + {base_path}_config.pkl + """ + if self.model is None: + raise ValueError("Model not trained. Call train() first.") + + if base_path is None: + base_path = str(Path(self.model_path).with_suffix('')) + + model_keras = f"{base_path}_model.keras" + self.model.save(model_keras) + print(f"Saved Keras model : {model_keras}") + + model_h5 = f"{base_path}_model.h5" + self.model.save(model_h5, save_format='h5') + print(f"Saved H5 model : {model_h5}") + + config = { + 'feature_variables': self.feature_variables, + 'target_variable': self.target_variable, + 'lstm_units': self.lstm_units, + 'dropout': self.dropout, + 'normalization_method': self.normalization_method, + 'use_rolling_windows': self.use_rolling_windows, + 'window_size': self.window_size, + 'window_stride': self.window_stride, + 'forecast_horizon': self.forecast_horizon, + 'target_log_transform': self.target_log_transform, + 'loss_fn': self.loss_fn, + 'max_timesteps': self.max_timesteps, + 'scaler_X': self.scaler_X, + 'scaler_y': self.scaler_y, + 'epochs': self.epochs, + 'batch_size': self.batch_size, + 'validation_split': self.validation_split, + } + cfg_file = f"{base_path}_config.pkl" + with open(cfg_file, 'wb') as fh: + pickle.dump(config, fh) + print(f"Saved config : {cfg_file}") + print(f"\nTo reload: ModelGeneration.load('{base_path}')") + + @classmethod + def load(cls, base_path): + """ + Load a saved ModelGeneration instance. + + Parameters + ---------- + base_path : str + Base path used when save() was called (without extension). + + Returns + ------- + ModelGeneration ready for prediction. + """ + cfg_file = f"{base_path}_config.pkl" + with open(cfg_file, 'rb') as fh: + cfg = pickle.load(fh) + + print(f"Loading model from : {base_path}") + print(f" Features : {cfg['feature_variables']}") + print(f" Targets : {cfg['target_variable']}") + print(f" Rolling windows : {cfg['use_rolling_windows']}") + print(f" Forecast horizon : {cfg.get('forecast_horizon', 0)}") + print(f" Log transform : {cfg.get('target_log_transform', False)}") + print(f" Loss function : {cfg.get('loss_fn', 'mse')}") + + instance = cls( + model_path=f"{base_path}_model.keras", + testing_path="", + training_path="", + text_file_name="", + feature_variables=cfg['feature_variables'], + target_variable=cfg['target_variable'], + epochs=cfg['epochs'], + lstm_units=cfg['lstm_units'], + dropout=cfg['dropout'], + normalization_method=cfg['normalization_method'], + batch_size=cfg['batch_size'], + validation_split=cfg['validation_split'], + use_rolling_windows=cfg['use_rolling_windows'], + window_size=cfg['window_size'], + window_stride=cfg['window_stride'], + forecast_horizon=cfg.get('forecast_horizon', 0), + target_log_transform=cfg.get('target_log_transform', False), + loss_fn=cfg.get('loss_fn', 'mse'), + ) + + instance.model = keras.models.load_model(f"{base_path}_model.keras") + instance.scaler_X = cfg['scaler_X'] + instance.scaler_y = cfg['scaler_y'] + instance.max_timesteps = cfg['max_timesteps'] + + print("Model loaded successfully!") + return instance + + # ───────────────────────────────────────────────────────────────────────── + # Logging + # ───────────────────────────────────────────────────────────────────────── + + def _save_model_info(self): + """Save model configuration and training history to text file.""" + with open(self.text_file_name, 'w') as f: + f.write("=" * 60 + "\n") + f.write("LSTM MODEL CONFIGURATION\n") + f.write("=" * 60 + "\n\n") + + f.write("Model Architecture:\n") + f.write(f" LSTM Units : {self.lstm_units}\n") + f.write(f" Dropout : {self.dropout}\n") + f.write(f" Total Parameters : {self.model.count_params():,}\n\n") + + f.write("Training Configuration:\n") + f.write(f" Epochs : {self.epochs}\n") + f.write(f" Batch Size : {self.batch_size}\n") + f.write(f" Normalization : {self.normalization_method}\n") + f.write(f" Validation Split : {self.validation_split}\n") + f.write(f" Loss Function : {self.loss_fn}\n\n") + + f.write("Target Configuration:\n") + f.write(f" Log Transform : {self.target_log_transform}\n") + if self.target_log_transform: + f.write(f" Transform : log1p (forward) / expm1 (inverse)\n\n") + else: + f.write("\n") + + f.write("Data Mode:\n") + if self.use_rolling_windows: + f.write(f" Mode : Rolling Windows\n") + f.write(f" Window Size : {self.window_size}\n") + f.write(f" Window Stride : {self.window_stride}\n") + f.write(f" Forecast Horizon : {self.forecast_horizon}\n") + pred_type = ("Aligned (nowcasting)" if self.forecast_horizon == 0 + else f"Forecasting ({self.forecast_horizon} steps ahead)") + f.write(f" Prediction Type : {pred_type}\n") + f.write(f" Sequence Length : {self.max_timesteps} (fixed)\n") + else: + f.write(f" Mode : Variable-length sequences\n") + f.write(f" Max Seq Length : {self.max_timesteps}\n") + if self.forecast_horizon > 0: + f.write(f" Forecast Horizon : {self.forecast_horizon}\n") + f.write("\n") + + f.write("Data Configuration:\n") + f.write(f" Feature Variables : {self.feature_variables}\n") + f.write(f" Target Variables : {self.target_variable}\n") + f.write(f" Training Samples : {len(self.X_train)}\n") + f.write(f" Validation Samples: {len(self.X_val) if self.X_val else 0}\n") + f.write(f" Testing Samples : {len(self.X_test)}\n\n") + + f.write("Paths:\n") + f.write(f" Model : {self.model_path}\n") + f.write(f" Training Data : {self.training_path}\n") + f.write(f" Testing Data : {self.testing_path}\n\n") + + if self.history: + f.write("Training History (Final Epoch):\n") + f.write(f" Loss : {self.history.history['loss'][-1]:.6f}\n") + if 'val_loss' in self.history.history: + f.write(f" Val Loss : {self.history.history['val_loss'][-1]:.6f}\n") + f.write(f" MAE : {self.history.history['mae'][-1]:.6f}\n") + if 'val_mae' in self.history.history: + f.write(f" Val MAE : {self.history.history['val_mae'][-1]:.6f}\n") + + f.write("\n" + "=" * 60 + "\n") + f.write("MODEL SUMMARY\n") + f.write("=" * 60 + "\n") + self.model.summary(print_fn=lambda x: f.write(x + '\n')) + + print(f"\nModel information saved to: {self.text_file_name}") \ No newline at end of file diff --git a/examples/prometheus_torch/__init__.py b/examples/prometheus_torch/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/prometheus_torch/data/npm1_pwr_model_2024.keras b/examples/prometheus_torch/data/npm1_pwr_model_2024.keras new file mode 100644 index 0000000..5ffa27a Binary files /dev/null and b/examples/prometheus_torch/data/npm1_pwr_model_2024.keras differ diff --git a/examples/prometheus_torch/model.py b/examples/prometheus_torch/model.py new file mode 100644 index 0000000..322cbe5 --- /dev/null +++ b/examples/prometheus_torch/model.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +import gc +import json +import os +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +from torch import Tensor +from torch.optim import Optimizer +from torch.utils.data import ConcatDataset, DataLoader + +from config.configuration import Config +from examples.prometheus_torch.utils import ( + SequenceDataset, + compute_cumulative_stats, + load_csvs, + make_loader, + make_sequence_windows, +) +from model.torch_model_harness import BaseModelHarness + + +class TorchTemporalModel(nn.Module): + """PyTorch reproduction of the Keras TemporalPredict architecture. + + LSTM(128) -> Dropout -> LSTM(64) -> Dropout -> LSTM(32) -> Linear(n_targets) + Output shape: ``(batch, seq_len, n_targets)`` (full-sequence prediction). + """ + + def __init__(self, n_features: int, n_targets: int, dropout: float = 0.1): + super().__init__() + self.lstm1 = nn.LSTM(n_features, 128, batch_first=True) + self.drop1 = nn.Dropout(dropout) + self.lstm2 = nn.LSTM(128, 64, batch_first=True) + self.drop2 = nn.Dropout(dropout) + self.lstm3 = nn.LSTM(64, 32, batch_first=True) + self.head = nn.Linear(32, n_targets) + + def forward(self, x: Tensor) -> Tensor: + x, _ = self.lstm1(x) + x = self.drop1(x) + x, _ = self.lstm2(x) + x = self.drop2(x) + x, _ = self.lstm3(x) + return self.head(x) + +class TorchTemporalForecastModel(nn.Module): + """Encoder-decoder reproduction of the 2024 TemporalPredict Keras model. + + LSTM(128) -> Dropout -> LSTM(64) -> Dropout -> LSTM(32) + -> RepeatVector(horizon) -> LSTM(32) -> Linear(n_targets). + + Output shape: (batch, forecast_horizon, n_targets). + """ + + def __init__( + self, + n_features: int, + n_targets: int, + forecast_horizon: int = 30, + dropout: float = 0.2, + ): + super().__init__() + if forecast_horizon <= 0: + raise ValueError("forecast_horizon must be positive for forecasting") + + self.forecast_horizon = forecast_horizon + + # These explicit dropout layers approximate Keras LSTM(dropout=...). + # PyTorch ignores the dropout argument on a one-layer nn.LSTM. + self.input_drop1 = nn.Dropout(dropout) + self.lstm1 = nn.LSTM(n_features, 128, batch_first=True) + self.drop1 = nn.Dropout(dropout) + + self.input_drop2 = nn.Dropout(dropout) + self.lstm2 = nn.LSTM(128, 64, batch_first=True) + self.drop2 = nn.Dropout(dropout) + + self.input_drop3 = nn.Dropout(dropout) + self.encoder = nn.LSTM(64, 32, batch_first=True) + + self.input_drop_decoder = nn.Dropout(dropout) + self.decoder = nn.LSTM(32, 32, batch_first=True) + self.head = nn.Linear(32, n_targets) + + def forward(self, x: Tensor) -> Tensor: + """Forecast `forecast_horizon` targets from an input sequence.""" + x, _ = self.lstm1(self.input_drop1(x)) + x = self.drop1(x) + + x, _ = self.lstm2(self.input_drop2(x)) + x = self.drop2(x) + + # Equivalent to Keras's final encoder LSTM with + # return_sequences=False, followed by RepeatVector(horizon). + _, (hidden, _) = self.encoder(self.input_drop3(x)) + decoder_input = hidden[-1].unsqueeze(1).expand( + -1, self.forecast_horizon, -1 + ) + + decoded, _ = self.decoder(self.input_drop_decoder(decoder_input)) + return self.head(decoded) + + +class PrometheusHarness(BaseModelHarness): + """Continual-learning harness for the reproduced Prometheus temporal model. + + Each ``update_data_stream()`` call advances to the next training CSV, + recomputes cumulative normalization stats from all CSVs seen so far, + and builds windowed datasets with full-sequence targets matching the + ``TorchTemporalModel`` output shape ``(batch, seq_len, n_targets)``. + """ + + FEATURE_COLS: List[str] = [ + "NRAD_RX_REG_POS", + "NRAD_RX_SHIM1_POS", + "NRAD_RX_SHIM2_POS", + "total_rod_position", + "NRAD_RX_PERIOD_Inverse", + "NRAD_RX_REG_POS_dt", + "NRAD_RX_REG_POS_dt2", + "NRAD_RX_SHIM1_POS_dt", + "NRAD_RX_SHIM1_POS_dt2", + "NRAD_RX_SHIM2_POS_dt", + "NRAD_RX_SHIM2_POS_dt2", + "NRAD_RX_NMP1_PWR_integral", + ] + TARGET_COLS: List[str] = ["NRAD_RX_NMP1_PWR"] + SEQUENCE_LENGTH: int = 10 + FORECAST_HORIZON: int = 30 + VAL_RATIO: float = 0.2 + + def __init__(self, cfg: Config): + n_features = len(self.FEATURE_COLS) + n_targets = len(self.TARGET_COLS) + #model = TorchTemporalModel(n_features=n_features, n_targets=n_targets) + model = TorchTemporalForecastModel(n_features=n_features, + n_targets=n_targets, + forecast_horizon=self.FORECAST_HORIZON, + ) + super().__init__(cfg=cfg, model=model) + + self.eval_metrics: Dict[str, Any] = {"mse": self.get_criterion()} + self.higher_is_better: Dict[str, bool] = {"mse": False} + + # Load pretrained weights + pretrained_path = cfg.model.pretrained_path + if pretrained_path: + try: + state_dict = torch.load( + pretrained_path, map_location=cfg.device, weights_only=False + ) + self.model.load_state_dict(state_dict) + print(f"Loaded pretrained PrometheusV2 model from {pretrained_path}") + except FileNotFoundError: + print( + f"Warning: Pretrained model not found at {pretrained_path}, " + "using randomly initialised weights." + ) + except Exception as e: + print(f"Warning: Failed to load pretrained model: {e}") + + # Load all training CSVs (one per file, sequential) + train_dir = os.path.join(cfg.data.path, "train") + self.train_dfs: List = load_csvs(train_dir) + if not self.train_dfs: + raise ValueError(f"No CSV files found in '{train_dir}'.") + + # State tracking + self.task_counter: int = 0 + self._dfs_seen: List = [] # accumulates for cumulative stats + self._all_cols = self.FEATURE_COLS + self.TARGET_COLS + self._task_datasets: List[Tuple[SequenceDataset, SequenceDataset]] = [] + self._cur_train_loader: Optional[DataLoader] = None + self._cur_val_loader: Optional[DataLoader] = None + + # ------------------------------------------------------------------ # + # Private helpers # + # ------------------------------------------------------------------ # + + def _dispose_current_loaders(self) -> None: + if self._cur_train_loader is not None: + del self._cur_train_loader + self._cur_train_loader = None + if self._cur_val_loader is not None: + del self._cur_val_loader + self._cur_val_loader = None + gc.collect() + + def _make_loader(self, ds: SequenceDataset, shuffle: bool) -> DataLoader: + return make_loader( + ds, + batch_size=self.cfg.train.batch_size, + shuffle=shuffle, + num_workers=self.cfg.train.num_workers, + pin_memory=torch.cuda.is_available(), + ) + + # ------------------------------------------------------------------ # + # BaseModelHarness interface # + # ------------------------------------------------------------------ # + + def get_optmizer(self) -> Optimizer: + return torch.optim.Adam(self.model.parameters(), lr=self.cfg.train.init_lr) + + def get_criterion(self) -> nn.MSELoss: + return nn.MSELoss() + + def save_ckpt(self, event: int) -> str: + """Save model checkpoint with a stats sidecar for reproduce_prometheus.py compare.""" + ckpt_path = super().save_ckpt(event) + + # Write cumulative normalization stats alongside the checkpoint + stats = compute_cumulative_stats(self._dfs_seen, self._all_cols) + sidecar = Path(ckpt_path).with_suffix(".stats.json") + payload = {k: [float(mu), float(std)] for k, (mu, std) in stats.items()} + with open(sidecar, "w") as f: + json.dump(payload, f, indent=2) + print(f" stats sidecar -> {sidecar}") + return ckpt_path + + def update_data_stream(self) -> None: + """Advance to the next training CSV with cumulative normalization.""" + self._dispose_current_loaders() + + csv_idx = self.task_counter % len(self.train_dfs) + current_df = self.train_dfs[csv_idx] + self._dfs_seen.append(current_df) + + # Recompute cumulative stats from all CSVs seen so far + stats = compute_cumulative_stats(self._dfs_seen, self._all_cols) + + label = current_df.attrs.get("source", f"csv{csv_idx:02d}") + print( + f"Prometheus: loading CSV {csv_idx + 1}/{len(self.train_dfs)} " + f"[{label}] (stats from {len(self._dfs_seen)} file(s))" + ) + + # Build windowed dataset for the current CSV + X, Y = make_sequence_windows( + [current_df], + self.FEATURE_COLS, + self.TARGET_COLS, + self.SEQUENCE_LENGTH, + stats, + forecast_horizon=self.FORECAST_HORIZON, + ) + + # 80/20 temporal split + # This may need to be modified since it biases the start. + n = len(X) + n_val = max(1, int(n * self.VAL_RATIO)) + n_train = n - n_val + + ds_train = SequenceDataset(X[:n_train], Y[:n_train]) + ds_val = SequenceDataset(X[n_train:], Y[n_train:]) + self._task_datasets.append((ds_train, ds_val)) + + self._cur_train_loader = self._make_loader(ds_train, shuffle=True) + self._cur_val_loader = self._make_loader(ds_val, shuffle=False) + + self.task_counter += 1 + + def get_cur_data_loaders( + self, + ) -> Tuple[DataLoader, DataLoader]: + return self._cur_train_loader, self._cur_val_loader + + def get_hist_data_loaders( + self, + ) -> Tuple[Optional[DataLoader], Optional[DataLoader]]: + """Return loaders over all prior task datasets. ``(None, None)`` if no history.""" + if self.task_counter <= 1: + return None, None + + prior = self._task_datasets[:-1] + hist_train: ConcatDataset = ConcatDataset([ds[0] for ds in prior]) + hist_val: ConcatDataset = ConcatDataset([ds[1] for ds in prior]) + + return ( + self._make_loader(hist_train, shuffle=True), # type: ignore[arg-type] + self._make_loader(hist_val, shuffle=False), # type: ignore[arg-type] + ) diff --git a/examples/prometheus_torch/prometheus.toml b/examples/prometheus_torch/prometheus.toml new file mode 100644 index 0000000..7106d3c --- /dev/null +++ b/examples/prometheus_torch/prometheus.toml @@ -0,0 +1,32 @@ +seed = 1337 +device = "auto" + +[model] +name = "prometheus_torch" +pretrained_path = "examples/prometheus_torch/output/prometheus_torch/checkpoints/torch_case_00.pt" +max_ckpts = 1 +ckpts_path = "examples/prometheus_torch/output/apeiron/" + +[data] +name = "prometheus_torch" +path = "examples/prometheus_torch/data" + +[train] +batch_size = 64 +num_workers = 0 +init_lr = 0.001 +max_iter = 600 + +[drift_detection] +detector_name = "ADWINDetector" +detection_interval = 1 +adwin_delta = 0.05 +metric_index = 0 +max_stream_updates = 23 + +[continual_learning] +update_mode = "base" + +[logging] +backend = "wandb" +experiment_name = "prometheus" diff --git a/examples/prometheus_torch/reproduce_prometheus.py b/examples/prometheus_torch/reproduce_prometheus.py new file mode 100644 index 0000000..c841779 --- /dev/null +++ b/examples/prometheus_torch/reproduce_prometheus.py @@ -0,0 +1,784 @@ +""" +reproduce_prometheus.py + +Concise reimplementation of the Prometheus temporal-prediction model +(originally TemporalPredict), with Keras and PyTorch backends. + +Framework is auto-detected from file extension: + .keras / .h5 -> keras + .pt / .pth -> torch +(or override with --framework keras|torch) + +Modes: + # Eval a saved model (framework detected from extension) + python reproduce_prometheus.py eval --model ./data/npm1_pwr_model.keras + python reproduce_prometheus.py eval --model ./output/reproduce_prometheus/retrained.pt + + # Train a fresh model (framework detected from --save extension) + python reproduce_prometheus.py train --save ./output/reproduce_prometheus/retrained.keras + python reproduce_prometheus.py train --save ./output/reproduce_prometheus/retrained.pt + + # Eval baseline then train + python reproduce_prometheus.py both --model ./data/npm1_pwr_model.keras \\ + --save ./output/reproduce_prometheus/retrained.pt + + # Overlay a Keras baseline and a PyTorch reproduction on the test set + python reproduce_prometheus.py compare --model ./data/npm1_pwr_model.keras \\ + --torch-model ./output/reproduce_prometheus/retrained.pt + +Training uses a continual / per-case scheme: one fit() per training file, +with normalization statistics recomputed from cases seen so far (0..N) at +the start of each case. Each per-case checkpoint is written alongside a +JSON stats sidecar so eval/compare can reproduce the exact normalization +used at training time. +""" + +import argparse +import json +import os +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import tensorflow as tf +import torch +import torch.nn as nn +import torch.optim as optim +from sklearn.metrics import mean_absolute_error, r2_score + + +# ---------- Config -------------------------------------------------------- + + +@dataclass +class Config: + train_dir: str = "./data/train" + test_dir: str = "./data/test" + out_dir: str = "./output/prometheus_torch/" + sequence_length: int = 10 + epochs: int = 1 + batch_size: int = 64 + learning_rate: float = 1e-3 + feature_cols: List[str] = field( + default_factory=lambda: [ + "NRAD_RX_REG_POS", + "NRAD_RX_SHIM1_POS", + "NRAD_RX_SHIM2_POS", + "total_rod_position", + "NRAD_RX_PERIOD_Inverse", + "NRAD_RX_REG_POS_dt", + "NRAD_RX_REG_POS_dt2", + "NRAD_RX_SHIM1_POS_dt", + "NRAD_RX_SHIM1_POS_dt2", + "NRAD_RX_SHIM2_POS_dt", + "NRAD_RX_SHIM2_POS_dt2", + "NRAD_RX_NMP1_PWR_integral", + ] + ) + target_cols: List[str] = field(default_factory=lambda: ["NRAD_RX_NMP1_PWR"]) + + +# ---------- Data ---------------------------------------------------------- + + +def load_csvs(folder: str) -> List[pd.DataFrame]: + """Read every .csv in folder (sorted) and stash the filename stem on + df.attrs['source'] so downstream plots can label the data by when it + was generated (file names encode the collection date). + """ + files = sorted(f for f in os.listdir(folder) if f.endswith(".csv")) + dfs = [] + for f in files: + df = pd.read_csv(os.path.join(folder, f)) + df.attrs["source"] = os.path.splitext(f)[0] + dfs.append(df) + return dfs + + +def df_label(df: pd.DataFrame, fallback: str = "") -> str: + """Return the source-file stem recorded by load_csvs, or a fallback.""" + return df.attrs.get("source", fallback) + + +def compute_train_stats(dfs: List[pd.DataFrame], cols: List[str]) -> dict: + """Return {col: (mu, std)} from concatenated training dataframes.""" + full = pd.concat([df[cols] for df in dfs], ignore_index=True) + return {c: (float(full[c].mean()), float(full[c].std())) for c in cols} + + +def normalize(df: pd.DataFrame, stats: dict) -> pd.DataFrame: + out = df.copy() + for c, (mu, std) in stats.items(): + if c in out.columns: + out[c] = (out[c] - mu) / std + return out + + +def make_windows( + dfs: List[pd.DataFrame], + feature_cols: List[str], + target_cols: List[str], + seq_len: int, + stats: dict, +) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """Sliding-window sequences. Returns per-file lists of arrays: + X[k] has shape (n_windows, seq_len, n_features), + Y[k] has shape (n_windows, seq_len, n_targets). + """ + X_list, Y_list = [], [] + for df in dfs: + if len(df) <= seq_len: + continue + data = normalize(df[feature_cols + target_cols], stats) + feat = data[feature_cols].to_numpy(dtype=np.float32) + targ = data[target_cols].to_numpy(dtype=np.float32) + n = len(data) - seq_len + X = np.stack([feat[i : i + seq_len] for i in range(n)]) + Y = np.stack([targ[i : i + seq_len] for i in range(n)]) + X_list.append(X) + Y_list.append(Y) + return X_list, Y_list + + +# ---------- Models (Keras + Torch) ---------------------------------------- + + +def build_keras_model( + seq_len: int, n_features: int, n_targets: int, lr: float +) -> tf.keras.Model: + """LSTM(128) -> Dropout -> LSTM(64) -> Dropout -> LSTM(32) -> TimeDistributed(Dense). + Output shape: (batch, seq_len, n_targets). + """ + inp = tf.keras.Input(shape=(seq_len, n_features)) + x = tf.keras.layers.LSTM(128, return_sequences=True)(inp) + x = tf.keras.layers.Dropout(0.1)(x) + x = tf.keras.layers.LSTM(64, return_sequences=True)(x) + x = tf.keras.layers.Dropout(0.1)(x) + x = tf.keras.layers.LSTM(32, return_sequences=True)(x) + out = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(n_targets))(x) + model = tf.keras.Model(inp, out) + model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=lr), + loss=tf.keras.losses.MeanSquaredError(), + ) + return model + + +class TorchTemporalModel(nn.Module): + """PyTorch mirror of the Keras architecture above. A Linear layer applied + to the (batch, seq_len, hidden) tensor is equivalent to Keras' + TimeDistributed(Dense(...)). + """ + + def __init__(self, n_features: int, n_targets: int, dropout: float = 0.1): + super().__init__() + self.lstm1 = nn.LSTM(n_features, 128, batch_first=True) + self.drop1 = nn.Dropout(dropout) + self.lstm2 = nn.LSTM(128, 64, batch_first=True) + self.drop2 = nn.Dropout(dropout) + self.lstm3 = nn.LSTM(64, 32, batch_first=True) + self.head = nn.Linear(32, n_targets) + + def forward(self, x): + x, _ = self.lstm1(x) + x = self.drop1(x) + x, _ = self.lstm2(x) + x = self.drop2(x) + x, _ = self.lstm3(x) + return self.head(x) + + +# ---------- Framework dispatch -------------------------------------------- + + +def framework_from_path(path: str) -> str: + ext = os.path.splitext(path)[1].lower() + if ext in (".keras", ".h5"): + return "keras" + if ext in (".pt", ".pth"): + return "torch" + raise ValueError( + f"Cannot detect framework from extension: {ext!r}. " + f"Use --framework keras|torch to override." + ) + + +def load_saved_model(path: str, n_features: int, n_targets: int, framework: str): + if framework == "keras": + return tf.keras.models.load_model(path) + if framework == "torch": + model = TorchTemporalModel(n_features, n_targets) + state = torch.load(path, map_location="cpu") + model.load_state_dict(state) + model.eval() + return model + raise ValueError(f"Unknown framework: {framework}") + + +def save_trained_model(model, path: str, framework: str) -> None: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + if framework == "keras": + model.save(path) + elif framework == "torch": + torch.save(model.state_dict(), path) + else: + raise ValueError(f"Unknown framework: {framework}") + + +def predict_np(model, X: np.ndarray, framework: str) -> np.ndarray: + """Return model predictions as a numpy array of shape (n_windows, seq_len, n_targets).""" + if framework == "keras": + return model.predict(X, verbose=0) + if framework == "torch": + model.eval() + with torch.no_grad(): + t = torch.from_numpy(X.astype(np.float32)) + return model(t).cpu().numpy() + raise ValueError(f"Unknown framework: {framework}") + + +# ---------- Stats sidecar (per-checkpoint persistence) ------------------- + + +def stats_sidecar_path(model_path: str) -> str: + base, _ = os.path.splitext(model_path) + return base + ".stats.json" + + +def save_stats(stats: dict, model_path: str) -> str: + path = stats_sidecar_path(model_path) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + payload = {k: [float(mu), float(std)] for k, (mu, std) in stats.items()} + with open(path, "w") as f: + json.dump(payload, f, indent=2) + return path + + +def load_stats( + model_path: str, required_cols: Optional[List[str]] = None +) -> Optional[dict]: + """Load a stats sidecar next to model_path. Returns None if missing.""" + path = stats_sidecar_path(model_path) + if not os.path.exists(path): + return None + with open(path) as f: + payload = json.load(f) + stats = {k: (float(v[0]), float(v[1])) for k, v in payload.items()} + if required_cols is not None: + missing = [c for c in required_cols if c not in stats] + if missing: + raise ValueError(f"Stats sidecar {path} is missing columns: {missing}") + return stats + + +def stats_for_model( + model_path: str, train_dfs: List[pd.DataFrame], cols: List[str] +) -> dict: + """Return per-checkpoint stats if a sidecar exists, else fall back to + global training stats computed over all train_dfs. Logs which path was taken. + """ + stats = load_stats(model_path, required_cols=cols) + if stats is not None: + print(f" using stats sidecar: {stats_sidecar_path(model_path)}") + return stats + print(f" no stats sidecar for {model_path}; falling back to global train stats") + return compute_train_stats(train_dfs, cols) + + +# ---------- Denormalize + plot helpers ------------------------------------ + + +def denorm_last_step( + pred_norm: np.ndarray, stats: dict, target_cols: List[str] +) -> pd.DataFrame: + """Take the last timestep of each window and invert the train-stats normalization.""" + last = pred_norm[:, -1, :] + df = pd.DataFrame(last, columns=target_cols) + for c in target_cols: + mu, std = stats[c] + df[c] = df[c] * std + mu + return df + + +def plot_pred_vs_truth( + preds: Union[pd.DataFrame, Dict[str, pd.DataFrame]], + truth: pd.DataFrame, + target_cols: List[str], + title: str, + save_path: str, +) -> None: + """preds can be a single DataFrame or a dict {label: DataFrame} for overlays.""" + if isinstance(preds, pd.DataFrame): + preds = {"Prediction": preds} + colors = [ + "tab:red", + "tab:blue", + "tab:green", + "tab:orange", + "tab:purple", + "tab:brown", + "tab:pink", + "tab:cyan", + ] + n = len(target_cols) + fig, axes = plt.subplots(n, 1, figsize=(10, 3 * n), squeeze=False) + for i, var in enumerate(target_cols): + ax = axes[i, 0] + truth.loc[:, var] = np.where(truth[var].values == 0, 1e-2, truth[var]) + ax.plot( + truth.index, + truth[var].values, + color="black", + linewidth=1.5, + label="Ground Truth", + ) + for j, (label, pred) in enumerate(preds.items()): + pred.loc[:, var] = np.where(pred[var].values == 0, 1e-2, pred[var]) + ax.plot( + pred.index, + pred[var].values, + color=colors[j % len(colors)], + linewidth=1.2, + label=label, + ) + #ax.set_yscale('log') + #ax.set_ylim(bottom=1e-20) + ax.set_xlabel("time step") + ax.set_ylabel(var) + ax.legend(loc="best") + ax.grid(True, alpha=0.3) + fig.suptitle(title) + fig.tight_layout() + os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) + fig.savefig(save_path, dpi=150) + plt.close(fig) + + +# ---------- Evaluation ---------------------------------------------------- + + +def evaluate( + model, + test_dfs: List[pd.DataFrame], + X_test: List[np.ndarray], + stats: dict, + cfg: Config, + tag: str, + framework: str, +) -> None: + """Predict on each test file, compute metrics, save plots.""" + out_subdir = os.path.join(cfg.out_dir, f"eval_{tag}") + os.makedirs(out_subdir, exist_ok=True) + print(f"\n=== Evaluation [{tag}] ({framework}) ===") + for i, (df, X) in enumerate(zip(test_dfs, X_test)): + label = df_label(df, fallback=f"test{i:02d}") + pred_norm = predict_np(model, X, framework) + pred = denorm_last_step(pred_norm, stats, cfg.target_cols).reset_index( + drop=True + ) + truth = df[cfg.target_cols].iloc[cfg.sequence_length :].reset_index(drop=True) + m = min(len(pred), len(truth)) + pred, truth = pred.iloc[:m], truth.iloc[:m] + r2 = r2_score(truth, pred) + mae = mean_absolute_error(truth, pred) + print(f" test {i} [{label}]: R2={r2:.4f} MAE={mae:.3f}") + plot_pred_vs_truth( + pred, + truth, + cfg.target_cols, + title=f"{tag} — {label} (R²={r2:.3f}, MAE={mae:.3f})", + save_path=os.path.join(out_subdir, f"test{i:02d}_{label}.png"), + ) + + +# ---------- Torch per-case training loop --------------------------------- + + +def torch_fit_case( + model: nn.Module, + X: np.ndarray, + Y: np.ndarray, + epochs: int, + batch_size: int, + lr: float, + val_split: float = 0.2, + patience: int = 8, +) -> nn.Module: + """Mirror of keras model.fit(...) for one case, with early stopping and + best-weight restoration on val loss. + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + + n = len(X) + n_val = max(1, int(n * val_split)) + perm = np.random.permutation(n) + val_idx, tr_idx = perm[:n_val], perm[n_val:] + + X_tr = torch.from_numpy(X[tr_idx]).to(device) + Y_tr = torch.from_numpy(Y[tr_idx]).to(device) + X_val = torch.from_numpy(X[val_idx]).to(device) + Y_val = torch.from_numpy(Y[val_idx]).to(device) + + optimizer = optim.Adam(model.parameters(), lr=lr) + loss_fn = nn.MSELoss() + + best_val = float("inf") + best_state = {k: v.detach().clone() for k, v in model.state_dict().items()} + bad_epochs = 0 + + for epoch in range(epochs): + model.train() + idx = torch.randperm(len(X_tr), device=device) + total = 0.0 + for start in range(0, len(idx), batch_size): + b = idx[start : start + batch_size] + optimizer.zero_grad() + pred = model(X_tr[b]) + loss = loss_fn(pred, Y_tr[b]) + loss.backward() + optimizer.step() + total += loss.item() * len(b) + train_loss = total / len(X_tr) + + model.eval() + with torch.no_grad(): + val_loss = loss_fn(model(X_val), Y_val).item() + + print( + f" Epoch {epoch + 1:3d}/{epochs} train_loss={train_loss:.4f} val_loss={val_loss:.4f}" + ) + + if val_loss < best_val - 1e-6: + best_val = val_loss + best_state = {k: v.detach().clone() for k, v in model.state_dict().items()} + bad_epochs = 0 + else: + bad_epochs += 1 + if bad_epochs >= patience: + print(f" Early stopping at epoch {epoch + 1}") + break + + model.load_state_dict(best_state) + model.to("cpu") + return model + + +# ---------- Training ------------------------------------------------------ + + +def train( + cfg: Config, + train_dfs: List[pd.DataFrame], + test_dfs: List[pd.DataFrame], + save_path: str, + framework: str, +): + """Per-case training with stats-so-far normalization. + + At the start of each case N: + * stats_N = compute_train_stats(train_dfs[:N+1]) + * The case N windows and the test windows are (re)built with stats_N. + * The model (weights carried over from case N-1) is fit on case N. + * Predictions are denormalized with stats_N for the post-fit plot and + the test eval pass. + * A per-case checkpoint is saved alongside a stats sidecar so later + eval/compare runs can reproduce the same normalization. + + This is an intentional continual-learning setup where the input scale + drifts between cases, and the "final" model is just the model after the + last case (saved to save_path with its own stats sidecar). + """ + all_cols = cfg.feature_cols + cfg.target_cols + n_features, n_targets = len(cfg.feature_cols), len(cfg.target_cols) + + if framework == "keras": + model = build_keras_model( + cfg.sequence_length, n_features, n_targets, cfg.learning_rate + ) + model.summary() + else: + model = TorchTemporalModel(n_features, n_targets) + print(model) + + train_plot_dir = os.path.join(cfg.out_dir, "train_fits") + ckpt_dir = os.path.join(cfg.out_dir, "checkpoints") + os.makedirs(ckpt_dir, exist_ok=True) + ckpt_ext = os.path.splitext(save_path)[1] or ( + ".keras" if framework == "keras" else ".pt" + ) + + stats_n = None + for case_idx in range(len(train_dfs)): + # Stats use only cases seen so far (inclusive of the current one) + stats_n = compute_train_stats(train_dfs[: case_idx + 1], all_cols) + + # Rebuild just this case's windows with stats_n + Xc_list, Yc_list = make_windows( + [train_dfs[case_idx]], + cfg.feature_cols, + cfg.target_cols, + cfg.sequence_length, + stats_n, + ) + if not Xc_list: + print(f"\n--- Case {case_idx}: empty after windowing, skipping ---") + continue + Xc, Yc = Xc_list[0], Yc_list[0] + + # Rebuild test windows with the same stats so the eval pass is consistent + X_test, _ = make_windows( + test_dfs, cfg.feature_cols, cfg.target_cols, cfg.sequence_length, stats_n + ) + + case_label = df_label(train_dfs[case_idx], fallback=f"case{case_idx:02d}") + print( + f"\n--- Case {case_idx} [{case_label}] ({framework}): X={Xc.shape}, Y={Yc.shape} " + f"(stats from {case_idx + 1} case(s)) ---" + ) + if framework == "keras": + early_stop = tf.keras.callbacks.EarlyStopping( + monitor="val_loss", patience=8, restore_best_weights=True + ) + model.fit( + Xc, + Yc, + epochs=cfg.epochs, + batch_size=cfg.batch_size, + validation_split=0.2, + callbacks=[early_stop], + verbose=1, + ) + else: + model = torch_fit_case( + model, Xc, Yc, cfg.epochs, cfg.batch_size, cfg.learning_rate + ) + + # Single post-fit plot: predictions on the case we just trained on + pred_norm = predict_np(model, Xc, framework) + pred = denorm_last_step(pred_norm, stats_n, cfg.target_cols).reset_index( + drop=True + ) + truth = ( + train_dfs[case_idx][cfg.target_cols] + .iloc[cfg.sequence_length :] + .reset_index(drop=True) + ) + m = min(len(pred), len(truth)) + plot_pred_vs_truth( + pred.iloc[:m], + truth.iloc[:m], + cfg.target_cols, + title=f"train case {case_idx} — {case_label} — post-fit ({framework})", + save_path=os.path.join( + train_plot_dir, f"{framework}_case_{case_idx:02d}_{case_label}.png" + ), + ) + + # Per-case checkpoint + stats sidecar + ckpt_path = os.path.join(ckpt_dir, f"{framework}_case_{case_idx:02d}{ckpt_ext}") + save_trained_model(model, ckpt_path, framework) + save_stats(stats_n, ckpt_path) + print(f" checkpoint -> {ckpt_path}") + print(f" stats -> {stats_sidecar_path(ckpt_path)}") + + evaluate( + model, + test_dfs, + X_test, + stats_n, + cfg, + tag=f"{framework}_case_{case_idx:02d}", + framework=framework, + ) + + # Final save uses the last case's stats + save_trained_model(model, save_path, framework) + if stats_n is not None: + save_stats(stats_n, save_path) + print(f"\nSaved final {framework} model to {save_path}") + if stats_n is not None: + print(f"Saved final stats sidecar to {stats_sidecar_path(save_path)}") + return model + + +# ---------- Compare mode -------------------------------------------------- + + +def compare( + cfg: Config, + train_dfs: List[pd.DataFrame], + test_dfs: List[pd.DataFrame], + keras_path: str, + torch_paths: List[str], +) -> None: + """Load a Keras baseline and one or more PyTorch models, then overlay their + predictions against the ground truth on each test file. Each model uses + its own stats sidecar (if present) for normalization and denormalization. + Labels in plots and titles use the .pt filename stem. + """ + n_features, n_targets = len(cfg.feature_cols), len(cfg.target_cols) + all_cols = cfg.feature_cols + cfg.target_cols + + # --- Keras baseline --- + print(f"Loading keras model: {keras_path}") + keras_model = load_saved_model(keras_path, n_features, n_targets, "keras") + k_stats = stats_for_model(keras_path, train_dfs, all_cols) + X_test_k, _ = make_windows( + test_dfs, cfg.feature_cols, cfg.target_cols, cfg.sequence_length, k_stats + ) + keras_stem = os.path.splitext(os.path.basename(keras_path))[0] + + # --- Torch models --- + torch_entries: List[Tuple[str, object, dict, List[np.ndarray]]] = [] + for tp in torch_paths: + stem = os.path.splitext(os.path.basename(tp))[0] + print(f"Loading torch model: {tp}") + model = load_saved_model(tp, n_features, n_targets, "torch") + stats = stats_for_model(tp, train_dfs, all_cols) + X_test_t, _ = make_windows( + test_dfs, cfg.feature_cols, cfg.target_cols, cfg.sequence_length, stats + ) + torch_entries.append((stem, model, stats, X_test_t)) + + stems = [keras_stem] + [e[0] for e in torch_entries] + pair_name = "__vs__".join(stems) + out_subdir = os.path.join(cfg.out_dir, "compare", pair_name) + os.makedirs(out_subdir, exist_ok=True) + print(f" writing comparison plots to {out_subdir}") + print(f"\n=== Comparison: {' vs '.join(stems)} ===") + + for i, df in enumerate(test_dfs): + label = df_label(df, fallback=f"test{i:02d}") + truth = df[cfg.target_cols].iloc[cfg.sequence_length :].reset_index(drop=True) + + # Keras prediction + k_pred = denorm_last_step( + predict_np(keras_model, X_test_k[i], "keras"), k_stats, cfg.target_cols + ).reset_index(drop=True) + + # Torch predictions + overlay: Dict[str, pd.DataFrame] = {keras_stem: k_pred} + min_len = min(len(k_pred), len(truth)) + metrics_parts: List[str] = [] + + k_r2 = r2_score(truth.iloc[:min_len], k_pred.iloc[:min_len]) + k_mae = mean_absolute_error(truth.iloc[:min_len], k_pred.iloc[:min_len]) + metrics_parts.append(f"{keras_stem} R²={k_r2:.3f}") + print( + f" test {i} [{label}]: {keras_stem} R2={k_r2:.4f} MAE={k_mae:.3f}", end="" + ) + + for stem, model, stats, X_test_t in torch_entries: + t_pred = denorm_last_step( + predict_np(model, X_test_t[i], "torch"), stats, cfg.target_cols + ).reset_index(drop=True) + overlay[stem] = t_pred + min_len = min(min_len, len(t_pred)) + t_r2 = r2_score(truth.iloc[:min_len], t_pred.iloc[:min_len]) + t_mae = mean_absolute_error(truth.iloc[:min_len], t_pred.iloc[:min_len]) + metrics_parts.append(f"{stem} R²={t_r2:.3f}") + print(f" | {stem} R2={t_r2:.4f} MAE={t_mae:.3f}", end="") + print() + + # Trim all to common length + truth_trimmed = truth.iloc[:min_len] + overlay_trimmed = {k: v.iloc[:min_len] for k, v in overlay.items()} + + plot_pred_vs_truth( + overlay_trimmed, + truth_trimmed, + cfg.target_cols, + title=f"compare — {label} ({', '.join(metrics_parts)})", + save_path=os.path.join(out_subdir, f"test{i:02d}_{label}.png"), + ) + + +# ---------- Main ---------------------------------------------------------- + + +def resolve_framework(explicit: str, path: str) -> str: + return explicit if explicit != "auto" else framework_from_path(path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("mode", choices=["eval", "train", "both", "compare"]) + parser.add_argument( + "--model", + default="./data/npm1_pwr_model.keras", + help="Saved model path for eval/both, or the Keras baseline for compare.", + ) + parser.add_argument( + "--save", + default="./output/reproduce_prometheus/retrained.keras", + help="Where to save the retrained model (for train/both). " + "Extension determines framework (.keras/.h5 or .pt/.pth).", + ) + parser.add_argument( + "--torch-model", + nargs="+", + default=None, + help="Path(s) to PyTorch model(s) for compare mode. " + "Multiple paths overlay all models on the same plot.", + ) + parser.add_argument( + "--framework", + choices=["keras", "torch", "auto"], + default="auto", + help="Override framework detection.", + ) + parser.add_argument("--epochs", type=int, default=None) + args = parser.parse_args() + + cfg = Config() + if args.epochs is not None: + cfg.epochs = args.epochs + os.makedirs(cfg.out_dir, exist_ok=True) + + print("Loading data ...") + train_dfs = load_csvs(cfg.train_dir) + test_dfs = load_csvs(cfg.test_dir) + print(f" {len(train_dfs)} training files, {len(test_dfs)} test files") + + n_features, n_targets = len(cfg.feature_cols), len(cfg.target_cols) + all_cols = cfg.feature_cols + cfg.target_cols + + if args.mode in ("eval", "both"): + framework = resolve_framework(args.framework, args.model) + print(f"Loading {framework} model: {args.model}") + model = load_saved_model(args.model, n_features, n_targets, framework) + stats = stats_for_model(args.model, train_dfs, all_cols) + X_test, _ = make_windows( + test_dfs, cfg.feature_cols, cfg.target_cols, cfg.sequence_length, stats + ) + evaluate( + model, + test_dfs, + X_test, + stats, + cfg, + tag=f"{framework}_baseline", + framework=framework, + ) + + if args.mode in ("train", "both"): + framework = resolve_framework(args.framework, args.save) + train(cfg, train_dfs, test_dfs, args.save, framework) + + if args.mode == "compare": + if not args.torch_model: + parser.error("compare mode requires --torch-model") + compare( + cfg, + train_dfs, + test_dfs, + keras_path=args.model, + torch_paths=args.torch_model, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/prometheus_torch/utils.py b/examples/prometheus_torch/utils.py new file mode 100644 index 0000000..20047b0 --- /dev/null +++ b/examples/prometheus_torch/utils.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import os +from typing import Dict, List, Tuple + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader, Dataset + + +def load_csvs(folder: str) -> List[pd.DataFrame]: + """Read every .csv in *folder* (sorted) and stash the filename stem on + ``df.attrs['source']`` for downstream labelling. + """ + files = sorted(f for f in os.listdir(folder) if f.endswith(".csv")) + dfs: List[pd.DataFrame] = [] + for f in files: + df = pd.read_csv(os.path.join(folder, f)) + df.attrs["source"] = os.path.splitext(f)[0] + dfs.append(df) + return dfs + + +def compute_cumulative_stats( + dfs: List[pd.DataFrame], cols: List[str] +) -> Dict[str, Tuple[float, float]]: + """Return ``{col: (mean, std)}`` computed over the concatenation of *dfs*.""" + full = pd.concat([df[cols] for df in dfs], ignore_index=True) + return {c: (float(full[c].mean()), float(full[c].std())) for c in cols} + + +def normalize(df: pd.DataFrame, stats: Dict[str, Tuple[float, float]]) -> pd.DataFrame: + """Z-score normalize *df* using pre-computed *stats*.""" + out = df.copy() + for c, (mu, std) in stats.items(): + if c in out.columns: + out[c] = (out[c] - mu) / std + return out + + +def make_sequence_windows( + dfs: List[pd.DataFrame], + feature_cols: List[str], + target_cols: List[str], + seq_len: int, + stats: Dict[str, Tuple[float, float]], + forecast_horizon: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Sliding-window sequences with full-sequence targets. + + Returns + ------- + X : Tensor of shape ``(N, seq_len, n_features)`` + Y : Tensor of shape ``(N, seq_len, n_targets)`` + """ + all_cols = feature_cols + target_cols + x_parts: List[np.ndarray] = [] + y_parts: List[np.ndarray] = [] + + for df in dfs: + if len(df) < seq_len +forecast_horizon: + continue + data = normalize(df[all_cols], stats) + feat = data[feature_cols].to_numpy(dtype=np.float32) + targ = data[target_cols].to_numpy(dtype=np.float32) + n = len(data) - seq_len - forecast_horizon + 1 + x_parts.append(np.stack([feat[i : i + seq_len] for i in range(n)])) + if forecast_horizon == 0: + y_parts.append(np.stack([targ[i : i + seq_len] for i in range(n)])) + else: + y_parts.append( + np.stack( + [ + targ[ + i + seq_len: + i + seq_len + forecast_horizon + ] + for i in range(n) + ] + ) + ) + + if not x_parts: + raise ValueError( + "No sequences created. Check data path and column names " + f"(features={feature_cols}, targets={target_cols})." + ) + + X = torch.from_numpy(np.concatenate(x_parts)) + Y = torch.from_numpy(np.concatenate(y_parts)) + return X, Y + + +class SequenceDataset(Dataset): + """Wraps pre-built ``(X, Y)`` tensors.""" + + def __init__(self, x: torch.Tensor, y: torch.Tensor): + self.x = x + self.y = y + + def __len__(self) -> int: + return len(self.x) + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + return self.x[idx], self.y[idx] + + +def make_loader( + ds: Dataset, + batch_size: int, + shuffle: bool, + num_workers: int = 0, + pin_memory: bool = False, +) -> DataLoader: + """Build a DataLoader from a Dataset.""" + kwargs: dict = dict(batch_size=batch_size, shuffle=shuffle, drop_last=False) + if num_workers > 0: + kwargs.update( + dict( + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=True, + prefetch_factor=2, + ) + ) + return DataLoader(ds, **kwargs) # type: ignore[arg-type] diff --git a/examples/utils.py b/examples/utils.py index 1a1c413..77351df 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -15,6 +15,10 @@ def get_example(cfg: Config) -> BaseModelHarness: from examples.imagenet.model import IMAGENET_VISION return IMAGENET_VISION(cfg=cfg) + elif cfg.data.name == "prometheus_torch": + from examples.prometheus_torch.model import PrometheusHarness + + return PrometheusHarness(cfg=cfg) else: raise NotImplementedError( f"Example for dataset {cfg.data.name} is not implemented."