JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor DeepHealth indices around disease expression

Browse Source
This commit is contained in:
Jiarui Li
2026-06-26 16:25:36 +08:00
parent 9ec3921bed
commit 6df0435f65
15 changed files with 4168 additions and 4280 deletions
Download Patch File Download Diff File Expand all files Collapse all files

364
burden_index.py
View File

@@ -2,12 +2,13 @@ from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Sequence
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,
@@ -16,15 +17,11 @@ from evaluate_auc_v2 import (
resolve_dist_mode_for_checkpoint,
validate_dataset_metadata,
)
from eval_data import load_sequence_eval_dataset
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
FormedBurdenMode = Literal["observed", "model_weighted"]
@dataclass(frozen=True)
class BurdenContext:
class DeepHealthContext:
model: torch.nn.Module
dataset: Any
cfg: dict[str, Any]
@@ -34,33 +31,35 @@ class BurdenContext:
@dataclass(frozen=True)
class DiseaseBurdenResult:
disease_id: int
formed: float
future: float
total: float
formed_mode: str
horizon: float
class DiseaseExpressionResult:
disease_ids: np.ndarray
expression: np.ndarray
t_query: float
@dataclass(frozen=True)
class BurdenIndexResult:
historical: np.ndarray
future: np.ndarray
total: np.ndarray
class OrganInvolvementResult:
organ_ids: list[str]
involvement: np.ndarray
disease_ids: np.ndarray
formed: np.ndarray
disease_future: np.ndarray
disease_total: np.ndarray
formed_mode: str
horizon: float
expression: np.ndarray
t_query: float
def load_burden_context(
@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,
) -> BurdenContext:
) -> DeepHealthContext:
run_path = Path(run_path)
config_path = run_path / "train_config.json"
model_ckpt_path = run_path / "best_model.pt"
@@ -73,13 +72,13 @@ def load_burden_context(
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
if model_target_mode == "next_token":
raise RuntimeError(
"Burden Index computation requires an all_future checkpoint because "
"it uses p_d(h, Delta). The provided run is model_target_mode='next_token'."
"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 for burden "
f"computation, got {model_target_mode!r}."
"train_config.json model_target_mode must be all_future, got "
f"{model_target_mode!r}."
)
device_obj = torch.device(
@@ -88,20 +87,17 @@ def load_burden_context(
if device_obj.type == "cuda" and not torch.cuda.is_available():
raise RuntimeError(f"Requested device {device_obj}, but CUDA is not available.")
data_prefix = cfg.get("data_prefix", "ukb")
labels_file = cfg.get("labels_file", "labels.csv")
extra_info_types = cfg.get("extra_info_types", None)
dataset = load_sequence_eval_dataset(
model_target_mode="all_future",
data_prefix=data_prefix,
labels_file=labels_file,
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=extra_info_types,
extra_info_types=cfg.get("extra_info_types", None),
)
validate_dataset_metadata(dataset, cfg)
@@ -115,12 +111,11 @@ def load_burden_context(
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
args = _ConfigNamespace()
model = build_model_from_dataset(args, cfg_model, dataset)
model = build_model_from_dataset(_ConfigNamespace(), cfg_model, dataset)
load_model_state(model, state_dict)
model.eval().to(device_obj)
return BurdenContext(
return DeepHealthContext(
model=model,
dataset=dataset,
cfg=cfg,
@@ -131,56 +126,9 @@ def load_burden_context(
@torch.inference_mode()
def compute_disease_burden(
def compute_disease_expression(
*,
run_path: str | Path,
disease_id: int,
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,
horizon: float,
formed_mode: FormedBurdenMode,
device: str | torch.device | None = None,
context: BurdenContext | None = None,
) -> DiseaseBurdenResult:
ctx = context or load_burden_context(run_path, device=device)
disease_ids = np.asarray([int(disease_id)], dtype=np.int64)
formed, future_prob = _compute_disease_components(
ctx=ctx,
disease_ids=disease_ids,
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),
horizon=float(horizon),
formed_mode=formed_mode,
)
future = (1.0 - formed) * future_prob
total = formed + future
return DiseaseBurdenResult(
disease_id=int(disease_id),
formed=float(formed[0]),
future=float(future[0]),
total=float(total[0]),
formed_mode=str(formed_mode),
horizon=float(horizon),
)
@torch.inference_mode()
def compute_burden_index(
*,
run_path: str | Path,
burden_matrix: np.ndarray,
disease_ids: Sequence[int] | np.ndarray,
event_seq: Sequence[int] | np.ndarray,
time_seq: Sequence[float] | np.ndarray,
@@ -190,26 +138,12 @@ def compute_burden_index(
other_value_kind: Sequence[int] | np.ndarray,
other_time: Sequence[float] | np.ndarray,
t_query: float,
horizon: float,
formed_mode: FormedBurdenMode,
device: str | torch.device | None = None,
context: BurdenContext | None = None,
) -> BurdenIndexResult:
ctx = context or load_burden_context(run_path, device=device)
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)
if disease_ids_arr.ndim != 1 or disease_ids_arr.size == 0:
raise ValueError("disease_ids must be a non-empty 1D sequence.")
A = np.asarray(burden_matrix, dtype=np.float64)
if A.ndim != 2:
raise ValueError(f"burden_matrix must be 2D, got shape {A.shape}")
if A.shape[1] != disease_ids_arr.size:
raise ValueError(
"burden_matrix columns must match disease_ids length, got "
f"{A.shape[1]} columns vs {disease_ids_arr.size} disease ids."
)
formed, future_prob = _compute_disease_components(
expression = model_implied_disease_expression(
ctx=ctx,
disease_ids=disease_ids_arr,
event_seq=event_seq,
@@ -220,36 +154,48 @@ def compute_burden_index(
other_value_kind=other_value_kind,
other_time=other_time,
t_query=float(t_query),
horizon=float(horizon),
formed_mode=formed_mode,
)
disease_future = (1.0 - formed) * future_prob
disease_total = formed + disease_future
historical = A @ formed
future = A @ disease_future
total = A @ disease_total
return BurdenIndexResult(
historical=historical,
future=future,
total=total,
return DiseaseExpressionResult(
disease_ids=disease_ids_arr.copy(),
formed=formed,
disease_future=disease_future,
disease_total=disease_total,
formed_mode=str(formed_mode),
horizon=float(horizon),
expression=expression,
t_query=float(t_query),
)
class _ConfigNamespace:
def __getattr__(self, _name: str) -> None:
return None
def _compute_disease_components(
def compute_organ_involvement_from_expression(
*,
ctx: BurdenContext,
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,
@@ -259,41 +205,30 @@ def _compute_disease_components(
other_value_kind: Sequence[int] | np.ndarray,
other_time: Sequence[float] | np.ndarray,
t_query: float,
horizon: float,
formed_mode: FormedBurdenMode,
) -> tuple[np.ndarray, np.ndarray]:
formed_mode = _validate_formed_mode(formed_mode)
) -> 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)
)
if horizon < 0:
raise ValueError(f"horizon must be non-negative, got {horizon}")
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)
if formed_mode == "observed":
formed = _observed_formed_burden(
disease_ids=disease_ids,
event_seq=event_seq_arr,
time_seq=time_seq_arr,
t_query=t_query,
)
else:
formed = _model_weighted_formed_burden(
ctx=ctx,
disease_ids=disease_ids,
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,
t_query=t_query,
)
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 = _query_hidden(
hidden = query_hidden(
ctx=ctx,
event_seq=event_seq_arr,
time_seq=time_seq_arr,
@@ -302,84 +237,19 @@ def _compute_disease_components(
other_value=other_value_arr,
other_value_kind=other_value_kind_arr,
other_time=other_time_arr,
query_times=np.asarray([t_query], dtype=np.float32),
query_times=grid[valid].astype(np.float32),
)
future_prob = _probabilities_from_hidden(
ctx=ctx,
hidden=hidden_query,
disease_ids=disease_ids,
deltas=np.asarray([horizon], dtype=np.float32),
)[0]
return formed.astype(np.float64), future_prob.astype(np.float64)
def _model_weighted_formed_burden(
*,
ctx: BurdenContext,
disease_ids: np.ndarray,
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,
t_query: float,
) -> np.ndarray:
grid = _build_readout_grid(
event_seq=event_seq,
time_seq=time_seq,
other_type=other_type,
other_time=other_time,
t_query=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)
if np.all(deltas <= 0):
return np.zeros(disease_ids.size, dtype=np.float64)
hidden = _query_hidden(
ctx=ctx,
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,
query_times=grid.astype(np.float32),
)
interval_prob = _probabilities_from_hidden(
interval_prob = probabilities_from_hidden(
ctx=ctx,
hidden=hidden,
disease_ids=disease_ids,
deltas=deltas,
).astype(np.float64)
survival = np.prod(1.0 - np.clip(interval_prob, 0.0, 1.0), axis=0)
return 1.0 - survival
def _observed_formed_burden(
*,
disease_ids: np.ndarray,
event_seq: np.ndarray,
time_seq: np.ndarray,
t_query: float,
) -> np.ndarray:
valid = (
(time_seq <= np.float32(t_query))
& (event_seq > PAD_IDX)
& (event_seq != CHECKUP_IDX)
& (event_seq != NO_EVENT_IDX)
deltas=deltas[valid],
)
observed = set(int(x) for x in event_seq[valid].tolist())
return np.asarray([1.0 if int(d) in observed else 0.0 for d in disease_ids])
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(
def build_readout_grid(
*,
event_seq: np.ndarray,
time_seq: np.ndarray,
@@ -400,9 +270,10 @@ def _build_readout_grid(
return np.unique(times)
def _query_hidden(
@torch.inference_mode()
def query_hidden(
*,
ctx: BurdenContext,
ctx: DeepHealthContext,
event_seq: np.ndarray,
time_seq: np.ndarray,
sex: int,
@@ -430,31 +301,24 @@ def _query_hidden(
tq = torch.from_numpy(query_times.astype(np.float32, copy=False)).float()
event = event.to(ctx.device)
times = times.to(ctx.device)
other_t = other_t.to(ctx.device)
other_v = other_v.to(ctx.device)
other_k = other_k.to(ctx.device)
other_tm = other_tm.to(ctx.device)
sex_t = sex_t.to(ctx.device)
tq = tq.to(ctx.device)
return ctx.model(
event_seq=event,
time_seq=times,
sex=sex_t,
time_seq=times.to(ctx.device),
sex=sex_t.to(ctx.device),
padding_mask=event > PAD_IDX,
t_query=tq,
other_type=other_t,
other_value=other_v,
other_value_kind=other_k,
other_time=other_tm,
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",
)
def _probabilities_from_hidden(
@torch.inference_mode()
def probabilities_from_hidden(
*,
ctx: BurdenContext,
ctx: DeepHealthContext,
hidden: torch.Tensor,
disease_ids: np.ndarray,
deltas: np.ndarray,
@@ -489,16 +353,12 @@ def _probabilities_from_hidden(
return prob.detach().cpu().numpy().astype(np.float64, copy=False)
def _validate_formed_mode(mode: str) -> FormedBurdenMode:
if mode not in {"observed", "model_weighted"}:
raise ValueError(
"formed_mode must be either 'observed' or 'model_weighted', "
f"got {mode!r}."
)
return mode # type: ignore[return-value]
class _ConfigNamespace:
def __getattr__(self, _name: str) -> None:
return None
def _validate_disease_ids(ctx: BurdenContext, disease_ids: np.ndarray) -> 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))