diff --git a/docs/api/models.rst b/docs/api/models.rst index e881e7894..de1e2e871 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -254,3 +254,4 @@ API Reference models/pyhealth.models.califorest models/pyhealth.models.XGBoostModel models/pyhealth.models.survival_mdn + models/pyhealth.models.personalized_survival_moe diff --git a/docs/api/models/pyhealth.models.personalized_survival_moe.rst b/docs/api/models/pyhealth.models.personalized_survival_moe.rst new file mode 100644 index 000000000..8a3708d6c --- /dev/null +++ b/docs/api/models/pyhealth.models.personalized_survival_moe.rst @@ -0,0 +1,9 @@ +pyhealth.models.PersonalizedSurvivalMoE +======================================= + +Personalized mixture-of-experts model for discrete-time survival prediction. + +.. autoclass:: pyhealth.models.PersonalizedSurvivalMoE + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/support2_survival_personalized_moe.py b/examples/support2_survival_personalized_moe.py new file mode 100644 index 000000000..cbf00d3c3 --- /dev/null +++ b/examples/support2_survival_personalized_moe.py @@ -0,0 +1,430 @@ +"""SUPPORT2 personalized MoE dropout ablation following Morrill et al. + + pip install SurvivalEVAL + python examples/support2_survival_personalized_moe.py --root data/support2 + python examples/support2_survival_personalized_moe.py --synthetic --epochs 2 + +Downloads missing data from UCI; loads patients through Support2Dataset in CSV +order. Uses the reference's 42 features, physiological/median imputation, missing +indicators, scaling, categorical embeddings and 20-quantile hday buckets. All +preprocessing fits the fixed seed-42 80/10/10 training split only. + +Defaults: seeds 42–46, one 128-unit layer, eight experts, 4-dimensional embeddings, +Adam (lr=0.0005, eps=1e-7), batch 64, up to 1,000 epochs with patience 10 on +validation negative log-likelihood (NLL). Each run keeps the checkpoint with the +lowest validation NLL. Each seed pairs initialization and batch order across dropout +0/0.1/0.3/0.5. The grid has 100 training-range points plus a final overflow point. + +Reports median-survival Harrell C-index, grid-averaged IPCW Brier/ECE (ten equal- +mass risk bins, reverse KM floored at 0.01), train/validation/test NLL and NLL gaps, +mean/sample SD and paired changes from dropout=0. Synthetic mode checks execution +only. Repeated training +seeds measure optimization variability on one split, not population uncertainty. +SUPPORT2 results (7283/911/911 patients; mean +/- sample SD over five seeds): + Dropout C-index IPCW Brier IPCW ECE NLL gap + 0.0 0.8085+/-0.0025 0.1433+/-0.0031 0.0444+/-0.0043 0.1949+/-0.0403 + 0.1 0.8105+/-0.0027 0.1409+/-0.0027 0.0413+/-0.0029 0.1881+/-0.0069 + 0.3 0.8075+/-0.0037 0.1424+/-0.0017 0.0426+/-0.0059 0.2046+/-0.0411 + 0.5 0.8043+/-0.0022 0.1431+/-0.0016 0.0473+/-0.0034 0.1871+/-0.0204 +Early stopping ended runs after 14–38 epochs. Dropout 0.1 had the best mean +discrimination and calibration; larger dropout did not consistently improve them. +Mean NLL at each run's selected checkpoint (lower is better): + Dropout Training Validation Test Validation - training + 0.0 2.0034 2.1984 2.1801 0.1949 + 0.1 2.0016 2.1896 2.1662 0.1881 + 0.3 1.9885 2.1932 2.1695 0.2046 + 0.5 2.0144 2.2015 2.1729 0.1871 +Dropout 0.1 also had the lowest mean validation and test NLL. Dropout 0.5 had the +smallest gap, showing that a smaller gap alone does not imply better predictions. +NLL excludes the expert-balancing penalty; test NLL does not select checkpoints. +Reference: https://github.com/ToddMorrill/survival-moe (7ce58fa). +""" + +import argparse +import copy +import json +from pathlib import Path +from tempfile import TemporaryDirectory +from urllib.request import urlopen + +import numpy as np +import pandas as pd +import torch +from sklearn.impute import SimpleImputer +from sklearn.model_selection import train_test_split +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import KBinsDiscretizer, StandardScaler +from torch.utils.data import DataLoader + +from pyhealth.datasets import Support2Dataset +from pyhealth.models import PersonalizedSurvivalMoE + +PHYSIOLOGICAL = { + "alb": 3.5, + "pafi": 333.3, + "bili": 1.01, + "crea": 1.01, + "bun": 6.51, + "wblc": 9.0, + "urine": 2502, +} +NUMERIC = list(PHYSIOLOGICAL) + [ + "age", + "num.co", + "scoma", + "charges", + "totcst", + "totmcst", + "avtisst", + "sps", + "aps", + "surv2m", + "surv6m", + "diabetes", + "dementia", + "prg2m", + "dnrday", + "meanbp", + "hrt", + "resp", + "temp", + "sod", + "ph", + "glucose", + "adls", + "adlsc", +] +CATEGORICAL = [ + "sex", + "dzgroup", + "dzclass", + "edu", + "income", + "race", + "ca", + "prg6m", + "dnr", + "adlp", +] + + +def preprocess(frame: pd.DataFrame, train: np.ndarray): + """Fit reference preprocessing on training rows; return tensors and vocab sizes.""" + numeric, categorical, cardinalities = [], [], [] + for col in NUMERIC: + values = ( + pd.to_numeric(frame[col], errors="coerce") + .replace([np.inf, -np.inf], np.nan) + .to_numpy() + .reshape(-1, 1) + ) + imputer = SimpleImputer( + strategy="constant" if col in PHYSIOLOGICAL else "median", + fill_value=PHYSIOLOGICAL.get(col), + keep_empty_features=True, + ) + transform = make_pipeline(imputer, StandardScaler()).fit(values[train]) + numeric.extend([transform.transform(values), np.isnan(values)]) + for col in CATEGORICAL: + values = frame[col].fillna("missing").astype(str) + vocabulary = ["missing"] + sorted(set(values.iloc[train]) - {"missing"}) + codes = values.map({v: i for i, v in enumerate(vocabulary)}).fillna(0) + categorical.append(codes.to_numpy(int)) + cardinalities.append(len(vocabulary)) + hday = pd.to_numeric(frame.hday, errors="coerce").to_numpy().reshape(-1, 1) + discretizer = KBinsDiscretizer(n_bins=20, encode="ordinal", strategy="quantile") + if "quantile_method" in discretizer.get_params(): + discretizer.set_params(quantile_method="linear") + bucket = make_pipeline(SimpleImputer(strategy="median"), discretizer) + bucket.fit(hday[train]) + categorical.append(bucket.transform(hday).ravel().astype(int)) + cardinalities.append(len(bucket[-1].bin_edges_[0])) + return ( + torch.tensor(np.concatenate(numeric, axis=1), dtype=torch.float32), + torch.tensor(np.stack(categorical, axis=1), dtype=torch.long), + tuple(cardinalities), + ) + + +def censoring_km(time: np.ndarray, event: np.ndarray): + """Reverse KM: tied events leave the risk set before censoring; floor at .01.""" + grid, probabilities, current = np.unique(time), [], 1.0 + for t in grid: + censored = np.sum((time == t) & (event == 0)) + at_risk = np.sum(time >= t) - np.sum((time == t) & (event == 1)) + if censored: + current *= 1 - censored / at_risk + probabilities.append(max(current, 0.01)) + return grid, np.asarray(probabilities) + + +def km_at(km, t): + """Evaluate the right-continuous censoring distribution.""" + grid, probability = km + index = np.searchsorted(grid, t, side="right") - 1 + return np.where(index < 0, 1.0, probability[np.maximum(index, 0)]) + + +def calibration_metrics(train_time, train_event, time, event, survival, grid): + """Reference IPCW Brier/ECE, averaged uniformly over the full time grid.""" + km = censoring_km(train_time, train_event) + clamped = np.minimum(time, train_time.max()) + event_weight = event / km_at(km, clamped) + horizon_weight = 1 / km_at(km, np.minimum(grid, train_time.max())) + brier, ece = [], [] + for j, horizon in enumerate(grid): + risk = 1 - survival[:, j] + target = time <= horizon + weights = np.where(clamped <= horizon, event_weight, horizon_weight[j]) + brier.append(np.mean((risk - target) ** 2 * weights)) + events = event_weight * target + survivors = (time > horizon) * horizon_weight[j] + error = 0.0 + for idx in np.array_split(np.argsort(risk, kind="stable"), min(10, len(time))): + numerator = events[idx].sum() + denominator = numerator + survivors[idx].sum() + empirical = numerator / denominator if denominator else 0.0 + error += len(idx) / len(time) * abs(risk[idx].mean() - empirical) + ece.append(error) + return {"ipcw_brier": float(np.mean(brier)), "ipcw_ece": float(np.mean(ece))} + + +def ensure_support2(path: Path) -> None: + """Download a missing CSV and normalize its ID column for Support2Dataset.""" + if path.exists(): + return + url = "https://archive.ics.uci.edu/static/public/880/data.csv" + with urlopen(url, timeout=30) as response: + frame = pd.read_csv(response).rename(columns={"id": "sno"}) + if not {"sno", "d.time", "death"}.issubset(frame.columns) or frame.empty: + raise ValueError( + "Downloaded SUPPORT2 CSV is missing patients or outcome columns" + ) + if frame.sno.isna().any() or frame.sno.duplicated().any(): + raise ValueError("Downloaded SUPPORT2 CSV has invalid patient IDs") + path.parent.mkdir(parents=True, exist_ok=True) + with TemporaryDirectory(dir=path.parent) as tmp: + pending = Path(tmp) / "support2.csv" + frame.to_csv(pending, index=False) + pending.replace(path) + + +def load_data(root: str | None, synthetic: bool, seed: int) -> pd.DataFrame: + """Load baseline fields via PyHealth, or generate independent synthetic data.""" + if synthetic: + rng = np.random.default_rng(seed) + frame = pd.DataFrame( + rng.normal(size=(240, len(NUMERIC) + 1)), columns=NUMERIC + ["hday"] + ) + for col in CATEGORICAL: + frame[col] = rng.choice(["a", "b"], len(frame)) + event_time = rng.exponential(100 * np.exp(-0.6 * frame.age)) + 1 + censor_time = rng.uniform(20, 250, len(frame)) + frame["d.time"] = np.minimum(event_time, censor_time) + frame["death"] = (event_time <= censor_time).astype(int) + return frame + ensure_support2(Path(root) / "support2.csv") + dataset = Support2Dataset( + root=root, tables=["support2"], cache_dir=Path(root) / ".pyhealth_cache" + ) + rows = [] + for pid in pd.read_csv(Path(root) / "support2.csv", dtype={"sno": str}).sno: + events = dataset.get_patient(pid).get_events(event_type="support2") + if len(events) != 1: + raise ValueError(f"Expected exactly one SUPPORT2 row for {pid}") + rows.append( + { + key: getattr(events[0], key, None) + for key in NUMERIC + CATEGORICAL + ["hday", "d.time", "death"] + } + ) + return pd.DataFrame(rows) + + +def summarize(rows): + """Summarize seed variability and paired changes from the no-dropout arm.""" + metrics = [ + "c_index", + "ipcw_brier", + "ipcw_ece", + "train_nll", + "val_nll", + "test_nll", + "gap", + ] + baseline = {row["seed"]: row for row in rows if row["dropout"] == 0} + summary = [] + for dropout in [0.0, 0.1, 0.3, 0.5]: + arm = [row for row in rows if row["dropout"] == dropout] + record = {"dropout": dropout} + for metric in metrics: + values = np.array([row[metric] for row in arm], dtype=float) + delta = values - [baseline[row["seed"]][metric] for row in arm] + record[metric] = { + "mean": float(values.mean()), + "sd": float(values.std(ddof=1)) if len(values) > 1 else None, + "paired_delta_mean": float(delta.mean()), + "paired_delta_sd": float(delta.std(ddof=1)) if len(delta) > 1 else None, + } + summary.append(record) + return summary + + +def run(args: argparse.Namespace) -> dict: + """Run paired dropout arms and report held-out results for every seed.""" + from SurvivalEVAL import LifelinesEvaluator + + if min(args.epochs, args.batch_size, args.patience) < 1 or args.time_bins < 3: + raise ValueError("epochs/batch_size/patience must be positive; time_bins >= 3") + if not args.seeds or len(set(args.seeds)) != len(args.seeds): + raise ValueError("Provide distinct training seeds") + torch.set_num_threads(1) + frame = load_data(args.root, args.synthetic, 42) + times = pd.to_numeric(frame["d.time"]).to_numpy(float) + events = pd.to_numeric(frame.death).to_numpy(float) + if ( + not np.isfinite(times).all() + or (times < 0).any() + or not np.isin(events, [0, 1]).all() + ): + raise ValueError("Invalid survival times or event indicators") + train, test = train_test_split( + np.arange(len(frame)), test_size=0.1, random_state=42 + ) + train, val = train_test_split(train, test_size=1 / 9, random_state=42) + features, categories, cardinalities = preprocess(frame, train) + boundaries = torch.linspace(0, times[train].max(), args.time_bins) + grid = torch.cat([boundaries, boundaries[-1:] + 1]).numpy() + time_bin = torch.tensor(np.searchsorted(grid, times).clip(max=len(grid) - 1)) + event_tensor = torch.tensor(events, dtype=torch.float32) + rows = [] + for seed in args.seeds: + for dropout in [0.0, 0.1, 0.3, 0.5]: + torch.manual_seed(seed) + model = PersonalizedSurvivalMoE( + None, + features.shape[1], + len(grid), + hidden_dim=args.hidden_dim, + num_experts=args.experts, + dropout=dropout, + num_hidden_layers=1, + categorical_cardinalities=cardinalities, + embedding_dim=4, + ) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, eps=1e-7) + loader = DataLoader( + train, + batch_size=args.batch_size, + shuffle=True, + generator=torch.Generator().manual_seed(42), + ) + + def predict(idx, model=model): + return model( + features=features[idx], + categories=categories[idx], + time_bin=time_bin[idx], + event=event_tensor[idx], + ) + + best, state, best_epoch, stale = float("inf"), None, 0, 0 + for epoch in range(args.epochs): + model.train() + for idx in loader: + optimizer.zero_grad() + predict(idx)["loss"].backward() + optimizer.step() + model.eval() + with torch.no_grad(): + score = float(predict(val)["nll"]) + if score < best: + best, state, best_epoch, stale = ( + score, + copy.deepcopy(model.state_dict()), + epoch + 1, + 0, + ) + else: + stale += 1 + if stale >= args.patience: + break + model.load_state_dict(state) + model.eval() + with torch.no_grad(): + out = predict(test) + train_nll = float(predict(train)["nll"]) + survival = np.column_stack([out["survival"].numpy(), np.zeros(len(test))]) + evaluator = LifelinesEvaluator( + pd.DataFrame(survival.T, index=grid), + times[test], + events[test], + times[train], + events[train], + predict_time_method="Median", + interpolation="Linear", + ) + try: + c_index = float(evaluator.concordance()[0]) + except ZeroDivisionError: + c_index = float("nan") + rows.append( + dict( + seed=seed, + dropout=dropout, + best_epoch=best_epoch, + epochs_run=epoch + 1, + train_nll=train_nll, + val_nll=best, + test_nll=float(out["nll"]), + gap=best - train_nll, + c_index=c_index, + **calibration_metrics( + times[train], + events[train], + times[test], + events[test], + survival, + grid, + ), + ) + ) + print(json.dumps(rows[-1]), flush=True) + result = { + "config": vars(args), + "sizes": list(map(len, [train, val, test])), + "time_grid": grid.tolist(), + "results": rows, + "summary": summarize(rows), + } + + def finite(value): + if isinstance(value, dict): + return {key: finite(v) for key, v in value.items()} + if isinstance(value, list): + return [finite(v) for v in value] + return None if isinstance(value, float) and not np.isfinite(value) else value + + return finite(result) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--root", help="Data directory; downloads support2.csv if absent" + ) + source.add_argument("--synthetic", action="store_true") + parser.add_argument("--epochs", type=int, default=1000) + parser.add_argument("--patience", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--time-bins", type=int, default=100) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--experts", type=int, default=8) + parser.add_argument("--lr", type=float, default=0.0005) + parser.add_argument("--seeds", type=int, nargs="+", default=[42, 43, 44, 45, 46]) + parser.add_argument("--output", default="personalized_moe_results.json") + args = parser.parse_args() + results = run(args) + Path(args.output).write_text(json.dumps(results, indent=2, allow_nan=False) + "\n") + print(json.dumps(results["summary"], indent=2, allow_nan=False)) diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index a9e6f90b2..e4212b6e7 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -57,3 +57,4 @@ from .generators.synthea import Synthea as Synthea from .generators.corgan import CorGAN as CorGAN from .survival_mdn import SurvivalMDN as SurvivalMDN +from .personalized_survival_moe import PersonalizedSurvivalMoE as PersonalizedSurvivalMoE diff --git a/pyhealth/models/personalized_survival_moe.py b/pyhealth/models/personalized_survival_moe.py new file mode 100644 index 000000000..4bf1e9597 --- /dev/null +++ b/pyhealth/models/personalized_survival_moe.py @@ -0,0 +1,220 @@ +"""Personalized discrete-time survival MoE (Morrill et al., ML4H). + +Paper: https://proceedings.mlr.press/v297/morrill26a.html +Reference: https://github.com/ToddMorrill/survival-moe. +""" + +import math +from typing import Any + +import torch +from torch import Tensor, nn +from torch.nn import functional as F + +from pyhealth.datasets import SampleDataset +from pyhealth.models.base_model import BaseModel + + +class PersonalizedSurvivalMoE(BaseModel): + """Mix patient-specific MTLR probability mass functions in log space. + + Args: + dataset: Optional PyHealth dataset with dense, preprocessed features. + input_dim: Number of continuous input features. + num_time_bins: Number of categories, including the final overflow bin. + hidden_dim: Encoder width. + num_experts: Number of personalized experts. + dropout: Dropout before each encoder linear layer. + num_hidden_layers: Number of encoder layers. + categorical_cardinalities: Vocabulary sizes for optional categorical inputs. + embedding_dim: Width of each categorical embedding. + + The encoder uses disjoint expert value subspaces. Routing + temperature is learned from 2.0; load-balancing regularization weight is 0.01. + ``hidden_dim`` must be divisible by ``num_experts``. + + Labels use zero-based categories. An event contributes -log p[category]; + censoring contributes -log sum(p[category:]), matching the reference's + inclusive-bin convention. The final category represents later/no event. + ``survival[:, j]`` is P(category > j), for finite categories only. + Fit preprocessing and time boundaries on training patients only. + + Survival is not a standard classification mode: use a custom training loop + and censoring-aware metrics, as in the accompanying SUPPORT2 example. + + Examples: + >>> model = PersonalizedSurvivalMoE(None, input_dim=3, num_time_bins=5) + >>> model(features=torch.zeros(2, 3))["survival"].shape + torch.Size([2, 4]) + """ + + def __init__( + self, + dataset: SampleDataset | None, + input_dim: int, + num_time_bins: int, + hidden_dim: int = 128, + num_experts: int = 8, + dropout: float = 0.0, + num_hidden_layers: int = 1, + categorical_cardinalities: tuple[int, ...] = (), + embedding_dim: int = 4, + ) -> None: + super().__init__(dataset) + for name, value in { + "input_dim": input_dim, + "num_time_bins": num_time_bins, + "hidden_dim": hidden_dim, + "num_experts": num_experts, + "num_hidden_layers": num_hidden_layers, + "embedding_dim": embedding_dim, + **{f"cardinality_{i}": n for i, n in enumerate(categorical_cardinalities)}, + }.items(): + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + if num_time_bins < 2: + raise ValueError("num_time_bins must include a finite and overflow bin") + if not math.isfinite(dropout) or not 0 <= dropout < 1: + raise ValueError("dropout must be in [0, 1)") + if hidden_dim % num_experts: + raise ValueError("hidden_dim must be divisible by num_experts") + self.mode = None + self.input_dim, self.num_time_bins = input_dim, num_time_bins + self.num_experts = num_experts + self.embeddings = nn.ModuleList( + [nn.Embedding(n, embedding_dim) for n in categorical_cardinalities] + ) + encoder_dim = input_dim + len(self.embeddings) * embedding_dim + layers = [] + for i in range(num_hidden_layers): + layers.extend( + [ + nn.Dropout(dropout), + nn.Linear(encoder_dim if i == 0 else hidden_dim, hidden_dim), + nn.ReLU(), + ] + ) + self.encoder = nn.Sequential(*layers) + self.query = nn.Linear(hidden_dim, hidden_dim) + self.keys = nn.Parameter(torch.empty(num_experts, hidden_dim)) + nn.init.kaiming_uniform_(self.keys, a=math.sqrt(5)) + self.value = nn.Linear(hidden_dim, hidden_dim) + width = hidden_dim // num_experts + self.experts = nn.ModuleList( + [nn.Linear(width, num_time_bins) for _ in range(num_experts)] + ) + self.raw_temperature = nn.Parameter( + torch.tensor(2.0 + math.log(-math.expm1(-2.0))) + ) + + @staticmethod + def negative_log_likelihood( + log_pmf: Tensor, time_bin: Tensor, event: Tensor + ) -> Tensor: + """Return mean event/censoring NLL with inclusive censor categories. + + Args: + log_pmf: Normalized log masses of shape (batch, categories). + time_bin: Integer categories, shape (batch,) or (batch, 1). + event: Binary observed-event flags, same number of rows. + + Returns: + Scalar mean negative log likelihood. + + Raises: + ValueError: For invalid shapes, categories or event flags. + """ + if log_pmf.ndim != 2 or not log_pmf.shape[0]: + raise ValueError("log_pmf must be a nonempty matrix") + n, bins = log_pmf.shape + time_bin = torch.as_tensor(time_bin, device=log_pmf.device) + event = torch.as_tensor(event, device=log_pmf.device) + if time_bin.shape not in [(n,), (n, 1)] or event.shape not in [(n,), (n, 1)]: + raise ValueError("labels must have shape (batch,) or (batch, 1)") + time_bin, event = time_bin.reshape(-1), event.reshape(-1) + if not torch.isfinite(time_bin).all() or not torch.all( + time_bin == time_bin.long() + ): + raise ValueError("time_bin must contain finite integers") + if torch.any((time_bin < 0) | (time_bin >= bins)): + raise ValueError("time_bin is outside the modeled categories") + if not torch.all((event == 0) | (event == 1)): + raise ValueError("event must contain only 0 or 1") + idx = time_bin.long().unsqueeze(1) + log_tail = log_pmf.flip(-1).logcumsumexp(-1).flip(-1) + log_likelihood = torch.where( + event.bool(), + log_pmf.gather(1, idx).squeeze(1), + log_tail.gather(1, idx).squeeze(1), + ) + return -log_likelihood.mean() + + def forward(self, **kwargs: Any) -> dict[str, Tensor]: + """Predict PMFs, survival and routing weights; optionally compute loss. + + Args: + **kwargs: Dense features (batch, input_dim), optional paired + time_bin/event labels, integer categories (batch, number of + embeddings), and ignored sample metadata. + + Returns: + logit (log mixture PMF), pmf, survival/y_prob, expert_pmf, + mixture_weights, and, when labeled, loss, nll and y_true. + """ + x = torch.as_tensor( + kwargs["features"], device=self.device, dtype=self.query.weight.dtype + ) + if x.ndim != 2 or x.shape[1] != self.input_dim or not x.shape[0]: + raise ValueError("features must have shape (nonempty batch, input_dim)") + if not torch.isfinite(x).all(): + raise ValueError("features must be finite; impute missing values first") + if self.embeddings: + categories = torch.as_tensor(kwargs["categories"], device=self.device) + if categories.shape != (len(x), len(self.embeddings)): + raise ValueError("categories must have one column per embedding") + if not torch.isfinite(categories).all() or not torch.all( + categories == categories.long() + ): + raise ValueError("categories must contain finite integers") + embedded = [] + for i, embedding in enumerate(self.embeddings): + col = categories[:, i].long() + if torch.any((col < 0) | (col >= embedding.num_embeddings)): + raise ValueError("category is outside its vocabulary") + embedded.append(embedding(col)) + x = torch.cat([x, *embedded], dim=-1) + h = self.encoder(x) + router_logits = self.query(h) @ self.keys.T + temp = F.softplus(self.raw_temperature).clamp_min(torch.finfo(h.dtype).eps) + log_weights = F.log_softmax(router_logits / temp, dim=-1) + values = self.value(h) + chunks = values.chunk(self.num_experts, dim=-1) + increments = torch.stack([head(v) for head, v in zip(self.experts, chunks)], 1) + scores = increments.flip(-1).cumsum(-1).flip(-1) + log_experts = F.log_softmax(scores, dim=-1) + log_pmf = torch.logsumexp(log_weights.unsqueeze(-1) + log_experts, dim=1) + survival = log_pmf.flip(-1).logcumsumexp(-1).flip(-1)[:, 1:].exp().clamp(max=1) + result = { + "logit": log_pmf, + "pmf": log_pmf.exp(), + "survival": survival, + "y_prob": survival, + "expert_pmf": log_experts.exp(), + "mixture_weights": log_weights.exp(), + } + labeled = [key in kwargs for key in ("time_bin", "event")] + if any(labeled) and not all(labeled): + raise ValueError("time_bin and event must be provided together") + if all(labeled): + time = torch.as_tensor(kwargs["time_bin"], device=self.device) + event = torch.as_tensor(kwargs["event"], device=self.device) + nll = self.negative_log_likelihood(log_pmf, time, event) + # Balance untempered routing probabilities, as in the reference. + avg = router_logits.softmax(-1).mean(0) + balance = self.num_experts * avg.square().sum() + result.update( + nll=nll, + loss=nll + 0.01 * balance, + y_true=torch.stack([time.reshape(-1), event.reshape(-1)], -1), + ) + return result diff --git a/tests/core/test_personalized_survival_moe.py b/tests/core/test_personalized_survival_moe.py new file mode 100644 index 000000000..eec5ffe6a --- /dev/null +++ b/tests/core/test_personalized_survival_moe.py @@ -0,0 +1,215 @@ +"""Tiny synthetic correctness tests; no dataset downloads or persistent files.""" + +import math +from pathlib import Path +from tempfile import TemporaryDirectory + +import pytest +import torch + +from pyhealth.models import PersonalizedSurvivalMoE + + +@pytest.fixture +def model(): + torch.manual_seed(7) + return PersonalizedSurvivalMoE(None, 3, 5, hidden_dim=8, num_experts=2) + + +def test_instantiation(model): + from pyhealth.models import BaseModel + + assert isinstance(model, BaseModel) + assert model.input_dim == 3 + assert model.num_time_bins == 5 + assert model.num_experts == 2 + assert model.mode is None + + +def test_distribution_and_gradients(model): + out = model(features=torch.randn(4, 3), time_bin=[0, 1, 3, 4], event=[1, 0, 1, 0]) + assert out["pmf"].shape == (4, 5) + assert out["logit"].shape == (4, 5) + assert out["mixture_weights"].shape == (4, 2) + assert out["loss"].ndim == 0 + assert torch.isfinite(out["loss"]) + assert out["survival"].shape == (4, 4) + assert out["expert_pmf"].shape == (4, 2, 5) + assert out["y_true"].shape == (4, 2) + for key in ["pmf", "expert_pmf", "mixture_weights"]: + assert torch.allclose( + out[key].sum(-1), torch.ones_like(out[key].sum(-1)), atol=1e-6 + ) + assert torch.all((out["survival"] >= 0) & (out["survival"] <= 1)) + assert torch.all(out["survival"][:, 1:] <= out["survival"][:, :-1]) + mixture = (out["mixture_weights"].unsqueeze(-1) * out["expert_pmf"]).sum(1) + assert torch.allclose(out["pmf"], mixture, atol=1e-6) + out["loss"].backward() + assert any( + param.grad is not None and torch.any(param.grad != 0) + for param in model.parameters() + ) + for name, param in model.named_parameters(): + if name != "_dummy_param": + assert param.grad is not None, name + assert torch.isfinite(param.grad).all(), name + + +def test_hand_computed_censoring_likelihood(): + log_p = torch.tensor([[0.2, 0.3, 0.5]] * 3, dtype=torch.double).log() + loss = PersonalizedSurvivalMoE.negative_log_likelihood(log_p, [0, 1, 2], [1, 0, 0]) + expected = -(math.log(0.2) + math.log(0.8) + math.log(0.5)) / 3 + assert loss.item() == pytest.approx(expected) + assert PersonalizedSurvivalMoE.negative_log_likelihood( + log_p, [0] * 3, [0] * 3 + ).item() == pytest.approx(0) + + +@pytest.mark.parametrize("event", [[0] * 4, [1] * 4]) +def test_all_censored_or_observed(model, event): + out = model(features=torch.zeros(4, 3), time_bin=[0, 1, 2, 4], event=event) + assert torch.isfinite(out["loss"]) + assert torch.all((out["survival"] >= 0) & (out["survival"] <= 1)) + out["loss"].backward() + + +@pytest.mark.parametrize( + "options", + [ + {"dropout": -0.1}, + {"dropout": 1}, + {"num_experts": 3}, + {"num_time_bins": 1}, + {"input_dim": 0}, + ], +) +def test_invalid_configuration(options): + config = { + "dataset": None, + "input_dim": 3, + "num_time_bins": 5, + "hidden_dim": 8, + "num_experts": 2, + } + config.update(options) + with pytest.raises(ValueError): + PersonalizedSurvivalMoE(**config) + + +@pytest.mark.parametrize( + "times,events", + [ + ([-1, 0], [1, 0]), + ([0, 5], [1, 0]), + ([0.5, 1], [1, 0]), + ([0, float("nan")], [1, 0]), + ([0, 1], [2, 0]), + ([0], [1, 0]), + ], +) +def test_invalid_labels(model, times, events): + with pytest.raises(ValueError): + model(features=torch.zeros(2, 3), time_bin=times, event=events) + + +def test_inference_validation_and_serialization(model): + model.eval() + x = torch.randn(2, 3) + out = model(features=x) + assert "loss" not in out + with pytest.raises(ValueError): + model(features=x, event=[1, 0]) + with pytest.raises(ValueError): + model(features=torch.full((2, 3), float("nan"))) + with pytest.raises(ValueError): + model(features=torch.zeros(2, 4)) + with TemporaryDirectory() as directory: + checkpoint = Path(directory) / "model.pt" + torch.save(model.state_dict(), checkpoint) + restored = PersonalizedSurvivalMoE(None, 3, 5, hidden_dim=8, num_experts=2) + restored.load_state_dict(torch.load(checkpoint, weights_only=True)) + restored.eval() + assert torch.equal(restored(features=x)["pmf"], out["pmf"]) + assert not Path(directory).exists() + + +def test_dropout_train_and_eval(): + model = PersonalizedSurvivalMoE( + None, + 3, + 5, + hidden_dim=8, + num_experts=2, + dropout=0.5, + ) + x = torch.randn(4, 3) + model.eval() + assert torch.equal(model(features=x)["pmf"], model(features=x)["pmf"]) + model.train() + assert not torch.equal(model(features=x)["pmf"], model(features=x)["pmf"]) + + +def test_extreme_scores_stable(model): + with torch.no_grad(): + for head in model.experts: + head.bias.copy_(torch.tensor([1000.0, -1000.0, 1000.0, -1000.0, 0.0])) + out = model(features=torch.zeros(2, 3), time_bin=[1, 4], event=[1, 0]) + assert torch.isfinite(out["loss"]) + assert torch.all((out["survival"] >= 0) & (out["survival"] <= 1)) + out["loss"].backward() + assert torch.isfinite(model.experts[0].bias.grad).all() + + +def test_pyhealth_sample_dataset_and_dataloader(): + from pyhealth.datasets import create_sample_dataset, get_dataloader + + dataset = create_sample_dataset( + samples=[ + { + "patient_id": str(i), + "features": [0.0, float(i), 1.0], + "time_bin": [i], + "event": [i % 2], + } + for i in range(4) + ], + input_schema={"features": "tensor"}, + output_schema={"time_bin": "tensor", "event": "tensor"}, + in_memory=True, + ) + model = PersonalizedSurvivalMoE(dataset, 3, 5, hidden_dim=8, num_experts=2) + batch = next(iter(get_dataloader(dataset, batch_size=4, shuffle=False))) + assert model.mode is None + assert torch.isfinite(model(**batch)["loss"]) + + +def test_categorical_embeddings_and_gradients(): + model = PersonalizedSurvivalMoE( + None, + 3, + 5, + hidden_dim=8, + num_experts=2, + num_hidden_layers=1, + categorical_cardinalities=(3, 2), + embedding_dim=4, + ) + features = torch.randn(4, 3) + categories = torch.tensor([[0, 0], [1, 1], [2, 0], [1, 1]]) + out = model( + features=features, + categories=categories, + time_bin=[0, 1, 3, 4], + event=[1, 0, 1, 0], + ) + assert out["pmf"].shape == (4, 5) + assert out["loss"].ndim == 0 + assert torch.isfinite(out["loss"]) + out["loss"].backward() + for embedding in model.embeddings: + assert embedding.weight.grad is not None + assert torch.isfinite(embedding.weight.grad).all() + assert embedding.weight.grad.abs().sum() > 0 + for invalid in [categories[:, :1], categories + 5, categories.float() + 0.5]: + with pytest.raises(ValueError): + model(features=features, categories=invalid)