Export Weibull parameters to one HDF5 file

export_weibull_parameters.py

@@ -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
+The output is one HDF5 file.  It contains the token table, age summary and one
-contains aligned ``shape`` and ``scale`` matrices with rows for patients and
+group per landmark age.  Each age group stores aligned patient metadata plus
-columns for tokens.  ``tokens.csv`` defines the column order and
+chunked, compressed ``shape`` and ``scale`` matrices.  Chunking is internal to
-``manifest.csv`` defines the shard/row order for each age.
+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 = [
+FORMAT_VERSION = 2
-    "age",
+DEFAULT_COMPRESSION_LEVEL = 4
-    "shard",
-    "file",
-    "row_start",
-    "row_stop",
-    "n_rows",
-    "n_tokens",
-    "nonfinite_shape_values",
-    "nonfinite_scale_values",
-]
 
 
 def parse_args() -> argparse.Namespace:
@@ -91,8 +83,8 @@ def parse_args() -> argparse.Namespace:
         "--output_path",
         default=None,
         help=(
-            "Output directory. Defaults to "
+            "Single output HDF5 file. Defaults to "
-            "<run_path>/weibull_parameters_test_age40_80_step2."
+            "<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 "
+            "Patient rows per internal HDF5 chunk for shape/scale matrices. "
-            "independent of inference batch size."
+            "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(
+def require_h5py() -> Any:
-    path: Path,
+    """Import h5py lazily so ``--help`` remains available without it."""
-    fieldnames: Sequence[str],
+    try:
-    rows: Iterable[Dict[str, Any]],
+        return importlib.import_module("h5py")
-) -> None:
+    except ImportError as exc:
-    temporary = path.with_name(f".{path.name}.tmp")
+        raise RuntimeError(
-    with temporary.open("w", newline="", encoding="utf-8-sig") as handle:
+            "This exporter writes one HDF5 file and requires h5py. Install "
-        writer = csv.DictWriter(handle, fieldnames=fieldnames)
+            "h5py in the same Python environment used for DeepHealth."
-        writer.writeheader()
+        ) from exc
-        writer.writerows(rows)
-    temporary.replace(path)
 
 
-def write_json(path: Path, payload: Dict[str, Any]) -> None:
+def age_group_name(age: float) -> str:
-    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:
     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."""
+    """Represent an age with no eligible patients inside the HDF5 file."""
-    shard_path = age_dir / "shard_000000.npz"
+    age_group.attrs["age"] = float(age)
-    np.savez_compressed(
+    age_group.attrs["n_rows"] = 0
-        shard_path,
+    age_group.attrs["n_tokens"] = int(token_count)
-        eid=np.empty(0, dtype=np.int64),
+    age_group.attrs["nonfinite_shape_values"] = 0
-        dataset_index=np.empty(0, dtype=np.int64),
+    age_group.attrs["nonfinite_scale_values"] = 0
-        sex=np.empty(0, dtype=np.int8),
+    age_group.create_dataset("eid", shape=(0,), dtype=np.int64)
-        age=np.empty(0, dtype=np.float32),
+    age_group.create_dataset("dataset_index", shape=(0,), dtype=np.int64)
-        shape=np.empty((0, token_count), dtype=np.float32),
+    age_group.create_dataset("sex", shape=(0,), dtype=np.int8)
-        scale=np.empty((0, token_count), dtype=np.float32),
+    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,
+    selected_eids: np.ndarray,
-    output_path: Path,
+    age_group: Any,
-    rows_per_shard: int,
+    rows_per_chunk: int,
-) -> List[Dict[str, Any]]:
+) -> Dict[str, Any]:
-    """Export all patient-by-token matrices for one age in row-order shards."""
+    """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]] = []
+    n_rows_expected = int(len(loader.dataset))
-    row_start = 0
+    n_tokens = int(len(tokens))
-    shape_parts: List[np.ndarray] = []
+    age_group.attrs["age"] = float(age)
-    scale_parts: List[np.ndarray] = []
+    age_group.attrs["n_rows"] = n_rows_expected
-    eid_parts: List[np.ndarray] = []
+    age_group.attrs["n_tokens"] = n_tokens
-    dataset_index_parts: List[np.ndarray] = []
-    sex_parts: List[np.ndarray] = []
-    age_parts: List[np.ndarray] = []
-    buffered_rows = 0
 
-    def flush_shard() -> None:
+    if n_rows_expected == 0:
-        nonlocal row_start, buffered_rows
+        return write_empty_age_group(
-        if buffered_rows == 0:
+            age=age,
-            return
+            age_group=age_group,
-        shape_matrix = np.concatenate(shape_parts, axis=0)
+            token_count=n_tokens,
-        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,
         )
