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 }} 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/run_semisynthetic_study.py b/corebehrt/azure/components/run_semisynthetic_study.py new file mode 100644 index 00000000..2190a1b1 --- /dev/null +++ b/corebehrt/azure/components/run_semisynthetic_study.py @@ -0,0 +1,51 @@ +"""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"}, + "cohort": {"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, + "--cohort", + cfg.paths.cohort, + "--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/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/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/configs/causal/simulate_semisynthetic.yaml b/corebehrt/configs/causal/simulate_semisynthetic.yaml new file mode 100644 index 00000000..c8992987 --- /dev/null +++ b/corebehrt/configs/causal/simulate_semisynthetic.yaml @@ -0,0 +1,72 @@ +# -------------------------------------------------------------------------- +# 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 + +# -------------------------------------------------------------------------- +# Outcomes +# -------------------------------------------------------------------------- +outcomes: + OUTCOME: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + 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 + 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 + coefficients: + disease_burden: 0.2 + age: 0.4 + noise_scale: 0.1 + treatment_effect: + mode: constant + delta: 0.0 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/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/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/calibrate_semisynthetic.py b/corebehrt/main_causal/calibrate_semisynthetic.py new file mode 100644 index 00000000..6275007c --- /dev/null +++ b/corebehrt/main_causal/calibrate_semisynthetic.py @@ -0,0 +1,229 @@ +"""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 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.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) + + # 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": []} + all_tau = {} # outcome_name -> [] + + for shard, _ in shard_loader(): + result = simulator.extract_features_and_probabilities(shard) + if result is None: + continue + 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 in sim_config.outcomes: + all_probas.setdefault(outcome_name, {"P0": [], "P1": []}) + 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.") + 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/finetune_exp_y.py b/corebehrt/main_causal/finetune_exp_y.py index b70df4a2..b0a08f9d 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, @@ -14,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 ( @@ -33,6 +37,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() @@ -103,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() @@ -115,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" @@ -142,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" ) @@ -186,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 @@ -206,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 9cf77e10..e3665df6 100644 --- a/corebehrt/main_causal/helper/train_baseline.py +++ b/corebehrt/main_causal/helper/train_baseline.py @@ -28,13 +28,16 @@ from corebehrt.constants.paths import ( FOLDS_FILE, ) +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 @@ -177,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 @@ -246,6 +251,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 +338,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 +365,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!") @@ -443,17 +451,18 @@ 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=42, + random_state=cfg.get("seed", 42), ) logger.info( 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( @@ -474,6 +483,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 +530,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 +551,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, ) @@ -641,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() @@ -688,6 +703,7 @@ def nested_cv_loop( target_name, i, prediction_storage, + cfg.get("seed", 42) + i, ) all_unbiased_scores.append(unbiased_auc) @@ -714,11 +730,27 @@ 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") + 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/main_causal/simulate_semisynthetic.py b/corebehrt/main_causal/simulate_semisynthetic.py new file mode 100644 index 00000000..2ededf7c --- /dev/null +++ b/corebehrt/main_causal/simulate_semisynthetic.py @@ -0,0 +1,124 @@ +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 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") + + +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) + + # 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) + + # Pass 2: simulate outcomes using globally standardized features + 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. + + 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) + 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 ---") + + 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__": + args = get_args(CONFIG_PATH) + main_simulate(args.config_path) diff --git a/corebehrt/modules/causal/estimate.py b/corebehrt/modules/causal/estimate.py index 3f4e01ba..e294f5e2 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,17 @@ 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) + effect_df[EffectColumns.effect_type] = self.effect_type - 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 +182,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 +227,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, @@ -326,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: @@ -381,8 +443,19 @@ 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) + 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( + "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 +536,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 +548,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/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/corebehrt/modules/simulation/config_semisynthetic.py b/corebehrt/modules/simulation/config_semisynthetic.py new file mode 100644 index 00000000..1c97591c --- /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 + cohort: str = None + index_dates: str = None + + +@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 + + +@dataclass +class OutcomeModelConfig: + """Outcome model: eta^(0) = beta_0 + f(r_i).""" + + run_in_days: int = 1 + beta_0: float = -2.0 + 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..0006bf97 --- /dev/null +++ b/corebehrt/modules/simulation/oracle_features.py @@ -0,0 +1,267 @@ +"""Extract hand-crafted oracle features from pre-index patient histories.""" + +import logging + +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") + +# 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, + pids: np.ndarray, + index_dates: pd.Series, + feature_config: FeatureConfig, +) -> 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 + 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 + """ + prefixes = feature_config.code_prefixes + + 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) + + 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 + return features_df + + +def standardize_features(features_df, means=None, stds=None): + """Z-score features using provided or computed statistics. + + Args: + features_df: raw features DataFrame + means: per-feature means (computed from df if None) + stds: per-feature stds (computed from df if None) + + 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 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _filter_by_prefix_and_window(history_df, index_dates, prefix, window_days): + """Filter events matching code prefix within lookback window per patient.""" + 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 + ) + 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, None, 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, 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() + 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) + return age_series.fillna(mean_age) + + +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() + 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() + 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) + + +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, None, burst_window_days + ) + lookback_filtered = _filter_by_prefix_and_window( + history_df, index_dates, None, 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) + + # 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) + 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 new file mode 100644 index 00000000..520a30d8 --- /dev/null +++ b/corebehrt/modules/simulation/semisynthetic_simulator.py @@ -0,0 +1,679 @@ +"""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, + EXPOSURES_FILE, + INDEX_DATE_MATCHING_FILE, +) +from corebehrt.constants.paths import INDEX_DATES_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 ( + ORACLE_FEATURE_NAMES, + extract_oracle_features, + standardize_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) + self._global_means = None + 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. + + 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.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): + 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}") + return pids + + 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. 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 or self.config.paths.cohort + 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. + + 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.""" + 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. + + 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 + ) + 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 + ) + 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 + ) + + # 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, + ) + + 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 + ) + if self._global_means is not None: + features_df, _, _ = standardize_features( + features_df, self._global_means, self._global_stds + ) + + 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 + # ------------------------------------------------------------------ + + 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.""" + 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]"), + ) + + # 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]) + + 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]: + n_patients = len(pids) + ite_records = {PID_COL: pids} + cf_records = {PID_COL: pids, EXPOSURE_COL: is_exposed.astype(int)} + all_factual_events = [] + + 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) + + 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 + + def _compute_eta_0( + self, + features_df: pd.DataFrame, + outcome_model: OutcomeModelConfig, + ) -> np.ndarray: + """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.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, + ) -> 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...") + self._calculate_theoretical_roc_auc(cf_records, is_exposed, output_dir) + + 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) + + 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, + ) + + # ------------------------------------------------------------------ + # 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/experiments/semisynthetic_simulation/README.md b/experiments/semisynthetic_simulation/README.md new file mode 100644 index 00000000..db2d97b4 --- /dev/null +++ b/experiments/semisynthetic_simulation/README.md @@ -0,0 +1,38 @@ +# 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 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) | +| 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/__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/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 new file mode 100644 index 00000000..b8a65770 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/estimate.yaml @@ -0,0 +1,31 @@ +# 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"] + effect_type: "ATE" + n_bootstrap: 100 + common_support_threshold: 0.001 + clip_percentile: 1.0 + +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/estimate_baseline.yaml b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml new file mode 100644 index 00000000..c19def8f --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/estimate_baseline.yaml @@ -0,0 +1,29 @@ +# Stage 2c (baseline): estimate causal effects from calibrated baseline predictions. +# n_bootstrap and its seed are set by the study runner. +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"] + effect_type: "ATE" + n_bootstrap: 100 + common_support_threshold: 0.001 + clip_percentile: 1.0 + +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..e46a9550 --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/prepare.yaml @@ -0,0 +1,45 @@ +# 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 + +paths: + ## INPUTS + features: "{{FEATURES}}" + tokenized: "{{TOKENIZED}}" + 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 + + # Exposure is read from the cohort (cohort/exposures.csv), so it is not set here. + + ## 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/simulate.yaml b/experiments/semisynthetic_simulation/base_configs/simulate.yaml new file mode 100644 index 00000000..949aa2cb --- /dev/null +++ b/experiments/semisynthetic_simulation/base_configs/simulate.yaml @@ -0,0 +1,69 @@ +# 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 must cover the patients in the cohort (filtered by index_dates). + splits: ["tuning", "train"] + outcomes: "{{RUN_DIR}}/simulated_outcomes" + cohort: "{{COHORT}}" # requires index_dates.csv and exposures.csv + +min_num_codes: 3 +# 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: + 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 + +# 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: &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.0 + treatment_effect: + mode: constant + delta: 0.0 + + OUTCOME_MEDIUM: + outcome_model: + run_in_days: 1 + beta_0: -2.0 + 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 new file mode 100755 index 00000000..4fd9702f --- /dev/null +++ b/experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh @@ -0,0 +1,51 @@ +#!/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 + 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 # 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 ... +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=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 +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-bootstrap $N_BOOTSTRAP ${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" \ + --bash-args "$BASH_ARGS" +done 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/docs/study.md b/experiments/semisynthetic_simulation/docs/study.md new file mode 100644 index 00000000..60e6992b --- /dev/null +++ b/experiments/semisynthetic_simulation/docs/study.md @@ -0,0 +1,97 @@ +# Running the Study on Azure + +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`). + +``` +Fixed shared cohort → index_dates.csv + cohort_config.yaml + pids + │ + per outer run s (own seed): + simulate outcomes → prepare + → 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 +``` + +- **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. +- **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) + +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__`. + +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. Phase 1 — first shot (does it recover the effect?) + +One run, single fit, method + baseline: + +```bash +./experiments/semisynthetic_simulation/bash_scripts/submit_runs.sh +``` + +Then summarize and check θ̂ vs θ* and CI coverage: + +```bash +python -m experiments.semisynthetic_simulation.python_scripts.summarize \ + --study-dir /run_01 +``` + +### 2. Full run + +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 +``` +(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` | outer runs (parallel jobs) | 1 | +| `-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 ...` +(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, 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 + └── bootstrap_results.csv ← diagnostic B post-estimation draws +``` + +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/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() diff --git a/experiments/semisynthetic_simulation/job_config_template.yaml b/experiments/semisynthetic_simulation/job_config_template.yaml new file mode 100644 index 00000000..3b8a1605 --- /dev/null +++ b/experiments/semisynthetic_simulation/job_config_template.yaml @@ -0,0 +1,15 @@ +# 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 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/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__" 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..24677d26 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/run_study.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +""" +Runner for the semi-synthetic simulation study. + +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: +- 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 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. +""" + +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.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 + +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" +DEFAULT_BOOTSTRAPS = 100 + + +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 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) + + shared = { + "{{MEDS}}": args.meds, + "{{FEATURES}}": args.features, + "{{TOKENIZED}}": args.tokenized, + "{{PRETRAIN_MODEL}}": args.pretrain_model, + "{{COHORT}}": args.cohort, + "{{RUN_DIR}}": str(run_dir), + } + logger.info("=" * 70) + 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) ---- + main_simulate( + fill_config( + 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 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} / refit {k}/{args.inner_runs} ({inner_id}) ---") + + refit_seed = seed * 1000 + k + + def configure_refit(config, _seed=refit_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", {}) + 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( + "bert", + base, + repl, + config_dir, + inner_id, + configure_refit, + configure_estimation, + main_finetune, + "finetune.yaml", + "calibrate.yaml", + "estimate.yaml", + ) + if not args.bert_only: + _run_model( + "baseline", + base, + repl, + config_dir, + inner_id, + configure_refit, + configure_estimation, + main_baseline, + "train_baseline.yaml", + "calibrate_baseline.yaml", + "estimate_baseline.yaml", + ) + + logger.info(f"OUTER RUN {run_id} complete") + + +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) + ) + + parser.add_argument( + "--run-id", + dest="run_id", + default=None, + 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=1, + help="Patient-bootstrap model refits per outer run.", + ) + parser.add_argument( + "--n-bootstrap", + "-b", + dest="n_bootstrap", + type=int, + default=DEFAULT_BOOTSTRAPS, + help="Diagnostic post-estimation bootstrap samples per fitted model.", + ) + parser.add_argument("--base-seed", dest="base_seed", type=int, default=42) + + parser.add_argument("--bert-only", action="store_true") + parser.add_argument("--baseline-only", action="store_true") + + 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 + + +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() 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..93295631 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/study_summary.py @@ -0,0 +1,155 @@ +"""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, E.effect_type] + + +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) -> 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)) + 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) + point = uncertainty.pop("effect") + else: + 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, + "n_refits": group["inner_id"].nunique(), + E.true_effect: true_effect, + E.effect: point, + **uncertainty, + "covered": covered, + } + ) + + return pd.DataFrame(rows).sort_values([OUTCOME, "model", "method", "run_id"]) + + +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) -> 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 = 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() + ): + 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 new file mode 100644 index 00000000..294781d9 --- /dev/null +++ b/experiments/semisynthetic_simulation/python_scripts/summarize.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +"""Summarize bootstrap-refit study results across simulation runs. + +Usage: + python -m experiments.semisynthetic_simulation.python_scripts.summarize \ + --study-dir [--out ] +""" + +import argparse +from pathlib import Path + +import pandas as pd + +from experiments.semisynthetic_simulation.python_scripts.study_summary import ( + aggregate_replicates, + load_results, + summarize_performance, +) + + +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) + replicates = aggregate_replicates(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}") + print(table.to_string(index=False)) + print(f"\nSaved summary to {out}") + + +if __name__ == "__main__": + main() 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..dc020b81 --- /dev/null +++ b/tests/test_main_causal/test_baseline_folds.py @@ -0,0 +1,63 @@ +"""Tests for baseline bootstrap-refit folds.""" + +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 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)) + 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, 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__": + unittest.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_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/__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..5ae4966b --- /dev/null +++ b/tests/test_modules/test_simulation/test_oracle_features.py @@ -0,0 +1,224 @@ +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 ( + extract_oracle_features, + standardize_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(**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 = extract_oracle_features(history_df, pids, index_dates, config) + self.assertEqual(features_df.shape[0], 5) + self.assertEqual(features_df.shape[1], 10) + + +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() + 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(standardized_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..18503d34 --- /dev/null +++ b/tests/test_modules/test_simulation/test_semisynthetic.py @@ -0,0 +1,325 @@ +"""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.oracle_features import ( + ORACLE_FEATURE_NAMES, + extract_oracle_features, +) +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, + coefficients={"disease_burden": 0.3}, + noise_scale=noise_scale, + ), + treatment_effect=TreatmentEffectConfig(mode=mode, delta=delta), + ) + return SemiSyntheticSimulationConfig( + paths=paths, + features=FeatureConfig(), + 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) + + +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())) + + +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..6d92ebc1 --- /dev/null +++ b/tests/test_semisynthetic_study_summary.py @@ -0,0 +1,72 @@ +"""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_type: "ATE", + E.effect: effects, + E.effect_1: np.asarray(effects) + 0.2, + E.effect_0: 0.2, + E.true_effect: true_effect, + } + ) + + +class TestNestedStudySummary(unittest.TestCase): + def test_aggregates_bootstrap_refits_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, + ) + + 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.10, 0.12], 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_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] + results.loc[0, E.effect_1] = 0.39 + results.loc[0, E.effect_0] = 0.19 + + 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) + + +if __name__ == "__main__": + unittest.main()