"""Evaluate extra-info attribution to death parameters and future disease risks. For each landmark query, this script scans selected extra-info types that are available at or before the query age. For each such type it re-runs the model with that extra-info type removed and summarizes: * death distribution parameters before and after ablation; * disease distribution parameters before and after ablation, by ICD-10 chapter-derived organ/system groups. Death is always token vocab_size - 1. """ from __future__ import annotations import argparse import json import re from pathlib import Path from typing import Any, Sequence import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils.data import DataLoader from tqdm.auto import tqdm from evaluate_auc_v2 import ( build_model_from_dataset, cfg_get, load_checkpoint_state_dict, load_json_config, load_model_state, resolve_dist_mode_for_checkpoint, resolve_eval_device, validate_dataset_metadata, ) from landmark_eval_utils import ( IndexedLandmarkDataset, LandmarkDataset, build_first_occurrence_maps_for_landmarks, collate_indexed_landmark_fn, infer_landmark_hidden, load_eval_sequence_dataset, load_organ_groups, make_landmark_ages, ) EXTRA_KEY_COLUMNS = [ "selected_extra_info_type_id", "selected_extra_info_var_name", "selected_extra_info_full_name", "landmark_age", "sex", ] DEATH_PARAMETER_COLUMNS = [ "original_death_lambda", "ablated_death_lambda", "original_death_scale", "ablated_death_scale", "original_death_shape", "ablated_death_shape", ] DISEASE_PARAMETER_KEY_COLUMNS = [ *EXTRA_KEY_COLUMNS, "target_group", "target_group_label", ] DISEASE_PARAMETER_COLUMNS = [ "original_disease_lambda", "ablated_disease_lambda", "original_disease_scale", "ablated_disease_scale", "original_disease_shape", "ablated_disease_shape", ] 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 text == "": return None if text.startswith("["): raw = json.loads(text) if not isinstance(raw, list): raise ValueError("Expected JSON list for integer list") return [int(x) for x in raw] return [int(x.strip()) for x in re.split(r"[,;\s]+", text) if x.strip()] def load_extra_info_metadata( *, dataset_extra_info_types: Sequence[int], search_root: Path = Path("."), ) -> dict[int, dict[str, Any]]: metadata: dict[int, dict[str, Any]] = { int(type_id): { "type_id": int(type_id), "var_name": f"extra_info_{int(type_id)}", "full_name": f"extra-info type {int(type_id)}", } for type_id in dataset_extra_info_types } line_re = re.compile(r"^\s*(\d+)\s*#\s*([^|#]+?)(?:\s*\|\s*(.*?))?\s*$") for path in sorted(search_root.glob("extra_info_types*.txt")): for line in path.read_text(encoding="utf-8").splitlines(): match = line_re.match(line) if not match: continue type_id = int(match.group(1)) if type_id not in metadata: continue var_name = match.group(2).strip() full_name = (match.group(3) or var_name).strip() metadata[type_id] = { "type_id": type_id, "var_name": var_name, "full_name": full_name, } return metadata def resolve_extra_info_types( value: str | None, *, dataset_extra_info_types: Sequence[int], metadata: dict[int, dict[str, Any]], ) -> list[int]: available = [int(x) for x in dataset_extra_info_types] if value is None or str(value).strip() == "": return available out: list[int] = [] seen: set[int] = set() by_var = { str(meta.get("var_name", "")).lower(): int(type_id) for type_id, meta in metadata.items() } by_full = { str(meta.get("full_name", "")).lower(): int(type_id) for type_id, meta in metadata.items() } for part in re.split(r"[,;]+", str(value)): text = part.strip() if not text: continue if text.isdigit() or (text.startswith("-") and text[1:].isdigit()): type_id = int(text) else: lower = text.lower() if lower in by_var: type_id = by_var[lower] elif lower in by_full: type_id = by_full[lower] else: matches = [ int(t) for t, meta in metadata.items() if lower in str(meta.get("var_name", "")).lower() or lower in str(meta.get("full_name", "")).lower() ] if len(matches) != 1: raise ValueError( f"--extra_info={text!r} matched {len(matches)} types; " "use a type id or exact variable name." ) type_id = matches[0] if type_id not in available: raise ValueError( f"extra-info type {type_id} is not available in this dataset/run" ) if type_id not in seen: out.append(type_id) seen.add(type_id) return out def death_distribution_parameters( model, hidden: torch.Tensor, *, dist_mode: str, eps: float = 1e-8, ) -> tuple[str, torch.Tensor]: logits = model.calc_risk(hidden) death_idx = int(logits.shape[1]) - 1 death_lambda = F.softplus(logits[:, death_idx]) + float(eps) if dist_mode == "exponential": nan = torch.full_like(death_lambda, float("nan")) return "exponential", torch.stack([death_lambda, nan, nan], dim=1) if dist_mode == "weibull": rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype) elif dist_mode == "mixed": rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype) else: raise ValueError(f"Unsupported dist_mode={dist_mode!r}") shape = rho.clamp_min(float(eps)) scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape) nan = torch.full_like(death_lambda, float("nan")) return "weibull", torch.stack([nan, scale, shape], dim=1) def disease_distribution_parameters( model, hidden: torch.Tensor, *, token_ids: Sequence[int], dist_mode: str, logits: torch.Tensor | None = None, rho: torch.Tensor | None = None, eps: float = 1e-8, ) -> tuple[str, torch.Tensor]: ids = [int(x) for x in token_ids] if not ids: empty = hidden.new_empty((hidden.shape[0], 0, 3)) return "none", empty all_logits = model.calc_risk(hidden) if logits is None else logits disease_lambda = F.softplus(all_logits[:, ids]) + float(eps) if dist_mode in {"exponential", "mixed"}: nan = torch.full_like(disease_lambda, float("nan")) return "exponential", torch.stack([disease_lambda, nan, nan], dim=2) if dist_mode == "weibull": all_rho = model.calc_weibull_rho(hidden) if rho is None else rho shape = all_rho[:, ids].to(dtype=disease_lambda.dtype).clamp_min(float(eps)) scale = torch.pow(disease_lambda.clamp_min(float(eps)), -1.0 / shape) nan = torch.full_like(disease_lambda, float("nan")) return "weibull", torch.stack([nan, scale, shape], dim=2) raise ValueError(f"Unsupported dist_mode={dist_mode!r}") def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: return torch.stack( [ original[:, 0], ablated[:, 0], original[:, 1], ablated[:, 1], original[:, 2], ablated[:, 2], ], dim=1, ) def disease_parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: return torch.stack( [ original[:, :, 0], ablated[:, :, 0], original[:, :, 1], ablated[:, :, 1], original[:, :, 2], ablated[:, :, 2], ], dim=2, ) def build_extra_info_ablated_slice( batch: dict[str, torch.Tensor], *, row_indices: torch.Tensor, extra_info_type_id: int, ) -> dict[str, torch.Tensor]: out: dict[str, torch.Tensor] = {} repeated_keys = ( "event_seq", "time_seq", "padding_mask", "readout_mask", "sex", "landmark_pos", "t_query", "patient_id", "landmark_age", "followup_end_time", "death_time", "row_idx", ) for key in repeated_keys: out[key] = batch[key][row_indices] out["other_type"] = batch["other_type"][row_indices].clone() out["other_value"] = batch["other_value"][row_indices].clone() out["other_value_kind"] = batch["other_value_kind"][row_indices].clone() out["other_time"] = batch["other_time"][row_indices].clone() remove = out["other_type"] == int(extra_info_type_id) out["other_type"][remove] = 0 out["other_value"][remove] = 0 out["other_value_kind"][remove] = 0 out["other_time"][remove] = 0 return out def concat_tensor_batches(chunks: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: return {key: torch.cat([chunk[key] for chunk in chunks], dim=0) for key in chunks[0]} def iter_extra_info_ablated_batches( batch: dict[str, torch.Tensor], *, selected_extra_info_types: Sequence[int], max_batch_size: int, ): pending_batches: list[dict[str, torch.Tensor]] = [] pending_types: list[int] = [] pending_rows: list[int] = [] pending_n = 0 other_type = batch["other_type"] visible = other_type > 0 visible &= batch["other_time"] <= batch["t_query"][:, None].to(batch["other_time"].dtype) for type_id in selected_extra_info_types: active_rows = torch.nonzero( ((other_type == int(type_id)) & visible).any(dim=1), as_tuple=False, ).flatten() if active_rows.numel() == 0: continue row_offset = 0 while row_offset < int(active_rows.numel()): capacity = int(max_batch_size) - pending_n row_stop = min(int(active_rows.numel()), row_offset + capacity) row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device) chunk = build_extra_info_ablated_slice( batch, row_indices=row_indices, extra_info_type_id=int(type_id), ) chunk_n = int(row_indices.numel()) pending_batches.append(chunk) pending_types.extend([int(type_id)] * chunk_n) pending_rows.extend(int(x) for x in row_indices.detach().cpu().tolist()) pending_n += chunk_n row_offset = row_stop if pending_n >= int(max_batch_size): yield concat_tensor_batches(pending_batches), pending_types, pending_rows pending_batches = [] pending_types = [] pending_rows = [] pending_n = 0 if pending_batches: yield concat_tensor_batches(pending_batches), pending_types, pending_rows def finite_float64(values: Any) -> np.ndarray: arr = np.asarray(values, dtype=np.float64) return arr[np.isfinite(arr)] def update_death_summary( summary: dict[tuple[Any, ...], dict[str, float]], *, key_rows: pd.DataFrame, values: np.ndarray, ) -> None: if key_rows.empty: return table = key_rows.copy() for idx, column in enumerate(DEATH_PARAMETER_COLUMNS): table[column] = values[:, idx] for key, group in table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False): if not isinstance(key, tuple): key = (key,) acc = summary.setdefault( key, { "n": 0.0, **{f"count__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, **{f"sum__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, **{f"sumsq__{col}": 0.0 for col in DEATH_PARAMETER_COLUMNS}, }, ) acc["n"] += float(len(group)) for column in DEATH_PARAMETER_COLUMNS: vals = finite_float64(pd.to_numeric(group[column], errors="coerce")) acc[f"count__{column}"] += float(vals.size) acc[f"sum__{column}"] += float(vals.sum()) acc[f"sumsq__{column}"] += float(np.square(vals).sum()) def update_disease_parameter_summary( summary: dict[tuple[Any, ...], dict[str, float]], *, key_rows: pd.DataFrame, target_group: str, target_group_label: str, values: np.ndarray, ) -> None: if key_rows.empty or values.size == 0: return table = key_rows.reset_index(drop=True).copy() grouped = table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False) for key, group in grouped: if not isinstance(key, tuple): key = (key,) full_key = (*key, str(target_group), str(target_group_label)) idx = group.index.to_numpy(dtype=np.int64) vals_2d = values[idx].reshape(-1, values.shape[-1]) acc = summary.setdefault( full_key, { "n": 0.0, **{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, **{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, }, ) acc["n"] += float(vals_2d.shape[0]) for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS): vals = finite_float64(vals_2d[:, col_idx]) acc[f"count__{column}"] += float(vals.size) acc[f"sum__{column}"] += float(vals.sum()) acc[f"sumsq__{column}"] += float(np.square(vals).sum()) def write_death_summary_csv( path: Path, summary: dict[tuple[Any, ...], dict[str, float]], *, death_distribution: str, ) -> int: rows: list[dict[str, Any]] = [] for key, acc in summary.items(): n = int(acc["n"]) row = {column: value for column, value in zip(EXTRA_KEY_COLUMNS, key)} row["n"] = n row["death_distribution"] = death_distribution for column in DEATH_PARAMETER_COLUMNS: count = int(acc[f"count__{column}"]) mean = acc[f"sum__{column}"] / count if count > 0 else np.nan second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan row[f"mean__{column}"] = mean row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan rows.append(row) columns = [ *EXTRA_KEY_COLUMNS, "n", "death_distribution", *[ name for column in DEATH_PARAMETER_COLUMNS for name in (f"mean__{column}", f"var__{column}") ], ] pd.DataFrame(rows, columns=columns).sort_values( ["selected_extra_info_type_id", "landmark_age", "sex"], kind="mergesort", ).to_csv(path, index=False) return len(rows) def write_disease_parameter_summary_csv( path: Path, summary: dict[tuple[Any, ...], dict[str, float]], ) -> int: rows: list[dict[str, Any]] = [] for key, acc in summary.items(): n = int(acc["n"]) row = {column: value for column, value in zip(DISEASE_PARAMETER_KEY_COLUMNS, key)} row["n"] = n for column in DISEASE_PARAMETER_COLUMNS: count = int(acc[f"count__{column}"]) mean = acc[f"sum__{column}"] / count if count > 0 else np.nan second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan row[f"mean__{column}"] = mean row[f"var__{column}"] = second - mean * mean if count > 0 else np.nan rows.append(row) columns = [ *DISEASE_PARAMETER_KEY_COLUMNS, "n", *[ name for column in DISEASE_PARAMETER_COLUMNS for name in (f"mean__{column}", f"var__{column}") ], ] pd.DataFrame(rows, columns=columns).sort_values( ["selected_extra_info_type_id", "target_group", "landmark_age", "sex"], kind="mergesort", ).to_csv(path, index=False) return len(rows) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Compute extra-info ablation attribution for death parameters and future disease risks." ) parser.add_argument("--run_path", type=str, required=True) parser.add_argument( "--extra_info", type=str, default=None, help=( "Optional type id, variable name, exact full name, or comma-separated list. " "If omitted, scan all extra-info types available in the run." ), ) parser.add_argument("--output_dir", type=str, default=None) parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv") parser.add_argument("--eval_split", type=str, default=None) parser.add_argument("--dataset_subset_size", type=int, default=None) parser.add_argument("--train_eid_file", type=str, default=None) parser.add_argument("--val_eid_file", type=str, default=None) parser.add_argument("--test_eid_file", type=str, default=None) parser.add_argument("--landmark_start", type=float, default=40.0) parser.add_argument("--landmark_stop", type=float, default=80.0) parser.add_argument("--landmark_step", type=float, default=5.0) parser.add_argument("--min_history_events", type=int, default=None) parser.add_argument("--batch_size", type=int, default=None) parser.add_argument( "--attribution_batch_size", type=int, default=None, help="Forward batch size for expanded extra-info ablation queries.", ) parser.add_argument("--num_workers", type=int, default=None) parser.add_argument("--device", type=str, default=None) return parser.parse_args() def main() -> None: args = parse_args() # Dataset extra-info types must reproduce the checkpoint training config. # --extra_info only filters which already-trained types are ablated. args.extra_info_types = None run_path = Path(args.run_path) config_path = run_path / "train_config.json" checkpoint_path = run_path / "best_model.pt" if not config_path.exists(): raise FileNotFoundError(f"train_config.json not found: {config_path}") if not checkpoint_path.exists(): raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}") cfg = load_json_config(config_path) model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() if model_target_mode not in {"next_token", "all_future"}: raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}") target_mode = str(cfg.get("target_mode", "uts")) attn_mask_mode = str( cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware") ) readout_name = str( cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token") ) readout_reduce = str(cfg.get("readout_reduce", "mean")) dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg) validate_dataset_metadata(dataset, cfg) extra_metadata = load_extra_info_metadata( dataset_extra_info_types=dataset.extra_info_types, search_root=Path("."), ) selected_extra_info_types = resolve_extra_info_types( args.extra_info, dataset_extra_info_types=dataset.extra_info_types, metadata=extra_metadata, ) if not selected_extra_info_types: raise ValueError("No extra-info types selected for attribution") landmark_ages = make_landmark_ages( float(args.landmark_start), float(args.landmark_stop), float(args.landmark_step), ) first_occurrence_by_token = build_first_occurrence_maps_for_landmarks( dataset, subset_indices, ) death_idx = int(dataset.vocab_size) - 1 landmark_dataset = LandmarkDataset( dataset=dataset, subset_indices=subset_indices, landmark_ages=landmark_ages, attn_mask_mode=attn_mask_mode, model_target_mode=model_target_mode, min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)), first_occurrence_by_token=first_occurrence_by_token, death_token_ids=[death_idx], ) organ_groups, organ_labels, _token_to_group = load_organ_groups( Path(args.organ_mapping_path), vocab_size=int(dataset.vocab_size), ) all_disease_tokens = sorted( { int(token) for tokens in organ_groups.values() for token in tokens if int(token) != death_idx } ) risk_groups = { "all_modeled_diseases": all_disease_tokens, **{group: tokens for group, tokens in sorted(organ_groups.items())}, } risk_group_labels = { "all_modeled_diseases": "All modeled diseases", **organ_labels, } 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) death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull" cfg_model = dict(cfg) cfg_model["dist_mode"] = dist_mode device = resolve_eval_device(args.device) model = build_model_from_dataset(args, cfg_model, dataset).to(device) load_model_state(model, state_dict) model.eval() batch_size = int(cfg_get(args, cfg, "batch_size", 128)) attribution_batch_size = int( cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096)) ) if attribution_batch_size <= 0: raise ValueError("attribution_batch_size must be positive") num_workers = int(cfg_get(args, cfg, "num_workers", 4)) loader = DataLoader( IndexedLandmarkDataset(landmark_dataset), batch_size=batch_size, shuffle=False, collate_fn=collate_indexed_landmark_fn, num_workers=num_workers, pin_memory=device.type == "cuda", persistent_workers=num_workers > 0, prefetch_factor=2 if num_workers > 0 else None, ) output_dir = ( Path(args.output_dir) if args.output_dir else run_path / f"extra_info_attribution_{eval_split}" ) output_dir.mkdir(parents=True, exist_ok=True) print(f"Eval split: {eval_split}") print(f"Split source: {split_source}") print(f"Selected patients: {len(subset_indices)}") print(f"Landmark ages: {landmark_ages.tolist()}") print(f"Dist mode: {dist_mode}") print(f"Device: {device}") print(f"Death token: {death_idx}") print(f"Extra-info types: {selected_extra_info_types}") print(f"Landmark rows: {len(landmark_dataset)}") print(f"Attribution batch size: {attribution_batch_size}") print(f"Output directory: {output_dir}") death_summary: dict[tuple[Any, ...], dict[str, float]] = {} disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {} for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True): batch_dev = { k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) for k, v in batch.items() } with torch.no_grad(): hidden = infer_landmark_hidden( model=model, batch=batch_dev, device=device, model_target_mode=model_target_mode, readout_name=readout_name, readout_reduce=readout_reduce, ) _death_distribution, original_death_params = death_distribution_parameters( model, hidden, dist_mode=dist_mode, ) original_logits = model.calc_risk(hidden) original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None original_disease_params_by_group = {} disease_distribution_name_by_group = {} for group, tokens in risk_groups.items(): disease_distribution, original_disease_params = disease_distribution_parameters( model, hidden, token_ids=tokens, dist_mode=dist_mode, logits=original_logits, rho=original_rho, ) disease_distribution_name_by_group[group] = disease_distribution original_disease_params_by_group[group] = original_disease_params for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches( batch_dev, selected_extra_info_types=selected_extra_info_types, max_batch_size=attribution_batch_size, ): row_tensor = torch.as_tensor(local_rows, dtype=torch.long, device=device) with torch.no_grad(): ablated_hidden = infer_landmark_hidden( model=model, batch=ablated_batch, device=device, model_target_mode=model_target_mode, readout_name=readout_name, readout_reduce=readout_reduce, ) _ablated_distribution, ablated_death_params = death_distribution_parameters( model, ablated_hidden, dist_mode=dist_mode, ) ablated_logits = model.calc_risk(ablated_hidden) ablated_rho = model.calc_weibull_rho(ablated_hidden) if dist_mode == "weibull" else None key_rows = [] for type_id, local_row in zip(type_ids, local_rows): meta = extra_metadata[int(type_id)] key_rows.append( { "selected_extra_info_type_id": int(type_id), "selected_extra_info_var_name": str(meta.get("var_name", "")), "selected_extra_info_full_name": str(meta.get("full_name", "")), "landmark_age": float(batch["landmark_age"][int(local_row)].item()), "sex": int(batch["sex"][int(local_row)].item()), } ) key_table = pd.DataFrame(key_rows, columns=EXTRA_KEY_COLUMNS) value_block = parameter_pair_block( original_death_params[row_tensor], ablated_death_params, ).detach().cpu().numpy() update_death_summary( death_summary, key_rows=key_table, values=value_block, ) for group, tokens in risk_groups.items(): disease_distribution, ablated_disease_params = disease_distribution_parameters( model, ablated_hidden, token_ids=tokens, dist_mode=dist_mode, logits=ablated_logits, rho=ablated_rho, ) if disease_distribution_name_by_group[group] != disease_distribution: raise RuntimeError( "Disease distribution changed between original and ablated passes " f"for group {group!r}: {disease_distribution_name_by_group[group]!r} " f"vs {disease_distribution!r}" ) value_block = disease_parameter_pair_block( original_disease_params_by_group[group][row_tensor], ablated_disease_params, ).detach().cpu().numpy() update_disease_parameter_summary( disease_parameter_summary, key_rows=key_table, target_group=group, target_group_label=risk_group_labels[group], values=value_block, ) death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv" death_rows = write_death_summary_csv( death_summary_path, death_summary, death_distribution=death_distribution_name, ) disease_rows = write_disease_parameter_summary_csv( disease_summary_path, disease_parameter_summary, ) manifest = { "death_summary_file": death_summary_path.name, "disease_parameter_summary_file": disease_summary_path.name, "death_summary_rows": int(death_rows), "disease_parameter_summary_rows": int(disease_rows), "eval_split": eval_split, "split_source": split_source, "dist_mode": dist_mode, "landmark_start": float(args.landmark_start), "landmark_stop": float(args.landmark_stop), "landmark_step": float(args.landmark_step), "selected_extra_info_types": [ extra_metadata[int(type_id)] for type_id in selected_extra_info_types ], } with (output_dir / "manifest.json").open("w", encoding="utf-8") as f: json.dump(manifest, f, ensure_ascii=False, indent=2) print(f"Wrote {death_rows} death summary rows to {death_summary_path}") print(f"Wrote {disease_rows} disease-parameter summary rows to {disease_summary_path}") if __name__ == "__main__": main()