revision/scripts/reanalysis_core.py

"""Core utilities for isolated, reviewer-requested reanalysis.

The functions in this module are deliberately independent of the original
``hybrid_oof_models`` directory.  They define molecular identity, group-aware
splits, cluster-aware uncertainty summaries, and classical feature matrices.
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, List, Mapping, Sequence, Tuple

import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import GroupKFold, StratifiedGroupKFold
from sklearn.preprocessing import OneHotEncoder


REQUIRED_DATA_COLUMNS = ("SMILES", "Lab", "RT")


def _validate_equal_lengths(**arrays: Sequence[object]) -> int:
    lengths = {name: len(values) for name, values in arrays.items()}
    if len(set(lengths.values())) != 1:
        raise ValueError(f"Arrays must have equal lengths; received {lengths}.")
    return next(iter(lengths.values()), 0)


def annotate_structures(frame: pd.DataFrame) -> pd.DataFrame:
    """Add canonical identity and Bemis-Murcko scaffold columns.

    The InChIKey connectivity block is the primary grouping key.  This is a
    conservative identity boundary that keeps stereoisomers and protonation
    variants with the same heavy-atom connectivity together.  Full InChIKey and
    both isomeric/non-isomeric canonical SMILES remain in the manifest for a
    sensitivity audit.  No parent-fragment selection, neutralization, or other
    salt/tautomer/protonation standardization is performed.  If InChI generation
    is unavailable for an otherwise valid molecule, canonical non-isomeric
    SMILES is used as a documented fallback.  Acyclic molecules receive their
    structure identity as the scaffold group so unrelated acyclic compounds are
    not collapsed into one giant group.
    """

    missing = [column for column in REQUIRED_DATA_COLUMNS if column not in frame.columns]
    if missing:
        raise ValueError(f"Dataset is missing required columns: {missing}.")

    output = frame.copy().reset_index(drop=True)
    canonical_isomeric: List[str] = []
    canonical_nonisomeric: List[str] = []
    inchi_keys_full: List[str] = []
    inchi_keys_connectivity: List[str] = []
    structure_groups: List[str] = []
    scaffold_groups: List[str] = []
    fragment_counts: List[int] = []
    formal_charges: List[int] = []
    stereo_flags: List[bool] = []

    for row_index, value in enumerate(output["SMILES"]):
        smiles = str(value)
        molecule = Chem.MolFromSmiles(smiles)
        if molecule is None:
            raise ValueError(f"Invalid SMILES at row {row_index}: {smiles}")

        canonical_iso = Chem.MolToSmiles(molecule, canonical=True, isomericSmiles=True)
        canonical_noniso = Chem.MolToSmiles(molecule, canonical=True, isomericSmiles=False)
        inchi_key_full = Chem.MolToInchiKey(molecule) or ""
        inchi_key_connectivity = inchi_key_full.split("-", 1)[0] if inchi_key_full else ""
        structure_group = (
            f"connectivity:{inchi_key_connectivity}"
            if inchi_key_connectivity
            else f"smiles_nonisomeric:{canonical_noniso}"
        )

        scaffold_molecule = MurckoScaffold.GetScaffoldForMol(molecule)
        if scaffold_molecule.GetNumAtoms() == 0:
            scaffold_group = f"acyclic:{structure_group}"
        else:
            scaffold = Chem.MolToSmiles(
                scaffold_molecule,
                canonical=True,
                isomericSmiles=False,
            )
            scaffold_group = f"murcko:{scaffold}"

        canonical_isomeric.append(canonical_iso)
        canonical_nonisomeric.append(canonical_noniso)
        inchi_keys_full.append(inchi_key_full)
        inchi_keys_connectivity.append(inchi_key_connectivity)
        structure_groups.append(structure_group)
        scaffold_groups.append(scaffold_group)
        fragment_counts.append(len(Chem.GetMolFrags(molecule)))
        formal_charges.append(sum(atom.GetFormalCharge() for atom in molecule.GetAtoms()))
        stereo_flags.append(canonical_iso != canonical_noniso)

    output.insert(0, "record_index", np.arange(len(output), dtype=int))
    output["canonical_smiles_isomeric"] = canonical_isomeric
    output["canonical_smiles_nonisomeric"] = canonical_nonisomeric
    output["inchi_key_full"] = inchi_keys_full
    output["inchi_key_connectivity"] = inchi_keys_connectivity
    output["structure_group"] = structure_groups
    output["scaffold_group"] = scaffold_groups

    # A normalized InChI identity can occasionally be represented by two
    # tautomeric SMILES that yield different raw Murcko scaffolds. Build
    # connected scaffold components so a scaffold holdout also preserves the
    # stricter molecular-identity boundary. Raw scaffold labels remain in the
    # manifest for auditability.
    parent = {scaffold: scaffold for scaffold in set(scaffold_groups)}

    def find(scaffold: str) -> str:
        while parent[scaffold] != scaffold:
            parent[scaffold] = parent[parent[scaffold]]
            scaffold = parent[scaffold]
        return scaffold

    def union(left: str, right: str) -> None:
        left_root = find(left)
        right_root = find(right)
        if left_root == right_root:
            return
        if left_root < right_root:
            parent[right_root] = left_root
        else:
            parent[left_root] = right_root

    scaffolds_by_structure: Dict[str, List[str]] = {}
    for structure_group, scaffold_group in zip(structure_groups, scaffold_groups):
        scaffolds_by_structure.setdefault(structure_group, []).append(scaffold_group)
    for related_scaffolds in scaffolds_by_structure.values():
        anchor = related_scaffolds[0]
        for related in related_scaffolds[1:]:
            union(anchor, related)

    component_members: Dict[str, List[str]] = {}
    for scaffold in parent:
        component_members.setdefault(find(scaffold), []).append(scaffold)
    component_label = {
        scaffold: f"scaffold_component:{min(component_members[find(scaffold)])}"
        for scaffold in parent
    }
    output["scaffold_component_group"] = [
        component_label[scaffold] for scaffold in scaffold_groups
    ]
    output["fragment_count"] = fragment_counts
    output["formal_charge"] = formal_charges
    output["has_explicit_stereo"] = stereo_flags
    return output


def _stratified_group_splitter(n_splits: int, seed: int) -> StratifiedGroupKFold:
    if n_splits < 2:
        raise ValueError("n_splits must be at least 2.")
    return StratifiedGroupKFold(n_splits=n_splits, shuffle=True, random_state=int(seed))


def make_grouped_holdout(
    frame: pd.DataFrame,
    *,
    group_column: str,
    seed: int,
    n_splits: int = 10,
    balance_group_sizes: bool = False,
) -> Tuple[np.ndarray, np.ndarray]:
    """Return one predeclared group-aware outer fold.

    Laboratory labels are used only for approximate balance; molecular or
    scaffold groups define the hard non-overlap boundary.  Taking the first fold
    is predeclared and deterministic, so no outcome-based fold selection occurs.
    """

    if group_column not in frame.columns:
        raise ValueError(f"Unknown group column: {group_column}")
    if "Lab" not in frame.columns:
        raise ValueError("Dataset must contain Lab for stratification.")

    index_values = frame.index.to_numpy(dtype=int)
    if balance_group_sizes:
        # GroupKFold's greedy assignment balances row counts substantially
        # better than a stratified splitter when one scaffold contains many
        # observations. Randomized surrogate group labels make tied group-size
        # assignments seed-dependent without using RT values or model results.
        original_groups = frame[group_column].astype(str).to_numpy()
        unique_groups = np.unique(original_groups)
        rng = np.random.default_rng(int(seed))
        shuffled_ranks = rng.permutation(len(unique_groups))
        rank_by_group = dict(zip(unique_groups, shuffled_ranks))
        randomized_groups = np.asarray(
            [rank_by_group[group] for group in original_groups],
            dtype=int,
        )
        candidates = list(
            GroupKFold(n_splits=n_splits).split(
                np.zeros((len(frame), 1), dtype=np.float32),
                groups=randomized_groups,
            )
        )
        all_labs = set(frame["Lab"].astype(str))
        target_rows = len(frame) / float(n_splits)

        def outcome_independent_score(
            candidate: Tuple[np.ndarray, np.ndarray],
        ) -> Tuple[int, float, int]:
            train_positions, test_positions = candidate
            train_labs = set(frame.iloc[train_positions]["Lab"].astype(str))
            test_labs = set(frame.iloc[test_positions]["Lab"].astype(str))
            missing_labs = len(all_labs - train_labs) + len(all_labs - test_labs)
            size_error = abs(len(test_positions) - target_rows)
            return missing_labs, float(size_error), int(test_positions.min())

        train_positions, test_positions = min(
            candidates,
            key=outcome_independent_score,
        )
        return np.sort(index_values[train_positions]), np.sort(index_values[test_positions])

    splitter = _stratified_group_splitter(n_splits=n_splits, seed=seed)
    train_positions, test_positions = next(
        splitter.split(
            np.zeros((len(frame), 1), dtype=np.float32),
            frame["Lab"].astype(str).to_numpy(),
            groups=frame[group_column].astype(str).to_numpy(),
        )
    )
    return np.sort(index_values[train_positions]), np.sort(index_values[test_positions])


def make_grouped_folds(
    frame: pd.DataFrame,
    development_indices: Sequence[int],
    *,
    group_column: str,
    seed: int,
    n_splits: int = 6,
) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Create group-aware inner folds expressed in original row indices."""

    development = np.asarray(development_indices, dtype=int)
    if len(np.unique(development)) != len(development):
        raise ValueError("development_indices contains duplicates.")
    subset = frame.loc[development]
    splitter = _stratified_group_splitter(n_splits=n_splits, seed=seed)
    folds: List[Tuple[np.ndarray, np.ndarray]] = []
    for train_positions, validation_positions in splitter.split(
        np.zeros((len(subset), 1), dtype=np.float32),
        subset["Lab"].astype(str).to_numpy(),
        groups=subset[group_column].astype(str).to_numpy(),
    ):
        folds.append(
            (
                np.sort(development[train_positions]),
                np.sort(development[validation_positions]),
            )
        )
    return folds


