diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index 345969a..8614cc1 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -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,