Jiarui Li
9ca8909e3a
- Introduced `prepare_data.R` for merging disease and other data from CSV files. - Added `prepare_data.py` for processing UK Biobank data, including: - Mapping field IDs to human-readable names. - Handling date variables and converting them to offsets. - Processing disease events and constructing tabular features. - Splitting data into training, validation, and test sets. - Saving processed data to binary and CSV formats.
212 lines
9.6 KiB
Python
212 lines
9.6 KiB
Python
import pandas as pd # Pandas for data manipulation
|
|
import tqdm # Progress bar for chunk processing
|
|
import numpy as np # Numerical operations
|
|
|
|
train_frac = 0.7 # Fraction of participants for training split
|
|
val_frac = 0.15 # Fraction of participants for validation split
|
|
test_frac = 0.15 # Fraction of participants for test split
|
|
|
|
# CSV mapping field IDs to human-readable names
|
|
field_map_file = "field_ids_enriched.csv"
|
|
field_dict = {} # Map original field ID -> new column name
|
|
tabular_fields = [] # List of tabular feature column names
|
|
with open(field_map_file, "r", encoding="utf-8") as f: # Open the field mapping file
|
|
next(f) # skip header line
|
|
for line in f: # Iterate through lines
|
|
parts = line.strip().split(",") # Split by CSV commas
|
|
if len(parts) >= 3: # Ensure we have at least id and name columns (fix: was >=2)
|
|
# Original field identifier (e.g., "34-0.0")
|
|
field_id = parts[0]
|
|
field_name = parts[2] # Human-readable column name
|
|
field_dict[field_id] = field_name # Record the mapping
|
|
# Track as a potential tabular feature
|
|
tabular_fields.append(field_name)
|
|
# Exclude raw date parts and target columns
|
|
exclude_fields = ['year', 'month', 'Death', 'age_at_assessment']
|
|
tabular_fields = [
|
|
# Filter out excluded columns
|
|
field for field in tabular_fields if field not in exclude_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, "r", 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, "r", 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
|
|
# Map code to index (0 for padding, 1 for checkup)
|
|
label_dict[parts[0]] = idx + 2
|
|
|
|
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)
|
|
|
|
# Construct date of birth from year and month (day fixed to 1)
|
|
ukb_chunk['dob'] = pd.to_datetime(
|
|
# Guard against malformed dates
|
|
ukb_chunk[['year', 'month']].assign(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 datetime and compute day offsets from dob
|
|
if present_date_vars:
|
|
date_cols = ukb_chunk[present_date_vars].apply(
|
|
pd.to_datetime, format="%Y-%m-%d", errors='coerce' # Parse dates safely
|
|
)
|
|
date_cols_days = date_cols.sub(
|
|
ukb_chunk['dob'], axis=0) # Timedelta relative to dob
|
|
ukb_chunk[present_date_vars] = date_cols_days.apply(
|
|
lambda x: x.dt.days) # Store days since dob
|
|
|
|
ukb_chunk = ukb_chunk.convert_dtypes()
|
|
|
|
# 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())
|
|
|
|
# Process disease events from ICD10-related date columns
|
|
# Take ICD date cols plus 'Death' if present by order
|
|
icd10_cols = present_date_vars[:len_icd + 1]
|
|
# Melt to long form: participant id, event code (column name), and days offset
|
|
melted_df = ukb_chunk.reset_index().melt(
|
|
id_vars=['eid'],
|
|
value_vars=icd10_cols,
|
|
var_name='event_code',
|
|
value_name='days',
|
|
)
|
|
# Require non-missing day offsets
|
|
melted_df.dropna(subset=['days'], inplace=True)
|
|
if not melted_df.empty:
|
|
melted_df['label'] = melted_df['event_code'].map(
|
|
label_dict) # Map event code to numeric label
|
|
# Fix: ensure labels exist before int cast
|
|
melted_df.dropna(subset=['label'], inplace=True)
|
|
if not melted_df.empty:
|
|
event_list.append(
|
|
melted_df[['eid', 'days', 'label']]
|
|
.astype(int) # Safe now since label and days are non-null
|
|
.to_numpy()
|
|
)
|
|
|
|
# Add assesment date as a "checkup" event (label=1)
|
|
if 'date_of_assessment' in ukb_chunk.columns:
|
|
assessment_array = (
|
|
ukb_chunk.reset_index()[['eid', 'date_of_assessment']]
|
|
.dropna()
|
|
.assign(label=1) # Checkup label
|
|
.astype(int)
|
|
.to_numpy()
|
|
)
|
|
if assessment_array.size > 0:
|
|
event_list.append(assessment_array) # Append checkup events
|
|
|
|
df_res = ukb_chunk.reset_index() # Bring participant ID out of index
|
|
# Simplify stub names for wide_to_long
|
|
# Rename date stubs
|
|
rename_dict = {f'cancer_date_{j}': f'cancerdate{j}' for j in range(17)}
|
|
rename_dict.update(
|
|
# Rename type stubs
|
|
{f'cancer_type_{j}': f'cancertype{j}' for j in range(17)})
|
|
df_renamed = df_res.rename(columns=rename_dict) # Apply renaming
|
|
stubs_to_use = [] # Collect available stubs
|
|
if any('cancerdate' in col for col in df_renamed.columns):
|
|
stubs_to_use.append('cancerdate') # Date stub present
|
|
if any('cancertype' in col for col in df_renamed.columns):
|
|
stubs_to_use.append('cancertype') # Type stub present
|
|
|
|
if len(stubs_to_use) == 2: # Only proceed if both date and type columns exist
|
|
long_cancer = pd.wide_to_long(
|
|
df_renamed,
|
|
stubnames=stubs_to_use,
|
|
i=['eid'], # Participant ID identifier
|
|
j='cancer_num' # Index over cancer record number (0..16)
|
|
).dropna() # Remove rows missing either date or type
|
|
if not long_cancer.empty:
|
|
long_cancer['cancer'] = long_cancer['cancertype'].str.slice(
|
|
0, 3) # Use first 3 chars as code
|
|
long_cancer['cancer_label'] = long_cancer['cancer'].map(
|
|
label_dict) # Map to label id
|
|
cancer_array = (
|
|
long_cancer.reset_index(
|
|
)[['eid', 'cancerdate', 'cancer_label']]
|
|
.dropna()
|
|
.astype(int)
|
|
.to_numpy()
|
|
)
|
|
if cancer_array.size > 0:
|
|
event_list.append(cancer_array) # Append cancer events
|
|
|
|
# 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.
|
|
data = pd.DataFrame(data).drop_duplicates([0, 2]).values
|
|
|
|
# Store compactly using unsigned 32-bit integers
|
|
data = data.astype(np.uint32)
|
|
|
|
# Split data into train/val/test sets by participant ID
|
|
unique_ids = np.unique(data[:, 0]) # Unique participant IDs
|
|
# ID cutoff for train
|
|
train_split_id = unique_ids[int(len(unique_ids) * train_frac)]
|
|
# ID cutoff for val
|
|
val_split_id = unique_ids[int(len(unique_ids) * (train_frac + val_frac))]
|
|
|
|
train_data = data[data[:, 0] <= train_split_id].tofile("ukb_train.bin")
|
|
val_data = data[(data[:, 0] > train_split_id) & (
|
|
data[:, 0] <= val_split_id)].tofile("ukb_val.bin")
|
|
test_data = data[data[:, 0] > val_split_id].tofile("ukb_test.bin")
|
|
train_tabular = final_tabular[final_tabular.index <= train_split_id]
|
|
val_tabular = final_tabular[(final_tabular.index > train_split_id) & (
|
|
final_tabular.index <= val_split_id)]
|
|
test_tabular = final_tabular[final_tabular.index > val_split_id]
|
|
train_tabular.to_csv("ukb_train_tabular.csv")
|
|
val_tabular.to_csv("ukb_val_tabular.csv")
|
|
test_tabular.to_csv("ukb_test_tabular.csv")
|