From 1246ea58b1519da646996544b2d2ecd2a428988d Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 11:19:48 +0200 Subject: [PATCH 01/20] add semi-synthetic causal simulator with observed treatment Introduces a new simulation DGP where treatment assignment is kept from real data and only the outcome is simulated via hand-crafted oracle features extracted from patient histories. Features are split into baseline risk (r_B) and longitudinal (r_L) groups, with configurable coefficients and treatment effects defined in YAML. New files: - config_semisynthetic.py: dataclass configs + factory function - oracle_features.py: 10 feature extractors (7 baseline, 3 longitudinal) - semisynthetic_simulator.py: simulator class with same output format - simulate_semisynthetic.py: pipeline entry point - calibrate_semisynthetic.py: standalone calibration diagnostics + plots - simulate_semisynthetic.yaml: default config - tests for feature extraction and simulator (17 tests) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../causal/simulate_semisynthetic.yaml | 77 +++ .../main_causal/calibrate_semisynthetic.py | 245 +++++++++ .../main_causal/simulate_semisynthetic.py | 60 ++ .../simulation/config_semisynthetic.py | 109 ++++ .../modules/simulation/oracle_features.py | 257 +++++++++ .../simulation/semisynthetic_simulator.py | 514 ++++++++++++++++++ .../test_modules/test_simulation/__init__.py | 0 .../test_simulation/test_oracle_features.py | 250 +++++++++ .../test_simulation/test_semisynthetic.py | 240 ++++++++ 9 files changed, 1752 insertions(+) create mode 100644 corebehrt/configs/causal/simulate_semisynthetic.yaml create mode 100644 corebehrt/main_causal/calibrate_semisynthetic.py create mode 100644 corebehrt/main_causal/simulate_semisynthetic.py create mode 100644 corebehrt/modules/simulation/config_semisynthetic.py create mode 100644 corebehrt/modules/simulation/oracle_features.py create mode 100644 corebehrt/modules/simulation/semisynthetic_simulator.py create mode 100644 tests/test_modules/test_simulation/__init__.py create mode 100644 tests/test_modules/test_simulation/test_oracle_features.py create mode 100644 tests/test_modules/test_simulation/test_semisynthetic.py diff --git a/corebehrt/configs/causal/simulate_semisynthetic.yaml b/corebehrt/configs/causal/simulate_semisynthetic.yaml new file mode 100644 index 00000000..1c9123b9 --- /dev/null +++ b/corebehrt/configs/causal/simulate_semisynthetic.yaml @@ -0,0 +1,77 @@ +# -------------------------------------------------------------------------- +# Semi-Synthetic Simulation: Observed Treatment, Simulated Outcome +# -------------------------------------------------------------------------- +# Treatment assignment (A_i) and index dates are taken from the real data. +# Only the outcome is simulated from hand-crafted oracle features. + +logging: + level: INFO + path: ./outputs/logs/causal + +paths: + data: ./example_data/synthea_meds_causal + splits: ["tuning"] + outcomes: ./outputs/causal/semisynthetic_outcomes + +seed: 42 +debug: false +min_num_codes: 3 +exposure_code: "EXPOSURE" + +# -------------------------------------------------------------------------- +# Oracle Feature Extraction +# -------------------------------------------------------------------------- +features: + code_prefixes: + diagnosis: "D/" + medication: "M/" + procedure: "P/" + admission: "ADM/" + lookback_days: 365 + recent_window_days: 90 + burst_window_days: 30 + motif_window_days: 30 + standardize: true + +# -------------------------------------------------------------------------- +# Outcomes +# -------------------------------------------------------------------------- +outcomes: + OUTCOME: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + # Baseline risk features (r_B) + baseline_coefficients: + recent_event_count: 0.3 + disease_burden: 0.2 + medication_count: 0.15 + utilization_intensity: 0.1 + age: 0.4 + chronic_disease_count: 0.25 + code_diversity: 0.1 + # Longitudinal features (r_L) + longitudinal_coefficients: + event_recency: -0.15 + recent_burst_ratio: 0.2 + sequence_motif_count: 0.1 + interactions: + - features: [disease_burden, age] + coefficient: 0.1 + noise_scale: 0.1 + treatment_effect: + mode: constant + delta: 1.0 + + OUTCOME_NULL: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + baseline_coefficients: + disease_burden: 0.2 + age: 0.4 + longitudinal_coefficients: {} + noise_scale: 0.1 + treatment_effect: + mode: constant + delta: 0.0 diff --git a/corebehrt/main_causal/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py new file mode 100644 index 00000000..6ba1e39a --- /dev/null +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -0,0 +1,245 @@ +"""Calibration diagnostics for the semi-synthetic simulation. + +Runs the feature extraction and probability computation pipeline +without Bernoulli sampling, then prints a calibration report and +saves diagnostic plots. +""" + +import logging +import os +from os.path import join + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from scipy.special import expit + +from corebehrt.functional.setup.args import get_args +from corebehrt.functional.utils.azure_save import save_figure_with_azure_copy +from corebehrt.modules.features.loader import ShardLoader +from corebehrt.modules.setup.causal.directory import CausalDirectoryPreparer +from corebehrt.modules.setup.config import load_config +from corebehrt.modules.simulation.config_semisynthetic import ( + create_semisynthetic_config, +) +from corebehrt.modules.simulation.oracle_features import extract_oracle_features +from corebehrt.modules.simulation.plot import plot_probability_distributions +from corebehrt.modules.simulation.semisynthetic_simulator import ( + SemiSyntheticCausalSimulator, +) + +logger = logging.getLogger("calibrate") + +CONFIG_PATH = "./corebehrt/configs/causal/simulate_semisynthetic.yaml" + + +def main_calibrate(config_path): + cfg = load_config(config_path) + CausalDirectoryPreparer(cfg).setup_simulate_from_sequence() + + shard_loader = ShardLoader(cfg.paths.data, cfg.paths.splits) + sim_config = create_semisynthetic_config(cfg) + simulator = SemiSyntheticCausalSimulator(sim_config) + + # Accumulate across shards + all_features = [] + all_is_exposed = [] + all_probas = {} # outcome_name -> {"P0": [], "P1": []} + all_tau = {} # outcome_name -> [] + + for shard, _ in shard_loader(): + pids, is_exposed, index_dates = simulator._extract_treatment_and_index_dates( + shard + ) + if len(pids) == 0: + continue + + history_df = simulator._filter_to_pre_index(shard, index_dates) + history_df, pids, is_exposed, index_dates = simulator._apply_min_num_codes( + history_df, pids, is_exposed, index_dates + ) + if len(pids) == 0: + continue + + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, sim_config.features + ) + all_features.append(features_df.assign(is_exposed=is_exposed)) + all_is_exposed.append(is_exposed) + + for outcome_name, outcome_cfg in sim_config.outcomes.items(): + eta_0 = simulator._compute_eta_0(features_df, outcome_cfg.outcome_model) + tau = simulator._compute_tau(features_df, outcome_cfg.treatment_effect) + p0 = expit(eta_0) + p1 = expit(eta_0 + tau) + + all_probas.setdefault(outcome_name, {"P0": [], "P1": []}) + all_probas[outcome_name]["P0"].append(p0) + all_probas[outcome_name]["P1"].append(p1) + all_tau.setdefault(outcome_name, []).append(tau) + + if not all_features: + logger.error("No patients found across shards.") + return + + features_combined = pd.concat(all_features, ignore_index=True) + is_exposed_combined = np.concatenate(all_is_exposed) + output_dir = sim_config.paths.outcomes + figs_dir = join(output_dir, "figs") + os.makedirs(figs_dir, exist_ok=True) + + # --- Feature diagnostics --- + _print_feature_diagnostics(features_combined, is_exposed_combined) + + # --- Probability and causal effect diagnostics --- + probas_for_plot = {} + for outcome_name in sim_config.outcomes: + p0 = np.concatenate(all_probas[outcome_name]["P0"]) + p1 = np.concatenate(all_probas[outcome_name]["P1"]) + tau_arr = np.concatenate(all_tau[outcome_name]) + probas_for_plot[outcome_name] = {"P0": p0, "P1": p1} + + _print_probability_diagnostics(outcome_name, p0, p1, is_exposed_combined) + _print_causal_diagnostics(outcome_name, p0, p1, is_exposed_combined) + _plot_ite_histogram(tau_arr, outcome_name, figs_dir) + + plot_probability_distributions(probas_for_plot, figs_dir) + + # --- SMD love plot --- + feature_cols = [c for c in features_combined.columns if c != "is_exposed"] + _plot_smd_love_plot(features_combined, feature_cols, is_exposed_combined, figs_dir) + + logger.info("Calibration complete. Plots saved to %s", figs_dir) + + +# --------------------------------------------------------------------------- +# Reporting helpers +# --------------------------------------------------------------------------- + + +def _print_feature_diagnostics(features_df: pd.DataFrame, is_exposed: np.ndarray): + feature_cols = [c for c in features_df.columns if c != "is_exposed"] + print("\n" + "=" * 80) + print("FEATURE DIAGNOSTICS") + print("=" * 80) + + quantiles = [0.05, 0.25, 0.5, 0.75, 0.95] + header = f"{'Feature':<30} {'mean':>8} {'std':>8} {'min':>8}" + for q in quantiles: + header += f" {'p' + str(int(q * 100)):>6}" + header += f" {'max':>8} {'SMD':>8}" + print(header) + print("-" * len(header)) + + for col in feature_cols: + vals = features_df[col].values + q_vals = np.quantile(vals, quantiles) + smd = _compute_smd(vals[is_exposed], vals[~is_exposed]) + row = ( + f"{col:<30} {np.mean(vals):>8.3f} {np.std(vals):>8.3f} {np.min(vals):>8.3f}" + ) + for qv in q_vals: + row += f" {qv:>6.3f}" + row += f" {np.max(vals):>8.3f} {smd:>8.3f}" + print(row) + print() + + +def _print_probability_diagnostics( + outcome_name: str, p0: np.ndarray, p1: np.ndarray, is_exposed: np.ndarray +): + print(f"\n--- Probability diagnostics: {outcome_name} ---") + for label, probs in [("P(Y(0))", p0), ("P(Y(1))", p1)]: + print( + f" {label}: mean={np.mean(probs):.4f}, std={np.std(probs):.4f}, " + f"min={np.min(probs):.4f}, max={np.max(probs):.4f}, " + f"median={np.median(probs):.4f}" + ) + frac_extreme_p0 = np.mean((p0 < 0.01) | (p0 > 0.80)) + frac_extreme_p1 = np.mean((p1 < 0.01) | (p1 > 0.80)) + print( + f" Extreme probabilities (<0.01 or >0.80): P0={frac_extreme_p0:.3f}, P1={frac_extreme_p1:.3f}" + ) + + # Expected prevalence under factual assignment + factual_prob = np.where(is_exposed, p1, p0) + print(f" Expected factual prevalence: {np.mean(factual_prob):.4f}") + + +def _print_causal_diagnostics( + outcome_name: str, p0: np.ndarray, p1: np.ndarray, is_exposed: np.ndarray +): + ite = p1 - p0 + ate = np.mean(ite) + att = np.mean(ite[is_exposed]) if np.any(is_exposed) else float("nan") + atc = np.mean(ite[~is_exposed]) if np.any(~is_exposed) else float("nan") + rr = np.mean(p1) / np.mean(p0) if np.mean(p0) > 0 else float("nan") + + print(f"\n--- Causal effect diagnostics: {outcome_name} ---") + print(f" True ATE = {ate:.4f}") + print(f" True ATT = {att:.4f}") + print(f" True ATC = {atc:.4f}") + print(f" True RR = {rr:.4f}") + + +def _compute_smd(treated: np.ndarray, control: np.ndarray) -> float: + """Standardized mean difference.""" + pooled_std = np.sqrt((np.var(treated) + np.var(control)) / 2) + if pooled_std == 0: + return 0.0 + return (np.mean(treated) - np.mean(control)) / pooled_std + + +# --------------------------------------------------------------------------- +# Plotting helpers +# --------------------------------------------------------------------------- + + +def _plot_ite_histogram(tau: np.ndarray, outcome_name: str, figs_dir: str): + fig, ax = plt.subplots(figsize=(8, 6)) + ax.hist(tau, bins=50, edgecolor="black", alpha=0.7) + ax.set_title(f"ITE Distribution: {outcome_name}") + ax.set_xlabel("Individual Treatment Effect (logit scale)") + ax.set_ylabel("Count") + ax.axvline( + np.mean(tau), color="red", linestyle="--", label=f"Mean={np.mean(tau):.3f}" + ) + ax.legend() + ax.spines["right"].set_visible(False) + ax.spines["top"].set_visible(False) + save_figure_with_azure_copy( + fig, join(figs_dir, f"ite_histogram_{outcome_name}.png") + ) + + +def _plot_smd_love_plot( + features_df: pd.DataFrame, + feature_cols, + is_exposed: np.ndarray, + figs_dir: str, +): + smds = [] + for col in feature_cols: + vals = features_df[col].values + smds.append(_compute_smd(vals[is_exposed], vals[~is_exposed])) + + fig, ax = plt.subplots(figsize=(8, max(4, len(feature_cols) * 0.4))) + y_pos = np.arange(len(feature_cols)) + ax.barh(y_pos, smds, color="#3498db", edgecolor="black", alpha=0.7) + ax.set_yticks(y_pos) + ax.set_yticklabels(feature_cols) + ax.set_xlabel("Standardized Mean Difference") + ax.set_title("Feature Balance (Treated vs Control)") + ax.axvline(0, color="black", linewidth=0.8) + ax.axvline(0.1, color="red", linestyle="--", alpha=0.5, label="SMD=0.1") + ax.axvline(-0.1, color="red", linestyle="--", alpha=0.5) + ax.legend() + ax.spines["right"].set_visible(False) + ax.spines["top"].set_visible(False) + plt.tight_layout() + save_figure_with_azure_copy(fig, join(figs_dir, "smd_love_plot.png")) + + +if __name__ == "__main__": + args = get_args(CONFIG_PATH) + main_calibrate(args.config_path) diff --git a/corebehrt/main_causal/simulate_semisynthetic.py b/corebehrt/main_causal/simulate_semisynthetic.py new file mode 100644 index 00000000..46fcb8d7 --- /dev/null +++ b/corebehrt/main_causal/simulate_semisynthetic.py @@ -0,0 +1,60 @@ +from corebehrt.functional.setup.args import get_args +from corebehrt.modules.setup.config import load_config +from corebehrt.modules.setup.causal.directory import CausalDirectoryPreparer +from corebehrt.modules.features.loader import ShardLoader +from corebehrt.modules.simulation.semisynthetic_simulator import ( + SemiSyntheticCausalSimulator as CausalSimulator, +) +from corebehrt.modules.simulation.config_semisynthetic import ( + create_semisynthetic_config, +) +from collections import defaultdict +import pandas as pd +from os.path import join +import logging +from tqdm import tqdm + +logger = logging.getLogger("simulate") + + +CONFIG_PATH = "./corebehrt/configs/causal/simulate_semisynthetic.yaml" + + +def main_simulate(config_path): + cfg = load_config(config_path) + + # Setup directories + CausalDirectoryPreparer(cfg).setup_simulate_from_sequence() + + shard_loader = ShardLoader(cfg.paths.data, cfg.paths.splits) + simulation_config = create_semisynthetic_config(cfg) + simulator = CausalSimulator(simulation_config) + simulate(shard_loader, simulator, cfg.paths.outcomes) + + +def simulate(shard_loader: ShardLoader, simulator: CausalSimulator, outcomes_dir: str): + """ + Simulates outcomes by processing data shards in a single pass. + + Iterates through each data shard, calls simulate_dataset, + aggregates the results, and saves each outcome type to a separate CSV file. + """ + logger.info("--- Starting semi-synthetic simulation ---") + simulated_outcomes = defaultdict(list) + for shard, _ in tqdm(shard_loader(), desc="Simulating from shards"): + simulated_temp = simulator.simulate_dataset(shard) + for k, df in simulated_temp.items(): + if not df.empty: + simulated_outcomes[k].append(df) + + logger.info("--- Simulation complete, saving results ---") + + for k, df_list in simulated_outcomes.items(): + if df_list: + df = pd.concat(df_list, ignore_index=True) + df.to_csv(join(outcomes_dir, f"{k}.csv"), index=False) + + +if __name__ == "__main__": + args = get_args(CONFIG_PATH) + main_simulate(args.config_path) diff --git a/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py new file mode 100644 index 00000000..bb86bacc --- /dev/null +++ b/corebehrt/modules/simulation/config_semisynthetic.py @@ -0,0 +1,109 @@ +from dataclasses import dataclass, field +from typing import Dict, List + + +@dataclass +class PathsConfig: + """File paths for the semi-synthetic simulation.""" + + data: str + splits: List[str] + outcomes: str + + +@dataclass +class CodePrefixConfig: + """Maps concept types to their code prefixes in the MEDS data.""" + + diagnosis: str = "D/" + medication: str = "M/" + procedure: str = "P/" + admission: str = "ADM/" + + +@dataclass +class FeatureConfig: + """Controls oracle feature extraction parameters.""" + + code_prefixes: CodePrefixConfig = field(default_factory=CodePrefixConfig) + lookback_days: int = 365 + recent_window_days: int = 90 + burst_window_days: int = 30 + motif_window_days: int = 30 + standardize: bool = True + + +@dataclass +class OutcomeModelConfig: + """Outcome model: eta^(0) = beta_0 + f_B(r_B) + f_L(r_L).""" + + run_in_days: int = 1 + beta_0: float = -2.0 + baseline_coefficients: Dict[str, float] = field(default_factory=dict) + longitudinal_coefficients: Dict[str, float] = field(default_factory=dict) + interactions: List[Dict] = field(default_factory=list) + noise_scale: float = 0.0 + + +@dataclass +class TreatmentEffectConfig: + """Treatment effect: constant (tau=delta) or heterogeneous (tau=delta_0 + g(r_B)).""" + + mode: str = "constant" + delta: float = 1.0 + delta_0: float = 0.5 + heterogeneous_coefficients: Dict[str, float] = field(default_factory=dict) + + +@dataclass +class SemiSyntheticOutcomeConfig: + """Bundles outcome model and treatment effect for one outcome.""" + + outcome_model: OutcomeModelConfig + treatment_effect: TreatmentEffectConfig + + +@dataclass +class SemiSyntheticSimulationConfig: + """Top-level configuration for the semi-synthetic simulation.""" + + paths: PathsConfig + features: FeatureConfig + outcomes: Dict[str, SemiSyntheticOutcomeConfig] + seed: int = 42 + debug: bool = False + min_num_codes: int = 5 + exposure_code: str = "EXPOSURE" + + +def create_semisynthetic_config(cfg) -> SemiSyntheticSimulationConfig: + """Parse a config object/dict into a SemiSyntheticSimulationConfig.""" + paths_config = PathsConfig(**cfg["paths"]) + + prefix_cfg = CodePrefixConfig(**cfg.get("features", {}).get("code_prefixes", {})) + feature_dict = dict(cfg.get("features", {})) + feature_dict.pop("code_prefixes", None) + feature_config = FeatureConfig(code_prefixes=prefix_cfg, **feature_dict) + + outcomes_config = {} + for name, outcome_data in cfg["outcomes"].items(): + om_data = dict(outcome_data.get("outcome_model", {})) + outcome_model = OutcomeModelConfig(**om_data) + + te_data = dict(outcome_data.get("treatment_effect", {})) + treatment_effect = TreatmentEffectConfig(**te_data) + + outcomes_config[name] = SemiSyntheticOutcomeConfig( + outcome_model=outcome_model, + treatment_effect=treatment_effect, + ) + + return SemiSyntheticSimulationConfig( + paths=paths_config, + features=feature_config, + outcomes=outcomes_config, + seed=cfg.get("seed", 42), + debug=cfg.get("debug", False), + min_num_codes=cfg.get("min_num_codes", 5), + exposure_code=cfg.get("exposure_code", "EXPOSURE"), + ) diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py new file mode 100644 index 00000000..233a6230 --- /dev/null +++ b/corebehrt/modules/simulation/oracle_features.py @@ -0,0 +1,257 @@ +"""Extract hand-crafted oracle features from pre-index patient histories.""" + +import logging +from typing import List, Tuple + +import numpy as np +import pandas as pd + +from corebehrt.constants.data import BIRTH_CODE, CONCEPT_COL, PID_COL, TIMESTAMP_COL +from corebehrt.modules.simulation.config_semisynthetic import FeatureConfig + +logger = logging.getLogger("oracle_features") + +BASELINE_FEATURES = [ + "recent_event_count", + "disease_burden", + "medication_count", + "utilization_intensity", + "age", + "chronic_disease_count", + "code_diversity", +] +LONGITUDINAL_FEATURES = [ + "event_recency", + "recent_burst_ratio", + "sequence_motif_count", +] + + +def extract_oracle_features( + history_df: pd.DataFrame, + pids: np.ndarray, + index_dates: pd.Series, + feature_config: FeatureConfig, +) -> Tuple[pd.DataFrame, List[str], List[str]]: + """Extract oracle features from pre-index patient histories. + + Args: + history_df: MEDS DataFrame already filtered to pre-index events + pids: array of patient IDs to extract features for + index_dates: per-patient index dates (Series: PID -> Timestamp) + feature_config: feature extraction configuration + + Returns: + features_df: DataFrame with PID_COL as index, one column per feature + baseline_feature_names: list of r_B feature names present in the output + longitudinal_feature_names: list of r_L feature names present in the output + """ + prefixes = feature_config.code_prefixes + + features = {} + # Baseline risk features + features["recent_event_count"] = _compute_recent_event_count( + history_df, pids, index_dates, feature_config.recent_window_days + ) + features["disease_burden"] = _compute_disease_burden( + history_df, pids, index_dates, prefixes.diagnosis, feature_config.lookback_days + ) + features["medication_count"] = _compute_medication_count( + history_df, pids, index_dates, prefixes.medication, feature_config.lookback_days + ) + features["utilization_intensity"] = _compute_utilization_intensity( + history_df, pids, index_dates, feature_config.lookback_days + ) + features["age"] = _compute_age(history_df, pids, index_dates) + features["chronic_disease_count"] = _compute_chronic_disease_count( + history_df, pids, prefixes.diagnosis + ) + features["code_diversity"] = _compute_code_diversity(history_df, pids) + + # Longitudinal features + features["event_recency"] = _compute_event_recency(history_df, pids, index_dates) + features["recent_burst_ratio"] = _compute_recent_burst_ratio( + history_df, + pids, + index_dates, + feature_config.burst_window_days, + feature_config.lookback_days, + ) + features["sequence_motif_count"] = _compute_sequence_motif_count( + history_df, + pids, + index_dates, + prefixes.diagnosis, + prefixes.medication, + feature_config.motif_window_days, + ) + + features_df = pd.DataFrame(features, index=pids) + features_df.index.name = PID_COL + + if feature_config.standardize: + features_df = _standardize(features_df) + + baseline_names = [n for n in BASELINE_FEATURES if n in features_df.columns] + longitudinal_names = [n for n in LONGITUDINAL_FEATURES if n in features_df.columns] + return features_df, baseline_names, longitudinal_names + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _filter_by_prefix_and_window(history_df, index_dates, prefix, window_days): + """Filter events matching code prefix within lookback window per patient.""" + mask = history_df[CONCEPT_COL].str.startswith(prefix) + filtered = history_df[mask].copy() + if filtered.empty: + logger.warning("No codes found with prefix '%s'", prefix) + return filtered + if window_days is not None: + cutoff = index_dates.reindex(filtered[PID_COL]).values - pd.Timedelta( + days=window_days + ) + filtered = filtered[filtered[TIMESTAMP_COL].values >= cutoff] + return filtered + + +def _count_per_patient(filtered_df, pids, count_col=None, unique=False): + """Count (total or unique) events per patient, filling missing with 0.""" + if filtered_df.empty: + return pd.Series(0, index=pids, dtype=int) + grouped = filtered_df.groupby(PID_COL) + if unique: + col = count_col if count_col else CONCEPT_COL + counts = grouped[col].nunique() + else: + counts = grouped.size() + return counts.reindex(pids, fill_value=0) + + +# --------------------------------------------------------------------------- +# Baseline risk features (r_B) +# --------------------------------------------------------------------------- + + +def _compute_recent_event_count(history_df, pids, index_dates, recent_window_days): + filtered = _filter_by_prefix_and_window( + history_df, index_dates, "", recent_window_days + ) + return _count_per_patient(filtered, pids) + + +def _compute_disease_burden(history_df, pids, index_dates, diag_prefix, lookback_days): + filtered = _filter_by_prefix_and_window( + history_df, index_dates, diag_prefix, lookback_days + ) + return _count_per_patient(filtered, pids, unique=True) + + +def _compute_medication_count(history_df, pids, index_dates, med_prefix, lookback_days): + filtered = _filter_by_prefix_and_window( + history_df, index_dates, med_prefix, lookback_days + ) + return _count_per_patient(filtered, pids, unique=True) + + +def _compute_utilization_intensity(history_df, pids, index_dates, lookback_days): + filtered = _filter_by_prefix_and_window(history_df, index_dates, "", lookback_days) + return _count_per_patient(filtered, pids) + + +def _compute_age(history_df, pids, index_dates): + dob_events = history_df[history_df[CONCEPT_COL] == BIRTH_CODE] + dob_per_patient = dob_events.groupby(PID_COL)[TIMESTAMP_COL].first() + age_series = pd.Series(index=pids, dtype=float) + for pid in pids: + if pid in dob_per_patient.index: + dob = dob_per_patient[pid] + idx_date = index_dates[pid] + age_series[pid] = (idx_date - dob).days / 365.25 + mean_age = age_series.mean() + age_series = age_series.fillna(mean_age) + return age_series + + +def _compute_chronic_disease_count(history_df, pids, diag_prefix): + diag_events = history_df[history_df[CONCEPT_COL].str.startswith(diag_prefix)] + if diag_events.empty: + logger.warning("No diagnosis codes found for chronic disease count") + return pd.Series(0, index=pids, dtype=int) + groups = diag_events.copy() + groups["_diag_group"] = groups[CONCEPT_COL].str[:5] + counts = groups.groupby(PID_COL)["_diag_group"].nunique() + return counts.reindex(pids, fill_value=0) + + +def _compute_code_diversity(history_df, pids): + return _count_per_patient(history_df, pids, unique=True) + + +# --------------------------------------------------------------------------- +# Longitudinal features (r_L) +# --------------------------------------------------------------------------- + + +def _compute_event_recency(history_df, pids, index_dates): + last_event = history_df.groupby(PID_COL)[TIMESTAMP_COL].max() + recency = pd.Series(index=pids, dtype=float) + for pid in pids: + if pid in last_event.index: + recency[pid] = (index_dates[pid] - last_event[pid]).days + else: + recency[pid] = np.nan + mean_recency = recency.mean() + recency = recency.fillna(mean_recency) + return recency + + +def _compute_recent_burst_ratio( + history_df, pids, index_dates, burst_window_days, lookback_days +): + burst_filtered = _filter_by_prefix_and_window( + history_df, index_dates, "", burst_window_days + ) + lookback_filtered = _filter_by_prefix_and_window( + history_df, index_dates, "", lookback_days + ) + burst_counts = _count_per_patient(burst_filtered, pids) + lookback_counts = _count_per_patient(lookback_filtered, pids) + return burst_counts / (lookback_counts + 1) + + +def _compute_sequence_motif_count( + history_df, pids, _index_dates, diag_prefix, med_prefix, motif_window_days +): + diag_events = history_df[history_df[CONCEPT_COL].str.startswith(diag_prefix)] + med_events = history_df[history_df[CONCEPT_COL].str.startswith(med_prefix)] + if diag_events.empty or med_events.empty: + logger.warning("Missing diagnosis or medication codes for motif counting") + return pd.Series(0, index=pids, dtype=int) + + motif_counts = {} + window = pd.Timedelta(days=motif_window_days) + for pid in pids: + pid_diags = diag_events[diag_events[PID_COL] == pid][TIMESTAMP_COL].values + pid_meds = med_events[med_events[PID_COL] == pid][TIMESTAMP_COL].values + count = 0 + for diag_time in pid_diags: + gaps = pid_meds - diag_time + count += int(np.sum((gaps >= np.timedelta64(0)) & (gaps <= window))) + motif_counts[pid] = count + + return pd.Series(motif_counts).reindex(pids, fill_value=0) + + +# --------------------------------------------------------------------------- +# Standardization +# --------------------------------------------------------------------------- + + +def _standardize(features_df): + means = features_df.mean() + stds = features_df.std() + stds = stds.replace(0, 1) + return (features_df - means) / stds diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py new file mode 100644 index 00000000..048d1155 --- /dev/null +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -0,0 +1,514 @@ +"""Semi-synthetic causal simulator using oracle features from real EHR data.""" + +import logging +import os +from os.path import join +from typing import Dict, Tuple + +import numpy as np +import pandas as pd +from scipy.special import expit +from sklearn.metrics import roc_auc_score + +from corebehrt.constants.causal.data import ( + CONTROL_PID_COL, + EXPOSED_PID_COL, + EXPOSURE_COL, + OUTCOME_COL, + SIMULATED_OUTCOME_CONTROL, + SIMULATED_OUTCOME_EXPOSED, + SIMULATED_PROBAS_CONTROL, + SIMULATED_PROBAS_EXPOSED, +) +from corebehrt.constants.causal.paths import ( + COUNTERFACTUALS_FILE, + INDEX_DATE_MATCHING_FILE, +) +from corebehrt.constants.data import ( + ABSPOS_COL, + CONCEPT_COL, + DEATH_CODE, + PID_COL, + TIMESTAMP_COL, +) +from corebehrt.functional.utils.filter import safe_control_pids +from corebehrt.functional.utils.time import get_hours_since_epoch +from corebehrt.modules.simulation.config_semisynthetic import ( + OutcomeModelConfig, + SemiSyntheticSimulationConfig, + TreatmentEffectConfig, +) +from corebehrt.modules.simulation.oracle_features import extract_oracle_features +from corebehrt.modules.simulation.plot import ( + plot_probability_distributions, + plot_true_effects_vs_risk_differences, +) + +logger = logging.getLogger("simulate") + +ASSIGNED_INDEX_DATE_COL = "assigned_index_date" +EXCLUDED_PREFIXES = ("OUTCOME",) + + +class SemiSyntheticCausalSimulator: + """Simulate causal outcomes from real EHR sequences using oracle features. + + Treatment assignment comes from the data (presence of exposure_code). + Outcome probabilities are computed via a parametric model over + hand-crafted oracle features extracted from pre-index history. + """ + + def __init__(self, config: SemiSyntheticSimulationConfig): + self.config = config + self.rng = np.random.default_rng(config.seed) + + def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: + """Orchestrate the semi-synthetic simulation for a single data shard. + + Returns a dict of DataFrames matching the format produced by + ``RealisticCausalSimulator.simulate_dataset``. + """ + pids, is_exposed, index_dates = self._extract_treatment_and_index_dates( + shard_df + ) + if len(pids) == 0: + return {} + + history_df = self._filter_to_pre_index(shard_df, index_dates) + history_df, pids, is_exposed, index_dates = self._apply_min_num_codes( + history_df, pids, is_exposed, index_dates + ) + if len(pids) == 0: + return {} + + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, self.config.features + ) + + ite_records, cf_records, all_factual_events, all_probas = ( + self._simulate_outcomes(features_df, pids, is_exposed, index_dates) + ) + + # Create exposure events for exposed patients + if np.any(is_exposed): + exposure_events = self._create_exposure_events( + pids[is_exposed], index_dates + ) + all_factual_events.append(exposure_events) + + return self._package_results( + pids, + ite_records, + cf_records, + all_factual_events, + all_probas, + is_exposed, + ) + + # ------------------------------------------------------------------ + # Data extraction + # ------------------------------------------------------------------ + + def _extract_treatment_and_index_dates( + self, shard_df: pd.DataFrame + ) -> Tuple[np.ndarray, np.ndarray, pd.Series]: + """Identify exposed/control patients and their index dates.""" + # Drop patients without an assigned index date + valid = shard_df.dropna(subset=[ASSIGNED_INDEX_DATE_COL]) + if valid.empty: + return ( + np.array([]), + np.array([], dtype=bool), + pd.Series(dtype="datetime64[ns]"), + ) + + # Per-patient index dates + index_dates = valid.groupby(PID_COL)[ASSIGNED_INDEX_DATE_COL].first() + + # Patients with at least one exposure event + exposed_pids = set( + valid.loc[valid[CONCEPT_COL] == self.config.exposure_code, PID_COL].unique() + ) + + pids = index_dates.index.values + is_exposed = np.array([pid in exposed_pids for pid in pids]) + + logger.info( + f"Extracted {len(pids)} patients ({np.sum(is_exposed)} exposed, " + f"{np.sum(~is_exposed)} control)" + ) + return pids, is_exposed, index_dates + + def _filter_to_pre_index( + self, df: pd.DataFrame, index_dates: pd.Series + ) -> pd.DataFrame: + """Keep only events before each patient's index date, excluding special codes.""" + df = df[df[PID_COL].isin(index_dates.index)].copy() + patient_index = index_dates.reindex(df[PID_COL]).values + df = df[df[TIMESTAMP_COL] <= patient_index] + + # Exclude special codes (keep BIRTH_CODE — needed for age computation) + excluded_exact = {self.config.exposure_code, DEATH_CODE, "GENDER"} + mask_exact = ~df[CONCEPT_COL].isin(excluded_exact) + mask_prefix = ~df[CONCEPT_COL].str.startswith(EXCLUDED_PREFIXES) + df = df[mask_exact & mask_prefix].copy() + + logger.info( + f"Pre-index history: {df[PID_COL].nunique()} patients, {len(df)} events" + ) + return df + + def _apply_min_num_codes( + self, + history_df: pd.DataFrame, + pids: np.ndarray, + is_exposed: np.ndarray, + index_dates: pd.Series, + ) -> Tuple[pd.DataFrame, np.ndarray, np.ndarray, pd.Series]: + """Drop patients with fewer than min_num_codes unique codes.""" + if self.config.min_num_codes <= 1: + return history_df, pids, is_exposed, index_dates + + code_counts = history_df.groupby(PID_COL)[CONCEPT_COL].nunique() + keep_pids = set(code_counts[code_counts >= self.config.min_num_codes].index) + before = len(pids) + mask = np.array([pid in keep_pids for pid in pids]) + pids = pids[mask] + is_exposed = is_exposed[mask] + index_dates = index_dates.loc[pids] + history_df = history_df[history_df[PID_COL].isin(keep_pids)].copy() + + logger.info( + f"After min_num_codes filter (>={self.config.min_num_codes}): " + f"{len(pids)} patients (dropped {before - len(pids)})" + ) + return history_df, pids, is_exposed, index_dates + + # ------------------------------------------------------------------ + # Outcome simulation + # ------------------------------------------------------------------ + + def _simulate_outcomes( + self, + features_df: pd.DataFrame, + pids: np.ndarray, + is_exposed: np.ndarray, + index_dates: pd.Series, + ) -> Tuple[Dict, Dict, list, Dict]: + n_patients = len(pids) + ite_records = {PID_COL: pids} + cf_records = {PID_COL: pids, EXPOSURE_COL: is_exposed.astype(int)} + all_factual_events = [] + all_probas = {} + + for outcome_name, outcome_cfg in self.config.outcomes.items(): + eta_0 = self._compute_eta_0(features_df, outcome_cfg.outcome_model) + tau = self._compute_tau(features_df, outcome_cfg.treatment_effect) + noise = self.rng.normal( + 0, outcome_cfg.outcome_model.noise_scale, n_patients + ) + + p0 = expit(eta_0 + noise) + p1 = expit(eta_0 + tau + noise) + + y1 = self.rng.binomial(1, p1) + y0 = self.rng.binomial(1, p0) + y_obs = np.where(is_exposed, y1, y0) + + all_probas[outcome_name] = {"P1": p1, "P0": p0} + ite_records[f"ite_{outcome_name}"] = p1 - p0 + + cf_records[f"{OUTCOME_COL}_{outcome_name}"] = y_obs + cf_records[f"{SIMULATED_OUTCOME_EXPOSED}_{outcome_name}"] = y1 + cf_records[f"{SIMULATED_OUTCOME_CONTROL}_{outcome_name}"] = y0 + cf_records[f"{SIMULATED_PROBAS_EXPOSED}_{outcome_name}"] = p1 + cf_records[f"{SIMULATED_PROBAS_CONTROL}_{outcome_name}"] = p0 + + # Create outcome events for patients with factual outcome == 1 + run_in_days = outcome_cfg.outcome_model.run_in_days + patients_with_outcome = pids[y_obs == 1] + if len(patients_with_outcome) > 0: + outcome_timestamps = index_dates.loc[ + patients_with_outcome + ].values + pd.Timedelta(days=run_in_days) + events = pd.DataFrame( + { + PID_COL: patients_with_outcome, + TIMESTAMP_COL: outcome_timestamps, + CONCEPT_COL: outcome_name, + } + ) + all_factual_events.append(events) + + return ite_records, cf_records, all_factual_events, all_probas + + def _compute_eta_0( + self, + features_df: pd.DataFrame, + outcome_model: OutcomeModelConfig, + ) -> np.ndarray: + """Compute the baseline log-odds: beta_0 + baseline terms + longitudinal terms + interactions.""" + n = len(features_df) + eta = np.full(n, outcome_model.beta_0) + + for name, coeff in outcome_model.baseline_coefficients.items(): + if name in features_df.columns: + eta += coeff * features_df[name].values + + for name, coeff in outcome_model.longitudinal_coefficients.items(): + if name in features_df.columns: + eta += coeff * features_df[name].values + + for interaction in outcome_model.interactions: + feat_a = interaction.get("features", [None, None])[0] + feat_b = interaction.get("features", [None, None])[1] + coeff = interaction.get("coefficient", 0.0) + if feat_a in features_df.columns and feat_b in features_df.columns: + eta += coeff * features_df[feat_a].values * features_df[feat_b].values + + return eta + + def _compute_tau( + self, + features_df: pd.DataFrame, + treatment_effect: TreatmentEffectConfig, + ) -> np.ndarray: + """Compute individual treatment effects on the logit scale.""" + n = len(features_df) + if treatment_effect.mode == "constant": + return np.full(n, treatment_effect.delta) + + # heterogeneous mode + tau = np.full(n, treatment_effect.delta_0) + for name, coeff in treatment_effect.heterogeneous_coefficients.items(): + if name in features_df.columns: + tau += coeff * features_df[name].values + return tau + + # ------------------------------------------------------------------ + # Result packaging (mirrors RealisticCausalSimulator._package_results) + # ------------------------------------------------------------------ + + def _package_results( + self, + pids, + ite_records, + cf_records, + all_factual_events, + all_probas, + is_exposed, + ) -> Dict[str, pd.DataFrame]: + output_dir = self.config.paths.outcomes + + logger.info("Calculating and saving simulation statistics...") + self._calculate_and_save_simulation_stats( + pids, is_exposed, cf_records, output_dir + ) + + logger.info("Calculating theoretical maximum ROC AUC...") + theoretical_aucs = self._calculate_theoretical_roc_auc( + cf_records, is_exposed, output_dir + ) + logger.info(f"Theoretical maximum ROC AUC: {theoretical_aucs}") + + # Plots + figs_dir = join(output_dir, "figs") + os.makedirs(figs_dir, exist_ok=True) + logger.info("Plotting ground truth probability distributions...") + plot_probability_distributions(all_probas, figs_dir) + + ite_df = pd.DataFrame(ite_records) + cf_df = pd.DataFrame(cf_records) + + # Build true_effects_config for the comparison plot + true_effects_config = {} + for outcome_name, outcome_cfg in self.config.outcomes.items(): + te = outcome_cfg.treatment_effect + om = outcome_cfg.outcome_model + true_effects_config[outcome_name] = { + "exposure_effect": te.delta if te.mode == "constant" else te.delta_0, + "p_base": expit(om.beta_0), + } + + logger.info("Plotting true effects vs observed risk differences...") + plot_true_effects_vs_risk_differences( + ite_df=ite_df, + cf_df=cf_df, + true_effects_config=true_effects_config, + output_dir=figs_dir, + ) + + # Build output DataFrames + output_dfs = {} + if all_factual_events: + events_df = pd.concat(all_factual_events, ignore_index=True) + events_df[ABSPOS_COL] = get_hours_since_epoch(events_df[TIMESTAMP_COL]) + for code, group in events_df.groupby(CONCEPT_COL): + output_dfs[str(code)] = group[ + [PID_COL, TIMESTAMP_COL, ABSPOS_COL] + ].copy() + + output_dfs["ite"] = ite_df + output_dfs[COUNTERFACTUALS_FILE.split(".")[0]] = cf_df + if EXPOSURE_COL in output_dfs: + output_dfs[INDEX_DATE_MATCHING_FILE.split(".")[0]] = ( + self._create_index_date_matching_df(output_dfs[EXPOSURE_COL], pids) + ) + + return output_dfs + + # ------------------------------------------------------------------ + # Helpers (same patterns as RealisticCausalSimulator) + # ------------------------------------------------------------------ + + def _create_exposure_events( + self, exposed_pids: np.ndarray, index_dates: pd.Series + ) -> pd.DataFrame: + timestamps = index_dates.loc[exposed_pids].values + df = pd.DataFrame( + { + PID_COL: exposed_pids, + TIMESTAMP_COL: timestamps, + CONCEPT_COL: EXPOSURE_COL, + } + ) + return df + + def _create_index_date_matching_df( + self, exposure_df: pd.DataFrame, all_pids: np.ndarray + ) -> pd.DataFrame: + exposed_pids = exposure_df[PID_COL].unique() + control_pids = safe_control_pids(all_pids, exposed_pids) + if len(exposed_pids) == 0 or len(control_pids) == 0: + return pd.DataFrame( + columns=[CONTROL_PID_COL, EXPOSED_PID_COL, TIMESTAMP_COL, ABSPOS_COL] + ) + + matched_exposed_pids = self.rng.choice( + exposed_pids, size=len(control_pids), replace=True + ) + match_df = pd.DataFrame( + {CONTROL_PID_COL: control_pids, EXPOSED_PID_COL: matched_exposed_pids} + ) + exposure_info = ( + exposure_df[[PID_COL, TIMESTAMP_COL, ABSPOS_COL]] + .drop_duplicates(subset=[PID_COL]) + .set_index(PID_COL) + ) + match_df = match_df.merge( + exposure_info, left_on=EXPOSED_PID_COL, right_index=True + ) + return match_df + + def _calculate_and_save_simulation_stats( + self, + pids: np.ndarray, + is_exposed: np.ndarray, + cf_records: Dict[str, np.ndarray], + output_dir: str, + ): + total_patients = len(pids) + num_exposed = int(np.sum(is_exposed)) + num_control = total_patients - num_exposed + + outcome_stats = {} + for outcome_name in self.config.outcomes: + outcome_col = f"{OUTCOME_COL}_{outcome_name}" + if outcome_col in cf_records: + num_with = int(np.sum(cf_records[outcome_col])) + outcome_stats[outcome_name] = { + "total_with_outcome": num_with, + "percentage_with_outcome": num_with / total_patients * 100, + } + + stats_rows = [ + ["Statistic", "Value"], + ["Total Patients", total_patients], + ["Number Exposed", num_exposed], + ["Number Control", num_control], + ["Exposure Rate (%)", f"{num_exposed / total_patients * 100:.2f}"], + ] + for outcome_name, data in outcome_stats.items(): + stats_rows.append( + [f"{outcome_name} - Total with Outcome", data["total_with_outcome"]] + ) + stats_rows.append( + [ + f"{outcome_name} - Percentage with Outcome (%)", + f"{data['percentage_with_outcome']:.2f}", + ] + ) + + stats_df = pd.DataFrame(stats_rows[1:], columns=stats_rows[0]) + os.makedirs(output_dir, exist_ok=True) + stats_path = join(output_dir, "simulation_stats.csv") + stats_df.to_csv(stats_path, index=False) + + logger.info(f"Simulation statistics saved to {stats_path}") + logger.info( + f"Total patients: {total_patients}, Exposed: {num_exposed}, Control: {num_control}" + ) + for outcome_name, data in outcome_stats.items(): + logger.info( + f"{outcome_name}: {data['total_with_outcome']} patients " + f"({data['percentage_with_outcome']:.2f}%)" + ) + + def _calculate_theoretical_roc_auc( + self, + cf_records: Dict[str, np.ndarray], + is_exposed: np.ndarray, + output_dir: str, + ) -> Dict[str, float]: + theoretical_aucs = {} + results_data = [] + + for outcome_name in self.config.outcomes: + outcome_col = f"{OUTCOME_COL}_{outcome_name}" + p_exposed_col = f"{SIMULATED_PROBAS_EXPOSED}_{outcome_name}" + p_control_col = f"{SIMULATED_PROBAS_CONTROL}_{outcome_name}" + + if outcome_col not in cf_records: + continue + if p_exposed_col not in cf_records or p_control_col not in cf_records: + continue + + y_true = cf_records[outcome_col] + p_treated = cf_records[p_exposed_col] + p_control = cf_records[p_control_col] + y_prob_factual = np.where(is_exposed, p_treated, p_control) + + if len(np.unique(y_true)) > 1: + auc_factual = roc_auc_score(y_true, y_prob_factual) + auc_treated = roc_auc_score(y_true, p_treated) + auc_control = roc_auc_score(y_true, p_control) + theoretical_aucs[outcome_name] = auc_factual + + results_data.append( + { + "outcome": outcome_name, + "auc_factual_dgp": auc_factual, + "auc_if_all_treated": auc_treated, + "auc_if_all_control": auc_control, + "n_positive": int(np.sum(y_true)), + "n_total": len(y_true), + "prevalence": np.mean(y_true), + } + ) + logger.info( + f"{outcome_name}: Theoretical max ROC AUC = {auc_factual:.4f}" + ) + else: + logger.warning( + f"Cannot calculate ROC AUC for {outcome_name}: only one class present." + ) + theoretical_aucs[outcome_name] = np.nan + + if results_data: + results_df = pd.DataFrame(results_data) + os.makedirs(output_dir, exist_ok=True) + results_path = join(output_dir, "theoretical_max_roc_auc.csv") + results_df.to_csv(results_path, index=False) + logger.info(f"Theoretical ROC AUC results saved to {results_path}") + + return theoretical_aucs diff --git a/tests/test_modules/test_simulation/__init__.py b/tests/test_modules/test_simulation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_modules/test_simulation/test_oracle_features.py b/tests/test_modules/test_simulation/test_oracle_features.py new file mode 100644 index 00000000..7f37c34f --- /dev/null +++ b/tests/test_modules/test_simulation/test_oracle_features.py @@ -0,0 +1,250 @@ +import unittest + +import numpy as np +import pandas as pd + +from corebehrt.constants.data import CONCEPT_COL, PID_COL, TIMESTAMP_COL +from corebehrt.modules.simulation.config_semisynthetic import FeatureConfig +from corebehrt.modules.simulation.oracle_features import ( + BASELINE_FEATURES, + LONGITUDINAL_FEATURES, + extract_oracle_features, +) + + +def _make_test_df(records): + """Build a MEDS DataFrame from (subject_id, time_str, code) tuples.""" + rows = [] + for pid, time_str, code in records: + rows.append( + {PID_COL: pid, TIMESTAMP_COL: pd.Timestamp(time_str), CONCEPT_COL: code} + ) + return pd.DataFrame(rows) + + +def _default_feature_config(**overrides): + kwargs = dict(standardize=False) + kwargs.update(overrides) + return FeatureConfig(**kwargs) + + +class TestExtractOracleFeaturesBasic(unittest.TestCase): + def test_extract_oracle_features_basic(self): + records = [ + (1, "1990-01-01", "DOB"), + (1, "2020-06-01", "D/11111"), + (1, "2020-09-01", "M/22222"), + (2, "1985-03-15", "DOB"), + (2, "2020-05-01", "D/33333"), + (3, "1975-07-20", "DOB"), + (3, "2020-08-01", "P/44444"), + (4, "2000-01-01", "DOB"), + (4, "2020-07-01", "D/55555"), + (4, "2020-08-15", "M/66666"), + (5, "1995-05-05", "DOB"), + (5, "2020-04-01", "D/77777"), + ] + history_df = _make_test_df(records) + pids = np.array([1, 2, 3, 4, 5]) + index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) + config = _default_feature_config() + features_df, baseline_names, longitudinal_names = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.shape[0], 5) + self.assertEqual( + features_df.shape[1], len(BASELINE_FEATURES) + len(LONGITUDINAL_FEATURES) + ) + self.assertEqual(baseline_names, BASELINE_FEATURES) + self.assertEqual(longitudinal_names, LONGITUDINAL_FEATURES) + + +class TestRecentEventCount(unittest.TestCase): + def test_recent_event_count(self): + # 3 events in last 90 days, 2 outside window + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-01-01", "DOB"), # outside 90-day window + (1, "2020-05-01", "D/11111"), # outside 90-day window + (1, "2020-10-15", "D/22222"), # within 90 days + (1, "2020-11-01", "M/33333"), # within 90 days + (1, "2020-12-01", "P/44444"), # within 90 days + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config(recent_window_days=90) + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "recent_event_count"], 3) + + +class TestDiseaseBurden(unittest.TestCase): + def test_disease_burden(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-03-01", "D/11111"), + (1, "2020-04-01", "D/22222"), + (1, "2020-05-01", "D/33333"), + (1, "2020-06-01", "D/44444"), + (1, "2020-07-01", "D/55555"), + (1, "2020-08-01", "M/66666"), # medication, not a diagnosis + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config() + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "disease_burden"], 5) + + +class TestMedicationCount(unittest.TestCase): + def test_medication_count(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-05-01", "M/11111"), + (1, "2020-06-01", "M/22222"), + (1, "2020-07-01", "M/33333"), + (1, "2020-08-01", "D/44444"), # diagnosis, not medication + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config() + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "medication_count"], 3) + + +class TestAgeComputation(unittest.TestCase): + def test_age_computation(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "1990-01-01", "DOB"), + (1, "2020-06-01", "D/11111"), + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config() + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + expected_age = (index - pd.Timestamp("1990-01-01")).days / 365.25 + self.assertAlmostEqual(features_df.loc[1, "age"], expected_age, places=2) + + +class TestEventRecency(unittest.TestCase): + def test_event_recency(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-06-01", "D/11111"), + (1, "2020-12-22", "M/22222"), # 10 days before index + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config() + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "event_recency"], 10) + + +class TestRecentBurstRatio(unittest.TestCase): + def test_recent_burst_ratio(self): + index = pd.Timestamp("2021-01-01") + # 5 events in burst window (30 days), 20 total in lookback (365 days) + records = [] + # 15 events outside burst window but inside lookback + for i in range(15): + records.append((1, f"2020-{3 + i % 9 + 1:02d}-01", f"D/{10000 + i}")) + # 5 events inside burst window (last 30 days of 2020) + for i in range(5): + records.append((1, f"2020-12-{5 + i * 5:02d}", f"M/{20000 + i}")) + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config(burst_window_days=30, lookback_days=365) + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + # burst=5, lookback=20, ratio = 5 / (20 + 1) + expected = 5.0 / 21.0 + self.assertAlmostEqual( + features_df.loc[1, "recent_burst_ratio"], expected, places=5 + ) + + +class TestNoMedicationCodes(unittest.TestCase): + def test_no_medication_codes(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-06-01", "D/11111"), + (2, "2020-07-01", "D/22222"), + ] + history_df = _make_test_df(records) + pids = np.array([1, 2]) + index_dates = pd.Series({1: index, 2: index}) + config = _default_feature_config() + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "medication_count"], 0) + self.assertEqual(features_df.loc[2, "medication_count"], 0) + + +class TestStandardization(unittest.TestCase): + def test_standardization(self): + records = [ + (1, "1990-01-01", "DOB"), + (1, "2020-06-01", "D/11111"), + (1, "2020-09-01", "M/22222"), + (2, "1985-03-15", "DOB"), + (2, "2020-05-01", "D/33333"), + (2, "2020-10-01", "M/44444"), + (3, "1975-07-20", "DOB"), + (3, "2020-08-01", "D/55555"), + (3, "2020-11-01", "M/66666"), + ] + history_df = _make_test_df(records) + pids = np.array([1, 2, 3]) + index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) + config = FeatureConfig(standardize=True) + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + for col in features_df.columns: + col_std = features_df[col].std() + # columns with zero variance stay at 0 after standardization + if col_std > 1e-10: + self.assertAlmostEqual(features_df[col].mean(), 0.0, places=10) + + +class TestSequenceMotifCount(unittest.TestCase): + def test_sequence_motif_count(self): + index = pd.Timestamp("2021-01-01") + records = [ + (1, "2020-12-01", "D/11111"), # diagnosis at day -31 + ( + 1, + "2020-12-20", + "M/22222", + ), # medication 19 days later -> within 30-day window + ] + history_df = _make_test_df(records) + pids = np.array([1]) + index_dates = pd.Series({1: index}) + config = _default_feature_config(motif_window_days=30) + features_df, _, _ = extract_oracle_features( + history_df, pids, index_dates, config + ) + self.assertEqual(features_df.loc[1, "sequence_motif_count"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modules/test_simulation/test_semisynthetic.py b/tests/test_modules/test_simulation/test_semisynthetic.py new file mode 100644 index 00000000..639c03f4 --- /dev/null +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -0,0 +1,240 @@ +"""Tests for the semi-synthetic causal simulator.""" + +import unittest + +import numpy as np +import pandas as pd + +from corebehrt.constants.data import CONCEPT_COL, PID_COL, TIMESTAMP_COL +from corebehrt.modules.simulation.config_semisynthetic import ( + FeatureConfig, + OutcomeModelConfig, + PathsConfig, + SemiSyntheticOutcomeConfig, + SemiSyntheticSimulationConfig, + TreatmentEffectConfig, +) +from corebehrt.modules.simulation.semisynthetic_simulator import ( + ASSIGNED_INDEX_DATE_COL, + SemiSyntheticCausalSimulator, +) + + +def _make_test_shard(n_patients=20, n_exposed=8, seed=42): + """Create a synthetic MEDS shard with known structure.""" + rng = np.random.RandomState(seed) + base_date = pd.Timestamp("2018-01-01") + index_date = pd.Timestamp("2020-01-01") + + rows = [] + for i in range(n_patients): + pid = i + 1 + # DOB event + dob = base_date - pd.Timedelta(days=rng.randint(10000, 25000)) + rows.append( + { + PID_COL: pid, + TIMESTAMP_COL: dob, + CONCEPT_COL: "DOB", + ASSIGNED_INDEX_DATE_COL: index_date, + } + ) + + # Generate several diagnosis codes (for min_num_codes to pass) + n_events = rng.randint(6, 15) + for _ in range(n_events): + days_before = rng.randint(1, 700) + code = f"D/{rng.choice(['A01', 'B02', 'C03', 'D04', 'E05', 'F06'])}" + rows.append( + { + PID_COL: pid, + TIMESTAMP_COL: index_date - pd.Timedelta(days=days_before), + CONCEPT_COL: code, + ASSIGNED_INDEX_DATE_COL: index_date, + } + ) + + # Add some medication codes + for _ in range(rng.randint(2, 5)): + days_before = rng.randint(1, 400) + code = f"M/{rng.choice(['X01', 'X02', 'X03'])}" + rows.append( + { + PID_COL: pid, + TIMESTAMP_COL: index_date - pd.Timedelta(days=days_before), + CONCEPT_COL: code, + ASSIGNED_INDEX_DATE_COL: index_date, + } + ) + + # Exposure event for first n_exposed patients + if i < n_exposed: + rows.append( + { + PID_COL: pid, + TIMESTAMP_COL: index_date, + CONCEPT_COL: "EXPOSURE", + ASSIGNED_INDEX_DATE_COL: index_date, + } + ) + + df = pd.DataFrame(rows) + df[TIMESTAMP_COL] = pd.to_datetime(df[TIMESTAMP_COL]) + df[ASSIGNED_INDEX_DATE_COL] = pd.to_datetime(df[ASSIGNED_INDEX_DATE_COL]) + return df + + +def _make_config( + tmpdir, + delta=1.0, + mode="constant", + seed=42, + min_num_codes=3, + noise_scale=0.0, +): + """Create a minimal SemiSyntheticSimulationConfig for testing.""" + paths = PathsConfig(data=".", splits=["test"], outcomes=tmpdir) + outcome = SemiSyntheticOutcomeConfig( + outcome_model=OutcomeModelConfig( + beta_0=-1.0, + baseline_coefficients={"disease_burden": 0.3}, + longitudinal_coefficients={}, + noise_scale=noise_scale, + ), + treatment_effect=TreatmentEffectConfig(mode=mode, delta=delta), + ) + return SemiSyntheticSimulationConfig( + paths=paths, + features=FeatureConfig(standardize=True), + outcomes={"OUTCOME_test": outcome}, + seed=seed, + min_num_codes=min_num_codes, + ) + + +class TestSemiSyntheticSimulatorOutputFormat(unittest.TestCase): + def setUp(self): + import tempfile + + self.tmpdir = tempfile.mkdtemp() + self.config = _make_config(self.tmpdir) + self.simulator = SemiSyntheticCausalSimulator(self.config) + self.shard = _make_test_shard() + self.results = self.simulator.simulate_dataset(self.shard) + + def test_simulate_dataset_output_format(self): + self.assertIn("counterfactuals", self.results) + self.assertIn("ite", self.results) + + cf_df = self.results["counterfactuals"] + self.assertIn(PID_COL, cf_df.columns) + self.assertIn("exposure", cf_df.columns) + self.assertIn("outcome_OUTCOME_test", cf_df.columns) + self.assertIn("Y1_OUTCOME_test", cf_df.columns) + self.assertIn("Y0_OUTCOME_test", cf_df.columns) + self.assertIn("P1_OUTCOME_test", cf_df.columns) + self.assertIn("P0_OUTCOME_test", cf_df.columns) + + ite_df = self.results["ite"] + self.assertIn(PID_COL, ite_df.columns) + self.assertIn("ite_OUTCOME_test", ite_df.columns) + + def test_probabilities_in_range(self): + cf_df = self.results["counterfactuals"] + p0 = cf_df["P0_OUTCOME_test"].values + p1 = cf_df["P1_OUTCOME_test"].values + self.assertTrue(np.all(p0 >= 0) and np.all(p0 <= 1)) + self.assertTrue(np.all(p1 >= 0) and np.all(p1 <= 1)) + + def test_outcome_consistency(self): + """Y_obs = A * Y1 + (1-A) * Y0 for all patients.""" + cf_df = self.results["counterfactuals"] + a = cf_df["exposure"].values + y1 = cf_df["Y1_OUTCOME_test"].values + y0 = cf_df["Y0_OUTCOME_test"].values + y_obs = cf_df["outcome_OUTCOME_test"].values + expected = a * y1 + (1 - a) * y0 + np.testing.assert_array_equal(y_obs, expected) + + +class TestSemiSyntheticSimulatorReproducibility(unittest.TestCase): + def test_reproducibility(self): + import tempfile + + tmpdir = tempfile.mkdtemp() + config = _make_config(tmpdir, seed=123) + shard = _make_test_shard() + + sim1 = SemiSyntheticCausalSimulator(config) + results1 = sim1.simulate_dataset(shard) + + # Re-create from scratch with same seed + config2 = _make_config(tmpdir, seed=123) + sim2 = SemiSyntheticCausalSimulator(config2) + results2 = sim2.simulate_dataset(shard) + + cf1 = results1["counterfactuals"] + cf2 = results2["counterfactuals"] + np.testing.assert_array_equal( + cf1["outcome_OUTCOME_test"].values, + cf2["outcome_OUTCOME_test"].values, + ) + np.testing.assert_array_almost_equal( + cf1["P0_OUTCOME_test"].values, + cf2["P0_OUTCOME_test"].values, + ) + + +class TestSemiSyntheticTreatmentEffect(unittest.TestCase): + def test_constant_treatment_effect(self): + """With constant positive delta, ITE should be positive for most patients.""" + import tempfile + + tmpdir = tempfile.mkdtemp() + config = _make_config(tmpdir, delta=2.0) + simulator = SemiSyntheticCausalSimulator(config) + results = simulator.simulate_dataset(_make_test_shard()) + + ite = results["ite"]["ite_OUTCOME_test"].values + # Most ITEs should be positive with delta=2.0 + self.assertGreater(np.mean(ite > 0), 0.5) + + def test_null_treatment_effect(self): + """With delta=0, mean ITE should be close to 0.""" + import tempfile + + tmpdir = tempfile.mkdtemp() + config = _make_config(tmpdir, delta=0.0) + simulator = SemiSyntheticCausalSimulator(config) + results = simulator.simulate_dataset(_make_test_shard()) + + ite = results["ite"]["ite_OUTCOME_test"].values + self.assertAlmostEqual(np.mean(ite), 0.0, places=5) + + +class TestSemiSyntheticExposureFromData(unittest.TestCase): + def test_exposure_from_data(self): + """Exposed patients should be exactly those with EXPOSURE events.""" + import tempfile + + tmpdir = tempfile.mkdtemp() + n_exposed = 8 + shard = _make_test_shard(n_patients=20, n_exposed=n_exposed) + config = _make_config(tmpdir) + simulator = SemiSyntheticCausalSimulator(config) + results = simulator.simulate_dataset(shard) + + cf_df = results["counterfactuals"] + exposed_pids_from_sim = set(cf_df.loc[cf_df["exposure"] == 1, PID_COL].values) + + # In the test shard, patients 1..n_exposed have EXPOSURE events + expected_exposed = set(range(1, n_exposed + 1)) + # The sim may have dropped some patients (min_num_codes), so check + # that the exposed set is a subset of the expected + remaining_pids = set(cf_df[PID_COL].values) + expected_in_remaining = expected_exposed & remaining_pids + self.assertEqual(exposed_pids_from_sim, expected_in_remaining) + + +if __name__ == "__main__": + unittest.main() From bcbbd41b1df2415e629f88fc013abd20a8d54598 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 11:25:33 +0200 Subject: [PATCH 02/20] merge baseline/longitudinal coefficients into single dict The updated simulation description drops the r_B/r_L feature grouping. The outcome model is now eta^(0) = beta_0 + f(r_i) with a single flat feature vector. Merges baseline_coefficients and longitudinal_coefficients into a single coefficients dict across config, simulator, YAML, and tests. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../causal/simulate_semisynthetic.yaml | 8 +-- .../main_causal/calibrate_semisynthetic.py | 2 +- .../simulation/config_semisynthetic.py | 5 +- .../modules/simulation/oracle_features.py | 24 ++------- .../simulation/semisynthetic_simulator.py | 10 ++-- .../test_simulation/test_oracle_features.py | 51 +++++-------------- .../test_simulation/test_semisynthetic.py | 3 +- 7 files changed, 25 insertions(+), 78 deletions(-) diff --git a/corebehrt/configs/causal/simulate_semisynthetic.yaml b/corebehrt/configs/causal/simulate_semisynthetic.yaml index 1c9123b9..1421ec0e 100644 --- a/corebehrt/configs/causal/simulate_semisynthetic.yaml +++ b/corebehrt/configs/causal/simulate_semisynthetic.yaml @@ -41,8 +41,7 @@ outcomes: outcome_model: run_in_days: 1 beta_0: -2.0 - # Baseline risk features (r_B) - baseline_coefficients: + coefficients: recent_event_count: 0.3 disease_burden: 0.2 medication_count: 0.15 @@ -50,8 +49,6 @@ outcomes: age: 0.4 chronic_disease_count: 0.25 code_diversity: 0.1 - # Longitudinal features (r_L) - longitudinal_coefficients: event_recency: -0.15 recent_burst_ratio: 0.2 sequence_motif_count: 0.1 @@ -67,10 +64,9 @@ outcomes: outcome_model: run_in_days: 1 beta_0: -2.0 - baseline_coefficients: + coefficients: disease_burden: 0.2 age: 0.4 - longitudinal_coefficients: {} noise_scale: 0.1 treatment_effect: mode: constant diff --git a/corebehrt/main_causal/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py index 6ba1e39a..566f645c 100644 --- a/corebehrt/main_causal/calibrate_semisynthetic.py +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -61,7 +61,7 @@ def main_calibrate(config_path): if len(pids) == 0: continue - features_df, _, _ = extract_oracle_features( + features_df, _ = extract_oracle_features( history_df, pids, index_dates, sim_config.features ) all_features.append(features_df.assign(is_exposed=is_exposed)) diff --git a/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py index bb86bacc..1ed8f1c8 100644 --- a/corebehrt/modules/simulation/config_semisynthetic.py +++ b/corebehrt/modules/simulation/config_semisynthetic.py @@ -35,12 +35,11 @@ class FeatureConfig: @dataclass class OutcomeModelConfig: - """Outcome model: eta^(0) = beta_0 + f_B(r_B) + f_L(r_L).""" + """Outcome model: eta^(0) = beta_0 + f(r_i).""" run_in_days: int = 1 beta_0: float = -2.0 - baseline_coefficients: Dict[str, float] = field(default_factory=dict) - longitudinal_coefficients: Dict[str, float] = field(default_factory=dict) + coefficients: Dict[str, float] = field(default_factory=dict) interactions: List[Dict] = field(default_factory=list) noise_scale: float = 0.0 diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index 233a6230..3b2daae7 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -11,28 +11,13 @@ logger = logging.getLogger("oracle_features") -BASELINE_FEATURES = [ - "recent_event_count", - "disease_burden", - "medication_count", - "utilization_intensity", - "age", - "chronic_disease_count", - "code_diversity", -] -LONGITUDINAL_FEATURES = [ - "event_recency", - "recent_burst_ratio", - "sequence_motif_count", -] - def extract_oracle_features( history_df: pd.DataFrame, pids: np.ndarray, index_dates: pd.Series, feature_config: FeatureConfig, -) -> Tuple[pd.DataFrame, List[str], List[str]]: +) -> Tuple[pd.DataFrame, List[str]]: """Extract oracle features from pre-index patient histories. Args: @@ -43,8 +28,7 @@ def extract_oracle_features( Returns: features_df: DataFrame with PID_COL as index, one column per feature - baseline_feature_names: list of r_B feature names present in the output - longitudinal_feature_names: list of r_L feature names present in the output + feature_names: list of feature names present in the output """ prefixes = feature_config.code_prefixes @@ -92,9 +76,7 @@ def extract_oracle_features( if feature_config.standardize: features_df = _standardize(features_df) - baseline_names = [n for n in BASELINE_FEATURES if n in features_df.columns] - longitudinal_names = [n for n in LONGITUDINAL_FEATURES if n in features_df.columns] - return features_df, baseline_names, longitudinal_names + return features_df, list(features_df.columns) # --------------------------------------------------------------------------- diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 048d1155..4b0a1d5d 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -81,7 +81,7 @@ def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: if len(pids) == 0: return {} - features_df, _, _ = extract_oracle_features( + features_df, _ = extract_oracle_features( history_df, pids, index_dates, self.config.features ) @@ -247,15 +247,11 @@ def _compute_eta_0( features_df: pd.DataFrame, outcome_model: OutcomeModelConfig, ) -> np.ndarray: - """Compute the baseline log-odds: beta_0 + baseline terms + longitudinal terms + interactions.""" + """Compute the baseline log-odds: beta_0 + f(r_i) + interactions.""" n = len(features_df) eta = np.full(n, outcome_model.beta_0) - for name, coeff in outcome_model.baseline_coefficients.items(): - if name in features_df.columns: - eta += coeff * features_df[name].values - - for name, coeff in outcome_model.longitudinal_coefficients.items(): + for name, coeff in outcome_model.coefficients.items(): if name in features_df.columns: eta += coeff * features_df[name].values diff --git a/tests/test_modules/test_simulation/test_oracle_features.py b/tests/test_modules/test_simulation/test_oracle_features.py index 7f37c34f..a5fd66da 100644 --- a/tests/test_modules/test_simulation/test_oracle_features.py +++ b/tests/test_modules/test_simulation/test_oracle_features.py @@ -5,11 +5,7 @@ from corebehrt.constants.data import CONCEPT_COL, PID_COL, TIMESTAMP_COL from corebehrt.modules.simulation.config_semisynthetic import FeatureConfig -from corebehrt.modules.simulation.oracle_features import ( - BASELINE_FEATURES, - LONGITUDINAL_FEATURES, - extract_oracle_features, -) +from corebehrt.modules.simulation.oracle_features import extract_oracle_features def _make_test_df(records): @@ -48,15 +44,12 @@ def test_extract_oracle_features_basic(self): pids = np.array([1, 2, 3, 4, 5]) index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) config = _default_feature_config() - features_df, baseline_names, longitudinal_names = extract_oracle_features( + features_df, feature_names = extract_oracle_features( history_df, pids, index_dates, config ) self.assertEqual(features_df.shape[0], 5) - self.assertEqual( - features_df.shape[1], len(BASELINE_FEATURES) + len(LONGITUDINAL_FEATURES) - ) - self.assertEqual(baseline_names, BASELINE_FEATURES) - self.assertEqual(longitudinal_names, LONGITUDINAL_FEATURES) + self.assertEqual(features_df.shape[1], 10) + self.assertEqual(feature_names, list(features_df.columns)) class TestRecentEventCount(unittest.TestCase): @@ -74,9 +67,7 @@ def test_recent_event_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(recent_window_days=90) - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "recent_event_count"], 3) @@ -95,9 +86,7 @@ def test_disease_burden(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "disease_burden"], 5) @@ -114,9 +103,7 @@ def test_medication_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "medication_count"], 3) @@ -131,9 +118,7 @@ def test_age_computation(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) expected_age = (index - pd.Timestamp("1990-01-01")).days / 365.25 self.assertAlmostEqual(features_df.loc[1, "age"], expected_age, places=2) @@ -149,9 +134,7 @@ def test_event_recency(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "event_recency"], 10) @@ -170,9 +153,7 @@ def test_recent_burst_ratio(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(burst_window_days=30, lookback_days=365) - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) # burst=5, lookback=20, ratio = 5 / (20 + 1) expected = 5.0 / 21.0 self.assertAlmostEqual( @@ -191,9 +172,7 @@ def test_no_medication_codes(self): pids = np.array([1, 2]) index_dates = pd.Series({1: index, 2: index}) config = _default_feature_config() - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "medication_count"], 0) self.assertEqual(features_df.loc[2, "medication_count"], 0) @@ -215,9 +194,7 @@ def test_standardization(self): pids = np.array([1, 2, 3]) index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) config = FeatureConfig(standardize=True) - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) for col in features_df.columns: col_std = features_df[col].std() # columns with zero variance stay at 0 after standardization @@ -240,9 +217,7 @@ def test_sequence_motif_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(motif_window_days=30) - features_df, _, _ = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "sequence_motif_count"], 1) diff --git a/tests/test_modules/test_simulation/test_semisynthetic.py b/tests/test_modules/test_simulation/test_semisynthetic.py index 639c03f4..988135bf 100644 --- a/tests/test_modules/test_simulation/test_semisynthetic.py +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -97,8 +97,7 @@ def _make_config( outcome = SemiSyntheticOutcomeConfig( outcome_model=OutcomeModelConfig( beta_0=-1.0, - baseline_coefficients={"disease_burden": 0.3}, - longitudinal_coefficients={}, + coefficients={"disease_burden": 0.3}, noise_scale=noise_scale, ), treatment_effect=TreatmentEffectConfig(mode=mode, delta=delta), From b42dca6d180ccd35c7e42501a411853d6e0cd5f4 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 11:40:38 +0200 Subject: [PATCH 03/20] fix edge cases, add azure components and docs Bug fixes from CodeRabbit review: - Handle all-missing-DOB with default age fallback (65 years) - Fix single-patient standardization: use ddof=0 to avoid NaN - Move stats/plots to finalize() method called after shard aggregation - Add clarifying comment re intentional noise omission in calibration New files: - Azure components for simulate and calibrate - experiments/semisynthetic_simulation/ docs and config generator Co-Authored-By: Claude Opus 4.6 (1M context) --- .../components/calibrate_semisynthetic.py | 20 ++++ .../components/simulate_semisynthetic.py | 20 ++++ .../main_causal/calibrate_semisynthetic.py | 2 + .../main_causal/simulate_semisynthetic.py | 7 ++ .../modules/simulation/oracle_features.py | 6 +- .../simulation/semisynthetic_simulator.py | 82 ++++++++-------- .../semisynthetic_simulation/README.md | 37 ++++++++ .../semisynthetic_simulation/docs/azure.md | 43 +++++++++ .../docs/calibration.md | 93 ++++++++++++++++++ .../semisynthetic_simulation/docs/features.md | 41 ++++++++ .../semisynthetic_simulation/docs/local.md | 77 +++++++++++++++ .../docs/multiple_runs.md | 86 +++++++++++++++++ .../generate_configs.py | 95 +++++++++++++++++++ 13 files changed, 569 insertions(+), 40 deletions(-) create mode 100644 corebehrt/azure/components/calibrate_semisynthetic.py create mode 100644 corebehrt/azure/components/simulate_semisynthetic.py create mode 100644 experiments/semisynthetic_simulation/README.md create mode 100644 experiments/semisynthetic_simulation/docs/azure.md create mode 100644 experiments/semisynthetic_simulation/docs/calibration.md create mode 100644 experiments/semisynthetic_simulation/docs/features.md create mode 100644 experiments/semisynthetic_simulation/docs/local.md create mode 100644 experiments/semisynthetic_simulation/docs/multiple_runs.md create mode 100644 experiments/semisynthetic_simulation/generate_configs.py diff --git a/corebehrt/azure/components/calibrate_semisynthetic.py b/corebehrt/azure/components/calibrate_semisynthetic.py new file mode 100644 index 00000000..ae127ec9 --- /dev/null +++ b/corebehrt/azure/components/calibrate_semisynthetic.py @@ -0,0 +1,20 @@ +from corebehrt.azure.util import job + +INPUTS = { + "data": {"type": "uri_folder"}, +} + +OUTPUTS = { + "outcomes": {"type": "uri_folder"}, +} + + +if __name__ == "__main__": + from corebehrt.main_causal import calibrate_semisynthetic + + job.run_main( + "calibrate_semisynthetic", + calibrate_semisynthetic.main_calibrate, + INPUTS, + OUTPUTS, + ) diff --git a/corebehrt/azure/components/simulate_semisynthetic.py b/corebehrt/azure/components/simulate_semisynthetic.py new file mode 100644 index 00000000..63295de5 --- /dev/null +++ b/corebehrt/azure/components/simulate_semisynthetic.py @@ -0,0 +1,20 @@ +from corebehrt.azure.util import job + +INPUTS = { + "data": {"type": "uri_folder"}, +} + +OUTPUTS = { + "outcomes": {"type": "uri_folder"}, +} + + +if __name__ == "__main__": + from corebehrt.main_causal import simulate_semisynthetic + + job.run_main( + "simulate_semisynthetic", + simulate_semisynthetic.main_simulate, + INPUTS, + OUTPUTS, + ) diff --git a/corebehrt/main_causal/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py index 566f645c..1adefa18 100644 --- a/corebehrt/main_causal/calibrate_semisynthetic.py +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -70,6 +70,8 @@ def main_calibrate(config_path): for outcome_name, outcome_cfg in sim_config.outcomes.items(): eta_0 = simulator._compute_eta_0(features_df, outcome_cfg.outcome_model) tau = simulator._compute_tau(features_df, outcome_cfg.treatment_effect) + # Noise omitted intentionally: calibration shows the deterministic + # risk surface, not the noisy realization used during sampling. p0 = expit(eta_0) p1 = expit(eta_0 + tau) diff --git a/corebehrt/main_causal/simulate_semisynthetic.py b/corebehrt/main_causal/simulate_semisynthetic.py index 46fcb8d7..2efeb644 100644 --- a/corebehrt/main_causal/simulate_semisynthetic.py +++ b/corebehrt/main_causal/simulate_semisynthetic.py @@ -38,6 +38,7 @@ def simulate(shard_loader: ShardLoader, simulator: CausalSimulator, outcomes_dir Iterates through each data shard, calls simulate_dataset, aggregates the results, and saves each outcome type to a separate CSV file. + Then computes stats and plots from the aggregated results. """ logger.info("--- Starting semi-synthetic simulation ---") simulated_outcomes = defaultdict(list) @@ -49,10 +50,16 @@ def simulate(shard_loader: ShardLoader, simulator: CausalSimulator, outcomes_dir logger.info("--- Simulation complete, saving results ---") + aggregated = {} for k, df_list in simulated_outcomes.items(): if df_list: df = pd.concat(df_list, ignore_index=True) df.to_csv(join(outcomes_dir, f"{k}.csv"), index=False) + aggregated[k] = df + + if "counterfactuals" in aggregated and "ite" in aggregated: + logger.info("--- Computing stats and plots from aggregated results ---") + simulator.finalize(aggregated["counterfactuals"], aggregated["ite"]) if __name__ == "__main__": diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index 3b2daae7..efe5a637 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -153,6 +153,9 @@ def _compute_age(history_df, pids, index_dates): idx_date = index_dates[pid] age_series[pid] = (idx_date - dob).days / 365.25 mean_age = age_series.mean() + if np.isnan(mean_age): + mean_age = 65.0 + logger.warning("No DOB events found, using default age %.0f", mean_age) age_series = age_series.fillna(mean_age) return age_series @@ -234,6 +237,5 @@ def _compute_sequence_motif_count( def _standardize(features_df): means = features_df.mean() - stds = features_df.std() - stds = stds.replace(0, 1) + stds = features_df.std(ddof=0).replace(0, 1).fillna(1) return (features_df - means) / stds diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 4b0a1d5d..c1972abb 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -85,8 +85,8 @@ def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: history_df, pids, index_dates, self.config.features ) - ite_records, cf_records, all_factual_events, all_probas = ( - self._simulate_outcomes(features_df, pids, is_exposed, index_dates) + ite_records, cf_records, all_factual_events = self._simulate_outcomes( + features_df, pids, is_exposed, index_dates ) # Create exposure events for exposed patients @@ -101,8 +101,6 @@ def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: ite_records, cf_records, all_factual_events, - all_probas, - is_exposed, ) # ------------------------------------------------------------------ @@ -194,12 +192,11 @@ def _simulate_outcomes( pids: np.ndarray, is_exposed: np.ndarray, index_dates: pd.Series, - ) -> Tuple[Dict, Dict, list, Dict]: + ) -> Tuple[Dict, Dict, list]: n_patients = len(pids) ite_records = {PID_COL: pids} cf_records = {PID_COL: pids, EXPOSURE_COL: is_exposed.astype(int)} all_factual_events = [] - all_probas = {} for outcome_name, outcome_cfg in self.config.outcomes.items(): eta_0 = self._compute_eta_0(features_df, outcome_cfg.outcome_model) @@ -215,7 +212,6 @@ def _simulate_outcomes( y0 = self.rng.binomial(1, p0) y_obs = np.where(is_exposed, y1, y0) - all_probas[outcome_name] = {"P1": p1, "P0": p0} ite_records[f"ite_{outcome_name}"] = p1 - p0 cf_records[f"{OUTCOME_COL}_{outcome_name}"] = y_obs @@ -240,7 +236,7 @@ def _simulate_outcomes( ) all_factual_events.append(events) - return ite_records, cf_records, all_factual_events, all_probas + return ite_records, cf_records, all_factual_events def _compute_eta_0( self, @@ -291,32 +287,61 @@ def _package_results( ite_records, cf_records, all_factual_events, - all_probas, - is_exposed, ) -> Dict[str, pd.DataFrame]: + ite_df = pd.DataFrame(ite_records) + cf_df = pd.DataFrame(cf_records) + + output_dfs = {} + if all_factual_events: + events_df = pd.concat(all_factual_events, ignore_index=True) + events_df[ABSPOS_COL] = get_hours_since_epoch(events_df[TIMESTAMP_COL]) + for code, group in events_df.groupby(CONCEPT_COL): + output_dfs[str(code)] = group[ + [PID_COL, TIMESTAMP_COL, ABSPOS_COL] + ].copy() + + output_dfs["ite"] = ite_df + output_dfs[COUNTERFACTUALS_FILE.split(".")[0]] = cf_df + if EXPOSURE_COL in output_dfs: + output_dfs[INDEX_DATE_MATCHING_FILE.split(".")[0]] = ( + self._create_index_date_matching_df(output_dfs[EXPOSURE_COL], pids) + ) + + return output_dfs + + def finalize(self, cf_df: pd.DataFrame, ite_df: pd.DataFrame): + """Compute stats and plots from aggregated results. Call after all shards.""" output_dir = self.config.paths.outcomes + is_exposed = cf_df[EXPOSURE_COL].values.astype(bool) + pids = cf_df[PID_COL].values + cf_records = {col: cf_df[col].values for col in cf_df.columns} + logger.info("Calculating and saving simulation statistics...") self._calculate_and_save_simulation_stats( pids, is_exposed, cf_records, output_dir ) logger.info("Calculating theoretical maximum ROC AUC...") - theoretical_aucs = self._calculate_theoretical_roc_auc( - cf_records, is_exposed, output_dir - ) - logger.info(f"Theoretical maximum ROC AUC: {theoretical_aucs}") + self._calculate_theoretical_roc_auc(cf_records, is_exposed, output_dir) - # Plots figs_dir = join(output_dir, "figs") os.makedirs(figs_dir, exist_ok=True) + + # Rebuild probability dicts for plotting + all_probas = {} + for outcome_name in self.config.outcomes: + p1_col = f"{SIMULATED_PROBAS_EXPOSED}_{outcome_name}" + p0_col = f"{SIMULATED_PROBAS_CONTROL}_{outcome_name}" + if p1_col in cf_df.columns and p0_col in cf_df.columns: + all_probas[outcome_name] = { + "P1": cf_df[p1_col].values, + "P0": cf_df[p0_col].values, + } + logger.info("Plotting ground truth probability distributions...") plot_probability_distributions(all_probas, figs_dir) - ite_df = pd.DataFrame(ite_records) - cf_df = pd.DataFrame(cf_records) - - # Build true_effects_config for the comparison plot true_effects_config = {} for outcome_name, outcome_cfg in self.config.outcomes.items(): te = outcome_cfg.treatment_effect @@ -334,25 +359,6 @@ def _package_results( output_dir=figs_dir, ) - # Build output DataFrames - output_dfs = {} - if all_factual_events: - events_df = pd.concat(all_factual_events, ignore_index=True) - events_df[ABSPOS_COL] = get_hours_since_epoch(events_df[TIMESTAMP_COL]) - for code, group in events_df.groupby(CONCEPT_COL): - output_dfs[str(code)] = group[ - [PID_COL, TIMESTAMP_COL, ABSPOS_COL] - ].copy() - - output_dfs["ite"] = ite_df - output_dfs[COUNTERFACTUALS_FILE.split(".")[0]] = cf_df - if EXPOSURE_COL in output_dfs: - output_dfs[INDEX_DATE_MATCHING_FILE.split(".")[0]] = ( - self._create_index_date_matching_df(output_dfs[EXPOSURE_COL], pids) - ) - - return output_dfs - # ------------------------------------------------------------------ # Helpers (same patterns as RealisticCausalSimulator) # ------------------------------------------------------------------ diff --git a/experiments/semisynthetic_simulation/README.md b/experiments/semisynthetic_simulation/README.md new file mode 100644 index 00000000..bd1a6589 --- /dev/null +++ b/experiments/semisynthetic_simulation/README.md @@ -0,0 +1,37 @@ +# Semi-Synthetic Simulation + +Semi-synthetic causal simulation where **treatment is kept from real data** and only the outcome is simulated from hand-crafted oracle features. + +## Quick Links + +| What | Where | +|------|-------| +| **How to run locally** | [docs/local.md](docs/local.md) | +| **How to run on Azure** | [docs/azure.md](docs/azure.md) | +| **How to run multiple replicates** | [docs/multiple_runs.md](docs/multiple_runs.md) | +| **Feature definitions** | [docs/features.md](docs/features.md) | +| **Calibrating parameters** | [docs/calibration.md](docs/calibration.md) | +| Config file | `corebehrt/configs/causal/simulate_semisynthetic.yaml` | +| Simulator code | `corebehrt/modules/simulation/semisynthetic_simulator.py` | +| Feature extraction | `corebehrt/modules/simulation/oracle_features.py` | +| Config dataclasses | `corebehrt/modules/simulation/config_semisynthetic.py` | +| Azure component | `corebehrt/azure/components/simulate_semisynthetic.py` | +| Tests | `tests/test_modules/test_simulation/` | + +## How It Works (30-Second Summary) + +1. Load real EHR data (MEDS format) with real treatment assignments +2. Extract oracle features from each patient's pre-index history (disease burden, age, recency, etc.) +3. Simulate outcome: `P(Y(0)=1) = sigmoid(beta_0 + f(features))`, `P(Y(1)=1) = sigmoid(beta_0 + f(features) + tau)` +4. Observed outcome: `Y = A * Y(1) + (1-A) * Y(0)` where A is the real treatment +5. True ATE is known by construction + +## Key Difference from Old Simulation + +| | Old (`realistic_simulator`) | New (`semisynthetic_simulator`) | +|---|---|---| +| Treatment | Simulated from latent factors | **Real** (from data) | +| Outcome model | Random latent factor weights | Hand-crafted interpretable features | +| Confounding | Synthetic (shared latent factors) | **Real** (from clinical practice) | +| Index dates | Single global date | Per-patient (from cohort) | +| Config complexity | Latent dims, sparsity, influence scales | Feature coefficients, beta_0, delta | diff --git a/experiments/semisynthetic_simulation/docs/azure.md b/experiments/semisynthetic_simulation/docs/azure.md new file mode 100644 index 00000000..eb6d1d94 --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/azure.md @@ -0,0 +1,43 @@ +# Running on Azure + +## Single Job + +```bash +python -m corebehrt.azure job simulate_semisynthetic CPU-20-LP \ + --config corebehrt/configs/causal/simulate_semisynthetic.yaml \ + -e semisynthetic_sim +``` + +The Azure component is at `corebehrt/azure/components/simulate_semisynthetic.py`. + +### Azure Config Paths + +Replace local paths with Azure datastore paths: + +```yaml +paths: + data: "researcher_data:path/to/MEDS/data" + splits: ["tuning"] + outcomes: "researcher_data:path/to/semisynthetic_outcomes" +``` + +## Calibration Job + +Run the calibration script to inspect probability distributions before committing to a full experiment: + +```bash +python -m corebehrt.azure job calibrate_semisynthetic CPU-20-LP \ + --config corebehrt/configs/causal/simulate_semisynthetic.yaml \ + -e semisynthetic_calibrate +``` + +## Pipeline Integration + +The semisynthetic simulator produces the same output format as `simulate_from_sequence`, so it can be swapped into the existing `FINETUNE_ESTIMATE_SIMULATED` pipeline by replacing the simulation component. + +To use it in a custom pipeline, the component signature is: + +```python +INPUTS = {"data": {"type": "uri_folder"}} +OUTPUTS = {"outcomes": {"type": "uri_folder"}} +``` diff --git a/experiments/semisynthetic_simulation/docs/calibration.md b/experiments/semisynthetic_simulation/docs/calibration.md new file mode 100644 index 00000000..566a962c --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/calibration.md @@ -0,0 +1,93 @@ +# Calibrating Simulation Parameters + +Before running a full experiment, use the calibration script to check that your coefficients and intercept produce realistic probability distributions. + +## Running the Calibration Script + +```bash +python -m corebehrt.main_causal.calibrate_semisynthetic \ + --config_path corebehrt/configs/causal/simulate_semisynthetic.yaml +``` + +On Azure: + +```bash +python -m corebehrt.azure job calibrate_semisynthetic CPU-20-LP \ + --config corebehrt/configs/causal/simulate_semisynthetic.yaml \ + -e calibrate_semisynthetic +``` + +## What It Reports + +### Feature diagnostics (printed table) + +Per feature: mean, std, min, p5, p25, p50, p75, p95, max, and **SMD** (standardized mean difference between treated and control groups). + +SMD tells you which features are associated with real treatment assignment. Features with |SMD| > 0.1 are meaningfully imbalanced and will create confounding in the simulation. + +### Probability diagnostics + +For each outcome: +- P(Y(0)) and P(Y(1)): mean, std, min, max, median +- Fraction of extreme probabilities (< 0.01 or > 0.80) +- Expected factual outcome prevalence + +### Causal effect diagnostics + +- True ATE = mean(P1 - P0) +- True ATT = mean(P1 - P0 | A=1) +- True ATC = mean(P1 - P0 | A=0) +- True RR = mean(P1) / mean(P0) + +### Plots (saved to `outcomes/figs/`) + +- P(Y(0)) and P(Y(1)) histograms +- ITE distribution per outcome +- SMD love plot (feature balance between treated/control) + +## Calibration Workflow + +### 1. Start with a target baseline prevalence + +Decide the untreated outcome rate: +- ~5% for rare outcomes +- ~10-20% for moderate +- ~30% for common + +### 2. Set beta_0 to approximate that prevalence + +`sigmoid(beta_0)` is roughly the baseline risk when all features are at their mean (which is 0 after standardization). So: +- beta_0 = -3.0 -> ~5% baseline +- beta_0 = -2.0 -> ~12% baseline +- beta_0 = -1.0 -> ~27% baseline +- beta_0 = 0.0 -> ~50% baseline + +### 3. Keep coefficients moderate + +With standardized features, a coefficient of: +- 0.1-0.2: small effect +- 0.3-0.5: moderate effect +- 0.8+: large effect (use sparingly) + +Multiple large coefficients will push logits to extremes and saturate probabilities near 0 or 1. + +### 4. Run calibration and check + +Good signs: +- Most P(Y(0)) values between 0.01 and 0.80 +- Mean prevalence close to your target +- ATE is detectable but not absurdly large +- Fraction of extreme probabilities < 10% + +Bad signs: +- Many probabilities near 0 or 1 (reduce coefficients) +- ATE too small to detect (increase delta) +- ATE too large (every method will succeed trivially) + +### 5. Check feature SMDs + +If the SMD love plot shows most features near zero, then the features don't create meaningful confounding and the simulation is too easy. Look for features with |SMD| > 0.1 — these are the ones that make the problem realistic. + +### 6. Iterate + +Adjust beta_0, coefficients, and delta, re-run calibration, until the distributions look sensible. diff --git a/experiments/semisynthetic_simulation/docs/features.md b/experiments/semisynthetic_simulation/docs/features.md new file mode 100644 index 00000000..9126ff20 --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/features.md @@ -0,0 +1,41 @@ +# Oracle Feature Definitions + +All features are extracted from the pre-index patient history by `corebehrt/modules/simulation/oracle_features.py`. + +Features are optionally z-scored (mean=0, std=1) when `features.standardize: true`, so each coefficient represents the effect of a 1-SD change. + +## Feature Table + +| Feature | Type | Description | Window | +|---------|------|-------------|--------| +| `recent_event_count` | count | Number of events (any code type) before index | `recent_window_days` (default: 90d) | +| `disease_burden` | unique count | Unique diagnosis codes before index | `lookback_days` (default: 365d) | +| `medication_count` | unique count | Unique medication codes before index (polypharmacy proxy) | `lookback_days` | +| `utilization_intensity` | count | Total events before index | `lookback_days` | +| `age` | continuous | Age in years at index date (from DOB events) | full history | +| `chronic_disease_count` | unique count | Distinct diagnosis code groups (first 5 chars) | full history | +| `code_diversity` | unique count | Total unique codes across full history | full history | +| `event_recency` | days | Days since most recent event before index | full history | +| `recent_burst_ratio` | ratio | events_in_burst_window / (events_in_lookback + 1) | `burst_window_days` / `lookback_days` | +| `sequence_motif_count` | count | (diagnosis -> medication) pairs within motif window | `motif_window_days` (default: 30d) | + +## Code Prefix Configuration + +Features that filter by medical concept type use configurable prefixes: + +```yaml +features: + code_prefixes: + diagnosis: "D/" # Used by: disease_burden, chronic_disease_count, sequence_motif_count + medication: "M/" # Used by: medication_count, sequence_motif_count + procedure: "P/" # Available but not currently used by any feature + admission: "ADM/" # Available but not currently used by any feature +``` + +## Notes + +- If no codes match a prefix (e.g., no medication codes in the data), the feature is filled with 0 and a warning is logged. +- `age` requires DOB events in the history. Missing DOB is filled with mean age. +- `event_recency` for patients with no events is filled with mean recency. +- `chronic_disease_count` uses the first 5 characters of diagnosis codes as group keys. The granularity depends on the coding system. +- `sequence_motif_count` counts (diagnosis, medication) pairs where the medication occurs within `motif_window_days` after the diagnosis. diff --git a/experiments/semisynthetic_simulation/docs/local.md b/experiments/semisynthetic_simulation/docs/local.md new file mode 100644 index 00000000..7cbb034d --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/local.md @@ -0,0 +1,77 @@ +# Running Locally + +## Single Run + +```bash +python -m corebehrt.main_causal.simulate_semisynthetic \ + --config_path corebehrt/configs/causal/simulate_semisynthetic.yaml +``` + +This reads the MEDS data, extracts real treatment assignments, computes oracle features, simulates outcomes, and saves results to the configured `paths.outcomes` directory. + +## Output Files + +``` +outputs/causal/semisynthetic_outcomes/ +├── counterfactuals.csv # subject_id, exposure, outcome_X, Y0_X, Y1_X, P0_X, P1_X +├── ite.csv # subject_id, ite_X (individual treatment effects) +├── OUTCOME.csv # Event records for patients with Y_obs=1 +├── index_date_matching.csv # Exposed/control matching +├── simulation_stats.csv # Patient counts, exposure rates, outcome prevalences +├── theoretical_max_roc_auc.csv # Best achievable AUC from the DGP +└── figs/ + ├── probability_distributions.png + └── true_effects_vs_risk_differences.png +``` + +## Config Structure + +```yaml +paths: + data: ./example_data/synthea_meds_causal # MEDS shard directory + splits: ["tuning"] # Which splits to use + outcomes: ./outputs/causal/semisynthetic_outcomes + +seed: 42 +min_num_codes: 3 +exposure_code: "EXPOSURE" # Code string identifying treatment events + +features: + code_prefixes: + diagnosis: "D/" + medication: "M/" + procedure: "P/" + admission: "ADM/" + lookback_days: 365 # General lookback window + recent_window_days: 90 # "Recent" events window + burst_window_days: 30 # Short-term burst window + motif_window_days: 30 # Max gap for sequential motifs + standardize: true # Z-score features before applying coefficients + +outcomes: + OUTCOME: # One block per simulated outcome + outcome_model: + run_in_days: 1 + beta_0: -2.0 # Intercept (controls baseline prevalence) + coefficients: # Feature name -> coefficient (logit scale) + disease_burden: 0.2 + age: 0.4 + event_recency: -0.15 + interactions: + - features: [disease_burden, age] + coefficient: 0.1 + noise_scale: 0.1 # Logit-scale noise std dev + treatment_effect: + mode: constant # "constant" or "heterogeneous" + delta: 1.0 # Constant treatment effect (logit scale) +``` + +## MEDS Data Requirements + +The simulator expects MEDS-format parquet files with columns: +- `subject_id`: Patient identifier +- `time`: Event timestamp (datetime) +- `code`: Medical code string (e.g., `D/12345`, `M/67890`, `DOB`, `EXPOSURE`) +- `assigned_index_date`: Per-patient index date (datetime, can be NaT for patients to exclude) + +Patients with `EXPOSURE` events are treated (A=1). Others are control (A=0). diff --git a/experiments/semisynthetic_simulation/docs/multiple_runs.md b/experiments/semisynthetic_simulation/docs/multiple_runs.md new file mode 100644 index 00000000..132f6efe --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/multiple_runs.md @@ -0,0 +1,86 @@ +# Running Multiple Replicates + +The simulation is deterministic given a seed. To generate S independent replicates (e.g., for computing bias, coverage, SE calibration), run with different seeds and output directories. + +## Using the Config Generator + +```bash +python experiments/semisynthetic_simulation/generate_configs.py \ + my_scenario \ + --n_runs 50 \ + --base_seed 42 \ + --experiments_dir ./outputs/causal/semisynthetic_study/runs \ + --meds ./path/to/meds/data +``` + +This generates one config per run under `generated_configs/`: + +``` +generated_configs/ +├── my_scenario_run_01.yaml # seed=43, outcomes → runs/run_01/my_scenario/ +├── my_scenario_run_02.yaml # seed=44, outcomes → runs/run_02/my_scenario/ +├── ... +└── my_scenario_run_50.yaml # seed=92, outcomes → runs/run_50/my_scenario/ +``` + +Each config is a copy of the base config with: +- `seed` set to `base_seed + run_number` +- `paths.outcomes` set to a unique per-run directory + +## Running All Replicates + +### Locally (sequential) + +```bash +for cfg in generated_configs/my_scenario_run_*.yaml; do + python -m corebehrt.main_causal.simulate_semisynthetic --config_path "$cfg" +done +``` + +### On Azure (parallel) + +Submit each run as a separate Azure job: + +```bash +for cfg in generated_configs/my_scenario_run_*.yaml; do + python -m corebehrt.azure job simulate_semisynthetic CPU-20-LP \ + --config "$cfg" -e semisynthetic_study +done +``` + +## Output Structure + +``` +outputs/causal/semisynthetic_study/runs/ +├── run_01/my_scenario/ +│ ├── counterfactuals.csv +│ ├── ite.csv +│ └── ... +├── run_02/my_scenario/ +│ └── ... +└── run_50/my_scenario/ + └── ... +``` + +## Evaluation Metrics + +Across S runs, compute: +- **Bias**: mean(theta_hat) - theta_true +- **Empirical SD**: std(theta_hat) +- **Mean estimated SE**: mean(SE_hat) +- **SE calibration**: SD_emp / mean(SE_hat) (target: 1) +- **Coverage**: fraction of runs where 95% CI contains theta_true + +The true effect theta_true is computed from the known P(Y(0)) and P(Y(1)) in each run's `counterfactuals.csv`. + +## Scenario Grid + +A typical experiment varies 3 dimensions: + +| Dimension | Levels | Config parameter | +|-----------|--------|-----------------| +| Baseline prevalence | low (~5%), moderate (~15%), high (~30%) | `beta_0` | +| Outcome complexity | simple (few features), rich (all features + interactions) | `coefficients`, `interactions` | +| Treatment effect | null (delta=0), non-null (delta=1) | `treatment_effect.delta` | + +Create one experiment config per scenario (e.g., `low_simple_null.yaml`, `low_rich_nonnull.yaml`), then run each scenario for S replicates. diff --git a/experiments/semisynthetic_simulation/generate_configs.py b/experiments/semisynthetic_simulation/generate_configs.py new file mode 100644 index 00000000..92691176 --- /dev/null +++ b/experiments/semisynthetic_simulation/generate_configs.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +"""Generate per-run config files for multi-replicate semisynthetic simulation. + +For each run, produces a YAML config with a unique seed and output directory. +Seed = base_seed + run_number (e.g., base_seed=42, run_01 -> seed=43). + +Usage: + python generate_configs.py my_scenario --n_runs 50 --base_seed 42 +""" + +import argparse +from pathlib import Path + +import yaml + + +def generate_configs( + scenario_name, + base_config_path, + n_runs, + base_seed=42, + experiments_dir="./outputs/causal/semisynthetic_study/runs", + meds_data=None, +): + """Generate one config file per run from a base config.""" + with open(base_config_path) as f: + base_config = yaml.safe_load(f) + + output_dir = Path("generated_configs") + output_dir.mkdir(parents=True, exist_ok=True) + + for run_number in range(1, n_runs + 1): + run_id = f"run_{run_number:02d}" + seed = base_seed + run_number + + config = _deep_copy_config(base_config) + config["seed"] = seed + config["paths"]["outcomes"] = f"{experiments_dir}/{run_id}/{scenario_name}" + if meds_data is not None: + config["paths"]["data"] = meds_data + + output_path = output_dir / f"{scenario_name}_{run_id}.yaml" + with open(output_path, "w") as f: + yaml.dump(config, f, default_flow_style=False, sort_keys=False) + + print(f"Generated: {output_path} (seed={seed})") + + print(f"\n{n_runs} configs generated for scenario: {scenario_name}") + + +def _deep_copy_config(config): + """Deep copy a config dict (handles nested dicts and lists).""" + if isinstance(config, dict): + return {k: _deep_copy_config(v) for k, v in config.items()} + elif isinstance(config, list): + return [_deep_copy_config(item) for item in config] + return config + + +def main(): + parser = argparse.ArgumentParser( + description="Generate per-run configs for semisynthetic simulation" + ) + parser.add_argument("scenario_name", help="Scenario identifier") + parser.add_argument( + "--base_config", + default="corebehrt/configs/causal/simulate_semisynthetic.yaml", + help="Path to base config (default: simulate_semisynthetic.yaml)", + ) + parser.add_argument( + "--n_runs", type=int, default=50, help="Number of runs (default: 50)" + ) + parser.add_argument( + "--base_seed", type=int, default=42, help="Base seed (default: 42)" + ) + parser.add_argument( + "--experiments_dir", + default="./outputs/causal/semisynthetic_study/runs", + help="Base output directory for runs", + ) + parser.add_argument("--meds", default=None, help="Override MEDS data path") + args = parser.parse_args() + + generate_configs( + args.scenario_name, + args.base_config, + args.n_runs, + args.base_seed, + args.experiments_dir, + args.meds, + ) + + +if __name__ == "__main__": + main() From 187bda6e4b51be6c067ce91fe55c4edd3ad85210 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 12:00:16 +0200 Subject: [PATCH 04/20] vectorize feature functions and add public calibration API - Vectorize _compute_age, _compute_event_recency, and _compute_sequence_motif_count using pandas merge/reindex instead of Python loops over patients - Use None sentinel instead of "" for prefix in _filter_by_prefix_and_window - Add public extract_features_and_probabilities() on simulator so calibration script no longer calls private methods Co-Authored-By: Claude Opus 4.6 (1M context) --- .../main_causal/calibrate_semisynthetic.py | 36 ++------- .../modules/simulation/oracle_features.py | 76 ++++++++++--------- .../simulation/semisynthetic_simulator.py | 41 ++++++++++ 3 files changed, 88 insertions(+), 65 deletions(-) diff --git a/corebehrt/main_causal/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py index 1adefa18..e60be594 100644 --- a/corebehrt/main_causal/calibrate_semisynthetic.py +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -12,7 +12,6 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd -from scipy.special import expit from corebehrt.functional.setup.args import get_args from corebehrt.functional.utils.azure_save import save_figure_with_azure_copy @@ -22,7 +21,6 @@ from corebehrt.modules.simulation.config_semisynthetic import ( create_semisynthetic_config, ) -from corebehrt.modules.simulation.oracle_features import extract_oracle_features from corebehrt.modules.simulation.plot import plot_probability_distributions from corebehrt.modules.simulation.semisynthetic_simulator import ( SemiSyntheticCausalSimulator, @@ -48,37 +46,17 @@ def main_calibrate(config_path): all_tau = {} # outcome_name -> [] for shard, _ in shard_loader(): - pids, is_exposed, index_dates = simulator._extract_treatment_and_index_dates( - shard - ) - if len(pids) == 0: - continue - - history_df = simulator._filter_to_pre_index(shard, index_dates) - history_df, pids, is_exposed, index_dates = simulator._apply_min_num_codes( - history_df, pids, is_exposed, index_dates - ) - if len(pids) == 0: + result = simulator.extract_features_and_probabilities(shard) + if result is None: continue - - features_df, _ = extract_oracle_features( - history_df, pids, index_dates, sim_config.features - ) + features_df, _, is_exposed, probas, tau = result all_features.append(features_df.assign(is_exposed=is_exposed)) all_is_exposed.append(is_exposed) - - for outcome_name, outcome_cfg in sim_config.outcomes.items(): - eta_0 = simulator._compute_eta_0(features_df, outcome_cfg.outcome_model) - tau = simulator._compute_tau(features_df, outcome_cfg.treatment_effect) - # Noise omitted intentionally: calibration shows the deterministic - # risk surface, not the noisy realization used during sampling. - p0 = expit(eta_0) - p1 = expit(eta_0 + tau) - + for outcome_name in sim_config.outcomes: all_probas.setdefault(outcome_name, {"P0": [], "P1": []}) - all_probas[outcome_name]["P0"].append(p0) - all_probas[outcome_name]["P1"].append(p1) - all_tau.setdefault(outcome_name, []).append(tau) + all_probas[outcome_name]["P0"].append(probas[outcome_name]["P0"]) + all_probas[outcome_name]["P1"].append(probas[outcome_name]["P1"]) + all_tau.setdefault(outcome_name, []).append(tau[outcome_name]) if not all_features: logger.error("No patients found across shards.") diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index efe5a637..4c1b280b 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -86,11 +86,14 @@ def extract_oracle_features( def _filter_by_prefix_and_window(history_df, index_dates, prefix, window_days): """Filter events matching code prefix within lookback window per patient.""" - mask = history_df[CONCEPT_COL].str.startswith(prefix) - filtered = history_df[mask].copy() - if filtered.empty: - logger.warning("No codes found with prefix '%s'", prefix) - return filtered + if prefix is not None: + mask = history_df[CONCEPT_COL].str.startswith(prefix) + filtered = history_df[mask].copy() + if filtered.empty: + logger.warning("No codes found with prefix '%s'", prefix) + return filtered + else: + filtered = history_df.copy() if window_days is not None: cutoff = index_dates.reindex(filtered[PID_COL]).values - pd.Timedelta( days=window_days @@ -119,7 +122,7 @@ def _count_per_patient(filtered_df, pids, count_col=None, unique=False): def _compute_recent_event_count(history_df, pids, index_dates, recent_window_days): filtered = _filter_by_prefix_and_window( - history_df, index_dates, "", recent_window_days + history_df, index_dates, None, recent_window_days ) return _count_per_patient(filtered, pids) @@ -139,25 +142,24 @@ def _compute_medication_count(history_df, pids, index_dates, med_prefix, lookbac def _compute_utilization_intensity(history_df, pids, index_dates, lookback_days): - filtered = _filter_by_prefix_and_window(history_df, index_dates, "", lookback_days) + filtered = _filter_by_prefix_and_window( + history_df, index_dates, None, lookback_days + ) return _count_per_patient(filtered, pids) def _compute_age(history_df, pids, index_dates): dob_events = history_df[history_df[CONCEPT_COL] == BIRTH_CODE] dob_per_patient = dob_events.groupby(PID_COL)[TIMESTAMP_COL].first() - age_series = pd.Series(index=pids, dtype=float) - for pid in pids: - if pid in dob_per_patient.index: - dob = dob_per_patient[pid] - idx_date = index_dates[pid] - age_series[pid] = (idx_date - dob).days / 365.25 + pid_index = pd.Index(pids) + dob_aligned = dob_per_patient.reindex(pid_index) + idx_aligned = index_dates.reindex(pid_index) + age_series = (idx_aligned - dob_aligned).dt.days / 365.25 mean_age = age_series.mean() if np.isnan(mean_age): mean_age = 65.0 logger.warning("No DOB events found, using default age %.0f", mean_age) - age_series = age_series.fillna(mean_age) - return age_series + return age_series.fillna(mean_age) def _compute_chronic_disease_count(history_df, pids, diag_prefix): @@ -182,25 +184,22 @@ def _compute_code_diversity(history_df, pids): def _compute_event_recency(history_df, pids, index_dates): last_event = history_df.groupby(PID_COL)[TIMESTAMP_COL].max() - recency = pd.Series(index=pids, dtype=float) - for pid in pids: - if pid in last_event.index: - recency[pid] = (index_dates[pid] - last_event[pid]).days - else: - recency[pid] = np.nan + pid_index = pd.Index(pids) + last_aligned = last_event.reindex(pid_index) + idx_aligned = index_dates.reindex(pid_index) + recency = (idx_aligned - last_aligned).dt.days.astype(float) mean_recency = recency.mean() - recency = recency.fillna(mean_recency) - return recency + return recency.fillna(mean_recency) def _compute_recent_burst_ratio( history_df, pids, index_dates, burst_window_days, lookback_days ): burst_filtered = _filter_by_prefix_and_window( - history_df, index_dates, "", burst_window_days + history_df, index_dates, None, burst_window_days ) lookback_filtered = _filter_by_prefix_and_window( - history_df, index_dates, "", lookback_days + history_df, index_dates, None, lookback_days ) burst_counts = _count_per_patient(burst_filtered, pids) lookback_counts = _count_per_patient(lookback_filtered, pids) @@ -216,18 +215,23 @@ def _compute_sequence_motif_count( logger.warning("Missing diagnosis or medication codes for motif counting") return pd.Series(0, index=pids, dtype=int) - motif_counts = {} + # Inner-join per patient produces all (diag, med) timestamp pairs. + # NOTE: this can be large if a patient has many diagnoses AND many medications + # (cross product per patient), but for typical EHR data this is fine. + diag_slim = diag_events[[PID_COL, TIMESTAMP_COL]].rename( + columns={TIMESTAMP_COL: "_diag_time"} + ) + med_slim = med_events[[PID_COL, TIMESTAMP_COL]].rename( + columns={TIMESTAMP_COL: "_med_time"} + ) + pairs = diag_slim.merge(med_slim, on=PID_COL) + + gap = pairs["_med_time"] - pairs["_diag_time"] window = pd.Timedelta(days=motif_window_days) - for pid in pids: - pid_diags = diag_events[diag_events[PID_COL] == pid][TIMESTAMP_COL].values - pid_meds = med_events[med_events[PID_COL] == pid][TIMESTAMP_COL].values - count = 0 - for diag_time in pid_diags: - gaps = pid_meds - diag_time - count += int(np.sum((gaps >= np.timedelta64(0)) & (gaps <= window))) - motif_counts[pid] = count - - return pd.Series(motif_counts).reindex(pids, fill_value=0) + valid = pairs[(gap >= pd.Timedelta(0)) & (gap <= window)] + + counts = valid.groupby(PID_COL).size() + return counts.reindex(pids, fill_value=0) # --------------------------------------------------------------------------- diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index c1972abb..dfc9ecc7 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -103,6 +103,47 @@ def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: all_factual_events, ) + def extract_features_and_probabilities( + self, shard_df: pd.DataFrame + ) -> tuple | None: + """Extract features and compute noiseless probabilities for calibration. + + Noise is omitted intentionally: calibration shows the deterministic + risk surface, not the noisy realization used during sampling. + + Returns None if no patients remain after filtering, otherwise: + (features_df, pids, is_exposed, probas_dict, tau_dict) + where probas_dict and tau_dict map outcome_name -> arrays. + """ + pids, is_exposed, index_dates = self._extract_treatment_and_index_dates( + shard_df + ) + if len(pids) == 0: + return None + + history_df = self._filter_to_pre_index(shard_df, index_dates) + history_df, pids, is_exposed, index_dates = self._apply_min_num_codes( + history_df, pids, is_exposed, index_dates + ) + if len(pids) == 0: + return None + + features_df, _ = extract_oracle_features( + history_df, pids, index_dates, self.config.features + ) + + probas_dict = {} + tau_dict = {} + for outcome_name, outcome_cfg in self.config.outcomes.items(): + eta_0 = self._compute_eta_0(features_df, outcome_cfg.outcome_model) + tau = self._compute_tau(features_df, outcome_cfg.treatment_effect) + p0 = expit(eta_0) + p1 = expit(eta_0 + tau) + probas_dict[outcome_name] = {"P0": p0, "P1": p1} + tau_dict[outcome_name] = tau + + return features_df, pids, is_exposed, probas_dict, tau_dict + # ------------------------------------------------------------------ # Data extraction # ------------------------------------------------------------------ From 515d8dffed7f36e607fd133ad0e4a239529291eb Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 12:14:25 +0200 Subject: [PATCH 05/20] fix per-shard standardization: use two-pass global z-scoring Features were z-scored per shard independently, making simulated probabilities depend on which shard a patient lands in. Now the simulator does two passes: pass 1 computes global mean/std across all shards, pass 2 standardizes using those global stats before computing logits. - extract_oracle_features now returns raw (unstandardized) features - New standardize_features(df, means, stds) function for explicit control - compute_global_feature_stats() iterates shards once to collect stats - Entry points (simulate + calibrate) call pass 1 then pass 2 - Removed standardize flag from FeatureConfig (always standardize globally) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../causal/simulate_semisynthetic.yaml | 1 - .../main_causal/calibrate_semisynthetic.py | 6 ++- .../main_causal/simulate_semisynthetic.py | 6 +++ .../simulation/config_semisynthetic.py | 1 - .../modules/simulation/oracle_features.py | 42 ++++++++------- .../simulation/semisynthetic_simulator.py | 53 +++++++++++++++++-- .../test_simulation/test_oracle_features.py | 41 +++++++------- .../test_simulation/test_semisynthetic.py | 2 +- 8 files changed, 102 insertions(+), 50 deletions(-) diff --git a/corebehrt/configs/causal/simulate_semisynthetic.yaml b/corebehrt/configs/causal/simulate_semisynthetic.yaml index 1421ec0e..c8992987 100644 --- a/corebehrt/configs/causal/simulate_semisynthetic.yaml +++ b/corebehrt/configs/causal/simulate_semisynthetic.yaml @@ -31,7 +31,6 @@ features: recent_window_days: 90 burst_window_days: 30 motif_window_days: 30 - standardize: true # -------------------------------------------------------------------------- # Outcomes diff --git a/corebehrt/main_causal/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py index e60be594..6275007c 100644 --- a/corebehrt/main_causal/calibrate_semisynthetic.py +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -39,7 +39,11 @@ def main_calibrate(config_path): sim_config = create_semisynthetic_config(cfg) simulator = SemiSyntheticCausalSimulator(sim_config) - # Accumulate across shards + # Pass 1: compute global feature statistics for standardization + logger.info("Computing global feature statistics...") + simulator.compute_global_feature_stats(shard_loader) + + # Pass 2: extract standardized features and probabilities all_features = [] all_is_exposed = [] all_probas = {} # outcome_name -> {"P0": [], "P1": []} diff --git a/corebehrt/main_causal/simulate_semisynthetic.py b/corebehrt/main_causal/simulate_semisynthetic.py index 2efeb644..24398e8e 100644 --- a/corebehrt/main_causal/simulate_semisynthetic.py +++ b/corebehrt/main_causal/simulate_semisynthetic.py @@ -29,6 +29,12 @@ def main_simulate(config_path): shard_loader = ShardLoader(cfg.paths.data, cfg.paths.splits) simulation_config = create_semisynthetic_config(cfg) simulator = CausalSimulator(simulation_config) + + # Pass 1: compute global feature means/stds for standardization + logger.info("--- Pass 1: computing global feature statistics ---") + simulator.compute_global_feature_stats(shard_loader) + + # Pass 2: simulate outcomes using globally standardized features simulate(shard_loader, simulator, cfg.paths.outcomes) diff --git a/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py index 1ed8f1c8..6b555183 100644 --- a/corebehrt/modules/simulation/config_semisynthetic.py +++ b/corebehrt/modules/simulation/config_semisynthetic.py @@ -30,7 +30,6 @@ class FeatureConfig: recent_window_days: int = 90 burst_window_days: int = 30 motif_window_days: int = 30 - standardize: bool = True @dataclass diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index 4c1b280b..b4764462 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -1,7 +1,6 @@ """Extract hand-crafted oracle features from pre-index patient histories.""" import logging -from typing import List, Tuple import numpy as np import pandas as pd @@ -17,8 +16,11 @@ def extract_oracle_features( pids: np.ndarray, index_dates: pd.Series, feature_config: FeatureConfig, -) -> Tuple[pd.DataFrame, List[str]]: - """Extract oracle features from pre-index patient histories. +) -> pd.DataFrame: + """Extract raw oracle features from pre-index patient histories. + + Returns unstandardized features. Use ``standardize_features`` with + global statistics to z-score across the full cohort. Args: history_df: MEDS DataFrame already filtered to pre-index events @@ -28,12 +30,10 @@ def extract_oracle_features( Returns: features_df: DataFrame with PID_COL as index, one column per feature - feature_names: list of feature names present in the output """ prefixes = feature_config.code_prefixes features = {} - # Baseline risk features features["recent_event_count"] = _compute_recent_event_count( history_df, pids, index_dates, feature_config.recent_window_days ) @@ -52,7 +52,6 @@ def extract_oracle_features( ) features["code_diversity"] = _compute_code_diversity(history_df, pids) - # Longitudinal features features["event_recency"] = _compute_event_recency(history_df, pids, index_dates) features["recent_burst_ratio"] = _compute_recent_burst_ratio( history_df, @@ -72,11 +71,25 @@ def extract_oracle_features( features_df = pd.DataFrame(features, index=pids) features_df.index.name = PID_COL + return features_df + + +def standardize_features(features_df, means=None, stds=None): + """Z-score features using provided or computed statistics. - if feature_config.standardize: - features_df = _standardize(features_df) + Args: + features_df: raw features DataFrame + means: per-feature means (computed from df if None) + stds: per-feature stds (computed from df if None) - return features_df, list(features_df.columns) + Returns: + standardized DataFrame, means Series, stds Series + """ + if means is None: + means = features_df.mean() + if stds is None: + stds = features_df.std(ddof=0).replace(0, 1).fillna(1) + return (features_df - means) / stds, means, stds # --------------------------------------------------------------------------- @@ -232,14 +245,3 @@ def _compute_sequence_motif_count( counts = valid.groupby(PID_COL).size() return counts.reindex(pids, fill_value=0) - - -# --------------------------------------------------------------------------- -# Standardization -# --------------------------------------------------------------------------- - - -def _standardize(features_df): - means = features_df.mean() - stds = features_df.std(ddof=0).replace(0, 1).fillna(1) - return (features_df - means) / stds diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index dfc9ecc7..22f00b6c 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -38,7 +38,10 @@ SemiSyntheticSimulationConfig, TreatmentEffectConfig, ) -from corebehrt.modules.simulation.oracle_features import extract_oracle_features +from corebehrt.modules.simulation.oracle_features import ( + extract_oracle_features, + standardize_features, +) from corebehrt.modules.simulation.plot import ( plot_probability_distributions, plot_true_effects_vs_risk_differences, @@ -61,12 +64,44 @@ class SemiSyntheticCausalSimulator: def __init__(self, config: SemiSyntheticSimulationConfig): self.config = config self.rng = np.random.default_rng(config.seed) + self._global_means = None + self._global_stds = None + + def compute_global_feature_stats(self, shard_loader): + """Pass 1: compute global mean/std across all shards for standardization.""" + all_features = [] + for shard, _ in shard_loader(): + pids, is_exposed, index_dates = self._extract_treatment_and_index_dates( + shard + ) + if len(pids) == 0: + continue + history_df = self._filter_to_pre_index(shard, index_dates) + history_df, pids, is_exposed, index_dates = self._apply_min_num_codes( + history_df, pids, is_exposed, index_dates + ) + if len(pids) == 0: + continue + features_df = extract_oracle_features( + history_df, pids, index_dates, self.config.features + ) + all_features.append(features_df) + + if not all_features: + logger.warning("No patients found during global stats computation") + return + + combined = pd.concat(all_features) + self._global_means = combined.mean() + self._global_stds = combined.std(ddof=0).replace(0, 1).fillna(1) + logger.info("Computed global feature stats from %d patients", len(combined)) def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: """Orchestrate the semi-synthetic simulation for a single data shard. - Returns a dict of DataFrames matching the format produced by - ``RealisticCausalSimulator.simulate_dataset``. + Call ``compute_global_feature_stats`` first to enable global + standardization. Returns a dict of DataFrames matching the format + produced by ``RealisticCausalSimulator.simulate_dataset``. """ pids, is_exposed, index_dates = self._extract_treatment_and_index_dates( shard_df @@ -81,9 +116,13 @@ def simulate_dataset(self, shard_df: pd.DataFrame) -> Dict[str, pd.DataFrame]: if len(pids) == 0: return {} - features_df, _ = extract_oracle_features( + features_df = extract_oracle_features( history_df, pids, index_dates, self.config.features ) + if self._global_means is not None: + features_df, _, _ = standardize_features( + features_df, self._global_means, self._global_stds + ) ite_records, cf_records, all_factual_events = self._simulate_outcomes( features_df, pids, is_exposed, index_dates @@ -128,9 +167,13 @@ def extract_features_and_probabilities( if len(pids) == 0: return None - features_df, _ = extract_oracle_features( + features_df = extract_oracle_features( history_df, pids, index_dates, self.config.features ) + if self._global_means is not None: + features_df, _, _ = standardize_features( + features_df, self._global_means, self._global_stds + ) probas_dict = {} tau_dict = {} diff --git a/tests/test_modules/test_simulation/test_oracle_features.py b/tests/test_modules/test_simulation/test_oracle_features.py index a5fd66da..5ae4966b 100644 --- a/tests/test_modules/test_simulation/test_oracle_features.py +++ b/tests/test_modules/test_simulation/test_oracle_features.py @@ -5,7 +5,10 @@ from corebehrt.constants.data import CONCEPT_COL, PID_COL, TIMESTAMP_COL from corebehrt.modules.simulation.config_semisynthetic import FeatureConfig -from corebehrt.modules.simulation.oracle_features import extract_oracle_features +from corebehrt.modules.simulation.oracle_features import ( + extract_oracle_features, + standardize_features, +) def _make_test_df(records): @@ -19,8 +22,7 @@ def _make_test_df(records): def _default_feature_config(**overrides): - kwargs = dict(standardize=False) - kwargs.update(overrides) + kwargs = dict(**overrides) return FeatureConfig(**kwargs) @@ -44,12 +46,9 @@ def test_extract_oracle_features_basic(self): pids = np.array([1, 2, 3, 4, 5]) index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) config = _default_feature_config() - features_df, feature_names = extract_oracle_features( - history_df, pids, index_dates, config - ) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.shape[0], 5) self.assertEqual(features_df.shape[1], 10) - self.assertEqual(feature_names, list(features_df.columns)) class TestRecentEventCount(unittest.TestCase): @@ -67,7 +66,7 @@ def test_recent_event_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(recent_window_days=90) - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "recent_event_count"], 3) @@ -86,7 +85,7 @@ def test_disease_burden(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "disease_burden"], 5) @@ -103,7 +102,7 @@ def test_medication_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "medication_count"], 3) @@ -118,7 +117,7 @@ def test_age_computation(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) expected_age = (index - pd.Timestamp("1990-01-01")).days / 365.25 self.assertAlmostEqual(features_df.loc[1, "age"], expected_age, places=2) @@ -134,7 +133,7 @@ def test_event_recency(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config() - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "event_recency"], 10) @@ -153,7 +152,7 @@ def test_recent_burst_ratio(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(burst_window_days=30, lookback_days=365) - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) # burst=5, lookback=20, ratio = 5 / (20 + 1) expected = 5.0 / 21.0 self.assertAlmostEqual( @@ -172,7 +171,7 @@ def test_no_medication_codes(self): pids = np.array([1, 2]) index_dates = pd.Series({1: index, 2: index}) config = _default_feature_config() - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "medication_count"], 0) self.assertEqual(features_df.loc[2, "medication_count"], 0) @@ -193,13 +192,13 @@ def test_standardization(self): history_df = _make_test_df(records) pids = np.array([1, 2, 3]) index_dates = pd.Series({p: pd.Timestamp("2021-01-01") for p in pids}) - config = FeatureConfig(standardize=True) - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) - for col in features_df.columns: - col_std = features_df[col].std() - # columns with zero variance stay at 0 after standardization + config = FeatureConfig() + features_df = extract_oracle_features(history_df, pids, index_dates, config) + standardized_df, _, _ = standardize_features(features_df) + for col in standardized_df.columns: + col_std = standardized_df[col].std() if col_std > 1e-10: - self.assertAlmostEqual(features_df[col].mean(), 0.0, places=10) + self.assertAlmostEqual(standardized_df[col].mean(), 0.0, places=10) class TestSequenceMotifCount(unittest.TestCase): @@ -217,7 +216,7 @@ def test_sequence_motif_count(self): pids = np.array([1]) index_dates = pd.Series({1: index}) config = _default_feature_config(motif_window_days=30) - features_df, _ = extract_oracle_features(history_df, pids, index_dates, config) + features_df = extract_oracle_features(history_df, pids, index_dates, config) self.assertEqual(features_df.loc[1, "sequence_motif_count"], 1) diff --git a/tests/test_modules/test_simulation/test_semisynthetic.py b/tests/test_modules/test_simulation/test_semisynthetic.py index 988135bf..6af7b96d 100644 --- a/tests/test_modules/test_simulation/test_semisynthetic.py +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -104,7 +104,7 @@ def _make_config( ) return SemiSyntheticSimulationConfig( paths=paths, - features=FeatureConfig(standardize=True), + features=FeatureConfig(), outcomes={"OUTCOME_test": outcome}, seed=seed, min_num_codes=min_num_codes, From 50dd4f5ec5be8f8d94cce2ca5f9c21e86d4a645a Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 12:26:42 +0200 Subject: [PATCH 06/20] fix counterfactual AUC: score against matching labels auc_if_all_treated now uses Y(1) as labels (not factual Y), and auc_if_all_control uses Y(0). Previously both scored against the factual outcome which mixes treated and untreated labels. Co-Authored-By: Claude Opus 4.6 (1M context) --- corebehrt/modules/simulation/semisynthetic_simulator.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 22f00b6c..e43342bf 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -560,14 +560,19 @@ def _calculate_theoretical_roc_auc( continue y_true = cf_records[outcome_col] + y1_col = f"{SIMULATED_OUTCOME_EXPOSED}_{outcome_name}" + y0_col = f"{SIMULATED_OUTCOME_CONTROL}_{outcome_name}" p_treated = cf_records[p_exposed_col] p_control = cf_records[p_control_col] y_prob_factual = np.where(is_exposed, p_treated, p_control) if len(np.unique(y_true)) > 1: auc_factual = roc_auc_score(y_true, y_prob_factual) - auc_treated = roc_auc_score(y_true, p_treated) - auc_control = roc_auc_score(y_true, p_control) + # Score counterfactual probabilities against their matching labels + y1 = cf_records[y1_col] + y0 = cf_records[y0_col] + auc_treated = roc_auc_score(y1, p_treated) if len(np.unique(y1)) > 1 else np.nan + auc_control = roc_auc_score(y0, p_control) if len(np.unique(y0)) > 1 else np.nan theoretical_aucs[outcome_name] = auc_factual results_data.append( From d7bb3916b9f1a1e911ce2314768108a8c6c609c4 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 12:28:49 +0200 Subject: [PATCH 07/20] revert counterfactual AUC change: score against factual outcome MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The auc_if_all_treated/control metrics intentionally score P(Y(1)) and P(Y(0)) against the factual observed outcome, matching the existing RealisticCausalSimulator. This measures how well counterfactual risk models discriminate observed outcomes — a more informative diagnostic than scoring against counterfactual labels (which would be near-perfect since Y(a) ~ Bernoulli(P(Y(a)))). Co-Authored-By: Claude Opus 4.6 (1M context) --- corebehrt/modules/simulation/semisynthetic_simulator.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index e43342bf..22f00b6c 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -560,19 +560,14 @@ def _calculate_theoretical_roc_auc( continue y_true = cf_records[outcome_col] - y1_col = f"{SIMULATED_OUTCOME_EXPOSED}_{outcome_name}" - y0_col = f"{SIMULATED_OUTCOME_CONTROL}_{outcome_name}" p_treated = cf_records[p_exposed_col] p_control = cf_records[p_control_col] y_prob_factual = np.where(is_exposed, p_treated, p_control) if len(np.unique(y_true)) > 1: auc_factual = roc_auc_score(y_true, y_prob_factual) - # Score counterfactual probabilities against their matching labels - y1 = cf_records[y1_col] - y0 = cf_records[y0_col] - auc_treated = roc_auc_score(y1, p_treated) if len(np.unique(y1)) > 1 else np.nan - auc_control = roc_auc_score(y0, p_control) if len(np.unique(y0)) > 1 else np.nan + auc_treated = roc_auc_score(y_true, p_treated) + auc_control = roc_auc_score(y_true, p_control) theoretical_aucs[outcome_name] = auc_factual results_data.append( From 351c22b6f6e62264d154a81f231317cecf9e7980 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 13:10:26 +0200 Subject: [PATCH 08/20] add NaN fallback for event recency when no events exist Same pattern as the DOB/age fallback: if all patients in a shard have no events, mean_recency is NaN. Fall back to 365 days. Co-Authored-By: Claude Opus 4.6 (1M context) --- corebehrt/modules/simulation/oracle_features.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index b4764462..fcdaedf6 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -202,6 +202,9 @@ def _compute_event_recency(history_df, pids, index_dates): idx_aligned = index_dates.reindex(pid_index) recency = (idx_aligned - last_aligned).dt.days.astype(float) mean_recency = recency.mean() + if np.isnan(mean_recency): + mean_recency = 365.0 + logger.warning("No events found for recency, using default %.0f days", mean_recency) return recency.fillna(mean_recency) From e962f8bf10d841d6d46ed0ef70a9fde96fcb0800 Mon Sep 17 00:00:00 2001 From: Kiril Klein Date: Fri, 3 Apr 2026 13:13:48 +0200 Subject: [PATCH 09/20] fix ruff format for oracle_features.py Co-Authored-By: Claude Opus 4.6 (1M context) --- corebehrt/modules/simulation/oracle_features.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index fcdaedf6..0e19f434 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -204,7 +204,9 @@ def _compute_event_recency(history_df, pids, index_dates): mean_recency = recency.mean() if np.isnan(mean_recency): mean_recency = 365.0 - logger.warning("No events found for recency, using default %.0f days", mean_recency) + logger.warning( + "No events found for recency, using default %.0f days", mean_recency + ) return recency.fillna(mean_recency) From 86d83dcd218957cacffc9a2675aead20bfbbbf7a Mon Sep 17 00:00:00 2001 From: kirilklein Date: Fri, 26 Jun 2026 16:34:32 +0200 Subject: [PATCH 10/20] add semi-synthetic multi-run study (N outer x K inner) for azure Adds a study runner that mirrors the resampling study but uses the semi-synthetic simulator: each outer run is an independent simulation (own seed) followed by K inner reshuffle fits (finetune -> calibrate -> estimate), reusing the existing causal mains unchanged. - simulate_semisynthetic: optional config-gated patient sampling (SampledShardLoader), default off so existing behaviour is unchanged - config_semisynthetic: optional `cohort` path for saving sampled pids - run_study.py: lean N x K runner with templated base configs - run_semisynthetic_study azure component + job CLI registration - submit_runs.sh: one command submits N parallel outer-run jobs - base_configs + job template + docs/study.md (smoke test + full study) Smoke-tested locally: sampling + null/medium outcomes recover true effects (NULL RD 0.000, MEDIUM delta=0.5 -> RD 0.063) and emit the exposure/outcome/counterfactual files the downstream steps consume. Co-Authored-By: Claude Opus 4.8 --- .../components/run_semisynthetic_study.py | 48 +++++ corebehrt/azure/main/job.py | 1 + .../main_causal/simulate_semisynthetic.py | 51 +++++ .../simulation/config_semisynthetic.py | 1 + .../semisynthetic_simulation/README.md | 1 + .../semisynthetic_simulation/__init__.py | 0 .../base_configs/calibrate.yaml | 11 + .../base_configs/estimate.yaml | 30 +++ .../base_configs/finetune.yaml | 67 ++++++ .../base_configs/prepare.yaml | 44 ++++ .../base_configs/select_cohort.yaml | 34 ++++ .../base_configs/simulate.yaml | 62 ++++++ .../bash_scripts/submit_runs.sh | 51 +++++ .../semisynthetic_simulation/docs/study.md | 83 ++++++++ .../job_config_template.yaml | 13 ++ .../python_scripts/__init__.py | 0 .../python_scripts/run_study.py | 192 ++++++++++++++++++ 17 files changed, 689 insertions(+) create mode 100644 corebehrt/azure/components/run_semisynthetic_study.py create mode 100644 experiments/semisynthetic_simulation/__init__.py create mode 100644 experiments/semisynthetic_simulation/base_configs/calibrate.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/estimate.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/finetune.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/prepare.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/select_cohort.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/simulate.yaml create mode 100755 experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh create mode 100644 experiments/semisynthetic_simulation/docs/study.md create mode 100644 experiments/semisynthetic_simulation/job_config_template.yaml create mode 100644 experiments/semisynthetic_simulation/python_scripts/__init__.py create mode 100644 experiments/semisynthetic_simulation/python_scripts/run_study.py diff --git a/corebehrt/azure/components/run_semisynthetic_study.py b/corebehrt/azure/components/run_semisynthetic_study.py new file mode 100644 index 00000000..c2d62b76 --- /dev/null +++ b/corebehrt/azure/components/run_semisynthetic_study.py @@ -0,0 +1,48 @@ +"""Azure component for the semi-synthetic simulation study (one outer run per job).""" + +import shlex + +from corebehrt.azure.util import job + +INPUTS = { + "meds": {"type": "uri_folder"}, + "features": {"type": "uri_folder"}, + "tokenized": {"type": "uri_folder"}, + "pretrain_model": {"type": "uri_folder"}, +} + +OUTPUTS = { + "results": {"type": "uri_folder"}, # Output dir for this outer run +} + + +def main_run_study(config_path): + """Translate the job config into run_study CLI arguments and run it.""" + from corebehrt.modules.setup.config import load_config + + cfg = load_config(config_path) + + args = [ + "--meds", + cfg.paths.meds, + "--features", + cfg.paths.features, + "--tokenized", + cfg.paths.tokenized, + "--pretrain-model", + cfg.paths.pretrain_model, + "--experiment-dir", + cfg.paths.results, + ] + + # Run-specific flags (run-id, inner-runs, sample-fraction, ...) come via --bash-args. + if hasattr(cfg, "bash_args") and cfg.bash_args: + args.extend(shlex.split(cfg.bash_args)) + + from experiments.semisynthetic_simulation.python_scripts.run_study import main + + main(args) + + +if __name__ == "__main__": + job.run_main("run_semisynthetic_study", main_run_study, INPUTS, OUTPUTS) diff --git a/corebehrt/azure/main/job.py b/corebehrt/azure/main/job.py index 0f957f1c..90f4f4b9 100644 --- a/corebehrt/azure/main/job.py +++ b/corebehrt/azure/main/job.py @@ -62,6 +62,7 @@ def add_parser(subparsers) -> None: "evaluate_xgboost", "get_pat_counts_by_code", "run_batch_experiments", + "run_semisynthetic_study", }, help="Job to run.", ) diff --git a/corebehrt/main_causal/simulate_semisynthetic.py b/corebehrt/main_causal/simulate_semisynthetic.py index 24398e8e..2ededf7c 100644 --- a/corebehrt/main_causal/simulate_semisynthetic.py +++ b/corebehrt/main_causal/simulate_semisynthetic.py @@ -8,11 +8,16 @@ from corebehrt.modules.simulation.config_semisynthetic import ( create_semisynthetic_config, ) +from corebehrt.functional.causal.cohort_sampler import sample_cohort +from corebehrt.constants.data import PID_COL +from corebehrt.constants.paths import PID_FILE from collections import defaultdict +import os import pandas as pd from os.path import join import logging from tqdm import tqdm +import torch logger = logging.getLogger("simulate") @@ -30,6 +35,14 @@ def main_simulate(config_path): simulation_config = create_semisynthetic_config(cfg) simulator = CausalSimulator(simulation_config) + # Optionally restrict to a sampled subset of patients (e.g. for smoke tests + # or per-run resampling). Disabled by default, so default behaviour is unchanged. + sampling_cfg = cfg.get("sampling", {}) + if sampling_cfg.get("enabled", False): + shard_loader = _sample_patients( + shard_loader, sampling_cfg, cfg.get("seed", 42), cfg.paths.get("cohort") + ) + # Pass 1: compute global feature means/stds for standardization logger.info("--- Pass 1: computing global feature statistics ---") simulator.compute_global_feature_stats(shard_loader) @@ -38,6 +51,44 @@ def main_simulate(config_path): simulate(shard_loader, simulator, cfg.paths.outcomes) +class SampledShardLoader: + """Wraps a ShardLoader and filters every shard to a fixed set of patient IDs.""" + + def __init__(self, shard_loader: ShardLoader, pids_set: set): + self.shard_loader = shard_loader + self.pids_set = pids_set + + def __call__(self): + for shard, meta in self.shard_loader(): + yield shard[shard[PID_COL].isin(self.pids_set)], meta + + +def _sample_patients( + shard_loader: ShardLoader, sampling_cfg: dict, seed: int, cohort_dir: str +) -> SampledShardLoader: + """Sample a subset of patient IDs and return a shard loader filtered to them.""" + all_pids = set() + for shard, _ in tqdm(shard_loader(), desc="Scanning shards for sampling"): + all_pids.update(shard[PID_COL].unique()) + full_pids = torch.tensor(sorted(all_pids)) + + sampled_pids = sample_cohort( + full_pids, + sample_fraction=sampling_cfg.get("fraction"), + sample_size=sampling_cfg.get("size"), + seed=seed, + ) + logger.info( + f"Sampled {len(sampled_pids)} of {len(full_pids)} patients (seed={seed})" + ) + + if cohort_dir: + os.makedirs(cohort_dir, exist_ok=True) + torch.save(sampled_pids, join(cohort_dir, PID_FILE)) + + return SampledShardLoader(shard_loader, set(sampled_pids.tolist())) + + def simulate(shard_loader: ShardLoader, simulator: CausalSimulator, outcomes_dir: str): """ Simulates outcomes by processing data shards in a single pass. diff --git a/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py index 6b555183..f6b5904b 100644 --- a/corebehrt/modules/simulation/config_semisynthetic.py +++ b/corebehrt/modules/simulation/config_semisynthetic.py @@ -9,6 +9,7 @@ class PathsConfig: data: str splits: List[str] outcomes: str + cohort: str = None @dataclass diff --git a/experiments/semisynthetic_simulation/README.md b/experiments/semisynthetic_simulation/README.md index bd1a6589..db2d97b4 100644 --- a/experiments/semisynthetic_simulation/README.md +++ b/experiments/semisynthetic_simulation/README.md @@ -8,6 +8,7 @@ Semi-synthetic causal simulation where **treatment is kept from real data** and |------|-------| | **How to run locally** | [docs/local.md](docs/local.md) | | **How to run on Azure** | [docs/azure.md](docs/azure.md) | +| **How to run the multi-run study (N outer × K inner)** | [docs/study.md](docs/study.md) | | **How to run multiple replicates** | [docs/multiple_runs.md](docs/multiple_runs.md) | | **Feature definitions** | [docs/features.md](docs/features.md) | | **Calibrating parameters** | [docs/calibration.md](docs/calibration.md) | diff --git a/experiments/semisynthetic_simulation/__init__.py b/experiments/semisynthetic_simulation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/semisynthetic_simulation/base_configs/calibrate.yaml b/experiments/semisynthetic_simulation/base_configs/calibrate.yaml new file mode 100644 index 00000000..c9d43bc0 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/calibrate.yaml @@ -0,0 +1,11 @@ +# Stage 2b: calibrate predictions (isotonic regression on validation folds). +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + ## INPUTS + finetune_model: "{{INNER_DIR}}/models/bert" + + ## OUTPUTS + calibrated_predictions: "{{INNER_DIR}}/models/bert/calibrated" diff --git a/experiments/semisynthetic_simulation/base_configs/estimate.yaml b/experiments/semisynthetic_simulation/base_configs/estimate.yaml new file mode 100644 index 00000000..c250bf4b --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/estimate.yaml @@ -0,0 +1,30 @@ +# Stage 2c: estimate causal effects from calibrated predictions. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + ## INPUTS + calibrated_predictions: "{{INNER_DIR}}/models/bert/calibrated/" + counterfactual_outcomes: "{{RUN_DIR}}/simulated_outcomes" + + ## OUTPUTS + estimate: "{{INNER_DIR}}/estimate/bert" + +estimator: + methods: ["IPW", "TMLE", "TMLE_TH"] + effect_type: "ATE" + n_bootstrap: 30 + clip_percentile: 0.99 + +plot: + contingency_table: + max_outcomes_per_figure: 10 + max_number_of_figures: 10 + effect_size: + max_outcomes_per_figure: 10 + max_number_of_figures: 10 + plot_individual_effects: false + adjustment: + max_outcomes_per_figure: 8 + max_number_of_figures: 10 diff --git a/experiments/semisynthetic_simulation/base_configs/finetune.yaml b/experiments/semisynthetic_simulation/base_configs/finetune.yaml new file mode 100644 index 00000000..ddf6c005 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/finetune.yaml @@ -0,0 +1,67 @@ +# Stage 2a: finetune the causal model (exposure + outcomes jointly). +# {{INNER_DIR}} is unique per inner reshuffle; the runner sets data.reshuffle. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + ## INPUTS + pretrain_model: "{{PRETRAIN_MODEL}}" + prepared_data: "{{RUN_DIR}}/prepared_data" + + ## OUTPUTS + model: "{{INNER_DIR}}/models/bert" + +save_encodings: false +visualize_encodings: false +visualize_weight_distributions: false + +model: + head: + shared_representation: false + bidirectional: true + bottleneck_dim: 128 + l1_lambda: 0 + temperature: 1 + pooling_strategy: gru + loss: + name: bce + initialize_sigmoid_bias: false + +trainer_args: + loss_weight_function: + _target_: corebehrt.modules.trainer.utils.PositiveWeight.sqrt + batch_size: 64 + val_batch_size: 256 + effective_batch_size: 64 + epochs: 10 + info: true + shuffle: true + checkpoint_frequency: 1 + early_stopping: 4 + stopping_criterion: val_loss + n_layers_to_freeze: 0 + freeze_encoder_on_plateau: true + freeze_encoder_on_plateau_threshold: 0.01 + freeze_encoder_on_plateau_patience: 5 + use_pcgrad: false + plot_histograms: true + save_curves: false + plot_gradients: false + plot_gradients_frequency: 3 + +optimizer: + lr: 7e-5 + eps: 1e-6 + weight_decay: 0 + +scheduler: + _target_: transformers.get_cosine_schedule_with_warmup + num_training_epochs: 20 + num_warmup_epochs: 3 + +metrics: + roc_auc: + _target_: corebehrt.modules.monitoring.metrics.ROC_AUC + pr_auc: + _target_: corebehrt.modules.monitoring.metrics.PR_AUC diff --git a/experiments/semisynthetic_simulation/base_configs/prepare.yaml b/experiments/semisynthetic_simulation/base_configs/prepare.yaml new file mode 100644 index 00000000..61e95873 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/prepare.yaml @@ -0,0 +1,44 @@ +# Stage 1c: prepare finetuning data + cross-validation folds. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + ## INPUTS + features: "{{FEATURES}}" + tokenized: "{{TOKENIZED}}" + cohort: "{{RUN_DIR}}/cohort" + + outcomes: "{{RUN_DIR}}/simulated_outcomes" + outcome_files: + - OUTCOME_NULL.csv + - OUTCOME_MEDIUM.csv + + exposures: "{{RUN_DIR}}/simulated_outcomes" + exposure: exposure.csv + + ## OUTPUTS + prepared_data: "{{RUN_DIR}}/prepared_data" + +data: + type: finetune + truncation_len: 64 + min_len: 2 + cv_folds: 2 + min_instances_per_class: 10 + max_items_per_plot: 15 + max_number_of_plots: 10 + number_subjects_to_plot: 10 + sample_num_patients: null + sample_seed: 42 + +exposure: + n_hours_censoring: -10 + n_hours_start_follow_up: -10 + n_hours_end_follow_up: 1000 + +outcome: + n_hours_start_follow_up: 0 + n_hours_end_follow_up: null + n_hours_compliance: null + group_wise_follow_up: false diff --git a/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml b/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml new file mode 100644 index 00000000..14bfae6f --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml @@ -0,0 +1,34 @@ +# Stage 1b: select the analysis cohort from the simulated exposure. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + ### Inputs + features: "{{FEATURES}}/" + meds: "{{MEDS}}" + splits: [tuning] + exposures: "{{RUN_DIR}}/simulated_outcomes/" + exposure: exposure.csv + criteria_config: ./corebehrt/configs/causal/select_cohort_full/definitions.yaml + + ### Outputs + cohort: "{{RUN_DIR}}/cohort/" + +time_windows: + data_end: + year: 2025 + month: 01 + day: 01 + data_start: + year: 1920 + month: 1 + day: 1 + min_follow_up: + days: 1 + min_lookback: + days: 365 + +cv_folds: 2 +val_ratio: 0.1 +test_ratio: 0 diff --git a/experiments/semisynthetic_simulation/base_configs/simulate.yaml b/experiments/semisynthetic_simulation/base_configs/simulate.yaml new file mode 100644 index 00000000..267c4b75 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/simulate.yaml @@ -0,0 +1,62 @@ +# Stage 1a: semi-synthetic simulation (real treatment, simulated outcome). +# The double-brace placeholders are filled in per run by run_study.py. +# `seed` and the `sampling` block are injected by the runner. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + data: "{{MEDS}}" + splits: ["tuning"] + outcomes: "{{RUN_DIR}}/simulated_outcomes" + cohort: "{{RUN_DIR}}/cohort" + +min_num_codes: 3 +exposure_code: "EXPOSURE" + +features: + code_prefixes: + diagnosis: "D/" + medication: "M/" + procedure: "P/" + admission: "ADM/" + lookback_days: 365 + recent_window_days: 90 + burst_window_days: 30 + motif_window_days: 30 + +# Two outcomes sharing identical confounding, differing only in the treatment +# effect: a null effect (delta=0) and a medium effect (delta=0.5, ~6pp risk +# difference at a ~12% baseline). Add OUTCOME_LARGE (delta=1.0) for the full study. +outcomes: + OUTCOME_NULL: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + coefficients: + recent_event_count: 0.3 + disease_burden: 0.2 + medication_count: 0.15 + age: 0.4 + chronic_disease_count: 0.25 + event_recency: -0.15 + noise_scale: 0.1 + treatment_effect: + mode: constant + delta: 0.0 + + OUTCOME_MEDIUM: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + coefficients: + recent_event_count: 0.3 + disease_burden: 0.2 + medication_count: 0.15 + age: 0.4 + chronic_disease_count: 0.25 + event_recency: -0.15 + noise_scale: 0.1 + treatment_effect: + mode: constant + delta: 0.5 diff --git a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh new file mode 100755 index 00000000..d2a3f1ba --- /dev/null +++ b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Submit N independent outer runs of the semi-synthetic study as parallel Azure jobs. +# Each job runs one outer simulation (own seed) + K inner reshuffle fits. +# +# Usage: +# ./submit_runs.sh # 1 run, K=2, no sampling (full population) +# ./submit_runs.sh -n 2 -k 2 -f 0.1 # SMOKE TEST: 2 runs, 2 inner fits, 10% sampled +# ./submit_runs.sh -n 10 -k 10 # full study: 10 runs, 10 inner fits each +# +# Override defaults via env: POOL=... EXPERIMENT=... TEMPLATE=... ./submit_runs.sh ... +set -euo pipefail + +POOL="${POOL:-CPU-20-LP}" +EXPERIMENT="${EXPERIMENT:-semisynthetic_study}" +TEMPLATE="${TEMPLATE:-experiments/semisynthetic_simulation/job_config_template.yaml}" +GENERATED_DIR="${GENERATED_DIR:-experiments/semisynthetic_simulation/generated_job_configs}" + +N_RUNS=1 +INNER_RUNS=2 +SAMPLE_FRACTION="" + +while [[ $# -gt 0 ]]; do + case "$1" in + -n) N_RUNS="$2"; shift 2 ;; + -k) INNER_RUNS="$2"; shift 2 ;; + -f) SAMPLE_FRACTION="$2"; shift 2 ;; + -h|--help) + grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + *) echo "Unknown arg: $1" >&2; exit 1 ;; + esac +done + +if [[ ! -f "$TEMPLATE" ]]; then + echo "Template not found: $TEMPLATE" >&2; exit 1 +fi +mkdir -p "$GENERATED_DIR" + +for ((i = 1; i <= N_RUNS; i++)); do + RUN=$(printf "run_%02d" "$i") + CFG="$GENERATED_DIR/$RUN.yaml" + sed "s|__RUN__|$RUN|g" "$TEMPLATE" > "$CFG" + + BASH_ARGS="--run-id $RUN --inner-runs $INNER_RUNS" + [[ -n "$SAMPLE_FRACTION" ]] && BASH_ARGS="$BASH_ARGS --sample-fraction $SAMPLE_FRACTION" + + echo "Submitting $RUN (inner_runs=$INNER_RUNS, sample_fraction=${SAMPLE_FRACTION:-none})" + python -m corebehrt.azure job run_semisynthetic_study "$POOL" \ + -e "$EXPERIMENT" \ + -c "$CFG" \ + --bash-args "$BASH_ARGS" +done diff --git a/experiments/semisynthetic_simulation/docs/study.md b/experiments/semisynthetic_simulation/docs/study.md new file mode 100644 index 00000000..649d3cc0 --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/study.md @@ -0,0 +1,83 @@ +# Running the Multi-Run Study on Azure + +A "study" runs **N independent outer runs** (each = one semi-synthetic simulation +with its own seed) and, within each, **K inner reshuffle fits** (re-drawn CV folds +for variance estimation): + +``` +for each outer run (run_01 … run_NN): # one parallel Azure job each + simulate → select_cohort → prepare # Stage 1, once + for k in 1 … K: # Stage 2, K times + finetune → calibrate → estimate # folds reshuffled each time +``` + +Each Azure job is **one outer run**; `submit_runs.sh` submits the N jobs. + +## What to set (once) + +Edit `experiments/semisynthetic_simulation/job_config_template.yaml` — the datastore +paths for `meds`, `features`, `tokenized`, `pretrain_model`. Leave `results` ending in +`__RUN__` (the submit script fills in `run_01`, `run_02`, …). + +The simulated outcomes (and effect sizes) live in +`experiments/semisynthetic_simulation/base_configs/simulate.yaml`: +- `OUTCOME_NULL` → `delta: 0.0` (no effect) +- `OUTCOME_MEDIUM` → `delta: 0.5` (~6pp risk difference at a ~12% baseline) +- For the full study add `OUTCOME_LARGE` with `delta: 1.0`, and add + `OUTCOME_LARGE.csv` to `outcome_files` in `base_configs/prepare.yaml`. + +## What to run + +### 1. Smoke test (do this first) + +Two outer runs, two inner fits each, 10% of patients sampled: + +```bash +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 2 -k 2 -f 0.1 +``` + +Check one job's `estimate/bert/estimate_results.csv` against the true effects in +`simulated_outcomes/simulation_stats.csv` (NULL ≈ 0, MEDIUM ≈ 0.06 risk difference). + +### 2. Full study (after the smoke test looks right) + +Ten outer runs, ten inner fits each, full population (no `-f`): + +```bash +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 10 -k 10 +``` + +### Flags + +| flag | meaning | default | +|------|---------|---------| +| `-n` | number of outer runs (parallel jobs) | 1 | +| `-k` | inner reshuffle fits per run | 2 | +| `-f` | sample this fraction of patients (smoke tests) | none (full) | + +Override the compute/experiment via env: `POOL=CPU-20-LP EXPERIMENT=my_exp ./submit_runs.sh ...` + +## What a single job runs (equivalent manual command) + +```bash +python -m corebehrt.azure job run_semisynthetic_study CPU-20-LP \ + -e semisynthetic_study \ + -c experiments/semisynthetic_simulation/generated_job_configs/run_01.yaml \ + --bash-args "--run-id run_01 --inner-runs 2 --sample-fraction 0.1" +``` + +## Outputs (per outer run) + +``` +/run_NN/ +├── simulated_outcomes/ # exposure.csv, OUTCOME_*.csv, counterfactuals.csv, ite.csv, stats, figs +├── cohort/ # selected cohort + folds +├── prepared_data/ # tokenized finetune data +├── _configs/ # exact per-step configs used (for reproducibility) +└── reshuffles/ + ├── k_01/{models/bert, estimate/bert}/ + └── k_02/... +``` + +Seeds are independent across runs (`seed = base_seed + run_number`), so the N outer +runs are genuinely independent replicates. diff --git a/experiments/semisynthetic_simulation/job_config_template.yaml b/experiments/semisynthetic_simulation/job_config_template.yaml new file mode 100644 index 00000000..8f420bd6 --- /dev/null +++ b/experiments/semisynthetic_simulation/job_config_template.yaml @@ -0,0 +1,13 @@ +# Azure job config for one outer run of the semi-synthetic study. +# submit_runs.sh clones this per run, replacing __RUN__ with run_01, run_02, ... +# Edit the paths below to your datastore locations, then run submit_runs.sh. + +paths: + ## INPUTS (shared across every run) + meds: "researcher_data:AKK/shared/MEDS/TRACE/v01/data" + features: "researcher_data:AKK/shared/features/trace/v01/features" + tokenized: "researcher_data:AKK/shared/features/trace/v01/tokenized" + pretrain_model: "researcher_data:AKK/shared/pretrain/models/trace/small/len_512/v01" + + ## OUTPUT (one dir per outer run; __RUN__ is filled in by submit_runs.sh) + results: "researcher_data:AKK/experiments/trace/simulation/semisynthetic/smoketest/__RUN__" diff --git a/experiments/semisynthetic_simulation/python_scripts/__init__.py b/experiments/semisynthetic_simulation/python_scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/experiments/semisynthetic_simulation/python_scripts/run_study.py b/experiments/semisynthetic_simulation/python_scripts/run_study.py new file mode 100644 index 00000000..b497f566 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/run_study.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +""" +Runner for the semi-synthetic simulation study. + +Structure (mirrors the resampling study, but with the semi-synthetic simulator): + + for each OUTER run (independent simulation, own seed): + Stage 1 (once): simulate -> select_cohort -> prepare + Stage 2 (K times): finetune -> calibrate -> estimate (folds reshuffled each time) + +Each Azure job runs a single outer run (pass --run-id run_NN); the outer loop is +the set of parallel jobs submitted by bash_scripts/submit_runs.sh. Run locally with +--n-runs for several outer runs in one process. +""" + +import argparse +import logging +import re +from pathlib import Path + +import yaml + +from corebehrt.main_causal.simulate_semisynthetic import main_simulate +from corebehrt.main_causal.select_cohort_full import main as main_select_cohort +from corebehrt.main_causal.prepare_ft_exp_y import main as main_prepare +from corebehrt.main_causal.finetune_exp_y import main_finetune +from corebehrt.main_causal.calibrate_exp_y import main_calibrate +from corebehrt.main_causal.estimate import main_estimate + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger("semisynthetic_study") + +DEFAULT_BASE_CONFIGS = Path(__file__).resolve().parent.parent / "base_configs" + + +def fill_config(base_path: Path, replacements: dict, out_path: Path, edit=None) -> str: + """Substitute {{...}} placeholders into a base config, optionally edit, and save.""" + text = base_path.read_text() + for key, value in replacements.items(): + text = text.replace(key, value) + config = yaml.safe_load(text) + if edit is not None: + edit(config) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(yaml.dump(config, sort_keys=False)) + return str(out_path) + + +def run_outer(args, run_id: str, seed: int): + """Run one outer simulation followed by K inner reshuffle fits.""" + run_dir = Path(args.experiment_dir) / run_id + config_dir = run_dir / "_configs" + base = Path(args.base_configs_dir) + + shared = { + "{{MEDS}}": args.meds, + "{{FEATURES}}": args.features, + "{{TOKENIZED}}": args.tokenized, + "{{PRETRAIN_MODEL}}": args.pretrain_model, + "{{RUN_DIR}}": str(run_dir), + } + + logger.info("=" * 70) + logger.info(f"OUTER RUN {run_id} (seed={seed})") + logger.info("=" * 70) + + # ---- Stage 1: simulate -> select_cohort -> prepare (once) ---- + def set_simulation(config): + config["seed"] = seed + if args.sample_fraction is not None or args.sample_size is not None: + config["sampling"] = { + "enabled": True, + "fraction": args.sample_fraction, + "size": args.sample_size, + } + + main_simulate( + fill_config( + base / "simulate.yaml", shared, config_dir / "simulate.yaml", set_simulation + ) + ) + main_select_cohort( + fill_config( + base / "select_cohort.yaml", shared, config_dir / "select_cohort.yaml" + ) + ) + main_prepare( + fill_config(base / "prepare.yaml", shared, config_dir / "prepare.yaml") + ) + + # ---- Stage 2: K inner reshuffle fits ---- + for k in range(1, args.inner_runs + 1): + inner_id = f"k_{k:02d}" + inner_dir = run_dir / "reshuffles" / inner_id + repl = {**shared, "{{INNER_DIR}}": str(inner_dir)} + logger.info(f"--- {run_id} / inner fit {k}/{args.inner_runs} ({inner_id}) ---") + + def enable_reshuffle(config): + config.setdefault("data", {})["reshuffle"] = True + + main_finetune( + fill_config( + base / "finetune.yaml", + repl, + config_dir / f"finetune_{inner_id}.yaml", + enable_reshuffle, + ) + ) + main_calibrate( + fill_config( + base / "calibrate.yaml", repl, config_dir / f"calibrate_{inner_id}.yaml" + ) + ) + main_estimate( + fill_config( + base / "estimate.yaml", repl, config_dir / f"estimate_{inner_id}.yaml" + ) + ) + + logger.info(f"OUTER RUN {run_id} complete") + + +def parse_arguments(argv=None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the semi-synthetic simulation study" + ) + parser.add_argument("--meds", required=True) + parser.add_argument("--features", required=True) + parser.add_argument("--tokenized", required=True) + parser.add_argument("--pretrain-model", dest="pretrain_model", required=True) + parser.add_argument("--experiment-dir", dest="experiment_dir", required=True) + parser.add_argument( + "--base-configs-dir", dest="base_configs_dir", default=str(DEFAULT_BASE_CONFIGS) + ) + + parser.add_argument( + "--run-id", + dest="run_id", + default=None, + help="Single outer run id, e.g. run_03 (seed = base-seed + 3). One Azure job = one run-id.", + ) + parser.add_argument( + "--n-runs", + dest="n_runs", + type=int, + default=1, + help="Number of outer runs in this process (local use; ignored if --run-id is given).", + ) + parser.add_argument( + "--inner-runs", + "-k", + dest="inner_runs", + type=int, + default=2, + help="Inner reshuffle fits per outer run (variance estimation).", + ) + parser.add_argument("--base-seed", dest="base_seed", type=int, default=42) + + parser.add_argument( + "--sample-fraction", + dest="sample_fraction", + type=float, + default=None, + help="Sample this fraction of patients per run (smoke tests).", + ) + parser.add_argument("--sample-size", dest="sample_size", type=int, default=None) + + args = parser.parse_args(argv) + if args.sample_fraction is not None and args.sample_size is not None: + parser.error("Specify at most one of --sample-fraction / --sample-size") + return args + + +def run_number_from_id(run_id: str) -> int: + match = re.match(r"run_(\d+)", run_id) + return int(match.group(1)) if match else 0 + + +def main(argv=None): + args = parse_arguments(argv) + if args.run_id: + run_outer(args, args.run_id, args.base_seed + run_number_from_id(args.run_id)) + else: + for run_number in range(1, args.n_runs + 1): + run_id = f"run_{run_number:02d}" + run_outer(args, run_id, args.base_seed + run_number) + + +if __name__ == "__main__": + main() From 5d187a5e62ac6ae2e5cce271f8935a19b3b4d919 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Fri, 26 Jun 2026 16:43:10 +0200 Subject: [PATCH 11/20] fail fast on unknown feature names in simulation config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A coefficient or interaction naming a feature the extractor doesn't produce was silently ignored (`if name in features_df.columns`), so a typo in the outcome config would silently zero-weight that term and change the true data-generating process with no warning — dangerous for a simulation study where the config defines ground truth. - add ORACLE_FEATURE_NAMES (canonical feature set) in oracle_features - validate all outcome coefficients/interactions/heterogeneous terms at simulator init, raising ValueError listing unknown names - tests: drift guard (extractor output == ORACLE_FEATURE_NAMES) + raises on unknown coefficient/interaction names Addresses CodeRabbit PR #173 finding (semisynthetic_simulator.py:343). Co-Authored-By: Claude Opus 4.8 --- .../modules/simulation/oracle_features.py | 15 +++++++ .../simulation/semisynthetic_simulator.py | 24 ++++++++++ .../test_simulation/test_semisynthetic.py | 44 +++++++++++++++++++ 3 files changed, 83 insertions(+) diff --git a/corebehrt/modules/simulation/oracle_features.py b/corebehrt/modules/simulation/oracle_features.py index 0e19f434..0006bf97 100644 --- a/corebehrt/modules/simulation/oracle_features.py +++ b/corebehrt/modules/simulation/oracle_features.py @@ -10,6 +10,21 @@ logger = logging.getLogger("oracle_features") +# Canonical names of every feature produced by extract_oracle_features. +# Keep in sync with the `features` dict built below (guarded by a test). +ORACLE_FEATURE_NAMES = ( + "recent_event_count", + "disease_burden", + "medication_count", + "utilization_intensity", + "age", + "chronic_disease_count", + "code_diversity", + "event_recency", + "recent_burst_ratio", + "sequence_motif_count", +) + def extract_oracle_features( history_df: pd.DataFrame, diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 22f00b6c..4b519ce0 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -39,6 +39,7 @@ TreatmentEffectConfig, ) from corebehrt.modules.simulation.oracle_features import ( + ORACLE_FEATURE_NAMES, extract_oracle_features, standardize_features, ) @@ -66,6 +67,29 @@ def __init__(self, config: SemiSyntheticSimulationConfig): self.rng = np.random.default_rng(config.seed) self._global_means = None self._global_stds = None + self._validate_feature_references() + + def _validate_feature_references(self): + """Fail fast if any outcome config references an unknown feature. + + A misspelled coefficient name would otherwise be silently dropped, + changing the true data-generating process without warning. + """ + known = set(ORACLE_FEATURE_NAMES) + unknown = set() + for outcome_cfg in self.config.outcomes.values(): + outcome_model = outcome_cfg.outcome_model + unknown |= set(outcome_model.coefficients) - known + for interaction in outcome_model.interactions: + unknown |= {f for f in interaction.get("features", []) if f} - known + unknown |= ( + set(outcome_cfg.treatment_effect.heterogeneous_coefficients) - known + ) + if unknown: + raise ValueError( + f"Unknown feature name(s) in simulation config: {sorted(unknown)}. " + f"Available oracle features: {sorted(known)}." + ) def compute_global_feature_stats(self, shard_loader): """Pass 1: compute global mean/std across all shards for standardization.""" diff --git a/tests/test_modules/test_simulation/test_semisynthetic.py b/tests/test_modules/test_simulation/test_semisynthetic.py index 6af7b96d..96c27395 100644 --- a/tests/test_modules/test_simulation/test_semisynthetic.py +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -14,6 +14,10 @@ SemiSyntheticSimulationConfig, TreatmentEffectConfig, ) +from corebehrt.modules.simulation.oracle_features import ( + ORACLE_FEATURE_NAMES, + extract_oracle_features, +) from corebehrt.modules.simulation.semisynthetic_simulator import ( ASSIGNED_INDEX_DATE_COL, SemiSyntheticCausalSimulator, @@ -235,5 +239,45 @@ def test_exposure_from_data(self): self.assertEqual(exposed_pids_from_sim, expected_in_remaining) +class TestFeatureNameValidation(unittest.TestCase): + def test_extractor_produces_canonical_names(self): + """ORACLE_FEATURE_NAMES must stay in sync with the extractor output.""" + shard = _make_test_shard(n_patients=10, n_exposed=4) + index_dates = pd.Series( + {pid: pd.Timestamp("2021-01-01") for pid in shard[PID_COL].unique()} + ) + pids = np.array(sorted(index_dates.index)) + features_df = extract_oracle_features(shard, pids, index_dates, FeatureConfig()) + self.assertEqual(set(features_df.columns), set(ORACLE_FEATURE_NAMES)) + + def test_unknown_coefficient_raises(self): + """A misspelled coefficient name must fail fast, not be silently dropped.""" + import tempfile + + config = _make_config(tempfile.mkdtemp()) + config.outcomes["OUTCOME_test"].outcome_model.coefficients = { + "disease_burden": 0.3, + "diseaze_burden": 0.5, # typo + } + with self.assertRaises(ValueError): + SemiSyntheticCausalSimulator(config) + + def test_unknown_interaction_feature_raises(self): + import tempfile + + config = _make_config(tempfile.mkdtemp()) + config.outcomes["OUTCOME_test"].outcome_model.interactions = [ + {"features": ["age", "not_a_feature"], "coefficient": 0.1} + ] + with self.assertRaises(ValueError): + SemiSyntheticCausalSimulator(config) + + def test_valid_config_does_not_raise(self): + import tempfile + + # _make_config uses only valid feature names → must construct cleanly. + SemiSyntheticCausalSimulator(_make_config(tempfile.mkdtemp())) + + if __name__ == "__main__": unittest.main() From 3ec9ca0d38bd7407aa5a61cc2c8ec66f774ccad9 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Fri, 26 Jun 2026 16:47:49 +0200 Subject: [PATCH 12/20] make coverage badge step non-blocking in CI The gist badge update fails with 401 (GIST_SECRET expired/unauthorized), which turned the Test/Docstring Coverage checks red. Every other step in these jobs is already continue-on-error; the badge step was missing it. Add continue-on-error to the badge step so a failed badge update no longer fails the check (the badge itself still updates once the gist token is refreshed). Co-Authored-By: Claude Opus 4.8 --- .github/workflows/docstring_coverage.yml | 1 + .github/workflows/utest_coverage_badge.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/docstring_coverage.yml b/.github/workflows/docstring_coverage.yml index ea8997e3..e19f2698 100644 --- a/.github/workflows/docstring_coverage.yml +++ b/.github/workflows/docstring_coverage.yml @@ -36,6 +36,7 @@ jobs: - name: Create Dynamic Coverage Badge if: always() # Always try to create badge + continue-on-error: true # Badge update (gist auth) must not fail the check uses: schneegans/dynamic-badges-action@v1.7.0 with: auth: ${{ secrets.GIST_SECRET }} # Your secret token with the gist scope diff --git a/.github/workflows/utest_coverage_badge.yml b/.github/workflows/utest_coverage_badge.yml index 0845c6f1..332e8262 100644 --- a/.github/workflows/utest_coverage_badge.yml +++ b/.github/workflows/utest_coverage_badge.yml @@ -48,6 +48,7 @@ jobs: - name: Create Dynamic Coverage Badge if: always() # Always try to create badge + continue-on-error: true # Badge update (gist auth) must not fail the check uses: schneegans/dynamic-badges-action@v1.7.0 with: auth: ${{ secrets.GIST_SECRET }} From 6b79409dc9acc32e0fc7bf26be1fdb4eacd9a531 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sat, 27 Jun 2026 11:11:59 +0200 Subject: [PATCH 13/20] read index dates from a cohort artifact instead of raw MEDS The simulator required an assigned_index_date column in the MEDS shards, which only the pre-processed example data has. Real MEDS is raw and has no such column (KeyError on TRACE). Index dates are an analysis artifact (cohort selection output), so add paths.index_dates: when set, load per-patient index dates from a cohort dir / index_dates.csv (subject_id, time) and use them; otherwise fall back to the assigned_index_date column. Treatment is still taken from the data (presence of the exposure code). Validated locally against example data via a constructed index_dates.csv: recovers the true effect (ATE 0.062 vs true 0.063), all outputs produced. Co-Authored-By: Claude Opus 4.8 --- .../simulation/config_semisynthetic.py | 1 + .../simulation/semisynthetic_simulator.py | 44 +++++++++++++++---- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py index f6b5904b..1c97591c 100644 --- a/corebehrt/modules/simulation/config_semisynthetic.py +++ b/corebehrt/modules/simulation/config_semisynthetic.py @@ -10,6 +10,7 @@ class PathsConfig: splits: List[str] outcomes: str cohort: str = None + index_dates: str = None @dataclass diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 4b519ce0..97aaa145 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -24,6 +24,7 @@ COUNTERFACTUALS_FILE, INDEX_DATE_MATCHING_FILE, ) +from corebehrt.constants.paths import INDEX_DATES_FILE from corebehrt.constants.data import ( ABSPOS_COL, CONCEPT_COL, @@ -68,6 +69,27 @@ def __init__(self, config: SemiSyntheticSimulationConfig): self._global_means = None self._global_stds = None self._validate_feature_references() + self._index_dates = self._load_index_dates() + + def _load_index_dates(self): + """Load per-patient index dates from a cohort artifact, if configured. + + Index dates are an analysis artifact (produced by cohort selection), + not raw MEDS. When ``paths.index_dates`` points to a cohort dir or an + ``index_dates.csv``, use it; otherwise fall back to an + ``assigned_index_date`` column in the data (e.g. the example data). + """ + path = self.config.paths.index_dates + if not path: + return None + if os.path.isdir(path): + path = join(path, INDEX_DATES_FILE) + index_dates = pd.read_csv( + path, usecols=[PID_COL, TIMESTAMP_COL], parse_dates=[TIMESTAMP_COL] + ) + index_dates = index_dates.set_index(PID_COL)[TIMESTAMP_COL] + logger.info(f"Loaded {len(index_dates)} index dates from {path}") + return index_dates def _validate_feature_references(self): """Fail fast if any outcome config references an unknown feature. @@ -219,21 +241,27 @@ def _extract_treatment_and_index_dates( self, shard_df: pd.DataFrame ) -> Tuple[np.ndarray, np.ndarray, pd.Series]: """Identify exposed/control patients and their index dates.""" - # Drop patients without an assigned index date - valid = shard_df.dropna(subset=[ASSIGNED_INDEX_DATE_COL]) - if valid.empty: + if self._index_dates is not None: + # Index dates come from the cohort artifact; restrict to this shard. + shard_pids = shard_df[PID_COL].unique() + index_dates = self._index_dates[self._index_dates.index.isin(shard_pids)] + else: + # Fall back to an assigned_index_date column in the data. + valid = shard_df.dropna(subset=[ASSIGNED_INDEX_DATE_COL]) + index_dates = valid.groupby(PID_COL)[ASSIGNED_INDEX_DATE_COL].first() + + if index_dates.empty: return ( np.array([]), np.array([], dtype=bool), pd.Series(dtype="datetime64[ns]"), ) - # Per-patient index dates - index_dates = valid.groupby(PID_COL)[ASSIGNED_INDEX_DATE_COL].first() - - # Patients with at least one exposure event + # Patients with at least one exposure event (treatment from the data) exposed_pids = set( - valid.loc[valid[CONCEPT_COL] == self.config.exposure_code, PID_COL].unique() + shard_df.loc[ + shard_df[CONCEPT_COL] == self.config.exposure_code, PID_COL + ].unique() ) pids = index_dates.index.values From b358ef6432fa5a8a4191da10a7519eeb0bc0c272 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sat, 27 Jun 2026 11:23:01 +0200 Subject: [PATCH 14/20] add study summarizer (bias / SD / SE-calibration / coverage) Reads every estimate_results.csv under a study dir (each already carries the point estimate, CI, and appended true_effect), groups by model x method x outcome, and reports bias, empirical SD, mean estimated SE, SE-calibration (SD_emp/mean_SE), and 95% CI coverage. - single estimate per group (Phase 1): bias + coverage meaningful, SD/calibration undefined (NaN) - many estimates (bootstrap refits / outcome redraws): full table Validated on synthetic estimate_results fixtures. Co-Authored-By: Claude Opus 4.8 --- .../python_scripts/summarize.py | 114 ++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 experiments/semisynthetic_simulation/python_scripts/summarize.py diff --git a/experiments/semisynthetic_simulation/python_scripts/summarize.py b/experiments/semisynthetic_simulation/python_scripts/summarize.py new file mode 100644 index 00000000..7f777ad0 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/summarize.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +""" +Summarize a semi-synthetic study into an estimator-performance table. + +Reads every ``estimate_results.csv`` under a study directory (each already +carries the point estimate, its CI, and the appended ``true_effect``), groups +by model x method x outcome, and reports bias, empirical SD, mean estimated SE, +SE-calibration, and 95% CI coverage. + +- One estimate per group (Phase 1, single run): bias + covered are meaningful; + SD/calibration are undefined (need >1). +- Many estimates per group (bootstrap refits / outcome redraws): full table. + +Usage: + python -m experiments.semisynthetic_simulation.python_scripts.summarize \ + --study-dir [--out ] +""" + +import argparse +import re +from pathlib import Path + +import numpy as np +import pandas as pd + +from corebehrt.constants.causal.data import EffectColumns as E +from corebehrt.constants.causal.data import OUTCOME +from corebehrt.constants.causal.paths import ESTIMATE_RESULTS_FILE + +MODEL_NAMES = ("bert", "baseline") + + +def _tag_from_path(path: Path) -> dict: + """Infer model / outer-run / inner-refit ids from a result file's path.""" + parts = [p.lower() for p in path.parts] + model = next((m for m in MODEL_NAMES if m in parts), "unknown") + run_match = re.search(r"run_\d+", str(path)) + inner_match = re.search(r"k_\d+", str(path)) + return { + "model": model, + "run_id": run_match.group(0) if run_match else "run_01", + "inner_id": inner_match.group(0) if inner_match else "k_01", + } + + +def load_results(study_dir: Path) -> pd.DataFrame: + """Load and tag every estimate_results.csv under study_dir.""" + files = sorted(study_dir.rglob(ESTIMATE_RESULTS_FILE)) + if not files: + raise FileNotFoundError(f"No {ESTIMATE_RESULTS_FILE} found under {study_dir}") + frames = [] + for path in files: + df = pd.read_csv(path) + for key, value in _tag_from_path(path).items(): + df[key] = value + frames.append(df) + return pd.concat(frames, ignore_index=True) + + +def summarize(results: pd.DataFrame) -> pd.DataFrame: + """Aggregate estimates into a per (model, method, outcome) performance table.""" + rows = [] + for (model, method, outcome), group in results.groupby( + ["model", E.method, OUTCOME] + ): + effects = group[E.effect] + ses = group[E.std_err] + true = group[E.true_effect].iloc[0] + covered = (group[E.CI95_lower] <= true) & (group[E.CI95_upper] >= true) + n = len(group) + sd_emp = effects.std(ddof=1) if n > 1 else np.nan + mean_se = ses.mean() + rows.append( + { + "model": model, + "method": method, + "outcome": outcome, + "n": n, + "true_effect": true, + "mean_effect": effects.mean(), + "bias": effects.mean() - true, + "sd_emp": sd_emp, + "mean_se": mean_se, + "se_calibration": (sd_emp / mean_se) + if n > 1 and mean_se > 0 + else np.nan, + "coverage": covered.mean(), + } + ) + return pd.DataFrame(rows).sort_values(["outcome", "model", "method"]) + + +def main(): + parser = argparse.ArgumentParser(description="Summarize a semi-synthetic study") + parser.add_argument("--study-dir", required=True) + parser.add_argument( + "--out", default=None, help="Output CSV (default: /summary.csv)" + ) + args = parser.parse_args() + + study_dir = Path(args.study_dir) + results = load_results(study_dir) + table = summarize(results) + + out = Path(args.out) if args.out else study_dir / "summary.csv" + table.to_csv(out, index=False) + + pd.set_option("display.float_format", lambda v: f"{v:.4f}") + print(table.to_string(index=False)) + print(f"\nSaved summary to {out}") + + +if __name__ == "__main__": + main() From 55470e847d99bf281124d0d4acf0dcdcf8ea8931 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sat, 27 Jun 2026 11:57:58 +0200 Subject: [PATCH 15/20] restructure study to fixed-cohort flow with bootstrap refits + baseline Reworks the semi-synthetic study per the design discussion: - Run on a FIXED, shared pre-existing cohort (e.g. a diabetes cohort). Drop per-run select_cohort and per-run sampling; the cohort's index_dates.csv + cohort_config.yaml define the fixed design, treatment is the real exposure. Add `cohort` as a component/runner input. - noise_scale=0 and three effect sizes (null/medium/large) -> a single fixed true effect per scenario. - Inner refits (-k) are the SE mechanism: k=1 single fit (estimate's internal-bootstrap CI), k>1 each refit trains on a bootstrap resample (reshuffle + top-level bootstrap), summarizer combines the K. - Add the CatBoost baseline path alongside BERT (method + baseline). - Update component inputs, job template (cohort path), submit_runs.sh (drop sampling, default k=1, forward extra runner flags), and docs. Config generation, outcome-file match, component import, bash syntax and lint all validated locally; downstream steps need Azure (pretrain + CausalEstimate) so the Phase-1 run is the first end-to-end check. Co-Authored-By: Claude Opus 4.8 --- .../components/run_semisynthetic_study.py | 3 + .../base_configs/calibrate_baseline.yaml | 8 + .../base_configs/estimate.yaml | 2 +- .../base_configs/estimate_baseline.yaml | 28 +++ .../base_configs/prepare.yaml | 6 +- .../base_configs/select_cohort.yaml | 34 ---- .../base_configs/simulate.yaml | 43 +++-- .../base_configs/train_baseline.yaml | 22 +++ .../bash_scripts/submit_runs.sh | 26 ++- .../semisynthetic_simulation/docs/study.md | 98 +++++----- .../job_config_template.yaml | 4 +- .../python_scripts/run_study.py | 174 +++++++++++------- 12 files changed, 260 insertions(+), 188 deletions(-) create mode 100644 experiments/semisynthetic_simulation/base_configs/calibrate_baseline.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml delete mode 100644 experiments/semisynthetic_simulation/base_configs/select_cohort.yaml create mode 100644 experiments/semisynthetic_simulation/base_configs/train_baseline.yaml diff --git a/corebehrt/azure/components/run_semisynthetic_study.py b/corebehrt/azure/components/run_semisynthetic_study.py index c2d62b76..2190a1b1 100644 --- a/corebehrt/azure/components/run_semisynthetic_study.py +++ b/corebehrt/azure/components/run_semisynthetic_study.py @@ -9,6 +9,7 @@ "features": {"type": "uri_folder"}, "tokenized": {"type": "uri_folder"}, "pretrain_model": {"type": "uri_folder"}, + "cohort": {"type": "uri_folder"}, } OUTPUTS = { @@ -31,6 +32,8 @@ def main_run_study(config_path): cfg.paths.tokenized, "--pretrain-model", cfg.paths.pretrain_model, + "--cohort", + cfg.paths.cohort, "--experiment-dir", cfg.paths.results, ] diff --git a/experiments/semisynthetic_simulation/base_configs/calibrate_baseline.yaml b/experiments/semisynthetic_simulation/base_configs/calibrate_baseline.yaml new file mode 100644 index 00000000..0d47dd23 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/calibrate_baseline.yaml @@ -0,0 +1,8 @@ +# Stage 2b (baseline): calibrate baseline predictions. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + finetune_model: "{{INNER_DIR}}/models/baseline" + calibrated_predictions: "{{INNER_DIR}}/models/baseline/calibrated" diff --git a/experiments/semisynthetic_simulation/base_configs/estimate.yaml b/experiments/semisynthetic_simulation/base_configs/estimate.yaml index c250bf4b..37359709 100644 --- a/experiments/semisynthetic_simulation/base_configs/estimate.yaml +++ b/experiments/semisynthetic_simulation/base_configs/estimate.yaml @@ -14,7 +14,7 @@ paths: estimator: methods: ["IPW", "TMLE", "TMLE_TH"] effect_type: "ATE" - n_bootstrap: 30 + n_bootstrap: 100 # set by the runner: 0 for bootstrap-refit studies, >0 for single-fit CIs clip_percentile: 0.99 plot: diff --git a/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml new file mode 100644 index 00000000..61736c15 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml @@ -0,0 +1,28 @@ +# Stage 2c (baseline): estimate causal effects from calibrated baseline predictions. +# n_bootstrap is set by the runner (0 for bootstrap-refit studies, >0 for single-fit CIs). +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + calibrated_predictions: "{{INNER_DIR}}/models/baseline/calibrated/" + counterfactual_outcomes: "{{RUN_DIR}}/simulated_outcomes" + estimate: "{{INNER_DIR}}/estimate/baseline" + +estimator: + methods: ["IPW", "TMLE", "TMLE_TH"] + effect_type: "ATE" + n_bootstrap: 100 + clip_percentile: 0.99 + +plot: + contingency_table: + max_outcomes_per_figure: 10 + max_number_of_figures: 10 + effect_size: + max_outcomes_per_figure: 10 + max_number_of_figures: 10 + plot_individual_effects: false + adjustment: + max_outcomes_per_figure: 8 + max_number_of_figures: 10 diff --git a/experiments/semisynthetic_simulation/base_configs/prepare.yaml b/experiments/semisynthetic_simulation/base_configs/prepare.yaml index 61e95873..c4de8774 100644 --- a/experiments/semisynthetic_simulation/base_configs/prepare.yaml +++ b/experiments/semisynthetic_simulation/base_configs/prepare.yaml @@ -1,4 +1,5 @@ -# Stage 1c: prepare finetuning data + cross-validation folds. +# Stage 1b: prepare finetuning data from the FIXED cohort + simulated outcomes. +# Uses the existing cohort ({{COHORT}}) directly — no per-run cohort selection. logging: level: INFO path: ./outputs/logs/sim_study @@ -7,12 +8,13 @@ paths: ## INPUTS features: "{{FEATURES}}" tokenized: "{{TOKENIZED}}" - cohort: "{{RUN_DIR}}/cohort" + cohort: "{{COHORT}}" # pre-existing cohort (pids + cohort_config.yaml) outcomes: "{{RUN_DIR}}/simulated_outcomes" outcome_files: - OUTCOME_NULL.csv - OUTCOME_MEDIUM.csv + - OUTCOME_LARGE.csv exposures: "{{RUN_DIR}}/simulated_outcomes" exposure: exposure.csv diff --git a/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml b/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml deleted file mode 100644 index 14bfae6f..00000000 --- a/experiments/semisynthetic_simulation/base_configs/select_cohort.yaml +++ /dev/null @@ -1,34 +0,0 @@ -# Stage 1b: select the analysis cohort from the simulated exposure. -logging: - level: INFO - path: ./outputs/logs/sim_study - -paths: - ### Inputs - features: "{{FEATURES}}/" - meds: "{{MEDS}}" - splits: [tuning] - exposures: "{{RUN_DIR}}/simulated_outcomes/" - exposure: exposure.csv - criteria_config: ./corebehrt/configs/causal/select_cohort_full/definitions.yaml - - ### Outputs - cohort: "{{RUN_DIR}}/cohort/" - -time_windows: - data_end: - year: 2025 - month: 01 - day: 01 - data_start: - year: 1920 - month: 1 - day: 1 - min_follow_up: - days: 1 - min_lookback: - days: 365 - -cv_folds: 2 -val_ratio: 0.1 -test_ratio: 0 diff --git a/experiments/semisynthetic_simulation/base_configs/simulate.yaml b/experiments/semisynthetic_simulation/base_configs/simulate.yaml index 267c4b75..dfdc1ce7 100644 --- a/experiments/semisynthetic_simulation/base_configs/simulate.yaml +++ b/experiments/semisynthetic_simulation/base_configs/simulate.yaml @@ -1,17 +1,21 @@ -# Stage 1a: semi-synthetic simulation (real treatment, simulated outcome). -# The double-brace placeholders are filled in per run by run_study.py. -# `seed` and the `sampling` block are injected by the runner. +# Stage 1a: semi-synthetic simulation on a FIXED, pre-existing cohort. +# Index dates + cohort membership come from the cohort artifact ({{COHORT}}); +# treatment is the real exposure (presence of `exposure_code` in MEDS). +# Only the outcome is simulated. noise_scale=0 -> a single fixed true effect. logging: level: INFO path: ./outputs/logs/sim_study paths: data: "{{MEDS}}" - splits: ["tuning"] + # Splits must cover the patients in the cohort (filtered by index_dates). + splits: ["tuning", "train"] outcomes: "{{RUN_DIR}}/simulated_outcomes" - cohort: "{{RUN_DIR}}/cohort" + index_dates: "{{COHORT}}" # dir or index_dates.csv; index dates + membership min_num_codes: 3 +# Code marking a treated patient at index. Set to whatever marks exposure +# in your MEDS (e.g. the semaglutide exposure code), not necessarily "EXPOSURE". exposure_code: "EXPOSURE" features: @@ -25,22 +29,21 @@ features: burst_window_days: 30 motif_window_days: 30 -# Two outcomes sharing identical confounding, differing only in the treatment -# effect: a null effect (delta=0) and a medium effect (delta=0.5, ~6pp risk -# difference at a ~12% baseline). Add OUTCOME_LARGE (delta=1.0) for the full study. +# Three effect sizes sharing identical confounding; only delta differs. +# delta is on the logit scale: ~0.5 -> medium (~6pp RD @12% base), 1.0 -> large. outcomes: OUTCOME_NULL: outcome_model: run_in_days: 1 beta_0: -2.0 - coefficients: + coefficients: &coeffs recent_event_count: 0.3 disease_burden: 0.2 medication_count: 0.15 age: 0.4 chronic_disease_count: 0.25 event_recency: -0.15 - noise_scale: 0.1 + noise_scale: 0.0 treatment_effect: mode: constant delta: 0.0 @@ -49,14 +52,18 @@ outcomes: outcome_model: run_in_days: 1 beta_0: -2.0 - coefficients: - recent_event_count: 0.3 - disease_burden: 0.2 - medication_count: 0.15 - age: 0.4 - chronic_disease_count: 0.25 - event_recency: -0.15 - noise_scale: 0.1 + coefficients: *coeffs + noise_scale: 0.0 treatment_effect: mode: constant delta: 0.5 + + OUTCOME_LARGE: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + coefficients: *coeffs + noise_scale: 0.0 + treatment_effect: + mode: constant + delta: 1.0 diff --git a/experiments/semisynthetic_simulation/base_configs/train_baseline.yaml b/experiments/semisynthetic_simulation/base_configs/train_baseline.yaml new file mode 100644 index 00000000..31852877 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/train_baseline.yaml @@ -0,0 +1,22 @@ +# Stage 2a (baseline): CatBoost baseline model. Runs on CPU. +logging: + level: INFO + path: ./outputs/logs/sim_study + +paths: + prepared_data: "{{RUN_DIR}}/prepared_data" + model: "{{INNER_DIR}}/models/baseline" + +multihot: false +include_age: true + +catboost: + n_estimators: 100 + l2_leaf_reg: 0.00 + early_stopping_rounds: 10 + +tuning: + tune_hyperparameters: true + n_trials: 10 + inner_val_size: 0.2 + reuse_hyperparameters: true diff --git a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh index d2a3f1ba..f48e2336 100755 --- a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh +++ b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh @@ -1,12 +1,13 @@ #!/usr/bin/env bash -# Submit N independent outer runs of the semi-synthetic study as parallel Azure jobs. -# Each job runs one outer simulation (own seed) + K inner reshuffle fits. +# Submit N outer runs of the semi-synthetic study as parallel Azure jobs, on a +# FIXED shared cohort. Each job = one outer simulation (own seed) + K refits. # # Usage: -# ./submit_runs.sh # 1 run, K=2, no sampling (full population) -# ./submit_runs.sh -n 2 -k 2 -f 0.1 # SMOKE TEST: 2 runs, 2 inner fits, 10% sampled -# ./submit_runs.sh -n 10 -k 10 # full study: 10 runs, 10 inner fits each +# ./submit_runs.sh # Phase 1: 1 run, K=1 (single fit + analytic CI) +# ./submit_runs.sh --baseline-only # Phase 1, baseline only (CPU, fast) +# ./submit_runs.sh -n 5 -k 10 # full: 5 outer runs, 10 bootstrap refits each # +# Anything after the known flags is forwarded to the runner (e.g. --bert-only). # Override defaults via env: POOL=... EXPERIMENT=... TEMPLATE=... ./submit_runs.sh ... set -euo pipefail @@ -16,17 +17,15 @@ TEMPLATE="${TEMPLATE:-experiments/semisynthetic_simulation/job_config_template.y GENERATED_DIR="${GENERATED_DIR:-experiments/semisynthetic_simulation/generated_job_configs}" N_RUNS=1 -INNER_RUNS=2 -SAMPLE_FRACTION="" +INNER_RUNS=1 +EXTRA_ARGS=() while [[ $# -gt 0 ]]; do case "$1" in -n) N_RUNS="$2"; shift 2 ;; -k) INNER_RUNS="$2"; shift 2 ;; - -f) SAMPLE_FRACTION="$2"; shift 2 ;; - -h|--help) - grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; - *) echo "Unknown arg: $1" >&2; exit 1 ;; + -h|--help) grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + *) EXTRA_ARGS+=("$1"); shift ;; esac done @@ -40,10 +39,9 @@ for ((i = 1; i <= N_RUNS; i++)); do CFG="$GENERATED_DIR/$RUN.yaml" sed "s|__RUN__|$RUN|g" "$TEMPLATE" > "$CFG" - BASH_ARGS="--run-id $RUN --inner-runs $INNER_RUNS" - [[ -n "$SAMPLE_FRACTION" ]] && BASH_ARGS="$BASH_ARGS --sample-fraction $SAMPLE_FRACTION" + BASH_ARGS="--run-id $RUN --inner-runs $INNER_RUNS ${EXTRA_ARGS[*]:-}" - echo "Submitting $RUN (inner_runs=$INNER_RUNS, sample_fraction=${SAMPLE_FRACTION:-none})" + echo "Submitting $RUN (inner_runs=$INNER_RUNS) ${EXTRA_ARGS[*]:-}" python -m corebehrt.azure job run_semisynthetic_study "$POOL" \ -e "$EXPERIMENT" \ -c "$CFG" \ diff --git a/experiments/semisynthetic_simulation/docs/study.md b/experiments/semisynthetic_simulation/docs/study.md index 649d3cc0..cac62dd7 100644 --- a/experiments/semisynthetic_simulation/docs/study.md +++ b/experiments/semisynthetic_simulation/docs/study.md @@ -1,83 +1,83 @@ -# Running the Multi-Run Study on Azure +# Running the Study on Azure -A "study" runs **N independent outer runs** (each = one semi-synthetic simulation -with its own seed) and, within each, **K inner reshuffle fits** (re-drawn CV folds -for variance estimation): +The study runs on a **fixed, pre-existing cohort** (a `select_cohort_full` output, +e.g. a diabetes cohort). The cohort's `index_dates.csv` defines the patients and +their index dates; **treatment is the real exposure**. Only the outcome is +simulated, with a single fixed true effect per scenario (`noise_scale = 0`). ``` -for each outer run (run_01 … run_NN): # one parallel Azure job each - simulate → select_cohort → prepare # Stage 1, once - for k in 1 … K: # Stage 2, K times - finetune → calibrate → estimate # folds reshuffled each time +Fixed shared cohort → index_dates.csv + cohort_config.yaml + pids + │ + per outer run s (own seed): + simulate outcomes → prepare + → K refits: fit → calibrate → estimate (BERT and/or CatBoost baseline) + │ + summarize.py → bias / SD / SE-calibration / coverage, per estimator × outcome ``` -Each Azure job is **one outer run**; `submit_runs.sh` submits the N jobs. +- **Outer runs** redraw the simulated outcomes (Monte Carlo over the DGP). 1 is + enough for "does it recover the effect"; ~10–20 for a coverage sanity check. +- **Inner refits (`-k`)** estimate the SE the way the real experiments do: + - `k=1` → single fit; `estimate` reports its internal-bootstrap CI (quick check). + - `k>1` → each refit trains on a **bootstrap resample** of the cohort; the + summarizer combines the K point estimates into the SE/CI. ## What to set (once) -Edit `experiments/semisynthetic_simulation/job_config_template.yaml` — the datastore -paths for `meds`, `features`, `tokenized`, `pretrain_model`. Leave `results` ending in -`__RUN__` (the submit script fills in `run_01`, `run_02`, …). +Edit `job_config_template.yaml` — the datastore paths for `meds`, `features`, +`tokenized`, `pretrain_model`, and **`cohort`** (your diabetes cohort dir). Leave +`results` ending in `__RUN__`. -The simulated outcomes (and effect sizes) live in -`experiments/semisynthetic_simulation/base_configs/simulate.yaml`: -- `OUTCOME_NULL` → `delta: 0.0` (no effect) -- `OUTCOME_MEDIUM` → `delta: 0.5` (~6pp risk difference at a ~12% baseline) -- For the full study add `OUTCOME_LARGE` with `delta: 1.0`, and add - `OUTCOME_LARGE.csv` to `outcome_files` in `base_configs/prepare.yaml`. +In `base_configs/simulate.yaml`, set **`exposure_code`** to whatever marks a +treated patient in your MEDS, and check **`splits`** covers your cohort's patients. +Effect sizes are `OUTCOME_NULL` (δ=0), `OUTCOME_MEDIUM` (δ=0.5), `OUTCOME_LARGE` (δ=1.0). ## What to run -### 1. Smoke test (do this first) +### 1. Phase 1 — first shot (does it recover the effect?) -Two outer runs, two inner fits each, 10% of patients sampled: +One run, single fit, method + baseline: ```bash -./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 2 -k 2 -f 0.1 +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh ``` -Check one job's `estimate/bert/estimate_results.csv` against the true effects in -`simulated_outcomes/simulation_stats.csv` (NULL ≈ 0, MEDIUM ≈ 0.06 risk difference). +Then summarize and check θ̂ vs θ* and CI coverage: -### 2. Full study (after the smoke test looks right) +```bash +python -m experiments.semisynthetic_simulation.python_scripts.summarize \ + --study-dir /run_01 +``` -Ten outer runs, ten inner fits each, full population (no `-f`): +### 2. Full run + +A few outer runs, K bootstrap refits each (the real variance procedure): ```bash -./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 10 -k 10 +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 5 -k 10 ``` +(rename `smoketest`→`full` in the template's `results` first). Then +`summarize.py --study-dir ` over all runs for the full table. ### Flags | flag | meaning | default | |------|---------|---------| -| `-n` | number of outer runs (parallel jobs) | 1 | -| `-k` | inner reshuffle fits per run | 2 | -| `-f` | sample this fraction of patients (smoke tests) | none (full) | - -Override the compute/experiment via env: `POOL=CPU-20-LP EXPERIMENT=my_exp ./submit_runs.sh ...` +| `-n` | outer runs (parallel jobs) | 1 | +| `-k` | bootstrap refits per run | 1 | +| `--baseline-only` / `--bert-only` | restrict to one model | both | -## What a single job runs (equivalent manual command) - -```bash -python -m corebehrt.azure job run_semisynthetic_study CPU-20-LP \ - -e semisynthetic_study \ - -c experiments/semisynthetic_simulation/generated_job_configs/run_01.yaml \ - --bash-args "--run-id run_01 --inner-runs 2 --sample-fraction 0.1" -``` +Override compute/experiment via env: `POOL= EXPERIMENT=my_exp ./submit_runs.sh ...` +(BERT fits want a GPU pool; the CatBoost baseline runs on CPU.) ## Outputs (per outer run) ``` /run_NN/ -├── simulated_outcomes/ # exposure.csv, OUTCOME_*.csv, counterfactuals.csv, ite.csv, stats, figs -├── cohort/ # selected cohort + folds -├── prepared_data/ # tokenized finetune data -├── _configs/ # exact per-step configs used (for reproducibility) -└── reshuffles/ - ├── k_01/{models/bert, estimate/bert}/ - └── k_02/... +├── simulated_outcomes/ exposure.csv, OUTCOME_*.csv, counterfactuals.csv, stats, figs +├── prepared_data/ tokenized finetune data (+ folds) +├── _configs/ exact per-step configs used +└── reshuffles/k_NN/ + ├── models/{bert,baseline}/... + └── estimate/{bert,baseline}/estimate_results.csv ← point estimate + CI + true_effect ``` - -Seeds are independent across runs (`seed = base_seed + run_number`), so the N outer -runs are genuinely independent replicates. diff --git a/experiments/semisynthetic_simulation/job_config_template.yaml b/experiments/semisynthetic_simulation/job_config_template.yaml index 8f420bd6..b9d1d9b6 100644 --- a/experiments/semisynthetic_simulation/job_config_template.yaml +++ b/experiments/semisynthetic_simulation/job_config_template.yaml @@ -3,11 +3,13 @@ # Edit the paths below to your datastore locations, then run submit_runs.sh. paths: - ## INPUTS (shared across every run) + ## INPUTS (shared and FIXED across every run) meds: "researcher_data:AKK/shared/MEDS/TRACE/v01/data" features: "researcher_data:AKK/shared/features/trace/v01/features" tokenized: "researcher_data:AKK/shared/features/trace/v01/tokenized" pretrain_model: "researcher_data:AKK/shared/pretrain/models/trace/small/len_512/v01" + # Pre-existing cohort (its index_dates.csv + cohort_config.yaml define the fixed design) + cohort: "researcher_data:AKK/experiments/trace/semaglutide/cohort/adult_diab2" ## OUTPUT (one dir per outer run; __RUN__ is filled in by submit_runs.sh) results: "researcher_data:AKK/experiments/trace/simulation/semisynthetic/smoketest/__RUN__" diff --git a/experiments/semisynthetic_simulation/python_scripts/run_study.py b/experiments/semisynthetic_simulation/python_scripts/run_study.py index b497f566..f98234ee 100644 --- a/experiments/semisynthetic_simulation/python_scripts/run_study.py +++ b/experiments/semisynthetic_simulation/python_scripts/run_study.py @@ -2,15 +2,22 @@ """ Runner for the semi-synthetic simulation study. -Structure (mirrors the resampling study, but with the semi-synthetic simulator): - - for each OUTER run (independent simulation, own seed): - Stage 1 (once): simulate -> select_cohort -> prepare - Stage 2 (K times): finetune -> calibrate -> estimate (folds reshuffled each time) - -Each Azure job runs a single outer run (pass --run-id run_NN); the outer loop is -the set of parallel jobs submitted by bash_scripts/submit_runs.sh. Run locally with ---n-runs for several outer runs in one process. +Design (see docs/study.md): +- The cohort is FIXED and shared (a pre-existing select_cohort_full output). + Index dates and membership come from it; treatment is the real exposure. +- Stage 1 (once per outer run): simulate outcomes -> prepare finetune data. +- Stage 2 (K inner refits): fit -> calibrate -> estimate, for the causal model + (BERT) and/or the CatBoost baseline. + +Inner refits: +- K=1 -> a single plain fit; estimate reports its internal-bootstrap CI + (quick "does it recover the effect" check). +- K>1 -> each refit trains on a BOOTSTRAP resample of the cohort + (reshuffle + bootstrap); estimate gives a point per refit and the + summarizer combines the K into the SE/CI (the real-experiment variance). + +Each Azure job runs one outer run (--run-id run_NN); the outer loop = the set +of parallel jobs submitted by bash_scripts/submit_runs.sh. """ import argparse @@ -21,9 +28,9 @@ import yaml from corebehrt.main_causal.simulate_semisynthetic import main_simulate -from corebehrt.main_causal.select_cohort_full import main as main_select_cohort from corebehrt.main_causal.prepare_ft_exp_y import main as main_prepare from corebehrt.main_causal.finetune_exp_y import main_finetune +from corebehrt.main_causal.train_baseline import main_baseline from corebehrt.main_causal.calibrate_exp_y import main_calibrate from corebehrt.main_causal.estimate import main_estimate @@ -33,6 +40,7 @@ logger = logging.getLogger("semisynthetic_study") DEFAULT_BASE_CONFIGS = Path(__file__).resolve().parent.parent / "base_configs" +SINGLE_FIT_BOOTSTRAP = 100 # estimate's internal bootstrap when K=1 (no refit variance) def fill_config(base_path: Path, replacements: dict, out_path: Path, edit=None) -> str: @@ -49,7 +57,7 @@ def fill_config(base_path: Path, replacements: dict, out_path: Path, edit=None) def run_outer(args, run_id: str, seed: int): - """Run one outer simulation followed by K inner reshuffle fits.""" + """Run one outer simulation followed by K inner (bootstrap) refits.""" run_dir = Path(args.experiment_dir) / run_id config_dir = run_dir / "_configs" base = Path(args.base_configs_dir) @@ -59,77 +67,117 @@ def run_outer(args, run_id: str, seed: int): "{{FEATURES}}": args.features, "{{TOKENIZED}}": args.tokenized, "{{PRETRAIN_MODEL}}": args.pretrain_model, + "{{COHORT}}": args.cohort, "{{RUN_DIR}}": str(run_dir), } + do_bootstrap = args.inner_runs > 1 + n_bootstrap = 0 if do_bootstrap else SINGLE_FIT_BOOTSTRAP logger.info("=" * 70) - logger.info(f"OUTER RUN {run_id} (seed={seed})") + logger.info(f"OUTER RUN {run_id} (seed={seed}, bootstrap_refits={do_bootstrap})") logger.info("=" * 70) - # ---- Stage 1: simulate -> select_cohort -> prepare (once) ---- - def set_simulation(config): - config["seed"] = seed - if args.sample_fraction is not None or args.sample_size is not None: - config["sampling"] = { - "enabled": True, - "fraction": args.sample_fraction, - "size": args.sample_size, - } - + # ---- Stage 1: simulate -> prepare (once) ---- main_simulate( fill_config( - base / "simulate.yaml", shared, config_dir / "simulate.yaml", set_simulation - ) - ) - main_select_cohort( - fill_config( - base / "select_cohort.yaml", shared, config_dir / "select_cohort.yaml" + base / "simulate.yaml", + shared, + config_dir / "simulate.yaml", + lambda c: c.update(seed=seed), ) ) main_prepare( fill_config(base / "prepare.yaml", shared, config_dir / "prepare.yaml") ) - # ---- Stage 2: K inner reshuffle fits ---- + # ---- Stage 2: K inner refits ---- for k in range(1, args.inner_runs + 1): inner_id = f"k_{k:02d}" inner_dir = run_dir / "reshuffles" / inner_id repl = {**shared, "{{INNER_DIR}}": str(inner_dir)} - logger.info(f"--- {run_id} / inner fit {k}/{args.inner_runs} ({inner_id}) ---") + logger.info(f"--- {run_id} / refit {k}/{args.inner_runs} ({inner_id}) ---") - def enable_reshuffle(config): + def bootstrap_fit(config, _seed=seed * 1000 + k): + """Train this refit on a bootstrap resample of the cohort.""" config.setdefault("data", {})["reshuffle"] = True + config["data"]["reshuffle_seed"] = _seed + config["bootstrap"] = True + + def set_n_bootstrap(config): + config.setdefault("estimator", {})["n_bootstrap"] = n_bootstrap - main_finetune( - fill_config( - base / "finetune.yaml", + fit_edit = bootstrap_fit if do_bootstrap else None + + if not args.baseline_only: + _run_model( + "bert", + base, repl, - config_dir / f"finetune_{inner_id}.yaml", - enable_reshuffle, + config_dir, + inner_id, + fit_edit, + set_n_bootstrap, + main_finetune, + "finetune.yaml", + "calibrate.yaml", + "estimate.yaml", ) - ) - main_calibrate( - fill_config( - base / "calibrate.yaml", repl, config_dir / f"calibrate_{inner_id}.yaml" - ) - ) - main_estimate( - fill_config( - base / "estimate.yaml", repl, config_dir / f"estimate_{inner_id}.yaml" + if not args.bert_only: + _run_model( + "baseline", + base, + repl, + config_dir, + inner_id, + fit_edit, + set_n_bootstrap, + main_baseline, + "train_baseline.yaml", + "calibrate_baseline.yaml", + "estimate_baseline.yaml", ) - ) logger.info(f"OUTER RUN {run_id} complete") -def parse_arguments(argv=None) -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Run the semi-synthetic simulation study" +def _run_model( + name, + base, + repl, + config_dir, + inner_id, + fit_edit, + est_edit, + fit_main, + fit_cfg, + cal_cfg, + est_cfg, +): + """Run fit -> calibrate -> estimate for one model family on one refit.""" + fit_main( + fill_config( + base / fit_cfg, repl, config_dir / f"{name}_fit_{inner_id}.yaml", fit_edit + ) + ) + main_calibrate( + fill_config(base / cal_cfg, repl, config_dir / f"{name}_cal_{inner_id}.yaml") ) + main_estimate( + fill_config( + base / est_cfg, repl, config_dir / f"{name}_est_{inner_id}.yaml", est_edit + ) + ) + + +def parse_arguments(argv=None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the semi-synthetic study") parser.add_argument("--meds", required=True) parser.add_argument("--features", required=True) parser.add_argument("--tokenized", required=True) parser.add_argument("--pretrain-model", dest="pretrain_model", required=True) + parser.add_argument( + "--cohort", required=True, help="Pre-existing cohort dir (fixed across runs)" + ) parser.add_argument("--experiment-dir", dest="experiment_dir", required=True) parser.add_argument( "--base-configs-dir", dest="base_configs_dir", default=str(DEFAULT_BASE_CONFIGS) @@ -139,37 +187,25 @@ def parse_arguments(argv=None) -> argparse.Namespace: "--run-id", dest="run_id", default=None, - help="Single outer run id, e.g. run_03 (seed = base-seed + 3). One Azure job = one run-id.", - ) - parser.add_argument( - "--n-runs", - dest="n_runs", - type=int, - default=1, - help="Number of outer runs in this process (local use; ignored if --run-id is given).", + help="Single outer run id, e.g. run_03 (seed = base-seed + 3).", ) + parser.add_argument("--n-runs", dest="n_runs", type=int, default=1) parser.add_argument( "--inner-runs", "-k", dest="inner_runs", type=int, - default=2, - help="Inner reshuffle fits per outer run (variance estimation).", + default=1, + help="Bootstrap refits per outer run (K=1 -> single fit + internal-bootstrap CI).", ) parser.add_argument("--base-seed", dest="base_seed", type=int, default=42) - parser.add_argument( - "--sample-fraction", - dest="sample_fraction", - type=float, - default=None, - help="Sample this fraction of patients per run (smoke tests).", - ) - parser.add_argument("--sample-size", dest="sample_size", type=int, default=None) + parser.add_argument("--bert-only", action="store_true") + parser.add_argument("--baseline-only", action="store_true") args = parser.parse_args(argv) - if args.sample_fraction is not None and args.sample_size is not None: - parser.error("Specify at most one of --sample-fraction / --sample-size") + if args.bert_only and args.baseline_only: + parser.error("Cannot specify both --bert-only and --baseline-only") return args From 7500fe4b86ed9356d54a23d9f4481b6131722f32 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sat, 27 Jun 2026 12:34:10 +0200 Subject: [PATCH 16/20] point study cohort at adult_diab2/v01 Co-Authored-By: Claude Opus 4.8 --- experiments/semisynthetic_simulation/job_config_template.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/semisynthetic_simulation/job_config_template.yaml b/experiments/semisynthetic_simulation/job_config_template.yaml index b9d1d9b6..876f16ad 100644 --- a/experiments/semisynthetic_simulation/job_config_template.yaml +++ b/experiments/semisynthetic_simulation/job_config_template.yaml @@ -9,7 +9,7 @@ paths: tokenized: "researcher_data:AKK/shared/features/trace/v01/tokenized" pretrain_model: "researcher_data:AKK/shared/pretrain/models/trace/small/len_512/v01" # Pre-existing cohort (its index_dates.csv + cohort_config.yaml define the fixed design) - cohort: "researcher_data:AKK/experiments/trace/semaglutide/cohort/adult_diab2" + cohort: "researcher_data:AKK/experiments/trace/semaglutide/cohort/adult_diab2/v01" ## OUTPUT (one dir per outer run; __RUN__ is filled in by submit_runs.sh) results: "researcher_data:AKK/experiments/trace/simulation/semisynthetic/smoketest/__RUN__" From 1cc61c5799bc6549d439877654650b1c83095174 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sat, 27 Jun 2026 12:48:30 +0200 Subject: [PATCH 17/20] fix cohort path segment: cohorts/adult_diab2/v01 Co-Authored-By: Claude Opus 4.8 --- experiments/semisynthetic_simulation/job_config_template.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experiments/semisynthetic_simulation/job_config_template.yaml b/experiments/semisynthetic_simulation/job_config_template.yaml index 876f16ad..3b8a1605 100644 --- a/experiments/semisynthetic_simulation/job_config_template.yaml +++ b/experiments/semisynthetic_simulation/job_config_template.yaml @@ -9,7 +9,7 @@ paths: tokenized: "researcher_data:AKK/shared/features/trace/v01/tokenized" pretrain_model: "researcher_data:AKK/shared/pretrain/models/trace/small/len_512/v01" # Pre-existing cohort (its index_dates.csv + cohort_config.yaml define the fixed design) - cohort: "researcher_data:AKK/experiments/trace/semaglutide/cohort/adult_diab2/v01" + cohort: "researcher_data:AKK/experiments/trace/semaglutide/cohorts/adult_diab2/v01" ## OUTPUT (one dir per outer run; __RUN__ is filled in by submit_runs.sh) results: "researcher_data:AKK/experiments/trace/simulation/semisynthetic/smoketest/__RUN__" From 1369342ef83817ef4810a5c10f866a34e4c49fb4 Mon Sep 17 00:00:00 2001 From: kirilklein Date: Sun, 28 Jun 2026 20:43:47 +0200 Subject: [PATCH 18/20] take observed treatment from the cohort, not a MEDS exposure code On a real cohort there is no generic EXPOSURE code in raw MEDS, so the simulator found 0 exposed and never wrote exposure.csv -> prepare then failed its (bogus) exposure pre-check. Treatment is observed and already recorded by select_cohort_full in cohort/exposures.csv (same file prepare reads). So: - simulator: load exposed pids from cohort/exposures.csv when index_dates points at a cohort dir; exposure_code is now only a fallback. - prepare config: drop the exposures/exposure keys; prepare reads exposure from the cohort (load_cohort_data), so the sim-output pre-check was wrong. Validated locally with a bogus exposure_code: exposure still resolves from the cohort, exposure.csv is written, true effect recovered (ATE 0.064 vs 0.063). Co-Authored-By: Claude Opus 4.8 --- .../simulation/semisynthetic_simulator.py | 37 ++++++++++++++++--- .../base_configs/prepare.yaml | 3 +- .../base_configs/simulate.yaml | 4 +- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 97aaa145..6b6fac96 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -22,6 +22,7 @@ ) from corebehrt.constants.causal.paths import ( COUNTERFACTUALS_FILE, + EXPOSURES_FILE, INDEX_DATE_MATCHING_FILE, ) from corebehrt.constants.paths import INDEX_DATES_FILE @@ -70,6 +71,26 @@ def __init__(self, config: SemiSyntheticSimulationConfig): self._global_stds = None self._validate_feature_references() self._index_dates = self._load_index_dates() + self._exposed_pids = self._load_exposed_pids() + + def _load_exposed_pids(self): + """Load the set of exposed (treated) patient IDs from the cohort. + + Treatment is observed: when ``paths.index_dates`` points to a cohort + dir containing ``exposures.csv``, the exposed are exactly its patients + (same source select_cohort/prepare use). Otherwise fall back to the + presence of ``exposure_code`` in the MEDS data. + """ + path = self.config.paths.index_dates + if not path or not os.path.isdir(path): + return None + exposures_path = join(path, EXPOSURES_FILE) + if not os.path.exists(exposures_path): + return None + exposed = pd.read_csv(exposures_path, usecols=[PID_COL]) + pids = set(exposed[PID_COL].unique()) + logger.info(f"Loaded {len(pids)} exposed patient IDs from {exposures_path}") + return pids def _load_index_dates(self): """Load per-patient index dates from a cohort artifact, if configured. @@ -257,12 +278,16 @@ def _extract_treatment_and_index_dates( pd.Series(dtype="datetime64[ns]"), ) - # Patients with at least one exposure event (treatment from the data) - exposed_pids = set( - shard_df.loc[ - shard_df[CONCEPT_COL] == self.config.exposure_code, PID_COL - ].unique() - ) + # Observed treatment: from the cohort's exposed pids, or (fallback) + # the presence of the exposure code in the data. + if self._exposed_pids is not None: + exposed_pids = self._exposed_pids + else: + exposed_pids = set( + shard_df.loc[ + shard_df[CONCEPT_COL] == self.config.exposure_code, PID_COL + ].unique() + ) pids = index_dates.index.values is_exposed = np.array([pid in exposed_pids for pid in pids]) diff --git a/experiments/semisynthetic_simulation/base_configs/prepare.yaml b/experiments/semisynthetic_simulation/base_configs/prepare.yaml index c4de8774..e46a9550 100644 --- a/experiments/semisynthetic_simulation/base_configs/prepare.yaml +++ b/experiments/semisynthetic_simulation/base_configs/prepare.yaml @@ -16,8 +16,7 @@ paths: - OUTCOME_MEDIUM.csv - OUTCOME_LARGE.csv - exposures: "{{RUN_DIR}}/simulated_outcomes" - exposure: exposure.csv + # Exposure is read from the cohort (cohort/exposures.csv), so it is not set here. ## OUTPUTS prepared_data: "{{RUN_DIR}}/prepared_data" diff --git a/experiments/semisynthetic_simulation/base_configs/simulate.yaml b/experiments/semisynthetic_simulation/base_configs/simulate.yaml index dfdc1ce7..7004fe27 100644 --- a/experiments/semisynthetic_simulation/base_configs/simulate.yaml +++ b/experiments/semisynthetic_simulation/base_configs/simulate.yaml @@ -14,8 +14,8 @@ paths: index_dates: "{{COHORT}}" # dir or index_dates.csv; index dates + membership min_num_codes: 3 -# Code marking a treated patient at index. Set to whatever marks exposure -# in your MEDS (e.g. the semaglutide exposure code), not necessarily "EXPOSURE". +# Treatment is read from the cohort (cohort/exposures.csv). exposure_code is only +# a fallback used when no cohort exposures.csv is available (e.g. example data). exposure_code: "EXPOSURE" features: From 1edab37488d822e1957d7137dde8c83fa55945f7 Mon Sep 17 00:00:00 2001 From: kirilklein <49436907+kirilklein@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:05:18 +0200 Subject: [PATCH 19/20] fix semisynthetic bootstrap uncertainty workflow --- corebehrt/constants/causal/paths.py | 1 + corebehrt/functional/causal/effect.py | 4 +- corebehrt/functional/estimate/benchmarks.py | 19 +- corebehrt/main_causal/finetune_exp_y.py | 3 + .../main_causal/helper/train_baseline.py | 30 ++- corebehrt/modules/causal/estimate.py | 96 +++++++-- .../simulation/semisynthetic_simulator.py | 24 ++- .../base_configs/estimate.yaml | 7 +- .../base_configs/estimate_baseline.yaml | 7 +- .../base_configs/simulate.yaml | 6 +- .../bash_scripts/submit_runs.sh | 10 +- .../semisynthetic_simulation/docs/study.md | 28 ++- .../python_scripts/run_study.py | 64 +++--- .../python_scripts/study_summary.py | 192 ++++++++++++++++++ .../python_scripts/summarize.py | 90 +------- tests/test_main_causal/test_baseline_folds.py | 53 +++++ .../test_causal_estimate_bootstrap.py | 84 ++++++++ .../test_simulation/test_semisynthetic.py | 42 ++++ tests/test_semisynthetic_study_summary.py | 104 ++++++++++ 19 files changed, 702 insertions(+), 162 deletions(-) create mode 100644 experiments/semisynthetic_simulation/python_scripts/study_summary.py create mode 100644 tests/test_main_causal/test_baseline_folds.py create mode 100644 tests/test_modules/test_causal_estimate_bootstrap.py create mode 100644 tests/test_semisynthetic_study_summary.py diff --git a/corebehrt/constants/causal/paths.py b/corebehrt/constants/causal/paths.py index a26d6ec4..6b0ab47a 100644 --- a/corebehrt/constants/causal/paths.py +++ b/corebehrt/constants/causal/paths.py @@ -23,6 +23,7 @@ EXPERIMENT_DATA_FILE = "experiment_data.parquet" EXPERIMENT_STATS_FILE = "experiment_stats.csv" ESTIMATE_RESULTS_FILE = "estimate_results.csv" +BOOTSTRAP_RESULTS_FILE = "bootstrap_results.csv" SKS_DUMP_DIR = "corebehrt/main_causal/helper/data/sks_dumps" SKS_TREES_DIR = "corebehrt/main_causal/helper/data/sks_trees" diff --git a/corebehrt/functional/causal/effect.py b/corebehrt/functional/causal/effect.py index 0e0d3fb6..4236bf7d 100644 --- a/corebehrt/functional/causal/effect.py +++ b/corebehrt/functional/causal/effect.py @@ -52,7 +52,7 @@ def compute_effect_from_counterfactuals(df: pd.DataFrame, effect_type: str) -> f y1_mean = df[SIMULATED_PROBAS_EXPOSED].mean() y0_mean = df[SIMULATED_PROBAS_CONTROL].mean() - if effect_type == "ATE": + if effect_type in {"ATE", "ARR"}: effect = y1_mean - y0_mean elif effect_type in ["ATT", "ATC"]: treated_flag = 1 if effect_type == "ATT" else 0 @@ -62,7 +62,7 @@ def compute_effect_from_counterfactuals(df: pd.DataFrame, effect_type: str) -> f - subset[SIMULATED_PROBAS_CONTROL].mean() ) elif effect_type == "RR": - effect = (y1_mean + 1) / (y0_mean + 1) + effect = y1_mean / y0_mean elif effect_type == "OR": effect = (y1_mean / (1 - y1_mean)) / (y0_mean / (1 - y0_mean)) else: diff --git a/corebehrt/functional/estimate/benchmarks.py b/corebehrt/functional/estimate/benchmarks.py index e892ca09..a33b6f66 100644 --- a/corebehrt/functional/estimate/benchmarks.py +++ b/corebehrt/functional/estimate/benchmarks.py @@ -11,7 +11,10 @@ EffectColumns, ) from corebehrt.constants.data import PID_COL -from corebehrt.functional.causal.effect import compute_effect_from_ite +from corebehrt.functional.causal.effect import ( + compute_effect_from_counterfactuals, + compute_effect_from_ite, +) from corebehrt.functional.causal.estimate import ( calculate_risk_difference, calculate_risk_ratio, @@ -42,6 +45,8 @@ def append_true_effect( ite_df: pd.DataFrame, outcome_name: str, analysis_pids: np.ndarray, + effect_type: str = "ATE", + counterfactual_df: pd.DataFrame | None = None, ) -> pd.DataFrame: """ Add ground truth effect estimates from simulated counterfactual outcomes. @@ -51,9 +56,15 @@ def append_true_effect( Adds the true effect to the effect_df. (TRUE_EFFECT_COL) """ - effect_df[EffectColumns.true_effect] = compute_effect_from_ite( - ite_df, analysis_pids, outcome_name - ) + if counterfactual_df is not None: + counterfactuals = prepare_counterfactual_data_for_outcome( + counterfactual_df, outcome_name + ) + counterfactuals = counterfactuals[counterfactuals[PID_COL].isin(analysis_pids)] + true_effect = compute_effect_from_counterfactuals(counterfactuals, effect_type) + else: + true_effect = compute_effect_from_ite(ite_df, analysis_pids, outcome_name) + effect_df[EffectColumns.true_effect] = true_effect return effect_df diff --git a/corebehrt/main_causal/finetune_exp_y.py b/corebehrt/main_causal/finetune_exp_y.py index b70df4a2..7908b73d 100644 --- a/corebehrt/main_causal/finetune_exp_y.py +++ b/corebehrt/main_causal/finetune_exp_y.py @@ -6,6 +6,7 @@ import torch import time import numpy as np +from transformers import set_seed from corebehrt.constants.paths import ( FOLDS_FILE, @@ -33,6 +34,8 @@ def main_finetune(config_path): cfg = load_config(config_path) + if (seed := cfg.get("seed")) is not None: + set_seed(seed) # Setup directories DirectoryPreparer(cfg).setup_finetune() diff --git a/corebehrt/main_causal/helper/train_baseline.py b/corebehrt/main_causal/helper/train_baseline.py index 9cf77e10..f16da443 100644 --- a/corebehrt/main_causal/helper/train_baseline.py +++ b/corebehrt/main_causal/helper/train_baseline.py @@ -28,6 +28,7 @@ from corebehrt.constants.paths import ( FOLDS_FILE, ) +from corebehrt.functional.features.split import create_folds from corebehrt.functional.preparation.causal.one_hot import ( create_features_from_patients, ) @@ -246,6 +247,7 @@ def run_hyperparameter_tuning( config_params: Dict[str, Any], n_trials: int, scale_pos_weight: float, + random_seed: int = 42, ) -> Dict[str, Any]: """ INNER LOOP: Performs hyperparameter tuning using Optuna on a given train/val split. @@ -332,7 +334,7 @@ def objective(trial: optuna.Trial): model = CatBoostClassifier( n_estimators=base_params["n_estimators"], scale_pos_weight=scale_pos_weight, - random_state=42, + random_state=random_seed, verbose=0, **device_params, **prepared_trial_params, @@ -359,7 +361,9 @@ def objective(trial: optuna.Trial): logging.info( f" Running Optuna optimization for {len(params_to_tune)} parameters..." ) - study = optuna.create_study(direction="maximize") + study = optuna.create_study( + direction="maximize", sampler=optuna.samplers.TPESampler(seed=random_seed) + ) study.optimize(objective, n_trials=n_trials) logging.info(f" Hyperparameter tuning completed!") @@ -446,7 +450,7 @@ def _get_best_params_for_fold( inner_train_pids, inner_val_pids = train_test_split( outer_train_data.get_pids(), test_size=inner_val_size, - random_state=42, + random_state=cfg.get("seed", 42), ) logger.info( @@ -474,6 +478,7 @@ def _get_best_params_for_fold( config_params, n_trials, scale_pos_weight, + cfg.get("seed", 42), ) logger.info(" Hyperparameter tuning completed for this fold") @@ -520,6 +525,7 @@ def _train_and_evaluate_fold( target_name: str, fold_idx: int, prediction_storage: List[FoldPredictionData], + random_seed: int = 42, ) -> float: """Trains a final model and evaluates it on the holdout test set.""" logger.info(f" Training final model with outer train set size: {len(X_train)}") @@ -540,7 +546,7 @@ def _train_and_evaluate_fold( final_model = CatBoostClassifier( scale_pos_weight=scale_pos_weight, - random_state=42, + random_state=random_seed, **device_params, **prepared_best_params, ) @@ -688,6 +694,7 @@ def nested_cv_loop( target_name, i, prediction_storage, + cfg.get("seed", 42) + i, ) all_unbiased_scores.append(unbiased_auc) @@ -714,11 +721,22 @@ def nested_cv_loop( def handle_folds(cfg: Config, logger: logging.Logger) -> list: """ - Load predefined folds, log and persist them into the model directory, and return. + Load predefined folds and optionally reshuffle patients across them. """ folds_path = join(cfg.paths.prepared_data, FOLDS_FILE) folds = torch.load(folds_path) n_folds = len(folds) - logger.info(f"Using {n_folds} predefined folds") + data_cfg = cfg.get("data", {}) + if data_cfg.get("reshuffle", False): + pids = sorted( + {pid for fold in folds for split in fold.values() for pid in split} + ) + seed = data_cfg.get("reshuffle_seed", 42) + folds = create_folds(pids, n_folds, seed) + logger.info( + f"Reshuffled {len(pids)} patients into {n_folds} folds (seed={seed})" + ) + else: + logger.info(f"Using {n_folds} predefined folds") torch.save(folds, join(cfg.paths.model, FOLDS_FILE)) return folds diff --git a/corebehrt/modules/causal/estimate.py b/corebehrt/modules/causal/estimate.py index 3f4e01ba..9787e5c7 100644 --- a/corebehrt/modules/causal/estimate.py +++ b/corebehrt/modules/causal/estimate.py @@ -2,6 +2,7 @@ import os from os.path import join +import numpy as np import pandas as pd import torch from CausalEstimate import MultiEstimator @@ -19,6 +20,7 @@ EFFECT_ROUND_DIGIT, ) from corebehrt.constants.causal.paths import ( + BOOTSTRAP_RESULTS_FILE, COMBINED_CALIBRATED_PREDICTIONS_FILE, COUNTERFACTUALS_FILE, PATIENTS_FILE, @@ -78,6 +80,7 @@ def __init__(self, cfg: Config, logger: logging.Logger): ) self.estimator_cfg: dict = self.cfg.estimator self.init_estimator_args(self.estimator_cfg) + self.bootstrap_records = [] self._init_plot_configs() self.effect_type: str = self.cfg.estimator.effect_type self.df = pd.read_csv(self.predictions_file) @@ -144,14 +147,16 @@ def run_standard_estimation(self) -> None: df_for_outcome = prepare_data_for_outcome(self.analysis_df, outcome_name) # 2. Estimate effects using the new logic - effect_df = self._estimate_effects(df_for_outcome) + effect_df = self._estimate_effects(df_for_outcome, outcome_name) - if self.ite_df is not None: + if self.ite_df is not None or self.counterfactual_df is not None: effect_df = append_true_effect( effect_df, self.ite_df, outcome_name, self.analysis_df[PID_COL].values, + effect_type=self.effect_type, + counterfactual_df=self.counterfactual_df, ) effect_df = append_unadjusted_effect(df_for_outcome, effect_df) @@ -176,10 +181,13 @@ def run_standard_estimation(self) -> None: final_results_df, combined_stats_df, tmle_analysis_df = ( self._process_and_save_results(all_effects, all_stats, initial_estimates) ) + self._save_bootstrap_results() self._visualize_effects(final_results_df, combined_stats_df, tmle_analysis_df) self.logger.info("Effect estimation complete for all outcomes.") - def _estimate_effects(self, df: pd.DataFrame) -> pd.DataFrame: + def _estimate_effects( + self, df: pd.DataFrame, outcome_name: str | None = None + ) -> pd.DataFrame: """ Estimate effects, separating bootstrap methods from theoretical (single-shot) methods. """ @@ -218,27 +226,78 @@ def _estimate_effects(self, df: pd.DataFrame) -> pd.DataFrame: ) if bootstrap_estimator_list: + observed_results = ( + { + estimator.__class__.__name__: estimator.compute_effect(df) + for estimator in bootstrap_estimator_list + } + if self.use_observed_point_estimate + else {} + ) multi_estimator = MultiEstimator( estimators=bootstrap_estimator_list, verbose=False ) - # Run with bootstrapping only if n_bootstrap is greater than 0 - n_boot = self.n_bootstrap if self.n_bootstrap > 0 else 0 - if n_boot == 0: - self.logger.warning( - "n_bootstrap is 0, running bootstrap methods without CI." - ) - effect_dict_bs = multi_estimator.compute_effects( df, - n_bootstraps=n_boot, + n_bootstraps=self.n_bootstrap, apply_common_support=False, common_support_threshold=None, - return_bootstrap_samples=False, + return_bootstrap_samples=self.save_bootstrap_samples, ) + if self.save_bootstrap_samples: + self._collect_bootstrap_results(effect_dict_bs, outcome_name) + for method, observed in observed_results.items(): + for key, value in observed.items(): + if key not in { + EffectColumns.std_err, + EffectColumns.CI95_lower, + EffectColumns.CI95_upper, + }: + effect_dict_bs[method][key] = value all_effect_dicts.update(effect_dict_bs) return convert_effect_to_dataframe(all_effect_dicts) + def _collect_bootstrap_results( + self, effect_results: dict, outcome_name: str | None + ) -> None: + """Collect patient-bootstrap draws for study-level aggregation.""" + if outcome_name is None: + raise ValueError("outcome_name is required when saving bootstrap samples") + + for method, result in effect_results.items(): + samples = result.pop("bootstrap_samples", None) + if samples is None: + continue + for bootstrap_id, effect in enumerate(samples[EffectColumns.effect], 1): + self.bootstrap_records.append( + { + EffectColumns.method: method, + OUTCOME: outcome_name, + "effect_type": self.effect_type, + "bootstrap_id": bootstrap_id, + EffectColumns.effect: effect, + EffectColumns.effect_1: samples[EffectColumns.effect_1][ + bootstrap_id - 1 + ], + EffectColumns.effect_0: samples[EffectColumns.effect_0][ + bootstrap_id - 1 + ], + } + ) + + def _save_bootstrap_results(self) -> None: + """Persist raw patient-bootstrap estimates when requested.""" + if not self.save_bootstrap_samples: + return + if not self.bootstrap_records: + raise RuntimeError( + "Bootstrap samples were requested but none were produced" + ) + pd.DataFrame(self.bootstrap_records).to_csv( + join(self.exp_dir, BOOTSTRAP_RESULTS_FILE), index=False + ) + def _visualize_effects( self, final_results_df: pd.DataFrame, @@ -381,8 +440,15 @@ def init_estimator_args(self, cfg) -> None: """ self.common_support_threshold: float = cfg.get("common_support_threshold", None) self.common_support: bool = True if self.common_support_threshold else False - self.n_bootstrap: int = cfg.get("n_bootstrap", 0) + self.n_bootstrap: int = cfg.get("n_bootstrap", 1) self.clip_percentile: float = cfg.get("clip_percentile", 1) + self.save_bootstrap_samples: bool = cfg.get("save_bootstrap_samples", False) + self.use_observed_point_estimate: bool = cfg.get( + "use_observed_point_estimate", False + ) + bootstrap_seed = cfg.get("bootstrap_seed") + if bootstrap_seed is not None: + np.random.seed(bootstrap_seed) def _get_analysis_cohort(self, df: pd.DataFrame) -> pd.DataFrame: """ @@ -463,6 +529,8 @@ def _run_bias_simulation(self) -> None: self.ite_df, outcome_name, self.analysis_df[PID_COL].values, + effect_type=self.effect_type, + counterfactual_df=self.counterfactual_df, ) true_effect_value = true_effect_df[EffectColumns.true_effect].iloc[0] @@ -473,7 +541,7 @@ def _run_bias_simulation(self) -> None: ) # UPDATED: Call the refactored estimation method - effect_df = self._estimate_effects(df_biased) + effect_df = self._estimate_effects(df_biased, outcome_name) effect_df[EffectColumns.ps_bias] = ps_bias effect_df[EffectColumns.y_bias] = y_bias diff --git a/corebehrt/modules/simulation/semisynthetic_simulator.py b/corebehrt/modules/simulation/semisynthetic_simulator.py index 6b6fac96..520a30d8 100644 --- a/corebehrt/modules/simulation/semisynthetic_simulator.py +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -76,17 +76,19 @@ def __init__(self, config: SemiSyntheticSimulationConfig): def _load_exposed_pids(self): """Load the set of exposed (treated) patient IDs from the cohort. - Treatment is observed: when ``paths.index_dates`` points to a cohort - dir containing ``exposures.csv``, the exposed are exactly its patients - (same source select_cohort/prepare use). Otherwise fall back to the - presence of ``exposure_code`` in the MEDS data. + A configured cohort must contain ``exposures.csv``. Without a cohort, + example data may still use the presence of ``exposure_code`` in MEDS. """ - path = self.config.paths.index_dates - if not path or not os.path.isdir(path): + path = self.config.paths.cohort + if not path: return None + if not os.path.isdir(path): + raise NotADirectoryError(f"Cohort path is not a directory: {path}") exposures_path = join(path, EXPOSURES_FILE) if not os.path.exists(exposures_path): - return None + raise FileNotFoundError( + f"Configured cohort is missing observed treatments: {exposures_path}" + ) exposed = pd.read_csv(exposures_path, usecols=[PID_COL]) pids = set(exposed[PID_COL].unique()) logger.info(f"Loaded {len(pids)} exposed patient IDs from {exposures_path}") @@ -96,11 +98,11 @@ def _load_index_dates(self): """Load per-patient index dates from a cohort artifact, if configured. Index dates are an analysis artifact (produced by cohort selection), - not raw MEDS. When ``paths.index_dates`` points to a cohort dir or an - ``index_dates.csv``, use it; otherwise fall back to an - ``assigned_index_date`` column in the data (e.g. the example data). + not raw MEDS. Use ``paths.index_dates`` when provided, otherwise use + ``paths.cohort/index_dates.csv``. Example data may fall back to an + ``assigned_index_date`` column in MEDS. """ - path = self.config.paths.index_dates + path = self.config.paths.index_dates or self.config.paths.cohort if not path: return None if os.path.isdir(path): diff --git a/experiments/semisynthetic_simulation/base_configs/estimate.yaml b/experiments/semisynthetic_simulation/base_configs/estimate.yaml index 37359709..b8a65770 100644 --- a/experiments/semisynthetic_simulation/base_configs/estimate.yaml +++ b/experiments/semisynthetic_simulation/base_configs/estimate.yaml @@ -12,10 +12,11 @@ paths: estimate: "{{INNER_DIR}}/estimate/bert" estimator: - methods: ["IPW", "TMLE", "TMLE_TH"] + methods: ["IPW", "TMLE"] effect_type: "ATE" - n_bootstrap: 100 # set by the runner: 0 for bootstrap-refit studies, >0 for single-fit CIs - clip_percentile: 0.99 + n_bootstrap: 100 + common_support_threshold: 0.001 + clip_percentile: 1.0 plot: contingency_table: diff --git a/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml index 61736c15..c19def8f 100644 --- a/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml +++ b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml @@ -1,5 +1,5 @@ # Stage 2c (baseline): estimate causal effects from calibrated baseline predictions. -# n_bootstrap is set by the runner (0 for bootstrap-refit studies, >0 for single-fit CIs). +# n_bootstrap and its seed are set by the study runner. logging: level: INFO path: ./outputs/logs/sim_study @@ -10,10 +10,11 @@ paths: estimate: "{{INNER_DIR}}/estimate/baseline" estimator: - methods: ["IPW", "TMLE", "TMLE_TH"] + methods: ["IPW", "TMLE"] effect_type: "ATE" n_bootstrap: 100 - clip_percentile: 0.99 + common_support_threshold: 0.001 + clip_percentile: 1.0 plot: contingency_table: diff --git a/experiments/semisynthetic_simulation/base_configs/simulate.yaml b/experiments/semisynthetic_simulation/base_configs/simulate.yaml index 7004fe27..949aa2cb 100644 --- a/experiments/semisynthetic_simulation/base_configs/simulate.yaml +++ b/experiments/semisynthetic_simulation/base_configs/simulate.yaml @@ -11,11 +11,11 @@ paths: # Splits must cover the patients in the cohort (filtered by index_dates). splits: ["tuning", "train"] outcomes: "{{RUN_DIR}}/simulated_outcomes" - index_dates: "{{COHORT}}" # dir or index_dates.csv; index dates + membership + cohort: "{{COHORT}}" # requires index_dates.csv and exposures.csv min_num_codes: 3 -# Treatment is read from the cohort (cohort/exposures.csv). exposure_code is only -# a fallback used when no cohort exposures.csv is available (e.g. example data). +# Treatment is read from cohort/exposures.csv. exposure_code is only used by +# standalone example-data configs that do not configure a cohort. exposure_code: "EXPOSURE" features: diff --git a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh index f48e2336..4c4e055f 100755 --- a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh +++ b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh @@ -3,9 +3,9 @@ # FIXED shared cohort. Each job = one outer simulation (own seed) + K refits. # # Usage: -# ./submit_runs.sh # Phase 1: 1 run, K=1 (single fit + analytic CI) +# ./submit_runs.sh # Phase 1: 1 run, K=1, B=100 # ./submit_runs.sh --baseline-only # Phase 1, baseline only (CPU, fast) -# ./submit_runs.sh -n 5 -k 10 # full: 5 outer runs, 10 bootstrap refits each +# ./submit_runs.sh -n 5 -k 10 -b 100 # full: 5 runs, 10 refits, 100 bootstraps # # Anything after the known flags is forwarded to the runner (e.g. --bert-only). # Override defaults via env: POOL=... EXPERIMENT=... TEMPLATE=... ./submit_runs.sh ... @@ -18,12 +18,14 @@ GENERATED_DIR="${GENERATED_DIR:-experiments/semisynthetic_simulation/generated_j N_RUNS=1 INNER_RUNS=1 +N_BOOTSTRAP=100 EXTRA_ARGS=() while [[ $# -gt 0 ]]; do case "$1" in -n) N_RUNS="$2"; shift 2 ;; -k) INNER_RUNS="$2"; shift 2 ;; + -b) N_BOOTSTRAP="$2"; shift 2 ;; -h|--help) grep '^#' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; *) EXTRA_ARGS+=("$1"); shift ;; esac @@ -39,9 +41,9 @@ for ((i = 1; i <= N_RUNS; i++)); do CFG="$GENERATED_DIR/$RUN.yaml" sed "s|__RUN__|$RUN|g" "$TEMPLATE" > "$CFG" - BASH_ARGS="--run-id $RUN --inner-runs $INNER_RUNS ${EXTRA_ARGS[*]:-}" + BASH_ARGS="--run-id $RUN --inner-runs $INNER_RUNS --n-bootstrap $N_BOOTSTRAP ${EXTRA_ARGS[*]:-}" - echo "Submitting $RUN (inner_runs=$INNER_RUNS) ${EXTRA_ARGS[*]:-}" + echo "Submitting $RUN (refits=$INNER_RUNS, bootstraps=$N_BOOTSTRAP) ${EXTRA_ARGS[*]:-}" python -m corebehrt.azure job run_semisynthetic_study "$POOL" \ -e "$EXPERIMENT" \ -c "$CFG" \ diff --git a/experiments/semisynthetic_simulation/docs/study.md b/experiments/semisynthetic_simulation/docs/study.md index cac62dd7..def6e095 100644 --- a/experiments/semisynthetic_simulation/docs/study.md +++ b/experiments/semisynthetic_simulation/docs/study.md @@ -10,17 +10,19 @@ Fixed shared cohort → index_dates.csv + cohort_config.yaml + pids │ per outer run s (own seed): simulate outcomes → prepare - → K refits: fit → calibrate → estimate (BERT and/or CatBoost baseline) + → K refits: fit → calibrate → B patient bootstraps + (BERT and/or CatBoost baseline) │ summarize.py → bias / SD / SE-calibration / coverage, per estimator × outcome ``` - **Outer runs** redraw the simulated outcomes (Monte Carlo over the DGP). 1 is enough for "does it recover the effect"; ~10–20 for a coverage sanity check. -- **Inner refits (`-k`)** estimate the SE the way the real experiments do: - - `k=1` → single fit; `estimate` reports its internal-bootstrap CI (quick check). - - `k>1` → each refit trains on a **bootstrap resample** of the cohort; the - summarizer combines the K point estimates into the SE/CI. +- **Inner refits (`-k`)** use different model seeds and reshuffled CV folds, + yielding K independently fitted propensity-score models. +- **Patient bootstraps (`-b`, default 100)** resample the original analysis + cohort with replacement conditional on each fitted model. The summarizer + combines the K × B draws into one SE and CI per outer simulation replicate. ## What to set (once) @@ -51,10 +53,10 @@ python -m experiments.semisynthetic_simulation.python_scripts.summarize \ ### 2. Full run -A few outer runs, K bootstrap refits each (the real variance procedure): +A few outer runs, K seeded refits and B patient bootstraps each: ```bash -./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 5 -k 10 +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 5 -k 10 -b 100 ``` (rename `smoketest`→`full` in the template's `results` first). Then `summarize.py --study-dir ` over all runs for the full table. @@ -64,7 +66,8 @@ A few outer runs, K bootstrap refits each (the real variance procedure): | flag | meaning | default | |------|---------|---------| | `-n` | outer runs (parallel jobs) | 1 | -| `-k` | bootstrap refits per run | 1 | +| `-k` | independently seeded model refits per run | 1 | +| `-b` | patient bootstrap samples per refit | 100 | | `--baseline-only` / `--bert-only` | restrict to one model | both | Override compute/experiment via env: `POOL= EXPERIMENT=my_exp ./submit_runs.sh ...` @@ -79,5 +82,12 @@ Override compute/experiment via env: `POOL= EXPERIMENT=my_exp ./submit_runs ├── _configs/ exact per-step configs used └── reshuffles/k_NN/ ├── models/{bert,baseline}/... - └── estimate/{bert,baseline}/estimate_results.csv ← point estimate + CI + true_effect + └── estimate/{bert,baseline}/ + ├── estimate_results.csv + └── bootstrap_results.csv ← B patient-level bootstrap estimates ``` + +Running `summarize.py` writes `replicate_estimates.csv` (one K × B aggregate +per outer run) and `summary.csv` (bias, empirical SD, mean SE, SE calibration, +and coverage across outer runs). Common support is trimmed at the 0.1st and +99.9th percentiles within treatment arms before patient resampling. diff --git a/experiments/semisynthetic_simulation/python_scripts/run_study.py b/experiments/semisynthetic_simulation/python_scripts/run_study.py index f98234ee..4c72c566 100644 --- a/experiments/semisynthetic_simulation/python_scripts/run_study.py +++ b/experiments/semisynthetic_simulation/python_scripts/run_study.py @@ -10,11 +10,9 @@ (BERT) and/or the CatBoost baseline. Inner refits: -- K=1 -> a single plain fit; estimate reports its internal-bootstrap CI - (quick "does it recover the effect" check). -- K>1 -> each refit trains on a BOOTSTRAP resample of the cohort - (reshuffle + bootstrap); estimate gives a point per refit and the - summarizer combines the K into the SE/CI (the real-experiment variance). +- Each refit uses a distinct model seed and fold reshuffling. +- Conditional on each fitted model, estimate draws B patient-level bootstrap + samples. The summarizer combines all K x B estimates. Each Azure job runs one outer run (--run-id run_NN); the outer loop = the set of parallel jobs submitted by bash_scripts/submit_runs.sh. @@ -40,7 +38,7 @@ logger = logging.getLogger("semisynthetic_study") DEFAULT_BASE_CONFIGS = Path(__file__).resolve().parent.parent / "base_configs" -SINGLE_FIT_BOOTSTRAP = 100 # estimate's internal bootstrap when K=1 (no refit variance) +DEFAULT_BOOTSTRAPS = 100 def fill_config(base_path: Path, replacements: dict, out_path: Path, edit=None) -> str: @@ -57,8 +55,10 @@ def fill_config(base_path: Path, replacements: dict, out_path: Path, edit=None) def run_outer(args, run_id: str, seed: int): - """Run one outer simulation followed by K inner (bootstrap) refits.""" - run_dir = Path(args.experiment_dir) / run_id + """Run one outer simulation followed by K independently seeded refits.""" + run_dir = ( + Path(args.experiment_dir) if args.run_id else Path(args.experiment_dir) / run_id + ) config_dir = run_dir / "_configs" base = Path(args.base_configs_dir) @@ -70,11 +70,11 @@ def run_outer(args, run_id: str, seed: int): "{{COHORT}}": args.cohort, "{{RUN_DIR}}": str(run_dir), } - do_bootstrap = args.inner_runs > 1 - n_bootstrap = 0 if do_bootstrap else SINGLE_FIT_BOOTSTRAP - logger.info("=" * 70) - logger.info(f"OUTER RUN {run_id} (seed={seed}, bootstrap_refits={do_bootstrap})") + logger.info( + f"OUTER RUN {run_id} (seed={seed}, refits={args.inner_runs}, " + f"bootstraps={args.n_bootstrap})" + ) logger.info("=" * 70) # ---- Stage 1: simulate -> prepare (once) ---- @@ -97,16 +97,20 @@ def run_outer(args, run_id: str, seed: int): repl = {**shared, "{{INNER_DIR}}": str(inner_dir)} logger.info(f"--- {run_id} / refit {k}/{args.inner_runs} ({inner_id}) ---") - def bootstrap_fit(config, _seed=seed * 1000 + k): - """Train this refit on a bootstrap resample of the cohort.""" + refit_seed = seed * 1000 + k + + def configure_refit(config, _seed=refit_seed): + """Apply the model seed and reshuffle folds without resampling patients.""" config.setdefault("data", {})["reshuffle"] = True config["data"]["reshuffle_seed"] = _seed - config["bootstrap"] = True + config["seed"] = _seed - def set_n_bootstrap(config): - config.setdefault("estimator", {})["n_bootstrap"] = n_bootstrap - - fit_edit = bootstrap_fit if do_bootstrap else None + def configure_estimation(config, _seed=refit_seed): + estimator = config.setdefault("estimator", {}) + estimator["n_bootstrap"] = args.n_bootstrap + estimator["bootstrap_seed"] = _seed + estimator["save_bootstrap_samples"] = True + estimator["use_observed_point_estimate"] = True if not args.baseline_only: _run_model( @@ -115,8 +119,8 @@ def set_n_bootstrap(config): repl, config_dir, inner_id, - fit_edit, - set_n_bootstrap, + configure_refit, + configure_estimation, main_finetune, "finetune.yaml", "calibrate.yaml", @@ -129,8 +133,8 @@ def set_n_bootstrap(config): repl, config_dir, inner_id, - fit_edit, - set_n_bootstrap, + configure_refit, + configure_estimation, main_baseline, "train_baseline.yaml", "calibrate_baseline.yaml", @@ -196,7 +200,15 @@ def parse_arguments(argv=None) -> argparse.Namespace: dest="inner_runs", type=int, default=1, - help="Bootstrap refits per outer run (K=1 -> single fit + internal-bootstrap CI).", + help="Independently seeded model refits per outer run.", + ) + parser.add_argument( + "--n-bootstrap", + "-b", + dest="n_bootstrap", + type=int, + default=DEFAULT_BOOTSTRAPS, + help="Patient-level bootstrap samples per fitted propensity model.", ) parser.add_argument("--base-seed", dest="base_seed", type=int, default=42) @@ -206,6 +218,10 @@ def parse_arguments(argv=None) -> argparse.Namespace: args = parser.parse_args(argv) if args.bert_only and args.baseline_only: parser.error("Cannot specify both --bert-only and --baseline-only") + if args.inner_runs < 1: + parser.error("--inner-runs must be at least 1") + if args.n_bootstrap < 2: + parser.error("--n-bootstrap must be at least 2") return args diff --git a/experiments/semisynthetic_simulation/python_scripts/study_summary.py b/experiments/semisynthetic_simulation/python_scripts/study_summary.py new file mode 100644 index 00000000..a02856f3 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/study_summary.py @@ -0,0 +1,192 @@ +"""Aggregate nested semi-synthetic study estimates.""" + +from pathlib import Path + +import numpy as np +import pandas as pd + +from corebehrt.constants.causal.data import EffectColumns as E +from corebehrt.constants.causal.data import OUTCOME +from corebehrt.constants.causal.paths import ( + BOOTSTRAP_RESULTS_FILE, + ESTIMATE_RESULTS_FILE, +) + +MODEL_NAMES = ("bert", "baseline") +GROUP_COLUMNS = ["model", "method", OUTCOME] +REPLICATE_GROUP_COLUMNS = ["model", "run_id", "method", OUTCOME] + + +def _tag_from_path(path: Path) -> dict: + """Infer model and nested-resampling IDs from a result path.""" + parts = [part.lower() for part in path.parts] + model = next((name for name in MODEL_NAMES if name in parts), "unknown") + run_id = next((part for part in parts if part.startswith("run_")), "run_01") + inner_id = next((part for part in parts if part.startswith("k_")), "k_01") + return {"model": model, "run_id": run_id, "inner_id": inner_id} + + +def load_tagged_results(study_dir: Path, filename: str) -> pd.DataFrame: + """Recursively load result files and add model/run/refit identifiers.""" + files = sorted(study_dir.rglob(filename)) + if not files: + raise FileNotFoundError(f"No {filename} found under {study_dir}") + + frames = [] + for path in files: + frame = pd.read_csv(path) + for key, value in _tag_from_path(path).items(): + frame[key] = value + frames.append(frame) + return pd.concat(frames, ignore_index=True) + + +def load_results(study_dir: Path) -> pd.DataFrame: + return load_tagged_results(study_dir, ESTIMATE_RESULTS_FILE) + + +def load_bootstrap_results(study_dir: Path) -> pd.DataFrame: + return load_tagged_results(study_dir, BOOTSTRAP_RESULTS_FILE) + + +def aggregate_replicates( + results: pd.DataFrame, bootstrap_results: pd.DataFrame +) -> pd.DataFrame: + """Combine K point estimates and K x B patient-bootstrap estimates per run.""" + available = bootstrap_results[GROUP_COLUMNS].drop_duplicates() + results = results.merge(available, on=GROUP_COLUMNS, how="inner") + rows = [] + + for keys, group in results.groupby(REPLICATE_GROUP_COLUMNS): + tags = dict(zip(REPLICATE_GROUP_COLUMNS, keys)) + samples = _select_bootstrap_group(bootstrap_results, tags) + _validate_nested_samples(group, samples, tags) + + effect_type = _single_value(samples["effect_type"], "effect_type", tags) + point = group[E.effect].mean() + true_effect = group[E.true_effect].mean() + + if effect_type in {"RR", "RRT"}: + uncertainty = _risk_ratio_uncertainty(group, samples) + point = uncertainty.pop("effect") + else: + uncertainty = _difference_uncertainty(point, samples[E.effect]) + + rows.append( + { + **tags, + "effect_type": effect_type, + "n_refits": group["inner_id"].nunique(), + "n_bootstrap_per_refit": samples.groupby("inner_id").size().iloc[0], + "n_bootstrap": len(samples), + E.true_effect: true_effect, + E.effect: point, + **uncertainty, + "covered": ( + uncertainty[E.CI95_lower] + <= true_effect + <= uncertainty[E.CI95_upper] + ), + } + ) + + return pd.DataFrame(rows).sort_values([OUTCOME, "model", "method", "run_id"]) + + +def _select_bootstrap_group(samples: pd.DataFrame, tags: dict) -> pd.DataFrame: + mask = pd.Series(True, index=samples.index) + for column, value in tags.items(): + mask &= samples[column] == value + return samples[mask] + + +def _validate_nested_samples( + estimates: pd.DataFrame, samples: pd.DataFrame, tags: dict +) -> None: + expected_refits = set(estimates["inner_id"]) + actual_refits = set(samples["inner_id"]) + if actual_refits != expected_refits: + raise ValueError( + f"Bootstrap refits do not match point-estimate refits for {tags}: " + f"{sorted(actual_refits)} != {sorted(expected_refits)}" + ) + + counts = samples.groupby("inner_id").size() + if counts.nunique() != 1: + raise ValueError(f"Unequal bootstrap counts across refits for {tags}: {counts}") + if not np.isfinite(samples[[E.effect, E.effect_1, E.effect_0]]).all().all(): + raise ValueError(f"Non-finite bootstrap estimates found for {tags}") + + +def _single_value(series: pd.Series, name: str, tags: dict): + values = series.drop_duplicates() + if len(values) != 1: + raise ValueError(f"Expected one {name} for {tags}, found {values.tolist()}") + return values.iloc[0] + + +def _difference_uncertainty(point: float, samples: pd.Series) -> dict: + standard_error = samples.std(ddof=1) + margin = 1.96 * standard_error + return { + E.std_err: standard_error, + "std_err_log": np.nan, + E.CI95_lower: point - margin, + E.CI95_upper: point + margin, + } + + +def _risk_ratio_uncertainty(estimates: pd.DataFrame, samples: pd.DataFrame) -> dict: + p1 = estimates[E.effect_1].mean() + p0 = estimates[E.effect_0].mean() + if not 0 < p1 < 1 or not 0 < p0 < 1: + raise ValueError(f"Risk-ratio probabilities must lie in (0, 1): {p1=}, {p0=}") + + sample_p1 = samples[E.effect_1] + sample_p0 = samples[E.effect_0] + if ( + not sample_p1.between(0, 1, inclusive="neither").all() + or not sample_p0.between(0, 1, inclusive="neither").all() + ): + raise ValueError("Bootstrap risk-ratio probabilities must lie in (0, 1)") + + eta_1 = np.log(sample_p1 / (1 - sample_p1)) + eta_0 = np.log(sample_p0 / (1 - sample_p0)) + variance_log_rr = (1 - p1) ** 2 * eta_1.var(ddof=1) + (1 - p0) ** 2 * eta_0.var( + ddof=1 + ) + std_err_log = np.sqrt(variance_log_rr) + risk_ratio = p1 / p0 + margin = 1.96 * std_err_log + return { + E.effect: risk_ratio, + E.std_err: risk_ratio * std_err_log, + "std_err_log": std_err_log, + E.CI95_lower: np.exp(np.log(risk_ratio) - margin), + E.CI95_upper: np.exp(np.log(risk_ratio) + margin), + } + + +def summarize_performance(replicates: pd.DataFrame) -> pd.DataFrame: + """Summarize bias, empirical SD, SE calibration, and coverage across runs.""" + rows = [] + for keys, group in replicates.groupby(GROUP_COLUMNS + ["effect_type"]): + tags = dict(zip(GROUP_COLUMNS + ["effect_type"], keys)) + empirical_sd = group[E.effect].std(ddof=1) if len(group) > 1 else np.nan + mean_se = group[E.std_err].mean() + rows.append( + { + **tags, + "n": len(group), + E.true_effect: group[E.true_effect].mean(), + "mean_effect": group[E.effect].mean(), + "bias": (group[E.effect] - group[E.true_effect]).mean(), + "sd_emp": empirical_sd, + "mean_se": mean_se, + "se_calibration": ( + empirical_sd / mean_se if len(group) > 1 and mean_se > 0 else np.nan + ), + "coverage": group["covered"].mean(), + } + ) + return pd.DataFrame(rows).sort_values([OUTCOME, "model", "method"]) diff --git a/experiments/semisynthetic_simulation/python_scripts/summarize.py b/experiments/semisynthetic_simulation/python_scripts/summarize.py index 7f777ad0..d385320c 100644 --- a/experiments/semisynthetic_simulation/python_scripts/summarize.py +++ b/experiments/semisynthetic_simulation/python_scripts/summarize.py @@ -1,15 +1,5 @@ #!/usr/bin/env python3 -""" -Summarize a semi-synthetic study into an estimator-performance table. - -Reads every ``estimate_results.csv`` under a study directory (each already -carries the point estimate, its CI, and the appended ``true_effect``), groups -by model x method x outcome, and reports bias, empirical SD, mean estimated SE, -SE-calibration, and 95% CI coverage. - -- One estimate per group (Phase 1, single run): bias + covered are meaningful; - SD/calibration are undefined (need >1). -- Many estimates per group (bootstrap refits / outcome redraws): full table. +"""Summarize nested model-refit and patient-bootstrap study results. Usage: python -m experiments.semisynthetic_simulation.python_scripts.summarize \ @@ -17,77 +7,16 @@ """ import argparse -import re from pathlib import Path -import numpy as np import pandas as pd -from corebehrt.constants.causal.data import EffectColumns as E -from corebehrt.constants.causal.data import OUTCOME -from corebehrt.constants.causal.paths import ESTIMATE_RESULTS_FILE - -MODEL_NAMES = ("bert", "baseline") - - -def _tag_from_path(path: Path) -> dict: - """Infer model / outer-run / inner-refit ids from a result file's path.""" - parts = [p.lower() for p in path.parts] - model = next((m for m in MODEL_NAMES if m in parts), "unknown") - run_match = re.search(r"run_\d+", str(path)) - inner_match = re.search(r"k_\d+", str(path)) - return { - "model": model, - "run_id": run_match.group(0) if run_match else "run_01", - "inner_id": inner_match.group(0) if inner_match else "k_01", - } - - -def load_results(study_dir: Path) -> pd.DataFrame: - """Load and tag every estimate_results.csv under study_dir.""" - files = sorted(study_dir.rglob(ESTIMATE_RESULTS_FILE)) - if not files: - raise FileNotFoundError(f"No {ESTIMATE_RESULTS_FILE} found under {study_dir}") - frames = [] - for path in files: - df = pd.read_csv(path) - for key, value in _tag_from_path(path).items(): - df[key] = value - frames.append(df) - return pd.concat(frames, ignore_index=True) - - -def summarize(results: pd.DataFrame) -> pd.DataFrame: - """Aggregate estimates into a per (model, method, outcome) performance table.""" - rows = [] - for (model, method, outcome), group in results.groupby( - ["model", E.method, OUTCOME] - ): - effects = group[E.effect] - ses = group[E.std_err] - true = group[E.true_effect].iloc[0] - covered = (group[E.CI95_lower] <= true) & (group[E.CI95_upper] >= true) - n = len(group) - sd_emp = effects.std(ddof=1) if n > 1 else np.nan - mean_se = ses.mean() - rows.append( - { - "model": model, - "method": method, - "outcome": outcome, - "n": n, - "true_effect": true, - "mean_effect": effects.mean(), - "bias": effects.mean() - true, - "sd_emp": sd_emp, - "mean_se": mean_se, - "se_calibration": (sd_emp / mean_se) - if n > 1 and mean_se > 0 - else np.nan, - "coverage": covered.mean(), - } - ) - return pd.DataFrame(rows).sort_values(["outcome", "model", "method"]) +from experiments.semisynthetic_simulation.python_scripts.study_summary import ( + aggregate_replicates, + load_bootstrap_results, + load_results, + summarize_performance, +) def main(): @@ -100,9 +29,12 @@ def main(): study_dir = Path(args.study_dir) results = load_results(study_dir) - table = summarize(results) + bootstrap_results = load_bootstrap_results(study_dir) + replicates = aggregate_replicates(results, bootstrap_results) + table = summarize_performance(replicates) out = Path(args.out) if args.out else study_dir / "summary.csv" + replicates.to_csv(study_dir / "replicate_estimates.csv", index=False) table.to_csv(out, index=False) pd.set_option("display.float_format", lambda v: f"{v:.4f}") diff --git a/tests/test_main_causal/test_baseline_folds.py b/tests/test_main_causal/test_baseline_folds.py new file mode 100644 index 00000000..427e89e3 --- /dev/null +++ b/tests/test_main_causal/test_baseline_folds.py @@ -0,0 +1,53 @@ +"""Tests for baseline refit fold reshuffling.""" + +import logging +import os +import sys +import tempfile +import types +import unittest + +import torch + +try: + import optuna # noqa: F401 +except ImportError: + optuna_stub = types.ModuleType("optuna") + optuna_stub.Trial = object + sys.modules["optuna"] = optuna_stub + +from corebehrt.constants.data import TRAIN_KEY, VAL_KEY +from corebehrt.constants.paths import FOLDS_FILE +from corebehrt.main_causal.helper.train_baseline import handle_folds +from corebehrt.modules.setup.config import Config + + +class TestBaselineFoldReshuffling(unittest.TestCase): + def test_reshuffle_uses_configured_seed_and_all_patients(self): + prepared_dir = tempfile.mkdtemp() + model_dir = tempfile.mkdtemp() + pids = list(range(20)) + original = [ + {TRAIN_KEY: pids[10:], VAL_KEY: pids[:10]}, + {TRAIN_KEY: pids[:10], VAL_KEY: pids[10:]}, + ] + torch.save(original, os.path.join(prepared_dir, FOLDS_FILE)) + cfg = Config( + { + "paths": {"prepared_data": prepared_dir, "model": model_dir}, + "data": {"reshuffle": True, "reshuffle_seed": 123}, + } + ) + + reshuffled = handle_folds(cfg, logging.getLogger("test_baseline_folds")) + + self.assertNotEqual(reshuffled, original) + for fold in reshuffled: + self.assertEqual(set(fold[TRAIN_KEY]) | set(fold[VAL_KEY]), set(pids)) + self.assertTrue( + set(fold[TRAIN_KEY]).isdisjoint(set(fold[VAL_KEY])) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modules/test_causal_estimate_bootstrap.py b/tests/test_modules/test_causal_estimate_bootstrap.py new file mode 100644 index 00000000..a6ade399 --- /dev/null +++ b/tests/test_modules/test_causal_estimate_bootstrap.py @@ -0,0 +1,84 @@ +"""Tests for patient-bootstrap persistence in causal estimation.""" + +import logging +import unittest + +import numpy as np +import pandas as pd + +from corebehrt.constants.causal.data import ( + EXPOSURE_COL, + OUTCOME, + PROBAS, + PROBAS_CONTROL, + PROBAS_EXPOSED, + PS_COL, + EffectColumns, +) +from corebehrt.constants.data import PID_COL +from corebehrt.functional.estimate.benchmarks import append_true_effect +from corebehrt.modules.causal.estimate import EffectEstimator +from corebehrt.modules.setup.config import Config + + +class TestEffectEstimatorBootstrap(unittest.TestCase): + def test_saves_draws_but_keeps_observed_point_estimate(self): + rng = np.random.default_rng(42) + exposure = np.repeat([0, 1], 30) + outcome = rng.binomial(1, 0.3 + 0.2 * exposure) + frame = pd.DataFrame( + { + EXPOSURE_COL: exposure, + OUTCOME: outcome, + PS_COL: np.where(exposure, 0.55, 0.45), + PROBAS: 0.3 + 0.2 * exposure, + PROBAS_CONTROL: 0.3, + PROBAS_EXPOSED: 0.5, + } + ) + estimator = EffectEstimator.__new__(EffectEstimator) + estimator.estimator_cfg = Config({"methods": ["IPW"]}) + estimator.effect_type = "ATE" + estimator.clip_percentile = 0.99 + estimator.n_bootstrap = 5 + estimator.save_bootstrap_samples = True + estimator.use_observed_point_estimate = True + estimator.bootstrap_records = [] + estimator.logger = logging.getLogger("test_bootstrap") + + result = estimator._estimate_effects(frame, "OUTCOME") + + self.assertEqual(len(estimator.bootstrap_records), 5) + self.assertEqual( + {row["bootstrap_id"] for row in estimator.bootstrap_records}, + set(range(1, 6)), + ) + observed = outcome[exposure == 1].mean() - outcome[exposure == 0].mean() + self.assertAlmostEqual(result.iloc[0][EffectColumns.effect], observed) + self.assertGreater(result.iloc[0][EffectColumns.std_err], 0) + + def test_true_risk_ratio_uses_counterfactual_probabilities(self): + counterfactuals = pd.DataFrame( + { + PID_COL: [1, 2], + "Y1_OUTCOME": [0, 1], + "Y0_OUTCOME": [0, 0], + "P1_OUTCOME": [0.4, 0.6], + "P0_OUTCOME": [0.2, 0.3], + EXPOSURE_COL: [0, 1], + } + ) + result = append_true_effect( + pd.DataFrame({EffectColumns.method: ["IPW"]}), + ite_df=None, + outcome_name="OUTCOME", + analysis_pids=np.array([1, 2]), + effect_type="RR", + counterfactual_df=counterfactuals, + ) + + self.assertAlmostEqual(result.iloc[0][EffectColumns.true_effect], 2.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modules/test_simulation/test_semisynthetic.py b/tests/test_modules/test_simulation/test_semisynthetic.py index 96c27395..18503d34 100644 --- a/tests/test_modules/test_simulation/test_semisynthetic.py +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -279,5 +279,47 @@ def test_valid_config_does_not_raise(self): SemiSyntheticCausalSimulator(_make_config(tempfile.mkdtemp())) +class TestCohortTreatmentLoading(unittest.TestCase): + def test_configured_cohort_requires_exposures(self): + import os + import tempfile + + cohort_dir = tempfile.mkdtemp() + pd.DataFrame( + { + PID_COL: [1], + TIMESTAMP_COL: pd.to_datetime(["2020-01-01"]), + } + ).to_csv(os.path.join(cohort_dir, "index_dates.csv"), index=False) + config = _make_config(tempfile.mkdtemp()) + config.paths.cohort = cohort_dir + with self.assertRaisesRegex(FileNotFoundError, "observed treatments"): + SemiSyntheticCausalSimulator(config) + + def test_exposed_patients_come_from_cohort(self): + import os + import tempfile + + cohort_dir = tempfile.mkdtemp() + pd.DataFrame({PID_COL: [1, 3]}).to_csv( + os.path.join(cohort_dir, "exposures.csv"), index=False + ) + pd.DataFrame( + { + PID_COL: [1, 2, 3], + TIMESTAMP_COL: pd.to_datetime(["2020-01-01"] * 3), + } + ).to_csv(os.path.join(cohort_dir, "index_dates.csv"), index=False) + config = _make_config(tempfile.mkdtemp()) + config.paths.cohort = cohort_dir + + simulator = SemiSyntheticCausalSimulator(config) + pids, is_exposed, _ = simulator._extract_treatment_and_index_dates( + _make_test_shard(n_patients=3, n_exposed=0) + ) + + self.assertEqual(dict(zip(pids, is_exposed)), {1: True, 2: False, 3: True}) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_semisynthetic_study_summary.py b/tests/test_semisynthetic_study_summary.py new file mode 100644 index 00000000..ccd1544f --- /dev/null +++ b/tests/test_semisynthetic_study_summary.py @@ -0,0 +1,104 @@ +"""Tests for nested semi-synthetic study aggregation.""" + +import unittest + +import numpy as np +import pandas as pd + +from corebehrt.constants.causal.data import EffectColumns as E +from experiments.semisynthetic_simulation.python_scripts.study_summary import ( + aggregate_replicates, + summarize_performance, +) + + +def _estimate_rows(run_id, effects, true_effect): + return pd.DataFrame( + { + "model": "bert", + "run_id": run_id, + "inner_id": ["k_01", "k_02"], + E.method: "IPW", + "outcome": "OUTCOME", + E.effect: effects, + E.effect_1: np.asarray(effects) + 0.2, + E.effect_0: 0.2, + E.true_effect: true_effect, + } + ) + + +def _bootstrap_rows(run_id, samples): + rows = [] + for inner_id, effects in zip(["k_01", "k_02"], samples): + for bootstrap_id, effect in enumerate(effects, 1): + rows.append( + { + "model": "bert", + "run_id": run_id, + "inner_id": inner_id, + E.method: "IPW", + "outcome": "OUTCOME", + "effect_type": "ATE", + "bootstrap_id": bootstrap_id, + E.effect: effect, + E.effect_1: effect + 0.2, + E.effect_0: 0.2, + } + ) + return pd.DataFrame(rows) + + +class TestNestedStudySummary(unittest.TestCase): + def test_aggregates_k_times_b_with_run_specific_truth(self): + results = pd.concat( + [ + _estimate_rows("run_01", [0.10, 0.12], 0.11), + _estimate_rows("run_02", [0.20, 0.22], 0.19), + ], + ignore_index=True, + ) + bootstraps = pd.concat( + [ + _bootstrap_rows("run_01", [[0.08, 0.12], [0.10, 0.14]]), + _bootstrap_rows("run_02", [[0.18, 0.22], [0.20, 0.24]]), + ], + ignore_index=True, + ) + + replicates = aggregate_replicates(results, bootstraps) + self.assertEqual(list(replicates["n_bootstrap"]), [4, 4]) + self.assertAlmostEqual(replicates.iloc[0][E.effect], 0.11) + self.assertAlmostEqual( + replicates.iloc[0][E.std_err], np.std([0.08, 0.12, 0.10, 0.14], ddof=1) + ) + + summary = summarize_performance(replicates).iloc[0] + self.assertEqual(summary["n"], 2) + self.assertAlmostEqual(summary["bias"], 0.01) + self.assertAlmostEqual(summary[E.true_effect], 0.15) + + def test_rejects_unequal_bootstrap_counts(self): + results = _estimate_rows("run_01", [0.10, 0.12], 0.11) + bootstraps = _bootstrap_rows("run_01", [[0.08], [0.10, 0.14]]) + + with self.assertRaisesRegex(ValueError, "Unequal bootstrap counts"): + aggregate_replicates(results, bootstraps) + + def test_risk_ratio_uses_log_scale_interval(self): + results = _estimate_rows("run_01", [2.0, 2.0], 2.0) + results[E.effect_1] = [0.4, 0.4] + results[E.effect_0] = [0.2, 0.2] + bootstraps = _bootstrap_rows("run_01", [[2.0, 2.1], [1.9, 2.0]]) + bootstraps["effect_type"] = "RR" + bootstraps[E.effect_1] = [0.39, 0.41, 0.38, 0.40] + bootstraps[E.effect_0] = [0.20, 0.20, 0.19, 0.20] + + replicate = aggregate_replicates(results, bootstraps).iloc[0] + self.assertAlmostEqual(replicate[E.effect], 2.0) + self.assertGreater(replicate["std_err_log"], 0) + self.assertGreater(replicate[E.CI95_lower], 0) + + +if __name__ == "__main__": + unittest.main() From 18e2fa5f14ab373abd1c2f67c7bd652f5dbedb40 Mon Sep 17 00:00:00 2001 From: kirilklein <49436907+kirilklein@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:05:18 +0200 Subject: [PATCH 20/20] add training-fold bootstrap for k-refit uncertainty - resample training folds with replacement per refit; validation folds stay fixed so every patient keeps one out-of-fold score - preserve duplicate pids through dataset resampling for bert and baseline - collect per-refit patient-bootstrap draws in the estimate step and aggregate the k refit estimates in the study summary - fail fast when estimator n_bootstrap < 1: the previously pushed runner set 0 for k>1 and crashed only after all refits had trained, and no unit test exercises that wiring Co-Authored-By: Claude Fable 5 --- corebehrt/constants/causal/data.py | 1 + corebehrt/functional/features/split.py | 21 ++++++ corebehrt/main_causal/finetune_exp_y.py | 52 ++++++------- .../main_causal/helper/finetune_exp_y.py | 6 +- .../main_causal/helper/train_baseline.py | 26 +++++-- corebehrt/modules/causal/estimate.py | 13 +++- .../modules/preparation/causal/dataset.py | 8 ++ .../bash_scripts/submit_runs.sh | 4 +- .../semisynthetic_simulation/docs/study.md | 32 ++++---- .../python_scripts/run_study.py | 16 ++-- .../python_scripts/study_summary.py | 75 +++++-------------- .../python_scripts/summarize.py | 6 +- tests/test_main_causal/test_baseline_folds.py | 34 ++++++--- .../test_finetune_bootstrap_folds.py | 58 ++++++++++++++ .../test_causal_dataset_resampling.py | 34 +++++++++ tests/test_semisynthetic_study_summary.py | 58 ++++---------- 16 files changed, 268 insertions(+), 176 deletions(-) create mode 100644 tests/test_main_causal/test_finetune_bootstrap_folds.py create mode 100644 tests/test_modules/test_causal_dataset_resampling.py diff --git a/corebehrt/constants/causal/data.py b/corebehrt/constants/causal/data.py index 0ba8b8cf..d9a2a1a3 100644 --- a/corebehrt/constants/causal/data.py +++ b/corebehrt/constants/causal/data.py @@ -71,6 +71,7 @@ class EffectColumns: method = "method" + effect_type = "effect_type" effect = "effect" true_effect = "true_effect" ps_bias = "ps_bias" diff --git a/corebehrt/functional/features/split.py b/corebehrt/functional/features/split.py index 1f1365df..78780800 100644 --- a/corebehrt/functional/features/split.py +++ b/corebehrt/functional/features/split.py @@ -187,3 +187,24 @@ def create_folds( folds[i][VAL_KEY] = val_pids return folds + + +def bootstrap_training_folds( + folds: List[Dict[str, list]], seed: int = 42 +) -> List[Dict[str, list]]: + """Resample each training fold with replacement, keeping validation fixed.""" + rng = np.random.default_rng(seed) + bootstrapped = [] + for fold in folds: + train_pids = fold[TRAIN_KEY] + if not train_pids: + raise ValueError("Cannot bootstrap an empty training fold") + bootstrapped.append( + { + TRAIN_KEY: rng.choice( + train_pids, size=len(train_pids), replace=True + ).tolist(), + VAL_KEY: list(fold[VAL_KEY]), + } + ) + return bootstrapped diff --git a/corebehrt/main_causal/finetune_exp_y.py b/corebehrt/main_causal/finetune_exp_y.py index 7908b73d..b0a08f9d 100644 --- a/corebehrt/main_causal/finetune_exp_y.py +++ b/corebehrt/main_causal/finetune_exp_y.py @@ -15,7 +15,10 @@ TEST_PIDS_FILE, ) from corebehrt.functional.setup.args import get_args -from corebehrt.functional.features.split import create_folds +from corebehrt.functional.features.split import ( + bootstrap_training_folds, + create_folds, +) from corebehrt.main.helper.finetune_cv import check_for_overlap from corebehrt.main_causal.helper.finetune_exp_y import cv_loop from corebehrt.modules.monitoring.causal.metric_aggregation import ( @@ -106,9 +109,8 @@ def validate_folds( Validate fold structure for correctness. Checks: - - All PIDs are present (no loss) - - Each fold has unique PIDs (no duplicates within) unless bootstrap=True - - Validation sets don't overlap across folds + - Each validation PID appears exactly once across folds + - Training PIDs may repeat only when bootstrap=True - Train/val PIDs sum to total PIDs in each fold """ all_val_pids = set() @@ -118,11 +120,10 @@ def validate_folds( val_pids = set(fold[VAL_KEY]) if not bootstrap: - # Check: No duplicates within fold - assert len(train_pids) == len(fold[TRAIN_KEY]), ( - f"Fold {i}: Duplicate train PIDs" - ) - assert len(val_pids) == len(fold[VAL_KEY]), f"Fold {i}: Duplicate val PIDs" + assert len(train_pids) == len( + fold[TRAIN_KEY] + ), f"Fold {i}: Duplicate train PIDs" + assert len(val_pids) == len(fold[VAL_KEY]), f"Fold {i}: Duplicate val PIDs" # Check: No overlap between train and val (unique PIDs) assert train_pids.isdisjoint(val_pids), f"Fold {i}: Train/val overlap" @@ -145,16 +146,12 @@ def validate_folds( fold_total = train_pids | val_pids assert fold_total == expected_pids, f"Fold {i}: Missing or extra PIDs" - # Track validation PIDs across folds - if not bootstrap: - assert val_pids.isdisjoint(all_val_pids), ( - f"Fold {i}: Val PIDs overlap with other folds" - ) + assert val_pids.isdisjoint( + all_val_pids + ), f"Fold {i}: Val PIDs overlap with other folds" all_val_pids.update(val_pids) - if not bootstrap: - # Check: All PIDs appear in exactly one validation set - assert all_val_pids == expected_pids, "Not all PIDs covered in validation sets" + assert all_val_pids == expected_pids, "Not all PIDs covered in validation sets" logger.info( f"✓ Folds validated: {len(folds)} folds, {len(expected_pids)} unique PIDs" ) @@ -189,10 +186,9 @@ def handle_folds( expected_pids = set(train_val_pids) bootstrap = cfg.get("bootstrap", False) - # Validate loaded folds — pass bootstrap flag so that duplicate PIDs - # (expected with bootstrap sampling) don't cause validation to fail + # Prepared folds are always unresampled. logger.info("Validating loaded folds...") - validate_folds(folds, expected_pids, logger, bootstrap=bootstrap) + validate_folds(folds, expected_pids, logger) data_cfg = cfg.get("data", {}) # Check if we should reshuffle @@ -209,20 +205,24 @@ def handle_folds( # Recreate folds with new seed using existing create_folds function # This ensures proper handling of uneven fold sizes and correct KFold splitting - folds = create_folds( - train_val_pids, n_folds, reshuffle_seed, bootstrap=bootstrap - ) + folds = create_folds(train_val_pids, n_folds, reshuffle_seed) logger.info( f"Reshuffled {len(train_val_pids)} unique PIDs across {n_folds} folds" ) - # Validate reshuffled folds - logger.info("Validating reshuffled folds...") - validate_folds(folds, expected_pids, logger, bootstrap=bootstrap) else: logger.info("Using folds as loaded (no reshuffling)") + if bootstrap: + bootstrap_seed = data_cfg.get( + "bootstrap_seed", data_cfg.get("reshuffle_seed", 42) + ) + folds = bootstrap_training_folds(folds, bootstrap_seed) + logger.info(f"Bootstrapped training folds with seed={bootstrap_seed}") + + logger.info("Validating refit folds...") + validate_folds(folds, expected_pids, logger, bootstrap=bootstrap) check_for_overlap(folds, test_pids, logger) torch.save(folds, join(cfg.paths.model, FOLDS_FILE)) return folds diff --git a/corebehrt/main_causal/helper/finetune_exp_y.py b/corebehrt/main_causal/helper/finetune_exp_y.py index e00c086e..ac029568 100644 --- a/corebehrt/main_causal/helper/finetune_exp_y.py +++ b/corebehrt/main_causal/helper/finetune_exp_y.py @@ -43,7 +43,11 @@ def cv_loop( val_pids = fold_dict[VAL_KEY] logger.info(f"Training fold {fold}/{len(folds)}") - train_data = data.filter_by_pids(train_pids) + train_data = ( + data.resample_by_pids(train_pids) + if cfg.get("bootstrap", False) + else data.filter_by_pids(train_pids) + ) val_data = data.filter_by_pids(val_pids) with setup_metrics_dir(f"Fold {fold}"): diff --git a/corebehrt/main_causal/helper/train_baseline.py b/corebehrt/main_causal/helper/train_baseline.py index f16da443..e3665df6 100644 --- a/corebehrt/main_causal/helper/train_baseline.py +++ b/corebehrt/main_causal/helper/train_baseline.py @@ -28,14 +28,16 @@ from corebehrt.constants.paths import ( FOLDS_FILE, ) -from corebehrt.functional.features.split import create_folds +from corebehrt.functional.features.split import ( + bootstrap_training_folds, + create_folds, +) from corebehrt.functional.preparation.causal.one_hot import ( create_features_from_patients, ) from corebehrt.modules.preparation.causal.dataset import CausalPatientDataset from corebehrt.modules.setup.config import Config - # Cache for GPU detection to avoid repeated logging _CATBOOST_DEVICE_PARAMS_CACHE = None @@ -178,7 +180,9 @@ def save_combined_predictions( return # Group predictions by fold and organize by target - fold_data = {} # fold_idx -> {target_name: (pids, predictions, targets, cf_predictions)} + fold_data = ( + {} + ) # fold_idx -> {target_name: (pids, predictions, targets, cf_predictions)} for pred_data in prediction_storage: fold_idx = pred_data.fold_idx @@ -447,8 +451,9 @@ def _get_best_params_for_fold( logger.info(f" Inner validation size: {inner_val_size}") logger.info(f" Number of tuning trials: {n_trials}") + unique_outer_pids = list(dict.fromkeys(outer_train_data.get_pids())) inner_train_pids, inner_val_pids = train_test_split( - outer_train_data.get_pids(), + unique_outer_pids, test_size=inner_val_size, random_state=cfg.get("seed", 42), ) @@ -457,7 +462,7 @@ def _get_best_params_for_fold( f" Split outer train into inner train ({len(inner_train_pids)} patients) and inner val ({len(inner_val_pids)} patients)" ) - inner_train_data = data.filter_by_pids(inner_train_pids) + inner_train_data = outer_train_data.filter_by_pids(inner_train_pids) inner_val_data = data.filter_by_pids(inner_val_pids) X_inner_train, y_inner_train, X_inner_val, y_inner_val = _prepare_data_for_modeling( @@ -647,7 +652,11 @@ def nested_cv_loop( logger.info(f"Train patients in this fold: {len(fold_dict['train'])}") logger.info(f"Test patients in this fold: {len(fold_dict['val'])}") - outer_train_data = data.filter_by_pids(fold_dict["train"]) + outer_train_data = ( + data.resample_by_pids(fold_dict["train"]) + if cfg.get("bootstrap", False) + else data.filter_by_pids(fold_dict["train"]) + ) outer_test_data = data.filter_by_pids(fold_dict["val"]) test_pids = outer_test_data.get_pids() @@ -738,5 +747,10 @@ def handle_folds(cfg: Config, logger: logging.Logger) -> list: ) else: logger.info(f"Using {n_folds} predefined folds") + seed = data_cfg.get("bootstrap_seed", 42) + if cfg.get("bootstrap", False): + bootstrap_seed = data_cfg.get("bootstrap_seed", seed) + folds = bootstrap_training_folds(folds, bootstrap_seed) + logger.info(f"Bootstrapped training folds with seed={bootstrap_seed}") torch.save(folds, join(cfg.paths.model, FOLDS_FILE)) return folds diff --git a/corebehrt/modules/causal/estimate.py b/corebehrt/modules/causal/estimate.py index 9787e5c7..e294f5e2 100644 --- a/corebehrt/modules/causal/estimate.py +++ b/corebehrt/modules/causal/estimate.py @@ -148,6 +148,7 @@ def run_standard_estimation(self) -> None: # 2. Estimate effects using the new logic effect_df = self._estimate_effects(df_for_outcome, outcome_name) + effect_df[EffectColumns.effect_type] = self.effect_type if self.ite_df is not None or self.counterfactual_df is not None: effect_df = append_true_effect( @@ -385,9 +386,11 @@ def _process_and_save_results( if any(m.upper() in ["TMLE", "TMLE_TH"] for m in self.estimator_cfg.methods) else None ) - save_tmle_analysis( - tmle_analysis_df, self.exp_dir - ) if tmle_analysis_df is not None else None + ( + save_tmle_analysis(tmle_analysis_df, self.exp_dir) + if tmle_analysis_df is not None + else None + ) return final_results_df, combined_stats_df, tmle_analysis_df def _build_estimators_for_methods(self, methods: list) -> list: @@ -441,6 +444,10 @@ def init_estimator_args(self, cfg) -> None: self.common_support_threshold: float = cfg.get("common_support_threshold", None) self.common_support: bool = True if self.common_support_threshold else False self.n_bootstrap: int = cfg.get("n_bootstrap", 1) + if self.n_bootstrap < 1: + raise ValueError( + f"estimator.n_bootstrap must be >= 1, got {self.n_bootstrap}" + ) self.clip_percentile: float = cfg.get("clip_percentile", 1) self.save_bootstrap_samples: bool = cfg.get("save_bootstrap_samples", False) self.use_observed_point_estimate: bool = cfg.get( diff --git a/corebehrt/modules/preparation/causal/dataset.py b/corebehrt/modules/preparation/causal/dataset.py index 87ab540b..f891c8f7 100644 --- a/corebehrt/modules/preparation/causal/dataset.py +++ b/corebehrt/modules/preparation/causal/dataset.py @@ -62,6 +62,14 @@ def filter_by_pids(self, pids: List[str]) -> "CausalPatientDataset": [p for p in self.patients if p.pid in pids_set], self.vocab ) + def resample_by_pids(self, pids: List[str]) -> "CausalPatientDataset": + """Select patients in the requested order while preserving duplicate IDs.""" + patients_by_pid = {patient.pid: patient for patient in self.patients} + missing = set(pids) - patients_by_pid.keys() + if missing: + raise KeyError(f"Patient IDs not found in dataset: {sorted(missing)}") + return CausalPatientDataset([patients_by_pid[pid] for pid in pids], self.vocab) + def get_exposures(self): return [p.exposure for p in self.patients] diff --git a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh index 4c4e055f..4fd9702f 100755 --- a/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh +++ b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh @@ -1,11 +1,11 @@ #!/usr/bin/env bash # Submit N outer runs of the semi-synthetic study as parallel Azure jobs, on a -# FIXED shared cohort. Each job = one outer simulation (own seed) + K refits. +# FIXED shared cohort. Each job = one outer simulation + K bootstrap refits. # # Usage: # ./submit_runs.sh # Phase 1: 1 run, K=1, B=100 # ./submit_runs.sh --baseline-only # Phase 1, baseline only (CPU, fast) -# ./submit_runs.sh -n 5 -k 10 -b 100 # full: 5 runs, 10 refits, 100 bootstraps +# ./submit_runs.sh -n 5 -k 10 -b 100 # 5 runs, 10 refits, 100 diagnostics/refit # # Anything after the known flags is forwarded to the runner (e.g. --bert-only). # Override defaults via env: POOL=... EXPERIMENT=... TEMPLATE=... ./submit_runs.sh ... diff --git a/experiments/semisynthetic_simulation/docs/study.md b/experiments/semisynthetic_simulation/docs/study.md index def6e095..60e6992b 100644 --- a/experiments/semisynthetic_simulation/docs/study.md +++ b/experiments/semisynthetic_simulation/docs/study.md @@ -10,7 +10,8 @@ Fixed shared cohort → index_dates.csv + cohort_config.yaml + pids │ per outer run s (own seed): simulate outcomes → prepare - → K refits: fit → calibrate → B patient bootstraps + → K bootstrap refits: fit → calibrate → estimate + → optional B post-estimation bootstraps for diagnostics (BERT and/or CatBoost baseline) │ summarize.py → bias / SD / SE-calibration / coverage, per estimator × outcome @@ -18,11 +19,14 @@ Fixed shared cohort → index_dates.csv + cohort_config.yaml + pids - **Outer runs** redraw the simulated outcomes (Monte Carlo over the DGP). 1 is enough for "does it recover the effect"; ~10–20 for a coverage sanity check. -- **Inner refits (`-k`)** use different model seeds and reshuffled CV folds, - yielding K independently fitted propensity-score models. -- **Patient bootstraps (`-b`, default 100)** resample the original analysis - cohort with replacement conditional on each fitted model. The summarizer - combines the K × B draws into one SE and CI per outer simulation replicate. +- **Bootstrap refits (`-k`)** use different model seeds and resample each + training fold with replacement. The prepared validation folds are reused + and not resampled, so every original patient still receives one out-of-fold + propensity score. + The K resulting point estimates provide the main uncertainty estimate. +- **Post-estimation bootstraps (`-b`, default 100)** resample the analysis + cohort conditional on each fitted model. These cheap B draws are retained + for single-run diagnostics, but are not used by the main study summary. ## What to set (once) @@ -53,7 +57,7 @@ python -m experiments.semisynthetic_simulation.python_scripts.summarize \ ### 2. Full run -A few outer runs, K seeded refits and B patient bootstraps each: +A few outer runs, K bootstrap refits and B optional diagnostic bootstraps: ```bash ./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh -n 5 -k 10 -b 100 @@ -66,8 +70,8 @@ A few outer runs, K seeded refits and B patient bootstraps each: | flag | meaning | default | |------|---------|---------| | `-n` | outer runs (parallel jobs) | 1 | -| `-k` | independently seeded model refits per run | 1 | -| `-b` | patient bootstrap samples per refit | 100 | +| `-k` | bootstrap model refits per run | 1 | +| `-b` | diagnostic post-estimation bootstraps per refit | 100 | | `--baseline-only` / `--bert-only` | restrict to one model | both | Override compute/experiment via env: `POOL= EXPERIMENT=my_exp ./submit_runs.sh ...` @@ -84,10 +88,10 @@ Override compute/experiment via env: `POOL= EXPERIMENT=my_exp ./submit_runs ├── models/{bert,baseline}/... └── estimate/{bert,baseline}/ ├── estimate_results.csv - └── bootstrap_results.csv ← B patient-level bootstrap estimates + └── bootstrap_results.csv ← diagnostic B post-estimation draws ``` -Running `summarize.py` writes `replicate_estimates.csv` (one K × B aggregate -per outer run) and `summary.csv` (bias, empirical SD, mean SE, SE calibration, -and coverage across outer runs). Common support is trimmed at the 0.1st and -99.9th percentiles within treatment arms before patient resampling. +Running `summarize.py` writes `replicate_estimates.csv` (one aggregate of the K +bootstrap-refit point estimates per outer run) and `summary.csv` (bias, +empirical SD, mean SE, SE calibration, and coverage across outer runs). Common +support is trimmed at the 0.1st and 99.9th percentiles within treatment arms. diff --git a/experiments/semisynthetic_simulation/python_scripts/run_study.py b/experiments/semisynthetic_simulation/python_scripts/run_study.py index 4c72c566..24677d26 100644 --- a/experiments/semisynthetic_simulation/python_scripts/run_study.py +++ b/experiments/semisynthetic_simulation/python_scripts/run_study.py @@ -10,9 +10,10 @@ (BERT) and/or the CatBoost baseline. Inner refits: -- Each refit uses a distinct model seed and fold reshuffling. +- Each refit uses a distinct model seed and patient-bootstrap training sample. +- Validation patients are not resampled, so every patient receives one OOF score. - Conditional on each fitted model, estimate draws B patient-level bootstrap - samples. The summarizer combines all K x B estimates. + samples for diagnostics. The main summary uses only the K refit estimates. Each Azure job runs one outer run (--run-id run_NN); the outer loop = the set of parallel jobs submitted by bash_scripts/submit_runs.sh. @@ -100,10 +101,11 @@ def run_outer(args, run_id: str, seed: int): refit_seed = seed * 1000 + k def configure_refit(config, _seed=refit_seed): - """Apply the model seed and reshuffle folds without resampling patients.""" - config.setdefault("data", {})["reshuffle"] = True - config["data"]["reshuffle_seed"] = _seed + """Apply the model seed and bootstrap each fold's training patients.""" + config.setdefault("data", {})["reshuffle"] = False + config["data"]["bootstrap_seed"] = _seed config["seed"] = _seed + config["bootstrap"] = True def configure_estimation(config, _seed=refit_seed): estimator = config.setdefault("estimator", {}) @@ -200,7 +202,7 @@ def parse_arguments(argv=None) -> argparse.Namespace: dest="inner_runs", type=int, default=1, - help="Independently seeded model refits per outer run.", + help="Patient-bootstrap model refits per outer run.", ) parser.add_argument( "--n-bootstrap", @@ -208,7 +210,7 @@ def parse_arguments(argv=None) -> argparse.Namespace: dest="n_bootstrap", type=int, default=DEFAULT_BOOTSTRAPS, - help="Patient-level bootstrap samples per fitted propensity model.", + help="Diagnostic post-estimation bootstrap samples per fitted model.", ) parser.add_argument("--base-seed", dest="base_seed", type=int, default=42) diff --git a/experiments/semisynthetic_simulation/python_scripts/study_summary.py b/experiments/semisynthetic_simulation/python_scripts/study_summary.py index a02856f3..93295631 100644 --- a/experiments/semisynthetic_simulation/python_scripts/study_summary.py +++ b/experiments/semisynthetic_simulation/python_scripts/study_summary.py @@ -14,7 +14,7 @@ MODEL_NAMES = ("bert", "baseline") GROUP_COLUMNS = ["model", "method", OUTCOME] -REPLICATE_GROUP_COLUMNS = ["model", "run_id", "method", OUTCOME] +REPLICATE_GROUP_COLUMNS = ["model", "run_id", "method", OUTCOME, E.effect_type] def _tag_from_path(path: Path) -> dict: @@ -49,82 +49,45 @@ def load_bootstrap_results(study_dir: Path) -> pd.DataFrame: return load_tagged_results(study_dir, BOOTSTRAP_RESULTS_FILE) -def aggregate_replicates( - results: pd.DataFrame, bootstrap_results: pd.DataFrame -) -> pd.DataFrame: - """Combine K point estimates and K x B patient-bootstrap estimates per run.""" - available = bootstrap_results[GROUP_COLUMNS].drop_duplicates() - results = results.merge(available, on=GROUP_COLUMNS, how="inner") +def aggregate_replicates(results: pd.DataFrame) -> pd.DataFrame: + """Combine K bootstrap-refit point estimates into one estimate per run.""" + results = results.dropna(subset=[E.true_effect, E.effect_type]) rows = [] for keys, group in results.groupby(REPLICATE_GROUP_COLUMNS): tags = dict(zip(REPLICATE_GROUP_COLUMNS, keys)) - samples = _select_bootstrap_group(bootstrap_results, tags) - _validate_nested_samples(group, samples, tags) - - effect_type = _single_value(samples["effect_type"], "effect_type", tags) + effect_type = tags[E.effect_type] point = group[E.effect].mean() true_effect = group[E.true_effect].mean() if effect_type in {"RR", "RRT"}: - uncertainty = _risk_ratio_uncertainty(group, samples) + uncertainty = _risk_ratio_uncertainty(group) point = uncertainty.pop("effect") else: - uncertainty = _difference_uncertainty(point, samples[E.effect]) + uncertainty = _difference_uncertainty(point, group[E.effect]) + + lower = uncertainty[E.CI95_lower] + upper = uncertainty[E.CI95_upper] + covered = ( + lower <= true_effect <= upper + if np.isfinite(lower) and np.isfinite(upper) + else np.nan + ) rows.append( { **tags, - "effect_type": effect_type, "n_refits": group["inner_id"].nunique(), - "n_bootstrap_per_refit": samples.groupby("inner_id").size().iloc[0], - "n_bootstrap": len(samples), E.true_effect: true_effect, E.effect: point, **uncertainty, - "covered": ( - uncertainty[E.CI95_lower] - <= true_effect - <= uncertainty[E.CI95_upper] - ), + "covered": covered, } ) return pd.DataFrame(rows).sort_values([OUTCOME, "model", "method", "run_id"]) -def _select_bootstrap_group(samples: pd.DataFrame, tags: dict) -> pd.DataFrame: - mask = pd.Series(True, index=samples.index) - for column, value in tags.items(): - mask &= samples[column] == value - return samples[mask] - - -def _validate_nested_samples( - estimates: pd.DataFrame, samples: pd.DataFrame, tags: dict -) -> None: - expected_refits = set(estimates["inner_id"]) - actual_refits = set(samples["inner_id"]) - if actual_refits != expected_refits: - raise ValueError( - f"Bootstrap refits do not match point-estimate refits for {tags}: " - f"{sorted(actual_refits)} != {sorted(expected_refits)}" - ) - - counts = samples.groupby("inner_id").size() - if counts.nunique() != 1: - raise ValueError(f"Unequal bootstrap counts across refits for {tags}: {counts}") - if not np.isfinite(samples[[E.effect, E.effect_1, E.effect_0]]).all().all(): - raise ValueError(f"Non-finite bootstrap estimates found for {tags}") - - -def _single_value(series: pd.Series, name: str, tags: dict): - values = series.drop_duplicates() - if len(values) != 1: - raise ValueError(f"Expected one {name} for {tags}, found {values.tolist()}") - return values.iloc[0] - - def _difference_uncertainty(point: float, samples: pd.Series) -> dict: standard_error = samples.std(ddof=1) margin = 1.96 * standard_error @@ -136,14 +99,14 @@ def _difference_uncertainty(point: float, samples: pd.Series) -> dict: } -def _risk_ratio_uncertainty(estimates: pd.DataFrame, samples: pd.DataFrame) -> dict: +def _risk_ratio_uncertainty(estimates: pd.DataFrame) -> dict: p1 = estimates[E.effect_1].mean() p0 = estimates[E.effect_0].mean() if not 0 < p1 < 1 or not 0 < p0 < 1: raise ValueError(f"Risk-ratio probabilities must lie in (0, 1): {p1=}, {p0=}") - sample_p1 = samples[E.effect_1] - sample_p0 = samples[E.effect_0] + sample_p1 = estimates[E.effect_1] + sample_p0 = estimates[E.effect_0] if ( not sample_p1.between(0, 1, inclusive="neither").all() or not sample_p0.between(0, 1, inclusive="neither").all() diff --git a/experiments/semisynthetic_simulation/python_scripts/summarize.py b/experiments/semisynthetic_simulation/python_scripts/summarize.py index d385320c..294781d9 100644 --- a/experiments/semisynthetic_simulation/python_scripts/summarize.py +++ b/experiments/semisynthetic_simulation/python_scripts/summarize.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Summarize nested model-refit and patient-bootstrap study results. +"""Summarize bootstrap-refit study results across simulation runs. Usage: python -m experiments.semisynthetic_simulation.python_scripts.summarize \ @@ -13,7 +13,6 @@ from experiments.semisynthetic_simulation.python_scripts.study_summary import ( aggregate_replicates, - load_bootstrap_results, load_results, summarize_performance, ) @@ -29,8 +28,7 @@ def main(): study_dir = Path(args.study_dir) results = load_results(study_dir) - bootstrap_results = load_bootstrap_results(study_dir) - replicates = aggregate_replicates(results, bootstrap_results) + replicates = aggregate_replicates(results) table = summarize_performance(replicates) out = Path(args.out) if args.out else study_dir / "summary.csv" diff --git a/tests/test_main_causal/test_baseline_folds.py b/tests/test_main_causal/test_baseline_folds.py index 427e89e3..dc020b81 100644 --- a/tests/test_main_causal/test_baseline_folds.py +++ b/tests/test_main_causal/test_baseline_folds.py @@ -1,4 +1,4 @@ -"""Tests for baseline refit fold reshuffling.""" +"""Tests for baseline bootstrap-refit folds.""" import logging import os @@ -22,8 +22,8 @@ from corebehrt.modules.setup.config import Config -class TestBaselineFoldReshuffling(unittest.TestCase): - def test_reshuffle_uses_configured_seed_and_all_patients(self): +class TestBaselineBootstrapFolds(unittest.TestCase): + def test_bootstrap_resamples_training_and_keeps_validation_complete(self): prepared_dir = tempfile.mkdtemp() model_dir = tempfile.mkdtemp() pids = list(range(20)) @@ -35,18 +35,28 @@ def test_reshuffle_uses_configured_seed_and_all_patients(self): cfg = Config( { "paths": {"prepared_data": prepared_dir, "model": model_dir}, - "data": {"reshuffle": True, "reshuffle_seed": 123}, + "data": { + "reshuffle": False, + "bootstrap_seed": 123, + }, + "bootstrap": True, } ) - reshuffled = handle_folds(cfg, logging.getLogger("test_baseline_folds")) - - self.assertNotEqual(reshuffled, original) - for fold in reshuffled: - self.assertEqual(set(fold[TRAIN_KEY]) | set(fold[VAL_KEY]), set(pids)) - self.assertTrue( - set(fold[TRAIN_KEY]).isdisjoint(set(fold[VAL_KEY])) - ) + folds = handle_folds(cfg, logging.getLogger("test_baseline_folds")) + + self.assertNotEqual(folds, original) + validation_pids = [] + duplicate_training_found = False + for fold, original_fold in zip(folds, original): + validation_pids.extend(fold[VAL_KEY]) + self.assertEqual(fold[VAL_KEY], original_fold[VAL_KEY]) + self.assertEqual(len(fold[VAL_KEY]), len(set(fold[VAL_KEY]))) + self.assertTrue(set(fold[TRAIN_KEY]).isdisjoint(set(fold[VAL_KEY]))) + duplicate_training_found |= len(fold[TRAIN_KEY]) > len(set(fold[TRAIN_KEY])) + self.assertEqual(set(validation_pids), set(pids)) + self.assertEqual(len(validation_pids), len(pids)) + self.assertTrue(duplicate_training_found) if __name__ == "__main__": diff --git a/tests/test_main_causal/test_finetune_bootstrap_folds.py b/tests/test_main_causal/test_finetune_bootstrap_folds.py new file mode 100644 index 00000000..950b6fa5 --- /dev/null +++ b/tests/test_main_causal/test_finetune_bootstrap_folds.py @@ -0,0 +1,58 @@ +"""Tests for BERT bootstrap-refit fold construction.""" + +import logging +import os +import sys +import tempfile +import types +import unittest + +import torch + +sys.modules.setdefault("umap", types.ModuleType("umap")) + +from corebehrt.constants.data import TRAIN_KEY, VAL_KEY +from corebehrt.constants.paths import FOLDS_FILE +from corebehrt.main_causal.finetune_exp_y import handle_folds +from corebehrt.modules.setup.config import Config + + +class TestFinetuneBootstrapFolds(unittest.TestCase): + def test_training_is_bootstrapped_and_validation_covers_cohort(self): + prepared_dir = tempfile.mkdtemp() + model_dir = tempfile.mkdtemp() + pids = list(range(20)) + original = [ + {TRAIN_KEY: pids[10:], VAL_KEY: pids[:10]}, + {TRAIN_KEY: pids[:10], VAL_KEY: pids[10:]}, + ] + torch.save(original, os.path.join(prepared_dir, FOLDS_FILE)) + cfg = Config( + { + "paths": {"prepared_data": prepared_dir, "model": model_dir}, + "data": { + "reshuffle": False, + "bootstrap_seed": 123, + }, + "bootstrap": True, + } + ) + + folds = handle_folds( + cfg, [], pids, logging.getLogger("test_finetune_bootstrap_folds") + ) + + validation_pids = [pid for fold in folds for pid in fold[VAL_KEY]] + self.assertEqual( + [fold[VAL_KEY] for fold in folds], + [fold[VAL_KEY] for fold in original], + ) + self.assertEqual(set(validation_pids), set(pids)) + self.assertEqual(len(validation_pids), len(pids)) + self.assertTrue( + any(len(fold[TRAIN_KEY]) > len(set(fold[TRAIN_KEY])) for fold in folds) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modules/test_causal_dataset_resampling.py b/tests/test_modules/test_causal_dataset_resampling.py new file mode 100644 index 00000000..d25c43da --- /dev/null +++ b/tests/test_modules/test_causal_dataset_resampling.py @@ -0,0 +1,34 @@ +"""Tests for multiplicity-preserving patient resampling.""" + +import unittest + +from corebehrt.modules.preparation.causal.dataset import ( + CausalPatientData, + CausalPatientDataset, +) + + +class TestCausalPatientResampling(unittest.TestCase): + def test_resample_preserves_order_and_duplicates(self): + patients = [ + CausalPatientData(pid=1, concepts=[], abspos=[], segments=[], ages=[]), + CausalPatientData(pid=2, concepts=[], abspos=[], segments=[], ages=[]), + CausalPatientData(pid=3, concepts=[], abspos=[], segments=[], ages=[]), + ] + dataset = CausalPatientDataset(patients) + + resampled = dataset.resample_by_pids([2, 1, 2, 3, 2]) + + self.assertEqual(resampled.get_pids(), [2, 1, 2, 3, 2]) + + def test_resample_rejects_unknown_patient(self): + dataset = CausalPatientDataset( + [CausalPatientData(pid=1, concepts=[], abspos=[], segments=[], ages=[])] + ) + + with self.assertRaises(KeyError): + dataset.resample_by_pids([1, 2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_semisynthetic_study_summary.py b/tests/test_semisynthetic_study_summary.py index ccd1544f..6d92ebc1 100644 --- a/tests/test_semisynthetic_study_summary.py +++ b/tests/test_semisynthetic_study_summary.py @@ -20,6 +20,7 @@ def _estimate_rows(run_id, effects, true_effect): "inner_id": ["k_01", "k_02"], E.method: "IPW", "outcome": "OUTCOME", + E.effect_type: "ATE", E.effect: effects, E.effect_1: np.asarray(effects) + 0.2, E.effect_0: 0.2, @@ -28,29 +29,8 @@ def _estimate_rows(run_id, effects, true_effect): ) -def _bootstrap_rows(run_id, samples): - rows = [] - for inner_id, effects in zip(["k_01", "k_02"], samples): - for bootstrap_id, effect in enumerate(effects, 1): - rows.append( - { - "model": "bert", - "run_id": run_id, - "inner_id": inner_id, - E.method: "IPW", - "outcome": "OUTCOME", - "effect_type": "ATE", - "bootstrap_id": bootstrap_id, - E.effect: effect, - E.effect_1: effect + 0.2, - E.effect_0: 0.2, - } - ) - return pd.DataFrame(rows) - - class TestNestedStudySummary(unittest.TestCase): - def test_aggregates_k_times_b_with_run_specific_truth(self): + def test_aggregates_bootstrap_refits_with_run_specific_truth(self): results = pd.concat( [ _estimate_rows("run_01", [0.10, 0.12], 0.11), @@ -58,19 +38,12 @@ def test_aggregates_k_times_b_with_run_specific_truth(self): ], ignore_index=True, ) - bootstraps = pd.concat( - [ - _bootstrap_rows("run_01", [[0.08, 0.12], [0.10, 0.14]]), - _bootstrap_rows("run_02", [[0.18, 0.22], [0.20, 0.24]]), - ], - ignore_index=True, - ) - replicates = aggregate_replicates(results, bootstraps) - self.assertEqual(list(replicates["n_bootstrap"]), [4, 4]) + replicates = aggregate_replicates(results) + self.assertEqual(list(replicates["n_refits"]), [2, 2]) self.assertAlmostEqual(replicates.iloc[0][E.effect], 0.11) self.assertAlmostEqual( - replicates.iloc[0][E.std_err], np.std([0.08, 0.12, 0.10, 0.14], ddof=1) + replicates.iloc[0][E.std_err], np.std([0.10, 0.12], ddof=1) ) summary = summarize_performance(replicates).iloc[0] @@ -78,24 +51,19 @@ def test_aggregates_k_times_b_with_run_specific_truth(self): self.assertAlmostEqual(summary["bias"], 0.01) self.assertAlmostEqual(summary[E.true_effect], 0.15) - def test_rejects_unequal_bootstrap_counts(self): - results = _estimate_rows("run_01", [0.10, 0.12], 0.11) - bootstraps = _bootstrap_rows("run_01", [[0.08], [0.10, 0.14]]) - - with self.assertRaisesRegex(ValueError, "Unequal bootstrap counts"): - aggregate_replicates(results, bootstraps) - def test_risk_ratio_uses_log_scale_interval(self): results = _estimate_rows("run_01", [2.0, 2.0], 2.0) + results[E.effect_type] = "RR" results[E.effect_1] = [0.4, 0.4] results[E.effect_0] = [0.2, 0.2] - bootstraps = _bootstrap_rows("run_01", [[2.0, 2.1], [1.9, 2.0]]) - bootstraps["effect_type"] = "RR" - bootstraps[E.effect_1] = [0.39, 0.41, 0.38, 0.40] - bootstraps[E.effect_0] = [0.20, 0.20, 0.19, 0.20] + results.loc[0, E.effect_1] = 0.39 + results.loc[0, E.effect_0] = 0.19 - replicate = aggregate_replicates(results, bootstraps).iloc[0] - self.assertAlmostEqual(replicate[E.effect], 2.0) + replicate = aggregate_replicates(results).iloc[0] + self.assertAlmostEqual( + replicate[E.effect], + results[E.effect_1].mean() / results[E.effect_0].mean(), + ) self.assertGreater(replicate["std_err_log"], 0) self.assertGreater(replicate[E.CI95_lower], 0)