diff --git a/export_weibull_parameters.py b/export_weibull_parameters.py index f87c2a7..c419366 100644 --- a/export_weibull_parameters.py +++ b/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 -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 " - "/weibull_parameters_test_age40_80_step2." + "Single output HDF5 file. Defaults to " + "/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_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() + + 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 = { @@ -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}." + ) + 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_parts.append( + 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,89 +622,175 @@ def main() -> None: "exported, including tokens already prevalent by the landmark." ), } - write_json(output_path / "metadata.json", metadata) - manifest_rows: List[Dict[str, Any]] = [] - 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) - 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 - age_rows = [ - write_empty_age_export( + 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_dir=age_dir, + age_group=age_group, token_count=len(tokens), ) - ] - eligible_count = 0 - 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, - ) - age_rows = 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), - ) - eligible_count = len(landmark_dataset) + 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), + ) - 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) - metadata["eligible_queries_by_age"] = eligible_queries_by_age - write_json(output_path / "metadata.json", metadata) - print(f"Age {age:g}: exported {eligible_count} patient 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_metadata_json(metadata_dataset, metadata) + output_file.flush() + print(f"Age {age:g}: exported {eligible_count} patient rows") - nonfinite_shape = sum( - int(row["nonfinite_shape_values"]) for row in manifest_rows - ) - nonfinite_scale = sum( - int(row["nonfinite_scale_values"]) for row in manifest_rows - ) - metadata["total_exported_query_rows"] = sum( - int(row["n_rows"]) for row in manifest_rows - ) - metadata["nonfinite_shape_values"] = nonfinite_shape - metadata["nonfinite_scale_values"] = nonfinite_scale - metadata["complete"] = True - write_json(output_path / "metadata.json", metadata) + 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 individual Weibull parameter matrices to: {output_path}") + print(f"Saved unified Weibull parameter file to: {output_path}") if __name__ == "__main__":