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``.
* ``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)