def compute_regression_metrics(y_true: Sequence[float], y_pred: Sequence[float]) -> Dict[str, float]:
    """Compute the common RT metrics with explicit signed bias and calibration."""

    y_true_array = np.asarray(y_true, dtype=float).reshape(-1)
    y_pred_array = np.asarray(y_pred, dtype=float).reshape(-1)
    n = _validate_equal_lengths(y_true=y_true_array, y_pred=y_pred_array)
    if n == 0:
        raise ValueError("At least one prediction is required.")

    residual = y_pred_array - y_true_array
    r2 = float(r2_score(y_true_array, y_pred_array)) if n >= 2 and np.ptp(y_true_array) > 0 else float("nan")
    if n >= 2 and np.ptp(y_true_array) > 0:
        slope, intercept = np.polyfit(y_true_array, y_pred_array, 1)
    else:
        slope, intercept = float("nan"), float("nan")
    return {
        "n": int(n),
        "r2": r2,
        "mae": float(mean_absolute_error(y_true_array, y_pred_array)),
        "rmse": float(np.sqrt(mean_squared_error(y_true_array, y_pred_array))),
        "bias": float(np.mean(residual)),
        "calibration_slope": float(slope),
        "calibration_intercept": float(intercept),
    }


def per_lab_metrics(
    y_true: Sequence[float],
    y_pred: Sequence[float],
    lab_labels: Sequence[object],
    normalization_ranges: Mapping[str, float] | None = None,
) -> pd.DataFrame:
    """Return per-laboratory metrics with optional development-range normalization."""

    y_true_array = np.asarray(y_true, dtype=float).reshape(-1)
    y_pred_array = np.asarray(y_pred, dtype=float).reshape(-1)
    labs = np.asarray(lab_labels).astype(str).reshape(-1)
    _validate_equal_lengths(y_true=y_true_array, y_pred=y_pred_array, lab_labels=labs)

    rows: List[Dict[str, object]] = []
    for lab in sorted(np.unique(labs)):
        mask = labs == lab
        row: Dict[str, object] = {"Lab": lab}
        row.update(compute_regression_metrics(y_true_array[mask], y_pred_array[mask]))
        if normalization_ranges is None:
            normalization_range = float(np.ptp(y_true_array[mask]))
            normalization_source = "evaluation_subset"
        else:
            normalization_range = float(normalization_ranges.get(lab, float("nan")))
            normalization_source = "development_lab_rt_range"
        row["normalization_rt_range"] = normalization_range
        row["normalization_source"] = normalization_source
        if np.isfinite(normalization_range) and normalization_range > 0:
            row["nmae_by_rt_range"] = float(row["mae"]) / normalization_range
            row["nrmse_by_rt_range"] = float(row["rmse"]) / normalization_range
        else:
            row["nmae_by_rt_range"] = float("nan")
            row["nrmse_by_rt_range"] = float("nan")
        rows.append(row)
    return pd.DataFrame(rows)


