661 lines
24 KiB
Python
661 lines
24 KiB
Python
"""Build a random-access exposure index/cache from disease-level parquet files.
|
|
|
|
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.
|
|
|
|
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 default index directory contains:
|
|
|
|
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.
|
|
"""
|
|
|
|
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
|
|
|
|
|
|
DAILY_LENGTH = 1826
|
|
MONTHLY_LENGTH = 241
|
|
DAILY_CHANNELS = ("tmean", "tmax", "tmin", "rhmean")
|
|
MONTHLY_CHANNELS = ("tmean", "rhmean")
|
|
QUALITY_COLUMNS = (
|
|
"n_days_nonmissing",
|
|
"n_rh_days_nonmissing",
|
|
"n_months_nonmissing",
|
|
"n_rh_months_nonmissing",
|
|
)
|
|
|
|
|
|
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:
|
|
cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_LENGTH))
|
|
return cols
|
|
|
|
|
|
def _monthly_columns() -> list[str]:
|
|
cols: list[str] = []
|
|
for name in MONTHLY_CHANNELS:
|
|
cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_LENGTH))
|
|
return cols
|
|
|
|
|
|
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:
|
|
raise ImportError(
|
|
"prepare_exposure_cache.py requires pyarrow. Install requirements "
|
|
"or run `pip install pyarrow`."
|
|
) from exc
|
|
|
|
schema_names = set(pq.ParquetFile(path).schema.names)
|
|
return [col for col in columns if col in schema_names]
|
|
|
|
|
|
def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame:
|
|
return pd.read_parquet(path, columns=columns)
|
|
|
|
|
|
def _parquet_row_count(path: Path) -> int:
|
|
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
|
|
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)
|
|
|
|
|
|
def _load_summary(
|
|
exposure_dir: Path,
|
|
summary_file: str,
|
|
*,
|
|
show_progress: bool,
|
|
) -> pd.DataFrame:
|
|
summary_path = exposure_dir / summary_file
|
|
if not summary_path.is_file():
|
|
raise FileNotFoundError(f"summary.csv not found: {summary_path}")
|
|
|
|
summary = pd.read_csv(summary_path)
|
|
required = {"label_code", "daily_file", "monthly_file"}
|
|
missing = required - set(summary.columns)
|
|
if missing:
|
|
raise ValueError(f"{summary_path} is missing columns: {sorted(missing)}")
|
|
|
|
summary = summary.copy()
|
|
summary["daily_path"] = summary["daily_file"].map(lambda name: exposure_dir / str(name))
|
|
summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name))
|
|
|
|
counts: list[int] = []
|
|
iterator = summary.itertuples(index=False)
|
|
iterator = tqdm(
|
|
iterator,
|
|
total=len(summary),
|
|
desc="Counting exposure rows",
|
|
unit="file",
|
|
disable=not show_progress,
|
|
)
|
|
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_count = _parquet_row_count(daily_file)
|
|
monthly_count = _parquet_row_count(monthly_file)
|
|
if daily_count != monthly_count:
|
|
raise ValueError(
|
|
f"Daily/monthly row count mismatch for {row.label_code}: "
|
|
f"{daily_count} vs {monthly_count}"
|
|
)
|
|
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)
|
|
|
|
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):
|
|
raise ValueError(
|
|
f"Daily/monthly row count mismatch for {label_code}: "
|
|
f"{len(daily_df)} vs {len(monthly_df)}"
|
|
)
|
|
|
|
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,),
|
|
)
|
|
|
|
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",
|
|
)
|
|
return int(n_rows)
|
|
|
|
|
|
def build_exposure_cache(
|
|
*,
|
|
exposure_dir: str | Path,
|
|
output_dir: str | Path,
|
|
summary_file: str = "summary.csv",
|
|
overwrite: bool = False,
|
|
show_progress: bool = True,
|
|
) -> int:
|
|
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,
|
|
]
|
|
if any(path.exists() for path in outputs) and not overwrite:
|
|
raise FileExistsError(
|
|
f"{output_dir} already contains exposure cache files; pass --overwrite"
|
|
)
|
|
|
|
summary = _load_summary(
|
|
exposure_dir,
|
|
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(
|
|
onset_date_path,
|
|
mode="w+",
|
|
dtype="datetime64[D]",
|
|
shape=(n_rows,),
|
|
)
|
|
daily_mm = np.lib.format.open_memmap(
|
|
daily_path,
|
|
mode="w+",
|
|
dtype=np.float32,
|
|
shape=(n_rows, DAILY_LENGTH, len(DAILY_CHANNELS)),
|
|
)
|
|
monthly_mm = np.lib.format.open_memmap(
|
|
monthly_path,
|
|
mode="w+",
|
|
dtype=np.float32,
|
|
shape=(n_rows, MONTHLY_LENGTH, len(MONTHLY_CHANNELS)),
|
|
)
|
|
quality_mm = np.lib.format.open_memmap(
|
|
quality_path,
|
|
mode="w+",
|
|
dtype=np.float32,
|
|
shape=(n_rows, len(QUALITY_COLUMNS)),
|
|
)
|
|
|
|
daily_cols = _daily_columns()
|
|
monthly_cols = _monthly_columns()
|
|
offset = 0
|
|
|
|
rows = tqdm(
|
|
summary.itertuples(index=False),
|
|
total=len(summary),
|
|
desc="Materializing dense exposure cache",
|
|
unit="file",
|
|
disable=not show_progress,
|
|
)
|
|
for row in rows:
|
|
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",
|
|
"onset_date",
|
|
"token",
|
|
*_safe_columns(daily_file, daily_cols),
|
|
*_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]),
|
|
]
|
|
monthly_read_cols = [
|
|
"eid",
|
|
"onset_date",
|
|
"token",
|
|
*_safe_columns(monthly_file, monthly_cols),
|
|
*_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]),
|
|
]
|
|
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)}"
|
|
)
|
|
|
|
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")
|
|
|
|
keys[offset:end] = encode_exposure_key(
|
|
daily_df["eid"].to_numpy(dtype=np.int64),
|
|
daily_df["token"].to_numpy(dtype=np.int64),
|
|
)
|
|
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,
|
|
daily_cols,
|
|
DAILY_LENGTH,
|
|
len(DAILY_CHANNELS),
|
|
)
|
|
monthly_mm[offset:end] = _reshape_window(
|
|
monthly_df,
|
|
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
|
|
|
|
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."
|
|
)
|
|
|
|
manifest = {
|
|
"storage": "dense_npy",
|
|
"source_dir": str(exposure_dir),
|
|
"n_rows": int(n_rows),
|
|
"legacy_key": "(eid << 16) | raw_token",
|
|
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
|
|
"requires_basic_info_column": "date_of_birth",
|
|
"daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
|
|
"daily_channels": list(DAILY_CHANNELS),
|
|
"monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
|
|
"monthly_channels": list(MONTHLY_CHANNELS),
|
|
"quality_columns": list(QUALITY_COLUMNS),
|
|
"raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
|
|
}
|
|
manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
|
return int(n_rows)
|
|
|
|
|
|
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("--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",
|
|
help="Disable tqdm progress bars.",
|
|
)
|
|
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}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|