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

280
dataset.py
View File

@@ -2,7 +2,6 @@
from __future__ import annotations from __future__ import annotations
import json import json
from collections import OrderedDict
from pathlib import Path from pathlib import Path
from typing import Dict, Iterable, List, Literal, Optional, Tuple from typing import Dict, Iterable, List, Literal, Optional, Tuple
@@ -43,9 +42,9 @@ def _monthly_exposure_columns() -> list[str]:
class ExposureCache: class ExposureCache:
"""Random-access view over files produced by prepare_exposure_cache.py.""" """Eid-sequence-aligned exposure windows from prepare_exposure_cache.py."""
def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4): def __init__(self, cache_dir: str | Path):
cache_dir = Path(cache_dir) cache_dir = Path(cache_dir)
self.cache_dir = cache_dir self.cache_dir = cache_dir
manifest_path = cache_dir / "exposure_manifest.json" manifest_path = cache_dir / "exposure_manifest.json"
@@ -54,51 +53,48 @@ class ExposureCache:
if manifest_path.is_file() if manifest_path.is_file()
else {} else {}
) )
self.storage = self.manifest.get("storage", "dense_npy") self.storage = self.manifest.get("storage")
self._row_group_cache_size = int(row_group_cache_size) if self.storage != "eid_sequence_npy":
self._row_group_cache: OrderedDict[tuple[str, int, int], pd.DataFrame] = OrderedDict() raise ValueError(
self._parquet_files: dict[tuple[str, int], object] = {} "Exposure cache must use storage='eid_sequence_npy'. "
self._parquet_columns: dict[tuple[str, int], list[str]] = {}
eid_path = cache_dir / "exposure_eid.npy"
token_path = cache_dir / "exposure_token.npy"
onset_date_path = cache_dir / "exposure_onset_date.npy"
if not (eid_path.is_file() and token_path.is_file() and onset_date_path.is_file()):
raise FileNotFoundError(
"Exposure cache must contain exposure_eid.npy, "
"exposure_token.npy, and exposure_onset_date.npy. "
"Regenerate it with the current prepare_exposure_cache.py." "Regenerate it with the current prepare_exposure_cache.py."
) )
eid_path = cache_dir / "exposure_eid.npy"
token_path = cache_dir / "exposure_token.npy"
age_path = cache_dir / "exposure_age_days.npy"
onset_date_path = cache_dir / "exposure_onset_date.npy"
eid_index_path = cache_dir / "exposure_eid_index.npy"
eid_start_path = cache_dir / "exposure_eid_start.npy"
daily_path = cache_dir / "exposure_daily.npy"
monthly_path = cache_dir / "exposure_monthly.npy"
required_paths = [
eid_path,
token_path,
age_path,
onset_date_path,
eid_index_path,
eid_start_path,
daily_path,
monthly_path,
]
if any(not path.is_file() for path in required_paths):
raise FileNotFoundError(
"Exposure cache is missing one or more eid-sequence files. "
"Regenerate it with the current prepare_exposure_cache.py."
)
self.eids = np.load(eid_path, mmap_mode="r") self.eids = np.load(eid_path, mmap_mode="r")
self.raw_tokens = np.load(token_path, mmap_mode="r") self.raw_tokens = np.load(token_path, mmap_mode="r")
self.age_days = np.load(age_path, mmap_mode="r")
self.onset_dates = np.load(onset_date_path, mmap_mode="r") self.onset_dates = np.load(onset_date_path, mmap_mode="r")
self.daily = None self.eid_index = np.load(eid_index_path, mmap_mode="r")
self.monthly = None self.eid_start = np.load(eid_start_path, mmap_mode="r")
if self.storage == "dense_npy": self.daily = np.load(daily_path, mmap_mode="r")
self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r") self.monthly = np.load(monthly_path, mmap_mode="r")
self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r")
elif self.storage == "parquet_index":
self.daily_file_ids = np.load(cache_dir / "exposure_daily_file_id.npy", mmap_mode="r")
self.daily_row_groups = np.load(cache_dir / "exposure_daily_row_group.npy", mmap_mode="r")
self.daily_row_in_groups = np.load(
cache_dir / "exposure_daily_row_in_group.npy", mmap_mode="r"
)
self.monthly_file_ids = np.load(
cache_dir / "exposure_monthly_file_id.npy", mmap_mode="r"
)
self.monthly_row_groups = np.load(
cache_dir / "exposure_monthly_row_group.npy", mmap_mode="r"
)
self.monthly_row_in_groups = np.load(
cache_dir / "exposure_monthly_row_in_group.npy", mmap_mode="r"
)
self.daily_files = [Path(path) for path in self.manifest["daily_files"]]
self.monthly_files = [Path(path) for path in self.manifest["monthly_files"]]
else:
raise ValueError(f"Unknown exposure cache storage mode: {self.storage!r}")
quality_path = cache_dir / "exposure_quality.npy" quality_path = cache_dir / "exposure_quality.npy"
self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None
if self.storage == "dense_npy":
if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE: if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE:
raise ValueError( raise ValueError(
f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, " f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, "
@@ -110,67 +106,72 @@ class ExposureCache:
f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}" f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}"
) )
n_rows = len(self.eids) n_rows = len(self.eids)
if len(self.raw_tokens) != n_rows or len(self.onset_dates) != n_rows: if (
len(self.raw_tokens) != n_rows
or len(self.age_days) != n_rows
or len(self.onset_dates) != n_rows
or self.daily.shape[0] != n_rows
or self.monthly.shape[0] != n_rows
):
raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") raise ValueError("Exposure cache metadata/daily/monthly row counts do not match")
if self.storage == "dense_npy": if len(self.eid_start) != len(self.eid_index) + 1:
if self.daily.shape[0] != n_rows or self.monthly.shape[0] != n_rows: raise ValueError("exposure_eid_start.npy must have len(eid_index) + 1")
raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") if len(self.eid_start) and int(self.eid_start[-1]) != n_rows:
else: raise ValueError("Last exposure eid offset must equal exposure row count")
indexed_lengths = [
len(self.daily_file_ids),
len(self.daily_row_groups),
len(self.daily_row_in_groups),
len(self.monthly_file_ids),
len(self.monthly_row_groups),
len(self.monthly_row_in_groups),
]
if any(length != n_rows for length in indexed_lengths):
raise ValueError("Exposure parquet index row counts do not match metadata")
self._key_to_index: dict[tuple[int, int, int], int] | None = None self._eid_to_pos = {
int(eid): idx
for idx, eid in enumerate(np.asarray(self.eid_index, dtype=np.int64))
}
def locality_key(self, indices: np.ndarray) -> tuple[int, int]: def locality_key(self, indices: np.ndarray) -> tuple[int, int]:
"""Return a stable parquet locality key for sampler-side batching.""" """Return a stable locality key for sampler-side batching."""
indices = np.asarray(indices, dtype=np.int64) indices = np.asarray(indices, dtype=np.int64)
valid = indices[indices >= 0] valid = indices[indices >= 0]
if len(valid) == 0: if len(valid) == 0:
return (2**31 - 1, 2**31 - 1) return (2**31 - 1, 2**31 - 1)
if self.storage != "parquet_index": first = int(valid[0])
return (0, int(valid[0] // 1024)) return (first // 1024, first % 1024)
file_ids = np.asarray(self.daily_file_ids[valid], dtype=np.int64)
row_groups = np.asarray(self.daily_row_groups[valid], dtype=np.int64)
groups = np.stack([file_ids, row_groups], axis=1)
unique_groups, counts = np.unique(groups, axis=0, return_counts=True)
best = unique_groups[int(np.argmax(counts))]
return (int(best[0]), int(best[1]))
def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None: def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None:
keys: dict[tuple[int, int, int], int] = {} """Kept for the dataset constructor; the new cache already stores age days."""
eids = np.asarray(self.eids, dtype=np.int64) return None
tokens = np.asarray(self.raw_tokens, dtype=np.int64)
onset_dates = np.asarray(self.onset_dates, dtype="datetime64[D]")
for idx, (eid, token, onset_date) in enumerate(zip(eids, tokens, onset_dates)):
birth_date = birth_date_by_eid.get(int(eid))
if birth_date is None or np.isnat(onset_date) or np.isnat(birth_date):
continue
age_days = int((onset_date - birth_date).astype("timedelta64[D]").astype(np.int64))
if age_days < 0:
continue
keys[(int(eid), int(token), age_days)] = idx
self._key_to_index = keys
def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray: def lookup_indices(self, eid: int, raw_tokens: np.ndarray, age_days: np.ndarray) -> np.ndarray:
if self._key_to_index is None:
raise RuntimeError("ExposureCache.build_age_index must be called before lookup")
out = np.full(len(raw_tokens), -1, dtype=np.int64) out = np.full(len(raw_tokens), -1, dtype=np.int64)
real = raw_tokens > 1 real = raw_tokens > 1
if not np.any(real): if not np.any(real):
return out return out
eid_pos = self._eid_to_pos.get(int(eid))
if eid_pos is None:
return out
start = int(self.eid_start[eid_pos])
end = int(self.eid_start[eid_pos + 1])
if start == end:
return out
real_pos = np.nonzero(real)[0] real_pos = np.nonzero(real)[0]
out[real_pos] = [ n_take = min(len(real_pos), end - start)
self._key_to_index.get((int(eid), int(raw_tokens[pos]), int(round(float(age_days[pos])))), -1) if n_take == 0:
for pos in real_pos return out
]
out[real_pos[:n_take]] = np.arange(start, start + n_take, dtype=np.int64)
expected_tokens = np.asarray(self.raw_tokens[start:start + n_take], dtype=np.int64)
expected_age_days = np.asarray(self.age_days[start:start + n_take], dtype=np.int64)
actual_tokens = np.asarray(raw_tokens[real_pos[:n_take]], dtype=np.int64)
actual_age_days = np.rint(
np.asarray(age_days[real_pos[:n_take]], dtype=np.float64)
).astype(np.int64)
if (
not np.array_equal(expected_tokens, actual_tokens)
or not np.array_equal(expected_age_days, actual_age_days)
):
raise ValueError(
"Exposure cache is not aligned to the disease sequence for "
f"eid={eid}. Regenerate it with the same data_prefix and labels."
)
return out return out
def daily_window(self, index: int) -> np.ndarray: def daily_window(self, index: int) -> np.ndarray:
@@ -198,111 +199,10 @@ class ExposureCache:
return out return out
valid_indices = indices[valid_pos] valid_indices = indices[valid_pos]
if self.storage == "dense_npy":
source = self.daily if kind == "daily" else self.monthly source = self.daily if kind == "daily" else self.monthly
out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32) out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32)
return out return out
if kind == "daily":
file_ids = np.asarray(self.daily_file_ids[valid_indices], dtype=np.int64)
row_groups = np.asarray(self.daily_row_groups[valid_indices], dtype=np.int64)
row_in_groups = np.asarray(self.daily_row_in_groups[valid_indices], dtype=np.int64)
columns = _daily_exposure_columns()
else:
file_ids = np.asarray(self.monthly_file_ids[valid_indices], dtype=np.int64)
row_groups = np.asarray(self.monthly_row_groups[valid_indices], dtype=np.int64)
row_in_groups = np.asarray(
self.monthly_row_in_groups[valid_indices],
dtype=np.int64,
)
columns = _monthly_exposure_columns()
group_keys = np.stack([file_ids, row_groups], axis=1)
unique_groups, inverse = np.unique(group_keys, axis=0, return_inverse=True)
for group_idx, (file_id, row_group) in enumerate(unique_groups):
group_pos = np.nonzero(inverse == group_idx)[0]
frame = self._read_parquet_row_group(
kind,
int(file_id),
int(row_group),
columns,
)
row_values = frame.iloc[row_in_groups[group_pos]].reindex(columns=columns)
values = (
row_values.to_numpy(dtype=np.float32, copy=True)
.reshape(len(group_pos), shape[1], shape[0])
.transpose(0, 2, 1)
)
out[valid_pos[group_pos]] = values
return out
def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray:
if kind == "daily":
file_id = int(self.daily_file_ids[index])
row_group = int(self.daily_row_groups[index])
row_in_group = int(self.daily_row_in_groups[index])
shape = DAILY_EXPOSURE_SHAPE
columns = _daily_exposure_columns()
else:
file_id = int(self.monthly_file_ids[index])
row_group = int(self.monthly_row_groups[index])
row_in_group = int(self.monthly_row_in_groups[index])
shape = MONTHLY_EXPOSURE_SHAPE
columns = _monthly_exposure_columns()
frame = self._read_parquet_row_group(kind, file_id, row_group, columns)
row = frame.iloc[row_in_group].reindex(columns)
n_channels = shape[1]
return (
row.to_numpy(dtype=np.float32, copy=True)
.reshape(n_channels, shape[0])
.transpose(1, 0)
)
def _read_parquet_row_group(
self,
kind: Literal["daily", "monthly"],
file_id: int,
row_group: int,
columns: list[str],
) -> pd.DataFrame:
cache_key = (kind, file_id, row_group)
cached = self._row_group_cache.get(cache_key)
if cached is not None:
self._row_group_cache.move_to_end(cache_key)
return cached
try:
import pyarrow.parquet as pq
except ImportError as exc:
raise ImportError(
"Parquet exposure index loading requires pyarrow. Install requirements "
"or use a dense numpy exposure cache."
) from exc
parquet_key = (kind, file_id)
parquet_file = self._parquet_files.get(parquet_key)
if parquet_file is None:
path = self.daily_files[file_id] if kind == "daily" else self.monthly_files[file_id]
parquet_file = pq.ParquetFile(path)
self._parquet_files[parquet_key] = parquet_file
available_columns = self._parquet_columns.get(parquet_key)
if available_columns is None:
available = set(parquet_file.schema.names)
available_columns = [col for col in columns if col in available]
self._parquet_columns[parquet_key] = available_columns
table = parquet_file.read_row_group(row_group, columns=available_columns)
frame = table.to_pandas()
if available_columns != columns:
frame = frame.reindex(columns=columns)
self._row_group_cache[cache_key] = frame
self._row_group_cache.move_to_end(cache_key)
while len(self._row_group_cache) > self._row_group_cache_size:
self._row_group_cache.popitem(last=False)
return frame
def load_label_vocab( def load_label_vocab(
labels_file: str, labels_file: str,
@@ -372,17 +272,13 @@ class _ExpoBaseDataset(Dataset):
include_no_event_in_uts_target: bool = False, include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None, exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False, mask_onset_exposure: bool = False,
exposure_row_group_cache_size: int = 4,
) -> None: ) -> None:
self.data_prefix = data_prefix self.data_prefix = data_prefix
self.labels_file = labels_file self.labels_file = labels_file
self.no_event_interval_years = float(no_event_interval_years) self.no_event_interval_years = float(no_event_interval_years)
self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target) self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target)
self.exposure_cache = ( self.exposure_cache = (
ExposureCache( ExposureCache(exposure_cache_dir)
exposure_cache_dir,
row_group_cache_size=exposure_row_group_cache_size,
)
if exposure_cache_dir is not None if exposure_cache_dir is not None
else None else None
) )
@@ -553,7 +449,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
include_no_event_in_uts_target: bool = False, include_no_event_in_uts_target: bool = False,
exposure_cache_dir: str | Path | None = None, exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False, mask_onset_exposure: bool = False,
exposure_row_group_cache_size: int = 4,
) -> None: ) -> None:
super().__init__( super().__init__(
data_prefix=data_prefix, data_prefix=data_prefix,
@@ -562,7 +457,6 @@ class NextStepHealthDataset(_ExpoBaseDataset):
include_no_event_in_uts_target=include_no_event_in_uts_target, include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir, exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure, mask_onset_exposure=mask_onset_exposure,
exposure_row_group_cache_size=exposure_row_group_cache_size,
) )
self.samples: List[Dict] = [] self.samples: List[Dict] = []
@@ -651,7 +545,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
validation_query_seed: int = 42, validation_query_seed: int = 42,
exposure_cache_dir: str | Path | None = None, exposure_cache_dir: str | Path | None = None,
mask_onset_exposure: bool = False, mask_onset_exposure: bool = False,
exposure_row_group_cache_size: int = 4,
) -> None: ) -> None:
if split not in {"train", "valid", "test"}: if split not in {"train", "valid", "test"}:
raise ValueError(f"split must be train/valid/test, got {split!r}") raise ValueError(f"split must be train/valid/test, got {split!r}")
@@ -663,7 +556,6 @@ class AllFutureHealthDataset(_ExpoBaseDataset):
include_no_event_in_uts_target=include_no_event_in_uts_target, include_no_event_in_uts_target=include_no_event_in_uts_target,
exposure_cache_dir=exposure_cache_dir, exposure_cache_dir=exposure_cache_dir,
mask_onset_exposure=mask_onset_exposure, mask_onset_exposure=mask_onset_exposure,
exposure_row_group_cache_size=exposure_row_group_cache_size,
) )
self.split = split self.split = split

