| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| from numpy_model import NumpyStableGateDANN, load_npz_text |
|
|
|
|
| LABEL_COLUMNS = {"label", "disease", "class", "diagnosis", "group"} |
| BATCH_COLUMNS = {"batch", "dataset", "platform", "cohort", "source"} |
|
|
|
|
| @dataclass |
| class PredictionBundle: |
| table: pd.DataFrame |
| matched_genes: int |
| expected_genes: int |
| matched_selected_genes: int |
| expected_selected_genes: int |
| missing_genes: list[str] |
| missing_selected_genes: list[str] |
| mode: str |
|
|
|
|
| @dataclass |
| class FoldRuntime: |
| name: str |
| genes: list[str] |
| selected_genes: list[str] |
| imputer_statistics: np.ndarray |
| scaler_center: np.ndarray |
| scaler_scale: np.ndarray |
| pd_class_index: int |
| model: NumpyStableGateDANN |
|
|
| def predict_probability(self, features: pd.DataFrame) -> np.ndarray: |
| aligned = features.reindex(columns=self.genes).to_numpy(dtype=np.float32) |
| aligned = np.where(np.isfinite(aligned), aligned, 0.0) |
| transformed = (aligned - self.scaler_center) / self.scaler_scale |
| logits = self.model.predict_logits(transformed.astype(np.float32)) |
| logits = logits - logits.max(axis=1, keepdims=True) |
| probabilities = np.exp(logits) |
| probabilities /= probabilities.sum(axis=1, keepdims=True) |
| return probabilities[:, self.pd_class_index] |
|
|
|
|
| class StableGateDeploymentPredictor: |
| """Real four-checkpoint StableGate-Hybrid DANN LODO ensemble.""" |
|
|
| def __init__(self, artifact_dir: str | Path = "artifacts"): |
| self.artifact_dir = Path(artifact_dir) |
| self.ensemble_dir = self.artifact_dir / "production_ensemble" |
| self.manifest = json.loads((self.ensemble_dir / "manifest.json").read_text(encoding="utf-8")) |
| self.minimum_gene_coverage = float(self.manifest["minimum_gene_coverage"]) |
| self.expected_selected_genes = int(self.manifest["top_k"]) |
| self.folds = [self._load_fold(fold) for fold in self.manifest["folds"]] |
| self.expected_genes = self.folds[0].genes |
| self.mode = "real_lodo_ensemble" |
|
|
| def _load_fold(self, fold: dict) -> FoldRuntime: |
| preprocess = load_npz_text(self.ensemble_dir / fold["preprocess"]) |
| genes = [str(g).upper() for g in preprocess["genes"].tolist()] |
| weights_path = self.ensemble_dir / fold["model"] |
| weights = load_npz_text(weights_path) |
| selected_genes = [ |
| gene for gene, selected in zip(genes, weights["gate_mask"].astype(bool)) if selected |
| ] |
| label_classes = [str(label) for label in preprocess["label_classes"].tolist()] |
| pd_class_index = label_classes.index("1") |
| scaler_scale = preprocess["scaler_scale"].astype(np.float32) |
| scaler_scale = np.where(scaler_scale == 0, 1.0, scaler_scale) |
| return FoldRuntime( |
| name=fold["name"], |
| genes=genes, |
| selected_genes=selected_genes, |
| imputer_statistics=preprocess["imputer_statistics"].astype(np.float32), |
| scaler_center=preprocess["scaler_center"].astype(np.float32), |
| scaler_scale=scaler_scale, |
| pd_class_index=pd_class_index, |
| model=NumpyStableGateDANN(weights_path), |
| ) |
|
|
| def predict_file(self, file_path: str | Path) -> PredictionBundle: |
| return self.predict(self._read_expression_file(file_path)) |
|
|
| def predict(self, frame: pd.DataFrame) -> PredictionBundle: |
| clean = self._sanitize_table(frame) |
| sample_column = self._detect_sample_column(clean) |
| sample_ids = clean[sample_column].astype(str).tolist() |
|
|
| excluded = {sample_column} |
| excluded.update(column for column in clean.columns if column.lower() in LABEL_COLUMNS | BATCH_COLUMNS) |
| features = clean.drop(columns=list(excluded), errors="ignore") |
| features.columns = [str(column).strip().upper() for column in features.columns] |
| features = features.loc[:, ~features.columns.duplicated(keep="first")] |
| features = features.apply(pd.to_numeric, errors="coerce") |
|
|
| input_genes = set(features.columns) |
| missing_genes = [gene for gene in self.expected_genes if gene not in input_genes] |
| matched_genes = len(self.expected_genes) - len(missing_genes) |
| coverage = matched_genes / len(self.expected_genes) |
|
|
| fold_probabilities = np.column_stack([fold.predict_probability(features) for fold in self.folds]) |
| mean_probability = fold_probabilities.mean(axis=1) |
| probability_std = fold_probabilities.std(axis=1) |
| predicted_pd = mean_probability >= 0.5 |
| fold_votes = (fold_probabilities >= 0.5).astype(int) |
| agreement = np.where(predicted_pd, fold_votes.sum(axis=1), len(self.folds) - fold_votes.sum(axis=1)) |
|
|
| selected_coverage = [] |
| missing_selected_union: set[str] = set() |
| for fold in self.folds: |
| missing = set(fold.selected_genes) - input_genes |
| missing_selected_union.update(missing) |
| selected_coverage.append(self.expected_selected_genes - len(missing)) |
|
|
| result = pd.DataFrame( |
| { |
| "Sample": sample_ids, |
| "PD probability": [f"{value * 100:.1f}%" for value in mean_probability], |
| "Ensemble SD": [f"{value * 100:.1f}%" for value in probability_std], |
| "Predicted label": np.where(predicted_pd, "PD", "Control"), |
| "Model agreement": [f"{value}/{len(self.folds)}" for value in agreement], |
| "Confidence": [ |
| self._confidence(float(probability), float(std), coverage) |
| for probability, std in zip(mean_probability, probability_std) |
| ], |
| } |
| ) |
| result["_probability_raw"] = mean_probability |
| return PredictionBundle( |
| table=result, |
| matched_genes=matched_genes, |
| expected_genes=len(self.expected_genes), |
| matched_selected_genes=min(selected_coverage), |
| expected_selected_genes=self.expected_selected_genes, |
| missing_genes=missing_genes, |
| missing_selected_genes=sorted(missing_selected_union), |
| mode=self.mode, |
| ) |
|
|
| @staticmethod |
| def _read_expression_file(file_path: str | Path) -> pd.DataFrame: |
| path = Path(file_path) |
| if path.suffix.lower() == ".csv": |
| return pd.read_csv(path) |
| if path.suffix.lower() in {".xlsx", ".xls"}: |
| return pd.read_excel(path) |
| raise ValueError("Only .csv, .xlsx, and .xls files are supported.") |
|
|
| @staticmethod |
| def _sanitize_table(frame: pd.DataFrame) -> pd.DataFrame: |
| if frame.empty: |
| raise ValueError("The uploaded expression matrix is empty.") |
| if frame.shape[1] < 2: |
| raise ValueError("The expression matrix must contain one sample column and gene columns.") |
| clean = frame.copy() |
| clean.columns = [str(column).strip() for column in clean.columns] |
| return clean |
|
|
| @staticmethod |
| def _detect_sample_column(frame: pd.DataFrame) -> str: |
| preferred = {"sample", "sample_id", "samples", "id", "name"} |
| for column in frame.columns: |
| if column.lower() in preferred: |
| return column |
| return frame.columns[0] |
|
|
| @staticmethod |
| def _confidence(probability: float, ensemble_std: float, gene_coverage: float) -> str: |
| if gene_coverage < 0.90: |
| return "low" |
| margin = abs(probability - 0.5) |
| if margin >= 0.30 and ensemble_std <= 0.10: |
| return "high" |
| if margin >= 0.15 and ensemble_std <= 0.20: |
| return "moderate" |
| return "low" |
|
|
