JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use eid-aligned exposure cache

Browse Source
This commit is contained in:
Jiarui Li
2026-07-08 12:17:30 +08:00
parent 642e8aad67
commit 46d530fad0
3 changed files with 275 additions and 630 deletions
Download Patch File Download Diff File Expand all files Collapse all files

582
prepare_exposure_cache.py
View File

@@ -1,65 +1,49 @@
"""Build a random-access exposure index/cache from disease-level parquet files.
"""Build an eid-sequence-aligned exposure cache for DeepHealth training.
The README-described exposure dataset is stored as one daily and one monthly
parquet file per disease. That layout is good for disease-specific analysis but
too expensive for mini-batch training, where we need exposure windows aligned
to arbitrary event sequences.
The source exposure dataset is stored as one daily and one monthly parquet file
per disease. That layout is inconvenient for mini-batch training because the
model consumes per-participant disease sequences. This script materializes one
large numpy cache ordered exactly like ``{data_prefix}_event_data.npy`` after
sorting by ``eid, age_days, token``.
By default this script builds a lightweight parquet index. It does not copy the
daily/monthly exposure windows; it only records which source parquet file,
row-group, and row each exposure event lives in. Dataset loading then reads the
original parquet row groups on demand.
The output directory contains:
The default index directory contains:
exposure_eid.npy int64 eid per real disease event
exposure_token.npy int32 raw disease token per real disease event
exposure_age_days.npy int32 age in days per real disease event
exposure_onset_date.npy datetime64[D] onset date per real disease event
exposure_eid_index.npy int64 unique eids in cache order
exposure_eid_start.npy int64 start offsets, length len(eid_index) + 1
exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2)
channels: tmean, rhmean
exposure_quality.npy float32 memmap, shape (N, 4)
n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata
exposure_eid.npy int64 eid per exposure row
exposure_token.npy int32 raw disease token per exposure row
exposure_onset_date.npy datetime64[D] onset date per exposure row
exposure_daily_file_id.npy int32 source daily file id per row
exposure_daily_row_group.npy int32 source daily row group per row
exposure_daily_row_in_group.npy int32 row offset inside daily row group
exposure_monthly_file_id.npy int32 source monthly file id per row
exposure_monthly_row_group.npy int32 source monthly row group per row
exposure_monthly_row_in_group.npy
int32 row offset inside monthly row group
exposure_manifest.json metadata and source parquet paths
For faster but much larger training storage, ``--mode dense`` materializes a
full dense numpy cache:
exposure_keys.npy uint64 legacy keys, key = (eid << 16) | raw_token
exposure_eid.npy int64 eid per exposure row
exposure_token.npy int32 raw disease token per exposure row
exposure_onset_date.npy datetime64[D] onset date per exposure row
exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2)
channels: tmean, rhmean
exposure_quality.npy float32 memmap, shape (N, 4)
n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata
The raw token convention follows the exposure README: padding=0, checkup=1,
and the first row of labels.csv is token=2. The model dataset inserts
<NO_EVENT> at token 2 and shifts real disease tokens by +1 internally; dataset
lookup converts back to these raw tokens before reading this cache. Dataset
alignment uses (eid, raw_token, onset_date - date_of_birth) so that raw
calendar dates in the exposure files match the age-day event times used by the
model.
Rows without matching exposure parquet records are kept as NaN windows. The
raw token convention follows the exposure README: padding=0, checkup=1, and
the first row of labels.csv is token=2. The model dataset inserts <NO_EVENT> at
token 2 and shifts real disease tokens by +1 internally; dataset lookup
converts back to these raw tokens before reading this cache.
"""
from __future__ import annotations
import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
import json
import os
from pathlib import Path
from typing import Iterable
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
try:
from tqdm.auto import tqdm
except ImportError:
def tqdm(iterable=None, **kwargs):
return iterable if iterable is not None else range(kwargs.get("total", 0))
DAILY_LENGTH = 1826
@@ -74,14 +58,6 @@ QUALITY_COLUMNS = (
)
def encode_exposure_key(eid: np.ndarray, raw_token: np.ndarray) -> np.ndarray:
eid_u64 = np.asarray(eid, dtype=np.uint64)
token_u64 = np.asarray(raw_token, dtype=np.uint64)
if np.any(token_u64 >= (1 << 16)):
raise ValueError("raw_token must fit in 16 bits")
return (eid_u64 << np.uint64(16)) | token_u64
def _daily_columns() -> list[str]:
cols: list[str] = []
for name in DAILY_CHANNELS:
@@ -97,7 +73,6 @@ def _monthly_columns() -> list[str]:
def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
"""Return the subset of requested columns present in a parquet file."""
try:
import pyarrow.parquet as pq
except ImportError as exc:
@@ -125,28 +100,6 @@ def _parquet_row_count(path: Path) -> int:
return int(pq.ParquetFile(path).metadata.num_rows)
def _row_group_positions(path: Path) -> tuple[np.ndarray, np.ndarray]:
"""Return row_group and row-in-group vectors for every parquet row."""
try:
import pyarrow.parquet as pq
except ImportError as exc:
raise ImportError(
"prepare_exposure_cache.py requires pyarrow. Install requirements "
"or run `pip install pyarrow`."
) from exc
parquet_file = pq.ParquetFile(path)
row_groups: list[np.ndarray] = []
row_offsets: list[np.ndarray] = []
for row_group_idx in range(parquet_file.num_row_groups):
n = parquet_file.metadata.row_group(row_group_idx).num_rows
row_groups.append(np.full(n, row_group_idx, dtype=np.int32))
row_offsets.append(np.arange(n, dtype=np.int32))
if not row_groups:
return np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
return np.concatenate(row_groups), np.concatenate(row_offsets)
def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)
@@ -173,9 +126,8 @@ def _load_summary(
summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name))
counts: list[int] = []
iterator = summary.itertuples(index=False)
iterator = tqdm(
iterator,
summary.itertuples(index=False),
total=len(summary),
desc="Counting exposure rows",
unit="file",
@@ -198,237 +150,63 @@ def _load_summary(
counts.append(daily_count)
summary["n_rows"] = counts
summary["offset"] = np.cumsum([0, *counts[:-1]], dtype=np.int64)
return summary
def _process_index_file_pair(task: tuple[int, str, str, str]) -> dict:
file_id, label_code, daily_path, monthly_path = task
daily_file = Path(daily_path)
monthly_file = Path(monthly_path)
def _load_sequence_rows(data_prefix: str) -> pd.DataFrame:
event_data = np.load(f"{data_prefix}_event_data.npy")
if event_data.ndim != 2 or event_data.shape[1] < 3:
raise ValueError(f"event_data must have shape (N, 3+), got {event_data.shape}")
event_data = event_data[:, :3].copy()
order = np.lexsort((event_data[:, 2], event_data[:, 1], event_data[:, 0]))
event_data = event_data[order]
daily_df = _read_parquet_columns(daily_file, ["eid", "onset_date", "token"])
monthly_df = _read_parquet_columns(monthly_file, ["eid", "onset_date", "token"])
if len(daily_df) != len(monthly_df):
basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
basic_table.index = basic_table.index.astype(np.int64)
if "date_of_birth" not in basic_table.columns:
raise ValueError(
f"Daily/monthly row count mismatch for {label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}"
f"{data_prefix}_basic_info.csv must contain date_of_birth for exposure alignment"
)
daily_rg, daily_row = _row_group_positions(daily_file)
monthly_rg_all, monthly_row_all = _row_group_positions(monthly_file)
n = len(daily_df)
if len(daily_rg) != n or len(monthly_rg_all) != n:
raise ValueError(f"Parquet row-group metadata row count mismatch for {label_code}")
daily_index = pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
monthly_index = pd.MultiIndex.from_frame(monthly_df[["eid", "onset_date", "token"]])
monthly_pos = monthly_index.get_indexer(daily_index)
if np.any(monthly_pos < 0):
raise ValueError(f"Monthly parquet is missing daily exposure keys for {label_code}")
return {
"file_id": int(file_id),
"label_code": label_code,
"n_rows": int(n),
"eid": daily_df["eid"].to_numpy(dtype=np.int64),
"token": daily_df["token"].to_numpy(dtype=np.int32),
"onset_date": pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).to_numpy(dtype="datetime64[D]"),
"daily_row_group": daily_rg,
"daily_row_in_group": daily_row,
"monthly_row_group": monthly_rg_all[monthly_pos],
"monthly_row_in_group": monthly_row_all[monthly_pos],
}
def build_exposure_index(
*,
exposure_dir: str | Path,
output_dir: str | Path,
summary_file: str = "summary.csv",
overwrite: bool = False,
workers: int = 1,
show_progress: bool = True,
) -> int:
exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
output_paths = [
output_dir / "exposure_eid.npy",
output_dir / "exposure_token.npy",
output_dir / "exposure_onset_date.npy",
output_dir / "exposure_daily_file_id.npy",
output_dir / "exposure_daily_row_group.npy",
output_dir / "exposure_daily_row_in_group.npy",
output_dir / "exposure_monthly_file_id.npy",
output_dir / "exposure_monthly_row_group.npy",
output_dir / "exposure_monthly_row_in_group.npy",
output_dir / "exposure_manifest.json",
]
if any(path.exists() for path in output_paths) and not overwrite:
raise FileExistsError(
f"{output_dir} already contains exposure index files; pass --overwrite"
)
summary = _load_summary(
exposure_dir,
summary_file,
show_progress=show_progress,
)
n_rows = int(summary["n_rows"].sum())
eids_mm = np.lib.format.open_memmap(
output_dir / "exposure_eid.npy", mode="w+", dtype=np.int64, shape=(n_rows,)
)
tokens_mm = np.lib.format.open_memmap(
output_dir / "exposure_token.npy", mode="w+", dtype=np.int32, shape=(n_rows,)
)
onset_dates_mm = np.lib.format.open_memmap(
output_dir / "exposure_onset_date.npy",
mode="w+",
dtype="datetime64[D]",
shape=(n_rows,),
)
daily_file_id_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_file_id.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
daily_row_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_row_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
daily_row_in_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_daily_row_in_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_file_id_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_file_id.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_row_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_row_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
)
monthly_row_in_group_mm = np.lib.format.open_memmap(
output_dir / "exposure_monthly_row_in_group.npy",
mode="w+",
dtype=np.int32,
shape=(n_rows,),
rows = pd.DataFrame(
{
"eid": event_data[:, 0].astype(np.int64),
"age_days": np.rint(event_data[:, 1].astype(np.float64)).astype(np.int32),
"token": event_data[:, 2].astype(np.int32),
}
)
rows = rows[rows["token"] > 1].reset_index(drop=True)
rows["position"] = np.arange(len(rows), dtype=np.int64)
tasks = [
(
int(file_id),
str(row.label_code),
str(Path(row.daily_path)),
str(Path(row.monthly_path)),
)
for file_id, row in enumerate(summary.itertuples(index=False))
]
workers = max(1, int(workers))
def write_result(result: dict) -> None:
file_id = int(result["file_id"])
row = summary.iloc[file_id]
offset = int(row.offset)
expected_n = int(row.n_rows)
n = int(result["n_rows"])
if n != expected_n:
raise RuntimeError(
f"Expected {expected_n} rows for {result['label_code']} "
f"from metadata but indexed {n}"
)
end = offset + n
if end > n_rows:
raise RuntimeError("Exposure index row count exceeded preallocated size")
eids_mm[offset:end] = result["eid"]
tokens_mm[offset:end] = result["token"]
onset_dates_mm[offset:end] = result["onset_date"]
daily_file_id_mm[offset:end] = file_id
daily_row_group_mm[offset:end] = result["daily_row_group"]
daily_row_in_group_mm[offset:end] = result["daily_row_in_group"]
monthly_file_id_mm[offset:end] = file_id
monthly_row_group_mm[offset:end] = result["monthly_row_group"]
monthly_row_in_group_mm[offset:end] = result["monthly_row_in_group"]
if workers == 1:
iterator = map(_process_index_file_pair, tasks)
for result in tqdm(
iterator,
total=len(tasks),
desc="Indexing exposure parquet",
unit="file",
disable=not show_progress,
):
write_result(result)
else:
with ProcessPoolExecutor(max_workers=workers) as executor:
futures = [executor.submit(_process_index_file_pair, task) for task in tasks]
for future in tqdm(
as_completed(futures),
total=len(futures),
desc=f"Indexing exposure parquet ({workers} workers)",
unit="file",
disable=not show_progress,
):
write_result(future.result())
for memmap in (
eids_mm,
tokens_mm,
onset_dates_mm,
daily_file_id_mm,
daily_row_group_mm,
daily_row_in_group_mm,
monthly_file_id_mm,
monthly_row_group_mm,
monthly_row_in_group_mm,
):
memmap.flush()
manifest = {
"storage": "parquet_index",
"source_dir": str(exposure_dir.resolve()),
"n_rows": int(n_rows),
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
"requires_basic_info_column": "date_of_birth",
"daily_files": [
str(Path(path).resolve()) for path in summary["daily_path"].tolist()
],
"monthly_files": [
str(Path(path).resolve()) for path in summary["monthly_path"].tolist()
],
"daily_shape_per_row": [DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS),
"monthly_shape_per_row": [MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
"monthly_channels": list(MONTHLY_CHANNELS),
"raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
}
(output_dir / "exposure_manifest.json").write_text(
json.dumps(manifest, indent=2),
encoding="utf-8",
birth = pd.to_datetime(
basic_table.loc[rows["eid"].to_numpy(), "date_of_birth"].to_numpy(),
errors="coerce",
)
return int(n_rows)
if birth.isna().any():
raise ValueError("date_of_birth contains missing or invalid values")
rows["onset_date"] = (
birth.to_numpy(dtype="datetime64[D]")
+ rows["age_days"].to_numpy(dtype="timedelta64[D]")
)
rows["onset_date"] = pd.to_datetime(rows["onset_date"]).dt.normalize()
return rows
def _write_eid_offsets(rows: pd.DataFrame, output_dir: Path) -> None:
eids = rows["eid"].to_numpy(dtype=np.int64)
unique_eids, starts = np.unique(eids, return_index=True)
starts = starts.astype(np.int64)
ends = np.concatenate([starts[1:], np.asarray([len(rows)], dtype=np.int64)])
eid_start = np.concatenate([starts, ends[-1:]]).astype(np.int64)
np.save(output_dir / "exposure_eid_index.npy", unique_eids.astype(np.int64))
np.save(output_dir / "exposure_eid_start.npy", eid_start)
def build_exposure_cache(
*,
exposure_dir: str | Path,
output_dir: str | Path,
data_prefix: str = "ukb",
summary_file: str = "summary.csv",
overwrite: bool = False,
show_progress: bool = True,
@@ -436,25 +214,20 @@ def build_exposure_cache(
exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
keys_path = output_dir / "exposure_keys.npy"
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
onset_date_path = output_dir / "exposure_onset_date.npy"
daily_path = output_dir / "exposure_daily.npy"
monthly_path = output_dir / "exposure_monthly.npy"
quality_path = output_dir / "exposure_quality.npy"
manifest_path = output_dir / "exposure_manifest.json"
outputs = [
keys_path,
eid_path,
token_path,
onset_date_path,
daily_path,
monthly_path,
quality_path,
manifest_path,
output_paths = [
output_dir / "exposure_eid.npy",
output_dir / "exposure_token.npy",
output_dir / "exposure_age_days.npy",
output_dir / "exposure_onset_date.npy",
output_dir / "exposure_eid_index.npy",
output_dir / "exposure_eid_start.npy",
output_dir / "exposure_daily.npy",
output_dir / "exposure_monthly.npy",
output_dir / "exposure_quality.npy",
output_dir / "exposure_manifest.json",
]
if any(path.exists() for path in outputs) and not overwrite:
if any(path.exists() for path in output_paths) and not overwrite:
raise FileExistsError(
f"{output_dir} already contains exposure cache files; pass --overwrite"
)
@@ -464,16 +237,29 @@ def build_exposure_cache(
summary_file,
show_progress=show_progress,
)
n_rows = int(summary["n_rows"].sum())
keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,))
eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,))
tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,))
onset_dates_mm = np.lib.format.open_memmap(
sequence_rows = _load_sequence_rows(data_prefix)
n_rows = len(sequence_rows)
if n_rows == 0:
raise ValueError(f"{data_prefix}_event_data.npy contains no real disease events")
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
age_path = output_dir / "exposure_age_days.npy"
onset_date_path = output_dir / "exposure_onset_date.npy"
daily_path = output_dir / "exposure_daily.npy"
monthly_path = output_dir / "exposure_monthly.npy"
quality_path = output_dir / "exposure_quality.npy"
manifest_path = output_dir / "exposure_manifest.json"
np.save(eid_path, sequence_rows["eid"].to_numpy(dtype=np.int64))
np.save(token_path, sequence_rows["token"].to_numpy(dtype=np.int32))
np.save(age_path, sequence_rows["age_days"].to_numpy(dtype=np.int32))
np.save(
onset_date_path,
mode="w+",
dtype="datetime64[D]",
shape=(n_rows,),
sequence_rows["onset_date"].to_numpy(dtype="datetime64[D]"),
)
_write_eid_offsets(sequence_rows, output_dir)
daily_mm = np.lib.format.open_memmap(
daily_path,
mode="w+",
@@ -492,25 +278,28 @@ def build_exposure_cache(
dtype=np.float32,
shape=(n_rows, len(QUALITY_COLUMNS)),
)
daily_mm[:] = np.nan
monthly_mm[:] = np.nan
quality_mm[:] = np.nan
daily_cols = _daily_columns()
monthly_cols = _monthly_columns()
offset = 0
wanted_by_token = {
int(token): frame.reset_index(drop=True)
for token, frame in sequence_rows.groupby("token", sort=False)
}
matched = np.zeros(n_rows, dtype=bool)
rows = tqdm(
iterator = tqdm(
summary.itertuples(index=False),
total=len(summary),
desc="Materializing dense exposure cache",
desc="Writing eid-sequence exposure cache",
unit="file",
disable=not show_progress,
)
for row in rows:
for row in iterator:
daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_path)
if not daily_file.is_file():
raise FileNotFoundError(f"Missing daily parquet: {daily_file}")
if not monthly_file.is_file():
raise FileNotFoundError(f"Missing monthly parquet: {monthly_file}")
daily_read_cols = [
"eid",
@@ -528,70 +317,76 @@ def build_exposure_cache(
]
daily_df = _read_parquet_columns(daily_file, daily_read_cols)
monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols)
if len(daily_df) != len(monthly_df):
raise ValueError(
f"Daily/monthly row count mismatch for {row.label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}"
)
daily_df = daily_df.copy()
monthly_df = monthly_df.copy()
daily_df["_source_row"] = np.arange(len(daily_df), dtype=np.int64)
daily_df["onset_date"] = pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).dt.normalize()
monthly_df["onset_date"] = pd.to_datetime(
monthly_df["onset_date"],
errors="coerce",
).dt.normalize()
monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
).reset_index()
n = len(daily_df)
end = offset + n
if end > n_rows:
raise RuntimeError("Exposure cache row count exceeded preallocated size")
tokens = daily_df["token"].dropna().astype(np.int64).unique()
wanted = pd.concat(
[wanted_by_token[int(token)] for token in tokens if int(token) in wanted_by_token],
ignore_index=True,
) if len(tokens) else pd.DataFrame()
if wanted.empty:
continue
keys[offset:end] = encode_exposure_key(
daily_df["eid"].to_numpy(dtype=np.int64),
daily_df["token"].to_numpy(dtype=np.int64),
matches = daily_df[["eid", "onset_date", "token", "_source_row"]].merge(
wanted[["eid", "onset_date", "token", "position"]],
on=["eid", "onset_date", "token"],
how="inner",
sort=False,
)
eids_mm[offset:end] = daily_df["eid"].to_numpy(dtype=np.int64)
tokens_mm[offset:end] = daily_df["token"].to_numpy(dtype=np.int32)
onset_dates_mm[offset:end] = pd.to_datetime(
daily_df["onset_date"],
errors="coerce",
).to_numpy(dtype="datetime64[D]")
daily_mm[offset:end] = _reshape_window(
daily_df,
if matches.empty:
continue
source_rows = matches["_source_row"].to_numpy(dtype=np.int64)
positions = matches["position"].to_numpy(dtype=np.int64)
daily_mm[positions] = _reshape_window(
daily_df.iloc[source_rows],
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
)
monthly_mm[offset:end] = _reshape_window(
monthly_df,
monthly_mm[positions] = _reshape_window(
monthly_df.iloc[source_rows],
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
)
quality_mm[offset:end, 0] = daily_df.get("n_days_nonmissing", np.nan)
quality_mm[offset:end, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
quality_mm[offset:end, 2] = monthly_df.get("n_months_nonmissing", np.nan)
quality_mm[offset:end, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
offset = end
quality_mm[positions, 0] = daily_df.iloc[source_rows].get("n_days_nonmissing", np.nan)
quality_mm[positions, 1] = daily_df.iloc[source_rows].get("n_rh_days_nonmissing", np.nan)
quality_mm[positions, 2] = monthly_df.iloc[source_rows].get("n_months_nonmissing", np.nan)
quality_mm[positions, 3] = monthly_df.iloc[source_rows].get("n_rh_months_nonmissing", np.nan)
matched[positions] = True
if offset != n_rows:
keys.flush()
eids_mm.flush()
tokens_mm.flush()
onset_dates_mm.flush()
daily_mm.flush()
monthly_mm.flush()
quality_mm.flush()
keys = np.lib.format.open_memmap(keys_path, mode="r+", dtype=np.uint64, shape=(offset,))
raise RuntimeError(
f"Expected {n_rows} rows from summary but wrote {offset}. "
"Check parquet metadata and regenerate summary.csv before building."
)
daily_mm.flush()
monthly_mm.flush()
quality_mm.flush()
manifest = {
"storage": "dense_npy",
"source_dir": str(exposure_dir),
"storage": "eid_sequence_npy",
"source_dir": str(exposure_dir.resolve()),
"data_prefix": data_prefix,
"n_rows": int(n_rows),
"legacy_key": "(eid << 16) | raw_token",
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
"matched_rows": int(matched.sum()),
"missing_rows": int((~matched).sum()),
"alignment_key": "(eid, raw_token, date_of_birth + age_days)",
"requires_basic_info_column": "date_of_birth",
"daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS),
@@ -608,25 +403,8 @@ def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--exposure-dir", required=True)
parser.add_argument("--output-dir", default="ukb_exposure_cache")
parser.add_argument("--data-prefix", default="ukb")
parser.add_argument("--summary-file", default="summary.csv")
parser.add_argument(
"--mode",
choices=("index", "dense"),
default="index",
help=(
"index writes only lightweight parquet row pointers; dense copies "
"all exposure windows into numpy memmaps."
),
)
parser.add_argument(
"--workers",
type=int,
default=max(1, min(8, (os.cpu_count() or 1))),
help=(
"Number of worker processes for --mode index. Dense mode remains "
"single-writer to avoid concurrent writes to the same memmap."
),
)
parser.add_argument(
"--no-progress",
action="store_true",
@@ -634,26 +412,16 @@ def main() -> None:
)
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
show_progress = not args.no_progress
if args.mode == "index":
n_rows = build_exposure_index(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
summary_file=args.summary_file,
overwrite=args.overwrite,
workers=args.workers,
show_progress=show_progress,
)
print(f"Wrote {n_rows:,} exposure row pointers to {args.output_dir}")
else:
n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
summary_file=args.summary_file,
overwrite=args.overwrite,
show_progress=show_progress,
)
print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")
n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
data_prefix=args.data_prefix,
summary_file=args.summary_file,
overwrite=args.overwrite,
show_progress=not args.no_progress,
)
print(f"Wrote {n_rows:,} eid-sequence-aligned exposure rows to {args.output_dir}")
if __name__ == "__main__":