DeepHealthExpo/prepare_exposure_cache.py

"""Build an eid-sequence-aligned exposure cache for DeepHealth training.

The source exposure dataset is stored as one daily and one monthly parquet file
per disease. That layout is inconvenient for mini-batch training because the
model consumes per-participant disease sequences. This script materializes one
large numpy cache ordered exactly like ``{data_prefix}_event_data.npy`` after
sorting by ``eid, age_days, token``.

The output directory contains:

    exposure_eid.npy          int64 eid per real disease event
    exposure_token.npy        int32 raw disease token per real disease event
    exposure_age_days.npy     int32 age in days per real disease event
    exposure_onset_date.npy   datetime64[D] onset date per real disease event
    exposure_row_index.npy    int64 window row per real disease event, -1 if missing
    exposure_eid_index.npy    int64 unique eids in cache order
    exposure_eid_start.npy    int64 start offsets, length len(eid_index) + 1
    exposure_daily.npy        float32 memmap, capacity (N, 1826, 4);
                              first M rows are sequential matched windows
                              channels: tmean, tmax, tmin, rhmean
    exposure_monthly.npy      float32 memmap, capacity (N, 241, 2);
                              first M rows are sequential matched windows
                              channels: tmean, rhmean
    exposure_quality.npy      float32 memmap, capacity (N, 4);
                              first M rows are matched-window quality stats
                              n_days, n_rh_days, n_months, n_rh_months
    exposure_manifest.json    metadata

Rows without matching exposure parquet records are kept as NaN windows. 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.
"""

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

import numpy as np
import pandas as pd

try:
    from tqdm.auto import tqdm
except ImportError:
    def tqdm(iterable=None, **kwargs):
        return iterable if iterable is not None else range(kwargs.get("total", 0))


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 _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]:
    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_matching_parquet_rows(
    path: Path,
    columns: list[str],
    wanted: pd.DataFrame,
) -> pd.DataFrame:
    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)
    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:
    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 _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,
    *,
    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))
    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}")
    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:
        raise ValueError(f"event_data must have shape (N, 3+), got {event_data.shape}")
    event_data = event_data[:, :3].copy()
    order = np.lexsort((event_data[:, 2], event_data[:, 1], event_data[:, 0]))
    event_data = event_data[order]

    basic_table = pd.read_csv(f"{data_prefix}_basic_info.csv", index_col=0)
    basic_table.index = basic_table.index.astype(np.int64)
    if "date_of_birth" not in basic_table.columns:
        raise ValueError(
            f"{data_prefix}_basic_info.csv must contain date_of_birth for exposure alignment"
        )

    rows = pd.DataFrame(
        {
            "eid": event_data[:, 0].astype(np.int64),
            "age_days": np.rint(event_data[:, 1].astype(np.float64)).astype(np.int32),
            "token": event_data[:, 2].astype(np.int32),
        }
    )
    rows = rows[rows["token"] > 1].reset_index(drop=True)
    rows["position"] = np.arange(len(rows), dtype=np.int64)

    birth = pd.to_datetime(
        basic_table.loc[rows["eid"].to_numpy(), "date_of_birth"].to_numpy(),
        errors="coerce",
    )
    if birth.isna().any():
        raise ValueError("date_of_birth contains missing or invalid values")
    rows["onset_date"] = (
        birth.to_numpy(dtype="datetime64[D]")
        + rows["age_days"].to_numpy(dtype="timedelta64[D]")
    )
    rows["onset_date"] = pd.to_datetime(rows["onset_date"]).dt.normalize()
    return rows


def _write_eid_offsets(rows: pd.DataFrame, output_dir: Path) -> None:
    eids = rows["eid"].to_numpy(dtype=np.int64)
    unique_eids, starts = np.unique(eids, return_index=True)
    starts = starts.astype(np.int64)
    ends = np.concatenate([starts[1:], np.asarray([len(rows)], dtype=np.int64)])
    eid_start = np.concatenate([starts, ends[-1:]]).astype(np.int64)
    np.save(output_dir / "exposure_eid_index.npy", unique_eids.astype(np.int64))
    np.save(output_dir / "exposure_eid_start.npy", eid_start)


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",
    workers: int = 1,
    max_in_flight: int = 0,
    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)

    output_paths = [
        output_dir / "exposure_eid.npy",
        output_dir / "exposure_token.npy",
        output_dir / "exposure_age_days.npy",
        output_dir / "exposure_onset_date.npy",
        output_dir / "exposure_row_index.npy",
        output_dir / "exposure_eid_index.npy",
        output_dir / "exposure_eid_start.npy",
        output_dir / "exposure_daily.npy",
        output_dir / "exposure_monthly.npy",
        output_dir / "exposure_quality.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 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,
    )
    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"
    age_path = output_dir / "exposure_age_days.npy"
    onset_date_path = output_dir / "exposure_onset_date.npy"
    row_index_path = output_dir / "exposure_row_index.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"

    np.save(eid_path, sequence_rows["eid"].to_numpy(dtype=np.int64))
    np.save(token_path, sequence_rows["token"].to_numpy(dtype=np.int32))
    np.save(age_path, sequence_rows["age_days"].to_numpy(dtype=np.int32))
    np.save(
        onset_date_path,
        sequence_rows["onset_date"].to_numpy(dtype="datetime64[D]"),
    )
    _write_eid_offsets(sequence_rows, output_dir)
    row_index_mm = np.lib.format.open_memmap(
        row_index_path,
        mode="w+",
        dtype=np.int64,
        shape=(n_rows,),
    )
    row_index_mm[:] = -1

    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()
    wanted_by_token = {
        int(token): frame.reset_index(drop=True)
        for token, frame in sequence_rows.groupby("token", sort=False)
    }
    write_offset = 0

    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
        tasks.append(
            {
                "daily_path": str(row.daily_path),
                "monthly_path": str(row.monthly_path),
                "wanted": wanted,
                "daily_cols": daily_cols,
                "monthly_cols": monthly_cols,
            }
        )

    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] = 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()
    quality_mm.flush()

    manifest = {
        "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),
        "window_capacity_rows": int(n_rows),
        "matched_rows": int(write_offset),
        "missing_rows": int(n_rows - write_offset),
        "alignment_key": "(eid, raw_token, date_of_birth + age_days)",
        "requires_basic_info_column": "date_of_birth",
        "daily_shape": [int(n_rows), DAILY_LENGTH, len(DAILY_CHANNELS)],
        "active_daily_shape": [int(write_offset), DAILY_LENGTH, len(DAILY_CHANNELS)],
        "daily_channels": list(DAILY_CHANNELS),
        "monthly_shape": [int(n_rows), MONTHLY_LENGTH, len(MONTHLY_CHANNELS)],
        "active_monthly_shape": [int(write_offset), 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("--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",
        help="Disable tqdm progress bars.",
    )
    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,
        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,
    )
    print(f"Wrote {n_rows:,} eid-sequence-aligned exposure rows to {args.output_dir}")


if __name__ == "__main__":
    main()