research-tool-cli 0.1.0__py3-none-any.whl

tests/test_cox_ph_plan.py

@@ -0,0 +1,498 @@
+"""Tests for Cox PH plan parsing and validation."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import shutil
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from core.cli.main import cmd_plan, cmd_analyze, cmd_lock
+from core.database import DATA_ROOT, get_connection, init_db
+from core.planning.study_plan import StudyPlan
+from core.planning.lock import lock_plan
+from core.masking.gate import unmask_study
+from core.stats.inferential import _cox_ph_model
+
+STUDY_ID = "test_cox_plan"
+RAW = f"raw_{STUDY_ID}"
+
+
+def _setup_db(extra_rows: list[tuple] | None = None):
+    study_path = DATA_ROOT / STUDY_ID
+    if study_path.exists():
+        shutil.rmtree(study_path)
+    study_path.mkdir(parents=True)
+    conn = get_connection(STUDY_ID)
+    init_db(conn)
+    conn.execute(
+        "INSERT OR REPLACE INTO studies (id, name, created_at, data_dir, study_type) VALUES (?, ?, ?, ?, ?)",
+        (STUDY_ID, "Cox PH Plan Test", "2025-01-01", str(study_path), "cohort"),
+    )
+    conn.execute(
+        f"CREATE TABLE IF NOT EXISTS {RAW} "
+        f"(row_id INTEGER PRIMARY KEY, treatment_arm TEXT, pfs_days TEXT, pfs_event TEXT, age TEXT)"
+    )
+    conn.executemany(
+        f"INSERT INTO {RAW} (treatment_arm, pfs_days, pfs_event, age) VALUES (?, ?, ?, ?)",
+        extra_rows or [
+            ("A", "100", "1", "65"),
+            ("A", "200", "0", "70"),
+            ("B", "150", "1", "55"),
+            ("B", "300", "0", "60"),
+        ],
+    )
+    conn.executemany(
+        "INSERT OR REPLACE INTO variables (id, study_id, column_name, role, data_type) VALUES (?, ?, ?, ?, ?)",
+        [
+            (1, STUDY_ID, "treatment_arm", "baseline", "categorical"),
+            (2, STUDY_ID, "pfs_days", "outcome", "time_to_event"),
+            (3, STUDY_ID, "pfs_event", "outcome", "categorical"),
+            (4, STUDY_ID, "age", "baseline", "continuous"),
+        ],
+    )
+    conn.commit()
+    conn.close()
+
+
+def _plan_args(cox_ph_models=None, **kwargs):
+    values = dict(
+        study_id=STUDY_ID,
+        outcome_var_ids="3",
+        tests=[],
+        covariates="",
+        study_type="cohort",
+        comparison="PFS by treatment arm",
+        overrides=[],
+        matching_criteria="",
+    )
+    values["cox_ph_models"] = cox_ph_models or []
+    values.update(kwargs)
+    return argparse.Namespace(**values)
+
+
+# --- valid ID/name resolution ---
+
+
+def test_cox_ph_accepts_valid_column_names(capsys):
+    _setup_db()
+    cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm:age:test rationale"]))
+    prov = DATA_ROOT / STUDY_ID / "study_plan.provisional.json"
+    assert prov.exists()
+
+
+def test_cox_ph_accepts_4_part_format(capsys):
+    _setup_db()
+    cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    prov = DATA_ROOT / STUDY_ID / "study_plan.provisional.json"
+    assert prov.exists()
+
+
+def test_cox_ph_accepts_5_part_format_no_rationale(capsys):
+    _setup_db()
+    cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm:age"]))
+    prov = DATA_ROOT / STUDY_ID / "study_plan.provisional.json"
+    assert prov.exists()
+
+
+# --- invalid column name rejected at plan ---
+
+
+def test_cox_ph_rejects_unknown_column(capsys):
+    _setup_db()
+    with pytest.raises(SystemExit):
+        cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:nonexistent_col:treatment_arm"]))
+    err = capsys.readouterr().err
+    assert "nonexistent_col" in err
+    assert "not found" in err
+
+
+def test_cox_ph_rejects_unknown_covariate(capsys):
+    _setup_db()
+    with pytest.raises(SystemExit):
+        cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm:fake_cov"]))
+    err = capsys.readouterr().err
+    assert "fake_cov" in err
+    assert "not found" in err
+
+
+# --- non-binary event_col rejected at plan ---
+
+
+def test_cox_ph_rejects_non_binary_event(capsys):
+    _setup_db(extra_rows=[
+        ("A", "100", "maybe", "65"),
+        ("B", "150", "yes", "55"),
+    ])
+    with pytest.raises(SystemExit):
+        cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    err = capsys.readouterr().err
+    assert "binary" in err.lower()
+    assert "maybe" in err or "yes" in err
+
+
+def test_cox_ph_accepts_binary_event(capsys):
+    _setup_db(extra_rows=[
+        ("A", "100", "1", "65"),
+        ("B", "150", "0", "55"),
+    ])
+    cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    prov = DATA_ROOT / STUDY_ID / "study_plan.provisional.json"
+    assert prov.exists()
+
+
+# --- negative survival_time rejected at plan ---
+
+
+def test_cox_ph_rejects_negative_survival_time(capsys):
+    _setup_db(extra_rows=[
+        ("A", "-5", "1", "65"),
+        ("B", "150", "0", "55"),
+    ])
+    with pytest.raises(SystemExit):
+        cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    err = capsys.readouterr().err
+    assert "negative" in err.lower()
+
+
+def test_cox_ph_rejects_non_numeric_survival_time(capsys):
+    _setup_db(extra_rows=[
+        ("A", "abc", "1", "65"),
+        ("B", "150", "0", "55"),
+    ])
+    with pytest.raises(SystemExit):
+        cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    err = capsys.readouterr().err
+    assert "not numeric" in err.lower()
+
+
+# --- short-form 4-part format accepted ---
+
+
+def test_cox_ph_4_part_no_covariates_no_rationale(capsys):
+    _setup_db()
+    cmd_plan(_plan_args(cox_ph_models=["m1:pfs_days:pfs_event:treatment_arm"]))
+    prov = DATA_ROOT / STUDY_ID / "study_plan.provisional.json"
+    assert prov.exists()
+    import json
+    plan = json.loads(prov.read_text())
+    m = plan["cox_ph_models"][0]
+    assert m["model_name"] == "m1"
+    assert m["survival_time_col"] == "pfs_days"
+    assert m["event_col"] == "pfs_event"
+    assert m["primary_treatment_col"] == "treatment_arm"
+    assert m["covariate_cols"] == []
+    assert m["rationale"] == ""
+
+
+# --- regression: locked-plan reload with CoxPHModel dataclass ---
+
+
+def _seed_large_dataset(study_id: str, n: int = 25):
+    """Seed N rows into raw table so EPV check passes."""
+    import random
+    rng = random.Random(42)
+    conn = get_connection(study_id)
+    raw = f"raw_{study_id}"
+    rows = []
+    for i in range(n):
+        arm = rng.choice(["A", "B"])
+        pfs = int(max(30, rng.expovariate(1/200)))
+        evt = rng.choices(["0", "1"], weights=[30, 70])[0]
+        rows.append((arm, str(pfs), evt))
+    conn.executemany(
+        f"INSERT INTO {raw} (treatment_arm, pfs_days, pfs_event) "
+        f"VALUES (?, ?, ?)",
+        rows,
+    )
+    conn.commit()
+    conn.close()
+
+
+def test_cox_ph_regression_locked_plan_reload():
+    """Plan → lock → reload (StudyPlan.from_dict) → unmask → analyze.
+
+    Regression: after loading a locked plan from disk, cox_ph_models entries
+    are CoxPHModel dataclass instances. cmd_analyze must handle them via
+    _model_field helper, not .get().
+    """
+    sid = "test_cox_reload_regression"
+    study_path = DATA_ROOT / sid
+    if study_path.exists():
+        shutil.rmtree(study_path)
+    study_path.mkdir(parents=True)
+    conn = get_connection(sid)
+    init_db(conn)
+    conn.execute(
+        "INSERT INTO studies (id, name, created_at, data_dir, study_type) VALUES (?, ?, ?, ?, ?)",
+        (sid, "Cox PH Reload Regression", "2025-01-01", str(study_path), "cohort"),
+    )
+    raw = f"raw_{sid}"
+    conn.execute(f"CREATE TABLE {raw} "
+        f"(row_id INTEGER PRIMARY KEY, treatment_arm TEXT, pfs_days TEXT, "
+        f"pfs_event TEXT, age TEXT, high_risk_fish TEXT, prior_lines TEXT)")
+    conn.executemany(
+        "INSERT OR REPLACE INTO variables (id, study_id, column_name, role, data_type) VALUES (?, ?, ?, ?, ?)",
+        [
+            (1, sid, "treatment_arm", "baseline", "categorical"),
+            (2, sid, "pfs_days", "outcome", "time_to_event"),
+            (3, sid, "pfs_event", "outcome", "categorical"),
+        ],
+    )
+    conn.commit()
+    conn.close()
+
+    _seed_large_dataset(sid, n=25)
+
+    # Seal outcomes so the masked shadow table exists at plan time
+    from core.masking.gate import seal_outcomes
+    seal_outcomes(sid)
+
+    # Plan with a Cox PH model (treatment_arm only, no extra covariates — avoids singular matrix with sparse levels)
+    plan_ns = argparse.Namespace(
+        study_id=sid,
+        outcome_var_ids="3",
+        tests=[],
+        covariates="",
+        study_type="cohort",
+        comparison="PFS by treatment arm",
+        overrides=[],
+        matching_criteria="",
+        cox_ph_models=["pfs_simple:pfs_days:pfs_event:treatment_arm"],
+    )
+    cmd_plan(plan_ns)
+
+    lock_ns = argparse.Namespace(study_id=sid, allow_duplicate_ids=False)
+    cmd_lock(lock_ns)
+
+    # Now reload the locked plan from disk as a fresh StudyPlan
+    locked_path = sorted(study_path.glob("study_plan.v*.locked.json"))[-1]
+    plan_data = json.loads(locked_path.read_text())
+    reloaded_plan = StudyPlan.from_dict(plan_data)
+    # Verify cox_ph_models entries are CoxPHModel instances
+    for m in reloaded_plan.cox_ph_models:
+        assert not isinstance(m, dict), "cox_ph_models must be CoxPHModel after from_dict, not dict"
+        from core.planning.study_plan import CoxPHModel
+        assert isinstance(m, CoxPHModel)
+
+    # Manually lock the reloaded plan to mimic what cmd_analyze does (load_plan)
+    path = lock_plan(sid, reloaded_plan)
+    assert path.exists()
+
+    # Unmask
+    unmask_study(sid)
+
+    # Analyze with --force to bypass EPV warning (EPV < 10 with small N)
+    analyze_ns = argparse.Namespace(
+        study_id=sid,
+        force=True,
+        post_hoc=False,
+        rerun=False,
+    )
+    cmd_analyze(analyze_ns)
+    # If we get here without AttributeError, the regression is fixed.
+    # Check results in DB
+    conn = get_connection(sid)
+    cur = conn.execute(
+        "SELECT id, test_name, statistic, p_value FROM analysis_results "
+        "WHERE study_id=? AND test_name=? ORDER BY id DESC LIMIT 1",
+        (sid, "cox_ph_model"),
+    )
+    row = cur.fetchone()
+    conn.close()
+    assert row is not None, "No Cox PH result found in DB"
+    assert row["statistic"] is not None, "Cox PH HR should not be None"
+    assert row["p_value"] is not None, "Cox PH p-value should not be None"
+    assert row["statistic"] > 0, "HR must be positive"
+
+    shutil.rmtree(study_path)
+
+
+# ── Interaction term tests with adequate N ──────────────────────────────
+
+
+def _synthetic_interaction_data(n: int, seed: int,
+                                true_interaction: bool = True) -> pd.DataFrame:
+    """Generate synthetic survival data with a known treatment × subgroup interaction.
+
+    Parameters
+    ----------
+    n : int
+        Number of patients.
+    seed : int
+        Random seed for reproducibility.
+    true_interaction : bool
+        If True, treatment effect differs by subgroup (HR_subgroup_B = 1.0).
+        If False, treatment effect is same across subgroups (no interaction).
+
+    Returns
+    -------
+    pd.DataFrame with columns:
+        pfs_days, pfs_event, treatment_arm, high_risk_fish, age
+    """
+    rng = random.Random(seed)
+    np_rng = np.random.RandomState(seed)
+
+    data = []
+    for _ in range(n):
+        arm = rng.choice(["A", "B"])
+        fish = rng.choices(["yes", "no"], weights=[0.4, 0.6])[0]
+        age = rng.randint(40, 85)
+
+        # Baseline hazard: exp(-t/200)
+        # Base HR for treatment (arm B vs A) — always present
+        if arm == "B":
+            base_hr = 0.5  # treatment halves hazard
+        else:
+            base_hr = 1.0
+
+        # Interaction: in high-risk subgroup, treatment effect differs
+        if fish == "yes" and arm == "B":
+            if true_interaction:
+                fish_hr = 2.0  # treatment loses effectiveness in high-risk
+            else:
+                fish_hr = 1.0  # no interaction, same effect
+        else:
+            fish_hr = 1.0
+
+        hr = base_hr * fish_hr
+
+        # Generate survival time from exponential with rate=hr/200
+        t = np_rng.exponential(200.0 / hr)
+        pfs_days = max(1, int(t))
+
+        # Censor at 500 days
+        event = 1 if pfs_days < 500 else 0
+        pfs_days = min(pfs_days, 500)
+
+        data.append((pfs_days, event, arm, fish, age))
+
+    return pd.DataFrame(data, columns=["pfs_days", "pfs_event",
+                                       "treatment_arm", "high_risk_fish", "age"])
+
+
+class TestInteractionTerms:
+    """Verify Cox PH interaction term detection with adequate power."""
+
+    def test_true_interaction_detected(self):
+        """Model with interaction should recover significant interaction term (N=400)."""
+        df = _synthetic_interaction_data(400, seed=42, true_interaction=True)
+        result = _cox_ph_model(
+            df,
+            time_col="pfs_days",
+            event_col="pfs_event",
+            group_col="treatment_arm",
+            covariates=["high_risk_fish", "age"],
+            interaction_terms=[["treatment_arm", "high_risk_fish"]],
+        )
+        cov_results = result.get("params", {}).get("per_covariate_results", [])
+        int_rows = [c for c in cov_results if "(interaction)" in c.get("covariate", "")]
+        assert len(int_rows) > 0, "No interaction term found in results"
+        int_row = int_rows[0]
+        assert int_row["hr"] is not None, "Interaction HR should not be None"
+        assert int_row["wald_p"] is not None, "Interaction p-value should not be None"
+        # With true interaction built in and N=400, p should be significant
+        assert int_row["wald_p"] < 0.05, (
+            f"True interaction should be detected (p={int_row['wald_p']:.4f})"
+        )
+
+    def test_no_interaction_not_detected(self):
+        """Model with no interaction should have non-significant interaction term."""
+        df = _synthetic_interaction_data(400, seed=43, true_interaction=False)
+        result = _cox_ph_model(
+            df,
+            time_col="pfs_days",
+            event_col="pfs_event",
+            group_col="treatment_arm",
+            covariates=["high_risk_fish", "age"],
+            interaction_terms=[["treatment_arm", "high_risk_fish"]],
+        )
+        cov_results = result.get("params", {}).get("per_covariate_results", [])
+        int_rows = [c for c in cov_results if "(interaction)" in c.get("covariate", "")]
+        assert len(int_rows) > 0, "No interaction term found in results"
+        int_row = int_rows[0]
+        # With no interaction, CI should cross 1
+        ci_lo = int_row.get("ci_lower")
+        ci_hi = int_row.get("ci_upper")
+        if ci_lo is not None and ci_hi is not None:
+            assert ci_lo <= 1 <= ci_hi, (
+                f"No-interaction CI ({ci_lo:.3f}, {ci_hi:.3f}) should cross 1"
+            )
+
+    def test_epv_counts_interaction_as_predictor(self):
+        """EPV check should count interaction terms as additional predictors."""
+        from core.planning.test_selector import check_cox_ph_model_assumptions
+
+        # Build a small dataset so EPV < 10 is triggered
+        study_id = "test_epv_interaction"
+        study_path = DATA_ROOT / study_id
+        if study_path.exists():
+            shutil.rmtree(study_path)
+        study_path.mkdir(parents=True)
+        conn = get_connection(study_id)
+        init_db(conn)
+        conn.execute(
+            "INSERT OR REPLACE INTO studies (id, name, created_at, data_dir, study_type) VALUES (?, ?, ?, ?, ?)",
+            (study_id, "EPV Interaction Test", "2025-01-01", str(study_path), "cohort"),
+        )
+        conn.execute(
+            f"CREATE TABLE IF NOT EXISTS raw_masked_{study_id} AS "
+            "SELECT * FROM (SELECT 'A' AS treatment_arm, 'yes' AS high_risk_fish, "
+            "'1' AS pfs_event, 100 AS pfs_days, 50 AS age) WHERE 1=0"
+        )
+        # Insert 5 events across 4 patients → EPV=5/3≈1.7 with interaction
+        rows = [
+            ("A", "yes", "1", 100, 50),
+            ("A", "no", "0", 200, 55),
+            ("B", "yes", "1", 80, 60),
+            ("B", "no", "1", 150, 45),
+        ]
+        for r in rows:
+            conn.execute(
+                f"INSERT INTO raw_masked_{study_id} "
+                f"(treatment_arm, high_risk_fish, pfs_event, pfs_days, age) "
+                f"VALUES (?, ?, ?, ?, ?)", r
+            )
+        conn.execute(
+            "INSERT OR REPLACE INTO variables (id, study_id, column_name, role, data_type) "
+            "VALUES (1, ?, 'treatment_arm', 'baseline', 'categorical')",
+            (study_id,),
+        )
+        conn.execute(
+            "INSERT OR REPLACE INTO variables (id, study_id, column_name, role, data_type) "
+            "VALUES (2, ?, 'pfs_days', 'outcome', 'time_to_event')",
+            (study_id,),
+        )
+        conn.execute(
+            "INSERT OR REPLACE INTO variables (id, study_id, column_name, role, data_type) "
+            "VALUES (3, ?, 'pfs_event', 'outcome', 'categorical')",
+            (study_id,),
+        )
+        conn.commit()
+        conn.close()
+
+        model = {
+            "model_name": "test_interaction",
+            "survival_time_col": "pfs_days",
+            "event_col": "pfs_event",
+            "primary_treatment_col": "treatment_arm",
+            "covariate_cols": ["high_risk_fish"],
+            "interaction_terms": [["treatment_arm", "high_risk_fish"]],
+        }
+        warnings = check_cox_ph_model_assumptions(study_id, [model])
+        warning_text = " ".join(warnings)
+        # The warning should mention 3 predictors (treatment + 1 covariate + 1 interaction)
+        assert "3 predictor" in warning_text, (
+            f"EPV warning should count 3 predictors: {warning_text}"
+        )
+        assert "EPV=" in warning_text, (
+            "EPV warning should fire when interaction adds a predictor"
+        )
+
+        shutil.rmtree(study_path)