"""Battery-grouped sequence-model evaluation with train-only preprocessing.""" from __future__ import annotations import random from collections.abc import Callable import numpy as np import pandas as pd from sklearn.model_selection import GroupShuffleSplit from src.evaluation.metrics import regression_metrics from src.evaluation.protocol import grouped_train_val_test_folds TORCH_MODEL_IDS = ( "vanilla_lstm", "bidirectional_lstm", "gru", "attention_lstm", "battery_gpt", "temporal_fusion_transformer", "vae_lstm", ) TF_MODEL_IDS = ("itransformer", "physics_itransformer", "dynamic_graph_itransformer") class SequenceStandardizer: """Median-impute and standardize sequence channels using training rows only.""" def fit(self, X: np.ndarray) -> "SequenceStandardizer": flat = np.asarray(X, dtype=float).reshape(-1, X.shape[-1]) self.median_ = np.nanmedian(flat, axis=0) filled = np.where(np.isnan(flat), self.median_, flat) self.mean_ = filled.mean(axis=0) self.scale_ = filled.std(axis=0) self.scale_[self.scale_ == 0] = 1.0 return self def transform(self, X: np.ndarray) -> np.ndarray: values = np.asarray(X, dtype=float) filled = np.where(np.isnan(values), self.median_, values) return ((filled - self.mean_) / self.scale_).astype(np.float32) def _set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) try: import torch torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) except ImportError: pass def _torch_factories(n_features: int) -> dict[str, Callable[[], object]]: from src.models.deep.lstm import AttentionLSTM, BidirectionalLSTM, GRUModel, VanillaLSTM from src.models.deep.transformer import BatteryGPT, TemporalFusionTransformer from src.models.deep.vae_lstm import VAE_LSTM return { "vanilla_lstm": lambda: VanillaLSTM(n_features, hidden_dim=64, n_layers=2), "bidirectional_lstm": lambda: BidirectionalLSTM(n_features, hidden_dim=64, n_layers=2), "gru": lambda: GRUModel(n_features, hidden_dim=64, n_layers=2), "attention_lstm": lambda: AttentionLSTM(n_features, hidden_dim=64, n_layers=2), "battery_gpt": lambda: BatteryGPT(n_features, d_model=64, n_heads=4, n_layers=2), "temporal_fusion_transformer": lambda: TemporalFusionTransformer( n_features, d_model=64, n_heads=4, n_layers=2 ), "vae_lstm": lambda: VAE_LSTM( n_features, seq_len=64, hidden_dim=64, latent_dim=16, n_layers=2 ), } def _train_torch_model( model_id: str, model: object, X_train: np.ndarray, y_train: np.ndarray, X_val: np.ndarray, y_val: np.ndarray, X_test: np.ndarray, *, max_epochs: int, patience: int, batch_size: int, device: str, ) -> tuple[np.ndarray, np.ndarray, int]: import torch from torch.utils.data import DataLoader, TensorDataset from src.models.deep.lstm import train_loop from src.models.deep.vae_lstm import train_vae y_mean = float(y_train.mean()) y_scale = float(y_train.std()) or 1.0 normalize = lambda values: ((values - y_mean) / y_scale).astype(np.float32) train_loader = DataLoader(TensorDataset( torch.from_numpy(X_train), torch.from_numpy(normalize(y_train)) ), batch_size=batch_size, shuffle=True) val_loader = DataLoader(TensorDataset( torch.from_numpy(X_val), torch.from_numpy(normalize(y_val)) ), batch_size=batch_size, shuffle=False) if model_id == "vae_lstm": history = train_vae( model, train_loader, val_loader, max_epochs=max_epochs, patience=patience, device=device, warmup_epochs=min(30, max_epochs // 3), ) else: history = train_loop( model, train_loader, val_loader, max_epochs=max_epochs, patience=patience, device=device, ) model.eval() def predict(values: np.ndarray) -> np.ndarray: outputs = [] loader = DataLoader(TensorDataset(torch.from_numpy(values)), batch_size=batch_size) with torch.no_grad(): for (xb,) in loader: raw = model(xb.to(device)) if isinstance(raw, dict): raw = raw["health_pred"] outputs.append(raw.detach().cpu().numpy().reshape(-1)) return np.concatenate(outputs) * y_scale + y_mean epochs = len(history.get("train_losses", [])) return predict(X_train), predict(X_test), epochs def _train_tensorflow_model( model_id: str, X_train: np.ndarray, y_train: np.ndarray, X_val: np.ndarray, y_val: np.ndarray, X_test: np.ndarray, *, max_epochs: int, patience: int, batch_size: int, ) -> tuple[np.ndarray, np.ndarray, int]: import tensorflow as tf from src.models.deep.itransformer import ( build_dynamic_graph_itransformer, build_itransformer, build_physics_itransformer, ) builders = { "itransformer": build_itransformer, "physics_itransformer": build_physics_itransformer, "dynamic_graph_itransformer": build_dynamic_graph_itransformer, } y_mean = float(y_train.mean()) y_scale = float(y_train.std()) or 1.0 train_target = ((y_train - y_mean) / y_scale).astype(np.float32) val_target = ((y_val - y_mean) / y_scale).astype(np.float32) model = builders[model_id](X_train.shape[1], X_train.shape[2], d_model=32, n_heads=4, n_blocks=2) if model_id == "physics_itransformer": model.compile( optimizer=tf.keras.optimizers.Adam(1e-3), loss={"soh_ml": "mae", "soh_phy": "mae"}, loss_weights={"soh_ml": 1.0, "soh_phy": model.physics_loss_weight}, ) fit_y = {"soh_ml": train_target, "soh_phy": train_target} val_y = {"soh_ml": val_target, "soh_phy": val_target} else: model.compile(optimizer=tf.keras.optimizers.Adam(1e-3), loss="mae") fit_y, val_y = train_target, val_target history = model.fit( X_train, fit_y, validation_data=(X_val, val_y), epochs=max_epochs, batch_size=batch_size, verbose=0, callbacks=[tf.keras.callbacks.EarlyStopping( monitor="val_loss", patience=patience, restore_best_weights=True )], ) def predict(values: np.ndarray) -> np.ndarray: raw = model.predict(values, batch_size=batch_size, verbose=0) if isinstance(raw, list): raw = raw[0] return np.asarray(raw).reshape(-1) * y_scale + y_mean result = (predict(X_train), predict(X_test), len(history.history["loss"])) tf.keras.backend.clear_session() return result def run_grouped_sequence_benchmark( X: np.ndarray, index: pd.DataFrame, *, dataset_name: str, n_splits: int = 5, seeds: tuple[int, ...] = (17, 42, 2026), max_epochs: int = 200, patience: int = 20, batch_size: int = 64, model_ids: tuple[str, ...] | None = None, ) -> tuple[pd.DataFrame, pd.DataFrame]: """Evaluate all ten sequence models without using test data for stopping.""" if len(index) != len(X): raise ValueError("Sequence tensor and index have different row counts") requested = set(model_ids or (TORCH_MODEL_IDS + TF_MODEL_IDS)) unknown = requested.difference(TORCH_MODEL_IDS + TF_MODEL_IDS) if unknown: raise KeyError(f"Unknown sequence model IDs: {sorted(unknown)}") y = index["SoH"].to_numpy(dtype=float) groups = index["battery_id"].astype(str).to_numpy() available_splits = min(n_splits, np.unique(groups).size) if available_splits < 2: raise ValueError("At least two batteries are required") metric_rows: list[dict[str, object]] = [] prediction_rows: list[dict[str, object]] = [] for seed in seeds: _set_seed(seed) folds = grouped_train_val_test_folds( groups, n_splits=available_splits, validation_fraction=0.2, random_state=seed, ) for fold, (train_idx, val_idx, test_idx) in enumerate(folds, start=1): scaler = SequenceStandardizer().fit(X[train_idx]) X_train, X_val, X_test = ( scaler.transform(X[subset]) for subset in (train_idx, val_idx, test_idx) ) device = "cuda" if __import__("torch").cuda.is_available() else "cpu" for model_id, factory in _torch_factories(X.shape[-1]).items(): if model_id not in requested: continue print( f"[{dataset_name}] seed={seed} fold={fold}/{available_splits} " f"model={model_id} start", flush=True, ) _set_seed(seed) train_pred, test_pred, epochs = _train_torch_model( model_id, factory(), X_train, y[train_idx], X_val, y[val_idx], X_test, max_epochs=max_epochs, patience=patience, batch_size=batch_size, device=device, ) _append_results( metric_rows, prediction_rows, dataset_name, seed, fold, model_id, train_idx, test_idx, groups, index, y, train_pred, test_pred, epochs, ) print( f"[{dataset_name}] seed={seed} fold={fold}/{available_splits} " f"model={model_id} done epochs={epochs} mae={metric_rows[-1]['mae']:.4f}", flush=True, ) for model_id in TF_MODEL_IDS: if model_id not in requested: continue print( f"[{dataset_name}] seed={seed} fold={fold}/{available_splits} " f"model={model_id} start", flush=True, ) _set_seed(seed) train_pred, test_pred, epochs = _train_tensorflow_model( model_id, X_train, y[train_idx], X_val, y[val_idx], X_test, max_epochs=max_epochs, patience=patience, batch_size=batch_size, ) _append_results( metric_rows, prediction_rows, dataset_name, seed, fold, model_id, train_idx, test_idx, groups, index, y, train_pred, test_pred, epochs, ) print( f"[{dataset_name}] seed={seed} fold={fold}/{available_splits} " f"model={model_id} done epochs={epochs} mae={metric_rows[-1]['mae']:.4f}", flush=True, ) return pd.DataFrame(metric_rows), pd.DataFrame(prediction_rows) def _append_results( metric_rows: list[dict[str, object]], prediction_rows: list[dict[str, object]], dataset_name: str, seed: int, fold: int, model_id: str, train_idx: np.ndarray, test_idx: np.ndarray, groups: np.ndarray, index: pd.DataFrame, y: np.ndarray, train_pred: np.ndarray, test_pred: np.ndarray, epochs: int, ) -> None: metrics = regression_metrics(y[test_idx], test_pred) train_metrics = regression_metrics(y[train_idx], train_pred) metric_rows.append({ "dataset": dataset_name, "seed": seed, "fold": fold, "model": model_id, **metrics, "train_mae": train_metrics["mae"], "generalization_gap_mae": metrics["mae"] - train_metrics["mae"], "epochs": epochs, }) for local, row_idx in enumerate(test_idx): prediction_rows.append({ "dataset": dataset_name, "seed": seed, "fold": fold, "model": model_id, "row_index": int(row_idx), "battery_id": groups[row_idx], "cycle_number": int(index.iloc[row_idx]["cycle_number"]), "y_true": y[row_idx], "y_pred": float(test_pred[local]), "residual": float(y[row_idx] - test_pred[local]), }) def run_zero_shot_sequence( source_X: np.ndarray, source_index: pd.DataFrame, targets: dict[str, tuple[np.ndarray, pd.DataFrame]], *, source_name: str = "NASA", seeds: tuple[int, ...] = (17, 42, 2026), max_epochs: int = 200, patience: int = 20, batch_size: int = 64, model_ids: tuple[str, ...] | None = None, ) -> tuple[pd.DataFrame, pd.DataFrame]: """Train each sequence model on NASA train/validation batteries once per seed. All target-domain labels remain untouched until scoring, and the same frozen source model is used for every supplied external target. """ y_source = source_index["SoH"].to_numpy(dtype=float) source_groups = source_index["battery_id"].astype(str).to_numpy() lengths = {name: len(values[0]) for name, values in targets.items()} target_concat = np.concatenate([values[0] for values in targets.values()], axis=0) metric_rows: list[dict[str, object]] = [] prediction_rows: list[dict[str, object]] = [] requested = set(model_ids or (TORCH_MODEL_IDS + TF_MODEL_IDS)) unknown = requested.difference(TORCH_MODEL_IDS + TF_MODEL_IDS) if unknown: raise KeyError(f"Unknown sequence model IDs: {sorted(unknown)}") for seed in seeds: splitter = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=seed) train_idx, val_idx = next(splitter.split(source_X, y_source, source_groups)) scaler = SequenceStandardizer().fit(source_X[train_idx]) X_train = scaler.transform(source_X[train_idx]) X_val = scaler.transform(source_X[val_idx]) X_target = scaler.transform(target_concat) device = "cuda" if __import__("torch").cuda.is_available() else "cpu" factories = _torch_factories(source_X.shape[-1]) for model_id in TORCH_MODEL_IDS + TF_MODEL_IDS: if model_id not in requested: continue print( f"[{source_name}->external] seed={seed} model={model_id} start", flush=True, ) _set_seed(seed) if model_id in factories: _, combined_pred, epochs = _train_torch_model( model_id, factories[model_id](), X_train, y_source[train_idx], X_val, y_source[val_idx], X_target, max_epochs=max_epochs, patience=patience, batch_size=batch_size, device=device, ) else: _, combined_pred, epochs = _train_tensorflow_model( model_id, X_train, y_source[train_idx], X_val, y_source[val_idx], X_target, max_epochs=max_epochs, patience=patience, batch_size=batch_size, ) offset = 0 for target_name, (_, target_index) in targets.items(): count = lengths[target_name] pred = combined_pred[offset : offset + count] offset += count y_target = target_index["SoH"].to_numpy(dtype=float) metric_rows.append({ "source_dataset": source_name, "target_dataset": target_name, "seed": seed, "model": model_id, "epochs": epochs, **regression_metrics(y_target, pred), }) for row, value in enumerate(pred): prediction_rows.append({ "source_dataset": source_name, "target_dataset": target_name, "seed": seed, "model": model_id, "battery_id": str(target_index.iloc[row]["battery_id"]), "cycle_number": int(target_index.iloc[row]["cycle_number"]), "y_true": y_target[row], "y_pred": float(value), "residual": float(y_target[row] - value), }) print( f"[{source_name}->external] seed={seed} model={model_id} " f"done epochs={epochs}", flush=True, ) return pd.DataFrame(metric_rows), pd.DataFrame(prediction_rows)