"""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``. 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 and sex. 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() ) 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 + ["date_of_birth"] ) # 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", ) ukb_chunk["date_of_birth"] = dob.dt.strftime("%Y-%m-%d") # 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 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)