DeepHealthExpo/prepare_exposure_cache.py

"""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.

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
    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)
                            channels: tmean, tmax, tmin, rhmean
    exposure_monthly.npy    float32 memmap, shape (N, 241, 2)
                            channels: tmean, rhmean
    exposure_quality.npy    float32 memmap, shape (N, 4)
                            n_days, n_rh_days, n_months, n_rh_months
    exposure_manifest.json  metadata

The raw token convention follows the exposure README: padding=0, checkup=1,
and the first row of labels.csv is token=2. The model dataset inserts
<NO_EVENT> at token 2 and shifts real disease tokens by +1 internally; dataset
lookup converts back to these raw tokens before reading this cache. Dataset
alignment uses (eid, raw_token, onset_date - date_of_birth) so that raw
calendar dates in the exposure files match the age-day event times used by the
model.
"""

from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
import json
import os
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd
from tqdm.auto import tqdm


DAILY_LENGTH = 1826
MONTHLY_LENGTH = 241
DAILY_CHANNELS = ("tmean", "tmax", "tmin", "rhmean")
MONTHLY_CHANNELS = ("tmean", "rhmean")
QUALITY_COLUMNS = (
    "n_days_nonmissing",
    "n_rh_days_nonmissing",
    "n_months_nonmissing",
    "n_rh_months_nonmissing",
)


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]:
    cols: list[str] = []
    for name in DAILY_CHANNELS:
        cols.extend(f"{name}_d{idx:04d}" for idx in range(DAILY_LENGTH))
    return cols


def _monthly_columns() -> list[str]:
    cols: list[str] = []
    for name in MONTHLY_CHANNELS:
        cols.extend(f"{name}_m{idx:03d}" for idx in range(MONTHLY_LENGTH))
    return cols


def _safe_columns(path: Path, columns: Iterable[str]) -> list[str]:
    """Return the subset of requested columns present in a parquet file."""
    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

    schema_names = set(pq.ParquetFile(path).schema.names)
    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:
    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
    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:
    arr = df.reindex(columns=cols).to_numpy(dtype=np.float32, copy=True)
    return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1)


def _load_summary(
    exposure_dir: Path,
    summary_file: str,
    *,
    show_progress: bool,
) -> pd.DataFrame:
    summary_path = exposure_dir / summary_file
    if not summary_path.is_file():
        raise FileNotFoundError(f"summary.csv not found: {summary_path}")

    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))

    counts: list[int] = []
    iterator = summary.itertuples(index=False)
    iterator = tqdm(
        iterator,
        total=len(summary),
        desc="Counting exposure rows",
        unit="file",
        disable=not show_progress,
    )
    for row in iterator:
        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
    summary["offset"] = np.cumsum([0, *counts[:-1]], dtype=np.int64)
    return summary


