798 lines
28 KiB
Python
798 lines
28 KiB
Python
"""Export individual Weibull parameters to one HDF5 file.
|
|
|
|
The test population is read from ``test_eid_file`` in the specified run's
|
|
``train_config.json``. Each eligible patient is queried at ages
|
|
40, 42, ..., 80 using only disease history observed by that age.
|
|
|
|
The model parameterization is
|
|
|
|
S(t) = exp(-rate * t ** shape)
|
|
|
|
and the standard Weibull parameters exported by this script are
|
|
|
|
shape = rho
|
|
scale = rate ** (-1 / shape)
|
|
|
|
The output is one HDF5 file. It contains the token table, age summary and one
|
|
group per landmark age. Each age group stores aligned patient metadata plus
|
|
chunked, compressed ``shape`` and ``scale`` matrices. Chunking is internal to
|
|
HDF5, so callers receive one file without loading the full multi-gigabyte
|
|
export into memory.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import contextlib
|
|
import importlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.utils.data import DataLoader
|
|
from tqdm.auto import tqdm
|
|
|
|
from dataset import (
|
|
DISEASE_HISTORY_MODE_TIMED,
|
|
NO_EVENT_IDX,
|
|
PAD_IDX,
|
|
RESERVED_IDX,
|
|
normalize_disease_history_mode,
|
|
)
|
|
from eval_data import (
|
|
build_model_from_dataset,
|
|
load_json_config,
|
|
load_sequence_eval_dataset,
|
|
resolve_eval_device,
|
|
select_indices_by_eid_file,
|
|
validate_dataset_metadata,
|
|
validate_training_mode_config,
|
|
)
|
|
from evaluate_auc_v2 import (
|
|
LandmarkDataset,
|
|
collate_landmark_fn,
|
|
load_checkpoint_state_dict,
|
|
load_model_state,
|
|
resolve_dist_mode_for_checkpoint,
|
|
)
|
|
from model_architectures import resolve_model_architecture
|
|
|
|
|
|
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
|
|
FORMAT_VERSION = 2
|
|
DEFAULT_COMPRESSION_LEVEL = 4
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Export each test patient's complete disease/death Weibull "
|
|
"shape and scale matrices at two-year age landmarks."
|
|
)
|
|
)
|
|
parser.add_argument(
|
|
"--run_path",
|
|
required=True,
|
|
help="Run directory containing train_config.json and best_model.pt.",
|
|
)
|
|
parser.add_argument(
|
|
"--output_path",
|
|
default=None,
|
|
help=(
|
|
"Single output HDF5 file. Defaults to "
|
|
"<run_path>/weibull_parameters_test_age40_80_step2.h5."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--test_eid_file",
|
|
default=None,
|
|
help=(
|
|
"Optional override. By default use test_eid_file from the run's "
|
|
"train_config.json, or ukb_test_eid.csv if the field is absent."
|
|
),
|
|
)
|
|
parser.add_argument("--age_start", type=float, default=40.0)
|
|
parser.add_argument("--age_stop", type=float, default=80.0)
|
|
parser.add_argument("--age_step", type=float, default=2.0)
|
|
parser.add_argument("--batch_size", type=int, default=128)
|
|
parser.add_argument(
|
|
"--rows_per_chunk",
|
|
type=int,
|
|
default=256,
|
|
help=(
|
|
"Patient rows per internal HDF5 chunk for shape/scale matrices. "
|
|
"This does not create separate output files."
|
|
),
|
|
)
|
|
parser.add_argument("--num_workers", type=int, default=4)
|
|
parser.add_argument(
|
|
"--device",
|
|
default=None,
|
|
help="For example cpu, cuda, or cuda:1. Defaults to CUDA when available.",
|
|
)
|
|
parser.add_argument(
|
|
"--use_amp",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=False,
|
|
help="Use float16 autocast during CUDA inference.",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset_subset_size",
|
|
type=int,
|
|
default=None,
|
|
help="Use only the first N matched test patients for a smoke test.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def build_age_grid(start: float, stop: float, step: float) -> np.ndarray:
|
|
"""Build an inclusive age grid and reject a stop not aligned to the step."""
|
|
if not all(math.isfinite(x) for x in (start, stop, step)):
|
|
raise ValueError("Age start/stop/step must be finite.")
|
|
if step <= 0:
|
|
raise ValueError("age_step must be > 0.")
|
|
if stop < start:
|
|
raise ValueError("age_stop must be >= age_start.")
|
|
count = int(math.floor((stop - start) / step + 1e-9)) + 1
|
|
ages = start + step * np.arange(count, dtype=np.float64)
|
|
if ages.size == 0 or not math.isclose(
|
|
float(ages[-1]), stop, rel_tol=0.0, abs_tol=1e-7
|
|
):
|
|
raise ValueError(
|
|
"age_stop must lie on the grid defined by age_start and age_step."
|
|
)
|
|
ages[-1] = stop
|
|
return ages.astype(np.float32)
|
|
|
|
|
|
def parse_int_list(value: Any) -> List[int] | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, (list, tuple, np.ndarray)):
|
|
return [int(x) for x in value]
|
|
text = str(value).strip()
|
|
if not text:
|
|
return None
|
|
if text.startswith("["):
|
|
parsed = json.loads(text)
|
|
if not isinstance(parsed, list):
|
|
raise ValueError("extra_info_types must be a list of integers.")
|
|
return [int(x) for x in parsed]
|
|
return [int(x.strip()) for x in text.split(",") if x.strip()]
|
|
|
|
|
|
def select_outcome_tokens(dataset: Any) -> List[int]:
|
|
"""Select every real outcome token, including Death, in token-id order."""
|
|
tokens = sorted(
|
|
int(token)
|
|
for token, code in dataset.label_id_to_code.items()
|
|
if int(token) not in SPECIAL_TOKENS and not str(code).startswith("<")
|
|
)
|
|
if not tokens:
|
|
raise RuntimeError("The dataset contains no disease/death outcome tokens.")
|
|
return tokens
|
|
|
|
|
|
def resolve_project_file(path_value: str | Path) -> Path:
|
|
path = Path(path_value)
|
|
if path.is_absolute():
|
|
return path
|
|
direct = Path.cwd() / path
|
|
return direct if direct.is_file() else Path(__file__).resolve().parent / path
|
|
|
|
|
|
def load_label_text(labels_file: str | Path) -> Dict[str, str]:
|
|
result: Dict[str, str] = {}
|
|
with resolve_project_file(labels_file).open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
text = line.strip()
|
|
if text:
|
|
result[text.split()[0]] = text
|
|
return result
|
|
|
|
|
|
def require_h5py() -> Any:
|
|
"""Import h5py lazily so ``--help`` remains available without it."""
|
|
try:
|
|
return importlib.import_module("h5py")
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"This exporter writes one HDF5 file and requires h5py. Install "
|
|
"h5py in the same Python environment used for DeepHealth."
|
|
) from exc
|
|
|
|
|
|
def age_group_name(age: float) -> str:
|
|
text = f"{age:g}".replace("-", "minus_").replace(".", "p")
|
|
return f"age_{text}"
|
|
|
|
|
|
def write_metadata_json(
|
|
metadata_dataset: Any,
|
|
payload: Dict[str, Any],
|
|
) -> None:
|
|
metadata_dataset[()] = json.dumps(
|
|
payload,
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
|
|
def write_empty_age_group(
|
|
*,
|
|
age: float,
|
|
age_group: Any,
|
|
token_count: int,
|
|
) -> Dict[str, Any]:
|
|
"""Represent an age with no eligible patients inside the HDF5 file."""
|
|
age_group.attrs["age"] = float(age)
|
|
age_group.attrs["n_rows"] = 0
|
|
age_group.attrs["n_tokens"] = int(token_count)
|
|
age_group.attrs["nonfinite_shape_values"] = 0
|
|
age_group.attrs["nonfinite_scale_values"] = 0
|
|
age_group.create_dataset("eid", shape=(0,), dtype=np.int64)
|
|
age_group.create_dataset("dataset_index", shape=(0,), dtype=np.int64)
|
|
age_group.create_dataset("sex", shape=(0,), dtype=np.int8)
|
|
age_group.create_dataset("age", shape=(0,), dtype=np.float32)
|
|
age_group.create_dataset(
|
|
"shape", shape=(0, token_count), dtype=np.float32
|
|
)
|
|
age_group.create_dataset(
|
|
"scale", shape=(0, token_count), dtype=np.float32
|
|
)
|
|
return {
|
|
"age": float(age),
|
|
"n_rows": 0,
|
|
"n_tokens": int(token_count),
|
|
"nonfinite_shape_values": 0,
|
|
"nonfinite_scale_values": 0,
|
|
}
|
|
|
|
|
|
def validate_hdf5_export(
|
|
output_file: Any,
|
|
*,
|
|
ages: np.ndarray,
|
|
token_count: int,
|
|
summaries: Sequence[Dict[str, Any]],
|
|
) -> None:
|
|
"""Validate the unified file structure before it is finalized."""
|
|
for group_name in (
|
|
"tokens",
|
|
"test_population",
|
|
"landmarks",
|
|
"age_summary",
|
|
):
|
|
if group_name not in output_file:
|
|
raise RuntimeError(f"HDF5 export is missing /{group_name}.")
|
|
if output_file["tokens/token_id"].shape != (token_count,):
|
|
raise RuntimeError("HDF5 token table length does not match n_tokens.")
|
|
if len(summaries) != int(ages.size):
|
|
raise RuntimeError("HDF5 age summary length does not match age grid.")
|
|
|
|
for age_value, summary in zip(ages.tolist(), summaries):
|
|
age = float(age_value)
|
|
group_path = f"landmarks/{age_group_name(age)}"
|
|
if group_path not in output_file:
|
|
raise RuntimeError(f"HDF5 export is missing /{group_path}.")
|
|
group = output_file[group_path]
|
|
n_rows = int(summary["n_rows"])
|
|
matrix_shape = (n_rows, token_count)
|
|
if group["shape"].shape != matrix_shape:
|
|
raise RuntimeError(
|
|
f"/{group_path}/shape has {group['shape'].shape}, expected "
|
|
f"{matrix_shape}."
|
|
)
|
|
if group["scale"].shape != matrix_shape:
|
|
raise RuntimeError(
|
|
f"/{group_path}/scale has {group['scale'].shape}, expected "
|
|
f"{matrix_shape}."
|
|
)
|
|
for dataset_name in ("eid", "dataset_index", "sex", "age"):
|
|
if group[dataset_name].shape != (n_rows,):
|
|
raise RuntimeError(
|
|
f"/{group_path}/{dataset_name} length does not match "
|
|
"the parameter matrices."
|
|
)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def export_age(
|
|
*,
|
|
age: float,
|
|
model: Any,
|
|
loader: DataLoader,
|
|
tokens: Sequence[int],
|
|
device: torch.device,
|
|
use_amp: bool,
|
|
subset_indices: np.ndarray,
|
|
selected_eids: np.ndarray,
|
|
age_group: Any,
|
|
rows_per_chunk: int,
|
|
) -> Dict[str, Any]:
|
|
"""Write one age directly into its group in the unified HDF5 file."""
|
|
token_index = torch.as_tensor(tokens, dtype=torch.long, device=device)
|
|
amp_enabled = bool(use_amp and device.type == "cuda")
|
|
n_rows_expected = int(len(loader.dataset))
|
|
n_tokens = int(len(tokens))
|
|
age_group.attrs["age"] = float(age)
|
|
age_group.attrs["n_rows"] = n_rows_expected
|
|
age_group.attrs["n_tokens"] = n_tokens
|
|
|
|
if n_rows_expected == 0:
|
|
return write_empty_age_group(
|
|
age=age,
|
|
age_group=age_group,
|
|
token_count=n_tokens,
|
|
)
|
|
|
|
row_chunk = min(int(rows_per_chunk), n_rows_expected)
|
|
vector_options = {
|
|
"chunks": (row_chunk,),
|
|
"compression": "gzip",
|
|
"compression_opts": DEFAULT_COMPRESSION_LEVEL,
|
|
"shuffle": True,
|
|
}
|
|
matrix_options = {
|
|
"chunks": (row_chunk, n_tokens),
|
|
"compression": "gzip",
|
|
"compression_opts": DEFAULT_COMPRESSION_LEVEL,
|
|
"shuffle": True,
|
|
}
|
|
eid_dataset = age_group.create_dataset(
|
|
"eid", shape=(n_rows_expected,), dtype=np.int64, **vector_options
|
|
)
|
|
dataset_index_dataset = age_group.create_dataset(
|
|
"dataset_index",
|
|
shape=(n_rows_expected,),
|
|
dtype=np.int64,
|
|
**vector_options,
|
|
)
|
|
sex_dataset = age_group.create_dataset(
|
|
"sex", shape=(n_rows_expected,), dtype=np.int8, **vector_options
|
|
)
|
|
age_dataset = age_group.create_dataset(
|
|
"age", shape=(n_rows_expected,), dtype=np.float32, **vector_options
|
|
)
|
|
shape_dataset = age_group.create_dataset(
|
|
"shape",
|
|
shape=(n_rows_expected, n_tokens),
|
|
dtype=np.float32,
|
|
**matrix_options,
|
|
)
|
|
scale_dataset = age_group.create_dataset(
|
|
"scale",
|
|
shape=(n_rows_expected, n_tokens),
|
|
dtype=np.float32,
|
|
**matrix_options,
|
|
)
|
|
|
|
row_start = 0
|
|
nonfinite_shape = 0
|
|
nonfinite_scale = 0
|
|
|
|
for batch in tqdm(loader, desc=f"Age {age:g}", dynamic_ncols=True):
|
|
batch_device = {
|
|
key: (
|
|
value.to(device, non_blocking=True)
|
|
if isinstance(value, torch.Tensor)
|
|
else value
|
|
)
|
|
for key, value in batch.items()
|
|
}
|
|
amp_context = (
|
|
torch.autocast(device_type="cuda", dtype=torch.float16)
|
|
if amp_enabled
|
|
else contextlib.nullcontext()
|
|
)
|
|
with amp_context:
|
|
hidden = model(
|
|
event_seq=batch_device["event_seq"],
|
|
time_seq=batch_device["time_seq"],
|
|
sex=batch_device["sex"],
|
|
padding_mask=batch_device["padding_mask"],
|
|
t_query=batch_device["t_query"],
|
|
other_type=batch_device["other_type"],
|
|
other_value=batch_device["other_value"],
|
|
other_value_kind=batch_device["other_value_kind"],
|
|
other_time=batch_device["other_time"],
|
|
)
|
|
logits = model.calc_risk(hidden).index_select(1, token_index).float()
|
|
shape = (
|
|
model.calc_weibull_rho(hidden)
|
|
.index_select(1, token_index)
|
|
.float()
|
|
)
|
|
rate = F.softplus(logits) + 1e-8
|
|
scale = torch.exp(-torch.log(rate) / shape)
|
|
|
|
shape_np = shape.cpu().numpy().astype(np.float32, copy=False)
|
|
scale_np = scale.cpu().numpy().astype(np.float32, copy=False)
|
|
patient_id = batch["patient_id"].cpu().numpy().astype(np.int64)
|
|
dataset_index = subset_indices[patient_id]
|
|
n_rows = int(shape_np.shape[0])
|
|
row_stop = row_start + n_rows
|
|
if row_stop > n_rows_expected:
|
|
raise RuntimeError(
|
|
f"Age {age:g} produced more rows than expected: "
|
|
f"{row_stop} > {n_rows_expected}."
|
|
)
|
|
eid_dataset[row_start:row_stop] = selected_eids[patient_id]
|
|
dataset_index_dataset[row_start:row_stop] = dataset_index.astype(
|
|
np.int64, copy=False
|
|
)
|
|
sex_dataset[row_start:row_stop] = (
|
|
batch["sex"].cpu().numpy().astype(np.int8, copy=False)
|
|
)
|
|
age_dataset[row_start:row_stop] = (
|
|
batch["landmark_age"]
|
|
.cpu()
|
|
.numpy()
|
|
.astype(np.float32, copy=False)
|
|
)
|
|
shape_dataset[row_start:row_stop, :] = shape_np
|
|
scale_dataset[row_start:row_stop, :] = scale_np
|
|
nonfinite_shape += int((~np.isfinite(shape_np)).sum())
|
|
nonfinite_scale += int((~np.isfinite(scale_np)).sum())
|
|
row_start = row_stop
|
|
|
|
if row_start != n_rows_expected:
|
|
raise RuntimeError(
|
|
f"Age {age:g} row count mismatch: wrote {row_start}, "
|
|
f"expected {n_rows_expected}."
|
|
)
|
|
age_group.attrs["nonfinite_shape_values"] = nonfinite_shape
|
|
age_group.attrs["nonfinite_scale_values"] = nonfinite_scale
|
|
return {
|
|
"age": float(age),
|
|
"n_rows": n_rows_expected,
|
|
"n_tokens": n_tokens,
|
|
"nonfinite_shape_values": nonfinite_shape,
|
|
"nonfinite_scale_values": nonfinite_scale,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
run_path = Path(args.run_path).resolve()
|
|
config_path = run_path / "train_config.json"
|
|
checkpoint_path = run_path / "best_model.pt"
|
|
if not config_path.is_file():
|
|
raise FileNotFoundError(config_path)
|
|
if not checkpoint_path.is_file():
|
|
raise FileNotFoundError(checkpoint_path)
|
|
if args.batch_size <= 0:
|
|
raise ValueError("batch_size must be > 0.")
|
|
if args.rows_per_chunk <= 0:
|
|
raise ValueError("rows_per_chunk must be > 0.")
|
|
if args.num_workers < 0:
|
|
raise ValueError("num_workers must be >= 0.")
|
|
if args.dataset_subset_size is not None and args.dataset_subset_size <= 0:
|
|
raise ValueError("dataset_subset_size must be > 0.")
|
|
|
|
cfg = load_json_config(config_path)
|
|
validate_training_mode_config(cfg)
|
|
model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
|
|
if model_target_mode != "all_future":
|
|
raise ValueError(
|
|
"This exporter requires model_target_mode='all_future'; got "
|
|
f"{model_target_mode!r}."
|
|
)
|
|
|
|
ages = build_age_grid(args.age_start, args.age_stop, args.age_step)
|
|
output_path = (
|
|
Path(args.output_path).resolve()
|
|
if args.output_path
|
|
else run_path / "weibull_parameters_test_age40_80_step2.h5"
|
|
)
|
|
if output_path.exists():
|
|
raise FileExistsError(
|
|
f"Output file already exists: {output_path}. Choose a new "
|
|
"--output_path."
|
|
)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary_output = output_path.with_name(f".{output_path.name}.partial")
|
|
if temporary_output.exists():
|
|
raise FileExistsError(
|
|
f"Partial output already exists: {temporary_output}. Remove or "
|
|
"rename it before retrying."
|
|
)
|
|
h5py = require_h5py()
|
|
|
|
data_prefix = str(cfg.get("data_prefix", "ukb"))
|
|
labels_file = str(cfg.get("labels_file", "labels.csv"))
|
|
disease_history_mode = normalize_disease_history_mode(
|
|
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
|
|
)
|
|
min_history_events = int(
|
|
cfg.get("all_future_min_history_events", cfg.get("min_history_events", 1))
|
|
)
|
|
min_future_events = int(
|
|
cfg.get("all_future_min_future_events", cfg.get("min_future_events", 1))
|
|
)
|
|
print("Loading dataset...")
|
|
dataset = load_sequence_eval_dataset(
|
|
model_target_mode=model_target_mode,
|
|
data_prefix=data_prefix,
|
|
labels_file=labels_file,
|
|
no_event_interval_years=float(cfg.get("no_event_interval_years", 5.0)),
|
|
min_history_events=min_history_events,
|
|
min_future_events=min_future_events,
|
|
extra_info_types=parse_int_list(cfg.get("extra_info_types")),
|
|
disease_history_mode=disease_history_mode,
|
|
)
|
|
validate_dataset_metadata(dataset, cfg)
|
|
|
|
test_eid_file = args.test_eid_file or cfg.get(
|
|
"test_eid_file", "ukb_test_eid.csv"
|
|
)
|
|
if not test_eid_file:
|
|
raise ValueError(
|
|
"No test_eid_file is defined. Provide --test_eid_file explicitly."
|
|
)
|
|
subset_indices, resolved_eid_file = select_indices_by_eid_file(
|
|
dataset, str(test_eid_file)
|
|
)
|
|
if args.dataset_subset_size is not None:
|
|
subset_indices = subset_indices[: args.dataset_subset_size]
|
|
if subset_indices.size == 0:
|
|
raise RuntimeError("The selected test subset is empty.")
|
|
|
|
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
|
|
dist_mode = resolve_dist_mode_for_checkpoint(
|
|
str(cfg.get("dist_mode", "exponential")), state_dict
|
|
)
|
|
if dist_mode != "weibull":
|
|
raise ValueError(
|
|
f"The specified run uses dist_mode={dist_mode!r}, not 'weibull'."
|
|
)
|
|
cfg_model = dict(cfg)
|
|
cfg_model["dist_mode"] = dist_mode
|
|
cfg_model["model_architecture"] = resolve_model_architecture(
|
|
cfg_model, state_dict
|
|
)
|
|
device = resolve_eval_device(args.device)
|
|
model = build_model_from_dataset(
|
|
args, cfg_model, dataset, state_dict=state_dict
|
|
).to(device)
|
|
load_model_state(model, state_dict)
|
|
model.eval()
|
|
|
|
tokens = select_outcome_tokens(dataset)
|
|
label_text = load_label_text(labels_file)
|
|
death_tokens = [
|
|
token
|
|
for token in tokens
|
|
if str(dataset.label_id_to_code[token]).lower() == "death"
|
|
]
|
|
if not death_tokens:
|
|
raise RuntimeError("Death token was not found in the outcome vocabulary.")
|
|
|
|
token_codes = [str(dataset.label_id_to_code[token]) for token in tokens]
|
|
token_text = [label_text.get(code, code) for code in token_codes]
|
|
outcome_types = [
|
|
"death" if code.lower() == "death" else "disease"
|
|
for code in token_codes
|
|
]
|
|
selected_eids = np.asarray(
|
|
[int(dataset.samples[int(index)]["eid"]) for index in subset_indices],
|
|
dtype=np.int64,
|
|
)
|
|
metadata: Dict[str, Any] = {
|
|
"format_version": FORMAT_VERSION,
|
|
"complete": False,
|
|
"run_path": str(run_path),
|
|
"checkpoint": str(checkpoint_path),
|
|
"test_eid_file": str(resolved_eid_file),
|
|
"n_selected_test_patients": int(subset_indices.size),
|
|
"ages": [float(x) for x in ages],
|
|
"n_tokens": len(tokens),
|
|
"matrix_dtype": "float32",
|
|
"hdf5_layout": {
|
|
"tokens": (
|
|
"/tokens/{column,token_id,label_code,label_text,outcome_type}"
|
|
),
|
|
"test_population": "/test_population/{eid,dataset_index}",
|
|
"age_summary": (
|
|
"/age_summary/{age,n_rows,nonfinite_shape_values,"
|
|
"nonfinite_scale_values}"
|
|
),
|
|
"landmarks": (
|
|
"/landmarks/age_*/{eid,dataset_index,sex,age,shape,scale}"
|
|
),
|
|
},
|
|
"compression": "gzip",
|
|
"compression_level": DEFAULT_COMPRESSION_LEVEL,
|
|
"rows_per_chunk": int(args.rows_per_chunk),
|
|
"parameterization": {
|
|
"survival": "S(t) = exp(-rate * t^shape)",
|
|
"shape": "rho = softplus(rho_logit) + 1e-6",
|
|
"rate": "softplus(risk_logit) + 1e-8",
|
|
"scale": "rate^(-1/shape)",
|
|
"time_unit": "years after landmark age",
|
|
},
|
|
"eligibility": (
|
|
"At each age: follow-up extends beyond the landmark, the patient "
|
|
"is alive at the landmark, and the configured minimum disease "
|
|
"history is available. All disease and death token parameters are "
|
|
"exported, including tokens already prevalent by the landmark."
|
|
),
|
|
}
|
|
|
|
string_dtype = h5py.string_dtype(encoding="utf-8")
|
|
summaries: List[Dict[str, Any]] = []
|
|
with h5py.File(temporary_output, "w") as output_file:
|
|
output_file.attrs["format_version"] = FORMAT_VERSION
|
|
output_file.attrs["complete"] = False
|
|
output_file.attrs["run_path"] = str(run_path)
|
|
output_file.attrs["checkpoint"] = str(checkpoint_path)
|
|
output_file.attrs["test_eid_file"] = str(resolved_eid_file)
|
|
output_file.attrs["n_selected_test_patients"] = int(
|
|
subset_indices.size
|
|
)
|
|
output_file.attrs["n_tokens"] = int(len(tokens))
|
|
output_file.attrs["matrix_dtype"] = "float32"
|
|
output_file.attrs["compression"] = "gzip"
|
|
output_file.attrs["compression_level"] = DEFAULT_COMPRESSION_LEVEL
|
|
|
|
metadata_dataset = output_file.create_dataset(
|
|
"metadata_json", shape=(), dtype=string_dtype
|
|
)
|
|
write_metadata_json(metadata_dataset, metadata)
|
|
output_file.create_dataset("ages", data=ages.astype(np.float32))
|
|
|
|
token_group = output_file.create_group("tokens")
|
|
token_group.create_dataset(
|
|
"column", data=np.arange(len(tokens), dtype=np.int64)
|
|
)
|
|
token_group.create_dataset(
|
|
"token_id", data=np.asarray(tokens, dtype=np.int64)
|
|
)
|
|
token_group.create_dataset(
|
|
"label_code",
|
|
data=np.asarray(token_codes, dtype=object),
|
|
dtype=string_dtype,
|
|
)
|
|
token_group.create_dataset(
|
|
"label_text",
|
|
data=np.asarray(token_text, dtype=object),
|
|
dtype=string_dtype,
|
|
)
|
|
token_group.create_dataset(
|
|
"outcome_type",
|
|
data=np.asarray(outcome_types, dtype=object),
|
|
dtype=string_dtype,
|
|
)
|
|
|
|
population_group = output_file.create_group("test_population")
|
|
population_group.create_dataset(
|
|
"eid", data=selected_eids, compression="gzip", shuffle=True
|
|
)
|
|
population_group.create_dataset(
|
|
"dataset_index",
|
|
data=subset_indices.astype(np.int64, copy=False),
|
|
compression="gzip",
|
|
shuffle=True,
|
|
)
|
|
|
|
landmark_root = output_file.create_group("landmarks")
|
|
eligible_queries_by_age: Dict[str, int] = {}
|
|
for age_value in ages.tolist():
|
|
age = float(age_value)
|
|
age_group = landmark_root.create_group(age_group_name(age))
|
|
try:
|
|
landmark_dataset = LandmarkDataset(
|
|
dataset=dataset,
|
|
subset_indices=subset_indices,
|
|
landmark_ages=np.asarray([age], dtype=np.float32),
|
|
model_target_mode=model_target_mode,
|
|
min_history_events=min_history_events,
|
|
first_occurrence_by_token={},
|
|
death_token_ids=death_tokens,
|
|
disease_history_mode=disease_history_mode,
|
|
)
|
|
except RuntimeError as exc:
|
|
if "No eligible landmark query samples" not in str(exc):
|
|
raise
|
|
summary = write_empty_age_group(
|
|
age=age,
|
|
age_group=age_group,
|
|
token_count=len(tokens),
|
|
)
|
|
else:
|
|
loader = DataLoader(
|
|
landmark_dataset,
|
|
batch_size=int(args.batch_size),
|
|
shuffle=False,
|
|
collate_fn=collate_landmark_fn,
|
|
num_workers=int(args.num_workers),
|
|
pin_memory=device.type == "cuda",
|
|
persistent_workers=args.num_workers > 0,
|
|
prefetch_factor=2 if args.num_workers > 0 else None,
|
|
)
|
|
summary = export_age(
|
|
age=age,
|
|
model=model,
|
|
loader=loader,
|
|
tokens=tokens,
|
|
device=device,
|
|
use_amp=bool(args.use_amp),
|
|
subset_indices=subset_indices,
|
|
selected_eids=selected_eids,
|
|
age_group=age_group,
|
|
rows_per_chunk=int(args.rows_per_chunk),
|
|
)
|
|
|
|
summaries.append(summary)
|
|
eligible_count = int(summary["n_rows"])
|
|
eligible_queries_by_age[f"{age:g}"] = eligible_count
|
|
metadata["eligible_queries_by_age"] = eligible_queries_by_age
|
|
write_metadata_json(metadata_dataset, metadata)
|
|
output_file.flush()
|
|
print(f"Age {age:g}: exported {eligible_count} patient rows")
|
|
|
|
summary_group = output_file.create_group("age_summary")
|
|
summary_group.create_dataset(
|
|
"age",
|
|
data=np.asarray([row["age"] for row in summaries], dtype=np.float32),
|
|
)
|
|
summary_group.create_dataset(
|
|
"n_rows",
|
|
data=np.asarray([row["n_rows"] for row in summaries], dtype=np.int64),
|
|
)
|
|
summary_group.create_dataset(
|
|
"nonfinite_shape_values",
|
|
data=np.asarray(
|
|
[row["nonfinite_shape_values"] for row in summaries],
|
|
dtype=np.int64,
|
|
),
|
|
)
|
|
summary_group.create_dataset(
|
|
"nonfinite_scale_values",
|
|
data=np.asarray(
|
|
[row["nonfinite_scale_values"] for row in summaries],
|
|
dtype=np.int64,
|
|
),
|
|
)
|
|
|
|
nonfinite_shape = sum(
|
|
int(row["nonfinite_shape_values"]) for row in summaries
|
|
)
|
|
nonfinite_scale = sum(
|
|
int(row["nonfinite_scale_values"]) for row in summaries
|
|
)
|
|
metadata["total_exported_query_rows"] = sum(
|
|
int(row["n_rows"]) for row in summaries
|
|
)
|
|
metadata["nonfinite_shape_values"] = nonfinite_shape
|
|
metadata["nonfinite_scale_values"] = nonfinite_scale
|
|
validate_hdf5_export(
|
|
output_file,
|
|
ages=ages,
|
|
token_count=len(tokens),
|
|
summaries=summaries,
|
|
)
|
|
metadata["validated"] = True
|
|
metadata["complete"] = True
|
|
write_metadata_json(metadata_dataset, metadata)
|
|
output_file.attrs["validated"] = True
|
|
output_file.attrs.modify("complete", True)
|
|
output_file.flush()
|
|
|
|
temporary_output.replace(output_path)
|
|
|
|
if nonfinite_shape or nonfinite_scale:
|
|
print(
|
|
"WARNING: exported non-finite values: "
|
|
f"shape={nonfinite_shape}, scale={nonfinite_scale}."
|
|
)
|
|
print(f"Saved unified Weibull parameter file to: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|