Clean basic info preparation

prepare_data.py

@@ -5,7 +5,8 @@ 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_basic_info.csv``: basic patient table indexed by ``eid`` with ``sex``
+  and ``date_of_birth``.
 
 Processing steps
 ----------------
@@ -14,7 +15,8 @@ Processing steps
 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.
+5. Write event data plus basic patient information needed for exposure-date
+   alignment.
 
 Usage
 -----
@@ -69,22 +71,9 @@ field_dict = dict(zip(field_df["field_instance"], field_df["var_name"]))
 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"]
-)
+# 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"
@@ -155,8 +144,6 @@ for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
         ),
         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]
 
@@ -167,10 +154,16 @@ for ukb_chunk in tqdm.tqdm(ukb_iterator, desc="Processing UK Biobank data"):
                 ukb_chunk[col], format="%Y-%m-%d", errors="coerce") - dob
         ).dt.days
 
-    # Append tabular features (use only columns that exist)
+    # Append compact basic info without mutating the wide raw chunk.
     present_tabular_fields = [
-        c for c in tabular_fields if c in ukb_chunk.columns]
-    tabular_list.append(ukb_chunk[present_tabular_fields].copy())
+        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]