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

Parallelize exposure cache reads

Browse Source
This commit is contained in:
Jiarui Li
2026-07-08 16:34:28 +08:00
parent efba7ac306
commit a35f1117c0
1 changed files with 159 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

213
prepare_exposure_cache.py
View File

@@ -36,7 +36,9 @@ converts back to these raw tokens before reading this cache.
from __future__ import annotations
import argparse
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
import json
import os
from pathlib import Path
from typing import Iterable
@@ -162,6 +164,75 @@ def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels:
return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)
def _quality_column(df: pd.DataFrame, name: str, n_rows: int) -> np.ndarray:
if name not in df:
return np.full(n_rows, np.nan, dtype=np.float32)
return df[name].to_numpy(dtype=np.float32, copy=True)
def _process_exposure_task(task: dict) -> dict:
daily_file = Path(task["daily_path"])
monthly_file = Path(task["monthly_path"])
wanted = task["wanted"]
daily_cols = task["daily_cols"]
monthly_cols = task["monthly_cols"]
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_matching_parquet_rows(daily_file, daily_read_cols, wanted)
monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted)
if daily_df.empty:
return {"positions": np.empty(0, dtype=np.int64)}
common_positions = np.intersect1d(
daily_df["position"].to_numpy(dtype=np.int64),
monthly_df["position"].to_numpy(dtype=np.int64),
)
if len(common_positions) == 0:
return {"positions": np.empty(0, dtype=np.int64)}
daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
n_match = len(common_positions)
quality = np.stack(
[
_quality_column(daily_df, "n_days_nonmissing", n_match),
_quality_column(daily_df, "n_rh_days_nonmissing", n_match),
_quality_column(monthly_df, "n_months_nonmissing", n_match),
_quality_column(monthly_df, "n_rh_months_nonmissing", n_match),
],
axis=1,
)
return {
"positions": common_positions.astype(np.int64),
"daily": _reshape_window(
daily_df,
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
),
"monthly": _reshape_window(
monthly_df,
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
),
"quality": quality,
}
def _load_summary(
exposure_dir: Path,
summary_file: str,
@@ -257,6 +328,8 @@ def build_exposure_cache(
data_prefix: str = "ukb",
labels_file: str | Path = "labels.csv",
summary_file: str = "summary.csv",
workers: int = 1,
max_in_flight: int = 0,
overwrite: bool = False,
show_progress: bool = True,
) -> int:
@@ -359,71 +432,83 @@ def build_exposure_cache(
}
write_offset = 0
iterator = tqdm(
summary.itertuples(index=False),
total=len(summary),
desc="Writing eid-sequence exposure cache",
unit="file",
disable=not show_progress,
)
for row in iterator:
daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_path)
tasks: list[dict] = []
for row in summary.itertuples(index=False):
token = int(row.raw_token)
wanted = wanted_by_token.get(token)
if wanted is None or wanted.empty:
continue
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_matching_parquet_rows(daily_file, daily_read_cols, wanted)
monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted)
if daily_df.empty:
continue
common_positions = np.intersect1d(
daily_df["position"].to_numpy(dtype=np.int64),
monthly_df["position"].to_numpy(dtype=np.int64),
tasks.append(
{
"daily_path": str(row.daily_path),
"monthly_path": str(row.monthly_path),
"wanted": wanted,
"daily_cols": daily_cols,
"monthly_cols": monthly_cols,
}
)
if len(common_positions) == 0:
continue
daily_df = daily_df.set_index("position").loc[common_positions].reset_index()
monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index()
positions = common_positions.astype(np.int64)
workers = max(1, int(workers))
max_in_flight = int(max_in_flight)
def write_result(result: dict) -> None:
nonlocal write_offset
positions = result["positions"]
if len(positions) == 0:
return
n_match = len(positions)
end_offset = write_offset + n_match
daily_mm[write_offset:end_offset] = _reshape_window(
daily_df,
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
)
monthly_mm[write_offset:end_offset] = _reshape_window(
monthly_df,
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
)
quality_mm[write_offset:end_offset, 0] = daily_df.get("n_days_nonmissing", np.nan)
quality_mm[write_offset:end_offset, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
quality_mm[write_offset:end_offset, 2] = monthly_df.get("n_months_nonmissing", np.nan)
quality_mm[write_offset:end_offset, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
daily_mm[write_offset:end_offset] = result["daily"]
monthly_mm[write_offset:end_offset] = result["monthly"]
quality_mm[write_offset:end_offset] = result["quality"]
row_index_mm[positions] = np.arange(write_offset, end_offset, dtype=np.int64)
write_offset = end_offset
if workers == 1:
iterator = tqdm(
map(_process_exposure_task, tasks),
total=len(tasks),
desc="Reading exposure parquet and writing cache",
unit="file",
disable=not show_progress,
)
for result in iterator:
write_result(result)
else:
with ProcessPoolExecutor(max_workers=workers) as executor:
task_iter = iter(tasks)
iterator = tqdm(
total=len(tasks),
desc=f"Reading exposure parquet ({workers} workers)",
unit="file",
disable=not show_progress,
)
if max_in_flight <= 0:
in_flight = {
executor.submit(_process_exposure_task, task)
for task in task_iter
}
while in_flight:
done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED)
for future in done:
write_result(future.result())
iterator.update(1)
else:
in_flight = {
executor.submit(_process_exposure_task, task)
for task in [next(task_iter, None) for _ in range(max_in_flight)]
if task is not None
}
while in_flight:
done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED)
for future in done:
write_result(future.result())
iterator.update(1)
task = next(task_iter, None)
if task is not None:
in_flight.add(executor.submit(_process_exposure_task, task))
iterator.close()
row_index_mm.flush()
daily_mm.flush()
monthly_mm.flush()
@@ -460,6 +545,24 @@ def main() -> None:
parser.add_argument("--data-prefix", default="ukb")
parser.add_argument("--labels-file", default="labels.csv")
parser.add_argument("--summary-file", default="summary.csv")
parser.add_argument(
"--workers",
type=int,
default=max(1, os.cpu_count() or 1),
help=(
"Worker processes for parquet reading and window extraction. "
"The main process remains the only writer to the output memmaps."
),
)
parser.add_argument(
"--max-in-flight",
type=int,
default=0,
help=(
"Maximum submitted parquet tasks waiting/running at once. "
"Use 0 to submit all tasks, which is the default for high-memory servers."
),
)
parser.add_argument(
"--no-progress",
action="store_true",
@@ -474,6 +577,8 @@ def main() -> None:
data_prefix=args.data_prefix,
labels_file=args.labels_file,
summary_file=args.summary_file,
workers=args.workers,
max_in_flight=args.max_in_flight,
overwrite=args.overwrite,
show_progress=not args.no_progress,
)