"""Compute per-disease attribution to predicted mortality distribution parameters. For each selected patient and landmark age, this script keeps only rows where each scanned disease token has already occurred in the history. It then deletes that historical disease token, re-queries the model, and reports the original and ablated fitted death distribution parameters. If --disease is omitted, all disease tokens in the mapping are scanned. Death is always token vocab_size - 1. """ from __future__ import annotations import argparse import json from pathlib import Path from typing import Any, Dict 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 evaluate_event_free_survival import ( IndexedLandmarkDataset, LandmarkDataset, build_first_occurrence_maps_for_landmarks, collate_indexed_landmark_fn, historical_counts_by_group, infer_landmark_hidden, load_eval_sequence_dataset, load_organ_groups, make_landmark_ages, ) from targets import CHECKUP_IDX, PAD_IDX OUTPUT_COLUMNS = [ "patient_id", "dataset_index", "eid", "sex", "landmark_age", "followup_end_time", "history_disease_count", "selected_disease_history_count", "selected_disease_token_id", "selected_disease_code", "selected_disease_name", "selected_disease_organ_system", "selected_disease_organ_system_label", "history_count__selected_organ_system", "death_distribution", "original_death_lambda", "ablated_death_lambda", "original_death_scale", "ablated_death_scale", "original_death_shape", "ablated_death_shape", ] SUMMARY_KEY_COLUMNS = [ "selected_disease_token_id", "selected_disease_code", "selected_disease_name", "selected_disease_organ_system", "selected_disease_organ_system_label", "landmark_age", "sex", ] SUMMARY_MEAN_COLUMNS = [ "history_disease_count", "selected_disease_history_count", "history_count__selected_organ_system", ] SUMMARY_PARAMETER_COLUMNS = [ "original_death_lambda", "ablated_death_lambda", "original_death_scale", "ablated_death_scale", "original_death_shape", "ablated_death_shape", ] def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int: table = table.reindex(columns=OUTPUT_COLUMNS) arrays: dict[str, np.ndarray] = { "__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"), } for column in OUTPUT_COLUMNS: values = table[column] if column in table else pd.Series([], dtype=object) if values.dtype == object: arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U") else: arrays[column] = values.to_numpy() np.savez_compressed(path, **arrays) return int(len(table)) def normalize_output_dir(path: Path) -> Path: if path.suffix: return path.with_suffix(path.suffix + "_shards") return path def write_manifest( output_dir: Path, *, rows: int, shards: list[dict[str, Any]], summary_file: str, scanned_diseases: list[dict[str, Any]], eval_split: str, dist_mode: str, landmark_start: float, landmark_stop: float, landmark_step: float, ) -> None: payload = { "format": "compressed_npz_shards", "columns": OUTPUT_COLUMNS, "rows": int(rows), "shards": shards, "summary_file": summary_file, "scanned_diseases": scanned_diseases, "eval_split": eval_split, "dist_mode": str(dist_mode), "landmark_start": float(landmark_start), "landmark_stop": float(landmark_stop), "landmark_step": float(landmark_step), } with (output_dir / "manifest.json").open("w", encoding="utf-8") as f: json.dump(payload, f, ensure_ascii=False, indent=2) def update_summary_accumulator( summary: dict[tuple[Any, ...], dict[str, float]], table: pd.DataFrame, ) -> None: if table.empty: return grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False) for key, group in grouped: if not isinstance(key, tuple): key = (key,) acc = summary.setdefault( key, { "n": 0.0, **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS}, **{f"count__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS}, **{f"sum__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS}, **{f"sumsq__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS}, }, ) n = int(len(group)) acc["n"] += float(n) for column in SUMMARY_MEAN_COLUMNS: acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum()) for column in SUMMARY_PARAMETER_COLUMNS: values = pd.to_numeric(group[column], errors="coerce").dropna() acc[f"count__{column}"] += float(len(values)) acc[f"sum__{column}"] += float(values.sum()) acc[f"sumsq__{column}"] += float((values * values).sum()) def write_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"]) out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)} out["n"] = n for column in SUMMARY_MEAN_COLUMNS: out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan for column in SUMMARY_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 out[f"mean__{column}"] = mean out[f"var__{column}"] = second - mean * mean if count > 0 else np.nan rows.append(out) columns = [ *SUMMARY_KEY_COLUMNS, "n", *[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS], *[ name for column in SUMMARY_PARAMETER_COLUMNS for name in (f"mean__{column}", f"var__{column}") ], ] pd.DataFrame(rows, columns=columns).sort_values( ["selected_disease_token_id", "landmark_age", "sex"], kind="mergesort", ).to_csv(path, index=False) return len(rows) def build_disease_ablated_slice( batch: Dict[str, torch.Tensor], row_indices: torch.Tensor, token_ids: torch.Tensor, ) -> Dict[str, torch.Tensor]: """Build an ablated slice for aligned (row, disease_token) pairs.""" event_seq = batch["event_seq"] row_indices = row_indices.to(device=event_seq.device, dtype=torch.long) token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype) out: Dict[str, torch.Tensor] = {} out["event_seq"] = event_seq[row_indices].clone() out["time_seq"] = batch["time_seq"][row_indices] out["readout_mask"] = batch["readout_mask"][row_indices].clone() out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone() out["landmark_pos"] = batch["landmark_pos"][row_indices].clone() seq_len = int(event_seq.shape[1]) positions = torch.arange(seq_len, device=event_seq.device)[None, :] remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"] out["event_seq"] = torch.where( remove, torch.full_like(out["event_seq"], PAD_IDX), out["event_seq"], ) out["padding_mask"] &= ~remove out["readout_mask"] &= ~remove has_valid = out["padding_mask"].any(dim=1) empty_rows = ~has_valid out["event_seq"][empty_rows, 0] = CHECKUP_IDX out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to( dtype=out["time_seq"].dtype ) out["padding_mask"][empty_rows, 0] = True out["readout_mask"][empty_rows, 0] = True out["landmark_pos"][empty_rows] = 0 has_readout = out["readout_mask"].any(dim=1) missing_readout = ~has_readout local_valid = out["padding_mask"] last_pos = torch.where( local_valid, positions.expand(local_valid.shape[0], -1), torch.zeros_like(positions.expand(local_valid.shape[0], -1)), ).amax(dim=1) out["readout_mask"][missing_readout] = False out["readout_mask"][missing_readout, last_pos[missing_readout]] = True out["landmark_pos"][missing_readout] = last_pos[missing_readout].to( dtype=out["landmark_pos"].dtype ) repeated_keys = ( "sex", "other_type", "other_value", "other_value_kind", "other_time", "t_query", "patient_id", "landmark_age", "followup_end_time", "death_time", "row_idx", ) for key in repeated_keys: out[key] = batch[key][row_indices] return out def load_disease_metadata( mapping_path: Path, *, vocab_size: int, ) -> dict[int, dict[str, Any]]: if not mapping_path.exists(): raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}") table = pd.read_csv(mapping_path) required = {"token_id", "code", "name", "is_death"} missing = required - set(table.columns) if missing: raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}") death_idx = int(vocab_size) - 1 out: dict[int, dict[str, Any]] = {} for row in table.itertuples(index=False): token = int(getattr(row, "token_id")) if token < 0 or token >= int(vocab_size) or token == death_idx: continue if int(getattr(row, "is_death")) == 1: continue meta = { "token_id": token, "code": str(getattr(row, "code")), "name": str(getattr(row, "name")), } for column in ( "icd10_chapter", "icd10_chapter_title", "organ_system", "organ_system_label", ): if hasattr(row, column): meta[column] = str(getattr(row, column)) out[token] = meta return out def resolve_disease_token( value: str, metadata: dict[int, dict[str, Any]], ) -> tuple[int, dict[str, Any]]: text = str(value).strip() if text == "": raise ValueError("--disease must not be empty") if text.isdigit() or (text.startswith("-") and text[1:].isdigit()): token = int(text) if token not in metadata: raise ValueError(f"Disease token_id {token} was not found in the mapping") return token, metadata[token] lower = text.lower() exact = [ (token, meta) for token, meta in metadata.items() if str(meta.get("code", "")).lower() == lower or str(meta.get("name", "")).lower() == lower ] if len(exact) == 1: return exact[0] if len(exact) > 1: raise ValueError(f"--disease={value!r} matched multiple diseases exactly") contains = [ (token, meta) for token, meta in metadata.items() if lower in str(meta.get("code", "")).lower() or lower in str(meta.get("name", "")).lower() ] if len(contains) == 1: return contains[0] if not contains: raise ValueError(f"--disease={value!r} did not match any disease token") preview = ", ".join( f"{token}:{meta.get('code')} ({meta.get('name')})" for token, meta in contains[:10] ) raise ValueError( f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. " f"First matches: {preview}" ) def resolve_disease_tokens( value: str | None, metadata: dict[int, dict[str, Any]], ) -> list[tuple[int, dict[str, Any]]]: if value is None or str(value).strip() == "": return [(token, metadata[token]) for token in sorted(metadata)] out: list[tuple[int, dict[str, Any]]] = [] seen: set[int] = set() for part in str(value).split(","): token, meta = resolve_disease_token(part, metadata) if token not in seen: out.append((token, meta)) seen.add(token) return out def death_distribution_parameters( model, hidden: torch.Tensor, *, dist_mode: str, eps: float = 1e-8, ) -> tuple[str, torch.Tensor]: """Return death distribution parameters with columns matching PARAMETER_VALUE_COLUMNS.""" 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 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 output_name_for_run(run_path: Path, eval_split: str, *, all_diseases: bool) -> Path: scope = "all_diseases" if all_diseases else "selected_diseases" return run_path / f"single_disease_mortality_parameters_{eval_split}_{scope}" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Compute per-disease model attribution to mortality distribution parameters." ) parser.add_argument("--run_path", type=str, required=True) parser.add_argument( "--disease", type=str, default=None, help=( "Optional disease token_id, ICD-10 code, exact name, unambiguous name " "substring, or comma-separated list. If omitted, scan all disease tokens." ), ) parser.add_argument( "--output_path", type=str, default=None, help="Output directory for compressed .npz shards.", ) 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 disease-token ablation queries.", ) parser.add_argument("--num_workers", type=int, default=None) parser.add_argument("--device", type=str, default=None) parser.add_argument("--extra_info_types", type=str, default=None) parser.add_argument( "--shard_rows", type=int, default=200_000, help="Approximate number of detailed rows to buffer before writing one .npz shard.", ) return parser.parse_args() def main() -> None: args = parse_args() 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) metadata = load_disease_metadata( Path(args.organ_mapping_path), vocab_size=int(dataset.vocab_size), ) scanned_disease_items = resolve_disease_tokens(args.disease, metadata) if not scanned_disease_items: raise ValueError("No diseases selected for attribution") scanned_disease_tokens = [token for token, _meta in scanned_disease_items] 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), ) group_names = sorted(organ_groups) 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) selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool) selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True 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") if int(args.shard_rows) <= 0: raise ValueError("--shard_rows 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_path = ( Path(args.output_path) if args.output_path else output_name_for_run( run_path, eval_split, all_diseases=args.disease is None or str(args.disease).strip() == "", ) ) output_dir = normalize_output_dir(output_path) 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}") if len(scanned_disease_items) == len(metadata): print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})") else: preview = ", ".join( f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10] ) print(f"Diseases: {len(scanned_disease_items)} selected ({preview})") print(f"Landmark rows: {len(landmark_dataset)}") print(f"Attribution batch size: {attribution_batch_size}") print(f"Output directory: {output_dir}") written_rows = 0 shard_index = 0 shards: list[dict[str, Any]] = [] row_base_cache: dict[int, dict[str, Any]] = {} result_row_idx_chunks: list[np.ndarray] = [] result_disease_token_chunks: list[np.ndarray] = [] result_value_chunks: list[np.ndarray] = [] def get_row_base(row_idx: int) -> dict[str, Any]: cached = row_base_cache.get(row_idx) if cached is not None: return cached meta = landmark_dataset.rows[int(row_idx)] dataset_index = int(meta["dataset_index"]) sample = dataset.samples[dataset_index] hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64) unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True) total_count, group_counts = historical_counts_by_group( hist_tokens, death_idx=death_idx, token_to_group=token_to_group, group_names=group_names, ) cached = { "patient_id": int(meta["patient_id"]), "dataset_index": dataset_index, "eid": int(sample.get("eid", -1)), "sex": int(meta["sex"]), "landmark_age": float(meta["landmark_age"]), "followup_end_time": float(meta["followup_end_time"]), "history_disease_count": int(total_count), "_hist_tokens": hist_tokens, "_token_counts": { int(token): int(count) for token, count in zip(unique_tokens.tolist(), token_counts.tolist()) }, "_group_counts": group_counts, } row_base_cache[row_idx] = cached return cached for batch in tqdm(loader, desc="Per-disease mortality 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() } 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, ) with torch.no_grad(): _death_distribution, original_params = death_distribution_parameters( model, hidden, dist_mode=dist_mode, ) event_np = batch["event_seq"].numpy() valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size)) selected_event = np.zeros_like(valid_event, dtype=bool) selected_event[valid_event] = selected_token_mask[event_np[valid_event]] pair_row_np, pair_pos_np = np.nonzero(selected_event) if pair_row_np.size == 0: continue pair_disease_np = event_np[pair_row_np, pair_pos_np].astype(np.int64, copy=False) pair_offset = 0 while pair_offset < int(pair_row_np.shape[0]): pair_stop = min(int(pair_row_np.shape[0]), pair_offset + int(attribution_batch_size)) local_rows_np = pair_row_np[pair_offset:pair_stop].astype(np.int64, copy=False) disease_tokens_np = pair_disease_np[pair_offset:pair_stop] local_rows = torch.as_tensor(local_rows_np, dtype=torch.long, device=device) disease_token_ids = torch.as_tensor(disease_tokens_np, dtype=torch.long, device=device) ablated_chunk = build_disease_ablated_slice( batch=batch_dev, row_indices=local_rows, token_ids=disease_token_ids, ) with torch.no_grad(): ablated_hidden = infer_landmark_hidden( model=model, batch=ablated_chunk, device=device, model_target_mode=model_target_mode, readout_name=readout_name, readout_reduce=readout_reduce, ) _ablated_distribution, ablated_params = death_distribution_parameters( model, ablated_hidden, dist_mode=dist_mode, ) value_block = parameter_pair_block( original_params[local_rows], ablated_params, ).detach().cpu().numpy() row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False) disease_tokens_list = disease_tokens_np result_row_idx_chunks.append(row_ids) result_disease_token_chunks.append(disease_tokens_list) result_value_chunks.append(value_block) pair_offset = pair_stop if result_value_chunks: all_row_ids = np.concatenate(result_row_idx_chunks).astype(np.int64, copy=False) all_disease_tokens = np.concatenate(result_disease_token_chunks).astype( np.int64, copy=False, ) all_values = np.concatenate(result_value_chunks, axis=0) rows: list[dict[str, Any]] = [] for i, (row_idx, disease_token) in enumerate( zip(all_row_ids.tolist(), all_disease_tokens.tolist()) ): disease_token = int(disease_token) disease_meta = metadata[disease_token] row_base = get_row_base(int(row_idx)) group_counts = row_base["_group_counts"] disease_history_count = int(row_base["_token_counts"].get(disease_token, 0)) if disease_history_count <= 0: raise RuntimeError( "Internal mismatch: occurred mask selected disease " f"{disease_token} for row {row_idx}, but cached history has count 0" ) rows.append( { "patient_id": row_base["patient_id"], "dataset_index": row_base["dataset_index"], "eid": row_base["eid"], "sex": row_base["sex"], "landmark_age": row_base["landmark_age"], "followup_end_time": row_base["followup_end_time"], "history_disease_count": row_base["history_disease_count"], "selected_disease_history_count": disease_history_count, "selected_disease_token_id": int(disease_token), "selected_disease_code": str(disease_meta.get("code", "")), "selected_disease_name": str(disease_meta.get("name", "")), "selected_disease_organ_system": str(disease_meta.get("organ_system", "")), "selected_disease_organ_system_label": str( disease_meta.get("organ_system_label", "") ), "history_count__selected_organ_system": int( group_counts.get(str(disease_meta.get("organ_system", "")), 0) ), "death_distribution": death_distribution_name, "original_death_lambda": float(all_values[i, 0]), "ablated_death_lambda": float(all_values[i, 1]), "original_death_scale": float(all_values[i, 2]), "ablated_death_scale": float(all_values[i, 3]), "original_death_shape": float(all_values[i, 4]), "ablated_death_shape": float(all_values[i, 5]), } ) result_table = pd.DataFrame(rows).reindex(columns=OUTPUT_COLUMNS) written_rows = int(len(result_table)) summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {} update_summary_accumulator(summary_accumulator, result_table) for start in range(0, written_rows, int(args.shard_rows)): stop = min(written_rows, start + int(args.shard_rows)) shard_name = f"part-{shard_index:06d}.npz" shard_path = output_dir / shard_name shard_rows = write_compressed_npz_table( shard_path, result_table.iloc[start:stop], ) shards.append({"file": shard_name, "rows": int(shard_rows)}) shard_index += 1 else: result_table = pd.DataFrame(columns=OUTPUT_COLUMNS) summary_accumulator = {} if not shards: empty_path = output_dir / "part-000000.npz" write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS)) shards.append({"file": empty_path.name, "rows": 0}) summary_path = output_dir / "summary_by_disease_age_sex.csv" summary_rows = write_summary_csv(summary_path, summary_accumulator) write_manifest( output_dir, rows=written_rows, shards=shards, summary_file=summary_path.name, scanned_diseases=[ {"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}} for token, meta in scanned_disease_items ], eval_split=eval_split, dist_mode=dist_mode, landmark_start=float(args.landmark_start), landmark_stop=float(args.landmark_stop), landmark_step=float(args.landmark_step), ) print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}") print(f"Wrote {summary_rows} summary rows to {summary_path}") if __name__ == "__main__": main()