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"