def _metric_value(metric: str, y_true: np.ndarray, y_pred: np.ndarray) -> float:
    if metric == "mae":
        return float(mean_absolute_error(y_true, y_pred))
    if metric == "rmse":
        return float(np.sqrt(mean_squared_error(y_true, y_pred)))
    if metric == "r2":
        if len(y_true) < 2 or np.ptp(y_true) == 0:
            return float("nan")
        return float(r2_score(y_true, y_pred))
    raise ValueError(f"Unsupported metric: {metric}")


def paired_group_bootstrap(
    *,
    y_true: Sequence[float],
    candidate: Sequence[float],
    reference: Sequence[float],
    groups: Sequence[object],
    n_resamples: int,
    seed: int,
    confidence: float = 0.95,
) -> Dict[str, Dict[str, float]]:
    """Paired cluster bootstrap of candidate minus reference performance.

    Unique molecular groups, not individual rows, are resampled.  Negative
    MAE/RMSE differences favor the candidate; positive R2 differences favor it.
    """

    y_true_array = np.asarray(y_true, dtype=float).reshape(-1)
    candidate_array = np.asarray(candidate, dtype=float).reshape(-1)
    reference_array = np.asarray(reference, dtype=float).reshape(-1)
    group_array = np.asarray(groups).astype(str).reshape(-1)
    _validate_equal_lengths(
        y_true=y_true_array,
        candidate=candidate_array,
        reference=reference_array,
        groups=group_array,
    )
    if n_resamples < 1:
        raise ValueError("n_resamples must be positive.")
    if not 0 < confidence < 1:
        raise ValueError("confidence must lie between 0 and 1.")

    unique_groups = np.unique(group_array)
    positions_by_group = {group: np.flatnonzero(group_array == group) for group in unique_groups}
    rng = np.random.default_rng(int(seed))
    bootstrap_differences: Dict[str, List[float]] = {metric: [] for metric in ("mae", "rmse", "r2")}

    for _ in range(int(n_resamples)):
        sampled_groups = rng.choice(unique_groups, size=len(unique_groups), replace=True)
        sampled_positions = np.concatenate([positions_by_group[group] for group in sampled_groups])
        sampled_true = y_true_array[sampled_positions]
        sampled_candidate = candidate_array[sampled_positions]
        sampled_reference = reference_array[sampled_positions]
        for metric in bootstrap_differences:
            candidate_value = _metric_value(metric, sampled_true, sampled_candidate)
            reference_value = _metric_value(metric, sampled_true, sampled_reference)
            difference = candidate_value - reference_value
            if np.isfinite(difference):
                bootstrap_differences[metric].append(float(difference))

    alpha = 1.0 - confidence
    output: Dict[str, Dict[str, float]] = {}
    for metric, values in bootstrap_differences.items():
        point = _metric_value(metric, y_true_array, candidate_array) - _metric_value(
            metric, y_true_array, reference_array
        )
        distribution = np.asarray(values, dtype=float)
        if distribution.size == 0:
            low = high = float("nan")
        else:
            low, high = np.quantile(distribution, [alpha / 2.0, 1.0 - alpha / 2.0])
        output[metric] = {
            "difference_point": float(point),
            "ci_low": float(low),
            "ci_high": float(high),
            "confidence": float(confidence),
            "n_groups": int(len(unique_groups)),
            "n_valid_resamples": int(distribution.size),
        }
    return output


