408 lines
14 KiB
Python
408 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from eval_data import load_sequence_eval_dataset
|
|
from evaluate_auc_v2 import (
|
|
build_model_from_dataset,
|
|
load_checkpoint_state_dict,
|
|
load_json_config,
|
|
load_model_state,
|
|
resolve_dist_mode_for_checkpoint,
|
|
validate_dataset_metadata,
|
|
)
|
|
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DeepHealthContext:
|
|
model: torch.nn.Module
|
|
dataset: Any
|
|
cfg: dict[str, Any]
|
|
dist_mode: str
|
|
device: torch.device
|
|
run_path: Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiseaseExpressionResult:
|
|
disease_ids: np.ndarray
|
|
expression: np.ndarray
|
|
t_query: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OrganInvolvementResult:
|
|
organ_ids: list[str]
|
|
involvement: np.ndarray
|
|
disease_ids: np.ndarray
|
|
expression: np.ndarray
|
|
t_query: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FrailtyRiskResult:
|
|
frailty_risk_index: float
|
|
disease_ids: np.ndarray
|
|
expression: np.ndarray
|
|
weights: np.ndarray
|
|
t_query: float
|
|
|
|
|
|
def load_deephealth_context(
|
|
run_path: str | Path,
|
|
*,
|
|
device: str | torch.device | None = None,
|
|
) -> DeepHealthContext:
|
|
run_path = Path(run_path)
|
|
config_path = run_path / "train_config.json"
|
|
model_ckpt_path = run_path / "best_model.pt"
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"train_config.json not found in {run_path}")
|
|
if not model_ckpt_path.exists():
|
|
raise FileNotFoundError(f"best_model.pt not found in {run_path}")
|
|
|
|
cfg = load_json_config(config_path)
|
|
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
|
if model_target_mode == "next_token":
|
|
raise RuntimeError(
|
|
"Disease expression indices require an all_future checkpoint because "
|
|
"they use p_d(h, Delta). The provided run is model_target_mode='next_token'."
|
|
)
|
|
if model_target_mode != "all_future":
|
|
raise ValueError(
|
|
"train_config.json model_target_mode must be all_future, got "
|
|
f"{model_target_mode!r}."
|
|
)
|
|
|
|
device_obj = torch.device(
|
|
device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
|
|
)
|
|
if device_obj.type == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError(f"Requested device {device_obj}, but CUDA is not available.")
|
|
|
|
dataset = load_sequence_eval_dataset(
|
|
model_target_mode="all_future",
|
|
data_prefix=cfg.get("data_prefix", "ukb"),
|
|
labels_file=cfg.get("labels_file", "labels.csv"),
|
|
no_event_interval_years=float(cfg.get("no_event_interval_years", 5.0)),
|
|
include_no_event_in_uts_target=bool(
|
|
cfg.get("include_no_event_in_uts_target", False)
|
|
),
|
|
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
|
|
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
|
|
extra_info_types=cfg.get("extra_info_types", None),
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
|
|
state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
|
|
dist_mode = resolve_dist_mode_for_checkpoint(
|
|
str(cfg.get("dist_mode", "exponential")),
|
|
state_dict,
|
|
)
|
|
if dist_mode not in {"exponential", "weibull", "mixed"}:
|
|
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
|
|
|
|
cfg_model = dict(cfg)
|
|
cfg_model["dist_mode"] = dist_mode
|
|
model = build_model_from_dataset(_ConfigNamespace(), cfg_model, dataset)
|
|
load_model_state(model, state_dict)
|
|
model.eval().to(device_obj)
|
|
|
|
return DeepHealthContext(
|
|
model=model,
|
|
dataset=dataset,
|
|
cfg=cfg,
|
|
dist_mode=dist_mode,
|
|
device=device_obj,
|
|
run_path=run_path,
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def compute_disease_expression(
|
|
*,
|
|
run_path: str | Path,
|
|
disease_ids: Sequence[int] | np.ndarray,
|
|
event_seq: Sequence[int] | np.ndarray,
|
|
time_seq: Sequence[float] | np.ndarray,
|
|
sex: int,
|
|
other_type: Sequence[int] | np.ndarray,
|
|
other_value: Sequence[float] | np.ndarray,
|
|
other_value_kind: Sequence[int] | np.ndarray,
|
|
other_time: Sequence[float] | np.ndarray,
|
|
t_query: float,
|
|
device: str | torch.device | None = None,
|
|
context: DeepHealthContext | None = None,
|
|
) -> DiseaseExpressionResult:
|
|
ctx = context or load_deephealth_context(run_path, device=device)
|
|
disease_ids_arr = np.asarray(disease_ids, dtype=np.int64)
|
|
expression = model_implied_disease_expression(
|
|
ctx=ctx,
|
|
disease_ids=disease_ids_arr,
|
|
event_seq=event_seq,
|
|
time_seq=time_seq,
|
|
sex=sex,
|
|
other_type=other_type,
|
|
other_value=other_value,
|
|
other_value_kind=other_value_kind,
|
|
other_time=other_time,
|
|
t_query=float(t_query),
|
|
)
|
|
return DiseaseExpressionResult(
|
|
disease_ids=disease_ids_arr.copy(),
|
|
expression=expression,
|
|
t_query=float(t_query),
|
|
)
|
|
|
|
|
|
def compute_organ_involvement_from_expression(
|
|
*,
|
|
expression: Sequence[float] | np.ndarray,
|
|
organ_matrix: np.ndarray,
|
|
) -> np.ndarray:
|
|
z = np.asarray(expression, dtype=np.float64)
|
|
A = np.asarray(organ_matrix, dtype=np.float64)
|
|
if A.ndim != 2:
|
|
raise ValueError(f"organ_matrix must be 2D, got shape {A.shape}")
|
|
if z.ndim != 1 or A.shape[1] != z.size:
|
|
raise ValueError(
|
|
"expression must be 1D and match organ_matrix columns, got "
|
|
f"{z.shape} and {A.shape}"
|
|
)
|
|
intensity = -np.log1p(-np.clip(z, 0.0, 1.0 - 1e-7))
|
|
return -np.expm1(-(A @ intensity))
|
|
|
|
|
|
def compute_frailty_risk_from_expression(
|
|
*,
|
|
expression: Sequence[float] | np.ndarray,
|
|
hfrs_weights: Sequence[float] | np.ndarray,
|
|
) -> float:
|
|
z = np.asarray(expression, dtype=np.float64)
|
|
w = np.asarray(hfrs_weights, dtype=np.float64)
|
|
if z.shape != w.shape:
|
|
raise ValueError(f"expression and hfrs_weights shape mismatch: {z.shape} vs {w.shape}")
|
|
return float(np.dot(w, z))
|
|
|
|
|
|
@torch.inference_mode()
|
|
def model_implied_disease_expression(
|
|
*,
|
|
ctx: DeepHealthContext,
|
|
disease_ids: np.ndarray,
|
|
event_seq: Sequence[int] | np.ndarray,
|
|
time_seq: Sequence[float] | np.ndarray,
|
|
sex: int,
|
|
other_type: Sequence[int] | np.ndarray,
|
|
other_value: Sequence[float] | np.ndarray,
|
|
other_value_kind: Sequence[int] | np.ndarray,
|
|
other_time: Sequence[float] | np.ndarray,
|
|
t_query: float,
|
|
) -> np.ndarray:
|
|
disease_ids = np.asarray(disease_ids, dtype=np.int64)
|
|
_validate_disease_ids(ctx, disease_ids)
|
|
event_seq_arr, time_seq_arr = _validate_event_inputs(event_seq, time_seq)
|
|
other_type_arr, other_value_arr, other_value_kind_arr, other_time_arr = (
|
|
_validate_other_inputs(other_type, other_value, other_value_kind, other_time)
|
|
)
|
|
grid = build_readout_grid(
|
|
event_seq=event_seq_arr,
|
|
time_seq=time_seq_arr,
|
|
other_type=other_type_arr,
|
|
other_time=other_time_arr,
|
|
t_query=float(t_query),
|
|
)
|
|
if grid.size == 0:
|
|
return np.zeros(disease_ids.size, dtype=np.float64)
|
|
|
|
end_times = np.concatenate([grid[1:], np.asarray([t_query], dtype=np.float32)])
|
|
deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
|
|
valid = deltas > 0
|
|
if not np.any(valid):
|
|
return np.zeros(disease_ids.size, dtype=np.float64)
|
|
|
|
hidden = query_hidden(
|
|
ctx=ctx,
|
|
event_seq=event_seq_arr,
|
|
time_seq=time_seq_arr,
|
|
sex=sex,
|
|
other_type=other_type_arr,
|
|
other_value=other_value_arr,
|
|
other_value_kind=other_value_kind_arr,
|
|
other_time=other_time_arr,
|
|
query_times=grid[valid].astype(np.float32),
|
|
)
|
|
interval_prob = probabilities_from_hidden(
|
|
ctx=ctx,
|
|
hidden=hidden,
|
|
disease_ids=disease_ids,
|
|
deltas=deltas[valid],
|
|
)
|
|
survival = np.prod(1.0 - np.clip(interval_prob, 0.0, 1.0), axis=0)
|
|
return (1.0 - survival).astype(np.float64, copy=False)
|
|
|
|
|
|
def build_readout_grid(
|
|
*,
|
|
event_seq: np.ndarray,
|
|
time_seq: np.ndarray,
|
|
other_type: np.ndarray,
|
|
other_time: np.ndarray,
|
|
t_query: float,
|
|
) -> np.ndarray:
|
|
event_mask = (event_seq > PAD_IDX) & (time_seq <= np.float32(t_query))
|
|
other_mask = (other_type > 0) & (other_time <= np.float32(t_query))
|
|
times = np.concatenate(
|
|
[
|
|
time_seq[event_mask].astype(np.float32, copy=False),
|
|
other_time[other_mask].astype(np.float32, copy=False),
|
|
]
|
|
)
|
|
if times.size == 0:
|
|
return np.zeros(0, dtype=np.float32)
|
|
return np.unique(times)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def query_hidden(
|
|
*,
|
|
ctx: DeepHealthContext,
|
|
event_seq: np.ndarray,
|
|
time_seq: np.ndarray,
|
|
sex: int,
|
|
other_type: np.ndarray,
|
|
other_value: np.ndarray,
|
|
other_value_kind: np.ndarray,
|
|
other_time: np.ndarray,
|
|
query_times: np.ndarray,
|
|
) -> torch.Tensor:
|
|
if query_times.ndim != 1:
|
|
raise ValueError("query_times must be 1D.")
|
|
batch_size = int(query_times.size)
|
|
if batch_size == 0:
|
|
return torch.empty(0, ctx.model.n_embd, device=ctx.device)
|
|
|
|
event = torch.from_numpy(event_seq[None, :].repeat(batch_size, axis=0)).long()
|
|
times = torch.from_numpy(time_seq[None, :].repeat(batch_size, axis=0)).float()
|
|
other_t = torch.from_numpy(other_type[None, :].repeat(batch_size, axis=0)).long()
|
|
other_v = torch.from_numpy(other_value[None, :].repeat(batch_size, axis=0)).float()
|
|
other_k = torch.from_numpy(
|
|
other_value_kind[None, :].repeat(batch_size, axis=0)
|
|
).long()
|
|
other_tm = torch.from_numpy(other_time[None, :].repeat(batch_size, axis=0)).float()
|
|
sex_t = torch.full((batch_size,), int(sex), dtype=torch.long)
|
|
tq = torch.from_numpy(query_times.astype(np.float32, copy=False)).float()
|
|
|
|
event = event.to(ctx.device)
|
|
return ctx.model(
|
|
event_seq=event,
|
|
time_seq=times.to(ctx.device),
|
|
sex=sex_t.to(ctx.device),
|
|
padding_mask=event > PAD_IDX,
|
|
t_query=tq.to(ctx.device),
|
|
other_type=other_t.to(ctx.device),
|
|
other_value=other_v.to(ctx.device),
|
|
other_value_kind=other_k.to(ctx.device),
|
|
other_time=other_tm.to(ctx.device),
|
|
target_mode="all_future",
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def probabilities_from_hidden(
|
|
*,
|
|
ctx: DeepHealthContext,
|
|
hidden: torch.Tensor,
|
|
disease_ids: np.ndarray,
|
|
deltas: np.ndarray,
|
|
) -> np.ndarray:
|
|
if hidden.ndim != 2:
|
|
raise ValueError(f"hidden must have shape (N, H), got {tuple(hidden.shape)}")
|
|
if deltas.ndim != 1 or deltas.size != hidden.shape[0]:
|
|
raise ValueError(
|
|
"deltas must be 1D with the same length as hidden rows, got "
|
|
f"{deltas.shape} vs {tuple(hidden.shape)}"
|
|
)
|
|
ids = torch.as_tensor(disease_ids, dtype=torch.long, device=ctx.device)
|
|
logits = ctx.model.calc_risk(hidden)[:, ids]
|
|
rate = F.softplus(logits).clamp_min(1e-8)
|
|
delta_t = torch.as_tensor(deltas, dtype=rate.dtype, device=ctx.device).clamp_min(0)
|
|
|
|
if ctx.dist_mode == "weibull":
|
|
rho = ctx.model.calc_weibull_rho(hidden)[:, ids]
|
|
exposure = torch.pow(delta_t[:, None], rho)
|
|
elif ctx.dist_mode == "mixed":
|
|
exposure = delta_t[:, None].expand_as(rate)
|
|
death_idx = int(getattr(ctx.model, "death_idx", getattr(ctx.model, "vocab_size", 0) - 1))
|
|
death_cols = [j for j, token in enumerate(disease_ids.tolist()) if int(token) == death_idx]
|
|
if death_cols:
|
|
death_rho = ctx.model.calc_death_rho(hidden)
|
|
for col in death_cols:
|
|
exposure[:, int(col)] = torch.pow(delta_t, death_rho)
|
|
else:
|
|
exposure = delta_t[:, None].expand_as(rate)
|
|
|
|
prob = -torch.expm1(-rate * exposure)
|
|
return prob.detach().cpu().numpy().astype(np.float64, copy=False)
|
|
|
|
|
|
class _ConfigNamespace:
|
|
def __getattr__(self, _name: str) -> None:
|
|
return None
|
|
|
|
|
|
def _validate_disease_ids(ctx: DeepHealthContext, disease_ids: np.ndarray) -> None:
|
|
if disease_ids.ndim != 1 or disease_ids.size == 0:
|
|
raise ValueError("disease_ids must be a non-empty 1D array.")
|
|
vocab_size = int(getattr(ctx.model, "vocab_size", ctx.model.risk_head.out_features))
|
|
if np.any(disease_ids < 0) or np.any(disease_ids >= vocab_size):
|
|
raise ValueError(f"disease_ids must be in [0, {vocab_size}), got {disease_ids}")
|
|
|
|
|
|
def _validate_event_inputs(
|
|
event_seq: Sequence[int] | np.ndarray,
|
|
time_seq: Sequence[float] | np.ndarray,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
events = np.asarray(event_seq, dtype=np.int64)
|
|
times = np.asarray(time_seq, dtype=np.float32)
|
|
if events.ndim != 1 or times.ndim != 1:
|
|
raise ValueError("event_seq and time_seq must be 1D.")
|
|
if events.shape != times.shape:
|
|
raise ValueError(
|
|
f"event_seq and time_seq must have the same shape, got {events.shape} vs {times.shape}"
|
|
)
|
|
if events.size == 0:
|
|
raise ValueError("event_seq must contain at least one token.")
|
|
if not np.all(np.isfinite(times)):
|
|
raise ValueError("time_seq contains non-finite values.")
|
|
return events, times
|
|
|
|
|
|
def _validate_other_inputs(
|
|
other_type: Sequence[int] | np.ndarray,
|
|
other_value: Sequence[float] | np.ndarray,
|
|
other_value_kind: Sequence[int] | np.ndarray,
|
|
other_time: Sequence[float] | np.ndarray,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
|
typ = np.asarray(other_type, dtype=np.int64)
|
|
val = np.asarray(other_value, dtype=np.float32)
|
|
kind = np.asarray(other_value_kind, dtype=np.int64)
|
|
tm = np.asarray(other_time, dtype=np.float32)
|
|
if not (typ.shape == val.shape == kind.shape == tm.shape):
|
|
raise ValueError(
|
|
"other_type, other_value, other_value_kind, and other_time must "
|
|
f"have the same shape, got {typ.shape}, {val.shape}, {kind.shape}, {tm.shape}."
|
|
)
|
|
if typ.ndim != 1:
|
|
raise ValueError("other_* inputs must be 1D.")
|
|
if not np.all(np.isfinite(tm)):
|
|
raise ValueError("other_time contains non-finite values.")
|
|
return typ, val, kind, tm
|