scripts/diagnose_frozen_genre_outputs.py

#!/usr/bin/env python3
"""Post-hoc diagnostics for the frozen genre-representation benchmark.

This script never changes the frozen split, embeddings, or preregistered verdict.
It adds controls, low-data curves, and test-set slices intended to explain a
negative result.  All preprocessing and slice thresholds are fit on train data.
"""

from __future__ import annotations

import argparse
from collections import Counter
import json
import math
from pathlib import Path
from typing import Any

import numpy as np
from rdkit import Chem
from rdkit import RDLogger
from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
from sklearn.metrics import accuracy_score, balanced_accuracy_score
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

from pino.genre_benchmark import SOLVENTS, composition_features, load_jsonl, nearest_centroid_predict


ELEMENTS = ("C", "H", "N", "O", "F", "P", "S", "Cl", "Br", "I", "B", "Si")
RDLogger.DisableLog("rdApp.error")


def identity(component: dict[str, Any]) -> str:
    return str(component.get("cas") or component.get("smiles") or component.get("name") or "").strip().lower()


def structural_features(records: list[dict[str, Any]]) -> tuple[np.ndarray, list[dict[str, Any]]]:
    """Weighted element counts plus label-free structural-quality diagnostics."""
    matrix, metadata = [], []
    for row in records:
        elements: Counter[str] = Counter()
        active, parsed, missing, invalid, charged, fragments = 0, 0, 0, 0, 0, 0
        for component in row.get("formula", []):
            if identity(component) in SOLVENTS:
                continue
            active += 1
            smiles = str(component.get("smiles") or "").strip()
            weight = max(float(component.get("weight_fraction", 0.0)), 0.0)
            if not smiles:
                missing += 1
                continue
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                invalid += 1
                continue
            parsed += 1
            fragments += int(len(Chem.GetMolFrags(mol)) > 1)
            charged += int(any(atom.GetFormalCharge() for atom in mol.GetAtoms()))
            for atom in mol.GetAtoms():
                elements[atom.GetSymbol()] += weight
        other = sum(value for key, value in elements.items() if key not in ELEMENTS)
        vector = [elements[e] for e in ELEMENTS] + [other, active, parsed, missing, invalid, charged, fragments]
        matrix.append(vector)
        metadata.append({
            "formula_size": active,
            "missing_smiles": missing,
            "invalid_smiles": invalid,
            "charged_components": charged,
            "multifragment_components": fragments,
            "elements": sorted(elements),
        })
    return np.asarray(matrix, dtype=float), metadata


def scores(y: np.ndarray, pred: np.ndarray) -> dict[str, float]:
    return {
        "accuracy": float(accuracy_score(y, pred)),
        "balanced_accuracy": float(balanced_accuracy_score(y, pred)),
    }


def mlp_hidden(input_dim: int, classes: int, budget: int) -> int:
    # (d + 1)h + (h + 1)c; use a common parameter budget across representations.
    return max(1, int(round((budget - classes) / (input_dim + classes + 1))))


def fit_controls(
    train_x: dict[str, np.ndarray], test_x: dict[str, np.ndarray], y_train: np.ndarray, y_test: np.ndarray,
    *, seed: int, parameter_budget: int,
) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
    results: dict[str, Any] = {}
    predictions: dict[str, np.ndarray] = {}
    classes = len(np.unique(y_train))
    for name, x_train in train_x.items():
        pred = nearest_centroid_predict(x_train, y_train, test_x[name])
        predictions[f"{name}_nearest_centroid"] = pred
        results[f"{name}_nearest_centroid"] = scores(y_test, pred)
        hidden = mlp_hidden(x_train.shape[1], classes, parameter_budget)
        model = make_pipeline(
            StandardScaler(),
            MLPClassifier(hidden_layer_sizes=(hidden,), max_iter=1000, early_stopping=True,
                          validation_fraction=0.15, random_state=seed),
        )
        model.fit(x_train, y_train)
        pred = model.predict(test_x[name])
        predictions[f"{name}_mlp"] = pred
        results[f"{name}_mlp"] = {
            **scores(y_test, pred), "hidden_units": hidden,
            "nominal_parameters": (x_train.shape[1] + 1) * hidden + (hidden + 1) * classes,
        }
    for cls in (ExtraTreesClassifier, RandomForestClassifier):
        model = cls(n_estimators=400, min_samples_leaf=2, class_weight="balanced", n_jobs=-1, random_state=seed)
        model.fit(train_x["composition"], y_train)
        pred = model.predict(test_x["composition"])
        key = f"composition_{cls.__name__.replace('Classifier', '').lower()}"
        predictions[key] = pred
        results[key] = scores(y_test, pred)
    majority = Counter(y_train).most_common(1)[0][0]
    pred = np.repeat(majority, len(y_test))
    predictions["majority"] = pred
    results["majority"] = scores(y_test, pred)
    return results, predictions


def shuffled_control(train: np.ndarray, test: np.ndarray, y_train: np.ndarray, y_test: np.ndarray, seed: int) -> dict[str, float]:
    rng = np.random.default_rng(seed)
    return scores(y_test, nearest_centroid_predict(train[rng.permutation(len(train))], y_train, test[rng.permutation(len(test))]))


