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

Remove extra info token pathway

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 16:57:49 +08:00
parent 6dfeb5a696
commit a0379daf29
13 changed files with 18 additions and 1390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
dataset.py
View File

@@ -1,7 +1,7 @@
# dataset.py
from __future__ import annotations
from typing import Dict, Iterable, List, Literal, Optional, Tuple
from typing import Dict, List, Literal, Optional, Tuple
import numpy as np
import pandas as pd
@@ -87,17 +87,11 @@ class _ExpoBaseDataset(Dataset):
labels_file: str = "labels.csv",
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
extra_info_types: Iterable[int] | 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.requested_extra_info_types = (
None
if extra_info_types is None
else list(dict.fromkeys(int(t) for t in extra_info_types))
)
self.label_code_to_id, self.label_id_to_code = load_label_vocab(
labels_file,
@@ -112,26 +106,12 @@ class _ExpoBaseDataset(Dataset):
self.event_data = event_data[order]
basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
other_info = np.load(f"{data_prefix}_other_info.npy")
if other_info.ndim != 2 or other_info.shape[1] != 5:
raise ValueError(
f"other_info must have shape (N, 5), got {other_info.shape}"
)
cate_types = pd.read_csv("cate_types.csv")
required_cate_cols = {"type", "name", "n_categories"}
missing_cate_cols = required_cate_cols - set(cate_types.columns)
if missing_cate_cols:
raise ValueError(
f"cate_types.csv is missing columns: {sorted(missing_cate_cols)}"
)
basic_table.index = basic_table.index.astype(np.int64)
unique_eids = np.unique(self.event_data[:, 0].astype(np.int64))
basic_table = basic_table.loc[unique_eids]
self._prepare_sex(basic_table, unique_eids)
self._prepare_other_info(other_info, cate_types, unique_eids)
max_id_in_vocab = max(self.label_id_to_code.keys())
max_id_in_data = int(self.event_data[:, 2].max()) if len(self.event_data) > 0 else 0
@@ -160,95 +140,6 @@ class _ExpoBaseDataset(Dataset):
)
self.sex_mapping = {int(eid): int(s) for eid, s in zip(unique_eids, sex01)}
def _prepare_other_info(
self,
other_info: np.ndarray,
cate_types: pd.DataFrame,
unique_eids: np.ndarray,
) -> None:
other_info = other_info.copy()
other_info[:, 0] = other_info[:, 0].astype(np.int64)
other_info[:, 1] = other_info[:, 1].astype(np.int64)
other_info[:, 3] = other_info[:, 3].astype(np.int64)
available_types = sorted(
int(t) for t in np.unique(other_info[:, 1]) if int(t) > 0
)
if self.requested_extra_info_types is None:
selected_types = available_types
else:
selected_types = self.requested_extra_info_types
missing = sorted(set(selected_types) - set(available_types))
if missing:
raise ValueError(f"Requested extra_info_types not found: {missing}")
keep = np.isin(other_info[:, 0].astype(np.int64), unique_eids)
keep &= np.isin(other_info[:, 1].astype(np.int64), selected_types)
other_info = other_info[keep]
cate_counts = {
int(row["type"]): int(row["n_categories"])
for _, row in cate_types.iterrows()
}
cate_offsets: Dict[int, int] = {}
next_offset = 0
for type_id in selected_types:
if type_id in cate_counts:
cate_offsets[type_id] = next_offset
next_offset += cate_counts[type_id]
kinds = other_info[:, 3].astype(np.int64)
types = other_info[:, 1].astype(np.int64)
cate_rows = kinds == 2
for type_id in np.unique(types[cate_rows]):
type_id = int(type_id)
if type_id not in cate_offsets:
raise ValueError(
f"type {type_id} appears categorical but is missing from cate_types.csv"
)
row_mask = cate_rows & (types == type_id)
local_value = other_info[row_mask, 2].astype(np.int64)
other_info[row_mask, 2] = local_value + cate_offsets[type_id]
cont_type_ids = [
int(t)
for t in selected_types
if np.any((types == int(t)) & (kinds == 1))
]
self.extra_info_types = selected_types
self.cate_type_offsets = cate_offsets
self.n_types = (max(selected_types) + 1) if selected_types else 1
self.cont_type_ids = cont_type_ids
self.n_cont_types = len(cont_type_ids)
self.n_categories = next_offset + 1
order = np.lexsort((other_info[:, 4], other_info[:, 1], other_info[:, 0]))
other_info = other_info[order]
self.other_info_by_eid: Dict[int, Dict[str, np.ndarray]] = {}
for eid in unique_eids.astype(np.int64):
self.other_info_by_eid[int(eid)] = {
"other_type": np.zeros(0, dtype=np.int64),
"other_value": np.zeros(0, dtype=np.float32),
"other_value_kind": np.zeros(0, dtype=np.int64),
"other_time": np.zeros(0, dtype=np.float32),
}
if len(other_info) == 0:
return
eids, starts = np.unique(other_info[:, 0].astype(np.int64), return_index=True)
ends = np.concatenate([starts[1:], [len(other_info)]])
for eid_raw, start, end in zip(eids, starts, ends):
rows = other_info[start:end]
self.other_info_by_eid[int(eid_raw)] = {
"other_type": rows[:, 1].astype(np.int64),
"other_value": rows[:, 2].astype(np.float32),
"other_value_kind": rows[:, 3].astype(np.int64),
"other_time": (rows[:, 4].astype(np.float32) / DAYS_PER_YEAR),
}
def _iter_patient_events(
self,
*,
@@ -279,12 +170,10 @@ class _ExpoBaseDataset(Dataset):
yield eid, times_days, labels
def _split_features(self, eid: int) -> Optional[Dict]:
other_info = self.other_info_by_eid.get(eid)
if other_info is None:
if eid not in self.sex_mapping:
return None
return {
"sex": self.sex_mapping[eid],
**other_info,
}
class NextStepHealthDataset(_ExpoBaseDataset):
@@ -305,14 +194,12 @@ class NextStepHealthDataset(_ExpoBaseDataset):
labels_file: str = "labels.csv",
no_event_interval_years: float = 5.0,
include_no_event_in_uts_target: bool = False,
extra_info_types: Iterable[int] | None = None,
) -> None:
super().__init__(
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
extra_info_types=extra_info_types,
)
self.samples: List[Dict] = []
@@ -355,10 +242,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
"event_seq": torch.from_numpy(s["event_seq"]).long(),
"time_seq": torch.from_numpy(s["time_seq"]).float(),
"sex": torch.tensor(s["sex"], dtype=torch.long),
"other_type": torch.from_numpy(s["other_type"]).long(),
"other_value": torch.from_numpy(s["other_value"]).float(),
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
"other_time": torch.from_numpy(s["other_time"]).float(),
"target_event_seq": torch.from_numpy(s["target_event_seq"]).long(),
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
@@ -390,7 +273,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
min_history_events: int = 1,
min_future_events: int = 1,
validation_query_seed: int = 42,
extra_info_types: Iterable[int] | None = None,
) -> None:
if split not in {"train", "valid", "test"}:
raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -400,7 +282,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
extra_info_types=extra_info_types,
)
self.split = split
@@ -537,10 +418,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
"future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(),
"exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32),
"sex": torch.tensor(patient["sex"], dtype=torch.long),
"other_type": torch.from_numpy(patient["other_type"]).long(),
"other_value": torch.from_numpy(patient["other_value"]).float(),
"other_value_kind": torch.from_numpy(patient["other_value_kind"]).long(),
"other_time": torch.from_numpy(patient["other_time"]).float(),
}
def __len__(self) -> int:
@@ -561,26 +438,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
def _collate_common_static(batch: List[Dict]) -> Dict:
return {
"sex": torch.stack([s["sex"] for s in batch]),
"other_type": pad_sequence(
[s["other_type"] for s in batch],
batch_first=True,
padding_value=0,
),
"other_value": pad_sequence(
[s["other_value"] for s in batch],
batch_first=True,
padding_value=0.0,
),
"other_value_kind": pad_sequence(
[s["other_value_kind"] for s in batch],
batch_first=True,
padding_value=0,
),
"other_time": pad_sequence(
[s["other_time"] for s in batch],
batch_first=True,
padding_value=0.0,
),
}