From 09d7415b3a7f0d255ce1d904478ed6d47019299b Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Tue, 1 Sep 2026 13:47:51 +0800 Subject: [PATCH] Add conditional death burden export --- export_death_burden.py | 522 ++++++++++++++++++++++++++++++ tests/test_export_death_burden.py | 59 ++++ 2 files changed, 581 insertions(+) create mode 100644 export_death_burden.py create mode 100644 tests/test_export_death_burden.py diff --git a/export_death_burden.py b/export_death_burden.py new file mode 100644 index 0000000..4e7d146 --- /dev/null +++ b/export_death_burden.py @@ -0,0 +1,522 @@ +"""Compute conditional fixed-horizon death burden from Weibull exports. + +For patient ``i``, landmark age ``a`` and horizon ``H``: + + B_ia(H) = P(T_death <= H | alive at age a) + = 1 - exp(-((H / scale_ia) ** shape_ia)). + +The input is the unified HDF5 file written by +``export_weibull_parameters.py``. That source file contains only patients from +its configured test EID file, and each landmark group contains patients alive +and otherwise eligible at that age. The output preserves those rows exactly. +""" + +from __future__ import annotations + +import argparse +import importlib +import json +from pathlib import Path +from typing import Any, Dict, Iterable, List, Sequence + +import numpy as np + + +FORMAT_VERSION = 1 +DEFAULT_HORIZONS = (1.0, 5.0, 10.0) +DEFAULT_COMPRESSION_LEVEL = 4 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Compute conditional death probabilities at fixed horizons from " + "an exported Weibull shape/scale HDF5 file." + ) + ) + parser.add_argument( + "--run_path", + required=True, + help="Run directory containing the Weibull export.", + ) + parser.add_argument( + "--input_path", + default=None, + help=( + "HDF5 file produced by export_weibull_parameters.py. Defaults to " + "/weibull_parameters_test_age40_80_step2.h5." + ), + ) + parser.add_argument( + "--output_path", + default=None, + help=( + "Output HDF5 path. Defaults to " + "/death_burden_test_age40_80_step2.h5." + ), + ) + parser.add_argument( + "--horizons", + nargs="+", + type=float, + default=list(DEFAULT_HORIZONS), + help="Positive horizons in years. Default: 1 5 10.", + ) + parser.add_argument( + "--rows_per_chunk", + type=int, + default=256, + help="Patient rows processed and stored per HDF5 chunk. Default: 256.", + ) + parser.add_argument( + "--compression_level", + type=int, + default=DEFAULT_COMPRESSION_LEVEL, + help="Gzip compression level from 0 to 9. Default: 4.", + ) + return parser.parse_args() + + +def require_h5py() -> Any: + try: + return importlib.import_module("h5py") + except ImportError as exc: + raise RuntimeError( + "This script requires h5py in the project Miniconda environment." + ) from exc + + +def decode_strings(values: np.ndarray) -> List[str]: + result: List[str] = [] + for value in np.asarray(values).tolist(): + if isinstance(value, bytes): + result.append(value.decode("utf-8")) + else: + result.append(str(value)) + return result + + +def validate_horizons(values: Iterable[float]) -> np.ndarray: + horizons = np.asarray(list(values), dtype=np.float64) + if horizons.ndim != 1 or horizons.size == 0: + raise ValueError("At least one horizon is required.") + if not np.all(np.isfinite(horizons)) or np.any(horizons <= 0.0): + raise ValueError("Every horizon must be finite and greater than zero.") + if np.unique(horizons).size != horizons.size: + raise ValueError("Horizons must not contain duplicates.") + return horizons + + +def age_group_name(age: float) -> str: + text = f"{age:g}".replace("-", "minus_").replace(".", "p") + return f"age_{text}" + + +def iter_slices(n_rows: int, rows_per_chunk: int) -> Iterable[slice]: + for start in range(0, n_rows, rows_per_chunk): + yield slice(start, min(start + rows_per_chunk, n_rows)) + + +def find_death_column(source_file: Any) -> int: + required_paths = ( + "tokens/column", + "tokens/token_id", + "tokens/label_code", + "tokens/label_text", + "tokens/outcome_type", + ) + missing = [path for path in required_paths if path not in source_file] + if missing: + raise ValueError(f"Input HDF5 is missing required paths: {missing}") + + token_count = int(source_file["tokens/token_id"].shape[0]) + for path in ( + "tokens/column", + "tokens/label_code", + "tokens/label_text", + "tokens/outcome_type", + ): + if source_file[path].shape != (token_count,): + raise ValueError(f"/{path} is not aligned with /tokens/token_id.") + columns = np.asarray(source_file["tokens/column"][...], dtype=np.int64) + if not np.array_equal(columns, np.arange(token_count, dtype=np.int64)): + raise ValueError("Input HDF5 /tokens/column is not zero-based and ordered.") + + codes = decode_strings(source_file["tokens/label_code"][...]) + outcome_types = decode_strings(source_file["tokens/outcome_type"][...]) + death_columns = [ + column + for column, (code, outcome_type) in enumerate(zip(codes, outcome_types)) + if code.lower() == "death" and outcome_type.lower() == "death" + ] + if len(death_columns) != 1: + raise ValueError( + "Input HDF5 must contain exactly one token labelled Death with " + f"outcome_type='death'; found {len(death_columns)}." + ) + return int(death_columns[0]) + + +def death_probability( + shape: np.ndarray, + scale: np.ndarray, + horizons: np.ndarray, +) -> np.ndarray: + """Return P(death within H | alive at the landmark).""" + shape64 = np.asarray(shape, dtype=np.float64).reshape(-1) + scale64 = np.asarray(scale, dtype=np.float64).reshape(-1) + horizons64 = validate_horizons(horizons) + if shape64.shape != scale64.shape: + raise ValueError("Death shape and scale vectors are not aligned.") + + valid = ( + np.isfinite(shape64) + & np.isfinite(scale64) + & (shape64 > 0.0) + & (scale64 > 0.0) + ) + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + log_cumulative_hazard = shape64[:, None] * ( + np.log(horizons64)[None, :] - np.log(scale64)[:, None] + ) + cumulative_hazard = np.exp(np.clip(log_cumulative_hazard, -87.0, 40.0)) + probability = -np.expm1(-cumulative_hazard) + probability[~valid, :] = np.nan + return probability.astype(np.float32, copy=False) + + +def copy_vector_dataset( + source_group: Any, + output_group: Any, + name: str, + rows_per_chunk: int, + compression_level: int, +) -> None: + source = source_group[name] + n_rows = int(source.shape[0]) + options: Dict[str, Any] = {} + if n_rows > 0: + options = { + "chunks": (min(rows_per_chunk, n_rows),), + "compression": "gzip", + "compression_opts": compression_level, + "shuffle": True, + } + output_group.create_dataset(name, data=source[...], dtype=source.dtype, **options) + + +def process_age_group( + *, + source_group: Any, + output_group: Any, + death_column: int, + horizons: np.ndarray, + rows_per_chunk: int, + compression_level: int, +) -> Dict[str, Any]: + landmark_age = float(source_group.attrs["age"]) + n_rows = int(source_group["eid"].shape[0]) + n_tokens = int(source_group.attrs["n_tokens"]) + expected_shape = (n_rows, n_tokens) + for name in ("shape", "scale"): + if source_group[name].shape != expected_shape: + raise ValueError( + f"Age {landmark_age:g} {name} matrix has unexpected dimensions." + ) + + output_group.attrs["age"] = landmark_age + output_group.attrs["n_rows"] = n_rows + output_group.attrs["n_horizons"] = int(horizons.size) + for name in ("eid", "dataset_index", "sex", "age"): + copy_vector_dataset( + source_group, + output_group, + name, + rows_per_chunk, + compression_level, + ) + + if n_rows == 0: + output_group.create_dataset( + "death_burden", shape=(0, horizons.size), dtype=np.float32 + ) + return { + "age": landmark_age, + "n_rows": 0, + "nonfinite_parameter_values": 0, + "nonfinite_burden_values": 0, + } + + row_chunk = min(rows_per_chunk, n_rows) + burden_dataset = output_group.create_dataset( + "death_burden", + shape=(n_rows, horizons.size), + dtype=np.float32, + chunks=(row_chunk, 1), + compression="gzip", + compression_opts=compression_level, + shuffle=True, + ) + + nonfinite_parameters = 0 + nonfinite_burden = 0 + for row_slice in iter_slices(n_rows, rows_per_chunk): + shape = np.asarray( + source_group["shape"][row_slice, death_column], dtype=np.float32 + ) + scale = np.asarray( + source_group["scale"][row_slice, death_column], dtype=np.float32 + ) + burden = death_probability(shape, scale, horizons) + burden_dataset[row_slice, :] = burden + valid_parameters = ( + np.isfinite(shape) + & np.isfinite(scale) + & (shape > 0.0) + & (scale > 0.0) + ) + nonfinite_parameters += int((~valid_parameters).sum()) + nonfinite_burden += int((~np.isfinite(burden)).sum()) + + output_group.attrs["nonfinite_parameter_values"] = nonfinite_parameters + output_group.attrs["nonfinite_burden_values"] = nonfinite_burden + return { + "age": landmark_age, + "n_rows": n_rows, + "nonfinite_parameter_values": nonfinite_parameters, + "nonfinite_burden_values": nonfinite_burden, + } + + +def validate_source_file(source_file: Any) -> None: + required_paths = ( + "ages", + "test_population/eid", + "test_population/dataset_index", + "landmarks", + ) + missing = [path for path in required_paths if path not in source_file] + if missing: + raise ValueError(f"Input HDF5 is missing required paths: {missing}") + if not bool(source_file.attrs.get("complete", False)): + raise ValueError("Input HDF5 is not marked complete.") + + +def validate_output_file( + output_file: Any, + *, + ages: np.ndarray, + horizons: np.ndarray, + summaries: Sequence[Dict[str, Any]], +) -> None: + for path in ( + "ages", + "horizons", + "death_token", + "test_population", + "landmarks", + "age_summary", + ): + if path not in output_file: + raise RuntimeError(f"Output HDF5 is missing /{path}.") + if len(summaries) != int(ages.size): + raise RuntimeError("Output age summary length does not match the age grid.") + for age_value, summary in zip(ages.tolist(), summaries): + group = output_file[f"landmarks/{age_group_name(float(age_value))}"] + expected_shape = (int(summary["n_rows"]), int(horizons.size)) + if group["death_burden"].shape != expected_shape: + raise RuntimeError( + f"Death burden matrix has {group['death_burden'].shape}, " + f"expected {expected_shape}." + ) + + +def main() -> None: + args = parse_args() + run_path = Path(args.run_path).resolve() + input_path = ( + Path(args.input_path).resolve() + if args.input_path + else run_path / "weibull_parameters_test_age40_80_step2.h5" + ) + if not input_path.is_file(): + raise FileNotFoundError(input_path) + output_path = ( + Path(args.output_path).resolve() + if args.output_path + else run_path / "death_burden_test_age40_80_step2.h5" + ) + if output_path == input_path: + raise ValueError("The output path must differ from the input path.") + if output_path.exists(): + raise FileExistsError( + f"Output file already exists: {output_path}. Choose a new --output_path." + ) + horizons = validate_horizons(args.horizons) + if args.rows_per_chunk <= 0: + raise ValueError("rows_per_chunk must be greater than zero.") + if not 0 <= args.compression_level <= 9: + raise ValueError("compression_level must be between 0 and 9.") + + 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." + ) + + h5py = require_h5py() + with h5py.File(input_path, "r") as source_file: + validate_source_file(source_file) + death_column = find_death_column(source_file) + ages = np.asarray(source_file["ages"][...], dtype=np.float32) + if ages.ndim != 1 or ages.size == 0: + raise ValueError("The input age grid is empty or invalid.") + + death_token_id = int(source_file["tokens/token_id"][death_column]) + death_code = decode_strings( + source_file["tokens/label_code"][death_column : death_column + 1] + )[0] + death_text = decode_strings( + source_file["tokens/label_text"][death_column : death_column + 1] + )[0] + metadata: Dict[str, Any] = { + "format_version": FORMAT_VERSION, + "complete": False, + "definition": ( + "B_ia(H) = P(T_death <= H | alive at landmark age a) = " + "1 - exp(-((H / scale_ia)^shape_ia))" + ), + "source_weibull_path": str(input_path), + "source_run_path": str(run_path), + "source_population": "test_population", + "death_source_column": death_column, + "death_token_id": death_token_id, + "ages": [float(value) for value in ages.tolist()], + "horizons_years": [float(value) for value in horizons.tolist()], + "matrix_dtype": "float32", + "hdf5_layout": { + "death_token": "/death_token/{source_column,token_id,label_code,label_text}", + "test_population": "/test_population/{eid,dataset_index}", + "landmarks": ( + "/landmarks/age_*/{eid,dataset_index,sex,age,death_burden}" + ), + "death_burden_dimensions": ["landmark_row", "horizon"], + }, + "rows_per_chunk": int(args.rows_per_chunk), + "compression": "gzip", + "compression_level": int(args.compression_level), + } + + 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["source_weibull_path"] = str(input_path) + output_file.attrs["run_path"] = str(run_path) + output_file.attrs["source_population"] = "test" + output_file.attrs["matrix_dtype"] = "float32" + metadata_dataset = output_file.create_dataset( + "metadata_json", shape=(), dtype=string_dtype + ) + metadata_dataset[()] = json.dumps( + metadata, ensure_ascii=False, separators=(",", ":") + ) + output_file.create_dataset("ages", data=ages) + output_file.create_dataset( + "horizons", data=horizons.astype(np.float32) + ) + + death_group = output_file.create_group("death_token") + death_group.create_dataset( + "source_column", data=np.asarray(death_column, dtype=np.int64) + ) + death_group.create_dataset( + "token_id", data=np.asarray(death_token_id, dtype=np.int64) + ) + death_group.create_dataset( + "label_code", data=death_code, dtype=string_dtype + ) + death_group.create_dataset( + "label_text", data=death_text, dtype=string_dtype + ) + + population_group = output_file.create_group("test_population") + for name in ("eid", "dataset_index"): + copy_vector_dataset( + source_file["test_population"], + population_group, + name, + args.rows_per_chunk, + args.compression_level, + ) + + landmark_root = output_file.create_group("landmarks") + for age_value in ages.tolist(): + age = float(age_value) + group_name = age_group_name(age) + source_path = f"landmarks/{group_name}" + if source_path not in source_file: + raise ValueError(f"Input HDF5 is missing /{source_path}.") + output_group = landmark_root.create_group(group_name) + summary = process_age_group( + source_group=source_file[source_path], + output_group=output_group, + death_column=death_column, + horizons=horizons, + rows_per_chunk=args.rows_per_chunk, + compression_level=args.compression_level, + ) + summaries.append(summary) + metadata["completed_ages"] = [ + float(row["age"]) for row in summaries + ] + metadata_dataset[()] = json.dumps( + metadata, ensure_ascii=False, separators=(",", ":") + ) + output_file.flush() + print(f"Age {age:g}: wrote {summary['n_rows']} death-burden rows") + + summary_group = output_file.create_group("age_summary") + for name, dtype in ( + ("age", np.float32), + ("n_rows", np.int64), + ("nonfinite_parameter_values", np.int64), + ("nonfinite_burden_values", np.int64), + ): + summary_group.create_dataset( + name, + data=np.asarray([row[name] for row in summaries], dtype=dtype), + ) + + validate_output_file( + output_file, + ages=ages, + horizons=horizons, + summaries=summaries, + ) + metadata["total_exported_query_rows"] = sum( + int(row["n_rows"]) for row in summaries + ) + metadata["nonfinite_parameter_values"] = sum( + int(row["nonfinite_parameter_values"]) for row in summaries + ) + metadata["nonfinite_burden_values"] = sum( + int(row["nonfinite_burden_values"]) for row in summaries + ) + metadata["validated"] = True + metadata["complete"] = True + metadata_dataset[()] = json.dumps( + metadata, ensure_ascii=False, separators=(",", ":") + ) + output_file.attrs["validated"] = True + output_file.attrs.modify("complete", True) + output_file.flush() + + temporary_output.replace(output_path) + print(f"Saved conditional death-burden file to: {output_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_export_death_burden.py b/tests/test_export_death_burden.py new file mode 100644 index 0000000..257642e --- /dev/null +++ b/tests/test_export_death_burden.py @@ -0,0 +1,59 @@ +import unittest + +import numpy as np + +from export_death_burden import ( + death_probability, + find_death_column, + validate_horizons, +) + + +class DeathBurdenTests(unittest.TestCase): + def test_death_column_is_selected_by_code_and_outcome_type(self) -> None: + source = { + "tokens/column": np.asarray([0, 1, 2], dtype=np.int64), + "tokens/token_id": np.asarray([3, 4, 5], dtype=np.int64), + "tokens/label_code": np.asarray([b"A00", b"I10", b"Death"]), + "tokens/label_text": np.asarray( + [b"A00 cholera", b"I10 hypertension", b"Death"] + ), + "tokens/outcome_type": np.asarray( + [b"disease", b"disease", b"death"] + ), + } + + self.assertEqual(find_death_column(source), 2) + + def test_death_probability_combines_shape_and_scale(self) -> None: + result = death_probability( + shape=np.asarray([2.0, 1.0], dtype=np.float32), + scale=np.asarray([10.0, 4.0], dtype=np.float32), + horizons=np.asarray([5.0, 10.0], dtype=np.float64), + ) + expected = np.asarray( + [ + [1.0 - np.exp(-0.25), 1.0 - np.exp(-1.0)], + [1.0 - np.exp(-1.25), 1.0 - np.exp(-2.5)], + ], + dtype=np.float32, + ) + np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-7) + + def test_invalid_parameters_produce_nan(self) -> None: + result = death_probability( + shape=np.asarray([1.0, -1.0], dtype=np.float32), + scale=np.asarray([np.nan, 2.0], dtype=np.float32), + horizons=np.asarray([5.0], dtype=np.float64), + ) + self.assertTrue(np.isnan(result).all()) + + def test_horizons_must_be_positive_and_unique(self) -> None: + with self.assertRaises(ValueError): + validate_horizons([0.0, 5.0]) + with self.assertRaises(ValueError): + validate_horizons([5.0, 5.0]) + + +if __name__ == "__main__": + unittest.main()