def _process_index_file_pair(task: tuple[int, str, str, str]) -> dict:
    file_id, label_code, daily_path, monthly_path = task
    daily_file = Path(daily_path)
    monthly_file = Path(monthly_path)

    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 {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 {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 {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",
        ).to_numpy(dtype="datetime64[D]"),
        "daily_row_group": daily_rg,
        "daily_row_in_group": daily_row,
        "monthly_row_group": monthly_rg_all[monthly_pos],
        "monthly_row_in_group": monthly_row_all[monthly_pos],
    }


def build_exposure_index(
    *,
    exposure_dir: str | Path,
    output_dir: str | Path,
    summary_file: str = "summary.csv",
    overwrite: bool = False,
    workers: int = 1,
    show_progress: bool = True,
) -> 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(
    *,
    exposure_dir: str | Path,
    output_dir: str | Path,
    summary_file: str = "summary.csv",
    overwrite: bool = False,
    show_progress: bool = True,
) -> int:
    exposure_dir = Path(exposure_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    keys_path = output_dir / "exposure_keys.npy"
    eid_path = output_dir / "exposure_eid.npy"
    token_path = output_dir / "exposure_token.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"
    outputs = [
        keys_path,
        eid_path,
        token_path,
        onset_date_path,
        daily_path,
        monthly_path,
        quality_path,
        manifest_path,
    ]
    if any(path.exists() for path in outputs) and not overwrite:
        raise FileExistsError(
            f"{output_dir} already contains exposure cache files; pass --overwrite"
        )

    summary = _load_summary(
        exposure_dir,
        summary_file,
        show_progress=show_progress,
    )
    n_rows = int(summary["n_rows"].sum())
    keys = np.lib.format.open_memmap(keys_path, mode="w+", dtype=np.uint64, shape=(n_rows,))
    eids_mm = np.lib.format.open_memmap(eid_path, mode="w+", dtype=np.int64, shape=(n_rows,))
    tokens_mm = np.lib.format.open_memmap(token_path, mode="w+", dtype=np.int32, shape=(n_rows,))
    onset_dates_mm = np.lib.format.open_memmap(
        onset_date_path,
        mode="w+",
        dtype="datetime64[D]",
        shape=(n_rows,),
    )
    daily_mm = np.lib.format.open_memmap(
        daily_path,
        mode="w+",
        dtype=np.float32,
        shape=(n_rows, DAILY_LENGTH, len(DAILY_CHANNELS)),
    )
    monthly_mm = np.lib.format.open_memmap(
        monthly_path,
        mode="w+",
        dtype=np.float32,
        shape=(n_rows, MONTHLY_LENGTH, len(MONTHLY_CHANNELS)),
    )
    quality_mm = np.lib.format.open_memmap(
        quality_path,
        mode="w+",
        dtype=np.float32,
        shape=(n_rows, len(QUALITY_COLUMNS)),
    )

    daily_cols = _daily_columns()
    monthly_cols = _monthly_columns()
    offset = 0

    rows = tqdm(
        summary.itertuples(index=False),
        total=len(summary),
        desc="Materializing dense exposure cache",
        unit="file",
        disable=not show_progress,
    )
    for row in rows:
        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_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_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)}"
            )

        monthly_df = monthly_df.set_index(["eid", "onset_date", "token"]).reindex(
            pd.MultiIndex.from_frame(daily_df[["eid", "onset_date", "token"]])
        ).reset_index()

        n = len(daily_df)
        end = offset + n
        if end > n_rows:
            raise RuntimeError("Exposure cache row count exceeded preallocated size")

        keys[offset:end] = encode_exposure_key(
            daily_df["eid"].to_numpy(dtype=np.int64),
            daily_df["token"].to_numpy(dtype=np.int64),
        )
        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_mm[offset:end] = _reshape_window(
            daily_df,
            daily_cols,
            DAILY_LENGTH,
            len(DAILY_CHANNELS),
        )
        monthly_mm[offset:end] = _reshape_window(
            monthly_df,
            monthly_cols,
            MONTHLY_LENGTH,
            len(MONTHLY_CHANNELS),
        )
        quality_mm[offset:end, 0] = daily_df.get("n_days_nonmissing", np.nan)
        quality_mm[offset:end, 1] = daily_df.get("n_rh_days_nonmissing", np.nan)
        quality_mm[offset:end, 2] = monthly_df.get("n_months_nonmissing", np.nan)
        quality_mm[offset:end, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan)
        offset = end

    if offset != n_rows:
        keys.flush()
        eids_mm.flush()
        tokens_mm.flush()
        onset_dates_mm.flush()
        daily_mm.flush()
        monthly_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 = {
        "storage": "dense_npy",
        "source_dir": str(exposure_dir),
        "n_rows": int(n_rows),
        "legacy_key": "(eid << 16) | raw_token",
        "alignment_key": "(eid, raw_token, onset_date - date_of_birth)",
        "requires_basic_info_column": "date_of_birth",
        "daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
        "daily_channels": list(DAILY_CHANNELS),
        "monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
        "monthly_channels": list(MONTHLY_CHANNELS),
        "quality_columns": list(QUALITY_COLUMNS),
        "raw_token_convention": "padding=0, checkup=1, labels.csv first row token=2",
    }
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    return int(n_rows)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    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(
        "--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(
        "--no-progress",
        action="store_true",
        help="Disable tqdm progress bars.",
    )
    parser.add_argument("--overwrite", action="store_true")
    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(
            exposure_dir=args.exposure_dir,
            output_dir=args.output_dir,
            summary_file=args.summary_file,
            overwrite=args.overwrite,
            show_progress=show_progress,
        )
        print(f"Wrote {n_rows:,} dense exposure rows to {args.output_dir}")


if __name__ == "__main__":
    main()