add semi-synthetic causal simulator with observed treatment

.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

.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 }}

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,
+    )

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)

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,
+    )

corebehrt/azure/main/job.py

@@ -62,6 +62,9 @@ def add_parser(subparsers) -> None:
             "evaluate_xgboost",
             "get_pat_counts_by_code",
             "run_batch_experiments",
+            "simulate_semisynthetic",
+            "calibrate_semisynthetic",
+            "run_semisynthetic_study",
         },
         help="Job to run.",
     )

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

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"

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"

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:

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
 
 

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