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

Add target construction and training script for DeepHealth model

Browse Source 
- Implemented target construction in `targets.py` for next-token and unique-time set supervision.
- Added validation functions and utility methods for target building.
- Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging.
- Integrated loss functions and readout mechanisms based on target modes.
- Established dataset splitting and DataLoader configurations for training, validation, and testing.
This commit is contained in:
Jiarui Li
2026-06-12 10:28:16 +08:00
commit 5e979e061b
13 changed files with 9165 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

385
prepare_data.py Normal file
View File

@@ -0,0 +1,385 @@
"""ETL pipeline for UK Biobank data preparation.
This script converts raw UK Biobank CSV exports into the artefacts consumed by
DeepHealth:
* ``ukb_event_data.npy``: ``(N, 3)`` uint32 array of ``(eid, days, label)``
disease/death/checkup events sorted by patient then time.
* ``ukb_basic_info.csv``: basic patient table indexed by ``eid`` with ``sex``.
* ``ukb_other_info.npy``: ``(M, 5)`` float64 array of
``(eid, type, value, value_kind, time)`` rows. ``type=0`` is reserved for
padding, ``value_kind=1`` means continuous, and ``value_kind=2`` means
categorical.
* ``cate_types.csv``: categorical-variable metadata with
``type,name,n_categories``. Dataset code computes global category ids after
experiment-specific variable selection.
Processing steps
----------------
1. Stream the raw CSV in 10 000-row chunks to bound memory usage.
2. Convert date columns to integer day offsets from date of birth.
3. Extract ICD-10 date fields and cancer date/type fields into long-form
events and map codes to integer labels via ``labels.csv``.
4. De-duplicate events per ``(eid, label)`` keeping the earliest occurrence.
5. Convert available non-sex tabular fields into unified other-information
tokens timestamped by ``date_of_assessment``.
6. Write event data, sex, unified other-information tokens, and categorical
type metadata.
Usage
-----
::
python prepare_data.py
The script expects these files in the working directory:
* ``ukb_data.csv``: raw UK Biobank CSV export.
* ``field_ids_enriched.csv``: field-ID to column-name mapping.
* ``icd10_codes_mod.tsv``: ICD-10 field to date-column mapping.
* ``labels.csv``: disease-code to integer-label mapping.
"""
import numpy as np # Numerical operations
import pandas as pd # Pandas for data manipulation
import tqdm # Progress bar for chunk processing
CONT_VALUE_KIND = 1
CATE_VALUE_KIND = 2
def _sort_values(values):
"""Sort mixed pandas/numpy scalar values deterministically."""
try:
return sorted(values)
except TypeError:
return sorted(values, key=lambda x: str(x))
def _build_other_info_tokens(
table: pd.DataFrame,
feature_fields: list[str],
*,
time_col: str = "date_of_assessment",
) -> tuple[np.ndarray, pd.DataFrame]:
"""Convert wide tabular features into (eid, type, value, kind, time) rows."""
rows = []
cate_meta = []
present_features = [
col for col in feature_fields
if col in table.columns and col not in {time_col, "sex"}
]
if time_col not in table.columns:
raise ValueError(
f"{time_col!r} is required to timestamp other-information tokens"
)
token_times = pd.to_numeric(table[time_col], errors="coerce")
for type_idx, col in enumerate(present_features, start=1):
series = table[col]
n_unique = series.dropna().nunique()
valid_time = token_times.notna()
if n_unique > 10:
numeric = pd.to_numeric(series, errors="coerce")
valid = numeric.notna() & valid_time
if not valid.any():
continue
rows.append(
np.column_stack(
(
table.index[valid].to_numpy(dtype=np.float64),
np.full(valid.sum(), type_idx, dtype=np.float64),
numeric[valid].to_numpy(dtype=np.float64),
np.full(valid.sum(), CONT_VALUE_KIND, dtype=np.float64),
token_times[valid].to_numpy(dtype=np.float64),
)
)
)
else:
unique_vals = _sort_values(series.dropna().unique())
value_map = {val: idx + 1 for idx, val in enumerate(unique_vals)}
mapped = series.map(value_map)
valid = mapped.notna() & valid_time
n_categories = len(unique_vals)
cate_meta.append(
{
"type": type_idx,
"name": col,
"n_categories": n_categories,
}
)
if not valid.any():
continue
rows.append(
np.column_stack(
(
table.index[valid].to_numpy(dtype=np.float64),
np.full(valid.sum(), type_idx, dtype=np.float64),
mapped[valid].to_numpy(dtype=np.float64),
np.full(valid.sum(), CATE_VALUE_KIND, dtype=np.float64),
token_times[valid].to_numpy(dtype=np.float64),
)
)
)
cate_types = pd.DataFrame(
cate_meta,
columns=["type", "name", "n_categories"],
)
if not rows:
return np.empty((0, 5), dtype=np.float64), cate_types
out = np.vstack(rows)
return out[np.lexsort((out[:, 3], out[:, 1], out[:, 0]))], cate_types
def _unique_preserve_order(values):
"""Return unique values while preserving first-seen order."""
seen = set()
out = []
for value in values:
if value not in seen:
seen.add(value)
out.append(value)
return out
# CSV mapping field IDs to human-readable names
field_map_file = "field_ids_enriched.csv"
field_df = pd.read_csv(field_map_file, low_memory=False)
required_cols = {"field_instance", "var_name", "field_type"}
missing_cols = required_cols - set(field_df.columns)
if missing_cols:
raise ValueError(
f"{field_map_file} is missing required columns: {sorted(missing_cols)}"
)
field_df = field_df.dropna(subset=["field_instance", "var_name", "field_type"])
field_df["field_instance"] = field_df["field_instance"].astype(str)
field_df["var_name"] = field_df["var_name"].astype(str)
field_df["field_type"] = field_df["field_type"].astype(int)
# Map original field ID -> renamed output variable.
field_dict = dict(zip(field_df["field_instance"], field_df["var_name"]))
# Build source field groups from field_type.
basic_info_fields = _unique_preserve_order(
field_df.loc[field_df["field_type"] == 0, "var_name"].tolist()
)
assessment_fields = _unique_preserve_order(
field_df.loc[field_df["field_type"] == 1, "var_name"].tolist()
)
exposure_fields = _unique_preserve_order(
field_df.loc[field_df["field_type"] == 2, "var_name"].tolist()
)
# Keep only sex and enrollment time for the basic info table.
basic_info_fields = [
f for f in ["sex", "date_of_assessment"] if f in set(basic_info_fields)
]
# Fields needed for tabular extraction from raw CSV.
tabular_fields = _unique_preserve_order(
basic_info_fields + assessment_fields + exposure_fields
)
# TSV mapping field IDs to ICD10-related date columns
field_to_icd_map = "icd10_codes_mod.tsv"
# Date-like variables to be converted to offsets
date_vars = []
with open(field_to_icd_map, encoding="utf-8") as f: # Open ICD10 mapping
for line in f: # Iterate each mapping row
parts = line.strip().split() # Split on whitespace for TSV
if len(parts) >= 6: # Guard against malformed lines
# Map field ID to the date column name
field_dict[parts[0]] = parts[5]
date_vars.append(parts[5]) # Track date column names in order
for j in range(17): # Map up to 17 cancer entry slots (dates and types)
# Cancer diagnosis date slot j
field_dict[f"40005-{j}.0"] = f"cancer_date_{j}"
field_dict[f"40006-{j}.0"] = f"cancer_type_{j}" # Cancer type/code slot j
# Number of ICD-related date columns before adding extras
len_icd = len(date_vars)
date_vars.extend(
["Death", "date_of_assessment"] # Add outcome date and assessment date
+
# Add cancer date columns
[f"cancer_date_{j}" for j in range(17)]
)
labels_file = "labels.csv" # File listing label codes
label_dict = {} # Map code string -> integer label id
with open(labels_file, encoding="utf-8") as f: # Open labels file
for idx, line in enumerate(f): # Enumerate to assign incremental label IDs
parts = line.strip().split(" ") # Split by space
if parts and parts[0]: # Guard against empty lines
# Start labels from 1 to reserve 0 for padding, 1 for checkup
label_dict[parts[0]] = idx + 2
# Pre-build lookup: ICD/Death column name -> integer label for fast per-column extraction
icd_label_lookup = {}
for col_name in date_vars[:len_icd]:
if col_name in label_dict:
icd_label_lookup[col_name] = label_dict[col_name]
if "Death" in label_dict:
icd_label_lookup["Death"] = label_dict["Death"]
event_list = [] # Accumulator for event arrays across chunks
tabular_list = [] # Accumulator for tabular feature DataFrames across chunks
ukb_iterator = pd.read_csv( # Stream UK Biobank data in chunks
"ukb_data.csv",
sep=",",
chunksize=10000, # Stream file in manageable chunks to reduce memory footprint
# First column (participant ID) becomes DataFrame index
index_col=0,
low_memory=False, # Disable type inference optimization for consistent dtypes
)
# Iterate chunks with progress
for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
# Rename columns to friendly names
ukb_chunk = ukb_chunk.rename(columns=field_dict)
# Require sex to be present
ukb_chunk.dropna(subset=["sex"], inplace=True)
if ukb_chunk.empty:
continue
# Construct date of birth from year and month (day fixed to 1)
dob = pd.to_datetime(
pd.DataFrame(
{"year": ukb_chunk["year"], "month": ukb_chunk["month"], "day": 1}
),
errors="coerce",
)
# Use only date variables that actually exist in the current chunk
present_date_vars = [c for c in date_vars if c in ukb_chunk.columns]
# Convert date-like columns to day offsets from dob (per-column, avoids .apply overhead)
for col in present_date_vars:
ukb_chunk[col] = (
pd.to_datetime(
ukb_chunk[col], format="%Y-%m-%d", errors="coerce") - dob
).dt.days
# Append tabular features (use only columns that exist)
present_tabular_fields = [
c for c in tabular_fields if c in ukb_chunk.columns]
tabular_list.append(ukb_chunk[present_tabular_fields].copy())
# Extract ICD10 + Death events directly per column (avoids costly melt)
icd10_cols = present_date_vars[: len_icd + 1]
for col in icd10_cols:
if col not in icd_label_lookup:
continue
label_id = icd_label_lookup[col]
series = ukb_chunk[col]
valid_mask = series.notna()
if not valid_mask.any():
continue
event_list.append(
np.column_stack(
(
ukb_chunk.index[valid_mask].values,
series[valid_mask].values.astype(np.int64),
np.full(valid_mask.sum(), label_id, dtype=np.int64),
)
)
)
# Optimized cancer processing without wide_to_long
cancer_frames = []
for j in range(17):
d_col = f"cancer_date_{j}"
t_col = f"cancer_type_{j}"
if d_col in ukb_chunk.columns and t_col in ukb_chunk.columns:
# Filter rows where both date and type are present
mask = ukb_chunk[d_col].notna() & ukb_chunk[t_col].notna()
if mask.any():
subset_idx = ukb_chunk.index[mask]
subset_days = ukb_chunk.loc[mask, d_col]
subset_type = ukb_chunk.loc[mask, t_col]
# Map cancer type to label
# Use first 3 chars
cancer_codes = subset_type.str.slice(0, 3)
labels = cancer_codes.map(label_dict)
# Filter valid labels
valid_label_mask = labels.notna()
if valid_label_mask.any():
# Create array: eid, days, label
# Ensure types are correct for numpy
c_eids = subset_idx[valid_label_mask].values
c_days = subset_days[valid_label_mask].values
c_labels = labels[valid_label_mask].values
# Stack
chunk_cancer_data = np.column_stack(
(c_eids, c_days, c_labels))
cancer_frames.append(chunk_cancer_data)
if cancer_frames:
event_list.append(np.vstack(cancer_frames))
# Add checkup events with label=1 using date_of_assessment (already in days from dob)
if "date_of_assessment" in ukb_chunk.columns:
doa_series = ukb_chunk["date_of_assessment"].dropna()
if not doa_series.empty:
checkup_data = np.column_stack(
(
doa_series.index.values,
doa_series.values.astype(int),
np.ones(len(doa_series), dtype=int),
)
)
event_list.append(checkup_data)
# Combine tabular chunks
final_tabular = pd.concat(tabular_list, axis=0, ignore_index=False)
final_tabular.index.name = "eid" # Ensure index named consistently
data = np.vstack(event_list) # Stack all event arrays into one
# Sort by participant then day
data = data[np.lexsort((data[:, 1], data[:, 0]))]
# Keep only events with non-negative day offsets
data = data[data[:, 1] >= 0]
# Remove duplicate (participant_id, label) pairs keeping first occurrence (numpy-based).
_, first_idx = np.unique(data[:, [0, 2]], axis=0, return_index=True)
first_idx.sort() # Preserve original sorted order
data = data[first_idx]
# Store compactly using unsigned 32-bit integers
data = data.astype(np.uint32)
# Select eid in both data and tabular
valid_eids = np.intersect1d(data[:, 0], final_tabular.index)
data = data[np.isin(data[:, 0], valid_eids)]
final_tabular = final_tabular.loc[valid_eids]
final_tabular = final_tabular.convert_dtypes()
# Save basic sex information separately.
basic_info = final_tabular[["sex"]].copy()
basic_info.to_csv("ukb_basic_info.csv")
# Save unified other-information tokens. Missing values simply produce no token.
other_info_fields = _unique_preserve_order(
basic_info_fields + assessment_fields + exposure_fields
)
other_info, cate_types = _build_other_info_tokens(
final_tabular,
other_info_fields,
time_col="date_of_assessment",
)
np.save("ukb_other_info.npy", other_info)
cate_types.to_csv("cate_types.csv", index=False)
# Save event data
np.save("ukb_event_data.npy", data)