diff --git a/dataset.py b/dataset.py index fb6fe29..f8a33e1 100644 --- a/dataset.py +++ b/dataset.py @@ -1,6 +1,8 @@ # dataset.py from __future__ import annotations +import json +from collections import OrderedDict from pathlib import Path from typing import Dict, Iterable, List, Literal, Optional, Tuple @@ -22,14 +24,41 @@ from targets import ( ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR DAILY_EXPOSURE_SHAPE = (1826, 4) MONTHLY_EXPOSURE_SHAPE = (241, 2) +DAILY_EXPOSURE_CHANNELS = ("tmean", "tmax", "tmin", "rhmean") +MONTHLY_EXPOSURE_CHANNELS = ("tmean", "rhmean") + + +def _daily_exposure_columns() -> list[str]: + cols: list[str] = [] + for name in DAILY_EXPOSURE_CHANNELS: + cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_EXPOSURE_SHAPE[0])) + return cols + + +def _monthly_exposure_columns() -> list[str]: + cols: list[str] = [] + for name in MONTHLY_EXPOSURE_CHANNELS: + cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_EXPOSURE_SHAPE[0])) + return cols class ExposureCache: """Random-access view over files produced by prepare_exposure_cache.py.""" - def __init__(self, cache_dir: str | Path): + def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 16): cache_dir = Path(cache_dir) self.cache_dir = cache_dir + manifest_path = cache_dir / "exposure_manifest.json" + self.manifest = ( + json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest_path.is_file() + else {} + ) + self.storage = self.manifest.get("storage", "dense_npy") + self._row_group_cache_size = int(row_group_cache_size) + self._row_group_cache: OrderedDict[tuple[str, int, int], pd.DataFrame] = OrderedDict() + self._parquet_files: dict[tuple[str, int], object] = {} + 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" @@ -42,29 +71,61 @@ class ExposureCache: self.eids = np.load(eid_path, mmap_mode="r") self.raw_tokens = np.load(token_path, mmap_mode="r") self.onset_dates = np.load(onset_date_path, mmap_mode="r") - self.daily = np.load(cache_dir / "exposure_daily.npy", mmap_mode="r") - self.monthly = np.load(cache_dir / "exposure_monthly.npy", mmap_mode="r") + self.daily = None + self.monthly = None + if self.storage == "dense_npy": + self.daily = np.load(cache_dir / "exposure_daily.npy", 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" self.quality = np.load(quality_path, mmap_mode="r") if quality_path.is_file() else None - if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE: - raise ValueError( - f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, " - f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}" - ) - if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE: - raise ValueError( - f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, " - f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}" - ) + if self.storage == "dense_npy": + if self.daily.ndim != 3 or self.daily.shape[1:] != DAILY_EXPOSURE_SHAPE: + raise ValueError( + f"exposure_daily.npy must have shape (N, {DAILY_EXPOSURE_SHAPE[0]}, " + f"{DAILY_EXPOSURE_SHAPE[1]}), got {self.daily.shape}" + ) + if self.monthly.ndim != 3 or self.monthly.shape[1:] != MONTHLY_EXPOSURE_SHAPE: + raise ValueError( + f"exposure_monthly.npy must have shape (N, {MONTHLY_EXPOSURE_SHAPE[0]}, " + f"{MONTHLY_EXPOSURE_SHAPE[1]}), got {self.monthly.shape}" + ) n_rows = len(self.eids) - if ( - len(self.raw_tokens) != n_rows - or len(self.onset_dates) != n_rows - or self.daily.shape[0] != n_rows - or self.monthly.shape[0] != n_rows - ): + if len(self.raw_tokens) != n_rows or len(self.onset_dates) != n_rows: raise ValueError("Exposure cache metadata/daily/monthly row counts do not match") + if self.storage == "dense_npy": + if 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") + else: + 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 @@ -100,12 +161,83 @@ class ExposureCache: def daily_window(self, index: int) -> np.ndarray: if index < 0: return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) - return np.asarray(self.daily[index], dtype=np.float32) + if self.storage == "dense_npy": + return np.asarray(self.daily[index], dtype=np.float32) + return self._parquet_window("daily", index) def monthly_window(self, index: int) -> np.ndarray: if index < 0: return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) - return np.asarray(self.monthly[index], dtype=np.float32) + if self.storage == "dense_npy": + return np.asarray(self.monthly[index], dtype=np.float32) + return self._parquet_window("monthly", index) + + 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( diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index cc36e4d..b7acedb 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -1,11 +1,31 @@ -"""Build a random-access exposure cache from disease-level parquet files. +"""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. -This script converts those parquet files into a compact directory: +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 @@ -91,6 +111,28 @@ def _read_parquet_columns(path: Path, columns: list[str]) -> pd.DataFrame: return pd.read_parquet(path, columns=columns) +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) @@ -102,6 +144,190 @@ def _count_rows(summary: pd.DataFrame) -> int: return int(sum(pd.read_parquet(path, columns=["eid"]).shape[0] for path in summary["daily_path"])) +def build_exposure_index( + *, + exposure_dir: str | Path, + output_dir: str | Path, + summary_file: str = "summary.csv", + overwrite: bool = False, +) -> int: + exposure_dir = Path(exposure_dir) + output_dir = Path(output_dir) + summary_path = exposure_dir / summary_file + if not summary_path.is_file(): + raise FileNotFoundError(f"summary.csv not found: {summary_path}") + + 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 = 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)) + + n_rows = _count_rows(summary) + 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,), + ) + + daily_files: list[str] = [] + monthly_files: list[str] = [] + offset = 0 + for file_id, row in enumerate(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_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 {row.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 {row.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 {row.label_code}" + ) + monthly_rg = monthly_rg_all[monthly_pos] + monthly_row = monthly_row_all[monthly_pos] + + end = offset + n + if end > n_rows: + raise RuntimeError("Exposure index row count exceeded preallocated size") + + 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_file_id_mm[offset:end] = file_id + daily_row_group_mm[offset:end] = daily_rg + daily_row_in_group_mm[offset:end] = daily_row + monthly_file_id_mm[offset:end] = file_id + monthly_row_group_mm[offset:end] = monthly_rg + monthly_row_in_group_mm[offset:end] = monthly_row + daily_files.append(str(daily_file.resolve())) + monthly_files.append(str(monthly_file.resolve())) + offset = end + + if offset != n_rows: + raise RuntimeError( + f"Expected {n_rows} rows from summary but indexed {offset}. " + "Regenerate summary.csv or remove n_cases before building." + ) + + 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": daily_files, + "monthly_files": monthly_files, + "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, @@ -264,6 +490,7 @@ def build_exposure_cache( ) manifest = { + "storage": "dense_npy", "source_dir": str(exposure_dir), "n_rows": int(n_rows), "legacy_key": "(eid << 16) | raw_token", @@ -285,15 +512,33 @@ def main() -> None: 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("--overwrite", action="store_true") args = parser.parse_args() - n_rows = build_exposure_cache( - exposure_dir=args.exposure_dir, - output_dir=args.output_dir, - summary_file=args.summary_file, - overwrite=args.overwrite, - ) - print(f"Wrote {n_rows:,} exposure rows to {args.output_dir}") + 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, + ) + 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, + ) + print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}") if __name__ == "__main__":