"""Create a compact calendar-dated disease-event index from ``ukb_data.csv``. Unlike ``prepare_data.py``, this ETL does not create model-ready relative-time sequences or other-information tokens. It writes one structured ``.npy`` file with exactly three fields: eid int64 event_date datetime64[D] token int32 ``token`` follows the existing ``labels.csv`` convention used by ``prepare_data.py``: padding=0, checkup=1 (not emitted here), and the first label in ``labels.csv`` receives token 2. Each ``(eid, token)`` is deduplicated to the first known event date. The output is intended for calendar-indexed temperature and air-pollution queries. It contains no date of birth, sex, covariates, or checkup events. Usage ----- python prepare_event_dates.py --output ukb_disease_event_dates.npy """ from __future__ import annotations import argparse from pathlib import Path import numpy as np import pandas as pd try: from tqdm import tqdm except ImportError: # Keep the ETL runnable in minimal Python environments. tqdm = None EVENT_DTYPE = np.dtype( [ ("eid", " dict[str, int]: """Return the label-code -> token mapping shared with ``prepare_data.py``.""" token_map: dict[str, int] = {} with Path(labels_file).open(encoding="utf-8") as handle: for index, line in enumerate(handle): code = line.strip().split(" ", maxsplit=1)[0] if code: token_map[code] = index + 2 return token_map def build_raw_column_map( field_map_file: str | Path, icd_map_file: str | Path, ) -> tuple[dict[str, str], list[str]]: """Build raw-column renames and the calendar-date event columns to inspect.""" field_df = pd.read_csv(field_map_file, low_memory=False) required = {"field_instance", "var_name"} missing = required - set(field_df.columns) if missing: raise ValueError(f"{field_map_file} is missing columns: {sorted(missing)}") raw_to_name = dict( zip(field_df["field_instance"].astype(str), field_df["var_name"].astype(str)) ) icd_date_columns: list[str] = [] with Path(icd_map_file).open(encoding="utf-8") as handle: for line in handle: parts = line.strip().split() if len(parts) >= 6: raw_to_name[parts[0]] = parts[5] icd_date_columns.append(parts[5]) for slot in range(17): raw_to_name[f"40005-{slot}.0"] = f"cancer_date_{slot}" raw_to_name[f"40006-{slot}.0"] = f"cancer_type_{slot}" return raw_to_name, icd_date_columns def _records_from_icd_columns( chunk: pd.DataFrame, event_columns: list[str], token_map: dict[str, int], ) -> list[pd.DataFrame]: frames: list[pd.DataFrame] = [] for column in event_columns: token = token_map.get(column) if token is None or column not in chunk.columns: continue event_date = pd.to_datetime(chunk[column], format="%Y-%m-%d", errors="coerce") valid = event_date.notna() if valid.any(): frames.append( pd.DataFrame( { "eid": chunk.index[valid].astype("int64"), "event_date": event_date.loc[valid].to_numpy(), "token": np.full(valid.sum(), token, dtype=np.int32), } ) ) return frames def _records_from_cancer_columns(chunk: pd.DataFrame, token_map: dict[str, int]) -> list[pd.DataFrame]: frames: list[pd.DataFrame] = [] for slot in range(17): date_column = f"cancer_date_{slot}" type_column = f"cancer_type_{slot}" if date_column not in chunk.columns or type_column not in chunk.columns: continue event_date = pd.to_datetime(chunk[date_column], format="%Y-%m-%d", errors="coerce") code = chunk[type_column].astype("string").str.slice(0, 3) token = code.map(token_map) valid = event_date.notna() & token.notna() if valid.any(): frames.append( pd.DataFrame( { "eid": chunk.index[valid].astype("int64"), "event_date": event_date.loc[valid].to_numpy(), "token": token.loc[valid].astype("int32").to_numpy(), } ) ) return frames def prepare_event_dates( *, ukb_data_file: str | Path, field_map_file: str | Path, icd_map_file: str | Path, labels_file: str | Path, output_file: str | Path, chunksize: int = 10_000, ) -> int: """Stream the raw UKB export, then write a sorted structured event array.""" token_map = load_label_tokens(labels_file) raw_to_name, icd_date_columns = build_raw_column_map(field_map_file, icd_map_file) event_columns = [*icd_date_columns, "Death"] frames: list[pd.DataFrame] = [] reader = pd.read_csv( ukb_data_file, chunksize=chunksize, index_col=0, # UKB participant ID / eid low_memory=False, ) chunk_iterator = tqdm(reader, desc="Extracting calendar-dated disease events") if tqdm else reader for raw_chunk in chunk_iterator: chunk = raw_chunk.rename(columns=raw_to_name) frames.extend(_records_from_icd_columns(chunk, event_columns, token_map)) frames.extend(_records_from_cancer_columns(chunk, token_map)) if not frames: result = np.empty(0, dtype=EVENT_DTYPE) else: events = pd.concat(frames, ignore_index=True) events = events.dropna(subset=["eid", "event_date", "token"]) events = events.sort_values(["eid", "token", "event_date"], kind="stable") # Match prepare_data.py: first occurrence of each disease/death token. events = events.drop_duplicates(["eid", "token"], keep="first") events = events.sort_values(["eid", "event_date", "token"], kind="stable") result = np.empty(len(events), dtype=EVENT_DTYPE) result["eid"] = events["eid"].to_numpy(dtype=np.int64) result["event_date"] = events["event_date"].to_numpy(dtype="datetime64[D]") result["token"] = events["token"].to_numpy(dtype=np.int32) np.save(output_file, result) return len(result) def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--ukb-data", default="ukb_data.csv") parser.add_argument("--field-map", default="field_ids_enriched.csv") parser.add_argument("--icd-map", default="icd10_codes_mod.tsv") parser.add_argument("--labels", default="labels.csv") parser.add_argument("--output", default="ukb_disease_event_dates.npy") parser.add_argument("--chunksize", type=int, default=10_000) args = parser.parse_args() if args.chunksize <= 0: raise ValueError("chunksize must be positive") count = prepare_event_dates( ukb_data_file=args.ukb_data, field_map_file=args.field_map, icd_map_file=args.icd_map, labels_file=args.labels, output_file=args.output, chunksize=args.chunksize, ) print(f"Wrote {count:,} first disease/death events to {args.output}") if __name__ == "__main__": main()