"""Matched FPNN laboratory-encoding ablation on the frozen outer splits. The submitted/revised FPNN concatenates a learned 64-dimensional laboratory embedding with the 2,048-bit Morgan fingerprint. This script retains the same fingerprint, hidden layers, optimizer, target scaling, early stopping, inner folds, and outer test partitions, replacing only the learned embedding by a fixed 23-dimensional one-hot laboratory vector. """ from __future__ import annotations import argparse import json import random import sys import time from pathlib import Path import numpy as np import pandas as pd import torch from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score from sklearn.preprocessing import OneHotEncoder 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.trainers import NeuralNetworkTrainer from src.neural_models import FingerprintNN TASKS = { "canonical_grouped": ("Identity-grouped", "structure_group"), "scaffold_aware": ("Scaffold-aware", "scaffold_component_group"), } SEEDS = (123456, 123457, 123458) def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def scores(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]: return { "r2": float(r2_score(y_true, y_pred)), "mae": float(mean_absolute_error(y_true, y_pred)), "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))), } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--artifacts", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--device", default="auto") args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) first = pd.read_csv( args.artifacts / "canonical_grouped" / "seed_123456" / "split_assignments.csv" ).sort_values("record_index") extractor = MolecularFeatureExtractor() fingerprints = np.stack( [extractor.get_morgan_fingerprint(smiles) for smiles in first["SMILES"].astype(str)] ).astype(np.float32) device = torch.device( "cuda" if args.device == "auto" and torch.cuda.is_available() else "cpu" if args.device == "auto" else args.device ) model_base = { "hidden_dims": [768, 384, 192, 96], "dropout": 0.15, "use_batch_norm": True, "num_labs": None, "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, } rows: list[dict[str, object]] = [] for task_dir, (task_label, group_column) in TASKS.items(): for repeat, seed in enumerate(SEEDS, start=1): run_dir = args.output_dir / task_dir / f"seed_{seed}" run_dir.mkdir(parents=True, exist_ok=True) split = pd.read_csv( args.artifacts / task_dir / f"seed_{seed}" / "split_assignments.csv" ).sort_values("record_index") if not np.array_equal(split["record_index"].to_numpy(), first["record_index"].to_numpy()): raise RuntimeError("Record order changed across frozen split files.") development = np.flatnonzero(split["outer_split"].eq("development").to_numpy()) test = np.flatnonzero(split["outer_split"].eq("test").to_numpy()) targets = split["RT"].to_numpy(dtype=np.float32) encoder = OneHotEncoder(sparse_output=False, handle_unknown="error", dtype=np.float32) encoder.fit(split.loc[development, ["Lab"]]) one_hot = encoder.transform(split[["Lab"]]).astype(np.float32) fixed_features = np.column_stack([fingerprints, one_hot]).astype(np.float32) config = dict(model_base, input_dim=int(fixed_features.shape[1])) fold_test_predictions: list[np.ndarray] = [] fold_metrics: list[dict[str, object]] = [] started = time.perf_counter() for fold in range(6): fold_validation = np.flatnonzero( split["inner_validation_fold"].eq(fold).to_numpy() ) fold_train = np.setdiff1d(development, fold_validation, assume_unique=False) if len(fold_validation) == 0 or not np.isin(fold_validation, development).all(): raise RuntimeError(f"Invalid inner fold {fold} for {task_dir}, seed {seed}.") set_seed(seed + fold) model = FingerprintNN(**config) trainer = NeuralNetworkTrainer( model=model, device=device, lr=training["lr"], weight_decay=training["weight_decay"], ) metrics, _ = trainer.train_fold( train_features=fixed_features[fold_train], train_targets=targets[fold_train], val_features=fixed_features[fold_validation], val_targets=targets[fold_validation], epochs=training["epochs"], batch_size=training["batch_size"], patience=training["patience"], gradient_clip=training["gradient_clip"], warmup_epochs=training["warmup_epochs"], verbose=False, ) trainer.save(str(run_dir / f"fpnn_one_hot_fold_{fold}.pt")) fold_test_predictions.append(trainer.predict(fixed_features[test])) fold_metrics.append({"fold": fold, **metrics}) prediction = np.mean(np.stack(fold_test_predictions), axis=0) elapsed = time.perf_counter() - started learned = pd.read_csv( args.artifacts / task_dir / f"seed_{seed}" / "neural_stack" / "test_predictions.csv" ).sort_values("record_index") if not np.array_equal(learned["record_index"].to_numpy(), split.iloc[test]["record_index"].to_numpy()): raise RuntimeError("Test-prediction order does not match frozen split assignment.") learned_scores = scores(targets[test], learned["prediction_fpnn"].to_numpy(dtype=float)) one_hot_scores = scores(targets[test], prediction) prediction_table = split.iloc[test][ ["record_index", "SMILES", "Lab", "RT", "structure_group", "scaffold_component_group"] ].copy() prediction_table["prediction_fpnn_learned_embedding"] = learned[ "prediction_fpnn" ].to_numpy(dtype=float) prediction_table["prediction_fpnn_one_hot"] = prediction prediction_table.to_csv(run_dir / "test_predictions.csv", index=False) metadata = { "task": task_dir, "task_label": task_label, "repeat": repeat, "outer_seed": seed, "device": str(device), "comparison": "Same FPNN trunk; learned 64-dimensional lab embedding versus fixed 23-dimensional one-hot lab vector.", "fingerprint": {"radius": 2, "n_bits": 2048, "use_chirality": True}, "group_column": group_column, "n_test_rows": int(len(test)), "n_test_groups": int(split.iloc[test][group_column].nunique()), "model_config": config, "training_config": training, "fold_metrics": fold_metrics, "training_seconds": elapsed, "learned_embedding": learned_scores, "one_hot": one_hot_scores, } (run_dir / "metrics.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") rows.append( { "task": task_label, "repeat": repeat, "seed": seed, "encoding": "Learned embedding", **learned_scores, } ) rows.append( { "task": task_label, "repeat": repeat, "seed": seed, "encoding": "One-hot vector", **one_hot_scores, } ) per_run = pd.DataFrame(rows) per_run.to_csv(args.output_dir / "matched_lab_encoding_by_run.csv", index=False) aggregate = ( per_run.groupby(["task", "encoding"], sort=False)[["r2", "mae", "rmse"]] .agg(["mean", "std"]) .reset_index() ) aggregate.columns = ["_".join(str(value) for value in col if value) for col in aggregate.columns] aggregate.to_csv(args.output_dir / "matched_lab_encoding_aggregate.csv", index=False) if __name__ == "__main__": main()