DeepHealthExpo/prepare_data.py

"""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``
  and ``date_of_birth``.

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. Write event data plus basic patient information needed for exposure-date
   alignment.

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


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()
)
# Keep only patient-level fields that are still consumed downstream. Exposure
# windows are prepared separately by prepare_exposure_cache.py.
basic_info_fields = [f for f in ["sex"] if f in set(basic_info_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 compact basic info without mutating the wide raw chunk.
    present_tabular_fields = [
        c for c in basic_info_fields if c in ukb_chunk.columns]
    dob_frame = pd.DataFrame(
        {"date_of_birth": dob.dt.strftime("%Y-%m-%d").to_numpy()},
        index=ukb_chunk.index,
    )
    tabular_list.append(
        pd.concat([ukb_chunk[present_tabular_fields].copy(), dob_frame], axis=1)
    )

    # 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 information needed by the model and exposure-date alignment.
basic_cols = ["sex"]
if "date_of_birth" in final_tabular.columns:
    basic_cols.append("date_of_birth")
basic_info = final_tabular[basic_cols].copy()
basic_info.to_csv("ukb_basic_info.csv")

# Save event data
np.save("ukb_event_data.npy", data)