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

Optimize eid exposure cache build

Browse Source
This commit is contained in:
Jiarui Li
2026-07-08 12:50:39 +08:00
parent 46d530fad0
commit 14eccc91f9
1 changed files with 112 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

185
prepare_exposure_cache.py
View File

@@ -85,11 +85,11 @@ def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
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:
def _read_matching_parquet_rows(
path: Path,
columns: list[str],
wanted: pd.DataFrame,
) -> pd.DataFrame:
try:
import pyarrow.parquet as pq
except ImportError as exc:
@@ -97,7 +97,60 @@ def _parquet_row_count(path: Path) -> int:
"prepare_exposure_cache.py requires pyarrow. Install requirements "
"or run `pip install pyarrow`."
) from exc
return int(pq.ParquetFile(path).metadata.num_rows)
parquet_file = pq.ParquetFile(path)
available = [col for col in columns if col in set(parquet_file.schema.names)]
key_cols = ["eid", "onset_date", "token"]
missing_keys = [col for col in key_cols if col not in available]
if missing_keys:
raise ValueError(f"{path} is missing key columns: {missing_keys}")
wanted_keys = wanted[["eid", "onset_date", "token", "position"]].copy()
wanted_keys["eid"] = wanted_keys["eid"].astype(np.int64)
wanted_keys["token"] = wanted_keys["token"].astype(np.int64)
wanted_keys["onset_date"] = pd.to_datetime(
wanted_keys["onset_date"],
errors="coerce",
).dt.normalize()
chunks: list[pd.DataFrame] = []
for row_group_idx in range(parquet_file.num_row_groups):
key_frame = parquet_file.read_row_group(
row_group_idx,
columns=key_cols,
).to_pandas()
if key_frame.empty:
continue
key_frame = key_frame.copy()
key_frame["_row_in_group"] = np.arange(len(key_frame), dtype=np.int64)
key_frame["eid"] = key_frame["eid"].astype(np.int64)
key_frame["token"] = key_frame["token"].astype(np.int64)
key_frame["onset_date"] = pd.to_datetime(
key_frame["onset_date"],
errors="coerce",
).dt.normalize()
matches = key_frame.merge(
wanted_keys,
on=key_cols,
how="inner",
sort=False,
)
if matches.empty:
continue
row_values = parquet_file.read_row_group(
row_group_idx,
columns=available,
).to_pandas()
selected = row_values.iloc[
matches["_row_in_group"].to_numpy(dtype=np.int64)
].copy()
selected["position"] = matches["position"].to_numpy(dtype=np.int64)
chunks.append(selected)
if not chunks:
return pd.DataFrame(columns=[*columns, "position"])
return pd.concat(chunks, ignore_index=True)
def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: int) -> np.ndarray:
@@ -124,35 +177,26 @@ def _load_summary(
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 = tqdm(
summary.itertuples(index=False),
total=len(summary),
desc="Counting exposure rows",
unit="file",
disable=not show_progress,
)
for row in iterator:
for row in summary.itertuples(index=False):
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
return summary
def _load_label_token_map(labels_file: str | Path) -> dict[str, int]:
out: dict[str, int] = {}
with Path(labels_file).open(encoding="utf-8") as handle:
for idx, line in enumerate(handle):
parts = line.strip().split(" ")
if parts and parts[0]:
out[parts[0]] = idx + 2
return out
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:
@@ -207,6 +251,7 @@ def build_exposure_cache(
exposure_dir: str | Path,
output_dir: str | Path,
data_prefix: str = "ukb",
labels_file: str | Path = "labels.csv",
summary_file: str = "summary.csv",
overwrite: bool = False,
show_progress: bool = True,
@@ -232,15 +277,29 @@ def build_exposure_cache(
f"{output_dir} already contains exposure cache files; pass --overwrite"
)
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")
label_token_map = _load_label_token_map(labels_file)
summary = _load_summary(
exposure_dir,
summary_file,
show_progress=show_progress,
)
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")
summary["raw_token"] = summary["label_code"].map(label_token_map)
needed_tokens = set(sequence_rows["token"].astype(np.int64).unique().tolist())
summary = summary[
summary["raw_token"].notna()
& summary["raw_token"].astype(np.int64).isin(needed_tokens)
].copy()
summary["raw_token"] = summary["raw_token"].astype(np.int64)
if summary.empty:
raise ValueError(
"No exposure summary rows match disease tokens in "
f"{data_prefix}_event_data.npy. Check --summary-file and --labels-file."
)
eid_path = output_dir / "exposure_eid.npy"
token_path = output_dir / "exposure_token.npy"
@@ -300,6 +359,10 @@ def build_exposure_cache(
for row in iterator:
daily_file = Path(row.daily_path)
monthly_file = Path(row.monthly_path)
token = int(row.raw_token)
wanted = wanted_by_token.get(token)
if wanted is None or wanted.empty:
continue
daily_read_cols = [
"eid",
@@ -315,64 +378,37 @@ def build_exposure_cache(
*_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)}"
)
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()
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:
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
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,
common_positions = np.intersect1d(
daily_df["position"].to_numpy(dtype=np.int64),
monthly_df["position"].to_numpy(dtype=np.int64),
)
if matches.empty:
if len(common_positions) == 0:
continue
source_rows = matches["_source_row"].to_numpy(dtype=np.int64)
positions = matches["position"].to_numpy(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()
positions = common_positions.astype(np.int64)
daily_mm[positions] = _reshape_window(
daily_df.iloc[source_rows],
daily_df,
daily_cols,
DAILY_LENGTH,
len(DAILY_CHANNELS),
)
monthly_mm[positions] = _reshape_window(
monthly_df.iloc[source_rows],
monthly_df,
monthly_cols,
MONTHLY_LENGTH,
len(MONTHLY_CHANNELS),
)
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)
quality_mm[positions, 0] = daily_df.get("n_days_nonmissing", np.nan)
quality_mm[positions, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
quality_mm[positions, 2] = monthly_df.get("n_months_nonmissing", np.nan)
quality_mm[positions, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
matched[positions] = True
daily_mm.flush()
@@ -383,6 +419,7 @@ def build_exposure_cache(
"storage": "eid_sequence_npy",
"source_dir": str(exposure_dir.resolve()),
"data_prefix": data_prefix,
"labels_file": str(Path(labels_file).resolve()),
"n_rows": int(n_rows),
"matched_rows": int(matched.sum()),
"missing_rows": int((~matched).sum()),
@@ -404,6 +441,7 @@ def main() -> None:
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("--labels-file", default="labels.csv")
parser.add_argument("--summary-file", default="summary.csv")
parser.add_argument(
"--no-progress",
@@ -417,6 +455,7 @@ def main() -> None:
exposure_dir=args.exposure_dir,
output_dir=args.output_dir,
data_prefix=args.data_prefix,
labels_file=args.labels_file,
summary_file=args.summary_file,
overwrite=args.overwrite,
show_progress=not args.no_progress,