|
| 1 | +""" |
| 2 | +Round-3 Smart-tool coverage for Sprint-B: compare_estimators registers |
| 3 | +g_computation, sensitivity_dashboard reports the correct method label |
| 4 | +for Sprint-B results, and ROADMAP §6 documents the compare_estimators |
| 5 | +gap for the remaining hint-driven estimators. |
| 6 | +""" |
| 7 | + |
| 8 | +from pathlib import Path |
| 9 | + |
| 10 | +import numpy as np |
| 11 | +import pandas as pd |
| 12 | +import pytest |
| 13 | + |
| 14 | +import statspai as sp |
| 15 | + |
| 16 | + |
| 17 | +# --------------------------------------------------------------------- |
| 18 | +# compare_estimators — g_computation is now a first-class option |
| 19 | +# --------------------------------------------------------------------- |
| 20 | + |
| 21 | +@pytest.fixture |
| 22 | +def binary_dgp(): |
| 23 | + rng = np.random.default_rng(0) |
| 24 | + n = 500 |
| 25 | + X = rng.normal(0, 1, n) |
| 26 | + D = rng.binomial(1, 0.5, n).astype(float) |
| 27 | + Y = 1.0 * D + 0.5 * X + rng.normal(0, 0.3, n) |
| 28 | + return pd.DataFrame({'y': Y, 'd': D, 'x': X}) |
| 29 | + |
| 30 | + |
| 31 | +def test_compare_estimators_accepts_g_computation(binary_dgp): |
| 32 | + comp = sp.compare_estimators( |
| 33 | + data=binary_dgp, y='y', treatment='d', |
| 34 | + methods=['ols', 'g_computation'], |
| 35 | + covariates=['x'], |
| 36 | + ) |
| 37 | + # Both methods ran (agreement diagnostics populated) |
| 38 | + assert hasattr(comp, 'summary') |
| 39 | + txt = comp.summary() |
| 40 | + assert 'G-computation' in txt |
| 41 | + assert 'OLS' in txt |
| 42 | + # Sign agreement 100 % on a linear, binary-D, positive-effect DGP. |
| 43 | + assert 'Sign agreement: 100%' in txt |
| 44 | + |
| 45 | + |
| 46 | +def test_compare_estimators_g_computation_recovers_truth(binary_dgp): |
| 47 | + """Both OLS and g_computation should hug the true effect ≈ 1.0.""" |
| 48 | + comp = sp.compare_estimators( |
| 49 | + data=binary_dgp, y='y', treatment='d', |
| 50 | + methods=['g_computation'], |
| 51 | + covariates=['x'], |
| 52 | + ) |
| 53 | + r = comp.results['G-computation'] |
| 54 | + assert abs(r.estimate - 1.0) < 0.15 |
| 55 | + |
| 56 | + |
| 57 | +def test_compare_estimators_unsupported_method_does_not_crash(binary_dgp): |
| 58 | + """ |
| 59 | + Sprint-B methods with extra required args (proximal, msm, ...) are |
| 60 | + documented as NOT supported. Requesting one must not silently |
| 61 | + drop the supported methods too. |
| 62 | +
|
| 63 | + ``compare_estimators`` wraps each method in its own try/except + |
| 64 | + ``warnings.warn`` (see src/statspai/smart/compare.py), so the call |
| 65 | + itself MUST NOT raise to the caller. The contract we're pinning: |
| 66 | + the supported method (OLS) still appears in ``.results`` even |
| 67 | + when an unsupported method ('proximal') sits beside it in the |
| 68 | + ``methods=`` list. |
| 69 | + """ |
| 70 | + import warnings |
| 71 | + with warnings.catch_warnings(): |
| 72 | + # compare_estimators emits UserWarning for the unsupported |
| 73 | + # branch; suppress it so the test output is clean while still |
| 74 | + # exercising the real path. |
| 75 | + warnings.simplefilter('ignore') |
| 76 | + comp = sp.compare_estimators( |
| 77 | + data=binary_dgp, y='y', treatment='d', |
| 78 | + methods=['ols', 'proximal'], # 'proximal' unsupported |
| 79 | + covariates=['x'], |
| 80 | + ) |
| 81 | + # OLS must still be present regardless of what the unsupported |
| 82 | + # method did. No permissive try/except — a true regression |
| 83 | + # (OLS silently disappearing) should fail this assertion. |
| 84 | + assert 'OLS' in comp.results |
| 85 | + # 'proximal' must NOT appear (never dispatched) |
| 86 | + assert 'Proximal' not in comp.results |
| 87 | + |
| 88 | + |
| 89 | +# --------------------------------------------------------------------- |
| 90 | +# sensitivity_dashboard — method label now reads from .method first |
| 91 | +# --------------------------------------------------------------------- |
| 92 | + |
| 93 | +def test_sensitivity_dashboard_reports_proximal_method_label(): |
| 94 | + rng = np.random.default_rng(0) |
| 95 | + n = 500 |
| 96 | + U = rng.normal(0, 1, n) |
| 97 | + D = 0.8 * U + rng.normal(0, 0.5, n) |
| 98 | + Z = 0.9 * U + rng.normal(0, 0.3, n) |
| 99 | + W = 0.9 * U + rng.normal(0, 0.3, n) |
| 100 | + Y = 1.5 * D + 0.7 * U + rng.normal(0, 0.3, n) |
| 101 | + df = pd.DataFrame({'y': Y, 'd': D, 'z': Z, 'w': W}) |
| 102 | + r = sp.proximal(df, y='y', treat='d', proxy_z=['z'], proxy_w=['w']) |
| 103 | + dash = sp.sensitivity_dashboard(r, data=df, verbose=False, |
| 104 | + dimensions=['sample']) |
| 105 | + assert 'Proximal' in dash.method, ( |
| 106 | + f"Expected 'Proximal' in dash.method, got {dash.method!r}" |
| 107 | + ) |
| 108 | + |
| 109 | + |
| 110 | +def test_sensitivity_dashboard_reports_msm_method_label(): |
| 111 | + rng = np.random.default_rng(0) |
| 112 | + rows = [] |
| 113 | + for i in range(100): |
| 114 | + for t in range(3): |
| 115 | + rows.append({ |
| 116 | + 'id': i, 'time': t, |
| 117 | + 'A': rng.binomial(1, 0.5), |
| 118 | + 'L': rng.normal(), |
| 119 | + 'V': rng.normal(), |
| 120 | + }) |
| 121 | + panel = pd.DataFrame(rows) |
| 122 | + panel['Y'] = (0.5 * panel.groupby('id')['A'].cumsum() |
| 123 | + + 0.3 * panel['V'] |
| 124 | + + rng.normal(0, 0.3, len(panel))) |
| 125 | + r = sp.msm(panel, y='Y', treat='A', id='id', time='time', |
| 126 | + time_varying=['L'], baseline=['V']) |
| 127 | + dash = sp.sensitivity_dashboard(r, data=panel, verbose=False, |
| 128 | + dimensions=['sample']) |
| 129 | + assert 'Marginal Structural' in dash.method, ( |
| 130 | + f"Expected 'Marginal Structural' in label, got {dash.method!r}" |
| 131 | + ) |
| 132 | + |
| 133 | + |
| 134 | +def test_sensitivity_dashboard_legacy_econometric_result_label(): |
| 135 | + """Backward-compat: EconometricResults path still reads model_type.""" |
| 136 | + rng = np.random.default_rng(0) |
| 137 | + n = 200 |
| 138 | + df = pd.DataFrame({ |
| 139 | + 'y': rng.normal(0, 1, n), |
| 140 | + 'x': rng.normal(0, 1, n), |
| 141 | + }) |
| 142 | + r = sp.regress('y ~ x', data=df) |
| 143 | + dash = sp.sensitivity_dashboard(r, data=df, verbose=False, |
| 144 | + dimensions=['sample']) |
| 145 | + # Not Unknown — the model_info['model_type'] fallback still works |
| 146 | + assert dash.method != 'Unknown' |
| 147 | + |
| 148 | + |
| 149 | +# --------------------------------------------------------------------- |
| 150 | +# ROADMAP §6 documents the compare_estimators Sprint-B gap |
| 151 | +# --------------------------------------------------------------------- |
| 152 | + |
| 153 | +def test_compare_estimators_g_computation_binary_guard_compact(): |
| 154 | + """ |
| 155 | + Review-fix: the g_computation branch emits a CLEAN warning on |
| 156 | + non-binary D (≤ 500 chars) rather than dumping the full list of |
| 157 | + unique D values (which for n=400 continuous D was ~20 KB). |
| 158 | + """ |
| 159 | + import warnings |
| 160 | + rng = np.random.default_rng(0) |
| 161 | + n = 400 |
| 162 | + df = pd.DataFrame({ |
| 163 | + 'y': rng.normal(0, 1, n), |
| 164 | + 'd': rng.normal(0, 1, n), # continuous — not binary |
| 165 | + 'x': rng.normal(0, 1, n), |
| 166 | + }) |
| 167 | + with warnings.catch_warnings(record=True) as caught: |
| 168 | + warnings.simplefilter('always') |
| 169 | + comp = sp.compare_estimators( |
| 170 | + data=df, y='y', treatment='d', |
| 171 | + methods=['g_computation'], |
| 172 | + covariates=['x'], |
| 173 | + ) |
| 174 | + # g_computation branch should have emitted a compact warning. |
| 175 | + g_warns = [w for w in caught if 'g_computation' in str(w.message)] |
| 176 | + assert len(g_warns) >= 1, "expected a g_computation warning" |
| 177 | + msg = str(g_warns[0].message) |
| 178 | + # Compact: the warning fits on a few lines, not a 20 KB value dump |
| 179 | + assert len(msg) < 500, ( |
| 180 | + f"g_computation warning is {len(msg)} chars; expected <500. " |
| 181 | + f"First 200 chars: {msg[:200]}" |
| 182 | + ) |
| 183 | + # Useful: it actually mentions the issue and points to the fix |
| 184 | + assert 'binary' in msg.lower() |
| 185 | + |
| 186 | + |
| 187 | +def test_roadmap_documents_compare_estimators_sprint_b_gap(): |
| 188 | + root = Path(__file__).resolve().parent.parent |
| 189 | + roadmap = (root / 'docs' / 'ROADMAP.md').read_text() |
| 190 | + # Section 6 heading present |
| 191 | + assert 'compare_estimators' in roadmap |
| 192 | + # The four hint-driven methods are called out |
| 193 | + for name in ('proximal', 'msm', 'principal_strat', 'mediator'): |
| 194 | + assert name in roadmap.lower() |
0 commit comments