-        manifest_rows.append(
+
-            {
+    row_chunk = min(int(rows_per_chunk), n_rows_expected)
-                "age": float(age),
+    vector_options = {
-                "shard": int(shard_index),
+        "chunks": (row_chunk,),
-                "file": str(shard_path.relative_to(output_path)),
+        "compression": "gzip",
-                "row_start": int(row_start),
+        "compression_opts": DEFAULT_COMPRESSION_LEVEL,
-                "row_stop": int(row_stop),
+        "shuffle": True,
-                "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()
-                ),
     }
+    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
+    dataset_index_dataset = age_group.create_dataset(
-        buffered_rows = 0
+        "dataset_index",
-        shape_parts.clear()
+        shape=(n_rows_expected,),
-        scale_parts.clear()
+        dtype=np.int64,
-        eid_parts.clear()
+        **vector_options,
-        dataset_index_parts.clear()
+    )
-        sex_parts.clear()
+    sex_dataset = age_group.create_dataset(
-        age_parts.clear()
+        "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)
+        row_stop = row_start + n_rows
-        scale_parts.append(scale_np)
+        if row_stop > n_rows_expected:
-        eid_parts.append(patient_eids[patient_id])
+            raise RuntimeError(
-        dataset_index_parts.append(
+                f"Age {age:g} produced more rows than expected: "
-            dataset_index.astype(np.int64, copy=False)
+                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
+        shape_dataset[row_start:row_stop, :] = shape_np
-        if buffered_rows >= rows_per_shard:
+        scale_dataset[row_start:row_stop, :] = scale_np
-            flush_shard()
+        nonfinite_shape += int((~np.isfinite(shape_np)).sum())
+        nonfinite_scale += int((~np.isfinite(scale_np)).sum())
+        row_start = row_stop
 
-    flush_shard()
+    if row_start != n_rows_expected:
-    return manifest_rows
+        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:
+    if args.rows_per_chunk <= 0:
-        raise ValueError("rows_per_shard must be > 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 "
+            f"Output file already exists: {output_path}. Choose a new "
-            "--output_path so shards from different exports cannot be mixed."
+            "--output_path."
         )
-    output_path.mkdir(parents=True, exist_ok=True)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
-    shards_root = output_path / "shards"
+    temporary_output = output_path.with_name(f".{output_path.name}.partial")
-    shards_root.mkdir(parents=True, exist_ok=True)
+    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 = [
+    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]
-            "column": column,
+    outcome_types = [
-            "token_id": token,
+        "death" if code.lower() == "death" else "disease"
-            "label_code": str(dataset.label_id_to_code[token]),
+        for code in token_codes
-            "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)
     ]
-    write_csv(
+    selected_eids = np.asarray(
-        output_path / "tokens.csv",
+        [int(dataset.samples[int(index)]["eid"]) for index in subset_indices],
-        ["column", "token_id", "label_code", "label_text", "outcome_type"],
+        dtype=np.int64,
-        token_rows,
     )
-
     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": {
+        "hdf5_layout": {
-            "eid": "int64 [n_rows]",
+            "tokens": (
-            "dataset_index": "int64 [n_rows]",
+                "/tokens/{column,token_id,label_code,label_text,outcome_type}"
-            "sex": "int8 [n_rows], 0=female and 1=male",
+            ),
-            "age": "float32 [n_rows]",
+            "test_population": "/test_population/{eid,dataset_index}",
-            "shape": "float32 [n_rows, n_tokens]",
+            "age_summary": (
-            "scale": "float32 [n_rows, n_tokens]",
+                "/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_group = landmark_root.create_group(age_group_name(age))
-        age_dir.mkdir(parents=True, exist_ok=True)
             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 = [
+                summary = write_empty_age_group(
-                write_empty_age_export(
                     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,
+                    selected_eids=selected_eids,
-                output_path=output_path,
+                    age_group=age_group,
-                rows_per_shard=int(args.rows_per_shard),
+                    rows_per_chunk=int(args.rows_per_chunk),
                 )
-            eligible_count = len(landmark_dataset)
 
-        manifest_rows.extend(age_rows)
+            summaries.append(summary)
-        eligible_queries_by_age[f"{age:g}"] = int(eligible_count)
+            eligible_count = int(summary["n_rows"])
-        write_csv(output_path / "manifest.csv", MANIFEST_FIELDS, manifest_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__":