def slice_report(y: np.ndarray, predictions: dict[str, np.ndarray], masks: dict[str, np.ndarray]) -> dict[str, Any]:
    report = {}
    for name, mask in masks.items():
        count = int(mask.sum())
        if count < 5:
            continue
        report[name] = {"n": count, "models": {key: scores(y[mask], pred[mask]) for key, pred in predictions.items()}}
    return report


def low_data_curve(x_train: dict[str, np.ndarray], x_test: dict[str, np.ndarray], y_train: np.ndarray,
                   y_test: np.ndarray, seed: int) -> list[dict[str, Any]]:
    rng = np.random.default_rng(seed)
    by_label = {label: np.flatnonzero(y_train == label) for label in np.unique(y_train)}
    output = []
    for fraction in (0.1, 0.25, 0.5, 1.0):
        selected = np.concatenate([
            rng.choice(indices, size=max(1, int(math.ceil(len(indices) * fraction))), replace=False)
            for indices in by_label.values()
        ])
        row = {"fraction": fraction, "n_train": int(len(selected)), "models": {}}
        for name in ("learned", "composition", "elements"):
            pred = nearest_centroid_predict(x_train[name][selected], y_train[selected], x_test[name])
            row["models"][name] = scores(y_test, pred)
        output.append(row)
    return output


def diagnose_split(split: dict[str, Any], seed: int, budget: int) -> dict[str, Any]:
    train, test = load_jsonl(split["train_records"]), load_jsonl(split["test_records"])
    learned_train, learned_test = np.load(split["learned_train"]), np.load(split["learned_test"])
    structural_train, train_meta = structural_features(train)
    structural_test, test_meta = structural_features(test)
    train_x = {"learned": learned_train, "composition": composition_features(train), "elements": structural_train}
    test_x = {"learned": learned_test, "composition": composition_features(test), "elements": structural_test}
    y_train, y_test = np.asarray([r["genre"] for r in train]), np.asarray([r["genre"] for r in test])
    controls, predictions = fit_controls(train_x, test_x, y_train, y_test, seed=seed, parameter_budget=budget)
    controls["learned_shuffled_rows"] = shuffled_control(learned_train, learned_test, y_train, y_test, seed)

    sizes_train = np.asarray([m["formula_size"] for m in train_meta])
    sizes_test = np.asarray([m["formula_size"] for m in test_meta])
    q1, q2 = np.quantile(sizes_train, [1 / 3, 2 / 3])
    element_frequency = Counter(e for meta in train_meta for e in set(meta["elements"]))
    rare = {e for e, count in element_frequency.items() if count < max(5, int(0.01 * len(train)))}
    ood_masks, ood_thresholds = {}, {}
    for feature_name in ("composition", "learned"):
        mean, std = train_x[feature_name].mean(0), train_x[feature_name].std(0)
        std[std == 0] = 1
        train_z = (train_x[feature_name] - mean) / std
        test_z = (test_x[feature_name] - mean) / std
        distance = np.sqrt(np.square(test_z[:, None, :] - train_z[None, :, :]).sum(2).min(1))
        # Select by rank so tied distances (common for singleton formulas) do not
        # silently turn a top-quartile diagnostic into the entire test set.
        top = np.zeros(len(test), dtype=bool)
        top[np.argsort(distance, kind="stable")[-max(1, math.ceil(len(test) / 4)):]] = True
        ood_masks[f"{feature_name}_ood_top_quartile"] = top
        ood_thresholds[feature_name] = float(distance[top].min())
    masks = {
        f"size_small_le_{q1:g}": sizes_test <= q1,
        f"size_medium_{q1:g}_to_{q2:g}": (sizes_test > q1) & (sizes_test <= q2),
        f"size_large_gt_{q2:g}": sizes_test > q2,
        "contains_rare_element": np.asarray([bool(set(m["elements"]) & rare) for m in test_meta]),
        "contains_charged_component": np.asarray([m["charged_components"] > 0 for m in test_meta]),
        "structurally_ambiguous": np.asarray([m["missing_smiles"] + m["invalid_smiles"] + m["multifragment_components"] > 0 for m in test_meta]),
        **ood_masks,
    }
    for label in np.unique(y_test):
        masks[f"genre_{label}"] = y_test == label
    return {
        "name": split.get("name", "split"), "n_train": len(train), "n_test": len(test),
        "controls": controls, "slice_thresholds_fit_on_train": {"size_tertiles": [float(q1), float(q2)],
        "rare_element_record_threshold": max(5, int(0.01 * len(train))), "rare_elements": sorted(rare),
        "ood_top_quartile_min_distance": ood_thresholds},
        "slice_counts": {name: int(mask.sum()) for name, mask in masks.items()},
        "slices": slice_report(y_test, predictions, masks),
        "low_data_curve": low_data_curve(train_x, test_x, y_train, y_test, seed),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--seed", type=int, default=20260714)
    parser.add_argument("--parameter-budget", type=int, default=4096)
    args = parser.parse_args()
    manifest = json.loads(Path(args.manifest).read_text())
    report = {
        "analysis_status": "post_hoc_diagnostic; does not alter preregistered verdict",
        "parameter_budget": args.parameter_budget,
        "splits": [diagnose_split(split, args.seed + i, args.parameter_budget) for i, split in enumerate(manifest)],
    }
    Path(args.output).write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()