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

Use precomputed exposure embeddings

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 16:49:49 +08:00
parent 552e09aa01
commit e439c3c98c
7 changed files with 311 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
dataset.py
View File

@@ -50,7 +50,11 @@ def _load_readonly_npy(path: Path) -> np.ndarray:
class ExposureCache:
"""Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""
def __init__(self, cache_dir: str | Path):
def __init__(
self,
cache_dir: str | Path,
embeddings_file: str | Path | None = None,
):
cache_dir = Path(cache_dir)
self.cache_dir = cache_dir
manifest_path = cache_dir / "exposure_manifest.json"
@@ -101,6 +105,11 @@ class ExposureCache:
self.eid_start = _load_readonly_npy(eid_start_path)
self.daily = _load_readonly_npy(daily_path)
self.monthly = _load_readonly_npy(monthly_path)
self.embeddings = (
_load_readonly_npy(Path(embeddings_file))
if embeddings_file is not None
else None
)
quality_path = cache_dir / "exposure_quality.npy"
self.quality = _load_readonly_npy(quality_path) if quality_path.is_file() else None
@@ -132,6 +141,16 @@ class ExposureCache:
raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1")
if len(self.eid_start) and int(self.eid_start[-1]) != n_rows:
raise ValueError("Last exposure eid offset must equal exposure row count")
if self.embeddings is not None:
if self.embeddings.ndim != 2:
raise ValueError(
"Exposure embeddings must have shape (N, D), got "
f"{self.embeddings.shape}"
)
if max_window_index >= len(self.embeddings):
raise ValueError(
"Exposure row index points past exposure embeddings"
)
self._eid_to_pos = {
int(eid): idx
@@ -203,6 +222,20 @@ class ExposureCache:
def monthly_windows(self, indices: np.ndarray) -> np.ndarray:
return self._windows("monthly", indices)
def embedding_windows(self, indices: np.ndarray) -> np.ndarray:
if self.embeddings is None:
raise RuntimeError("Exposure embeddings are not enabled")
indices = np.asarray(indices, dtype=np.int64)
out = np.zeros(
(len(indices), self.embeddings.shape[1]), dtype=np.float32
)
valid_pos = np.nonzero(indices >= 0)[0]
if len(valid_pos):
out[valid_pos] = np.asarray(
self.embeddings[indices[valid_pos]], dtype=np.float32
)
return out
def _windows(
self,
kind: Literal["daily", "monthly"],
@@ -288,18 +321,21 @@ class _ExpoBaseDataset(Dataset):
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
exposure_embeddings_file: str | Path | None = None,
) -> None:
self.data_prefix = data_prefix
self.labels_file = labels_file
self.no_event_interval_years = float(no_event_interval_years)
self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
self.exposure_cache = (
ExposureCache(exposure_cache_dir)
ExposureCache(exposure_cache_dir, exposure_embeddings_file)
if exposure_cache_dir is not None
else None
)
self.mask_onset_exposure = bool(mask_onset_exposure)
if exposure_embeddings_file is not None and exposure_cache_dir is None:
raise ValueError(
"exposure_cache_dir is required with exposure_embeddings_file"
)
self.label_code_to_id, self.label_id_to_code = load_label_vocab(
labels_file,
@@ -427,24 +463,14 @@ class _ExpoBaseDataset(Dataset):
age_days=input_times_days,
)
def _load_exposure_windows(self, exposure_index: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
def _load_exposure_embeddings(
self, exposure_index: np.ndarray
) -> torch.Tensor:
if self.exposure_cache is None:
raise RuntimeError("Exposure cache is not enabled")
daily = self.exposure_cache.daily_windows(exposure_index).astype(
np.float32,
copy=False,
)
monthly = self.exposure_cache.monthly_windows(exposure_index).astype(
np.float32,
copy=False,
)
if self.mask_onset_exposure:
daily[:, 0, :] = np.nan
monthly[:, 0, :] = np.nan
return torch.from_numpy(daily).float(), torch.from_numpy(monthly).float()
return torch.from_numpy(
self.exposure_cache.embedding_windows(exposure_index)
).float()
class NextStepHealthDataset(_ExpoBaseDataset):
"""
@@ -465,7 +491,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
exposure_embeddings_file: str | Path | None = None,
) -> None:
super().__init__(
data_prefix=data_prefix,
@@ -473,7 +499,7 @@ class NextStepHealthDataset(_ExpoBaseDataset):
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
exposure_embeddings_file=exposure_embeddings_file,
)
self.samples: List[Dict] = []
@@ -531,9 +557,9 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"target_multi_hot": torch.from_numpy(s["target_multi_hot"]).bool(),
}
if "exposure_index" in s:
daily, monthly = self._load_exposure_windows(s["exposure_index"])
out["exposure_daily"] = daily
out["exposure_monthly"] = monthly
out["exposure_embedding"] = self._load_exposure_embeddings(
s["exposure_index"]
)
return out
@@ -561,7 +587,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
min_future_events: int = 1,
validation_query_seed: int = 42,
exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False,
exposure_embeddings_file: str | Path | None = None,
) -> None:
if split not in {"train", "valid", "test"}:
raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -572,7 +598,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure,
exposure_embeddings_file=exposure_embeddings_file,
)
self.split = split
@@ -719,9 +745,9 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
input_times_days=times_days[hist],
)
if exposure_index is not None:
daily, monthly = self._load_exposure_windows(exposure_index)
out["exposure_daily"] = daily
out["exposure_monthly"] = monthly
out["exposure_embedding"] = self._load_exposure_embeddings(
exposure_index
)
return out
def __len__(self) -> int:
@@ -745,19 +771,20 @@ def _collate_common_static(batch: List[Dict]) -> Dict:
}
def _pad_exposure(batch: List[Dict], key: str, shape: tuple[int, int]) -> torch.Tensor:
def _pad_exposure_embedding(batch: List[Dict]) -> torch.Tensor:
max_len = max(int(s["event_seq"].numel()) for s in batch)
out = torch.full(
(len(batch), max_len, shape[0], shape[1]),
float("nan"),
dtype=torch.float32,
embedding_dim = next(
int(s["exposure_embedding"].size(1))
for s in batch
if "exposure_embedding" in s
)
out = torch.zeros(
len(batch), max_len, embedding_dim, dtype=torch.float32
)
for idx, sample in enumerate(batch):
value = sample.get(key)
if value is None:
continue
seq_len = int(value.size(0))
out[idx, :seq_len] = value
value = sample.get("exposure_embedding")
if value is not None:
out[idx, :value.size(0)] = value
return out
@@ -809,9 +836,8 @@ def next_step_collate_fn(batch: List[Dict]) -> Dict:
"target_multi_hot": target_multi_hot,
}
out.update(_collate_common_static(batch))
if any("exposure_daily" in s for s in batch):
out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
if any("exposure_embedding" in s for s in batch):
out["exposure_embedding"] = _pad_exposure_embedding(batch)
return out
@@ -847,9 +873,8 @@ def all_future_collate_fn(batch: List[Dict]) -> Dict:
"exposure": torch.stack([s["exposure"] for s in batch]),
}
out.update(_collate_common_static(batch))
if any("exposure_daily" in s for s in batch):
out["exposure_daily"] = _pad_exposure(batch, "exposure_daily", DAILY_EXPOSURE_SHAPE)
out["exposure_monthly"] = _pad_exposure(batch, "exposure_monthly", MONTHLY_EXPOSURE_SHAPE)
if any("exposure_embedding" in s for s in batch):
out["exposure_embedding"] = _pad_exposure_embedding(batch)
return out