"""Benchmark warm-checkpoint inference for one frozen outer run.""" from __future__ import annotations import argparse import json import sys import time from pathlib import Path from typing import Any, Sequence import joblib import numpy as np import pandas as pd import torch from sklearn.preprocessing import LabelEncoder from torch_geometric.loader import DataLoader as PyGDataLoader PROJECT_ROOT = Path(__file__).resolve().parents[2] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from src.data import MolecularFeatureExtractor from src.models import GATModel, GraphConvModel, HybridModel from src.neural_models import FingerprintNN from revision.scripts.reanalysis_core import annotate_structures from revision.scripts.reanalysis_pipeline import _extract_features MODEL_CLASSES = {"gat": GATModel, "gcn": GraphConvModel} def summarize_durations(durations: Sequence[float], *, n_records: int) -> dict[str, float | int]: values = np.asarray(durations, dtype=float) if len(values) == 0 or n_records <= 0: raise ValueError("Durations and a positive record count are required.") mean = float(values.mean()) sd = float(values.std(ddof=1)) if len(values) > 1 else 0.0 return { "repeats": int(len(values)), "seconds_mean": mean, "seconds_sd": sd, "milliseconds_per_record_mean": mean / n_records * 1000.0, } def _load_hybrid_model( checkpoint_path: Path, model_name: str, descriptor_dim: int, device: torch.device, ) -> HybridModel: payload: dict[str, Any] = torch.load( checkpoint_path, map_location="cpu", weights_only=True ) config = payload["config"] model = HybridModel( graph_model_class=MODEL_CLASSES[model_name], descriptor_dim=descriptor_dim, graph_model_kwargs=config.get("graph_model_kwargs"), graph_feature_dim=config.get("graph_feature_dim"), descriptor_hidden_dims=config.get("descriptor_hidden_dims"), final_hidden_dims=config.get("final_hidden_dims"), dropout=config.get("dropout", 0.2), use_batch_norm=config.get("use_batch_norm", True), output_dim=config.get("output_dim", 1), ) model.load_state_dict(payload["model_state"]) model.target_mean = float(payload["target_mean"]) model.target_std = float(payload["target_std"]) model.eval() return model.to(device) def _load_fpnn_model(checkpoint_path: Path, device: torch.device) -> FingerprintNN: payload: dict[str, Any] = torch.load( checkpoint_path, map_location="cpu", weights_only=True ) if "config" in payload: model_config = payload["config"] else: state = payload["model_state"] laboratory_shape = state["lab_embedding.weight"].shape linear_keys = sorted( ( key for key, value in state.items() if key.startswith("network.") and key.endswith(".weight") and value.ndim == 2 ), key=lambda key: int(key.split(".")[1]), ) linear_shapes = [state[key].shape for key in linear_keys] model_config = { "input_dim": int(linear_shapes[0][1] - laboratory_shape[1]), "hidden_dims": [int(shape[0]) for shape in linear_shapes[:-1]], "dropout": 0.15, "use_batch_norm": any("running_mean" in key for key in state), "num_labs": int(laboratory_shape[0]), "lab_embed_dim": int(laboratory_shape[1]), "input_dropout": 0.1, } model = FingerprintNN(**model_config) model.load_state_dict(payload["model_state"]) model.target_mean = float(payload["target_mean"]) model.target_std = float(payload["target_std"]) model.eval() return model.to(device) def _predict_hybrid( model: HybridModel, 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: list[float] = [] 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), ) predictions.extend(values.detach().cpu().numpy().reshape(-1).tolist()) values = np.asarray(predictions, dtype=np.float32) return values * model.target_std + model.target_mean def _predict_fpnn( model: FingerprintNN, fingerprints: np.ndarray, labs: np.ndarray, device: torch.device, ) -> np.ndarray: features = torch.tensor(fingerprints, dtype=torch.float32) laboratory = torch.tensor(labs, dtype=torch.long) loader = torch.utils.data.DataLoader( torch.utils.data.TensorDataset(features, laboratory), batch_size=256, shuffle=False, ) predictions: list[float] = [] 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.tolist()) values = np.asarray(predictions, dtype=np.float32) return values * model.target_std + model.target_mean def _synchronize(device: torch.device) -> None: if device.type == "cuda": torch.cuda.synchronize(device) def main() -> int: parser = argparse.ArgumentParser( description="Measure loaded six-fold neural-stack inference cost." ) parser.add_argument("--data", required=True) parser.add_argument("--artifacts-root", required=True) parser.add_argument("--strategy", default="canonical_grouped") parser.add_argument("--seed", type=int, default=123456) parser.add_argument("--expected-folds", type=int, default=6) parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--device", default="auto") parser.add_argument("--output", required=True) arguments = parser.parse_args() if arguments.expected_folds <= 0 or arguments.repeats <= 0: raise ValueError("expected-folds and repeats must be positive.") device = torch.device( "cuda" if arguments.device == "auto" and torch.cuda.is_available() else "cpu" if arguments.device == "auto" else arguments.device ) artifacts_root = Path(arguments.artifacts_root).resolve() run_dir = artifacts_root / arguments.strategy / f"seed_{arguments.seed}" neural_dir = run_dir / "neural_stack" output_path = Path(arguments.output).resolve() if output_path.exists(): raise FileExistsError(f"Refusing to overwrite inference benchmark: {output_path}") output_path.parent.mkdir(parents=True, exist_ok=True) config = json.loads( (artifacts_root / "FROZEN_CONFIG.json").read_text(encoding="utf-8") ) annotated = annotate_structures(pd.read_csv(arguments.data)) indices = np.load(run_dir / "split_indices.npz", allow_pickle=True) test_indices = np.asarray(indices["test_indices"], dtype=int) test_frame = annotated.iloc[test_indices].reset_index(drop=True) feature_started = time.perf_counter() descriptors, fingerprints = _extract_features( test_frame, config["descriptor_features"], config["fingerprint"], ) extractor = MolecularFeatureExtractor() graphs = [extractor.smiles_to_graph(value) for value in test_frame["SMILES"].astype(str)] if any(graph is None for graph in graphs): raise ValueError("Graph generation failed during inference benchmark.") feature_seconds = time.perf_counter() - feature_started encoder_payload = json.loads( (neural_dir / "fold_preprocessing" / "lab_encoder.json").read_text( encoding="utf-8" ) ) lab_encoder = LabelEncoder() lab_encoder.classes_ = np.asarray(encoder_payload["classes_in_index_order"], dtype=object) laboratory_indices = lab_encoder.transform(test_frame["Lab"].astype(str)).astype(np.int64) hybrid_models: dict[str, list[tuple[HybridModel, np.ndarray]]] = { "gat": [], "gcn": [], } fpnn_models: list[FingerprintNN] = [] for fold in range(arguments.expected_folds): preprocessing = joblib.load( neural_dir / "fold_preprocessing" / f"fold_{fold}.joblib" ) scaled_descriptors = preprocessing["scaler"].transform(descriptors).astype( np.float32 ) for model_name in ("gat", "gcn"): checkpoint = ( neural_dir / "checkpoints" / model_name / f"{model_name}_fold_{fold}.pt" ) if not checkpoint.is_file(): raise FileNotFoundError(f"Incomplete checkpoint matrix: {checkpoint}") hybrid_models[model_name].append( ( _load_hybrid_model( checkpoint, model_name, descriptor_dim=descriptors.shape[1], device=device, ), scaled_descriptors, ) ) fpnn_checkpoint = ( neural_dir / "checkpoints" / "fpnn" / f"fp_nn_fold_{fold}.pt" ) if not fpnn_checkpoint.is_file(): raise FileNotFoundError(f"Incomplete checkpoint matrix: {fpnn_checkpoint}") fpnn_models.append(_load_fpnn_model(fpnn_checkpoint, device)) stack_models = joblib.load(neural_dir / "stack_models.joblib") primary_stack = stack_models["stack_all_plus_descriptors"] component_durations = {"gat": [], "gcn": [], "fpnn": [], "stack": []} total_durations: list[float] = [] for repeat in range(arguments.repeats + 1): _synchronize(device) total_started = time.perf_counter() averaged: dict[str, np.ndarray] = {} for model_name in ("gat", "gcn"): _synchronize(device) component_started = time.perf_counter() fold_predictions = [ _predict_hybrid( model, graphs, laboratory_indices, scaled_descriptors, device, ) for model, scaled_descriptors in hybrid_models[model_name] ] _synchronize(device) elapsed = time.perf_counter() - component_started averaged[model_name] = np.mean(np.vstack(fold_predictions), axis=0) if repeat > 0: component_durations[model_name].append(elapsed) _synchronize(device) fpnn_started = time.perf_counter() fpnn_predictions = [ _predict_fpnn(model, fingerprints, laboratory_indices, device) for model in fpnn_models ] _synchronize(device) fpnn_elapsed = time.perf_counter() - fpnn_started averaged["fpnn"] = np.mean(np.vstack(fpnn_predictions), axis=0) stack_features = np.column_stack( [averaged["gat"], averaged["gcn"], averaged["fpnn"], descriptors] ) stack_started = time.perf_counter() _ = primary_stack.predict(stack_features) stack_elapsed = time.perf_counter() - stack_started total_elapsed = time.perf_counter() - total_started if repeat > 0: component_durations["fpnn"].append(fpnn_elapsed) component_durations["stack"].append(stack_elapsed) total_durations.append(total_elapsed) payload = { "strategy": arguments.strategy, "seed": arguments.seed, "device": str(device), "gpu_name": torch.cuda.get_device_name(device) if device.type == "cuda" else None, "n_records": int(len(test_frame)), "n_fold_models_per_base": arguments.expected_folds, "scope": "warm loaded checkpoints; includes batching and device transfer; excludes checkpoint loading and RDKit feature generation", "feature_generation": { "seconds": float(feature_seconds), "milliseconds_per_record": float(feature_seconds / len(test_frame) * 1000.0), }, "full_bundle": summarize_durations(total_durations, n_records=len(test_frame)), "components": { name: summarize_durations(values, n_records=len(test_frame)) for name, values in component_durations.items() }, } output_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") print(f"Wrote inference benchmark to: {output_path}") return 0 if __name__ == "__main__": raise SystemExit(main())