From d728d8585c39bf9661f5193f57af51e2471d7488 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 21 Aug 2026 13:48:59 +0800 Subject: [PATCH] Improve all-future first-onset training --- dataset.py | 105 +++-- eval_data.py | 7 +- evaluate_calibration.py | 8 +- losses.py | 149 ++++++- models.py | 30 +- ...ll_future_first_onset_experiments_linux.sh | 390 ++++++++++++++++++ tests/test_all_future_likelihood.py | 173 ++++++++ train_all_future.py | 95 ++++- train_util.py | 129 ++++++ 9 files changed, 1002 insertions(+), 84 deletions(-) create mode 100644 run_all_future_first_onset_experiments_linux.sh create mode 100644 tests/test_all_future_likelihood.py diff --git a/dataset.py b/dataset.py index 40b278c..1a651e7 100644 --- a/dataset.py +++ b/dataset.py @@ -17,9 +17,6 @@ from targets import ( build_next_token_targets, ) - -ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR - DISEASE_HISTORY_MODE_TIMED = "timed" DISEASE_HISTORY_MODE_ORDERED = "ordered" DISEASE_HISTORY_MODE_SET = "set" @@ -518,10 +515,10 @@ class AllFutureHealthDataset(_ExpoBaseDataset): Dataset with unified other-info tokens and DeepHealthV2-style all-future targets. - Train samples one query time per patient at each __getitem__ call. - Valid/test use random-but-fixed query points. For each patient with N real - disease events, N - 2 query points are sampled from the eligible observed - time range, with at least one future event after every query. + Every split uses the same patient-equal query distribution: choose one + eligible inter-event interval uniformly, then choose a time uniformly in + that interval. Train resamples on every ``__getitem__`` call; valid/test + keep one deterministic draw per patient. """ CACHE_VERSION = 5 @@ -591,6 +588,11 @@ class AllFutureHealthDataset(_ExpoBaseDataset): **features, } + query_intervals = self._eligible_query_intervals(patient) + if not query_intervals: + continue + patient["query_intervals"] = query_intervals + pidx = len(self.patients) self.patients.append(patient) @@ -615,7 +617,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): labels, np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64), ) - n_hist = int((times <= t_query).sum()) + n_hist = int(((times <= t_query) & real_event_mask).sum()) n_future = int(((times > t_query) & real_event_mask).sum()) return ( n_hist >= self.min_history_events @@ -623,62 +625,59 @@ class AllFutureHealthDataset(_ExpoBaseDataset): and patient["t_obs"] > t_query ) - def _sample_fixed_validation_queries( - self, - patient: Dict, - rng: np.random.RandomState, - ) -> List[float]: + def _eligible_query_intervals(self, patient: Dict) -> List[Tuple[float, float]]: times = np.asarray(patient["times"], dtype=np.float32) labels = np.asarray(patient["labels"], dtype=np.int64) real_event_mask = ~np.isin( labels, np.array([PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64), ) - real_times = np.sort(times[real_event_mask].astype(np.float32, copy=False)) - n_real_events = int(real_times.size) - n_queries = max(0, n_real_events - 2) - if n_queries == 0: - return [] + unique_times = np.unique(times[real_event_mask]) + intervals: List[Tuple[float, float]] = [] + for j in range(1, len(unique_times)): + left = float(unique_times[j - 1]) + right = float(unique_times[j]) + probe = float( + np.nextafter(np.float32(right), np.float32(-np.inf)) + ) + if np.isfinite(left) and np.isfinite(probe) and probe >= left: + if self._is_valid_query(patient, probe): + intervals.append((left, right)) + return intervals - min_hist = int(self.min_history_events) - min_future = int(self.min_future_events) - if n_real_events < min_hist + min_future: - return [] + def sample_query( + self, + patient: Dict, + rng, + ) -> float: + intervals = patient.get("query_intervals") + if intervals is None: + intervals = self._eligible_query_intervals(patient) + if not intervals: + raise RuntimeError("Patient has no eligible all-future query interval.") - left = float(real_times[min_hist - 1]) - right_event_time = float(real_times[n_real_events - min_future]) - right = np.nextafter(np.float32(right_event_time), np.float32(-np.inf)) - if not np.isfinite(left) or not np.isfinite(right) or float(right) <= left: - return [] + interval_idx = int(rng.randint(0, len(intervals))) + left, right_event_time = intervals[interval_idx] + right = float( + np.nextafter(np.float32(right_event_time), np.float32(-np.inf)) + ) + if right <= left: + t_query = float(left) + else: + t_query = float(rng.uniform(left, right)) + if not self._is_valid_query(patient, t_query): + raise RuntimeError("Sampled an invalid all-future query time.") + return t_query - queries: List[float] = [] - max_attempts = max(100, n_queries * 50) - for _ in range(max_attempts): - if len(queries) >= n_queries: - break - t_query = float(rng.uniform(left, float(right))) - if self._is_valid_query(patient, t_query): - queries.append(t_query) - - return queries + def _sample_fixed_validation_queries( + self, + patient: Dict, + rng: np.random.RandomState, + ) -> List[float]: + return [self.sample_query(patient, rng)] def _sample_train_query(self, patient: Dict) -> float: - unique_times = np.unique(patient["times"]) - if len(unique_times) < 2: - raise RuntimeError("Training patient has fewer than two unique times.") - - j = np.random.randint(1, len(unique_times)) - left = float(unique_times[j - 1]) - right = float(unique_times[j]) - - if right - left <= ONE_DAY_YEARS: - t_query = right - ONE_DAY_YEARS - else: - t_query = np.random.uniform(left, right - ONE_DAY_YEARS) - - if not self._is_valid_query(patient, t_query): - t_query = right - 1e-6 - return float(t_query) + return self.sample_query(patient, np.random) def _build_item(self, patient: Dict, t_query: float) -> Dict: times = patient["times"] diff --git a/eval_data.py b/eval_data.py index 4c598e2..db85efc 100644 --- a/eval_data.py +++ b/eval_data.py @@ -66,7 +66,6 @@ def validate_training_mode_config(cfg: Dict[str, Any]) -> None: "model_target_mode must be next_token or all_future, got " f"{model_target_mode!r}" ) - disease_history_mode = normalize_disease_history_mode( cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED) ) @@ -183,6 +182,11 @@ def build_model_from_dataset( "model_target_mode must be next_token or all_future, got " f"{model_target_mode!r}" ) + risk_head_bias = bool(cfg_get(args, cfg, "risk_head_bias", False)) + if state_dict is not None: + # The checkpoint schema is authoritative. This keeps all older + # bias-free checkpoints loadable while restoring the new baseline bias. + risk_head_bias = "risk_head.bias" in state_dict model_architecture = resolve_model_architecture(cfg, state_dict) continuous_value_center = None continuous_value_scale = None @@ -231,6 +235,7 @@ def build_model_from_dataset( dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), model_architecture=model_architecture, + risk_head_bias=risk_head_bias, ) diff --git a/evaluate_calibration.py b/evaluate_calibration.py index 73eb210..d222a45 100644 --- a/evaluate_calibration.py +++ b/evaluate_calibration.py @@ -66,7 +66,7 @@ from evaluate_auc_v2 import ( ) from losses import build_loss from model_architectures import resolve_model_architecture -from targets import PAD_IDX, RESERVED_IDX +from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX from train_util import load_eid_file @@ -1483,7 +1483,7 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame: def _build_point_process_criterion( dist_mode: str, ) -> Any: - ignored = {PAD_IDX, RESERVED_IDX} + ignored = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX} if dist_mode == "exponential": return build_loss("exponential", ignored_idx=ignored) if dist_mode == "weibull": @@ -1545,6 +1545,8 @@ def evaluate_point_process_nll( logits=logits, targets=batch_device["future_targets"], exposure=batch_device["exposure"], + dt=batch_device["future_dt"], + history=batch_device["event_seq"], ) elif dist_mode == "weibull": loss = criterion( @@ -1553,6 +1555,7 @@ def evaluate_point_process_nll( targets=batch_device["future_targets"], dt=batch_device["future_dt"], exposure=batch_device["exposure"], + history=batch_device["event_seq"], ) else: raise ValueError(f"Unsupported dist_mode: {dist_mode!r}") @@ -1564,6 +1567,7 @@ def evaluate_point_process_nll( query_count += batch_size valid_targets = batch["future_targets"] > PAD_IDX valid_targets &= batch["future_targets"] != RESERVED_IDX + valid_targets &= batch["future_targets"] != NO_EVENT_IDX future_event_count += int(valid_targets.sum().item()) exposure_sum += float(batch["exposure"].sum().item()) diff --git a/losses.py b/losses.py index dfbc122..3936a90 100644 --- a/losses.py +++ b/losses.py @@ -37,6 +37,54 @@ def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor: return logits.sum() * 0.0 +def _all_future_at_risk_mask( + vocab_size: int, + ignored_idx: Iterable[int], + logits: torch.Tensor, + history: torch.Tensor | None, +) -> torch.Tensor: + """Return outcomes that have not occurred by the query time.""" + batch_size = logits.shape[0] + at_risk = _valid_vocab_mask( + vocab_size, + ignored_idx, + logits.device, + ).unsqueeze(0).expand(batch_size, -1).clone() + if history is None or history.numel() == 0: + return at_risk + if history.dim() != 2 or history.shape[0] != batch_size: + raise ValueError( + "history must be (B, L). " + f"Got logits={tuple(logits.shape)}, history={tuple(history.shape)}" + ) + + history = history.to(device=logits.device, dtype=torch.long) + history_valid = (history >= 0) & (history < vocab_size) + for idx in ignored_idx: + history_valid &= history != int(idx) + safe_history = history.clamp(min=0, max=vocab_size - 1) + prevalent = torch.zeros( + (batch_size, vocab_size), + dtype=torch.long, + device=logits.device, + ) + prevalent.scatter_add_(1, safe_history, history_valid.long()) + return at_risk & ~prevalent.bool() + + +def _all_future_target_mask( + targets: torch.Tensor, + at_risk: torch.Tensor, + ignored_idx: Iterable[int], +) -> tuple[torch.Tensor, torch.Tensor]: + vocab_size = at_risk.shape[1] + in_vocab = (targets >= 0) & (targets < vocab_size) + for idx in ignored_idx: + in_vocab &= targets != int(idx) + safe_targets = targets.clamp(min=0, max=vocab_size - 1) + return in_vocab & at_risk.gather(1, safe_targets), safe_targets + + class Delphi2MLoss(nn.Module): """Next-token plus exponential time-to-next-token supervision.""" @@ -146,7 +194,7 @@ class Delphi2MLoss(nn.Module): class ExponentialLoss(nn.Module): - """Query-conditioned all-future-event exponential point-process loss.""" + """First-onset all-future exponential survival likelihood.""" def __init__( self, @@ -162,24 +210,61 @@ class ExponentialLoss(nn.Module): logits: torch.Tensor, targets: torch.Tensor, exposure: torch.Tensor, + dt: torch.Tensor, + history: torch.Tensor | None = None, ) -> torch.Tensor: - _, vocab_size = logits.shape + batch_size, vocab_size = logits.shape + if targets.dim() != 2 or targets.shape[0] != batch_size: + raise ValueError( + "targets must be (B, M). " + f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}" + ) + if exposure.shape != (batch_size,): + raise ValueError( + f"exposure must be ({batch_size},), got {tuple(exposure.shape)}" + ) + if dt.shape != targets.shape: + raise ValueError( + f"dt must match targets, got dt={tuple(dt.shape)}, " + f"targets={tuple(targets.shape)}" + ) + rate = F.softplus(logits) + self.eps - valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device) + at_risk = _all_future_at_risk_mask( + vocab_size, + self.ignored_idx, + logits, + history, + ) + target_valid, safe_targets = _all_future_target_mask( + targets, + at_risk, + self.ignored_idx, + ) - penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1) - target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device) - for idx in self.ignored_idx: - target_valid &= targets != idx + censor_time = exposure.to(rate.dtype).clamp_min(self.eps) + cumulative_at_censor = rate * censor_time.unsqueeze(1) + penalty = (cumulative_at_censor * at_risk.to(rate.dtype)).sum(dim=-1) + + # A first-onset outcome leaves the risk set at its event time, not at + # the common end of follow-up. + if targets.numel() > 0: + event_time = dt.to(rate.dtype).clamp_min(self.eps) + event_time = torch.minimum(event_time, censor_time.unsqueeze(1)) + event_cumulative = rate.gather(1, safe_targets) * event_time + censor_cumulative = cumulative_at_censor.gather(1, safe_targets) + penalty = penalty + ( + (event_cumulative - censor_cumulative) + * target_valid.to(rate.dtype) + ).sum(dim=-1) - safe_targets = targets.clamp(min=0, max=vocab_size - 1) observed = rate.log().gather(1, safe_targets) observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1) return (-observed + penalty).mean() class WeibullLoss(nn.Module): - """Query-conditioned all-future-event Weibull point-process loss.""" + """First-onset all-future Weibull survival likelihood.""" def __init__( self, @@ -197,8 +282,9 @@ class WeibullLoss(nn.Module): targets: torch.Tensor, dt: torch.Tensor, exposure: torch.Tensor, + history: torch.Tensor | None = None, ) -> torch.Tensor: - _, vocab_size = logits.shape + batch_size, vocab_size = logits.shape if weibull_rho is None: raise ValueError("weibull_rho is required for WeibullLoss") if weibull_rho.shape != logits.shape: @@ -206,23 +292,50 @@ class WeibullLoss(nn.Module): "weibull_rho must have the same shape as logits. " f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}" ) + if targets.dim() != 2 or targets.shape[0] != batch_size: + raise ValueError( + "targets must be (B, M). " + f"Got logits={tuple(logits.shape)}, targets={tuple(targets.shape)}" + ) + if dt.shape != targets.shape: + raise ValueError( + f"dt must match targets, got dt={tuple(dt.shape)}, " + f"targets={tuple(targets.shape)}" + ) + if exposure.shape != (batch_size,): + raise ValueError( + f"exposure must be ({batch_size},), got {tuple(exposure.shape)}" + ) dtype = logits.dtype rate = F.softplus(logits) + self.eps rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps) - valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device) + at_risk = _all_future_at_risk_mask( + vocab_size, + self.ignored_idx, + logits, + history, + ) + target_valid, safe_targets = _all_future_target_mask( + targets, + at_risk, + self.ignored_idx, + ) - t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1) - penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1) + censor_time = exposure.to(dtype).clamp_min(self.eps) + cumulative_at_censor = rate * torch.pow(censor_time.unsqueeze(1), rho) + penalty = (cumulative_at_censor * at_risk.to(dtype)).sum(dim=-1) - target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device) - for idx in self.ignored_idx: - target_valid &= targets != idx - - safe_targets = targets.clamp(min=0, max=vocab_size - 1) target_rate = rate.gather(1, safe_targets) target_rho = rho.gather(1, safe_targets) target_dt = dt.to(dtype).clamp_min(self.eps) + target_dt = torch.minimum(target_dt, censor_time.unsqueeze(1)) + event_cumulative = target_rate * torch.pow(target_dt, target_rho) + censor_cumulative = cumulative_at_censor.gather(1, safe_targets) + penalty = penalty + ( + (event_cumulative - censor_cumulative) * target_valid.to(dtype) + ).sum(dim=-1) + log_intensity = ( target_rate.log() + target_rho.log() diff --git a/models.py b/models.py index c2c1b30..9578da8 100644 --- a/models.py +++ b/models.py @@ -221,6 +221,8 @@ class DeepHealth(nn.Module): extra_pool_reduce: str = "mean", dropout: float = 0.0, model_architecture: str | None = None, + risk_head_bias: bool = False, + risk_head_bias_init: torch.Tensor | list[float] | None = None, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -302,7 +304,33 @@ class DeepHealth(nn.Module): self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0) self.final_ln = nn.LayerNorm(n_embd) - self.risk_head = nn.Linear(n_embd, vocab_size, bias=False) + self.risk_head = nn.Linear( + n_embd, + vocab_size, + bias=bool(risk_head_bias), + ) + if risk_head_bias_init is not None: + if self.risk_head.bias is None: + raise ValueError( + "risk_head_bias_init requires risk_head_bias=True" + ) + initial_bias = torch.as_tensor( + risk_head_bias_init, + dtype=self.risk_head.bias.dtype, + device=self.risk_head.bias.device, + ) + if initial_bias.shape != (vocab_size,): + raise ValueError( + "risk_head_bias_init must have shape " + f"({vocab_size},), got {tuple(initial_bias.shape)}" + ) + if not torch.isfinite(initial_bias).all(): + raise ValueError("risk_head_bias_init must contain only finite values") + with torch.no_grad(): + # Start exactly at the fitted marginal baseline; covariate and + # history effects are learned away from zero during training. + self.risk_head.weight.zero_() + self.risk_head.bias.copy_(initial_bias) if target_mode == "next_token": self.risk_head.weight = self.token_embedding.weight self.query_token = nn.Parameter(torch.zeros(n_embd)) diff --git a/run_all_future_first_onset_experiments_linux.sh b/run_all_future_first_onset_experiments_linux.sh new file mode 100644 index 0000000..491c445 --- /dev/null +++ b/run_all_future_first_onset_experiments_linux.sh @@ -0,0 +1,390 @@ +#!/usr/bin/env bash +# +# Run the matched experiments required after the all-future first-onset update. +# +# Shared by every experiment: +# - corrected first-onset likelihood and outcome-specific risk exposure; +# - patient/interval/time-uniform query sampling in train/valid/test; +# - timed disease history and smoking/alcohol/BMI extra information. +# +# Per seed, the script trains: +# 1. TrajMixer + relative Weibull + fitted risk baseline (primary model); +# 2. the primary model without the fitted risk baseline (baseline ablation); +# 3. TrajMixer + relative exponential (time-distribution control); +# 4. FFN + relative Weibull (architecture control). +# +# After training, AUC and calibration/point-process-NLL evaluation run +# automatically unless --train-only is supplied. + +set -uo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +GPU_CSV="" +SEED_CSV="42,43,44" +NUM_WORKERS=4 +PYTHON_BIN="${PYTHON_BIN:-python}" +CAMPAIGN_NAME="all_future_first_onset_v2_multiseed" +DRY_RUN=0 +TRAIN_ONLY=0 + +BATCH_SIZE=256 +EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_smoking_alcohol_bmi.txt" +TRAIN_EID_FILE="$SCRIPT_DIR/ukb_train_eid.csv" +VAL_EID_FILE="$SCRIPT_DIR/ukb_val_eid.csv" +TEST_EID_FILE="$SCRIPT_DIR/ukb_test_eid.csv" + +usage() { + cat <<'EOF' +Usage: + bash run_all_future_first_onset_experiments_linux.sh --gpus LIST [options] + +Required: + --gpus LIST Comma-separated GPU ids, for example 0 or 0,1,2,3. + +Options: + --seeds LIST Comma-separated seeds (default: 42,43,44). + --num-workers N DataLoader workers per active GPU (default: 4). + --python PATH Python executable (default: $PYTHON_BIN or python). + --campaign NAME Output campaign directory name. + --train-only Skip AUC and calibration/NLL evaluation. + --dry-run Print commands without running them. + -h, --help Show this help message. + +Fixed settings: + batch_size 256 + extra information smoking + alcohol + BMI + patient split fixed train/validation/test EID files + disease history timed + experiments per seed 4 + +Outputs: + runs//seed_//... + batch_logs//seed_/.log +EOF +} + +while (($# > 0)); do + case "$1" in + --gpus) + [[ $# -ge 2 ]] || { + echo "ERROR: --gpus requires a value." >&2 + exit 2 + } + GPU_CSV="$2" + shift 2 + ;; + --seeds) + [[ $# -ge 2 ]] || { + echo "ERROR: --seeds requires a value." >&2 + exit 2 + } + SEED_CSV="$2" + shift 2 + ;; + --num-workers) + [[ $# -ge 2 ]] || { + echo "ERROR: --num-workers requires a value." >&2 + exit 2 + } + NUM_WORKERS="$2" + shift 2 + ;; + --python) + [[ $# -ge 2 ]] || { + echo "ERROR: --python requires a value." >&2 + exit 2 + } + PYTHON_BIN="$2" + shift 2 + ;; + --campaign) + [[ $# -ge 2 ]] || { + echo "ERROR: --campaign requires a value." >&2 + exit 2 + } + CAMPAIGN_NAME="$2" + shift 2 + ;; + --train-only) + TRAIN_ONLY=1 + shift + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "ERROR: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$GPU_CSV" ]] || { + echo "ERROR: --gpus is required." >&2 + usage >&2 + exit 2 +} +[[ -n "$SEED_CSV" ]] || { + echo "ERROR: --seeds must not be empty." >&2 + exit 2 +} +[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || { + echo "ERROR: --num-workers must be a non-negative integer." >&2 + exit 2 +} +[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || { + echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2 + exit 2 +} + +required_files=( + "$SCRIPT_DIR/train_all_future.py" + "$SCRIPT_DIR/evaluate_all_runs_linux.sh" + "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh" + "$EXTRA_INFO_TYPES_FILE" + "$TRAIN_EID_FILE" + "$VAL_EID_FILE" + "$TEST_EID_FILE" +) +for required_file in "${required_files[@]}"; do + [[ -f "$required_file" ]] || { + echo "ERROR: missing required file: $required_file" >&2 + exit 2 + } +done +command -v "$PYTHON_BIN" >/dev/null 2>&1 || { + echo "ERROR: Python executable not found: $PYTHON_BIN" >&2 + exit 2 +} + +IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV" +declare -A SEEN_GPUS=() +for gpu in "${GPU_IDS[@]}"; do + [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || { + echo "ERROR: invalid GPU id: $gpu" >&2 + exit 2 + } + [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || { + echo "ERROR: duplicate GPU id: $gpu" >&2 + exit 2 + } + SEEN_GPUS["$gpu"]=1 +done + +IFS=',' read -r -a SEEDS <<< "$SEED_CSV" +declare -A SEEN_SEEDS=() +for seed in "${SEEDS[@]}"; do + [[ "$seed" =~ ^[0-9]+$ ]] || { + echo "ERROR: invalid seed: $seed" >&2 + exit 2 + } + [[ -z "${SEEN_SEEDS[$seed]+x}" ]] || { + echo "ERROR: duplicate seed: $seed" >&2 + exit 2 + } + SEEN_SEEDS["$seed"]=1 +done + +RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME" +LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME" +if ((DRY_RUN == 0)); then + mkdir -p "$RUNS_ROOT" "$LOG_ROOT" +fi + +declare -a JOB_NAMES=() +declare -a JOB_SEEDS=() +declare -a JOB_ARCHITECTURES=() +declare -a JOB_DIST_MODES=() +declare -a JOB_BASELINES=() + +add_job() { + JOB_NAMES+=("$1") + JOB_SEEDS+=("$2") + JOB_ARCHITECTURES+=("$3") + JOB_DIST_MODES+=("$4") + JOB_BASELINES+=("$5") +} + +for seed in "${SEEDS[@]}"; do + add_job \ + "traj_mixer_relative_weibull_first_onset" \ + "$seed" \ + "traj_mixer_v5" \ + "weibull" \ + "on" + add_job \ + "traj_mixer_relative_weibull_no_rate_baseline" \ + "$seed" \ + "traj_mixer_v5" \ + "weibull" \ + "off" + add_job \ + "traj_mixer_relative_exponential_first_onset" \ + "$seed" \ + "traj_mixer_v5" \ + "exponential" \ + "on" + add_job \ + "ffn_relative_weibull_first_onset" \ + "$seed" \ + "transformer_ffn_v1" \ + "weibull" \ + "on" +done + +print_command() { + printf '%q ' "$@" + printf '\n' +} + +run_job() { + local job_index="$1" + local gpu="$2" + local job_name="${JOB_NAMES[$job_index]}" + local seed="${JOB_SEEDS[$job_index]}" + local architecture="${JOB_ARCHITECTURES[$job_index]}" + local dist_mode="${JOB_DIST_MODES[$job_index]}" + local baseline="${JOB_BASELINES[$job_index]}" + local seed_runs_root="$RUNS_ROOT/seed_$seed" + local seed_log_root="$LOG_ROOT/seed_$seed" + local log_file="$seed_log_root/$job_name.log" + local -a command=( + "$PYTHON_BIN" + -u + "$SCRIPT_DIR/train_all_future.py" + --runs_root "$seed_runs_root" + --seed "$seed" + --batch_size "$BATCH_SIZE" + --num_workers "$NUM_WORKERS" + --device cuda + --model_architecture "$architecture" + --time_mode relative + --dist_mode "$dist_mode" + --disease_history_mode timed + --extra_info_types_file "$EXTRA_INFO_TYPES_FILE" + --train_eid_file "$TRAIN_EID_FILE" + --val_eid_file "$VAL_EID_FILE" + --test_eid_file "$TEST_EID_FILE" + ) + + if [[ "$baseline" == "on" ]]; then + command+=(--risk_head_bias) + else + command+=(--no-risk_head_bias) + fi + + if ((DRY_RUN == 0)); then + mkdir -p "$seed_runs_root" "$seed_log_root" + fi + + echo "[$(date '+%F %T')] START seed=$seed job=$job_name gpu=$gpu" + echo " log=$log_file" + if ((DRY_RUN)); then + printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu" + print_command "${command[@]}" + return 0 + fi + + if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \ + "${command[@]}" >"$log_file" 2>&1; then + echo "[$(date '+%F %T')] DONE seed=$seed job=$job_name gpu=$gpu" + return 0 + else + local exit_code=$? + echo "[$(date '+%F %T')] FAIL seed=$seed job=$job_name gpu=$gpu exit=$exit_code" >&2 + echo " See: $log_file" >&2 + return "$exit_code" + fi +} + +worker() { + local slot="$1" + local gpu="${GPU_IDS[$slot]}" + local job_index + local failed=0 + + for ((job_index = slot; job_index < ${#JOB_NAMES[@]}; job_index += ${#GPU_IDS[@]})); do + run_job "$job_index" "$gpu" || failed=1 + done + return "$failed" +} + +echo "Campaign: $CAMPAIGN_NAME" +echo "Seeds: ${SEEDS[*]}" +echo "GPUs: ${GPU_IDS[*]}" +echo "Experiments per seed: 4" +echo "Total training tasks: ${#JOB_NAMES[@]}" +echo "Runs root: $RUNS_ROOT" +echo "Log root: $LOG_ROOT" +echo + +declare -a WORKER_PIDS=() +for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do + worker "$slot" & + WORKER_PIDS+=("$!") +done + +overall_status=0 +for pid in "${WORKER_PIDS[@]}"; do + wait "$pid" || overall_status=1 +done + +if ((overall_status != 0)); then + echo "One or more training tasks failed. Inspect: $LOG_ROOT" >&2 + exit 1 +fi + +if ((TRAIN_ONLY == 0)); then + evaluation_status=0 + auc_command=( + bash + "$SCRIPT_DIR/evaluate_all_runs_linux.sh" + --gpus "$GPU_CSV" + --runs-root "$RUNS_ROOT" + --log-root "$LOG_ROOT/evaluate_auc" + --python "$PYTHON_BIN" + --num-workers "$NUM_WORKERS" + --num-workers-auc "$NUM_WORKERS" + ) + calibration_command=( + bash + "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh" + --gpus "$GPU_CSV" + --runs-root "$RUNS_ROOT" + --log-root "$LOG_ROOT/evaluate_calibration" + --python "$PYTHON_BIN" + --num-workers "$NUM_WORKERS" + --num-workers-calibration "$NUM_WORKERS" + ) + + echo ">> AUC evaluation" + print_command "${auc_command[@]}" + if ((DRY_RUN == 0)); then + "${auc_command[@]}" || evaluation_status=1 + fi + + echo ">> Calibration and point-process NLL evaluation" + print_command "${calibration_command[@]}" + if ((DRY_RUN == 0)); then + "${calibration_command[@]}" || evaluation_status=1 + fi + + if ((evaluation_status != 0)); then + echo "One or more evaluation workflows failed. Inspect: $LOG_ROOT" >&2 + exit 1 + fi +fi + +if ((DRY_RUN)); then + echo "Dry run completed successfully." +else + echo "All required all-future experiments completed successfully." +fi diff --git a/tests/test_all_future_likelihood.py b/tests/test_all_future_likelihood.py new file mode 100644 index 0000000..7081cd9 --- /dev/null +++ b/tests/test_all_future_likelihood.py @@ -0,0 +1,173 @@ +import unittest + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import Subset + +from dataset import AllFutureHealthDataset +from losses import ExponentialLoss, WeibullLoss +from models import DeepHealth +from train_util import fit_all_future_baseline + + +def _inverse_softplus(value: torch.Tensor) -> torch.Tensor: + return torch.log(torch.expm1(value)) + + +class AllFutureLikelihoodTests(unittest.TestCase): + def test_exponential_uses_event_specific_first_onset_exposure(self): + desired_rate = torch.tensor( + [[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]], + dtype=torch.float64, + ) + logits = _inverse_softplus(desired_rate) + criterion = ExponentialLoss(ignored_idx={0, 1, 2}, eps=1e-12) + + loss = criterion( + logits=logits, + targets=torch.tensor([[4, 0]]), + exposure=torch.tensor([5.0], dtype=torch.float64), + dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64), + history=torch.tensor([[3, 0]]), + ) + + rate = F.softplus(logits) + criterion.eps + expected = -(rate[0, 4].log()) + rate[0, 4] * 2.0 + rate[0, 5] * 5.0 + torch.testing.assert_close(loss, expected) + + def test_weibull_uses_event_time_and_excludes_prevalent_outcome(self): + desired_rate = torch.tensor( + [[0.01, 0.01, 0.01, 0.20, 0.30, 0.40]], + dtype=torch.float64, + ) + logits = _inverse_softplus(desired_rate) + rho = torch.tensor( + [[1.0, 1.0, 1.0, 1.2, 1.5, 0.8]], + dtype=torch.float64, + ) + criterion = WeibullLoss(ignored_idx={0, 1, 2}, eps=1e-12) + + loss = criterion( + logits=logits, + weibull_rho=rho, + targets=torch.tensor([[4, 0]]), + dt=torch.tensor([[2.0, 0.0]], dtype=torch.float64), + exposure=torch.tensor([5.0], dtype=torch.float64), + history=torch.tensor([[3, 0]]), + ) + + rate = F.softplus(logits) + criterion.eps + log_hazard = ( + rate[0, 4].log() + + rho[0, 4].log() + + (rho[0, 4] - 1.0) * torch.tensor(2.0).log() + ) + expected = ( + -log_hazard + + rate[0, 4] * torch.pow(torch.tensor(2.0), rho[0, 4]) + + rate[0, 5] * torch.pow(torch.tensor(5.0), rho[0, 5]) + ) + torch.testing.assert_close(loss, expected) + + +class AllFutureQueryDistributionTests(unittest.TestCase): + @staticmethod + def _dataset() -> AllFutureHealthDataset: + dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset) + dataset.min_history_events = 1 + dataset.min_future_events = 1 + return dataset + + @staticmethod + def _patient(): + return { + "times": np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32), + "labels": np.asarray([3, 4, 5, 6], dtype=np.int64), + "t_obs": 4.0, + } + + def test_train_validation_and_test_share_one_query_sampler(self): + dataset = self._dataset() + patient = self._patient() + patient["query_intervals"] = dataset._eligible_query_intervals(patient) + self.assertEqual(len(patient["query_intervals"]), 3) + + rng = np.random.RandomState(123) + fixed = dataset._sample_fixed_validation_queries(patient, rng) + self.assertEqual(len(fixed), 1) + self.assertTrue(dataset._is_valid_query(patient, fixed[0])) + + seen_intervals = set() + rng = np.random.RandomState(456) + for _ in range(200): + query = dataset.sample_query(patient, rng) + self.assertTrue(dataset._is_valid_query(patient, query)) + seen_intervals.add(int(np.floor(query))) + self.assertEqual(seen_intervals, {1, 2, 3}) + + +class AllFutureBaselineTests(unittest.TestCase): + def test_training_baseline_matches_first_onset_exposure(self): + dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset) + dataset.vocab_size = 7 + dataset.min_history_events = 1 + dataset.min_future_events = 2 + patient = { + "times": np.asarray([1.0, 3.0, 5.0], dtype=np.float32), + "labels": np.asarray([3, 4, 5], dtype=np.int64), + "t_obs": 5.0, + "query_intervals": [(1.0, 3.0)], + } + dataset.patients = [patient] + subset = Subset(dataset, [0]) + + stats = fit_all_future_baseline(dataset, subset, seed=17) + query = dataset.sample_query(patient, np.random.RandomState(17)) + + self.assertEqual(stats.event_count[3], 0) + self.assertEqual(stats.at_risk_exposure[3], 0.0) + self.assertEqual(stats.event_count[4], 1) + self.assertEqual(stats.event_count[5], 1) + self.assertAlmostEqual(stats.at_risk_exposure[4], 3.0 - query, places=5) + self.assertAlmostEqual(stats.at_risk_exposure[5], 5.0 - query, places=5) + self.assertAlmostEqual( + float(F.softplus(torch.tensor(stats.bias[4]))), + float(stats.rate[4]), + places=6, + ) + + def test_model_starts_exactly_at_fitted_output_baseline(self): + baseline_rate = torch.tensor([0.0, 0.0, 0.0, 0.02, 0.05, 0.10]) + baseline_bias = torch.zeros_like(baseline_rate) + baseline_bias[3:] = _inverse_softplus(baseline_rate[3:]) + model = DeepHealth( + vocab_size=6, + n_embd=4, + n_head=1, + n_layer=1, + n_types=1, + n_cont_types=0, + n_categories=1, + cont_type_ids=[], + target_mode="all_future", + time_mode="absolute", + dist_mode="weibull", + model_architecture="transformer_ffn_v1", + risk_head_bias=True, + risk_head_bias_init=baseline_bias, + ) + + torch.testing.assert_close(model.risk_head.weight, torch.zeros_like(model.risk_head.weight)) + output_rate = F.softplus(model.risk_head(torch.randn(3, 4))) + torch.testing.assert_close(output_rate[:, 3:], baseline_rate[None, 3:].expand(3, -1)) + torch.testing.assert_close( + F.softplus(model.rho_head.bias), + torch.ones_like(model.rho_head.bias), + atol=2e-5, + rtol=0.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_all_future.py b/train_all_future.py index ab7ac3c..5517d6b 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -5,8 +5,8 @@ Training samples are patient-level. For each patient and each __getitem__ call, AllFutureHealthDataset randomly samples a query time t_query, uses events at or before t_query as history, and uses events after t_query as the future target set. -Validation/test samples are deterministic query points built from future event -times, then split by patient. +All splits use the same patient/interval/time-uniform query distribution. +Validation/test keep one deterministic query draw per patient. """ from __future__ import annotations @@ -37,13 +37,15 @@ from model_architectures import ( SUPPORTED_MODEL_ARCHITECTURES, ) from models import DeepHealth -from targets import PAD_IDX, RESERVED_IDX +from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX from train_util import ( + AllFutureBaselineStats, ContinuousRobustScalerStats, configure_torch_for_training, create_unique_run_dir, format_extra_info_types, fit_continuous_robust_scaler, + fit_all_future_baseline, get_lr, get_model_parameter_counts, load_extra_info_types_file, @@ -113,6 +115,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dist_mode", type=str, default="exponential", choices=["exponential", "weibull"]) parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument( + "--risk_head_bias", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Initialize a learnable all-future output bias from training-only " + "marginal first-onset rates" + ), + ) parser.add_argument( "--model_architecture", type=str, @@ -179,6 +190,7 @@ def build_model( args: argparse.Namespace, dataset: AllFutureHealthDataset, scaler_stats: ContinuousRobustScalerStats, + baseline_stats: AllFutureBaselineStats | None, ) -> DeepHealth: if tuple(int(x) for x in dataset.cont_type_ids) != scaler_stats.cont_type_ids: raise ValueError( @@ -186,6 +198,8 @@ def build_model( ) center = scaler_stats.center if dataset.n_cont_types > 0 else None scale = scaler_stats.scale if dataset.n_cont_types > 0 else None + if args.risk_head_bias and baseline_stats is None: + raise ValueError("baseline_stats is required when risk_head_bias is enabled") return DeepHealth( vocab_size=dataset.vocab_size, n_embd=args.n_embd, @@ -204,11 +218,17 @@ def build_model( dist_mode=args.dist_mode, dropout=args.dropout, model_architecture=args.model_architecture, + risk_head_bias=args.risk_head_bias, + risk_head_bias_init=( + baseline_stats.bias + if args.risk_head_bias and baseline_stats is not None + else None + ), ) def build_criterion(args: argparse.Namespace): - ignored_idx = {PAD_IDX, RESERVED_IDX} + ignored_idx = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX} if args.dist_mode == "exponential": return build_loss("exponential", ignored_idx=ignored_idx) if args.dist_mode == "weibull": @@ -224,9 +244,7 @@ def compute_all_future_loss( device: torch.device, ) -> torch.Tensor: required_keys = set(MODEL_INPUT_KEYS) - required_keys.update(("future_targets", "exposure")) - if args.dist_mode == "weibull": - required_keys.add("future_dt") + required_keys.update(("future_targets", "future_dt", "exposure")) batch = move_batch_to_device( {key: batch[key] for key in required_keys}, device, @@ -250,6 +268,8 @@ def compute_all_future_loss( logits=logits, targets=batch["future_targets"], exposure=batch["exposure"], + dt=batch["future_dt"], + history=batch["event_seq"], ) elif args.dist_mode == "weibull": loss = criterion( @@ -258,6 +278,7 @@ def compute_all_future_loss( targets=batch["future_targets"], dt=batch["future_dt"], exposure=batch["exposure"], + history=batch["event_seq"], ) else: raise ValueError(f"Unknown dist_mode: {args.dist_mode}") @@ -325,6 +346,7 @@ def build_metadata( val_subset, test_subset, scaler_stats: ContinuousRobustScalerStats, + baseline_stats: AllFutureBaselineStats | None, ) -> Dict[str, Any]: scaler_metadata = scaler_stats.as_metadata() return { @@ -342,6 +364,17 @@ def build_metadata( "all_future_min_history_events": int(args.min_history_events), "all_future_min_future_events": int(args.min_future_events), "all_future_validation_query_seed": int(args.validation_query_seed), + "all_future_query_distribution": "patient_interval_time_uniform", + "all_future_queries_per_validation_patient": 1, + "all_future_likelihood": "first_onset_survival_v2", + "all_future_prevalent_outcomes_excluded": True, + "risk_head_bias": bool(args.risk_head_bias), + "risk_head_bias_weight_decay": 0.0 if args.risk_head_bias else None, + "risk_head_baseline": ( + baseline_stats.as_metadata() + if args.risk_head_bias and baseline_stats is not None + else {"method": "none"} + ), "extra_info_types_file": ( Path(args.extra_info_types_file).name if args.extra_info_types_file is not None @@ -476,6 +509,27 @@ def main() -> None: f"max_per_feature={int(scaler_stats.observation_count.max()):,}" ) + baseline_stats = None + if args.risk_head_bias: + logger.info( + "Fitting all-future output baseline on one seeded query per " + f"training patient: patients={len(train_subset):,}" + ) + baseline_stats = fit_all_future_baseline( + train_dataset, + train_subset, + seed=args.seed, + ) + baseline_summary = baseline_stats.as_metadata() + rate_summary = baseline_summary["rate_per_year"] + logger.info( + "All-future baseline rates/year: " + f"min={rate_summary['min']:.8g}, " + f"median={rate_summary['median']:.8g}, " + f"max={rate_summary['max']:.8g}, " + f"first_onsets={baseline_summary['observed_first_onsets']:,}" + ) + train_loader = DataLoader( train_subset, batch_size=args.batch_size, @@ -507,15 +561,37 @@ def main() -> None: prefetch_factor=2 if args.num_workers > 0 else None, ) - model = build_model(args, train_dataset, scaler_stats=scaler_stats).to(device) + model = build_model( + args, + train_dataset, + scaler_stats=scaler_stats, + baseline_stats=baseline_stats, + ).to(device) parameter_counts = get_model_parameter_counts(model) logger.info( "Model parameters: " f"total={parameter_counts['model_parameter_count']:,}, " f"trainable={parameter_counts['trainable_parameter_count']:,}" ) + if model.risk_head.bias is None: + optimizer_parameters = model.parameters() + else: + optimizer_parameters = [ + { + "params": [ + parameter + for parameter in model.parameters() + if parameter is not model.risk_head.bias + ], + "weight_decay": args.weight_decay, + }, + { + "params": [model.risk_head.bias], + "weight_decay": 0.0, + }, + ] optimizer = AdamW( - model.parameters(), + optimizer_parameters, lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, @@ -531,6 +607,7 @@ def main() -> None: val_subset, test_subset, scaler_stats, + baseline_stats, ) train_metadata.update(parameter_counts) save_config( diff --git a/train_util.py b/train_util.py index 9d13ecf..c4b3095 100644 --- a/train_util.py +++ b/train_util.py @@ -18,6 +18,7 @@ from torch.utils.data import Subset from dataset import AllFutureHealthDataset, HealthDataset from models import DeepHealth +from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX @dataclass(frozen=True) @@ -42,6 +43,134 @@ class ContinuousRobustScalerStats: } +@dataclass(frozen=True) +class AllFutureBaselineStats: + """Training-only marginal first-onset rates for output initialization.""" + + event_count: np.ndarray + at_risk_exposure: np.ndarray + rate: np.ndarray + bias: np.ndarray + query_count: int + rate_floor: float + rate_ceiling: float + + def as_metadata(self) -> Dict[str, Any]: + positive = self.rate[self.at_risk_exposure > 0] + if positive.size: + rate_summary = { + "min": float(positive.min()), + "median": float(np.median(positive)), + "max": float(positive.max()), + } + else: + rate_summary = {"min": 0.0, "median": 0.0, "max": 0.0} + return { + "method": "training_marginal_first_onset_rate", + "fitted_on": "one_seeded_query_per_training_patient", + "query_distribution": "patient_interval_time_uniform", + "query_count": int(self.query_count), + "observed_first_onsets": int(self.event_count.sum()), + "rate_floor": float(self.rate_floor), + "rate_ceiling": float(self.rate_ceiling), + "rate_per_year": rate_summary, + "checkpoint_parameter": "risk_head.bias", + } + + +def fit_all_future_baseline( + dataset: AllFutureHealthDataset, + subset: Subset, + *, + seed: int, + ignored_idx: Iterable[int] = (PAD_IDX, RESERVED_IDX, NO_EVENT_IDX), + rate_floor: float = 1e-6, + rate_ceiling: float = 10.0, +) -> AllFutureBaselineStats: + """Fit marginal per-outcome rates under the training query distribution.""" + if subset.dataset is not dataset: + raise ValueError("subset must reference the all-future training dataset") + if rate_floor <= 0 or rate_ceiling <= rate_floor: + raise ValueError("Require 0 < rate_floor < rate_ceiling") + + subset_indices = np.asarray(subset.indices, dtype=np.int64) + if subset_indices.ndim != 1 or subset_indices.size == 0: + raise ValueError("training subset must contain at least one patient") + + vocab_size = int(dataset.vocab_size) + ignored = { + int(idx) + for idx in ignored_idx + if 0 <= int(idx) < vocab_size + } + event_count = np.zeros(vocab_size, dtype=np.int64) + exposure_adjustment = np.zeros(vocab_size, dtype=np.float64) + common_exposure = 0.0 + rng = np.random.RandomState(int(seed)) + + for patient_index in subset_indices.tolist(): + patient = dataset.patients[int(patient_index)] + t_query = dataset.sample_query(patient, rng) + times = np.asarray(patient["times"], dtype=np.float64) + labels = np.asarray(patient["labels"], dtype=np.int64) + censor_time = max(float(patient["t_obs"]) - t_query, 0.0) + common_exposure += censor_time + + prevalent = { + int(label) + for label in labels[times <= t_query].tolist() + if 0 <= int(label) < vocab_size and int(label) not in ignored + } + for label in prevalent: + exposure_adjustment[label] -= censor_time + + # Keep the earliest future occurrence defensively; prepared disease + # events are already deduplicated to their first occurrence. + first_future_dt: Dict[int, float] = {} + for label, event_time in zip(labels[times > t_query], times[times > t_query]): + label = int(label) + if label in ignored or label in prevalent or not (0 <= label < vocab_size): + continue + event_dt = max(float(event_time) - t_query, 0.0) + previous = first_future_dt.get(label) + if previous is None or event_dt < previous: + first_future_dt[label] = event_dt + for label, event_dt in first_future_dt.items(): + event_count[label] += 1 + exposure_adjustment[label] += event_dt - censor_time + + at_risk_exposure = common_exposure + exposure_adjustment + at_risk_exposure = np.maximum(at_risk_exposure, 0.0) + rate = np.full(vocab_size, float(rate_floor), dtype=np.float64) + has_exposure = at_risk_exposure > 0 + rate[has_exposure] = ( + event_count[has_exposure].astype(np.float64) + / at_risk_exposure[has_exposure] + ) + rate = np.clip(rate, float(rate_floor), float(rate_ceiling)) + + for idx in ignored: + at_risk_exposure[idx] = 0.0 + rate[idx] = 0.0 + bias = np.zeros(vocab_size, dtype=np.float64) + valid = np.ones(vocab_size, dtype=bool) + if ignored: + valid[np.asarray(sorted(ignored), dtype=np.int64)] = False + bias[valid] = np.log(np.expm1(rate[valid])) + + if not np.isfinite(bias).all(): + raise RuntimeError("All-future baseline initialization is non-finite") + return AllFutureBaselineStats( + event_count=event_count, + at_risk_exposure=at_risk_exposure.astype(np.float32), + rate=rate.astype(np.float32), + bias=bias.astype(np.float32), + query_count=int(subset_indices.size), + rate_floor=float(rate_floor), + rate_ceiling=float(rate_ceiling), + ) + + def fit_continuous_robust_scaler( dataset: HealthDataset | AllFutureHealthDataset, subset: Subset,