JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Export Weibull parameters to one HDF5 file

Browse Source
This commit is contained in:
Jiarui Li
2026-08-05 12:54:14 +08:00
parent d27eca3f3d
commit 75c9f06114
1 changed files with 386 additions and 250 deletions
Download Patch File Download Diff File Expand all files Collapse all files

538
export_weibull_parameters.py
View File

@@ -1,4 +1,4 @@
"""Export individual Weibull parameters for every disease/death token.
"""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
@@ -13,21 +13,22 @@ and the standard Weibull parameters exported by this script are
shape = rho
scale = rate ** (-1 / shape)
Output is sharded compressed NPZ rather than one monolithic matrix. Every NPZ
contains aligned ``shape`` and ``scale`` matrices with rows for patients and
columns for tokens. ``tokens.csv`` defines the column order and
``manifest.csv`` defines the shard/row order for each age.
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 csv
import importlib
import json
import math
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence
from typing import Any, Dict, List, Sequence
import numpy as np
import torch
@@ -62,17 +63,8 @@ from model_architectures import resolve_model_architecture
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
MANIFEST_FIELDS = [
"age",
"shard",
"file",
"row_start",
"row_stop",
"n_rows",
"n_tokens",
"nonfinite_shape_values",
"nonfinite_scale_values",
]
FORMAT_VERSION = 2
DEFAULT_COMPRESSION_LEVEL = 4
def parse_args() -> argparse.Namespace:
@@ -91,8 +83,8 @@ def parse_args() -> argparse.Namespace:
"--output_path",
default=None,
help=(
"Output directory. Defaults to "
"<run_path>/weibull_parameters_test_age40_80_step2."
"Single output HDF5 file. Defaults to "
"<run_path>/weibull_parameters_test_age40_80_step2.h5."
),
)
parser.add_argument(
@@ -108,12 +100,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--age_step", type=float, default=2.0)
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument(
"--rows_per_shard",
"--rows_per_chunk",
type=int,
default=4096,
default=256,
help=(
"Approximate patient rows per compressed NPZ shard. This is "
"independent of inference batch size."
"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)
@@ -203,56 +195,57 @@ def load_label_text(labels_file: str | Path) -> Dict[str, str]:
return result
def write_csv(
path: Path,
fieldnames: Sequence[str],
rows: Iterable[Dict[str, Any]],
) -> None:
temporary = path.with_name(f".{path.name}.tmp")
with temporary.open("w", newline="", encoding="utf-8-sig") as handle:
writer = csv.DictWriter(handle, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
temporary.replace(path)
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 write_json(path: Path, payload: Dict[str, Any]) -> None:
temporary = path.with_name(f".{path.name}.tmp")
temporary.write_text(
json.dumps(payload, ensure_ascii=False, indent=2),
encoding="utf-8",
)
temporary.replace(path)
def age_directory_name(age: float) -> str:
def age_group_name(age: float) -> str:
text = f"{age:g}".replace("-", "minus_").replace(".", "p")
return f"age_{text}"
def write_empty_age_export(
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_dir: Path,
age_group: Any,
token_count: int,
) -> Dict[str, Any]:
"""Represent an age with no eligible test patients by one empty shard."""
shard_path = age_dir / "shard_000000.npz"
np.savez_compressed(
shard_path,
eid=np.empty(0, dtype=np.int64),
dataset_index=np.empty(0, dtype=np.int64),
sex=np.empty(0, dtype=np.int8),
age=np.empty(0, dtype=np.float32),
shape=np.empty((0, token_count), dtype=np.float32),
scale=np.empty((0, token_count), dtype=np.float32),
"""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),
"shard": 0,
"file": str(shard_path.relative_to(age_dir.parent.parent)),
"row_start": 0,
"row_stop": 0,
"n_rows": 0,
"n_tokens": int(token_count),
"nonfinite_shape_values": 0,
@@ -260,6 +253,53 @@ def write_empty_age_export(
}
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(
*,
@@ -269,77 +309,71 @@ def export_age(
tokens: Sequence[int],
device: torch.device,
use_amp: bool,
dataset: Any,
subset_indices: np.ndarray,
age_dir: Path,
output_path: Path,
rows_per_shard: int,
) -> List[Dict[str, Any]]:
"""Export all patient-by-token matrices for one age in row-order shards."""
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)
patient_eids = np.asarray(
[int(dataset.samples[int(index)]["eid"]) for index in subset_indices],
dtype=np.int64,
)
amp_enabled = bool(use_amp and device.type == "cuda")
manifest_rows: List[Dict[str, Any]] = []
row_start = 0
shape_parts: List[np.ndarray] = []
scale_parts: List[np.ndarray] = []
eid_parts: List[np.ndarray] = []
dataset_index_parts: List[np.ndarray] = []
sex_parts: List[np.ndarray] = []
age_parts: List[np.ndarray] = []
buffered_rows = 0
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
def flush_shard() -> None:
nonlocal row_start, buffered_rows
if buffered_rows == 0:
return
shape_matrix = np.concatenate(shape_parts, axis=0)
scale_matrix = np.concatenate(scale_parts, axis=0)
eid = np.concatenate(eid_parts)
dataset_index = np.concatenate(dataset_index_parts)
sex = np.concatenate(sex_parts)
age_values = np.concatenate(age_parts)
shard_index = len(manifest_rows)
n_rows = int(shape_matrix.shape[0])
row_stop = row_start + n_rows
shard_path = age_dir / f"shard_{shard_index:06d}.npz"
np.savez_compressed(
shard_path,
eid=eid,
dataset_index=dataset_index,
sex=sex,
age=age_values,
shape=shape_matrix,
scale=scale_matrix,
if n_rows_expected == 0:
return write_empty_age_group(
age=age,
age_group=age_group,
token_count=n_tokens,
)
manifest_rows.append(
{
"age": float(age),
"shard": int(shard_index),
"file": str(shard_path.relative_to(output_path)),
"row_start": int(row_start),
"row_stop": int(row_stop),
"n_rows": n_rows,
"n_tokens": int(shape_matrix.shape[1]),
"nonfinite_shape_values": int(
(~np.isfinite(shape_matrix)).sum()
),
"nonfinite_scale_values": int(
(~np.isfinite(scale_matrix)).sum()
),
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
)
row_start = row_stop
buffered_rows = 0
shape_parts.clear()
scale_parts.clear()
eid_parts.clear()
dataset_index_parts.clear()
sex_parts.clear()
age_parts.clear()
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 = {
@@ -381,27 +415,45 @@ def export_age(
patient_id = batch["patient_id"].cpu().numpy().astype(np.int64)
dataset_index = subset_indices[patient_id]
n_rows = int(shape_np.shape[0])
shape_parts.append(shape_np)
scale_parts.append(scale_np)
eid_parts.append(patient_eids[patient_id])
dataset_index_parts.append(
dataset_index.astype(np.int64, copy=False)
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}."
)
sex_parts.append(
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_parts.append(
age_dataset[row_start:row_stop] = (
batch["landmark_age"]
.cpu()
.numpy()
.astype(np.float32, copy=False)
)
buffered_rows += n_rows
if buffered_rows >= rows_per_shard:
flush_shard()
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
flush_shard()
return manifest_rows
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:
@@ -415,8 +467,8 @@ def main() -> None:
raise FileNotFoundError(checkpoint_path)
if args.batch_size <= 0:
raise ValueError("batch_size must be > 0.")
if args.rows_per_shard <= 0:
raise ValueError("rows_per_shard 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:
@@ -435,16 +487,21 @@ def main() -> None:
output_path = (
Path(args.output_path).resolve()
if args.output_path
else run_path / "weibull_parameters_test_age40_80_step2"
else run_path / "weibull_parameters_test_age40_80_step2.h5"
)
if output_path.exists() and any(output_path.iterdir()):
if output_path.exists():
raise FileExistsError(
f"Output directory is not empty: {output_path}. Choose a new "
"--output_path so shards from different exports cannot be mixed."
f"Output file already exists: {output_path}. Choose a new "
"--output_path."
)
output_path.mkdir(parents=True, exist_ok=True)
shards_root = output_path / "shards"
shards_root.mkdir(parents=True, exist_ok=True)
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"))
@@ -515,31 +572,18 @@ def main() -> None:
if not death_tokens:
raise RuntimeError("Death token was not found in the outcome vocabulary.")
token_rows = [
{
"column": column,
"token_id": token,
"label_code": str(dataset.label_id_to_code[token]),
"label_text": label_text.get(
str(dataset.label_id_to_code[token]),
str(dataset.label_id_to_code[token]),
),
"outcome_type": (
"death"
if str(dataset.label_id_to_code[token]).lower() == "death"
else "disease"
),
}
for column, token in enumerate(tokens)
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
]
write_csv(
output_path / "tokens.csv",
["column", "token_id", "label_code", "label_text", "outcome_type"],
token_rows,
selected_eids = np.asarray(
[int(dataset.samples[int(index)]["eid"]) for index in subset_indices],
dtype=np.int64,
)
metadata: Dict[str, Any] = {
"format_version": 1,
"format_version": FORMAT_VERSION,
"complete": False,
"run_path": str(run_path),
"checkpoint": str(checkpoint_path),
@@ -547,17 +591,23 @@ def main() -> None:
"n_selected_test_patients": int(subset_indices.size),
"ages": [float(x) for x in ages],
"n_tokens": len(tokens),
"token_columns_file": "tokens.csv",
"manifest_file": "manifest.csv",
"matrix_dtype": "float32",
"npz_arrays": {
"eid": "int64 [n_rows]",
"dataset_index": "int64 [n_rows]",
"sex": "int8 [n_rows], 0=female and 1=male",
"age": "float32 [n_rows]",
"shape": "float32 [n_rows, n_tokens]",
"scale": "float32 [n_rows, n_tokens]",
"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",
@@ -572,14 +622,68 @@ def main() -> None:
"exported, including tokens already prevalent by the landmark."
),
}
write_json(output_path / "metadata.json", metadata)
manifest_rows: List[Dict[str, Any]] = []
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_dir = shards_root / age_directory_name(age)
age_dir.mkdir(parents=True, exist_ok=True)
age_group = landmark_root.create_group(age_group_name(age))
try:
landmark_dataset = LandmarkDataset(
dataset=dataset,
@@ -594,14 +698,11 @@ def main() -> None:
except RuntimeError as exc:
if "No eligible landmark query samples" not in str(exc):
raise
age_rows = [
write_empty_age_export(
summary = write_empty_age_group(
age=age,
age_dir=age_dir,
age_group=age_group,
token_count=len(tokens),
)
]
eligible_count = 0
else:
loader = DataLoader(
landmark_dataset,
@@ -613,48 +714,83 @@ def main() -> None:
persistent_workers=args.num_workers > 0,
prefetch_factor=2 if args.num_workers > 0 else None,
)
age_rows = export_age(
summary = export_age(
age=age,
model=model,
loader=loader,
tokens=tokens,
device=device,
use_amp=bool(args.use_amp),
dataset=dataset,
subset_indices=subset_indices,
age_dir=age_dir,
output_path=output_path,
rows_per_shard=int(args.rows_per_shard),
selected_eids=selected_eids,
age_group=age_group,
rows_per_chunk=int(args.rows_per_chunk),
)
eligible_count = len(landmark_dataset)
manifest_rows.extend(age_rows)
eligible_queries_by_age[f"{age:g}"] = int(eligible_count)
write_csv(output_path / "manifest.csv", MANIFEST_FIELDS, manifest_rows)
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_json(output_path / "metadata.json", metadata)
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 manifest_rows
int(row["nonfinite_shape_values"]) for row in summaries
)
nonfinite_scale = sum(
int(row["nonfinite_scale_values"]) for row in manifest_rows
int(row["nonfinite_scale_values"]) for row in summaries
)
metadata["total_exported_query_rows"] = sum(
int(row["n_rows"]) for row in manifest_rows
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_json(output_path / "metadata.json", metadata)
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 individual Weibull parameter matrices to: {output_path}")
print(f"Saved unified Weibull parameter file to: {output_path}")
if __name__ == "__main__":