DeepHealth/export_weibull_parameters.py

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