def build_classical_feature_sets(
    *,
    fingerprints: np.ndarray,
    descriptors: np.ndarray,
    lab_labels: Sequence[object],
) -> Mapping[str, np.ndarray]:
    """Construct transparent classical baselines and descriptor ablations."""

    fingerprints_array = np.asarray(fingerprints, dtype=np.float32)
    descriptors_array = np.asarray(descriptors, dtype=np.float32)
    labs = np.asarray(lab_labels).astype(str).reshape(-1, 1)
    if fingerprints_array.ndim != 2 or descriptors_array.ndim != 2:
        raise ValueError("fingerprints and descriptors must be two-dimensional.")
    _validate_equal_lengths(
        fingerprints=fingerprints_array,
        descriptors=descriptors_array,
        lab_labels=labs,
    )

    encoder = OneHotEncoder(handle_unknown="ignore", sparse_output=False, dtype=np.float32)
    lab_one_hot = encoder.fit_transform(labs)
    return {
        "descriptor_only": np.column_stack([descriptors_array, lab_one_hot]).astype(np.float32),
        "fingerprint_only": np.column_stack([fingerprints_array, lab_one_hot]).astype(np.float32),
        "fingerprint_plus_descriptors": np.column_stack(
            [fingerprints_array, descriptors_array, lab_one_hot]
        ).astype(np.float32),
        "lab_one_hot": lab_one_hot.astype(np.float32),
    }


def ensure_new_output_dir(path: Path | str) -> Path:
    """Create an isolated output directory and refuse any nonempty target."""

    target = Path(path)
    if target.exists() and any(target.iterdir()):
        raise FileExistsError(f"Refusing to overwrite nonempty output directory: {target}")
    target.mkdir(parents=True, exist_ok=True)
    return target.resolve()