#!/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()