386 lines
14 KiB
Python
386 lines
14 KiB
Python
|
"""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)
|