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

Export disease-state Xiad trajectories

Browse Source
This commit is contained in:
Jiarui Li
2026-09-01 10:59:09 +08:00
parent d728d8585c
commit d9d1398a3a
2 changed files with 1166 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

995
export_disease_state_xiad.py Normal file
View File

@@ -0,0 +1,995 @@
"""Compute disease-state trajectories X_iad(H) from Weibull exports.
For patient ``i``, landmark age ``a``, disease ``d`` and horizon ``H``:
X_iad(H) = I(T_id <= a) + I(T_id > a) * P_iad(H)
where
P_iad(H) = 1 - exp(-((H / scale_iad) ** shape_iad)).
Diseases already observed by the landmark therefore have state 1. Diseases
not yet observed retain their predicted first-onset probability. Death is not
a disease state and is excluded from the output.
The input is the unified HDF5 file written by
``export_weibull_parameters.py``. The output preserves its landmark-row
alignment and writes one three-dimensional ``x`` dataset per age with shape
``(n_patients, n_diseases, n_horizons)``. Processing is chunked by patient row
so the multi-gigabyte parameter matrices are never loaded in full.
"""
from __future__ import annotations
import argparse
import csv
import importlib
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping, Sequence, Tuple
import numpy as np
from dataset import (
DISEASE_HISTORY_MODE_TIMED,
NO_EVENT_IDX,
normalize_disease_history_mode,
)
from eval_data import (
load_json_config,
load_sequence_eval_dataset,
validate_dataset_metadata,
validate_training_mode_config,
)
FORMAT_VERSION = 1
DEFAULT_HORIZONS = (1.0, 5.0, 10.0)
DEFAULT_COMPRESSION_LEVEL = 4
LABEL_OFFSET = NO_EVENT_IDX + 1
AXIS_MAPPING_FIELDS = (
"column",
"source_column",
"label_index",
"token_id",
"code",
"name",
"label_text",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Compute all disease states X_iad(H) from an exported Weibull "
"shape/scale HDF5 file."
)
)
parser.add_argument(
"--run_path",
required=True,
help="Run directory containing train_config.json and the Weibull export.",
)
parser.add_argument(
"--input_path",
default=None,
help=(
"HDF5 file produced by export_weibull_parameters.py. Defaults to "
"<run_path>/weibull_parameters_test_age40_80_step2.h5."
),
)
parser.add_argument(
"--output_path",
default=None,
help=(
"Output HDF5 path. Defaults to "
"<run_path>/disease_state_xiad_test_age40_80_step2.h5."
),
)
parser.add_argument(
"--mapping_output_path",
default=None,
help=(
"CSV describing the exact disease-axis order. Defaults to "
"<output_stem>_icd10_columns.csv."
),
)
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 parse_int_list(value: Any) -> List[int] | None:
if value is None:
return None
if isinstance(value, (list, tuple, np.ndarray)):
return [int(item) for item 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(item) for item in parsed]
return [int(item.strip()) for item in text.split(",") if item.strip()]
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 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 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_axis_rows(labels_file: str | Path) -> List[Dict[str, Any]]:
"""Read labels.csv using the same line-index/token convention as dataset.py."""
rows: List[Dict[str, Any]] = []
labels_path = resolve_project_file(labels_file)
if not labels_path.is_file():
raise FileNotFoundError(f"Labels file not found: {labels_path}")
with labels_path.open("r", encoding="utf-8") as handle:
for label_index, raw in enumerate(handle):
label_text = raw.strip()
if not label_text:
continue
code = label_text.split()[0]
name = label_text[len(code):].strip()
if name.startswith("(") and name.endswith(")"):
name = name[1:-1]
rows.append(
{
"label_index": int(label_index),
"token_id": int(LABEL_OFFSET + label_index),
"code": code,
"name": name or code,
"label_text": label_text,
"outcome_type": (
"death" if code.lower() == "death" else "disease"
),
}
)
if not rows:
raise ValueError(f"Labels file contains no outcomes: {labels_path}")
return rows
def load_matching_dataset(run_path: Path) -> Tuple[Any, Dict[str, Any]]:
cfg = load_json_config(run_path / "train_config.json")
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(
"X_iad(H) requires an all-future export; the run configuration "
f"uses model_target_mode={model_target_mode!r}."
)
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))
)
disease_history_mode = normalize_disease_history_mode(
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
)
dataset = load_sequence_eval_dataset(
model_target_mode=model_target_mode,
data_prefix=str(cfg.get("data_prefix", "ukb")),
labels_file=str(cfg.get("labels_file", "labels.csv")),
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)
return dataset, cfg
def validate_source_file(source_file: Any) -> None:
required_paths = (
"ages",
"tokens/column",
"tokens/token_id",
"tokens/label_code",
"tokens/label_text",
"tokens/outcome_type",
"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.")
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.")
def validate_source_token_order(
source_file: Any,
label_rows: Sequence[Mapping[str, Any]],
) -> List[Dict[str, Any]]:
"""Validate source columns against labels.csv and return the disease axis."""
source_token_ids = np.asarray(
source_file["tokens/token_id"][...], dtype=np.int64
)
source_columns = np.asarray(
source_file["tokens/column"][...], dtype=np.int64
)
expected_columns = np.arange(source_token_ids.size, dtype=np.int64)
if not np.array_equal(source_columns, expected_columns):
raise ValueError("Input HDF5 /tokens/column is not zero-based and ordered.")
source_codes = decode_strings(source_file["tokens/label_code"][...])
source_text = decode_strings(source_file["tokens/label_text"][...])
outcome_types = decode_strings(source_file["tokens/outcome_type"][...])
if source_token_ids.size != len(label_rows):
raise ValueError(
"Input HDF5 token count does not match labels.csv: "
f"{source_token_ids.size} versus {len(label_rows)}."
)
disease_rows: List[Dict[str, Any]] = []
for source_column, label_row in enumerate(label_rows):
expected_token = int(label_row["token_id"])
expected_code = str(label_row["code"])
expected_text = str(label_row["label_text"])
expected_type = str(label_row["outcome_type"])
observed = (
int(source_token_ids[source_column]),
source_codes[source_column],
source_text[source_column],
outcome_types[source_column],
)
expected = (
expected_token,
expected_code,
expected_text,
expected_type,
)
if observed != expected:
raise ValueError(
"Input HDF5 disease order does not match labels.csv at source "
f"column {source_column}: observed={observed!r}, "
f"expected={expected!r}."
)
if expected_type == "disease":
disease_rows.append(
{
"column": len(disease_rows),
"source_column": int(source_column),
"label_index": int(label_row["label_index"]),
"token_id": expected_token,
"code": expected_code,
"name": str(label_row["name"]),
"label_text": expected_text,
"outcome_type": "disease",
}
)
if not disease_rows:
raise ValueError("The validated labels.csv contains no disease outcomes.")
return disease_rows
def validate_token_table(
dataset: Any,
disease_rows: Sequence[Mapping[str, Any]],
) -> None:
for row in disease_rows:
token = int(row["token_id"])
code = str(row["code"])
dataset_code = dataset.label_id_to_code.get(token)
if dataset_code is None:
raise ValueError(f"Disease token {token} is absent from the dataset vocabulary.")
if str(dataset_code) != str(code):
raise ValueError(
f"Disease token {token} mismatch: input HDF5 has {code!r}, "
f"dataset has {dataset_code!r}."
)
def write_axis_mapping_csv(
output_path: Path,
disease_rows: Sequence[Mapping[str, Any]],
) -> None:
with output_path.open("w", encoding="utf-8", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=list(AXIS_MAPPING_FIELDS))
writer.writeheader()
for row in disease_rows:
writer.writerow({field: row[field] for field in AXIS_MAPPING_FIELDS})
def validate_axis_mapping_csv(
mapping_path: Path,
disease_rows: Sequence[Mapping[str, Any]],
) -> None:
with mapping_path.open("r", encoding="utf-8", newline="") as handle:
observed_rows = list(csv.DictReader(handle))
if len(observed_rows) != len(disease_rows):
raise RuntimeError("Disease-axis CSV row count is incorrect.")
for expected, observed in zip(disease_rows, observed_rows):
expected_text = {
field: str(expected[field]) for field in AXIS_MAPPING_FIELDS
}
if observed != expected_text:
raise RuntimeError(
"Disease-axis CSV does not match the validated labels order at "
f"output column {expected['column']}."
)
def validate_population_alignment(source_file: Any, dataset: Any) -> None:
dataset_indices = np.asarray(
source_file["test_population/dataset_index"][...], dtype=np.int64
)
eids = np.asarray(source_file["test_population/eid"][...], dtype=np.int64)
if dataset_indices.shape != eids.shape:
raise ValueError("Input test-population EIDs and dataset indices are misaligned.")
if dataset_indices.size == 0:
raise ValueError("Input test population is empty.")
if np.any(dataset_indices < 0) or np.any(dataset_indices >= len(dataset.samples)):
raise ValueError("Input test population contains an invalid dataset index.")
expected_eids = np.asarray(
[int(dataset.samples[int(index)]["eid"]) for index in dataset_indices],
dtype=np.int64,
)
if not np.array_equal(eids, expected_eids):
mismatch = int(np.flatnonzero(eids != expected_eids)[0])
raise ValueError(
"Input test population does not match the dataset loaded from the "
f"run configuration; first mismatch is at row {mismatch}."
)
def full_disease_history(sample: Mapping[str, Any]) -> Tuple[np.ndarray, np.ndarray]:
events = np.asarray(sample["event_seq"], dtype=np.int64)
times = np.asarray(sample["time_seq"], dtype=np.float32)
target_events = np.asarray(sample["target_event_seq"], dtype=np.int64)
target_times = np.asarray(sample["target_time_seq"], dtype=np.float32)
if target_events.size > 0:
events = np.concatenate([events, target_events[-1:]])
times = np.concatenate([times, target_times[-1:]])
if events.shape != times.shape:
raise ValueError("Disease events and times are misaligned in the dataset sample.")
return events, times
def build_prevalent_mask(
*,
dataset: Any,
dataset_indices: np.ndarray,
landmark_age: float,
token_to_column: Mapping[int, int],
n_diseases: int,
) -> np.ndarray:
"""Return I(T_id <= a) for a chunk of landmark rows."""
prevalent = np.zeros((len(dataset_indices), n_diseases), dtype=bool)
for row, dataset_index in enumerate(np.asarray(dataset_indices).tolist()):
sample = dataset.samples[int(dataset_index)]
events, times = full_disease_history(sample)
historical_events = events[times <= np.float32(landmark_age)]
for token in np.unique(historical_events).tolist():
column = token_to_column.get(int(token))
if column is not None:
prevalent[row, column] = True
return prevalent
def weibull_probability(
shape: np.ndarray,
scale: np.ndarray,
horizons: np.ndarray,
) -> np.ndarray:
"""Calculate stable Weibull fixed-horizon probabilities."""
shape64 = np.asarray(shape, dtype=np.float64)
scale64 = np.asarray(scale, dtype=np.float64)
horizons64 = np.asarray(horizons, dtype=np.float64)
if shape64.shape != scale64.shape or shape64.ndim != 2:
raise ValueError("shape and scale must be aligned two-dimensional matrices.")
horizons64 = validate_horizons(horizons64)
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, 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 compute_xiad(
*,
shape: np.ndarray,
scale: np.ndarray,
horizons: np.ndarray,
prevalent: np.ndarray,
) -> np.ndarray:
probability = weibull_probability(shape, scale, horizons)
prevalent = np.asarray(prevalent, dtype=bool)
if prevalent.shape != probability.shape[:2]:
raise ValueError("The prevalent mask is not aligned with shape and scale.")
return np.where(prevalent[:, :, None], 1.0, probability).astype(
np.float32, copy=False
)
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 age_group_name(age: float) -> str:
text = f"{age:g}".replace("-", "minus_").replace(".", "p")
return f"age_{text}"
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 write_string_dataset(
group: Any,
name: str,
values: Sequence[str],
string_dtype: Any,
) -> None:
group.create_dataset(
name,
data=np.asarray(list(values), dtype=object),
dtype=string_dtype,
)
def process_age_group(
*,
source_group: Any,
output_group: Any,
dataset: Any,
disease_source_columns: np.ndarray,
token_to_column: Mapping[int, int],
horizons: np.ndarray,
rows_per_chunk: int,
compression_level: int,
) -> Dict[str, Any]:
n_rows = int(source_group["eid"].shape[0])
n_diseases = int(disease_source_columns.size)
n_horizons = int(horizons.size)
landmark_age = float(source_group.attrs["age"])
expected_matrix_shape = (n_rows, int(source_group.attrs["n_tokens"]))
if source_group["shape"].shape != expected_matrix_shape:
raise ValueError(
f"Age {landmark_age:g} shape matrix has unexpected dimensions."
)
if source_group["scale"].shape != expected_matrix_shape:
raise ValueError(
f"Age {landmark_age:g} scale matrix has unexpected dimensions."
)
output_group.attrs["age"] = landmark_age
output_group.attrs["n_rows"] = n_rows
output_group.attrs["n_diseases"] = n_diseases
output_group.attrs["n_horizons"] = n_horizons
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(
"prevalent", shape=(0, n_diseases), dtype=np.uint8
)
output_group.create_dataset(
"x", shape=(0, n_diseases, n_horizons), dtype=np.float32
)
return {
"age": landmark_age,
"n_rows": 0,
"prevalent_values": 0,
"nonfinite_parameter_values": 0,
"nonfinite_x_values": 0,
}
row_chunk = min(rows_per_chunk, n_rows)
prevalent_dataset = output_group.create_dataset(
"prevalent",
shape=(n_rows, n_diseases),
dtype=np.uint8,
chunks=(row_chunk, n_diseases),
compression="gzip",
compression_opts=compression_level,
shuffle=True,
)
x_dataset = output_group.create_dataset(
"x",
shape=(n_rows, n_diseases, n_horizons),
dtype=np.float32,
chunks=(row_chunk, n_diseases, 1),
compression="gzip",
compression_opts=compression_level,
shuffle=True,
)
prevalent_values = 0
nonfinite_parameter_values = 0
nonfinite_x_values = 0
for row_slice in iter_slices(n_rows, rows_per_chunk):
dataset_indices = np.asarray(
source_group["dataset_index"][row_slice], dtype=np.int64
)
if np.any(dataset_indices < 0) or np.any(
dataset_indices >= len(dataset.samples)
):
raise ValueError(
f"Age {landmark_age:g} contains an invalid dataset index."
)
shape = np.asarray(
source_group["shape"][row_slice, disease_source_columns],
dtype=np.float32,
)
scale = np.asarray(
source_group["scale"][row_slice, disease_source_columns],
dtype=np.float32,
)
prevalent = build_prevalent_mask(
dataset=dataset,
dataset_indices=dataset_indices,
landmark_age=landmark_age,
token_to_column=token_to_column,
n_diseases=n_diseases,
)
x = compute_xiad(
shape=shape,
scale=scale,
horizons=horizons,
prevalent=prevalent,
)
prevalent_dataset[row_slice, :] = prevalent.astype(np.uint8, copy=False)
x_dataset[row_slice, :, :] = x
prevalent_values += int(prevalent.sum())
valid_parameters = (
np.isfinite(shape)
& np.isfinite(scale)
& (shape > 0.0)
& (scale > 0.0)
)
nonfinite_parameter_values += int((~valid_parameters).sum())
nonfinite_x_values += int((~np.isfinite(x)).sum())
output_group.attrs["prevalent_values"] = prevalent_values
output_group.attrs["nonfinite_parameter_values"] = nonfinite_parameter_values
output_group.attrs["nonfinite_x_values"] = nonfinite_x_values
return {
"age": landmark_age,
"n_rows": n_rows,
"prevalent_values": prevalent_values,
"nonfinite_parameter_values": nonfinite_parameter_values,
"nonfinite_x_values": nonfinite_x_values,
}
def validate_output_file(
output_file: Any,
*,
ages: np.ndarray,
n_diseases: int,
n_horizons: int,
disease_rows: Sequence[Mapping[str, Any]],
summaries: Sequence[Mapping[str, Any]],
) -> None:
for path in ("tokens", "test_population", "landmarks", "age_summary"):
if path not in output_file:
raise RuntimeError(f"Output HDF5 is missing /{path}.")
if output_file["tokens/token_id"].shape != (n_diseases,):
raise RuntimeError("Output disease token table has an unexpected length.")
expected_integer_columns = {
"column": np.asarray(
[row["column"] for row in disease_rows], dtype=np.int64
),
"source_column": np.asarray(
[row["source_column"] for row in disease_rows], dtype=np.int64
),
"label_index": np.asarray(
[row["label_index"] for row in disease_rows], dtype=np.int64
),
"token_id": np.asarray(
[row["token_id"] for row in disease_rows], dtype=np.int64
),
}
for name, expected in expected_integer_columns.items():
observed = np.asarray(output_file[f"tokens/{name}"][...], dtype=np.int64)
if not np.array_equal(observed, expected):
raise RuntimeError(f"Output /tokens/{name} order is incorrect.")
expected_string_columns = {
"label_code": [str(row["code"]) for row in disease_rows],
"name": [str(row["name"]) for row in disease_rows],
"label_text": [str(row["label_text"]) for row in disease_rows],
"outcome_type": ["disease"] * len(disease_rows),
}
for name, expected in expected_string_columns.items():
observed = decode_strings(output_file[f"tokens/{name}"][...])
if observed != expected:
raise RuntimeError(f"Output /tokens/{name} order is incorrect.")
if len(summaries) != int(ages.size):
raise RuntimeError("Output age summary 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))}"]
n_rows = int(summary["n_rows"])
if group["prevalent"].shape != (n_rows, n_diseases):
raise RuntimeError("Output prevalent matrix has unexpected dimensions.")
if group["x"].shape != (n_rows, n_diseases, n_horizons):
raise RuntimeError("Output X_iad(H) matrix has unexpected dimensions.")
def main() -> None:
args = parse_args()
run_path = Path(args.run_path).resolve()
config_path = run_path / "train_config.json"
if not config_path.is_file():
raise FileNotFoundError(config_path)
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)
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 = (
Path(args.output_path).resolve()
if args.output_path
else run_path / "disease_state_xiad_test_age40_80_step2.h5"
)
mapping_output_path = (
Path(args.mapping_output_path).resolve()
if args.mapping_output_path
else output_path.with_name(f"{output_path.stem}_icd10_columns.csv")
)
if output_path == input_path:
raise ValueError("The output path must differ from the input path.")
if mapping_output_path in {input_path, output_path}:
raise ValueError("The ICD-10 mapping path must be a separate file.")
if output_path.exists():
raise FileExistsError(
f"Output file already exists: {output_path}. Choose a new --output_path."
)
if mapping_output_path.exists():
raise FileExistsError(
"Disease-axis mapping already exists: "
f"{mapping_output_path}. Choose a new --mapping_output_path."
)
output_path.parent.mkdir(parents=True, exist_ok=True)
mapping_output_path.parent.mkdir(parents=True, exist_ok=True)
temporary_output = output_path.with_name(f".{output_path.name}.partial")
temporary_mapping = mapping_output_path.with_name(
f".{mapping_output_path.name}.partial"
)
if temporary_output.exists():
raise FileExistsError(
f"Partial output already exists: {temporary_output}. Remove or rename it."
)
if temporary_mapping.exists():
raise FileExistsError(
f"Partial mapping already exists: {temporary_mapping}. Remove or rename it."
)
h5py = require_h5py()
print(f"Loading dataset from run configuration: {run_path}")
dataset, cfg = load_matching_dataset(run_path)
labels_file = str(cfg.get("labels_file", "labels.csv"))
labels_path = resolve_project_file(labels_file).resolve()
label_rows = load_label_axis_rows(labels_path)
with h5py.File(input_path, "r") as source_file:
validate_source_file(source_file)
validate_population_alignment(source_file, dataset)
disease_rows = validate_source_token_order(source_file, label_rows)
validate_token_table(dataset, disease_rows)
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.")
disease_source_columns = np.asarray(
[row["source_column"] for row in disease_rows], dtype=np.int64
)
disease_label_indices = np.asarray(
[row["label_index"] for row in disease_rows], dtype=np.int64
)
disease_token_ids = np.asarray(
[row["token_id"] for row in disease_rows], dtype=np.int64
)
disease_codes = [str(row["code"]) for row in disease_rows]
disease_names = [str(row["name"]) for row in disease_rows]
disease_text = [str(row["label_text"]) for row in disease_rows]
token_to_column = {
int(token): column
for column, token in enumerate(disease_token_ids.tolist())
}
metadata: Dict[str, Any] = {
"format_version": FORMAT_VERSION,
"complete": False,
"definition": (
"X_iad(H) = I(T_id <= a) + I(T_id > a) * "
"[1 - exp(-((H / scale_iad)^shape_iad))]"
),
"source_weibull_path": str(input_path),
"source_run_path": str(run_path),
"labels_file": str(labels_path),
"disease_axis_mapping_csv": str(mapping_output_path),
"disease_axis_order": (
"Exactly labels.csv line order with Death removed; output "
"column is the second dimension of every landmark x dataset."
),
"disease_axis_validated_against_labels": True,
"ages": [float(value) for value in ages.tolist()],
"horizons_years": [float(value) for value in horizons.tolist()],
"n_diseases": int(disease_token_ids.size),
"matrix_dtype": "float32",
"prevalent_dtype": "uint8",
"death_excluded": True,
"prevalence_boundary": "first observed disease time <= landmark age",
"hdf5_layout": {
"tokens": (
"/tokens/{column,source_column,label_index,token_id,"
"label_code,name,label_text,outcome_type}"
),
"horizons": "/horizons",
"test_population": "/test_population/{eid,dataset_index}",
"landmarks": (
"/landmarks/age_*/{eid,dataset_index,sex,age,prevalent,x}"
),
"x_dimensions": ["landmark_row", "disease", "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["n_diseases"] = int(disease_token_ids.size)
output_file.attrs["n_horizons"] = int(horizons.size)
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)
)
token_group = output_file.create_group("tokens")
token_group.create_dataset(
"column",
data=np.arange(disease_token_ids.size, dtype=np.int64),
)
token_group.create_dataset(
"source_column", data=disease_source_columns
)
token_group.create_dataset(
"label_index", data=disease_label_indices
)
token_group.create_dataset("token_id", data=disease_token_ids)
write_string_dataset(
token_group,
"label_code",
disease_codes,
string_dtype,
)
write_string_dataset(
token_group,
"name",
disease_names,
string_dtype,
)
write_string_dataset(
token_group,
"label_text",
disease_text,
string_dtype,
)
write_string_dataset(
token_group,
"outcome_type",
["disease"] * len(disease_rows),
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,
dataset=dataset,
disease_source_columns=disease_source_columns,
token_to_column=token_to_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']} rows, "
f"prevalent states={summary['prevalent_values']}"
)
age_summary = output_file.create_group("age_summary")
for name, dtype in (
("age", np.float32),
("n_rows", np.int64),
("prevalent_values", np.int64),
("nonfinite_parameter_values", np.int64),
("nonfinite_x_values", np.int64),
):
age_summary.create_dataset(
name,
data=np.asarray([row[name] for row in summaries], dtype=dtype),
)
metadata["total_exported_query_rows"] = sum(
int(row["n_rows"]) for row in summaries
)
metadata["total_prevalent_values"] = sum(
int(row["prevalent_values"]) for row in summaries
)
metadata["nonfinite_parameter_values"] = sum(
int(row["nonfinite_parameter_values"]) for row in summaries
)
metadata["nonfinite_x_values"] = sum(
int(row["nonfinite_x_values"]) for row in summaries
)
validate_output_file(
output_file,
ages=ages,
n_diseases=int(disease_token_ids.size),
n_horizons=int(horizons.size),
disease_rows=disease_rows,
summaries=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()
write_axis_mapping_csv(temporary_mapping, disease_rows)
validate_axis_mapping_csv(temporary_mapping, disease_rows)
temporary_output.replace(output_path)
temporary_mapping.replace(mapping_output_path)
print(f"Saved disease-state X_iad(H) file to: {output_path}")
print(f"Saved disease-axis ICD-10 mapping to: {mapping_output_path}")
if __name__ == "__main__":
main()

171
tests/test_export_disease_state_xiad.py Normal file
View File

@@ -0,0 +1,171 @@
import tempfile
import unittest
from pathlib import Path
import numpy as np
from export_disease_state_xiad import (
build_prevalent_mask,
compute_xiad,
load_label_axis_rows,
validate_axis_mapping_csv,
validate_horizons,
validate_source_token_order,
weibull_probability,
write_axis_mapping_csv,
)
class _DummyDataset:
def __init__(self) -> None:
self.samples = [
{
"eid": 101,
"event_seq": np.asarray([3, 4], dtype=np.int64),
"time_seq": np.asarray([1.0, 2.0], dtype=np.float32),
"target_event_seq": np.asarray([4, 5], dtype=np.int64),
"target_time_seq": np.asarray([2.0, 4.0], dtype=np.float32),
},
{
"eid": 102,
"event_seq": np.asarray([3], dtype=np.int64),
"time_seq": np.asarray([1.5], dtype=np.float32),
"target_event_seq": np.asarray([6], dtype=np.int64),
"target_time_seq": np.asarray([5.0], dtype=np.float32),
},
]
class DiseaseStateXiadTests(unittest.TestCase):
def test_label_order_defines_disease_axis_and_excludes_death(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
labels_path = Path(tmp_dir) / "labels.csv"
labels_path.write_text(
"A00 (cholera)\nI10 Essential hypertension\nDeath\n",
encoding="utf-8",
)
label_rows = load_label_axis_rows(labels_path)
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 Essential hypertension", b"Death"]
),
"tokens/outcome_type": np.asarray(
[b"disease", b"disease", b"death"]
),
}
disease_rows = validate_source_token_order(source, label_rows)
self.assertEqual([row["column"] for row in disease_rows], [0, 1])
self.assertEqual([row["source_column"] for row in disease_rows], [0, 1])
self.assertEqual([row["label_index"] for row in disease_rows], [0, 1])
self.assertEqual([row["token_id"] for row in disease_rows], [3, 4])
self.assertEqual([row["code"] for row in disease_rows], ["A00", "I10"])
with tempfile.TemporaryDirectory() as tmp_dir:
mapping_path = Path(tmp_dir) / "icd10_columns.csv"
write_axis_mapping_csv(mapping_path, disease_rows)
validate_axis_mapping_csv(mapping_path, disease_rows)
header = mapping_path.read_text(encoding="utf-8").splitlines()[0]
self.assertEqual(
header,
"column,source_column,label_index,token_id,code,name,label_text",
)
def test_source_order_mismatch_is_rejected(self) -> None:
label_rows = [
{
"label_index": 0,
"token_id": 3,
"code": "A00",
"name": "cholera",
"label_text": "A00 (cholera)",
"outcome_type": "disease",
}
]
source = {
"tokens/column": np.asarray([0], dtype=np.int64),
"tokens/token_id": np.asarray([3], dtype=np.int64),
"tokens/label_code": np.asarray([b"A01"]),
"tokens/label_text": np.asarray([b"A01 wrong order"]),
"tokens/outcome_type": np.asarray([b"disease"]),
}
with self.assertRaises(ValueError):
validate_source_token_order(source, label_rows)
def test_weibull_probability_combines_shape_and_scale(self) -> None:
shape = np.asarray([[2.0, 1.0]], dtype=np.float32)
scale = np.asarray([[10.0, 4.0]], dtype=np.float32)
result = weibull_probability(
shape,
scale,
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_prevalent_disease_state_is_one(self) -> None:
result = compute_xiad(
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], dtype=np.float64),
prevalent=np.asarray([[True, False]]),
)
self.assertEqual(float(result[0, 0, 0]), 1.0)
self.assertAlmostEqual(
float(result[0, 1, 0]),
1.0 - np.exp(-1.25),
places=6,
)
def test_invalid_nonprevalent_parameter_produces_nan(self) -> None:
result = compute_xiad(
shape=np.asarray([[1.0, 1.0]], dtype=np.float32),
scale=np.asarray([[np.nan, -1.0]], dtype=np.float32),
horizons=np.asarray([5.0], dtype=np.float64),
prevalent=np.asarray([[True, False]]),
)
self.assertEqual(float(result[0, 0, 0]), 1.0)
self.assertTrue(np.isnan(result[0, 1, 0]))
def test_prevalence_uses_disease_time_at_or_before_landmark(self) -> None:
result = build_prevalent_mask(
dataset=_DummyDataset(),
dataset_indices=np.asarray([0, 1], dtype=np.int64),
landmark_age=2.0,
token_to_column={3: 0, 4: 1, 5: 2, 6: 3},
n_diseases=4,
)
expected = np.asarray(
[
[True, True, False, False],
[True, False, False, False],
]
)
np.testing.assert_array_equal(result, expected)
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()