revision/scripts/run_matched_lab_encoding_ablation.py

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