"""Leakage-controlled neural reanalysis using the existing training functions.""" from __future__ import annotations import copy import json import time from pathlib import Path from typing import Any, Dict, Mapping, Sequence import joblib import numpy as np import pandas as pd import torch from sklearn.ensemble import ExtraTreesRegressor from sklearn.preprocessing import LabelEncoder, StandardScaler from torch_geometric.loader import DataLoader as PyGDataLoader from src.data import MolecularFeatureExtractor from src.models import GATModel, GraphConvModel from src.trainers import train_fingerprint_nn_fold, train_hybrid_model_fold from revision.scripts.reanalysis_core import ( compute_regression_metrics, ensure_new_output_dir, paired_group_bootstrap, per_lab_metrics, ) PRIMARY_STACK_MODEL = "stack_all_plus_descriptors" BASE_MODEL_ORDER = ("gat", "gcn", "fpnn") def estimate_neural_workload(config: Mapping[str, Any]) -> Dict[str, int]: """Count the predeclared model fits before launching long training.""" outer_runs = len(config.get("outer_seeds", [])) * len(config.get("split_strategies", {})) inner_folds = int(config.get("inner_folds", 6)) base_models = config.get("neural", {}).get("base_models", list(BASE_MODEL_ORDER)) per_outer = inner_folds * len(base_models) return { "outer_runs": int(outer_runs), "base_model_fits_per_outer_run": int(per_outer), "total_base_model_fits": int(outer_runs * per_outer), } def _tree_parameters(estimator_config: Mapping[str, Any], seed: int) -> Dict[str, Any]: parameters = { "n_estimators": int(estimator_config.get("n_estimators", 800)), "max_depth": estimator_config.get("max_depth", 12), "min_samples_split": int(estimator_config.get("min_samples_split", 5)), "min_samples_leaf": int(estimator_config.get("min_samples_leaf", 3)), "max_features": estimator_config.get("max_features", 0.8), "bootstrap": bool(estimator_config.get("bootstrap", True)), "random_state": int(seed), "n_jobs": int(estimator_config.get("n_jobs", -1)), } if parameters["bootstrap"]: parameters["max_samples"] = estimator_config.get("max_samples", 0.85) return parameters def build_stack_ablations( *, oof_predictions: Mapping[str, np.ndarray], test_predictions: Mapping[str, np.ndarray], development_descriptors: np.ndarray, test_descriptors: np.ndarray, development_targets: Sequence[float], seed: int, estimator_config: Mapping[str, Any], ) -> Dict[str, Any]: """Fit predeclared full, descriptor-free, and pairwise stacks.""" missing = [name for name in BASE_MODEL_ORDER if name not in oof_predictions or name not in test_predictions] if missing: raise ValueError(f"Missing base predictions for: {missing}") y = np.asarray(development_targets, dtype=float) oof_columns = {name: np.asarray(oof_predictions[name], dtype=float) for name in BASE_MODEL_ORDER} test_columns = {name: np.asarray(test_predictions[name], dtype=float) for name in BASE_MODEL_ORDER} base_oof = np.column_stack([oof_columns[name] for name in BASE_MODEL_ORDER]) base_test = np.column_stack([test_columns[name] for name in BASE_MODEL_ORDER]) development_descriptors = np.asarray(development_descriptors, dtype=np.float32) test_descriptors = np.asarray(test_descriptors, dtype=np.float32) feature_sets = { "stack_all_base_only": (base_oof, base_test), "stack_all_plus_descriptors": ( np.column_stack([base_oof, development_descriptors]), np.column_stack([base_test, test_descriptors]), ), "stack_gat_gcn": (base_oof[:, [0, 1]], base_test[:, [0, 1]]), "stack_gat_fpnn": (base_oof[:, [0, 2]], base_test[:, [0, 2]]), "stack_gcn_fpnn": (base_oof[:, [1, 2]], base_test[:, [1, 2]]), } predictions: Dict[str, np.ndarray] = { "arithmetic_mean": np.mean(base_test, axis=1), } models: Dict[str, ExtraTreesRegressor] = {} for name, (train_features, heldout_features) in feature_sets.items(): model = ExtraTreesRegressor(**_tree_parameters(estimator_config, seed)) model.fit(train_features, y) predictions[name] = np.asarray(model.predict(heldout_features), dtype=float) models[name] = model return { "primary_model": PRIMARY_STACK_MODEL, "predictions": predictions, "models": models, "feature_dimensions": {name: int(values[0].shape[1]) for name, values in feature_sets.items()}, } def _default_neural_config(num_labs: int, input_dim: int, fingerprint_dim: int) -> Dict[str, Any]: return { "base_models": ["gat", "gcn", "fpnn"], "gat": { "model": { "graph_model_kwargs": { "input_dim": input_dim, "hidden_dim": 320, "num_layers": 4, "dropout": 0.18, "num_labs": num_labs, "lab_embed_dim": 32, }, "graph_feature_dim": 320, "descriptor_hidden_dims": [160, 80], "final_hidden_dims": [160, 80], "dropout": 0.2, "use_batch_norm": True, }, "training": { "epochs": 600, "lr": 0.0003, "weight_decay": 0.000005, "patience": 60, "plateau_patience": 40, "factor": 0.6, "min_lr": 0.000001, "batch_size": 32, "gradient_clip": 1.0, }, }, "gcn": { "model": { "graph_model_kwargs": { "input_dim": input_dim, "hidden_dim": 256, "num_layers": 5, "dropout": 0.2, "num_labs": num_labs, "lab_embed_dim": 32, }, "graph_feature_dim": 256, "descriptor_hidden_dims": [128, 64], "final_hidden_dims": [128, 64], "dropout": 0.2, "use_batch_norm": True, }, "training": { "epochs": 600, "lr": 0.0004, "weight_decay": 0.00001, "patience": 60, "plateau_patience": 35, "factor": 0.6, "min_lr": 0.000001, "batch_size": 32, "gradient_clip": 1.0, }, }, "fpnn": { "model": { "input_dim": fingerprint_dim, "hidden_dims": [768, 384, 192, 96], "dropout": 0.15, "use_batch_norm": True, "num_labs": num_labs, "lab_embed_dim": 64, "input_dropout": 0.1, }, "training": { "epochs": 600, "lr": 0.0015, "weight_decay": 0.00001, "patience": 60, "warmup_epochs": 15, "batch_size": 128, "gradient_clip": 1.0, }, }, "meta_learner": { "n_estimators": 800, "max_depth": 12, "min_samples_split": 5, "min_samples_leaf": 3, "max_features": 0.8, "bootstrap": True, "max_samples": 0.85, }, } def _merge_neural_config(defaults: Mapping[str, Any], overrides: Mapping[str, Any]) -> Dict[str, Any]: merged = copy.deepcopy(dict(defaults)) for key, value in overrides.items(): if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping): merged[key] = _merge_neural_config(merged[key], value) else: merged[key] = copy.deepcopy(value) return merged def _set_seed(seed: int) -> None: np.random.seed(int(seed)) torch.manual_seed(int(seed)) if torch.cuda.is_available(): torch.cuda.manual_seed_all(int(seed)) def _predict_hybrid_model( model: torch.nn.Module, graphs: Sequence[Any], labs: np.ndarray, descriptors: np.ndarray, device: torch.device, ) -> np.ndarray: prepared = [] for index, graph in enumerate(graphs): graph_copy = graph.clone() graph_copy.lab_feature = torch.tensor([int(labs[index])], dtype=torch.long) graph_copy.descriptors = torch.tensor(descriptors[index], dtype=torch.float32).reshape(1, -1) prepared.append(graph_copy) loader = PyGDataLoader(prepared, batch_size=64, shuffle=False) predictions = [] model = model.to(device) model.eval() with torch.no_grad(): for batch in loader: batch = batch.to(device) lab_tensor = batch.lab_feature.squeeze(-1) if batch.lab_feature.dim() > 1 else batch.lab_feature descriptor_tensor = batch.descriptors.reshape(batch.num_graphs, -1) values = model( batch.x, batch.edge_index, batch.batch, lab_tensor, descriptor_tensor, getattr(batch, "edge_attr", None), ) values = values.detach().cpu().numpy().reshape(-1) predictions.extend( (values * float(getattr(model, "target_std", 1.0)) + float(getattr(model, "target_mean", 0.0))).tolist() ) model.to("cpu") return np.asarray(predictions, dtype=np.float32) def _predict_fpnn_model( model: torch.nn.Module, fingerprints: np.ndarray, labs: np.ndarray, device: torch.device, ) -> np.ndarray: feature_tensor = torch.tensor(fingerprints, dtype=torch.float32) lab_tensor = torch.tensor(labs, dtype=torch.long) loader = torch.utils.data.DataLoader( torch.utils.data.TensorDataset(feature_tensor, lab_tensor), batch_size=256, shuffle=False, ) predictions = [] model = model.to(device) model.eval() with torch.no_grad(): for batch_features, batch_labs in loader: values = model(batch_features.to(device), batch_labs.to(device)).detach().cpu().numpy().reshape(-1) predictions.extend( (values * float(getattr(model, "target_std", 1.0)) + float(getattr(model, "target_mean", 0.0))).tolist() ) model.to("cpu") return np.asarray(predictions, dtype=np.float32) def _write_json(path: Path, payload: Mapping[str, Any]) -> None: def default(value: Any) -> Any: if isinstance(value, np.integer): return int(value) if isinstance(value, np.floating): return float(value) if isinstance(value, np.ndarray): return value.tolist() raise TypeError(type(value).__name__) path.write_text(json.dumps(payload, indent=2, ensure_ascii=False, default=default) + "\n", encoding="utf-8") def _maximum_train_tanimoto( train_fingerprints: np.ndarray, test_fingerprints: np.ndarray, *, chunk_size: int = 64, ) -> np.ndarray: """Compute each test row's maximum binary Tanimoto to development rows.""" train = (np.asarray(train_fingerprints) > 0).astype(np.float32) test = (np.asarray(test_fingerprints) > 0).astype(np.float32) train_counts = train.sum(axis=1, dtype=np.float64) maxima = np.zeros(len(test), dtype=np.float32) for start in range(0, len(test), chunk_size): stop = min(start + chunk_size, len(test)) chunk = test[start:stop] intersections = chunk @ train.T unions = ( chunk.sum(axis=1, dtype=np.float64)[:, None] + train_counts[None, :] - intersections ) similarities = np.divide( intersections, unions, out=np.zeros_like(intersections, dtype=np.float32), where=unions > 0, ) maxima[start:stop] = similarities.max(axis=1) return maxima def _prospective_domain_summary( *, y_true: np.ndarray, y_pred: np.ndarray, maximum_similarity: np.ndarray, base_model_spread: np.ndarray, thresholds: Sequence[float], ) -> Dict[str, Any]: absolute_error = np.abs(np.asarray(y_pred) - np.asarray(y_true)) similarity = np.asarray(maximum_similarity, dtype=float) spread = np.asarray(base_model_spread, dtype=float) bins = np.asarray([0.0, 0.2, 0.4, 0.6, 0.8, 1.000001], dtype=float) bin_rows = [] for lower, upper in zip(bins[:-1], bins[1:]): mask = (similarity >= lower) & (similarity < upper) bin_rows.append( { "similarity_lower": float(lower), "similarity_upper": float(min(upper, 1.0)), "n": int(mask.sum()), "mae": float(absolute_error[mask].mean()) if mask.any() else float("nan"), } ) threshold_rows = [] for threshold in thresholds: accepted = similarity >= float(threshold) rejected = ~accepted threshold_rows.append( { "threshold": float(threshold), "accepted_n": int(accepted.sum()), "rejected_n": int(rejected.sum()), "accepted_coverage": float(accepted.mean()), "accepted_mae": float(absolute_error[accepted].mean()) if accepted.any() else float("nan"), "rejected_mae": float(absolute_error[rejected].mean()) if rejected.any() else float("nan"), } ) return { "interpretation": ( "Maximum training-set Tanimoto is a prospective structural-domain diagnostic. " "Base-model standard deviation is uncalibrated model disagreement, not a confidence interval." ), "similarity_bins": bin_rows, "threshold_sensitivity": threshold_rows, "spearman_similarity_vs_absolute_error": float( pd.Series(similarity).corr(pd.Series(absolute_error), method="spearman") ), "spearman_model_spread_vs_absolute_error": float( pd.Series(spread).corr(pd.Series(absolute_error), method="spearman") ), } def run_neural_stack( config: Mapping[str, Any], *, annotated: pd.DataFrame, descriptor_matrix: np.ndarray, fingerprint_matrix: np.ndarray, train_indices: Sequence[int], test_indices: Sequence[int], inner_folds: Sequence[tuple[np.ndarray, np.ndarray]], output_dir: Path, seed: int, smoke_only: bool = False, ) -> Path: """Retrain the fixed neural stack for one predeclared outer split. All checkpoints, scalers, OOF predictions, and metrics are written below ``output_dir``. The function never references ``hybrid_oof_models``. """ run_started = time.perf_counter() output_dir = ensure_new_output_dir(output_dir) extractor = MolecularFeatureExtractor() graphs = [] for row_index, smiles in enumerate(annotated["SMILES"].astype(str)): graph = extractor.smiles_to_graph(smiles) if graph is None: raise ValueError(f"Graph conversion failed at annotated row {row_index}.") graphs.append(graph) lab_encoder = LabelEncoder() lab_indices = lab_encoder.fit_transform(annotated["Lab"].astype(str)).astype(np.int64) development = np.asarray(train_indices, dtype=int) test = np.asarray(test_indices, dtype=int) untrained_test_labs = sorted( set(lab_indices[test]) - set(lab_indices[development]) ) if untrained_test_labs and not smoke_only: raise ValueError("A test laboratory is absent from development; categorical embeddings cannot extrapolate it.") targets = annotated["RT"].to_numpy(dtype=np.float32) default_config = _default_neural_config( num_labs=len(lab_encoder.classes_), input_dim=int(graphs[0].x.shape[1]), fingerprint_dim=int(fingerprint_matrix.shape[1]), ) neural_config = _merge_neural_config(default_config, config.get("neural", {})) if tuple(neural_config.get("base_models", BASE_MODEL_ORDER)) != BASE_MODEL_ORDER: raise ValueError(f"base_models must remain predeclared as {BASE_MODEL_ORDER}.") if smoke_only: for name in BASE_MODEL_ORDER: neural_config[name]["training"]["epochs"] = 1 neural_config[name]["training"]["patience"] = 1 neural_config["meta_learner"]["n_estimators"] = min( 8, int(neural_config["meta_learner"].get("n_estimators", 8)) ) requested_device = config.get("neural", {}).get("device", "auto") device = torch.device( "cuda" if requested_device == "auto" and torch.cuda.is_available() else "cpu" if requested_device == "auto" else requested_device ) oof_predictions = {name: np.full(len(annotated), np.nan, dtype=np.float32) for name in BASE_MODEL_ORDER} fold_test_predictions = {name: [] for name in BASE_MODEL_ORDER} cv_scores = {name: [] for name in BASE_MODEL_ORDER} parameter_counts: Dict[str, int] = {} training_seconds: Dict[str, float] = {name: 0.0 for name in BASE_MODEL_ORDER} checkpoint_root = output_dir / "checkpoints" preprocessing_root = output_dir / "fold_preprocessing" preprocessing_root.mkdir(parents=True, exist_ok=False) _write_json( preprocessing_root / "lab_encoder.json", {"classes_in_index_order": [str(value) for value in lab_encoder.classes_]}, ) for fold_index, (fold_train, fold_validation) in enumerate(inner_folds): fold_seed = int(seed) + int(fold_index) _set_seed(fold_seed) scaler = StandardScaler() descriptors_train = scaler.fit_transform(descriptor_matrix[fold_train]).astype(np.float32) descriptors_validation = scaler.transform(descriptor_matrix[fold_validation]).astype(np.float32) descriptors_test = scaler.transform(descriptor_matrix[test]).astype(np.float32) joblib.dump( { "scaler": scaler, "descriptor_features": list(config["descriptor_features"]), "fold_train_indices": np.asarray(fold_train, dtype=np.int32), "fold_validation_indices": np.asarray(fold_validation, dtype=np.int32), }, preprocessing_root / f"fold_{fold_index}.joblib", ) for model_name, model_class in (("gat", GATModel), ("gcn", GraphConvModel)): fit_started = time.perf_counter() model, validation_predictions, metrics = train_hybrid_model_fold( model_name=model_name, graph_model_class=model_class, config=neural_config[model_name]["model"], training_config=neural_config[model_name]["training"], fold_train_graphs=[graphs[index].clone() for index in fold_train], fold_val_graphs=[graphs[index].clone() for index in fold_validation], fold_train_lab=lab_indices[fold_train], fold_val_lab=lab_indices[fold_validation], fold_train_targets=targets[fold_train], fold_val_targets=targets[fold_validation], fold_train_descriptors=descriptors_train, fold_val_descriptors=descriptors_validation, fold_idx=fold_index, device=device, save_dir=str(checkpoint_root / model_name), verbose_interval=50, ) elapsed = time.perf_counter() - fit_started training_seconds[model_name] += elapsed metrics = dict(metrics) metrics["training_seconds"] = float(elapsed) parameter_counts.setdefault( model_name, int(sum(parameter.numel() for parameter in model.parameters())), ) oof_predictions[model_name][fold_validation] = validation_predictions cv_scores[model_name].append(metrics) fold_test_predictions[model_name].append( _predict_hybrid_model( model, [graphs[index].clone() for index in test], lab_indices[test], descriptors_test, device, ) ) fit_started = time.perf_counter() fpnn_model, validation_predictions, metrics = train_fingerprint_nn_fold( train_fingerprints=fingerprint_matrix[fold_train], val_fingerprints=fingerprint_matrix[fold_validation], train_targets=targets[fold_train], val_targets=targets[fold_validation], fold_idx=fold_index, device=device, config=neural_config["fpnn"]["model"], training_config=neural_config["fpnn"]["training"], save_dir=str(checkpoint_root / "fpnn"), train_lab_indices=lab_indices[fold_train], val_lab_indices=lab_indices[fold_validation], ) elapsed = time.perf_counter() - fit_started training_seconds["fpnn"] += elapsed metrics = dict(metrics) metrics["training_seconds"] = float(elapsed) parameter_counts.setdefault( "fpnn", int(sum(parameter.numel() for parameter in fpnn_model.parameters())), ) oof_predictions["fpnn"][fold_validation] = validation_predictions cv_scores["fpnn"].append(metrics) fold_test_predictions["fpnn"].append( _predict_fpnn_model(fpnn_model, fingerprint_matrix[test], lab_indices[test], device) ) for model_name in BASE_MODEL_ORDER: if np.isnan(oof_predictions[model_name][development]).any(): raise RuntimeError(f"Incomplete OOF predictions for {model_name}.") test_predictions = { name: np.mean(np.stack(values, axis=0), axis=0) for name, values in fold_test_predictions.items() } stack = build_stack_ablations( oof_predictions={name: oof_predictions[name][development] for name in BASE_MODEL_ORDER}, test_predictions=test_predictions, development_descriptors=descriptor_matrix[development], test_descriptors=descriptor_matrix[test], development_targets=targets[development], seed=seed, estimator_config=neural_config["meta_learner"], ) all_test_predictions = {**test_predictions, **stack["predictions"]} metrics = { name: compute_regression_metrics(targets[test], predictions) for name, predictions in all_test_predictions.items() } maximum_similarity = _maximum_train_tanimoto( fingerprint_matrix[development], fingerprint_matrix[test], ) base_model_spread = np.std( np.column_stack([test_predictions[name] for name in BASE_MODEL_ORDER]), axis=1, ddof=0, ) primary_predictions = stack["predictions"][stack["primary_model"]] domain_summary = _prospective_domain_summary( y_true=targets[test], y_pred=primary_predictions, maximum_similarity=maximum_similarity, base_model_spread=base_model_spread, thresholds=config.get("prospective_domain", {}).get( "similarity_thresholds", [0.3, 0.4, 0.5], ), ) prediction_frame = annotated.loc[ test, [ "record_index", "SMILES", "Lab", "RT", "structure_group", "scaffold_group", "scaffold_component_group", ], ].copy() for name, values in all_test_predictions.items(): prediction_frame[f"prediction_{name}"] = values prediction_frame["maximum_development_tanimoto"] = maximum_similarity prediction_frame["base_model_spread"] = base_model_spread prediction_frame["primary_absolute_error"] = np.abs( primary_predictions - targets[test] ) prediction_frame.to_csv(output_dir / "test_predictions.csv", index=False) development_ranges = ( annotated.loc[development] .groupby("Lab")["RT"] .agg(lambda values: float(values.max() - values.min())) .to_dict() ) per_lab_tables = [] for name, values in all_test_predictions.items(): table = per_lab_metrics( targets[test], values, annotated.loc[test, "Lab"].to_numpy(), normalization_ranges=development_ranges, ) table.insert(0, "model", name) per_lab_tables.append(table) pd.concat(per_lab_tables, ignore_index=True).to_csv(output_dir / "per_lab_metrics.csv", index=False) primary = primary_predictions paired = {} for name, reference in all_test_predictions.items(): if name == stack["primary_model"]: continue paired[name] = paired_group_bootstrap( y_true=targets[test], candidate=primary, reference=reference, groups=annotated.loc[test, "structure_group"].to_numpy(), n_resamples=int(config["bootstrap"]["n_resamples"]), confidence=float(config["bootstrap"].get("confidence", 0.95)), seed=seed, ) np.savez_compressed( output_dir / "oof_predictions.npz", development_indices=development.astype(np.int32), y_true=targets[development], **{name: oof_predictions[name][development] for name in BASE_MODEL_ORDER}, ) np.savez_compressed( output_dir / "test_predictions.npz", test_indices=test.astype(np.int32), y_true=targets[test], maximum_development_tanimoto=maximum_similarity, base_model_spread=base_model_spread, **all_test_predictions, ) joblib.dump(stack["models"], output_dir / "stack_models.joblib") _write_json( output_dir / "metrics.json", { "primary_model_predeclared": stack["primary_model"], "metrics": metrics, "feature_dimensions": stack["feature_dimensions"], "selection_rule": "No outer-test model selection; all fixed ablations are reported.", }, ) _write_json(output_dir / "paired_group_bootstrap.json", paired) _write_json(output_dir / "prospective_domain_diagnostics.json", domain_summary) _write_json(output_dir / "cv_scores.json", cv_scores) _write_json( output_dir / "RUN_METADATA.json", { "seed": int(seed), "device": str(device), "smoke_only": bool(smoke_only), "smoke_only_untrained_test_lab_indices": untrained_test_labs, "descriptor_features_predeclared": list(config["descriptor_features"]), "neural_config": neural_config, "parameter_counts": parameter_counts, "training_seconds_by_model": training_seconds, "total_run_seconds": float(time.perf_counter() - run_started), "gpu_name": torch.cuda.get_device_name(device) if device.type == "cuda" else None, "test_set_used_for_selection": False, "checkpoint_root": "checkpoints", }, ) return output_dir