Spaces:
Running on Zero
Running on Zero
| """ | |
| Single-person inference for the frozen CKB 17-variable model package. | |
| The package has two separate layers: | |
| 1. The cluster label is produced by the frozen CKB nearest-centroid model. | |
| Its SHAP files explain a multiclass random-forest surrogate trained to | |
| reproduce deployed C1/C2/C3 labels; they do not change the cluster label. | |
| 2. Each incident outcome is predicted by the selected first-stage classifier. | |
| The outcome SHAP files explain the corresponding second-stage regression | |
| surrogate for that classifier's probability score. The classifier's | |
| probability and the surrogate score are therefore reported separately. | |
| The command-line entry point is intentionally small, but the same file can | |
| be imported as a library: | |
| from predict_single import predict_one | |
| result = predict_one({ | |
| "sex": 1, "age": 55, "edu_level": 2, | |
| "marital_status": 1, "work": 1, "retire": 0, | |
| "hh_size": 3, "smoking": 0, "alcohol": 1, | |
| "height_cm": 165, "weight_kg": 65, "waist_cm": 82, | |
| "sbp_mmhg": 125, "dbp_mmhg": 78, "bp_drugs": 0, | |
| "self_health": 2, "chronic_pain": 0, | |
| }) | |
| The outcome high-risk flag is defined as probability >= the frozen full-CKB | |
| Youden threshold recorded in manifest.json. No missing-value imputation, | |
| re-standardisation, feature selection, or cohort-specific refitting is done. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import datetime as dt | |
| import json | |
| import warnings | |
| from pathlib import Path | |
| from typing import Any, Mapping | |
| import joblib | |
| import numpy as np | |
| import pandas as pd | |
| import shap | |
| PACKAGE_ROOT = Path(__file__).resolve().parent | |
| MANIFEST_PATH = PACKAGE_ROOT / "manifest.json" | |
| def _json_default(value: Any) -> Any: | |
| """Convert common NumPy/path values into JSON-safe values.""" | |
| if isinstance(value, (np.integer,)): | |
| return int(value) | |
| if isinstance(value, (np.floating,)): | |
| return float(value) | |
| if isinstance(value, (np.bool_,)): | |
| return bool(value) | |
| if isinstance(value, Path): | |
| return str(value) | |
| if isinstance(value, np.ndarray): | |
| return value.tolist() | |
| raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") | |
| def _read_json(path: Path) -> dict[str, Any]: | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| def _safe_float(value: Any) -> float | None: | |
| try: | |
| number = float(value) | |
| except (TypeError, ValueError): | |
| return None | |
| return number if np.isfinite(number) else None | |
| class CKBSingleSamplePredictor: | |
| """Load the packaged models and make one complete-case prediction.""" | |
| def __init__(self, package_root: str | Path | None = None) -> None: | |
| self.package_root = Path(package_root or PACKAGE_ROOT).resolve() | |
| self.manifest = _read_json(self.package_root / "manifest.json") | |
| self.features = list(self.manifest["features"]) | |
| self.feature_labels = dict(self.manifest.get("feature_labels", {})) | |
| self.outcome_meta = dict(self.manifest["outcomes"]) | |
| cluster_dir = self.package_root / "cluster" | |
| with np.load(cluster_dir / "frozen_model.npz", allow_pickle=True) as data: | |
| self.cluster_centers = np.asarray(data["centers"], dtype=np.float64) | |
| self.cluster_mu = np.asarray(data["mu"], dtype=np.float64) | |
| self.cluster_sigma = np.asarray(data["sigma"], dtype=np.float64) | |
| self.cluster_log_applied = np.asarray(data["log_applied"], dtype=bool) | |
| frozen_features = [str(x) for x in data["features"].tolist()] | |
| self.cluster_k = int(np.asarray(data["K"]).ravel()[0]) | |
| if frozen_features != self.features: | |
| raise ValueError( | |
| "The frozen cluster feature order does not match manifest.json: " | |
| f"{frozen_features} != {self.features}" | |
| ) | |
| if self.cluster_centers.shape != (self.cluster_k, len(self.features)): | |
| raise ValueError("Unexpected frozen cluster-centre shape") | |
| self.cluster_thresholds = _read_json(cluster_dir / "ood_thresholds.json") | |
| calibration_path = cluster_dir / "calibration_model.npz" | |
| self.cluster_calibration = None | |
| if calibration_path.exists(): | |
| with np.load(calibration_path, allow_pickle=True) as data: | |
| self.cluster_calibration = { | |
| "coef": np.asarray(data["coef"], dtype=np.float64), | |
| "intercept": np.asarray(data["intercept"], dtype=np.float64), | |
| "d_std": float(np.asarray(data["d_std"]).ravel()[0]), | |
| } | |
| self._cluster_shap_model = None | |
| self._cluster_shap_explainer = None | |
| self._outcome_cache: dict[str, dict[str, Any]] = {} | |
| def _validate_input(self, values: Mapping[str, Any]) -> tuple[pd.DataFrame, list[str]]: | |
| if not isinstance(values, Mapping): | |
| raise TypeError("Input must be a mapping with the 17 required feature names") | |
| missing = [feature for feature in self.features if feature not in values] | |
| if missing: | |
| raise ValueError(f"Missing required input variables: {missing}") | |
| ignored = sorted(set(values) - set(self.features)) | |
| frame = pd.DataFrame([{feature: values[feature] for feature in self.features}]) | |
| for feature in self.features: | |
| frame[feature] = pd.to_numeric(frame[feature], errors="coerce") | |
| if frame.isna().any(axis=None): | |
| bad = frame.columns[frame.isna().any()].tolist() | |
| raise ValueError(f"All 17 variables must be finite numeric values; invalid: {bad}") | |
| array = frame.to_numpy(dtype=np.float64) | |
| if not np.isfinite(array).all(): | |
| raise ValueError("All 17 variables must be finite numeric values") | |
| log_features = [ | |
| feature for feature, use_log in zip(self.features, self.cluster_log_applied) | |
| if use_log | |
| ] | |
| negative_log = [feature for feature in log_features if float(frame.iloc[0][feature]) < 0] | |
| if negative_log: | |
| raise ValueError( | |
| "The following variables are log1p-transformed by the frozen cluster " | |
| f"model and cannot be negative: {negative_log}" | |
| ) | |
| return frame, ignored | |
| def _cluster_preprocess(self, frame: pd.DataFrame) -> np.ndarray: | |
| raw = frame[self.features].to_numpy(dtype=np.float64, copy=True) | |
| for index, use_log in enumerate(self.cluster_log_applied): | |
| if use_log: | |
| # This is the locked external-assignment rule from the CKB | |
| # validation code. Input validation above prevents negatives. | |
| raw[:, index] = np.log1p(np.maximum(raw[:, index], 0.0)) | |
| denominator = np.where(self.cluster_sigma > 1e-8, self.cluster_sigma, 1.0) | |
| return (raw - self.cluster_mu) / denominator | |
| def _cluster_confidence(self, distances: np.ndarray) -> np.ndarray: | |
| calibration = self.cluster_calibration | |
| if calibration is None: | |
| return np.full(distances.shape[0], np.nan, dtype=np.float64) | |
| centered = distances - distances.min(axis=1, keepdims=True) | |
| centered = np.clip(centered / (calibration["d_std"] + 1e-8), -10, 10) | |
| logits = centered @ calibration["coef"].T + calibration["intercept"] | |
| logits -= logits.max(axis=1, keepdims=True) | |
| probabilities = np.exp(logits) | |
| probabilities /= probabilities.sum(axis=1, keepdims=True) | |
| return probabilities.max(axis=1) | |
| def _predict_cluster(self, frame: pd.DataFrame) -> dict[str, Any]: | |
| z = self._cluster_preprocess(frame) | |
| distances = np.linalg.norm( | |
| z[:, None, :] - self.cluster_centers[None, :, :], axis=2 | |
| ) | |
| nearest = distances.argmin(axis=1).astype(int) | |
| ordered = np.sort(distances, axis=1) | |
| d1 = ordered[:, 0] | |
| d2 = ordered[:, 1] | |
| margin = np.where(d2 > 0, (d2 - d1) / d2, 1.0) | |
| confidence = self._cluster_confidence(distances) | |
| index = int(nearest[0]) | |
| nearest_distance = float(d1[0]) | |
| second_distance = float(d2[0]) | |
| nearest_margin = float(margin[0]) | |
| nearest_confidence = _safe_float(confidence[0]) | |
| ood = int(nearest_distance > float(self.cluster_thresholds["p99_9_overall"])) | |
| per_cluster = self.cluster_thresholds.get("p99_5_per_cluster", {}) | |
| uncertain = int( | |
| nearest_distance > float(per_cluster.get(str(index), np.inf)) | |
| or nearest_margin < float(self.cluster_thresholds["margin_p5"]) | |
| ) | |
| confidence_p10 = self.cluster_thresholds.get("confidence_p10") | |
| low_conf = int( | |
| (confidence_p10 is not None and nearest_confidence is not None | |
| and nearest_confidence < float(confidence_p10)) | |
| or nearest_margin < float(self.cluster_thresholds["margin_p5"]) | |
| ) | |
| if ood: | |
| status = "unclassified_ood" | |
| elif uncertain: | |
| status = "uncertain" | |
| else: | |
| status = "assigned" | |
| label = f"C{index + 1}" | |
| result = { | |
| # provisional_cluster_label is always available; cluster_label is | |
| # only definitive when the frozen protocol calls it assigned. | |
| "provisional_cluster_label": label, | |
| "cluster_label": label if status == "assigned" else None, | |
| "cluster_index": index, | |
| "classification_status": status, | |
| "distance_nearest": nearest_distance, | |
| "distance_second": second_distance, | |
| "margin": nearest_margin, | |
| "calibrated_confidence": nearest_confidence, | |
| "ood_flag": bool(ood), | |
| "uncertain_flag": bool(uncertain), | |
| "low_confidence_flag": bool(low_conf), | |
| } | |
| self._ensure_cluster_shap() | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| surrogate_probability = self._cluster_shap_model.predict_proba( | |
| frame.to_numpy(dtype=np.float64) | |
| )[0] | |
| result["shap_surrogate_cluster_probability"] = { | |
| f"C{int(class_index) + 1}": float(probability) | |
| for class_index, probability in zip( | |
| self._cluster_shap_model.classes_, surrogate_probability | |
| ) | |
| } | |
| return result | |
| def _ensure_cluster_shap(self) -> None: | |
| if self._cluster_shap_model is not None: | |
| return | |
| cluster_dir = self.package_root / "cluster" | |
| self._cluster_shap_model = joblib.load( | |
| cluster_dir / "cluster_shap_surrogate_random_forest.joblib" | |
| ) | |
| self._cluster_shap_explainer = shap.TreeExplainer(self._cluster_shap_model) | |
| def _normalise_cluster_shap_values( | |
| values: Any, n_classes: int, n_features: int | |
| ) -> np.ndarray: | |
| """Return an array with shape (class, sample, feature).""" | |
| if isinstance(values, list): | |
| arrays = [np.asarray(item, dtype=np.float64) for item in values] | |
| return np.stack(arrays, axis=0) | |
| array = np.asarray(values, dtype=np.float64) | |
| if array.ndim == 2: | |
| return array[None, :, :] | |
| if array.ndim != 3: | |
| raise ValueError(f"Unexpected multiclass SHAP shape: {array.shape}") | |
| if array.shape[0] == 1 and array.shape[1] == n_features: | |
| return np.transpose(array, (2, 0, 1)) | |
| if array.shape[0] == n_classes and array.shape[2] == n_features: | |
| return array | |
| if array.shape[0] == n_classes and array.shape[1] == n_features: | |
| return array | |
| raise ValueError(f"Could not normalise multiclass SHAP shape: {array.shape}") | |
| def _cluster_force_plots( | |
| self, frame: pd.DataFrame, output_dir: Path | |
| ) -> dict[str, str]: | |
| self._ensure_cluster_shap() | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| raw_values = self._cluster_shap_explainer.shap_values( | |
| frame.to_numpy(dtype=np.float64) | |
| ) | |
| n_classes = len(self._cluster_shap_model.classes_) | |
| values = self._normalise_cluster_shap_values( | |
| raw_values, n_classes=n_classes, n_features=len(self.features) | |
| ) | |
| expected = np.asarray(self._cluster_shap_explainer.expected_value, dtype=float).reshape(-1) | |
| plot_paths: dict[str, str] = {} | |
| feature_values = frame.iloc[0].to_numpy(dtype=np.float64) | |
| for class_position, class_index in enumerate(self._cluster_shap_model.classes_): | |
| class_label = f"C{int(class_index) + 1}" | |
| path = output_dir / f"single_sample_cluster_{class_label}_force.html" | |
| force = shap.force_plot( | |
| float(expected[class_position]), | |
| values[class_position, 0, :], | |
| feature_values, | |
| feature_names=[self.feature_labels.get(f, f) for f in self.features], | |
| matplotlib=False, | |
| ) | |
| shap.save_html(str(path), force) | |
| plot_paths[class_label] = str(path) | |
| return plot_paths | |
| def _load_outcome_bundle(self, outcome: str) -> dict[str, Any]: | |
| if outcome in self._outcome_cache: | |
| return self._outcome_cache[outcome] | |
| meta = self.outcome_meta[outcome] | |
| outcome_dir = self.package_root / "outcomes" / outcome | |
| bundle = { | |
| "classifier": joblib.load(outcome_dir / meta["classifier_file"]), | |
| "shap_surrogate": joblib.load(outcome_dir / meta["shap_surrogate_file"]), | |
| } | |
| self._outcome_cache[outcome] = bundle | |
| return bundle | |
| def _outcome_force_plot( | |
| self, | |
| outcome: str, | |
| frame: pd.DataFrame, | |
| bundle: dict[str, Any], | |
| output_dir: Path, | |
| ) -> tuple[str, float]: | |
| # The bundled explanation pickles were created on a different Python | |
| # runtime. In particular, PermutationExplainer serialises a Numba | |
| # dispatcher whose bytecode is not portable between Python versions. | |
| # Rebuild the explainer in the active runtime from the frozen surrogate | |
| # and its stored background sample instead of invoking that dispatcher. | |
| # The prediction-only API still never pays this cost. | |
| if "shap_explainer" not in bundle: | |
| meta = self.outcome_meta[outcome] | |
| outcome_dir = self.package_root / "outcomes" / outcome | |
| surrogate = bundle["shap_surrogate"] | |
| if hasattr(surrogate, "estimators_"): | |
| # Tree SHAP is deterministic for the frozen RandomForest | |
| # surrogates and avoids another cross-version pickle edge case. | |
| bundle["shap_explainer"] = shap.TreeExplainer(surrogate) | |
| else: | |
| stored_explainer = joblib.load( | |
| outcome_dir / meta["shap_explainer_file"] | |
| ) | |
| background = getattr( | |
| getattr(stored_explainer, "masker", None), "data", None | |
| ) | |
| if background is None: | |
| raise RuntimeError( | |
| f"Missing SHAP background sample for outcome {outcome}." | |
| ) | |
| bundle["shap_explainer"] = shap.PermutationExplainer( | |
| surrogate.predict, | |
| np.asarray(background, dtype=np.float64), | |
| ) | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| explanation = bundle["shap_explainer"]( | |
| frame.to_numpy(dtype=np.float64) | |
| ) | |
| values = np.asarray(explanation.values, dtype=np.float64) | |
| if values.ndim == 3: | |
| # The outcome SHAP artifacts are univariate. This branch makes a | |
| # malformed or future artifact fail loudly rather than silently | |
| # choosing an arbitrary class. | |
| if values.shape[-1] != 1: | |
| raise ValueError(f"Unexpected SHAP output shape for {outcome}: {values.shape}") | |
| values = values[:, :, 0] | |
| base_values = np.asarray(explanation.base_values, dtype=np.float64).reshape(-1) | |
| if values.shape != (1, len(self.features)) or base_values.size < 1: | |
| raise ValueError( | |
| f"Unexpected SHAP explanation shape for {outcome}: " | |
| f"values={values.shape}, base={base_values.shape}" | |
| ) | |
| surrogate_score = float(bundle["shap_surrogate"].predict(frame)[0]) | |
| path = output_dir / f"single_sample_{outcome}_force.html" | |
| force = shap.force_plot( | |
| float(base_values[0]), | |
| values[0], | |
| frame.iloc[0].to_numpy(dtype=np.float64), | |
| feature_names=[self.feature_labels.get(f, f) for f in self.features], | |
| matplotlib=False, | |
| ) | |
| shap.save_html(str(path), force) | |
| return str(path), surrogate_score | |
| def predict_one( | |
| self, | |
| values: Mapping[str, Any], | |
| output_dir: str | Path | None = None, | |
| make_force_plots: bool = True, | |
| sample_id: str | None = None, | |
| ) -> dict[str, Any]: | |
| """Predict the cluster and all 13 outcomes for one complete case. | |
| Parameters | |
| ---------- | |
| values: | |
| Mapping containing the 17 required variables. Extra keys are | |
| ignored and reported in ``input.ignored_fields``. | |
| output_dir: | |
| Directory for JSON and HTML force plots. If omitted, a timestamped | |
| folder is created under ``1-output/model/predictions``. | |
| make_force_plots: | |
| Save one HTML force plot for each cluster class and outcome when | |
| true. Setting false only skips plots; prediction probabilities and | |
| labels are unchanged. | |
| sample_id: | |
| Optional identifier written to the result JSON. It is never used | |
| as a model feature. | |
| """ | |
| frame, ignored = self._validate_input(values) | |
| if output_dir is None: | |
| stamp = dt.datetime.now().strftime("%Y%m%d_%H%M%S") | |
| output_path = self.package_root / "predictions" / stamp | |
| else: | |
| output_path = Path(output_dir).expanduser().resolve() | |
| output_path.mkdir(parents=True, exist_ok=True) | |
| force_dir = output_path / "force_plots" | |
| if make_force_plots: | |
| force_dir.mkdir(parents=True, exist_ok=True) | |
| result: dict[str, Any] = { | |
| "package_version": self.manifest.get("package_version"), | |
| "generated_at": dt.datetime.now(dt.timezone.utc).isoformat(), | |
| "sample_id": sample_id, | |
| "input": { | |
| "features": { | |
| feature: float(frame.iloc[0][feature]) for feature in self.features | |
| }, | |
| "ignored_fields": ignored, | |
| }, | |
| "cluster": self._predict_cluster(frame), | |
| "outcomes": {}, | |
| "force_plots": {}, | |
| "notes": [ | |
| "The definitive cluster is the frozen nearest-centroid assignment.", | |
| "Cluster SHAP explains a multiclass RF surrogate of deployed cluster labels.", | |
| "Outcome SHAP explains a regression surrogate of the selected classifier probability.", | |
| "High risk is probability >= the frozen full-CKB Youden threshold.", | |
| ], | |
| } | |
| if make_force_plots: | |
| result["force_plots"]["cluster"] = self._cluster_force_plots(frame, force_dir) | |
| X = frame.to_numpy(dtype=np.float64) | |
| for outcome, meta in self.outcome_meta.items(): | |
| bundle = self._load_outcome_bundle(outcome) | |
| classifier = bundle["classifier"] | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| probabilities = np.asarray(classifier.predict_proba(X), dtype=np.float64) | |
| classes = np.asarray(getattr(classifier, "classes_", [0, 1])) | |
| positive = np.where(classes == 1)[0] | |
| if positive.size != 1: | |
| raise ValueError(f"Outcome {outcome} does not have a unique positive class") | |
| probability = float(probabilities[0, int(positive[0])]) | |
| threshold = float(meta["youden_threshold"]) | |
| risk = "high_risk" if probability >= threshold else "low_risk" | |
| row: dict[str, Any] = { | |
| "outcome": meta.get("display_name", outcome), | |
| "model": meta["model"], | |
| "probability": probability, | |
| "youden_threshold": threshold, | |
| "risk_class": risk, | |
| "threshold_rule": "high_risk if probability >= Youden threshold", | |
| "shap_explanation_target": "selected_classifier_probability_surrogate", | |
| } | |
| if make_force_plots: | |
| path, surrogate_score = self._outcome_force_plot( | |
| outcome, frame, bundle, force_dir | |
| ) | |
| row["shap_surrogate_score"] = surrogate_score | |
| result["force_plots"][outcome] = path | |
| else: | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| row["shap_surrogate_score"] = float( | |
| bundle["shap_surrogate"].predict(frame)[0] | |
| ) | |
| result["outcomes"][outcome] = row | |
| result_path = output_path / "single_sample_prediction.json" | |
| result["result_json"] = str(result_path) | |
| result_path.write_text( | |
| json.dumps(result, ensure_ascii=False, indent=2, default=_json_default), | |
| encoding="utf-8", | |
| ) | |
| return result | |
| def predict_one( | |
| values: Mapping[str, Any], | |
| package_dir: str | Path | None = None, | |
| output_dir: str | Path | None = None, | |
| make_force_plots: bool = True, | |
| sample_id: str | None = None, | |
| ) -> dict[str, Any]: | |
| """Convenience wrapper around :class:`CKBSingleSamplePredictor`.""" | |
| predictor = CKBSingleSamplePredictor(package_dir) | |
| return predictor.predict_one( | |
| values, | |
| output_dir=output_dir, | |
| make_force_plots=make_force_plots, | |
| sample_id=sample_id, | |
| ) | |
| def _load_input_json(path: Path) -> tuple[dict[str, Any], str | None]: | |
| data = json.loads(path.read_text(encoding="utf-8")) | |
| if not isinstance(data, dict): | |
| raise ValueError("Input JSON must contain an object") | |
| sample_id = data.get("sample_id") | |
| if isinstance(data.get("features"), dict): | |
| data = data["features"] | |
| return data, str(sample_id) if sample_id is not None else None | |
| def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description="Run CKB frozen cluster and 13-outcome single-sample inference." | |
| ) | |
| parser.add_argument( | |
| "--input-json", required=True, type=Path, | |
| help="JSON object containing the 17 variables; example_input.json is provided.", | |
| ) | |
| parser.add_argument( | |
| "--output-dir", type=Path, default=None, | |
| help="Output folder; default is a timestamped folder under predictions/.", | |
| ) | |
| parser.add_argument( | |
| "--sample-id", default=None, | |
| help="Optional identifier; it is written to output but not used for prediction.", | |
| ) | |
| parser.add_argument( | |
| "--no-force-plots", action="store_true", | |
| help="Skip HTML force plots while retaining predictions and SHAP surrogate scores.", | |
| ) | |
| args = parser.parse_args() | |
| values, input_sample_id = _load_input_json(args.input_json) | |
| sample_id = args.sample_id if args.sample_id is not None else input_sample_id | |
| result = predict_one( | |
| values, | |
| output_dir=args.output_dir, | |
| make_force_plots=not args.no_force_plots, | |
| sample_id=sample_id, | |
| ) | |
| print(json.dumps({ | |
| "result_json": result["result_json"], | |
| "cluster": result["cluster"], | |
| "outcomes": result["outcomes"], | |
| "force_plots": result["force_plots"], | |
| }, ensure_ascii=False, indent=2, default=_json_default)) | |
| if __name__ == "__main__": | |
| main() | |