548
prepare_exposure_cache.py
View File

@@ -1,36 +1,19 @@
"""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 The source exposure dataset is stored as one daily and one monthly parquet file
parquet file per disease. That layout is good for disease-specific analysis but per disease. That layout is inconvenient for mini-batch training because the
too expensive for mini-batch training, where we need exposure windows aligned model consumes per-participant disease sequences. This script materializes one
to arbitrary event sequences. 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 The output directory contains:
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 real disease event
exposure_token.npy int32 raw disease token per real disease event
exposure_eid.npy int64 eid per exposure row exposure_age_days.npy int32 age in days per real disease event
exposure_token.npy int32 raw disease token per exposure row exposure_onset_date.npy datetime64[D] onset date per real disease event
exposure_onset_date.npy datetime64[D] onset date per exposure row exposure_eid_index.npy int64 unique eids in cache order
exposure_daily_file_id.npy int32 source daily file id per row exposure_eid_start.npy int64 start offsets, length len(eid_index) + 1
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) exposure_daily.npy float32 memmap, shape (N, 1826, 4)
channels: tmean, tmax, tmin, rhmean channels: tmean, tmax, tmin, rhmean
exposure_monthly.npy float32 memmap, shape (N, 241, 2) exposure_monthly.npy float32 memmap, shape (N, 241, 2)
@@ -39,27 +22,28 @@ full dense numpy cache:
n_days, n_rh_days, n_months, n_rh_months n_days, n_rh_days, n_months, n_rh_months
exposure_manifest.json metadata exposure_manifest.json metadata
The raw token convention follows the exposure README: padding=0, checkup=1, Rows without matching exposure parquet records are kept as NaN windows. The
and the first row of labels.csv is token=2. The model dataset inserts raw token convention follows the exposure README: padding=0, checkup=1, and
<NO_EVENT> at token 2 and shifts real disease tokens by +1 internally; dataset the first row of labels.csv is token=2. The model dataset inserts <NO_EVENT> at
lookup converts back to these raw tokens before reading this cache. Dataset token 2 and shifts real disease tokens by +1 internally; dataset lookup
alignment uses (eid, raw_token, onset_date - date_of_birth) so that raw converts back to these raw tokens before reading this cache.
calendar dates in the exposure files match the age-day event times used by the
model.
""" """
from __future__ import annotations from __future__ import annotations
import argparse import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
import json import json
import os
from pathlib import Path from pathlib import Path
from typing import Iterable from typing import Iterable
import numpy as np import numpy as np
import pandas as pd 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 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]: def _daily_columns() -> list[str]:
cols: list[str] = [] cols: list[str] = []
for name in DAILY_CHANNELS: for name in DAILY_CHANNELS:
@@ -97,7 +73,6 @@ def _monthly_columns() -> list[str]:
def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]: def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
"""Return the subset of requested columns present in a parquet file."""
try: try:
import pyarrow.parquet as pq import pyarrow.parquet as pq
except ImportError as exc: except ImportError as exc:
@@ -125,28 +100,6 @@ def _parquet_row_count(path: Path) -> int:
return int(pq.ParquetFile(path).metadata.num_rows) 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: 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) arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1) 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)) summary["monthly_path"] = summary["monthly_file"].map(lambda name: exposure_dir / str(name))
counts: list[int] = [] counts: list[int] = []
iterator = summary.itertuples(index=False)
iterator = tqdm( iterator = tqdm(
iterator, summary.itertuples(index=False),
total=len(summary), total=len(summary),
desc="Counting exposure rows", desc="Counting exposure rows",
unit="file", unit="file",
@@ -198,237 +150,63 @@ def _load_summary(
counts.append(daily_count) counts.append(daily_count)
summary["n_rows"] = counts summary["n_rows"] = counts
summary["offset"] = np.cumsum([0, *counts[:-1]], dtype=np.int64)
return summary return summary
def _process_index_file_pair(task: tuple[int, str, str, str]) -> dict: def _load_sequence_rows(data_prefix: str) -> pd.DataFrame:
file_id, label_code, daily_path, monthly_path = task event_data = np.load(f"{data_prefix}_event_data.npy")
daily_file = Path(daily_path) if event_data.ndim != 2 or event_data.shape[1] < 3:
monthly_file = Path(monthly_path) 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"]) basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
monthly_df = _read_parquet_columns(monthly_file, ["eid", "onset_date", "token"]) basic_table.index = basic_table.index.astype(np.int64)
if len(daily_df) != len(monthly_df): if "date_of_birth" not in basic_table.columns:
raise ValueError( raise ValueError(
f"Daily/monthly row count mismatch for {label_code}: " f"{data_prefix}_basic_info.csv must contain date_of_birth for exposure alignment"
f"{len(daily_df)} vs {len(monthly_df)}"
) )
daily_rg, daily_row = _row_group_positions(daily_file) rows = pd.DataFrame(
monthly_rg_all, monthly_row_all = _row_group_positions(monthly_file) {
n = len(daily_df) "eid": event_data[:, 0].astype(np.int64),
if len(daily_rg) != n or len(monthly_rg_all) != n: "age_days": np.rint(event_data[:, 1].astype(np.float64)).astype(np.int32),
raise ValueError(f"Parquet row-group metadata row count mismatch for {label_code}") "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)
daily_index = pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]]) birth = pd.to_datetime(
monthly_index = pd.MultiIndex.from_frame(monthly_df[["eid", "onset_date", "token"]]) basic_table.loc[rows["eid"].to_numpy(), "date_of_birth"].to_numpy(),
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", errors="coerce",
).to_numpy(dtype="datetime64[D]"), )
"daily_row_group": daily_rg, if birth.isna().any():
"daily_row_in_group": daily_row, raise ValueError("date_of_birth contains missing or invalid values")
"monthly_row_group": monthly_rg_all[monthly_pos], rows["onset_date"] = (
"monthly_row_in_group": monthly_row_all[monthly_pos], 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 build_exposure_index( def _write_eid_offsets(rows: pd.DataFrame, output_dir: Path) -> None:
*, eids = rows["eid"].to_numpy(dtype=np.int64)
exposure_dir: str | Path, unique_eids, starts = np.unique(eids, return_index=True)
output_dir: str | Path, starts = starts.astype(np.int64)
summary_file: str = "summary.csv", ends = np.concatenate([starts[1:], np.asarray([len(rows)], dtype=np.int64)])
overwrite: bool = False, eid_start = np.concatenate([starts, ends[-1:]]).astype(np.int64)
workers: int = 1, np.save(output_dir / "exposure_eid_index.npy", unique_eids.astype(np.int64))
show_progress: bool = True, np.save(output_dir / "exposure_eid_start.npy", eid_start)
) -> 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( def build_exposure_cache(
*, *,
exposure_dir: str | Path, exposure_dir: str | Path,
output_dir: str | Path, output_dir: str | Path,
data_prefix: str = "ukb",
summary_file: str = "summary.csv", summary_file: str = "summary.csv",
overwrite: bool = False, overwrite: bool = False,
show_progress: bool = True, show_progress: bool = True,
@@ -436,25 +214,20 @@ def build_exposure_cache(
exposure_dir = Path(exposure_dir) exposure_dir = Path(exposure_dir)
output_dir = Path(output_dir) output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)
keys_path = output_dir / "exposure_keys.npy"
eid_path = output_dir / "exposure_eid.npy" output_paths = [
token_path = output_dir / "exposure_token.npy" output_dir / "exposure_eid.npy",
onset_date_path = output_dir / "exposure_onset_date.npy" output_dir / "exposure_token.npy",
daily_path = output_dir / "exposure_daily.npy" output_dir / "exposure_age_days.npy",
monthly_path = output_dir / "exposure_monthly.npy" output_dir / "exposure_onset_date.npy",
quality_path = output_dir / "exposure_quality.npy" output_dir / "exposure_eid_index.npy",
manifest_path = output_dir / "exposure_manifest.json" output_dir / "exposure_eid_start.npy",
outputs = [ output_dir / "exposure_daily.npy",
keys_path, output_dir / "exposure_monthly.npy",
eid_path, output_dir / "exposure_quality.npy",
token_path, output_dir / "exposure_manifest.json",
onset_date_path,
daily_path,
monthly_path,
quality_path,
manifest_path,
] ]
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( raise FileExistsError(
f"{output_dir} already contains exposure cache files; pass --overwrite" f"{output_dir} already contains exposure cache files; pass --overwrite"
) )
@@ -464,16 +237,29 @@ def build_exposure_cache(
summary_file, summary_file,
show_progress=show_progress, show_progress=show_progress,
) )
n_rows = int(summary["n_rows"].sum()) sequence_rows = _load_sequence_rows(data_prefix)
keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,)) n_rows = len(sequence_rows)
eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,)) if n_rows == 0:
tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,)) raise ValueError(f"{data_prefix}_event_data.npy contains no real disease events")
onset_dates_mm = np.lib.format.open_memmap(
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, onset_date_path,
mode="w+", sequence_rows["onset_date"].to_numpy(dtype="datetime64[D]"),
dtype="datetime64[D]",
shape=(n_rows,),
) )
_write_eid_offsets(sequence_rows, output_dir)
daily_mm = np.lib.format.open_memmap( daily_mm = np.lib.format.open_memmap(
daily_path, daily_path,
mode="w+", mode="w+",
@@ -492,25 +278,28 @@ def build_exposure_cache(
dtype=np.float32, dtype=np.float32,
shape=(n_rows, len(QUALITY_COLUMNS)), shape=(n_rows, len(QUALITY_COLUMNS)),
) )
daily_mm[:] = np.nan
monthly_mm[:] = np.nan
quality_mm[:] = np.nan
daily_cols = _daily_columns() daily_cols = _daily_columns()
monthly_cols = _monthly_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), summary.itertuples(index=False),
total=len(summary), total=len(summary),
desc="Materializing dense exposure cache", desc="Writing eid-sequence exposure cache",
unit="file", unit="file",
disable=not show_progress, disable=not show_progress,
) )
for row in rows: for row in iterator:
daily_file = Path(row.daily_path) daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_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 = [ daily_read_cols = [
"eid", "eid",
@@ -528,70 +317,76 @@ def build_exposure_cache(
] ]
daily_df = _read_parquet_columns(daily_file, daily_read_cols) daily_df = _read_parquet_columns(daily_file, daily_read_cols)
monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols) monthly_df = _read_parquet_columns(monthly_file, monthly_read_cols)
if len(daily_df) != len(monthly_df): if len(daily_df) != len(monthly_df):
raise ValueError( raise ValueError(
f"Daily/monthly row count mismatch for {row.label_code}: " f"Daily/monthly row count mismatch for {row.label_code}: "
f"{len(daily_df)} vs {len(monthly_df)}" 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( monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]]) pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
).reset_index() ).reset_index()
n = len(daily_df) tokens = daily_df["token"].dropna().astype(np.int64).unique()
end = offset + n wanted = pd.concat(
if end > n_rows: [wanted_by_token[int(token)] for token in tokens if int(token) in wanted_by_token],
raise RuntimeError("Exposure cache row count exceeded preallocated size") ignore_index=True,
) if len(tokens) else pd.DataFrame()
if wanted.empty:
continue
keys[offset:end] = encode_exposure_key( matches = daily_df[["eid", "onset_date", "token", "_source_row"]].merge(
daily_df["eid"].to_numpy(dtype=np.int64), wanted[["eid", "onset_date", "token", "position"]],
daily_df["token"].to_numpy(dtype=np.int64), on=["eid", "onset_date", "token"],
how="inner",
sort=False,
) )
eids_mm[offset:end] = daily_df["eid"].to_numpy(dtype=np.int64) if matches.empty:
tokens_mm[offset:end] = daily_df["token"].to_numpy(dtype=np.int32) continue
onset_dates_mm[offset:end] = pd.to_datetime(
daily_df["onset_date"], source_rows = matches["_source_row"].to_numpy(dtype=np.int64)
errors="coerce", positions = matches["position"].to_numpy(dtype=np.int64)
).to_numpy(dtype="datetime64[D]") daily_mm[positions] = _reshape_window(
daily_mm[offset:end] = _reshape_window( daily_df.iloc[source_rows],
daily_df,
daily_cols, daily_cols,
DAILY_LENGTH, DAILY_LENGTH,
len(DAILY_CHANNELS), len(DAILY_CHANNELS),
) )
monthly_mm[offset:end] = _reshape_window( monthly_mm[positions] = _reshape_window(
monthly_df, monthly_df.iloc[source_rows],
monthly_cols, monthly_cols,
MONTHLY_LENGTH, MONTHLY_LENGTH,
len(MONTHLY_CHANNELS), len(MONTHLY_CHANNELS),
) )
quality_mm[offset:end, 0] = daily_df.get("n_days_nonmissing", np.nan) quality_mm[positions, 0] = daily_df.iloc[source_rows].get("n_days_nonmissing", np.nan)
quality_mm[offset:end, 1] = daily_df.get("n_rh_days_nonmissing", np.nan) quality_mm[positions, 1] = daily_df.iloc[source_rows].get("n_rh_days_nonmissing", np.nan)
quality_mm[offset:end, 2] = monthly_df.get("n_months_nonmissing", np.nan) quality_mm[positions, 2] = monthly_df.iloc[source_rows].get("n_months_nonmissing", np.nan)
quality_mm[offset:end, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan) quality_mm[positions, 3] = monthly_df.iloc[source_rows].get("n_rh_months_nonmissing", np.nan)
offset = end matched[positions] = True
if offset != n_rows:
keys.flush()
eids_mm.flush()
tokens_mm.flush()
onset_dates_mm.flush()
daily_mm.flush() daily_mm.flush()
monthly_mm.flush() monthly_mm.flush()
quality_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 = { manifest = {
"storage": "dense_npy", "storage": "eid_sequence_npy",
"source_dir": str(exposure_dir), "source_dir": str(exposure_dir.resolve()),
"data_prefix": data_prefix,
"n_rows": int(n_rows), "n_rows": int(n_rows),
"legacy_key": "(eid << 16) | raw_token", "matched_rows": int(matched.sum()),
"alignment_key": "(eid, raw_token, onset_date - date_of_birth)", "missing_rows": int((~matched).sum()),
"alignment_key": "(eid, raw_token, date_of_birth + age_days)",
"requires_basic_info_column": "date_of_birth", "requires_basic_info_column": "date_of_birth",
"daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)], "daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
"daily_channels": list(DAILY_CHANNELS), "daily_channels": list(DAILY_CHANNELS),
@@ -608,25 +403,8 @@ def main() -> None:
parser = argparse.ArgumentParser(description=__doc__) parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--exposure-dir", required=True) parser.add_argument("--exposure-dir", required=True)
parser.add_argument("--output-dir", default="ukb_exposure_cache") 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("--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( parser.add_argument(
"--no-progress", "--no-progress",
action="store_true", action="store_true",
@@ -634,26 +412,16 @@ def main() -> None:
) )
parser.add_argument("--overwrite", action="store_true") parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args() 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( n_rows = build_exposure_cache(
exposure_dir=args.exposure_dir, exposure_dir=args.exposure_dir,
output_dir=args.output_dir, output_dir=args.output_dir,
data_prefix=args.data_prefix,
summary_file=args.summary_file, summary_file=args.summary_file,
overwrite=args.overwrite, overwrite=args.overwrite,
show_progress=show_progress, show_progress=not args.no_progress,
) )
print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}") print(f"Wrote {n_rows:,} eid-sequence-aligned exposure rows to {args.output_dir}")
if __name__ == "__main__": if __name__ == "__main__":

19
train_next_step.py
View File

@@ -102,8 +102,8 @@ class ExposureLocalityBatchSampler(Sampler[List[int]]):
exposure_cache = getattr(dataset, "exposure_cache", None) exposure_cache = getattr(dataset, "exposure_cache", None)
if exposure_index is None or exposure_cache is None: if exposure_index is None or exposure_cache is None:
return (2**31 - 1, 2**31 - 1, raw_idx) return (2**31 - 1, 2**31 - 1, raw_idx)
file_id, row_group = exposure_cache.locality_key(exposure_index) block_id, block_offset = exposure_cache.locality_key(exposure_index)
return (file_id, row_group, raw_idx) return (block_id, block_offset, raw_idx)
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
@@ -136,16 +136,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--exposure_conv_kernel_size", type=int, default=7) parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0) parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
parser.add_argument("--no_exposure_gate", action="store_true") parser.add_argument("--no_exposure_gate", action="store_true")
parser.add_argument(
"--exposure_row_group_cache_size",
type=int,
default=4,
help=(
"Number of parquet exposure row groups cached per DataLoader worker "
"when using indexed exposure storage."
),
)
parser.add_argument("--target_mode", type=str, default="uts", parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"]) choices=["delphi2m", "uts"])
parser.add_argument("--readout_name", type=str, default=None, parser.add_argument("--readout_name", type=str, default=None,
@@ -195,8 +185,6 @@ def parse_args() -> argparse.Namespace:
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
if args.num_workers > 0 and args.prefetch_factor <= 0: if args.num_workers > 0 and args.prefetch_factor <= 0:
raise ValueError("prefetch_factor must be positive when num_workers > 0") raise ValueError("prefetch_factor must be positive when num_workers > 0")
if args.exposure_row_group_cache_size < 0:
raise ValueError("exposure_row_group_cache_size must be non-negative")
if args.exposure_locality_buffer_size < 0: if args.exposure_locality_buffer_size < 0:
raise ValueError("exposure_locality_buffer_size must be non-negative") raise ValueError("exposure_locality_buffer_size must be non-negative")
if args.target_mode == "uts": if args.target_mode == "uts":
@@ -447,7 +435,6 @@ def build_metadata(
"exposure_conv_kernel_size": int(args.exposure_conv_kernel_size), "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
"exposure_mlp_ratio": float(args.exposure_mlp_ratio), "exposure_mlp_ratio": float(args.exposure_mlp_ratio),
"exposure_use_gate": not bool(args.no_exposure_gate), "exposure_use_gate": not bool(args.no_exposure_gate),
"exposure_row_group_cache_size": int(args.exposure_row_group_cache_size),
"num_workers": int(args.num_workers), "num_workers": int(args.num_workers),
"prefetch_factor": int(args.prefetch_factor), "prefetch_factor": int(args.prefetch_factor),
"exposure_locality_buffer_size": int(args.exposure_locality_buffer_size), "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
@@ -485,7 +472,6 @@ def main() -> None:
"DataLoader IO: " "DataLoader IO: "
f"num_workers={args.num_workers}, " f"num_workers={args.num_workers}, "
f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, " f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, "
f"exposure_row_group_cache_size={args.exposure_row_group_cache_size}, "
f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}" f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}"
) )
@@ -496,7 +482,6 @@ def main() -> None:
include_no_event_in_uts_target=args.include_no_event_in_uts_target, include_no_event_in_uts_target=args.include_no_event_in_uts_target,
exposure_cache_dir=args.exposure_cache_dir, exposure_cache_dir=args.exposure_cache_dir,
mask_onset_exposure=args.mask_onset_exposure, mask_onset_exposure=args.mask_onset_exposure,
exposure_row_group_cache_size=args.exposure_row_group_cache_size,
) )
if args.train_eid_file and args.val_eid_file and args.test_eid_file: if args.train_eid_file and args.val_eid_file and args.test_eid_file:
train_subset, val_subset, test_subset = split_dataset_by_eid_files( train_subset, val_subset, test_subset = split_dataset_by_eid_files(