| """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() |
|
|
