revision/scripts/reanalysis_neural.py

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