diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 068f1f9..234bb5c 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -41,7 +41,7 @@ This repository contains two distinct Python packages for biomedical analysis: - `discrete_optimization.py`, `utilities.py` - Optimization utilities (used by KM methods) - `datacard.py` - Main Datacard class (core data parsing) - `command_line_interface.py` - KM-specific CLI functions - - `yi_correction.py` - Yi correction method + - `comparisons/` - published alternatives (Yi weights, Küchenhoff MC-SIMEX) ### Test Structure - **`test/roc_picker/`**: ROC Picker tests, datacards, and reference data @@ -61,15 +61,16 @@ This repository contains two distinct Python packages for biomedical analysis: - `03_kaplan_meier_example.md` - Jupyter notebook showing Python API usage examples - `04_analysis_demo.md` - Jupyter notebook with analysis demonstration - `05_compare_lifelines_greenwood.md` - Jupyter notebook comparing to `lifelines` and Greenwood method -- `06_yi_method_comparison.md` - Jupyter notebook comparing Yi correction method -- `07_command_line_interface.md` - **Pure Markdown** (no Python cells) documenting all CLI options for `kombine` and `kombine_twogroups` commands (synced with Jupytext) +- `06_compare_thomas_grunkemeier.md` - Jupyter notebook comparing Thomas–Grunkemeier intervals +- `07_previous_methods_comparison.md` - Jupyter notebook comparing Yi, MC-SIMEX, and KoMbine +- `08_command_line_interface.md` - **Pure Markdown** (no Python cells) documenting all CLI options for `kombine` and `kombine_twogroups` commands (synced with Jupytext) **Documentation style guidelines**: -- Files `03_*.md` through `06_*.md` are Jupytext notebooks with Python cells for interactive examples -- Files `01_*.md` and `07_*.md` are pure Markdown with Jupytext headers but no Python cells +- Files `03_*.md` through `07_*.md` are Jupytext notebooks with Python cells for interactive examples +- Files `01_*.md` and `08_*.md` are pure Markdown with Jupytext headers but no Python cells - All documentation markdown files must have Jupytext headers to allow `jupytext --sync` to process them -- All CLI options must be documented in `07_command_line_interface.md` and verified by `test/kombine/test_ci_and_documentation.py` -- When adding new CLI arguments, update `07_command_line_interface.md` and run the documentation test +- All CLI options must be documented in `08_command_line_interface.md` and verified by `test/kombine/test_ci_and_documentation.py` +- When adding new CLI arguments, update `08_command_line_interface.md` and run the documentation test - The table of contents (`01_table_of_contents.md`) must list all numbered documentation files **Compiling KoMbine LaTeX documentation (`02_kombine.tex`)**: @@ -214,6 +215,7 @@ python -m test.kombine.test_km_likelihood # ~2 minutes, requires unres python -m test.kombine.test_hazard_ratio python -m test.kombine.test_km_plotting python -m test.kombine.test_yi_correction +python -m test.kombine.test_mc_simex python -m test.kombine.test_discrete_class_observable python -m test.kombine.test_ci_and_documentation ``` diff --git a/.github/workflows/ci-cd.yml b/.github/workflows/ci-cd.yml index 73ed5d6..843f10e 100644 --- a/.github/workflows/ci-cd.yml +++ b/.github/workflows/ci-cd.yml @@ -15,6 +15,7 @@ jobs: uses: actions/checkout@v4 - name: Set up Python 3.x + id: setup-python uses: actions/setup-python@v5 with: python-version: 3.11 @@ -28,17 +29,23 @@ jobs: run: | jupytext --sync docs/roc_picker/*.md docs/kombine/*.md + # success() || failure(): still run later linters if an earlier one fails + # (not if the job is cancelled). The job stays red if any step failed. - name: pyflakes (all python files) + if: success() || failure() run: | python -m pyflakes kombine roc_picker test docs .github/workflows/checklog.py .github/workflows/__init__.py - name: pylint (all python files) + if: success() || failure() run: | python -m pylint kombine roc_picker test docs .github/workflows/checklog.py .github/workflows/__init__.py - uses: jakebailey/pyright-action@v2.3.2 + if: success() || failure() with: version: 1.1.400 + python-path: ${{ steps.setup-python.outputs.python-path }} test-roc-picker: runs-on: ubuntu-latest @@ -125,16 +132,15 @@ jobs: pip install . ipykernel jupytext lifelines nbconvert pyflakes pylint texoutparse && rm -r build - - name: Install Gurobi + # The gurobipy wheel installed by `pip install .` bundles the solver library, + # so no separate distribution is needed. The 10.0.3 tarball this replaces was + # installed over the wheel and downgraded the solver: Gurobi 10 approximates + # every log/exp with a 1000-piece linearization, which on the KM MINLPs is + # 2.7x slower than the exact handling in 11+ and shifts 2NLL by ~0.8. + # See docs/kombine/compute_budget.md. + - name: Report Gurobi version run: | - wget https://packages.gurobi.com/10.0/gurobi10.0.3_linux64.tar.gz - tar -xvzf gurobi10.0.3_linux64.tar.gz - echo "export GUROBI_HOME=$PWD/gurobi1003/linux64" >> $GITHUB_ENV - echo "$PWD/gurobi1003/linux64/bin" >> $GITHUB_PATH - echo "LD_LIBRARY_PATH=$PWD/gurobi1003/linux64/lib" >> $GITHUB_ENV - cd gurobi1003/linux64 - python setup.py install - cd ../.. + python -c "import gurobipy; print(gurobipy.gurobi.version())" - name: Configure Gurobi WLS license run: | @@ -181,13 +187,21 @@ jobs: run: | python -m test.kombine.test_yi_correction + - name: MC-SIMEX unit tests + run: | + python -m test.kombine.test_mc_simex + - name: convert KoMbine notebooks run: | jupytext --sync docs/kombine/*.md + # Notebook 07 defaults to n=20 then n=50. CI sets + # KOMBINE_QUICK_COMPARISON=1 to skip n=50 (n=20 only, minutes). - name: test KoMbine notebooks + env: + KOMBINE_QUICK_COMPARISON: "1" run: | - jupyter nbconvert --execute --to html -- docs/kombine/*.ipynb + jupyter nbconvert --execute --to html --ExecutePreprocessor.timeout=3600 -- docs/kombine/*.ipynb - name: Compile KM plots run: | diff --git a/.gitignore b/.gitignore index 3f3d20e..03c041e 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,6 @@ test/kombine/datacards/lung/flatfielding_systematic/*.txt test/kombine/datacards/lung/poisson_and_flatfielding/*.txt test/kombine/datacards/lung/uncorrected_flatfielding_systematic/*.txt test/kombine/datacards/lung/poisson_and_uncorrected_flatfielding/*.txt + +# On-demand rebinned methods_comparison cards (from rebin_methods_comparison_times.py) +test/kombine/datacards/simple_examples/rebinned/ diff --git a/docs/__init__.py b/docs/__init__.py index e69de29..d595826 100644 --- a/docs/__init__.py +++ b/docs/__init__.py @@ -0,0 +1 @@ +"""Docs package (notebooks and helpers; not installed via pip).""" diff --git a/docs/kombine/01_table_of_contents.md b/docs/kombine/01_table_of_contents.md index 438155b..49d6af3 100644 --- a/docs/kombine/01_table_of_contents.md +++ b/docs/kombine/01_table_of_contents.md @@ -8,7 +8,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 --- ```python @@ -37,9 +37,9 @@ jupyter: - `06_compare_thomas_grunkemeier.html` - comparison of KoMbine binomial-only confidence intervals to Thomas and Grunkemeier as implemented in R package `km.ci` - - `07_yi_method_comparison.html` - - comprehensive demonstration of Yi's misclassification correction - method applied to Kaplan-Meier curves, p-values, and hazard ratios + - `07_previous_methods_comparison.html` + - comparison of Yi's weights, Küchenhoff MC-SIMEX, and KoMbine + on Kaplan-Meier curves, p-values, and hazard ratios - `08_command_line_interface.html` - documentation for the command line interface (`kombine` and `kombine_twogroups` commands) diff --git a/docs/kombine/03_kaplan_meier_example.md b/docs/kombine/03_kaplan_meier_example.md index 5685f23..294be42 100644 --- a/docs/kombine/03_kaplan_meier_example.md +++ b/docs/kombine/03_kaplan_meier_example.md @@ -6,7 +6,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 kernelspec: display_name: rocpicker language: python diff --git a/docs/kombine/04_analysis_demo.md b/docs/kombine/04_analysis_demo.md index a375d2b..2a77f4d 100644 --- a/docs/kombine/04_analysis_demo.md +++ b/docs/kombine/04_analysis_demo.md @@ -6,7 +6,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 kernelspec: display_name: rocpicker language: python @@ -69,9 +69,9 @@ print(f"Censored: {sum(1 for p in datacard.patients if p.censored)}") ### KoMbine p-value calculation -KoMbine provides a likelihood-based p-value that accounts for measurement uncertainty when comparing survival curves across a biomarker threshold. For context, we also compute the standard logrank p-value. +KoMbine provides a permutation likelihood-ratio p-value that accounts for measurement uncertainty when comparing survival curves across a biomarker threshold. For context, we also compute the standard logrank p-value. -For this dataset and threshold, the likelihood p-value is $2.67\times 10^{-2}$ and the logrank p-value is $1.61\times 10^{-2}$. +For this dataset and threshold, the KoMbine permutation p-value is $6.00\times 10^{-2}$ and the logrank p-value is $1.61\times 10^{-2}$. ```python # Define biomarker threshold to split patients into two groups @@ -83,7 +83,11 @@ km_p_value_minlp = datacard.km_p_value( parameter_min=-np.inf, parameter_max=np.inf, ) -kombine_p_value, _, _ = km_p_value_minlp.solve_and_pvalue(cox_only=False) +kombine_p_value, _, _ = km_p_value_minlp.solve_and_pvalue( + cox_only=False, + n_permutations=49, + rng=0, +) # Conventional logrank p-value logrank_p_value = km_p_value_minlp.survival_curves_pvalue_logrank() diff --git a/docs/kombine/05_compare_lifelines_greenwood.md b/docs/kombine/05_compare_lifelines_greenwood.md index add3c6e..2e4edc0 100644 --- a/docs/kombine/05_compare_lifelines_greenwood.md +++ b/docs/kombine/05_compare_lifelines_greenwood.md @@ -6,7 +6,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 kernelspec: display_name: rocpicker language: python @@ -75,7 +75,7 @@ We do, in fact, get 1:1 agreement with `lifelines`. For completeness, KoMbine also provides logrank test p-values and hazard ratio calculations that match `lifelines` for fixed observables. These comparisons demonstrate that our implementation of standard survival analysis methods is correct. -For measurement error corrections using Yi's method, see **notebook 07_yi_method_comparison.ipynb**. +For measurement error corrections using Yi's method and MC-SIMEX, see **notebook 07_previous_methods_comparison.ipynb**. ### Fixed Observable: P-value and Hazard Ratio vs lifelines diff --git a/docs/kombine/06_compare_thomas_grunkemeier.md b/docs/kombine/06_compare_thomas_grunkemeier.md index 301c8ef..e8a7f00 100644 --- a/docs/kombine/06_compare_thomas_grunkemeier.md +++ b/docs/kombine/06_compare_thomas_grunkemeier.md @@ -6,7 +6,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 kernelspec: display_name: rocpicker language: python @@ -17,7 +17,6 @@ jupyter: """Compare KoMbine binomial-only KM intervals against Thomas-Grunkemeier via R.""" # pylint: disable=wrong-import-position -# pyright: reportMissingImports=false ``` ```python diff --git a/docs/kombine/07_previous_methods_comparison.md b/docs/kombine/07_previous_methods_comparison.md new file mode 100644 index 0000000..c376e34 --- /dev/null +++ b/docs/kombine/07_previous_methods_comparison.md @@ -0,0 +1,989 @@ +--- +jupyter: + jupytext: + formats: ipynb,md,py:percent + kernelspec: + display_name: Python 3 + language: python + name: python3 + main_language: python + text_representation: + extension: .md + format_name: markdown + format_version: '1.3' + jupytext_version: 1.19.5 +--- + +```python +# pylint: disable=bad-indentation,line-too-long,missing-module-docstring,redefined-outer-name,trailing-whitespace,too-many-locals,too-many-statements,too-many-arguments,too-many-branches,too-many-positional-arguments,wrong-import-order,wrong-import-position +``` + +# Previous Methods vs KoMbine: Yi, MC-SIMEX, and Profile Likelihood + +This notebook compares two published recipes for discrete covariate misclassification — Yi's probability weights and Küchenhoff MC-SIMEX — to KoMbine's profile likelihood over group assignments. + +**Two families** (n=20 always; n=50 unless CI skips it): + +- **n=20** (`*_hr_example*`, ~7 distinct death times, discrete $e = 0.20$, $0.25$, $0.40$): minutes-scale. This is the discrete ladder where KoMbine KM bands widen / collapse at large $e$. Always run. +- **n=50** (`methods_comparison_*`, regenerated into a gitignored `rebinned/` dir, **4** quantile-bin death-time medians, discrete $e = 0.05$, $0.10$, $0.25$): hours-scale. KoMbine KM bands use `crossing_mode="feasibility"` (oracle bracketing + brentq polish). A 2026-09 local n=50 pass was **~85 min** (Analysis 1 ~63 min, Analysis 2 ~21 min with $B=99$, Analysis 3 negligible); Poisson (small counts) dominates Analysis 1 (~50 min). Default (unset env) runs this family **after** n=20. + +**CI / skip n=50** (`KOMBINE_QUICK_COMPARISON=1`): n=20 only. The n=50 cells still run, but they print a skip message instead of loading cards or solving MINLPs. + +Fixed and Poisson cards are split at `0.5001` (a density/value cut). Discrete-class cards are split at `1` (the boundary between class indices 0 and 1): +- Fixed Hazard Ratio (deterministic, no measurement error) +- Discrete classes with class probabilities (small/medium/large uncertainty) +- Poisson density with large counts (small relative error) +- Poisson density with moderate counts +- Poisson density with small counts (high relative error) + +All three methods use the same measurement model (`observable.probability_in_range`). They differ in what they do with those probabilities. + + +## Method Overview and Comparison + +| Aspect | Yi's Method | MC-SIMEX | KoMbine | +|--------|---|---|---| +| **Core idea** | Weighted KM/logrank using probabilistic group membership | Extra flips of hard labels, then extrapolate the naive estimator to zero error | Full likelihood with explicit group assignment variables | +| **Optimization** | 1-D scalar min of the weighted Breslow 2NLL | Monte Carlo average at each $\lambda$, quadratic fit | Mixed Integer Nonlinear Programming (Gurobi) | +| **Computational cost** | Low | Low–medium | Medium-high (~1 h for the n=50 family) | +| **Accuracy (within model)** | Approximate to the full likelihood | Approximate (simulation + extrapolation) | Exact maximizer within solver tolerance | +| **Uncertainty** | Likelihood-ratio interval of the weighted Breslow 2NLL (no KM bands here) | Sampling CI of the extrapolated number (Wald for HR) | Profile likelihood | +| **Core assumptions** | Known measurement error distribution; independent errors; fractional group membership is an adequate proxy for uncertain assignment | Known measurement error distribution; independent errors; quadratic extrapolation of the naive hard-label estimator is adequate | Known measurement error distribution; independent errors; patients belong to one group; event times treated as observed and discrete; likelihood model is correctly specified | + +### How the three recipes use the same $e_i$ +- **Yi averages with weights.** Each patient contributes to both groups in proportion to $P(G\mid\text{data})$. +- **MC-SIMEX extrapolates a naive hard-label estimator.** The observed label $G^*_i$ is flipped with Küchenhoff probability $\bigl(1-(1-2e_i)^{\lambda}\bigr)/2$, the usual KM/logrank/Cox estimator is averaged over simulations, and that curve is fit vs $\lambda$ and evaluated at $\lambda=-1$. +- **KoMbine profiles assignments.** A binary assignment is chosen for each patient, scored by the measurement model, and confidence sets come from the profile likelihood. + +### How Yi's Method Works +- Convert each patient's observed biomarker value into a probability of being below or above the threshold using the measurement error model. +- Instead of a global misclassification matrix as Yi describes, we extend her method to compute these probabilities on a per-patient basis (allowing uncertainty to vary by individual measurement). +- Use those per-patient probabilities as weights in the Kaplan-Meier estimator and logrank test. +- Every patient contributes to both groups, in proportion to their probability of belonging there. +- This yields fast, smooth estimates that tend to shrink group differences as measurement uncertainty grows. +- It is an approximation because it does not enforce a single, discrete group assignment for each patient. + +### How MC-SIMEX Works +- Keep the same $e_i=1-P(G=G^*_i\mid\text{data})$ that Yi uses, but then discard the probabilities and work with hard labels only. +- At $\lambda=0$ the estimator is the usual naive KM / logrank / Cox fit on $G^*$. +- At $\lambda>0$ additional flips are simulated; averaging the naive estimator traces how it changes as misclassification grows. +- A quadratic in $\lambda$ is extrapolated to $\lambda=-1$ (no misclassification). +- The HR interval is a Wald interval of that extrapolated $\log H$, not a likelihood-ratio interval. The $\Delta$2NLL curve plotted below is the Wald quadratic $(\log H-\widehat{\log H})^2/\hat\sigma^2$. + +### How KoMbine Works +- Introduce a binary assignment variable for each patient (low vs high group). +- Combine the survival likelihood with a measurement error penalty that scores how plausible each assignment is. +- Solve a constrained optimization problem that finds the most likely set of assignments and survival parameters together. +- Compute confidence intervals via profile likelihood, which naturally widens as uncertainty increases. +- This is exact for the specified likelihood model but requires heavier computation than Yi or MC-SIMEX. + +### Key Comparisons +1. **Kaplan-Meier Curves**: Visual comparison of survival estimates +2. **P-Values**: Logrank test results and statistical significance +3. **Hazard Ratios**: Point estimates and confidence intervals + + +## Discrete Classes with Class Probabilities + +Yi's Section 3.7.1 treats misclassification of a discrete covariate using a +misclassification matrix. We mimic that structure using two classes and a +homogeneous binary matrix, +$$ +\Pi = \begin{pmatrix} 1 - e & e \\ e & 1 - e \end{pmatrix} +$$ +where $e$ is the misclassification rate shared by all patients. + +Each family uses three error levels chosen so the mid panel keeps KoMbine near the observed-label HR while the large-$e$ panel shows assignment search: + +- n=20 (`*_hr_example*`, always): $e = 0.20$, $0.25$, $0.40$ — KM bands stay similar at $0.20$ vs $0.25$ and widen / collapse at $e=0.40$ +- n=50 (`methods_comparison_*`, default only): $e = 0.05$, $0.10$, $0.25$ — KM bands stay similar (observed-label basin); the high panel is the joint HR assignment-search cliff + +`KOMBINE_QUICK_COMPARISON=1` skips the n=50 family. + +Each patient keeps the same survival time and censoring as that family's fixed baseline. +Only the class probabilities change: patients in the low group get probabilities +$(1-e, e)$, and patients in the high group get $(e, 1-e)$. + +MC-SIMEX uses the same $e$ as the per-patient flip rate. + +```python +import os +import sys +import datetime +import numpy as np +import matplotlib.pyplot as plt +import pathlib + +# Repo root on path so docs.kombine imports work under nbconvert (cwd is docs/kombine). +_repo_root = pathlib.Path(".").resolve().parent.parent +if str(_repo_root) not in sys.path: + sys.path.insert(0, str(_repo_root)) + +from kombine.datacard import Datacard +from kombine.comparisons import YiCorrectionForCoxPH + +# Single notebook budget. Edit these if you want denser scans. +N_PERMUTATIONS = 99 +N_HR_SCAN = 25 +SIMEX_B = 20 +LABEL_WIDTH = 9 +HAZARD_RATIOS_SCAN = np.logspace(-2, 2, N_HR_SCAN) # 0.01 to 100 + +# Default: n=20 then n=50. Set KOMBINE_QUICK_COMPARISON=1 to skip n=50 (CI). +QUICK_COMPARISON = bool(os.environ.get("KOMBINE_QUICK_COMPARISON")) +SIMEX_RNG = 0 + +TITLE_QUICK = "n=20, discrete e=0.20/0.25/0.40" +TITLE_FULL = "n=50, discrete e=0.05/0.10/0.25" + +MOSAIC_LAYOUT = [ + ['.', 'fixed', '.'], + ['dc_small', 'dc_moderate', 'dc_large'], + ['pois_large', 'pois_moderate', 'pois_small'], +] +MOSAIC_TO_SCENARIO = { + 'fixed': 'fixed', + 'dc_small': 'misclass_small', + 'dc_moderate': 'misclass_moderate', + 'dc_large': 'misclass_large', + 'pois_large': 'large', + 'pois_moderate': 'moderate', + 'pois_small': 'small', +} + +SCENARIOS_QUICK = { + 'fixed': { + 'file': 'fixed_hr_example.txt', + 'label': 'Fixed Observable', + 'description': 'no measurement error', + 'threshold': 0.5001, + }, + 'misclass_small': { + 'file': 'discrete_classes_hr_example_moderate.txt', + 'label': 'Disc. Classes (e=0.20)', + 'description': 'e = 0.20', + 'threshold': 1.0, + }, + 'misclass_moderate': { + 'file': 'discrete_classes_hr_example_large.txt', + 'label': 'Disc. Classes (e=0.25)', + 'description': 'e = 0.25', + 'threshold': 1.0, + }, + 'misclass_large': { + 'file': 'discrete_classes_hr_example_very_large.txt', + 'label': 'Disc. Classes (e=0.40)', + 'description': 'e = 0.40', + 'threshold': 1.0, + }, + 'large': { + 'file': 'poisson_density_hr_example_large.txt', + 'label': 'Poisson (large counts)', + 'description': '~2-3% relative error', + 'threshold': 0.5001, + }, + 'moderate': { + 'file': 'poisson_density_hr_example_moderate.txt', + 'label': 'Poisson (moderate counts)', + 'description': '~5-7% relative error', + 'threshold': 0.5001, + }, + 'small': { + 'file': 'poisson_density_hr_example_small.txt', + 'label': 'Poisson (small counts)', + 'description': '~25-70% relative error', + 'threshold': 0.5001, + }, +} + +SCENARIOS_FULL = { + 'fixed': { + 'file': 'methods_comparison_fixed.txt', + 'label': 'Fixed Observable', + 'description': 'no measurement error', + 'threshold': 0.5001, + }, + 'misclass_small': { + 'file': 'methods_comparison_discrete_e05.txt', + 'label': 'Disc. Classes (e=0.05)', + 'description': 'e = 0.05', + 'threshold': 1.0, + }, + 'misclass_moderate': { + 'file': 'methods_comparison_discrete_e10.txt', + 'label': 'Disc. Classes (e=0.10)', + 'description': 'e = 0.10', + 'threshold': 1.0, + }, + 'misclass_large': { + 'file': 'methods_comparison_discrete_e25.txt', + 'label': 'Disc. Classes (e=0.25)', + 'description': 'e = 0.25', + 'threshold': 1.0, + }, + 'large': { + 'file': 'methods_comparison_poisson_large.txt', + 'label': 'Poisson (large counts)', + 'description': '~2-5% relative error', + 'threshold': 0.5001, + }, + 'moderate': { + 'file': 'methods_comparison_poisson_moderate.txt', + 'label': 'Poisson (moderate counts)', + 'description': '~10-20% relative error', + 'threshold': 0.5001, + }, + 'small': { + 'file': 'methods_comparison_poisson_small.txt', + 'label': 'Poisson (small counts)', + 'description': '~25-70% relative error', + 'threshold': 0.5001, + }, +} + +COLORS_PALETTE = { + ('fixed', 'low'): '#0d47a1', + ('fixed', 'high'): '#6d1c1e', + ('misclass_small', 'low'): '#1b5e20', + ('misclass_small', 'high'): '#8e0000', + ('misclass_moderate', 'low'): '#2e7d32', + ('misclass_moderate', 'high'): '#b71c1c', + ('misclass_large', 'low'): '#66bb6a', + ('misclass_large', 'high'): '#e57373', + ('large', 'low'): '#1976d2', + ('large', 'high'): '#e53935', + ('moderate', 'low'): '#26a69a', + ('moderate', 'high'): '#fb8c00', + ('small', 'low'): '#80cbc4', + ('small', 'high'): '#ffd54f', +} + + +def progress(msg: str) -> None: + """Print a progress line with a wall-clock timestamp.""" + print(f"[{datetime.datetime.now()}] {msg}", flush=True) + + +def format_pvalue(p: float) -> str: + """Ordinary decimals for p >= 1e-3; scientific notation below that.""" + if not np.isfinite(p): + return str(p) + if p >= 1e-3: + return f'{p:.4g}' + return f'{p:.3e}' + + +def skip_n50(what: str) -> None: + """Print why the n=50 family is omitted under KOMBINE_QUICK_COMPARISON.""" + progress( + f"Skipping n=50 {what} because KOMBINE_QUICK_COMPARISON is set. " + "Unset it to run both the n=20 (e=0.20/0.25/0.40) and n=50 families." + ) + + +def load_datacards(directory, scenarios): + """Parse each scenario datacard and print n / deaths / distinct times.""" + loaded = {} + for key, info in scenarios.items(): + datacard = Datacard.parse_datacard(directory / info['file']) + loaded[key] = datacard + n_patients = len(datacard.patients) + n_deaths = sum(1 for patient in datacard.patients if not patient.censored) + uniq_deaths = len({ + round(patient.time, 10) + for patient in datacard.patients if not patient.censored + }) + progress( + f"{info['label']}: {n_patients} patients, {n_deaths} deaths, " + f"{uniq_deaths} distinct death times" + ) + return loaded + + +def run_km_analysis(scenarios, datacards): + """Yi, MC-SIMEX, and KoMbine KM (with KoMbine 95% bands) for every scenario.""" + km_results = {} + for scenario_key, scenario_info in scenarios.items(): + progress(f"[{scenario_info['label']}] Analysis 1 starting…") + dc = datacards[scenario_key] + threshold = scenario_info['threshold'] + + km_low = dc.km_likelihood( + parameter_min=-np.inf, + parameter_max=threshold, + ) + km_high = dc.km_likelihood( + parameter_min=threshold, + parameter_max=np.inf, + ) + + times_low = sorted(km_low.patient_death_times) + times_high = sorted(km_high.patient_death_times) + times_low_plot = [0.0] + times_low + times_high_plot = [0.0] + times_high + progress( + f" death times: low={len(times_low)} high={len(times_high)}" + ) + + progress(" Yi…") + result_low_yi = dc.km_survival_yi( + parameter_min=-np.inf, + parameter_max=threshold, + times_for_plot=times_low_plot, + ) + result_high_yi = dc.km_survival_yi( + parameter_min=threshold, + parameter_max=np.inf, + times_for_plot=times_high_plot, + ) + + progress(f" MC-SIMEX (B={SIMEX_B})…") + result_low_simex = dc.km_survival_mc_simex( + parameter_min=-np.inf, + parameter_max=threshold, + times_for_plot=times_low_plot, + rng=SIMEX_RNG, + B=SIMEX_B, + ) + result_high_simex = dc.km_survival_mc_simex( + parameter_min=threshold, + parameter_max=np.inf, + times_for_plot=times_high_plot, + rng=SIMEX_RNG, + B=SIMEX_B, + ) + + progress(" KoMbine low arm (full NLL, CLs=[0.95], crossing_mode=feasibility)…") + best_low, ci_low = km_low.survival_probabilities_likelihood( + CLs=[0.95], + times_for_plot=times_low, + binomial_only=(scenario_key == 'fixed'), + print_progress=True, + crossing_mode="feasibility", + ) + progress(" KoMbine high arm (full NLL, CLs=[0.95], crossing_mode=feasibility)…") + best_high, ci_high = km_high.survival_probabilities_likelihood( + CLs=[0.95], + times_for_plot=times_high, + binomial_only=(scenario_key == 'fixed'), + print_progress=True, + crossing_mode="feasibility", + ) + + km_results[scenario_key] = { + 'yi': { + 'low': result_low_yi, + 'high': result_high_yi, + }, + 'simex': { + 'low': result_low_simex, + 'high': result_high_simex, + }, + 'kombine': { + 'low': { + 'times': times_low, + 'best': best_low, + 'ci': ci_low, + }, + 'high': { + 'times': times_high, + 'best': best_high, + 'ci': ci_high, + } + } + } + + print(f"\n{scenario_info['label']}:") + print(f" {'Yi':<{LABEL_WIDTH}} - Low group final survival: {result_low_yi['survival_probabilities'][-1]:.4f}") + print(f" {'Yi':<{LABEL_WIDTH}} - High group final survival: {result_high_yi['survival_probabilities'][-1]:.4f}") + print(f" {'MC-SIMEX':<{LABEL_WIDTH}} - Low group final survival: {result_low_simex['survival_probabilities'][-1]:.4f}") + print(f" {'MC-SIMEX':<{LABEL_WIDTH}} - High group final survival: {result_high_simex['survival_probabilities'][-1]:.4f}") + + if len(ci_low) > 0: + ci_low_lower = ci_low[-1, 0, 0] + ci_low_upper = ci_low[-1, 0, 1] + print(f" {'KoMbine':<{LABEL_WIDTH}} - Low group final survival: {best_low[-1]:.4f} [{ci_low_lower:.4f}, {ci_low_upper:.4f}]") + else: + print(f" {'KoMbine':<{LABEL_WIDTH}} - Low group final survival: {best_low[-1]:.4f}") + + if len(ci_high) > 0: + ci_high_lower = ci_high[-1, 0, 0] + ci_high_upper = ci_high[-1, 0, 1] + print(f" {'KoMbine':<{LABEL_WIDTH}} - High group final survival: {best_high[-1]:.4f} [{ci_high_lower:.4f}, {ci_high_upper:.4f}]") + else: + print(f" {'KoMbine':<{LABEL_WIDTH}} - High group final survival: {best_high[-1]:.4f}") + return km_results + + +def _plot_km_in_ax(ax, scenario_key, scenario_info, result): + """Plot KM curves (Yi dashed, SIMEX dotted, KoMbine solid + CI shading).""" + color_low = COLORS_PALETTE[(scenario_key, 'low')] + color_high = COLORS_PALETTE[(scenario_key, 'high')] + + times_low_yi = result['yi']['low']['times_for_plot'] + surv_low_yi = result['yi']['low']['survival_probabilities'] + ax.step(times_low_yi, surv_low_yi, where='post', linewidth=2.5, + color=color_low, alpha=0.7, linestyle='--', label='Yi: Low group') + + times_high_yi = result['yi']['high']['times_for_plot'] + surv_high_yi = result['yi']['high']['survival_probabilities'] + ax.step(times_high_yi, surv_high_yi, where='post', linewidth=2.5, + color=color_high, alpha=0.7, linestyle='--', label='Yi: High group') + + times_low_simex = result['simex']['low']['times_for_plot'] + surv_low_simex = result['simex']['low']['survival_probabilities'] + ax.step(times_low_simex, surv_low_simex, where='post', linewidth=2.0, + color=color_low, alpha=0.85, linestyle=':', label='MC-SIMEX: Low group') + + times_high_simex = result['simex']['high']['times_for_plot'] + surv_high_simex = result['simex']['high']['survival_probabilities'] + ax.step(times_high_simex, surv_high_simex, where='post', linewidth=2.0, + color=color_high, alpha=0.85, linestyle=':', label='MC-SIMEX: High group') + + times_low_kombine = result['kombine']['low']['times'] + best_low_kombine = result['kombine']['low']['best'] + ci_low_kombine = result['kombine']['low']['ci'] + times_plot_low = [times_low_kombine[0]] + best_plot_low = [1.0] + for i, t in enumerate(times_low_kombine): + times_plot_low.append(t) + best_plot_low.append(best_low_kombine[i]) + ax.step(times_plot_low, best_plot_low, where='post', linewidth=2.5, + color=color_low, alpha=0.9, label='KoMbine: Low group', zorder=3) + if getattr(ci_low_kombine, 'size', 0): + ci_lower_plot_low = [1.0] + ci_upper_plot_low = [1.0] + for i, _t in enumerate(times_low_kombine): + ci_lower_plot_low.append(ci_low_kombine[i, 0, 0]) + ci_upper_plot_low.append(ci_low_kombine[i, 0, 1]) + ax.fill_between(times_plot_low, ci_lower_plot_low, ci_upper_plot_low, + step='post', alpha=0.15, color=color_low, + label='KoMbine: Low 95% CI', zorder=2) + + times_high_kombine = result['kombine']['high']['times'] + best_high_kombine = result['kombine']['high']['best'] + ci_high_kombine = result['kombine']['high']['ci'] + times_plot_high = [times_high_kombine[0]] + best_plot_high = [1.0] + for i, t in enumerate(times_high_kombine): + times_plot_high.append(t) + best_plot_high.append(best_high_kombine[i]) + ax.step(times_plot_high, best_plot_high, where='post', linewidth=2.5, + color=color_high, alpha=0.9, label='KoMbine: High group', zorder=3) + if getattr(ci_high_kombine, 'size', 0): + ci_lower_plot_high = [1.0] + ci_upper_plot_high = [1.0] + for i, _t in enumerate(times_high_kombine): + ci_lower_plot_high.append(ci_high_kombine[i, 0, 0]) + ci_upper_plot_high.append(ci_high_kombine[i, 0, 1]) + ax.fill_between(times_plot_high, ci_lower_plot_high, ci_upper_plot_high, + step='post', alpha=0.15, color=color_high, + label='KoMbine: High 95% CI', zorder=2) + + ax.set_xlabel('Time', fontsize=10) + ax.set_ylabel('Survival Probability', fontsize=10) + ax.set_title(scenario_info['label'], fontsize=11, fontweight='bold') + ax.legend(fontsize=8, loc='lower left') + ax.grid(True, alpha=0.3) + ax.set_ylim([0, 1.05]) + + +def _annotate_comparison_mosaic(axes_dict): + """Column headers and row labels shared by KM and HR mosaics.""" + col_headers_list = ['Small Uncertainty', 'Medium Uncertainty', 'Large Uncertainty'] + for panel_key, header in zip(['dc_small', 'dc_moderate', 'dc_large'], col_headers_list): + axes_dict[panel_key].annotate( + header, xy=(0.5, 1.0), xytext=(0, 30), + xycoords='axes fraction', textcoords='offset points', + ha='center', va='bottom', fontsize=12, fontweight='bold', + color='#333333', annotation_clip=False, + ) + for panel_key, row_label in zip(['dc_small', 'pois_large'], + ['Discrete\nClasses', 'Poisson\nCounts']): + axes_dict[panel_key].annotate( + row_label, xy=(0, 0.5), xytext=(-52, 0), + xycoords='axes fraction', textcoords='offset points', + ha='center', va='center', fontsize=11, fontweight='bold', + color='#333333', rotation=90, annotation_clip=False, + ) + + +def plot_km_mosaic(scenarios, km_results, title): + """3-row comparison mosaic of KM curves.""" + _, axes_dict = plt.subplot_mosaic(MOSAIC_LAYOUT, figsize=(14, 13), # pyright: ignore[reportCallIssue, reportArgumentType] + gridspec_kw={'hspace': 0.52, 'wspace': 0.35}, + ) + for panel_key, scenario_key in MOSAIC_TO_SCENARIO.items(): + _plot_km_in_ax(axes_dict[panel_key], scenario_key, + scenarios[scenario_key], km_results[scenario_key]) + _annotate_comparison_mosaic(axes_dict) + plt.suptitle(title, fontsize=14, fontweight='bold') + plt.tight_layout() + plt.show() + + +def run_pvalue_analysis(scenarios, datacards): + """Yi / MC-SIMEX logrank p-values and KoMbine permutation LRT.""" + pvalue_results = {} + for scenario_key, scenario_info in scenarios.items(): + dc = datacards[scenario_key] + threshold = scenario_info['threshold'] + progress(f"[{scenario_info['label']}] Analysis 2 starting…") + + progress(" Yi logrank…") + yi_result = dc.km_p_value_logrank_yi( + parameter_threshold=threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + + progress(f" MC-SIMEX logrank (B={SIMEX_B})…") + simex_result = dc.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + rng=SIMEX_RNG, + B=SIMEX_B, + ) + + progress(f" KoMbine permutation LRT (B={N_PERMUTATIONS})…") + kombine_calc = dc.km_p_value( + parameter_threshold=threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + pval_kombine, _, _ = kombine_calc.solve_and_pvalue( + cox_only=False, + n_permutations=N_PERMUTATIONS, + rng=SIMEX_RNG, + print_progress=True, + ) + + pvalue_results[scenario_key] = { + 'yi': yi_result['p_value'], + 'simex': simex_result['p_value'], + 'kombine': pval_kombine + } + + print(f"\n{scenario_info['label']}:") + print(f" {'Yi':<{LABEL_WIDTH}} p-value: {format_pvalue(yi_result['p_value'])}") + print(f" {'MC-SIMEX':<{LABEL_WIDTH}} p-value: {format_pvalue(simex_result['p_value'])}") + print(f" {'KoMbine':<{LABEL_WIDTH}} p-value: {format_pvalue(pval_kombine)}") + return pvalue_results + + +def plot_pvalue_bars(scenarios, pvalue_results, title): + """Grouped bar chart of Yi / MC-SIMEX / KoMbine p-values.""" + _, ax1 = plt.subplots(figsize=(14, 5)) + scenario_keys = list(scenarios.keys()) + scenario_labels = [scenarios[k]['label'] for k in scenario_keys] + yi_pvals = [pvalue_results[k]['yi'] for k in scenario_keys] + simex_pvals = [pvalue_results[k]['simex'] for k in scenario_keys] + kombine_pvals = [pvalue_results[k]['kombine'] for k in scenario_keys] + + x = np.arange(len(scenario_labels)) + width = 0.25 + bars1 = ax1.bar(x - width, yi_pvals, width, label="Yi's Method", color='steelblue') + bars2 = ax1.bar(x, simex_pvals, width, label='MC-SIMEX', color='mediumpurple') + bars3 = ax1.bar(x + width, kombine_pvals, width, label='KoMbine', color='coral') + + ax1.set_ylabel('P-value', fontsize=12) + ax1.set_title(title, fontsize=13, fontweight='bold') + ax1.set_xticks(x) + ax1.set_xticklabels(scenario_labels, rotation=15, ha='right') + ax1.legend(fontsize=11) + ax1.grid(True, alpha=0.3, axis='y') + + for bar_group in [bars1, bars2, bars3]: + for rect in bar_group: + height = rect.get_height() + ax1.text(rect.get_x() + rect.get_width()/2., height, + format_pvalue(height), ha='center', va='bottom', fontsize=8) + + plt.tight_layout() + plt.show() + + +def run_hr_analysis(scenarios, datacards): + """Yi Breslow profile, MC-SIMEX Wald, and KoMbine full-model HR profile.""" + hr_results = {} + for scenario_key, scenario_info in scenarios.items(): + dc = datacards[scenario_key] + hr_threshold = scenario_info['threshold'] + progress(f"[{scenario_info['label']}] Analysis 3 starting…") + + progress(" Yi HR CI…") + yi_calc = YiCorrectionForCoxPH( + patients=dc.patients, + parameter_min=-np.inf, + parameter_max=np.inf, + parameter_threshold=hr_threshold, + ) + best_hr_yi, yi_lower_ci, yi_upper_ci, yi_best_fit = ( + yi_calc.hazard_ratio_confidence_interval( + confidence_level=0.95, + hazard_ratio_min=0.01, + hazard_ratio_max=100.0, + ) + ) + yi_2nlls = [ + yi_calc.compute_2nll_at_hazard_ratio(hr).x + for hr in HAZARD_RATIOS_SCAN + ] + + progress(f" MC-SIMEX HR (B={SIMEX_B})…") + simex_calc = dc.km_hazard_ratio_mc_simex( + parameter_threshold=hr_threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + rng=SIMEX_RNG, + B=SIMEX_B, + ) + simex_estimate = simex_calc.estimate_hazard_ratio() + simex_2nlls = [ + simex_calc.compute_2nll_at_hazard_ratio(hr).x + for hr in HAZARD_RATIOS_SCAN + ] + + hr_calc = dc.km_hazard_ratio( + parameter_threshold=hr_threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + + progress(" KoMbine HR profile CI…") + best_hr_kombine, lower_ci, upper_ci, _ = hr_calc.hazard_ratio_confidence_interval( + cox_only=False, + confidence_level=0.95, + hazard_ratio_min=0.01, + hazard_ratio_max=100.0, + ) + + progress(f" KoMbine HR 2NLL scan ({N_HR_SCAN} points)…") + kombine_2nlls = [] + for hr in HAZARD_RATIOS_SCAN: + result = hr_calc.compute_2nll_at_hazard_ratio( + hr, cox_only=False, verbose=False, print_progress=True, + ) + kombine_2nlls.append(result.x) + + hr_results[scenario_key] = { + 'yi_best': best_hr_yi, + 'yi_2nlls': yi_2nlls, + 'yi_min_2nll': yi_best_fit.x, + 'yi_lower': yi_lower_ci, + 'yi_upper': yi_upper_ci, + 'simex_best': simex_estimate['hazard_ratio'], + 'simex_2nlls': simex_2nlls, + 'simex_lower': simex_estimate['ci_lower'], + 'simex_upper': simex_estimate['ci_upper'], + 'kombine_best': best_hr_kombine, + 'kombine_2nlls': kombine_2nlls, + 'kombine_lower': lower_ci, + 'kombine_upper': upper_ci, + } + + print(f"\n{scenario_info['label']}:") + print(f" {'Yi':<{LABEL_WIDTH}} best-fit HR: {best_hr_yi:.3f} [{yi_lower_ci:.3f}, {yi_upper_ci:.3f}]") + print(f" {'MC-SIMEX':<{LABEL_WIDTH}} best-fit HR: {simex_estimate['hazard_ratio']:.3f} [{simex_estimate['ci_lower']:.3f}, {simex_estimate['ci_upper']:.3f}]") + print(f" {'KoMbine':<{LABEL_WIDTH}} best-fit HR: {best_hr_kombine:.3f} [{lower_ci:.3f}, {upper_ci:.3f}]") + return hr_results + + +def plot_hr_mosaic(scenarios, hr_results, title): + """3-row comparison mosaic of HR Δ2NLL scans.""" + _, axes_dict = plt.subplot_mosaic(MOSAIC_LAYOUT, figsize=(14, 13), # pyright: ignore[reportCallIssue, reportArgumentType] + gridspec_kw={'hspace': 0.52, 'wspace': 0.35}, + ) + for panel_key, scenario_key in MOSAIC_TO_SCENARIO.items(): + ax = axes_dict[panel_key] + result = hr_results[scenario_key] + info = scenarios[scenario_key] + + yi_2nlls = result['yi_2nlls'] + delta_yi = np.array(yi_2nlls) - result['yi_min_2nll'] + delta_simex = np.array(result['simex_2nlls']) + kombine_2nlls = result['kombine_2nlls'] + delta_kombine = np.array(kombine_2nlls) - min(kombine_2nlls) + + ax.plot(HAZARD_RATIOS_SCAN, delta_yi, color='#1976d2', linewidth=2.5, marker='o', markersize=3, + label="Yi's Method", zorder=3) + ax.plot(HAZARD_RATIOS_SCAN, delta_simex, color='#7b1fa2', linewidth=2.5, marker='^', markersize=3, + linestyle=':', label='MC-SIMEX (Wald)', zorder=3) + ax.plot(HAZARD_RATIOS_SCAN, delta_kombine, color='#d32f2f', linewidth=2.5, marker='s', markersize=3, + label='KoMbine', zorder=3) + ax.axvline(result['yi_best'], color='#1976d2', linestyle='--', alpha=0.6, linewidth=1.5, zorder=2) + ax.axvline(result['simex_best'], color='#7b1fa2', linestyle=':', alpha=0.6, linewidth=1.5, zorder=2) + ax.axvline(result['kombine_best'], color='#d32f2f', linestyle='--', alpha=0.6, linewidth=1.5, zorder=2) + ax.axhline(3.84, color='gray', linestyle=':', alpha=0.6, linewidth=2.0, + label='95% CL (χ²=3.84)', zorder=1) + + ax.set_xlabel('Hazard Ratio', fontsize=10) + ax.set_ylabel(r'$-2 \Delta \ln L$', fontsize=10) + ax.set_title(info['label'], fontsize=11, fontweight='bold') + ax.legend(fontsize=8, loc='upper left') + ax.grid(True, alpha=0.3, which='both') + ax.set_xscale('log') + ax.set_xlim([0.01, 100.0]) + ax.set_ylim([0, 10]) + + _annotate_comparison_mosaic(axes_dict) + plt.suptitle(title, fontsize=14, fontweight='bold') + plt.tight_layout() + plt.show() + + +def print_summary_table(scenarios, pvalue_results, hr_results): + """Live p-value / HR table for one card family.""" + header = (f"{'Scenario':<36} {'Yi p':>9} {'SIMEX p':>9} {'KoMbine p':>11}" + f" {'Yi HR [95% CI]':>20} {'SIMEX HR [Wald]':>20} {'KoMbine HR [95% CI]':>22}") + print(header) + print('-' * len(header)) + for key, info in scenarios.items(): + pv = pvalue_results[key] + hr = hr_results[key] + yi_ci = (f"[{hr['yi_lower']:.3f}, {hr['yi_upper']:.3f}]" + if not np.isnan(hr['yi_lower']) else '[n/a]') + simex_ci = f"[{hr['simex_lower']:.3f}, {hr['simex_upper']:.3f}]" + ko_ci = f"[{hr['kombine_lower']:.3f}, {hr['kombine_upper']:.3f}]" + yi_hr_str = f"{hr['yi_best']:.3f} {yi_ci}" + simex_hr_str = f"{hr['simex_best']:.3f} {simex_ci}" + ko_hr_str = f"{hr['kombine_best']:.3f} {ko_ci}" + ko_p_str = 'n/a' if pv['kombine'] is None else format_pvalue(pv['kombine']) + print(f"{info['label']:<36} {format_pvalue(pv['yi']):>9} {format_pvalue(pv['simex']):>9} {ko_p_str:>11}" + f" {yi_hr_str:>20} {simex_hr_str:>20} {ko_hr_str:>22}") +``` + +```python +here = pathlib.Path(".").resolve() +test_dir = here.parent.parent / "test" / "kombine" +datacards_dir_quick = test_dir / "datacards" / "simple_examples" + +if QUICK_COMPARISON: + mode_name = "n=20 only (KOMBINE_QUICK_COMPARISON=1; n=50 cells will skip)" +else: + mode_name = "n=20 then n=50 (default)" + +progress(f"Notebook 07 mode: {mode_name}") +progress(f"Loading n=20 hr_example cards ({TITLE_QUICK})") +datacards_quick = load_datacards(datacards_dir_quick, SCENARIOS_QUICK) +``` + +## Analysis 1: Kaplan-Meier Curves + +Compare the Kaplan-Meier survival curves between Yi's method (dashed), MC-SIMEX (dotted), and KoMbine (solid lines with shaded 95% confidence intervals) across all scenarios. Yi and MC-SIMEX are point estimates only; they have no fill bands. + +### n=20 (`*_hr_example*`) + +```python +progress(f"Analysis 1 — {TITLE_QUICK}") +km_results_quick = run_km_analysis(SCENARIOS_QUICK, datacards_quick) +plot_km_mosaic( + SCENARIOS_QUICK, + km_results_quick, + f"Kaplan-Meier Curves: Yi, MC-SIMEX, and KoMbine ({TITLE_QUICK})", +) +``` + +### Understanding the Small Counts Anomaly: Why the High Group Can Look Better + +In the small-counts scenario, KoMbine’s best-fit *KM curves* can show the “high” group outperforming the “low” group. This is not necessarily a bug: it reflects how a **discrete-assignment likelihood** behaves when measurement error is large and many patients are effectively borderline relative to the threshold. + +**What is happening conceptually** +- With high uncertainty, several patients have substantial probability mass on both sides of the threshold. +- KoMbine must choose a single group per patient (within each optimization problem) and will pick the assignment(s) that maximize the likelihood. +- If early events are borderline, the likelihood can increase by assigning them to the group that best explains the observed survival times, which can make the *apparent* group ordering flip. + +**A subtle but important detail** +- In this notebook, the two KoMbine KM curves are fit **separately** (one KoMbine optimization for the low range and one for the high range). With large uncertainty, the best-fit assignments for those two separate optimizations need not form a perfectly complementary partition of patients. +- The KoMbine **hazard ratio / p-value** calculations, in contrast, are *joint* two-group fits. Those joint fits are the self-consistent way to summarize “which group is worse” under the KoMbine model. + +**Why Yi and MC-SIMEX look different** +- Yi’s method does not pick a single group; it spreads each borderline patient across both groups using probabilistic weights, which tends to shrink the group contrast as $e_i$ grows. +- MC-SIMEX keeps hard labels, adds extra flips, and extrapolates the naive curve; that can move the point estimate *away* from the naive (attenuated) curve rather than averaging the two groups. +- Neither method is a profile over discrete assignments, so both can disagree with KoMbine when membership is weakly identified. + +**Takeaway** +When measurement error is large, probabilistic weights (Yi), extrapolated hard-label estimators (MC-SIMEX), and discrete assignment (KoMbine) can tell qualitatively different stories. The right question is not “which is correct?”, but “how sensitive are the conclusions to how uncertain group membership is modeled?” + + +## Analysis 2: P-Values (Logrank / Likelihood-Ratio Test) + +We compare p-values from: + +- **Yi**: a *weighted* logrank-style calculation using per-patient probabilistic group membership (fractional membership). +- **MC-SIMEX**: the usual logrank statistic on extra-flipped hard labels, averaged vs $\lambda$ and extrapolated to $\lambda=-1$, then converted to a $\chi^2_1$ p-value. +- **KoMbine**: a *permutation* LRT of HR $=1$ using the full model (**`cox_only=False`**). Assignments are profiled on the observed data and on each shuffle of `(time, censored)`, so the null has the same reassignment freedom as the alternative. + +On the **fixed** card, Yi and MC-SIMEX therefore report the same hard-label logrank $\chi^2$ $p$, while KoMbine’s $p$ is a permutation LRT (here $B=99$) of the Cox alternative — they need not match even with zero measurement error. + +**What to expect** + +- As measurement error grows, **Yi's p-values typically increase** because fractional membership blurs the difference between groups. +- **MC-SIMEX p-values** are those of an extrapolated hard-label statistic; they need not match Yi even though both start from the same $e_i$. +- **KoMbine's permutation p-values** do not get smaller just because assignments become cheaper: the null can re-label too. +- Large disagreements among the three p-values indicate that inference is being driven by how group-membership uncertainty is modeled, not just by sampling noise. + +### n=20 (`*_hr_example*`) + +```python +progress(f"Analysis 2 — {TITLE_QUICK}") +pvalue_results_quick = run_pvalue_analysis(SCENARIOS_QUICK, datacards_quick) +plot_pvalue_bars( + SCENARIOS_QUICK, + pvalue_results_quick, + f"P-values: Yi / MC-SIMEX logrank χ² vs KoMbine permutation LRT ({TITLE_QUICK})", +) +``` + +## Analysis 3: Hazard Ratios + +We compare hazard ratios estimated using: + +- **Yi**: the weighted Breslow partial likelihood. The table reports a **continuous** MLE in $\log H$ and a likelihood-ratio interval (same `minimize_scalar` / `brentq` recipe as KoMbine). The curve is that 2NLL on an `N_HR_SCAN`-point log grid, recentered at the continuous MLE. +- **MC-SIMEX**: extrapolated $\widehat{\log H}$ with a Wald CI. The curve below is the Wald quadratic, **not** a profile likelihood. On the **fixed** panel membership is exact and the point HR matches Yi/KoMbine, but the purple scan is still Wald, so it will not overlay the Breslow profiles. +- **KoMbine**: the full profile likelihood (**`cox_only=False`**), which jointly optimizes discrete assignments and survival parameters. + +### Why the confidence intervals behave differently + +As noted in the paper text, Yi’s approach can reduce bias in the *point estimate* of the hazard ratio compared to ignoring measurement error, but the **confidence interval does not necessarily reflect loss of identifiability** when measurement error becomes very large. In the extreme-uncertainty limit we would expect the data to place almost no constraint on the hazard ratio, but Yi-style weighting does not automatically produce that behavior. + +MC-SIMEX is in the same family: the Wald interval is the sampling interval of the extrapolated number. It stays finite even when labels are nearly uninformative. + +KoMbine’s likelihood framework, by contrast, can naturally widen the profile-likelihood confidence interval as patient-wise uncertainty increases, because the model explicitly accounts for the possibility that the discrete group assignment itself is uncertain. + +### n=20 (`*_hr_example*`) + +```python +progress(f"Analysis 3 — {TITLE_QUICK}") +hr_results_quick = run_hr_analysis(SCENARIOS_QUICK, datacards_quick) +plot_hr_mosaic( + SCENARIOS_QUICK, + hr_results_quick, + f"Hazard Ratio: Yi and KoMbine Profiles vs MC-SIMEX Wald ({TITLE_QUICK})", +) +``` + + +## n=50 family (`methods_comparison_*`) + +Default (unset `KOMBINE_QUICK_COMPARISON`) repeats Analyses 1–3 on the stronger n=50 cards. CI sets the env var, so these cells print a skip message instead of solving. + +```python +datacards_full = None +km_results_full = None +pvalue_results_full = None +hr_results_full = None + +if QUICK_COMPARISON: + skip_n50("datacard load") +else: + from docs.kombine.rebin_methods_comparison_times import ensure_rebinned + datacards_dir_full = ensure_rebinned(n_bins=4) + progress(f"Loading n=50 rebinned methods_comparison cards ({TITLE_FULL})") + datacards_full = load_datacards(datacards_dir_full, SCENARIOS_FULL) +``` + +### Analysis 1 (n=50 KM) + +```python +if QUICK_COMPARISON: + skip_n50("KM analysis") +else: + progress(f"Analysis 1 — {TITLE_FULL}") + km_results_full = run_km_analysis(SCENARIOS_FULL, datacards_full) + plot_km_mosaic( + SCENARIOS_FULL, + km_results_full, + f"Kaplan-Meier Curves: Yi, MC-SIMEX, and KoMbine ({TITLE_FULL})", + ) +``` + +### Analysis 2 (n=50 p-values) + +```python +if QUICK_COMPARISON: + skip_n50("p-value analysis") +else: + progress(f"Analysis 2 — {TITLE_FULL}") + pvalue_results_full = run_pvalue_analysis(SCENARIOS_FULL, datacards_full) + plot_pvalue_bars( + SCENARIOS_FULL, + pvalue_results_full, + f"P-values: Yi / MC-SIMEX logrank χ² vs KoMbine permutation LRT ({TITLE_FULL})", + ) +``` + +### Analysis 3 (n=50 hazard ratios) + +```python +if QUICK_COMPARISON: + skip_n50("HR analysis") +else: + progress(f"Analysis 3 — {TITLE_FULL}") + hr_results_full = run_hr_analysis(SCENARIOS_FULL, datacards_full) + plot_hr_mosaic( + SCENARIOS_FULL, + hr_results_full, + f"Hazard Ratio: Yi and KoMbine Profiles vs MC-SIMEX Wald ({TITLE_FULL})", + ) +``` + + +## Summary of Findings + +The key modeling difference is **fractional weights vs extrapolated hard labels vs discrete assignment** under measurement uncertainty: +Yi’s method assigns each patient to both groups with weights, MC-SIMEX extrapolates a naive hard-label estimator, and KoMbine enforces one group per +patient and scores assignments using an explicit measurement-error model. + +The two card families are calibrated for different stories. The n=20 high-$e$ row is where KoMbine KM bands widen; the n=50 low-$e$ row is where the joint HR profile jumps into assignment search while the separate KM bands stay similar. + +### Kaplan–Meier Curves (Qualitative) +- **Fixed observable** (both families): Yi, MC-SIMEX, and KoMbine produce identical curves because group membership is exact. +- **Discrete classes, n=20** ($e=0.20$, $0.25$, $0.40$): KoMbine KM bands stay similar at $e=0.20$ vs $0.25$ (observed-label basin; the HR interval may already hit the scan bound). At $e=0.40$ the separate KM fits widen / collapse. That is the KM-widening figure. +- **Discrete classes, n=50** ($e=0.05$, $0.10$, $0.25$): At $e=0.05$–$0.10$ the KoMbine KM curves and point HR stay near the fixed baseline (HR $\approx 2.08$). At $e=0.25$ the KM bands remain similar (still not the $e=0.40$ collapse); Yi’s curves and HR drift toward 1 ($\widehat H \approx 1.40$), MC-SIMEX extrapolates the hard-label KM ($\widehat H \approx 2.86$), and the **joint** HR profile hits the scan bound ($\widehat H \approx 100$) with a wide interval. +- **Poisson (large/moderate counts)**: Yi shrinks the group gap ($\widehat H$ from 2.03 to 1.79 on n=50); MC-SIMEX is an extrapolated hard-label + curve with growing Wald intervals; KoMbine confidence bands and HR intervals widen as assignment uncertainty increases. +- **Poisson (small counts)**: Differences become qualitative (including apparent reversals in the + separate KM fits). On n=50, Yi is most conservative ($\widehat H \approx 1.58$, $p \approx 0.13$); MC-SIMEX can show the strongest separation ($\widehat H \approx 4.0$); KoMbine sits between them ($\widehat H \approx 3.0$) with a wide profile interval. + +### P-Values and Hazard Ratios +- Yi’s p-values generally increase as uncertainty grows; its best-fit HR drifts toward 1. + The printed Yi HR and CI are a continuous Breslow profile, not a grid argmin. +- MC-SIMEX p-values and HRs are those of an extrapolated hard-label statistic; the Wald HR interval stays finite. + The plotted MC-SIMEX curve is that Wald quadratic even when $e_i=0$. +- KoMbine’s plotted $p$ is a permutation LRT with $B=99$. + The profile interval can widen while the point HR remains moderate; on n=50 at $e=0.25$ the HR scan can pin at the upper bound. +- Large disagreements among the three indicate that inference is driven by how group-membership + uncertainty is modeled, not just by sampling noise. + +### Runtime +| Family | Analysis | Wall time | Notes | +|---|---|---:|---| +| n=20 | 1–3 | minutes | Always run; CI path | +| n=50 | 1 — KM bands | ~63 min (2026-09) | Dominated by Poisson (small counts) ~50 min | +| n=50 | 2 — P-values | ~21 min | Permutation LRT ($B=99$) | +| n=50 | 3 — HR profiles | ~0.1 min | Grid scan ($N\_HR\_SCAN=25$) | +| n=50 | **Total** | **~85 min** | Skipped when `KOMBINE_QUICK_COMPARISON=1` | + +### Practical Takeaways +1. When measurement error is tiny, KM curves and the Cox **point** HR agree. The HR *scan* + is still Wald (MC-SIMEX) vs two Breslow profiles (Yi, KoMbine), and KoMbine’s $p$ is a + permutation LRT rather than logrank $\chi^2$. +2. When measurement error is moderate/large, treat the conclusion as model-dependent and + report sensitivity to the modeling choice. +3. Yi’s method is fast and often conservative (it blurs separation as uncertainty grows). +4. MC-SIMEX is also fast; its Wald interval is a sampling interval of the extrapolated number. +5. KoMbine is likelihood-principled for the specified error model and can reveal when + parameters become weakly identified via widening profile-likelihood intervals. + +```python +progress(f"Summary table — {TITLE_QUICK}") +print_summary_table(SCENARIOS_QUICK, pvalue_results_quick, hr_results_quick) + +if QUICK_COMPARISON: + skip_n50("summary table") +else: + progress(f"Summary table — {TITLE_FULL}") + print_summary_table(SCENARIOS_FULL, pvalue_results_full, hr_results_full) +``` diff --git a/docs/kombine/07_yi_method_comparison.md b/docs/kombine/07_yi_method_comparison.md deleted file mode 100644 index ff694a0..0000000 --- a/docs/kombine/07_yi_method_comparison.md +++ /dev/null @@ -1,710 +0,0 @@ ---- -jupyter: - jupytext: - formats: ipynb,md,py:percent - text_representation: - extension: .md - format_name: markdown - format_version: '1.3' - jupytext_version: 1.18.1 - kernelspec: - display_name: rocpicker - language: python - name: python3 ---- - -```python -# pylint: disable=bad-indentation,line-too-long,missing-module-docstring,redefined-outer-name,trailing-whitespace,too-many-locals,wrong-import-order -``` - -# Yi's Method vs KoMbine: Comprehensive Comparison - -This notebook provides a comprehensive side-by-side comparison between Yi's method for Kaplan-Meier likelihood estimation and KoMbine's approach (MINLP) across multiple measurement scenarios: -- Fixed Hazard Ratio (deterministic, no measurement error) -- Discrete classes with class probabilities (small/medium/large uncertainty) -- Poisson density with large effect size (small relative error ~2-3%) -- Poisson density with moderate effect size (larger relative error ~5-7%) -- Poisson density with small counts (high relative error ~25-70%) -- Discrete class probabilities (controllable uncertainty via class distributions) - -Each analysis directly compares both methods to understand how they handle measurement uncertainty differently. - - -## Method Overview and Comparison - -| Aspect | Yi's Method | KoMbine | -|--------|---|---| -| **Core idea** | Weighted KM/logrank using probabilistic group membership | Full likelihood with explicit group assignment variables | -| **Optimization** | Direct probability calculation (no optimization) | Mixed Integer Nonlinear Programming (Gurobi) | -| **Computational cost** | Low | Medium-high | -| **Accuracy (within model)** | Approximate to the full likelihood | Exact maximizer within solver tolerance | -| **Core assumptions** | Known measurement error distribution; independent errors; fractional group membership is an adequate proxy for uncertain assignment | Known measurement error distribution; independent errors; patients belong to one group; event times treated as observed and discrete; likelihood model is correctly specified | - -### How Yi's Method Works (Intuition) -- Convert each patient's observed biomarker value into a probability of being below or above the threshold using the measurement error model. -- Instead of a global misclassification matrix as Yi describes, we extend her method to compute these probabilities on a per-patient basis (allowing uncertainty to vary by individual measurement). -- Use those per-patient probabilities as weights in the Kaplan-Meier estimator and logrank test. -- Every patient contributes to both groups, in proportion to their probability of belonging there. -- This yields fast, smooth estimates that tend to shrink group differences as measurement uncertainty grows. -- It is an approximation because it does not enforce a single, discrete group assignment for each patient. - -### How KoMbine Works (Intuition) -- Introduce a binary assignment variable for each patient (low vs high group). -- Combine the survival likelihood with a measurement error penalty that scores how plausible each assignment is. -- Solve a constrained optimization problem that finds the most likely set of assignments and survival parameters together. -- Compute confidence intervals via profile likelihood, which naturally widens as uncertainty increases. -- This is exact for the specified likelihood model but requires heavier computation. - -### What the Assumptions Mean (Plain Language) -- **Known measurement error distribution**: You have a reasonable model for how observed biomarker values deviate from the true value (e.g., Poisson noise). -- **Independent errors**: One patient's measurement error does not influence another patient's measurement error. -- **Fractional membership (Yi)**: It is acceptable to treat a patient as partly in each group, rather than forcing a single group. -- **Discrete membership (KoMbine)**: Each patient ultimately belongs to one group, but the model allows uncertainty in that assignment. -- **Correct likelihood model**: The survival model used (KM/logrank or Cox) is an appropriate description of the data-generating process. - -### Key Comparisons -1. **Kaplan-Meier Curves**: Visual comparison of survival estimates -2. **P-Values**: Logrank test results and statistical significance -3. **Hazard Ratios**: Point estimates and confidence intervals - - -## Discrete Classes with Class Probabilities - -Yi's Section 3.7.1 treats misclassification of a discrete covariate using a -misclassification matrix. We mimic that structure using two classes and a -homogeneous binary matrix, -$$ -\Pi = \begin{pmatrix} 1 - e & e \\ e & 1 - e \end{pmatrix} -$$ -where $e$ is the misclassification rate shared by all patients. - -To match the binary examples and exercises, we use three error levels: -- Small error: $e = 0.05$ -- Medium error: $e = 0.15$ -- Large error: $e = 0.30$ - -Each patient keeps the same survival time and censoring as the fixed baseline. -Only the class probabilities change: patients in the low group get probabilities -$(1-e, e)$, and patients in the high group get $(e, 1-e)$. - -```python -import numpy as np -import matplotlib.pyplot as plt -import pathlib -from kombine.datacard import Datacard -``` - -```python -# Setup - Load the comparison datacards -here = pathlib.Path(".").resolve() -test_dir = here.parent.parent / "test" / "kombine" -datacards_dir = test_dir / "datacards" / "simple_examples" - -# Define the scenarios -scenarios = { - 'fixed': { - 'file': 'fixed_hr_example.txt', - 'label': 'Fixed Observable', - 'description': 'no measurement error', - }, - 'misclass_small': { - 'file': 'discrete_classes_hr_example_small.txt', - 'label': 'Disc. Classes (e=0.05)', - 'description': 'e = 0.05', - }, - 'misclass_moderate': { - 'file': 'discrete_classes_hr_example_moderate.txt', - 'label': 'Disc. Classes (e=0.15)', - 'description': 'e = 0.15', - }, - 'misclass_large': { - 'file': 'discrete_classes_hr_example_large.txt', - 'label': 'Disc. Classes (e=0.30)', - 'description': 'e = 0.30', - }, - 'large': { - 'file': 'poisson_density_hr_example_large.txt', - 'label': 'Poisson (large counts)', - 'description': '~2-3% relative error', - }, - 'moderate': { - 'file': 'poisson_density_hr_example_moderate.txt', - 'label': 'Poisson (moderate counts)', - 'description': '~5-7% relative error', - }, - 'small': { - 'file': 'poisson_density_hr_example_small.txt', - 'label': 'Poisson (small counts)', - 'description': '~25-70% relative error', - }, -} - -# Load all datacards -datacards = {} -for key, info in scenarios.items(): - filepath = datacards_dir / info['file'] - datacard = Datacard.parse_datacard(filepath) - datacards[key] = datacard - n_patients = len(datacard.patients) - n_deaths = sum(1 for p in datacard.patients if not p.censored) - print(f"{info['label']}: {n_patients} patients, {n_deaths} deaths") - -threshold = 0.5001 -``` - -## Analysis 1: Kaplan-Meier Curves - -Compare the Kaplan-Meier survival curves between Yi's method (dashed lines) and KoMbine's approach (solid lines with shaded 95% confidence intervals) across all scenarios. This visualization directly shows how measurement error affects the survival curve estimates and their uncertainties. - -```python -# Calculate both Yi's weighted KM and KoMbine KM for each scenario -km_results = {} -label_width = 7 - -for scenario_key, scenario_info in scenarios.items(): - dc = datacards[scenario_key] - - # KoMbine method with confidence bands - km_low = dc.km_likelihood( - parameter_min=-np.inf, - parameter_max=threshold, - ) - - km_high = dc.km_likelihood( - parameter_min=threshold, - parameter_max=np.inf, - ) - - # Use the same time grids for Yi and KoMbine so curves align - times_low = sorted(km_low.patient_death_times) - times_high = sorted(km_high.patient_death_times) - times_low_plot = [0.0] + times_low - times_high_plot = [0.0] + times_high - - # Yi's method (each curve is a parameter range, weighted by membership probability) - result_low_yi = dc.km_survival_yi( - parameter_min=-np.inf, - parameter_max=threshold, - times_for_plot=times_low_plot, - ) - - result_high_yi = dc.km_survival_yi( - parameter_min=threshold, - parameter_max=np.inf, - times_for_plot=times_high_plot, - ) - - # Calculate best-fit and 95% CI for KoMbine - # Use full likelihood (not binomial_only) to include measurement uncertainty! - best_low, ci_low = km_low.survival_probabilities_likelihood( - CLs=[0.95], - times_for_plot=times_low, - binomial_only=(scenario_key == 'fixed'), # Only use binomial for fixed observable - ) - - best_high, ci_high = km_high.survival_probabilities_likelihood( - CLs=[0.95], - times_for_plot=times_high, - binomial_only=(scenario_key == 'fixed'), # Only use binomial for fixed observable - ) - - km_results[scenario_key] = { - 'yi': { - 'low': result_low_yi, - 'high': result_high_yi, - }, - 'kombine': { - 'low': { - 'times': times_low, - 'best': best_low, - 'ci': ci_low, - }, - 'high': { - 'times': times_high, - 'best': best_high, - 'ci': ci_high, - } - } - } - - print(f"\n{scenario_info['label']}:") - print(f" {'Yi':<{label_width}} - Low group final survival: {result_low_yi['survival_probabilities'][-1]:.4f}") - print(f" {'Yi':<{label_width}} - High group final survival: {result_high_yi['survival_probabilities'][-1]:.4f}") - - if len(ci_low) > 0: - ci_low_lower = ci_low[-1, 0, 0] - ci_low_upper = ci_low[-1, 0, 1] - print(f" {'KoMbine':<{label_width}} - Low group final survival: {best_low[-1]:.4f} [{ci_low_lower:.4f}, {ci_low_upper:.4f}]") - else: - print(f" {'KoMbine':<{label_width}} - Low group final survival: {best_low[-1]:.4f}") - - if len(ci_high) > 0: - ci_high_lower = ci_high[-1, 0, 0] - ci_high_upper = ci_high[-1, 0, 1] - print(f" {'KoMbine':<{label_width}} - High group final survival: {best_high[-1]:.4f} [{ci_high_lower:.4f}, {ci_high_upper:.4f}]") - else: - print(f" {'KoMbine':<{label_width}} - High group final survival: {best_high[-1]:.4f}") -``` - - -```python -# Define consistent color palette for all plots -colors_palette = { - ('fixed', 'low'): '#0d47a1', # Deep blue - ('fixed', 'high'): '#6d1c1e', # Deep red - ('misclass_small', 'low'): '#1b5e20', # Deep green - ('misclass_small', 'high'): '#8e0000', # Dark brick - ('misclass_moderate', 'low'): '#2e7d32', # Green - ('misclass_moderate', 'high'): '#b71c1c', # Dark red - ('misclass_large', 'low'): '#66bb6a', # Light green - ('misclass_large', 'high'): '#e57373', # Light red - ('large', 'low'): '#1976d2', # Strong blue - ('large', 'high'): '#e53935', # Strong red - ('moderate', 'low'): '#26a69a', # Teal - ('moderate', 'high'): '#fb8c00', # Orange - ('small', 'low'): '#80cbc4', # Light teal - ('small', 'high'): '#ffd54f', # Light amber -} - - -def _plot_km_in_ax(ax, scenario_key, scenario_info, result): - """Plot KM curves (Yi dashed, KoMbine solid + CI shading) in a single axes.""" - color_low = colors_palette[(scenario_key, 'low')] - color_high = colors_palette[(scenario_key, 'high')] - - times_low_yi = result['yi']['low']['times_for_plot'] - surv_low_yi = result['yi']['low']['survival_probabilities'] - ax.step(times_low_yi, surv_low_yi, where='post', linewidth=2.5, - color=color_low, alpha=0.7, linestyle='--', label='Yi: Low group') - - times_high_yi = result['yi']['high']['times_for_plot'] - surv_high_yi = result['yi']['high']['survival_probabilities'] - ax.step(times_high_yi, surv_high_yi, where='post', linewidth=2.5, - color=color_high, alpha=0.7, linestyle='--', label='Yi: High group') - - times_low_kombine = result['kombine']['low']['times'] - best_low_kombine = result['kombine']['low']['best'] - ci_low_kombine = result['kombine']['low']['ci'] - times_plot_low = [times_low_kombine[0]] - best_plot_low = [1.0] - ci_lower_plot_low = [1.0] - ci_upper_plot_low = [1.0] - for i, t in enumerate(times_low_kombine): - times_plot_low.append(t) - best_plot_low.append(best_low_kombine[i]) - ci_lower_plot_low.append(ci_low_kombine[i, 0, 0]) - ci_upper_plot_low.append(ci_low_kombine[i, 0, 1]) - ax.step(times_plot_low, best_plot_low, where='post', linewidth=2.5, - color=color_low, alpha=0.9, label='KoMbine: Low group', zorder=3) - ax.fill_between(times_plot_low, ci_lower_plot_low, ci_upper_plot_low, - step='post', alpha=0.15, color=color_low, label='KoMbine: Low 95% CI', zorder=2) - - times_high_kombine = result['kombine']['high']['times'] - best_high_kombine = result['kombine']['high']['best'] - ci_high_kombine = result['kombine']['high']['ci'] - times_plot_high = [times_high_kombine[0]] - best_plot_high = [1.0] - ci_lower_plot_high = [1.0] - ci_upper_plot_high = [1.0] - for i, t in enumerate(times_high_kombine): - times_plot_high.append(t) - best_plot_high.append(best_high_kombine[i]) - ci_lower_plot_high.append(ci_high_kombine[i, 0, 0]) - ci_upper_plot_high.append(ci_high_kombine[i, 0, 1]) - ax.step(times_plot_high, best_plot_high, where='post', linewidth=2.5, - color=color_high, alpha=0.9, label='KoMbine: High group', zorder=3) - ax.fill_between(times_plot_high, ci_lower_plot_high, ci_upper_plot_high, - step='post', alpha=0.15, color=color_high, label='KoMbine: High 95% CI', zorder=2) - - ax.set_xlabel('Time', fontsize=10) - ax.set_ylabel('Survival Probability', fontsize=10) - ax.set_title(scenario_info['label'], fontsize=11, fontweight='bold') - ax.legend(fontsize=9, loc='lower left') - ax.grid(True, alpha=0.3) - ax.set_ylim([0, 1.05]) - - -# Layout: fixed centered in row 0; cols aligned by uncertainty level -# Row 0: Fixed (baseline, centered) -# Row 1: Discrete Classes — small / moderate / large error -# Row 2: Poisson Counts — large / moderate / small counts -mosaic_layout = [ - ['.', 'fixed', '.'], - ['dc_small', 'dc_moderate', 'dc_large'], - ['pois_large', 'pois_moderate', 'pois_small'], -] -mosaic_to_scenario = { - 'fixed': 'fixed', - 'dc_small': 'misclass_small', - 'dc_moderate': 'misclass_moderate', - 'dc_large': 'misclass_large', - 'pois_large': 'large', - 'pois_moderate': 'moderate', - 'pois_small': 'small', -} - -fig, axes_dict = plt.subplot_mosaic(mosaic_layout, figsize=(14, 13), # pyright: ignore[reportCallIssue, reportArgumentType] - gridspec_kw={'hspace': 0.52, 'wspace': 0.35}, -) - -for panel_key, scenario_key in mosaic_to_scenario.items(): - _plot_km_in_ax(axes_dict[panel_key], scenario_key, - scenarios[scenario_key], km_results[scenario_key]) - -# Column headers: uncertainty level above row 1 -col_headers_list = ['Small Uncertainty', 'Medium Uncertainty', 'Large Uncertainty'] -for panel_key, header in zip(['dc_small', 'dc_moderate', 'dc_large'], col_headers_list): - axes_dict[panel_key].annotate( - header, xy=(0.5, 1.0), xytext=(0, 30), - xycoords='axes fraction', textcoords='offset points', - ha='center', va='bottom', fontsize=12, fontweight='bold', - color='#333333', annotation_clip=False, - ) - -# Row labels: observable type on the left of the first column -for panel_key, row_label in zip(['dc_small', 'pois_large'], - ['Discrete\nClasses', 'Poisson\nCounts']): - axes_dict[panel_key].annotate( - row_label, xy=(0, 0.5), xytext=(-52, 0), - xycoords='axes fraction', textcoords='offset points', - ha='center', va='center', fontsize=11, fontweight='bold', - color='#333333', rotation=90, annotation_clip=False, - ) - -plt.suptitle('Kaplan-Meier Curves: Yi vs KoMbine Across Measurement Scenarios', - fontsize=14, fontweight='bold') -plt.tight_layout() -plt.show() -``` - -### Understanding the Small Counts Anomaly: Why the High Group Can Look Better - -In the small-counts scenario, KoMbine’s best-fit *KM curves* can show the “high” group outperforming the “low” group. This is not necessarily a bug: it reflects how a **discrete-assignment likelihood** behaves when measurement error is large and many patients are effectively borderline relative to the threshold. - -**What is happening conceptually** -- With high uncertainty, several patients have substantial probability mass on both sides of the threshold. -- KoMbine must choose a single group per patient (within each optimization problem) and will pick the assignment(s) that maximize the likelihood. -- If early events are borderline, the likelihood can increase by assigning them to the group that best explains the observed survival times, which can make the *apparent* group ordering flip. - -**A subtle but important detail** -- In this notebook, the two KoMbine KM curves are fit **separately** (one KoMbine optimization for the low range and one for the high range). With large uncertainty, the best-fit assignments for those two separate optimizations need not form a perfectly complementary partition of patients. -- The KoMbine **hazard ratio / p-value** calculations, in contrast, are *joint* two-group fits. Those joint fits are the self-consistent way to summarize “which group is worse” under the KoMbine model. - -**Why Yi looks different** -- Yi’s method does not pick a single group; it spreads each borderline patient across both groups using probabilistic weights. -- That softens the group contrast and tends to avoid hard reversals; as uncertainty grows, Yi’s estimates drift toward each other. - -**Takeaway** -When measurement error is large, probabilistic assignment (Yi) and discrete assignment (KoMbine) can tell qualitatively different stories. The right question is not “which is correct?”, but “how sensitive are the conclusions to how uncertain group membership is modeled?” - - -## Analysis 2: P-Values (Logrank / Likelihood-Ratio Test) - -We compare p-values from: - -- **Yi**: a *weighted* logrank-style calculation using per-patient probabilistic group membership (fractional membership). -- **KoMbine**: a *likelihood-based* p-value using the full model (**`cox_only=False`**), which allows discrete group assignments to change in the fit when measurements are uncertain. - -**What to expect** - -- As measurement error grows, **Yi's p-values typically increase** because fractional membership blurs the difference between groups. -- **KoMbine's p-values need not increase**: because it enforces discrete membership, it can keep (or even increase) apparent separation by choosing the most likely global assignment consistent with the assumed measurement-error model. -- Large disagreements between the two p-values indicate that inference is being driven by how group-membership uncertainty is modeled, not just by sampling noise. - -```python -# Calculate p-values (Yi vs KoMbine) for all scenarios -pvalue_results = {} -label_width = 7 - -for scenario_key, scenario_info in scenarios.items(): - dc = datacards[scenario_key] - - # Yi's method - yi_result = dc.km_p_value_logrank_yi( - parameter_threshold=threshold, - parameter_min=-np.inf, - parameter_max=np.inf, - ) - - # KoMbine (full likelihood; includes patient-wise uncertainty) - kombine_calc = dc.km_p_value( - parameter_threshold=threshold, - parameter_min=-np.inf, - parameter_max=np.inf, - ) - pval_kombine, _, _ = kombine_calc.solve_and_pvalue(cox_only=False) - - pvalue_results[scenario_key] = { - 'yi': yi_result['p_value'], - 'kombine': pval_kombine - } - - print(f"\n{scenario_info['label']}:") - print(f" {'Yi':<{label_width}} p-value: {yi_result['p_value']:.4e}") - print(f" {'KoMbine':<{label_width}} p-value: {pval_kombine:.4e}") - rel_diff = abs(yi_result['p_value'] - pval_kombine) / min(yi_result['p_value'], pval_kombine) * 100 - print(f" Relative diff: {rel_diff:.1f}%") -``` - -```python -# Plot p-value comparison -fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) - -# Prepare data -scenario_keys = list(scenarios.keys()) -scenario_labels = [scenarios[k]['label'] for k in scenario_keys] -yi_pvals = [pvalue_results[k]['yi'] for k in scenario_keys] -kombine_pvals = [pvalue_results[k]['kombine'] for k in scenario_keys] - -# Bar plot -x = np.arange(len(scenario_labels)) -width = 0.35 - -bars1 = ax1.bar(x - width/2, yi_pvals, width, label="Yi's Method", color='steelblue') -bars2 = ax1.bar(x + width/2, kombine_pvals, width, label='KoMbine', color='coral') - -ax1.set_ylabel('P-value', fontsize=12) -ax1.set_title('Logrank Test P-Values Comparison', fontsize=13, fontweight='bold') -ax1.set_xticks(x) -ax1.set_xticklabels(scenario_labels, rotation=15, ha='right') -ax1.legend(fontsize=11) -ax1.grid(True, alpha=0.3, axis='y') - -# Add value labels on bars -for bars in [bars1, bars2]: - for bar in bars: - height = bar.get_height() - ax1.text(bar.get_x() + bar.get_width()/2., height, - f'{height:.3e}', ha='center', va='bottom', fontsize=9) - -# Relative difference plot -differences = [abs(yi_pvals[i] - kombine_pvals[i]) for i in range(len(scenario_labels))] -rel_diffs = [differences[i] / min(yi_pvals[i], kombine_pvals[i]) * 100 for i in range(len(scenario_labels))] - -ax2.bar(scenario_labels, rel_diffs, color='green', alpha=0.7) -ax2.set_ylabel('Relative Difference (%)', fontsize=12) -ax2.set_title('Relative Difference: |Yi - KoMbine| / min(Yi, KoMbine)', fontsize=13, fontweight='bold') -ax2.set_xticklabels(scenario_labels, rotation=15, ha='right') -ax2.grid(True, alpha=0.3, axis='y') - -# Add value labels -for i, (label, val) in enumerate(zip(scenario_labels, rel_diffs)): - ax2.text(i, val, f'{val:.1f}%', ha='center', va='bottom', fontsize=10) - -plt.tight_layout() -plt.show() -``` - -## Analysis 3: Hazard Ratios - -We compare hazard ratios estimated using: - -- **Yi**: a weighted likelihood / weighted logrank-style construction based on fractional group membership. -- **KoMbine**: the full profile likelihood (**`cox_only=False`**), which jointly optimizes discrete assignments and survival parameters. - -### Why the confidence intervals behave differently - -As noted in the paper text, Yi’s approach can reduce bias in the *point estimate* of the hazard ratio compared to ignoring measurement error, but the **confidence interval does not necessarily reflect loss of identifiability** when measurement error becomes very large. In the extreme-uncertainty limit we would expect the data to place almost no constraint on the hazard ratio, but Yi-style weighting does not automatically produce that behavior. - -KoMbine’s likelihood framework, by contrast, can naturally widen the profile-likelihood confidence interval as patient-wise uncertainty increases, because the model explicitly accounts for the possibility that the discrete group assignment itself is uncertain. - -```python -# Calculate hazard ratios (Yi vs KoMbine) for all scenarios -hr_results = {} -label_width = 7 -hazard_ratios_scan = np.logspace(-2, 2, 80) # Match 04: 0.01 to 100 with 80 points -chi2_95 = 3.84 # chi2.ppf(0.95, df=1) for a 95% two-sided CI - -for scenario_key, scenario_info in scenarios.items(): - dc = datacards[scenario_key] - hr_threshold = threshold - - # Yi's method - profile likelihood scan - yi_2nlls = [] - for hr in hazard_ratios_scan: - result = dc.km_hazard_ratio_yi( - parameter_threshold=hr_threshold, - hazard_ratio=hr, - parameter_min=-np.inf, - parameter_max=np.inf, - ) - yi_2nlls.append(result.x) - - best_idx_yi = np.argmin(yi_2nlls) - best_hr_yi = hazard_ratios_scan[best_idx_yi] - - # Calculate Yi's 95% CI from profile likelihood - min_yi_2nll = min(yi_2nlls) - delta_yi_2nll = np.array(yi_2nlls) - min_yi_2nll - # Find HR values where delta crosses the 95% threshold - below_threshold = delta_yi_2nll < chi2_95 - if np.any(below_threshold): - yi_ci_indices = np.where(below_threshold)[0] - yi_lower_ci = hazard_ratios_scan[yi_ci_indices[0]] - yi_upper_ci = hazard_ratios_scan[yi_ci_indices[-1]] - else: - yi_lower_ci = yi_upper_ci = np.nan - - # KoMbine (full likelihood; allows assignments to change with HR) - hr_calc = dc.km_hazard_ratio( - parameter_threshold=hr_threshold, - parameter_min=-np.inf, - parameter_max=np.inf, - ) - - best_hr_kombine, lower_ci, upper_ci, _ = hr_calc.hazard_ratio_confidence_interval( - cox_only=False, - confidence_level=0.95, - hazard_ratio_min=0.01, - hazard_ratio_max=100.0, - ) - - # KoMbine profile likelihood scan - kombine_2nlls = [] - for hr in hazard_ratios_scan: - result = hr_calc.compute_2nll_at_hazard_ratio(hr, cox_only=False, verbose=False) - kombine_2nlls.append(result.x) - - hr_results[scenario_key] = { - 'yi_best': best_hr_yi, - 'yi_2nlls': yi_2nlls, - 'yi_lower': yi_lower_ci, - 'yi_upper': yi_upper_ci, - 'kombine_best': best_hr_kombine, - 'kombine_2nlls': kombine_2nlls, - 'kombine_lower': lower_ci, - 'kombine_upper': upper_ci, - } - - print(f"\n{scenario_info['label']}:") - print(f" {'Yi':<{label_width}} best-fit HR: {best_hr_yi:.3f} [{yi_lower_ci:.3f}, {yi_upper_ci:.3f}]") - print(f" {'KoMbine':<{label_width}} best-fit HR: {best_hr_kombine:.3f} [{lower_ci:.3f}, {upper_ci:.3f}]") - rel_hr_diff = abs(best_hr_yi - best_hr_kombine) / best_hr_kombine * 100 - print(f" Relative HR diff: {rel_hr_diff:.1f}%") -``` - - -```python -# Plot hazard ratio profiles for all scenarios in a grid -mosaic_layout = [ - ['.', 'fixed', '.'], - ['dc_small', 'dc_moderate', 'dc_large'], - ['pois_large', 'pois_moderate', 'pois_small'], -] -mosaic_to_scenario = { - 'fixed': 'fixed', - 'dc_small': 'misclass_small', - 'dc_moderate': 'misclass_moderate', - 'dc_large': 'misclass_large', - 'pois_large': 'large', - 'pois_moderate': 'moderate', - 'pois_small': 'small', -} - -fig, axes_dict = plt.subplot_mosaic(mosaic_layout, figsize=(14, 13), # pyright: ignore[reportCallIssue, reportArgumentType] - gridspec_kw={'hspace': 0.52, 'wspace': 0.35}, -) - -for panel_key, scenario_key in mosaic_to_scenario.items(): - ax = axes_dict[panel_key] - result = hr_results[scenario_key] - info = scenarios[scenario_key] - - yi_2nlls = result['yi_2nlls'] - delta_yi = np.array(yi_2nlls) - min(yi_2nlls) - - kombine_2nlls = result['kombine_2nlls'] - delta_kombine = np.array(kombine_2nlls) - min(kombine_2nlls) - - ax.plot(hazard_ratios_scan, delta_yi, color='#1976d2', linewidth=2.5, marker='o', markersize=3, - label="Yi's Method", zorder=3) - ax.plot(hazard_ratios_scan, delta_kombine, color='#d32f2f', linewidth=2.5, marker='s', markersize=3, - label='KoMbine', zorder=3) - ax.axvline(result['yi_best'], color='#1976d2', linestyle='--', alpha=0.6, linewidth=1.5, zorder=2) - ax.axvline(result['kombine_best'], color='#d32f2f', linestyle='--', alpha=0.6, linewidth=1.5, zorder=2) - ax.axhline(3.84, color='gray', linestyle=':', alpha=0.6, linewidth=2.0, - label='95% CL (χ²=3.84)', zorder=1) - - ax.set_xlabel('Hazard Ratio', fontsize=10) - ax.set_ylabel(r'$-2 \Delta \ln L$', fontsize=10) - ax.set_title(info['label'], fontsize=11, fontweight='bold') - ax.legend(fontsize=9, loc='upper left') - ax.grid(True, alpha=0.3, which='both') - ax.set_xscale('log') - ax.set_xlim([0.01, 100.0]) - ax.set_ylim([0, 10]) - -# Column headers: uncertainty level above row 1 -col_headers_list = ['Small Uncertainty', 'Medium Uncertainty', 'Large Uncertainty'] -for panel_key, header in zip(['dc_small', 'dc_moderate', 'dc_large'], col_headers_list): - axes_dict[panel_key].annotate( - header, xy=(0.5, 1.0), xytext=(0, 30), - xycoords='axes fraction', textcoords='offset points', - ha='center', va='bottom', fontsize=12, fontweight='bold', - color='#333333', annotation_clip=False, - ) - -# Row labels: observable type on the left of the first column -for panel_key, row_label in zip(['dc_small', 'pois_large'], - ['Discrete\nClasses', 'Poisson\nCounts']): - axes_dict[panel_key].annotate( - row_label, xy=(0, 0.5), xytext=(-52, 0), - xycoords='axes fraction', textcoords='offset points', - ha='center', va='center', fontsize=11, fontweight='bold', - color='#333333', rotation=90, annotation_clip=False, - ) - -plt.suptitle('Profile Likelihood for Hazard Ratio: Yi vs KoMbine', - fontsize=14, fontweight='bold') -plt.tight_layout() -plt.show() -``` - - -## Summary of Findings - -The key modeling difference is **fractional vs discrete assignment** under measurement uncertainty: -Yi’s method assigns each patient to both groups with weights, while KoMbine enforces one group per -patient and scores assignments using an explicit measurement-error model. - -### Kaplan–Meier Curves (Qualitative) -- **Fixed observable**: Yi and KoMbine produce identical curves because group membership is exact. -- **Discrete classes (small e)**: Curves remain close to the fixed baseline. As e grows, Yi’s curves - drift toward each other while KoMbine can maintain separation using the most likely discrete assignment. -- **Poisson (large/moderate counts)**: Similar behavior — Yi shrinks the group gap; KoMbine confidence - bands widen as assignment uncertainty increases. -- **Poisson (small counts)**: Differences can become qualitative (including apparent reversals in the - individual KM curve fits) because group membership is weakly identified under the error model. - -### P-Values and Hazard Ratios -- Yi’s p-values generally increase as uncertainty grows; its best-fit HR drifts toward 1. -- KoMbine’s p-values are more stable; its best-fit HR stays near the baseline, but the - **confidence interval widens** as uncertainty grows. -- Large disagreements between the two indicate that inference is driven by how group-membership - uncertainty is modeled, not just by sampling noise. - -### Practical Takeaways -1. When measurement error is tiny, both methods agree and the choice is less critical. -2. When measurement error is moderate/large, treat the conclusion as model-dependent and - report sensitivity to the modeling choice. -3. Yi’s method is fast and often conservative (it blurs separation as uncertainty grows). -4. KoMbine is likelihood-principled for the specified error model and can reveal when - parameters become weakly identified via widening profile-likelihood intervals. - -```python -# Summary tables — computed live from pvalue_results and hr_results -header = (f"{'Scenario':<36} {'Yi p-val':>10} {'KoMbine p-val':>14}" - f" {'Yi HR [95% CI]':>22} {'KoMbine HR [95% CI]':>25}") -print(header) -print('-' * len(header)) -for key, info in scenarios.items(): - pv = pvalue_results[key] - hr = hr_results[key] - yi_ci = (f"[{hr['yi_lower']:.3f}, {hr['yi_upper']:.3f}]" - if not np.isnan(hr['yi_lower']) else '[n/a]') - ko_ci = f"[{hr['kombine_lower']:.3f}, {hr['kombine_upper']:.3f}]" - yi_hr_str = f"{hr['yi_best']:.3f} {yi_ci}" - ko_hr_str = f"{hr['kombine_best']:.3f} {ko_ci}" - print(f"{info['label']:<36} {pv['yi']:>10.3e} {pv['kombine']:>14.3e}" - f" {yi_hr_str:>22} {ko_hr_str:>25}") -``` diff --git a/docs/kombine/08_command_line_interface.md b/docs/kombine/08_command_line_interface.md index 6225da8..09aed41 100644 --- a/docs/kombine/08_command_line_interface.md +++ b/docs/kombine/08_command_line_interface.md @@ -8,7 +8,7 @@ jupyter: extension: .md format_name: markdown format_version: '1.3' - jupytext_version: 1.18.1 + jupytext_version: 1.19.5 --- ```python diff --git a/docs/kombine/__init__.py b/docs/kombine/__init__.py index e69de29..43707ef 100644 --- a/docs/kombine/__init__.py +++ b/docs/kombine/__init__.py @@ -0,0 +1 @@ +"""KoMbine documentation notebooks and helpers.""" diff --git a/docs/kombine/rebin_methods_comparison_times.py b/docs/kombine/rebin_methods_comparison_times.py new file mode 100644 index 0000000..d9bd735 --- /dev/null +++ b/docs/kombine/rebin_methods_comparison_times.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Rebin survival times on the methods_comparison_* card family. + +Snaps every patient's survival_time onto K quantile-bin medians taken from the +fixed baseline card, then writes the same survival_time / censored rows into +every sibling methods_comparison_*.txt under a gitignored ``rebinned/`` +directory (observables unchanged; tracked sources stay fine-grained). + +Notebook 07's default mode expects K=4 (chosen after capped probes: +K=6 censored a KM arm past 15 min; K=4 finished a hard-card arm in ~8 min). + +Usage (from repo root):: + + python docs/kombine/rebin_methods_comparison_times.py + python docs/kombine/rebin_methods_comparison_times.py --n-bins 4 + +Or from Python (package import):: + + from docs.kombine.rebin_methods_comparison_times import ensure_rebinned + datacards_dir = ensure_rebinned(n_bins=4) +""" + +from __future__ import annotations + +import argparse +import pathlib + +import numpy as np + +SCRIPT_DIR = pathlib.Path(__file__).resolve().parent +CARDS_DIR = ( + SCRIPT_DIR.parent.parent / "test" / "kombine" / "datacards" / "simple_examples" +) +BASELINE = CARDS_DIR / "methods_comparison_fixed.txt" +REBINNED_DIR = CARDS_DIR / "rebinned" + + +def rebin_times(times: np.ndarray, n_bins: int) -> np.ndarray: + """Snap each time onto the median of its quantile bin.""" + if n_bins < 1: + raise ValueError(f"n_bins must be >= 1, got {n_bins}") + if n_bins >= len(times): + return times.copy() + edges = np.quantile(times, np.linspace(0.0, 1.0, n_bins + 1)) + edges[-1] += 1e-9 + idx = np.clip(np.searchsorted(edges, times, side="right") - 1, 0, n_bins - 1) + return np.array([float(np.median(times[idx == b])) for b in idx]) + + +def parse_rows(text: str) -> dict[str, list[str]]: + """Extract survival_time and censored columns from a datacard text blob.""" + rows: dict[str, list[str]] = {} + for line in text.splitlines(): + parts = line.split("\t") + if len(parts) > 1 and parts[0] in {"survival_time", "censored"}: + rows[parts[0]] = parts[1:] + if "survival_time" not in rows or "censored" not in rows: + raise ValueError("datacard missing survival_time or censored row") + return rows + + +def rewrite_card( # pylint: disable=too-many-arguments + source: pathlib.Path, + dest: pathlib.Path, + times: np.ndarray, + censored: list[str], + *, + n_bins: int, + n_unique: int, +) -> None: + """Write one card to ``dest`` with rebinned times and updated header comments.""" + text = source.read_text(encoding="utf-8") + lines = text.splitlines() + out: list[str] = [] + header_done = False + for line in lines: + if line.startswith("survival_time"): + out.append("\t".join(["survival_time"] + [f"{t:.4f}" for t in times])) + continue + if line.startswith("censored"): + out.append("\t".join(["censored"] + list(censored))) + continue + if ( + not header_done + and line.startswith("#") + and ("n=50" in line or "deaths=" in line or "Notebook 07" in line) + ): + # Replace the stats comment once with the rebinned description. + if "n=50" in line or "deaths=" in line: + out.append( + f"# n=50, n_death_times={n_unique} (quantile-bin medians, K={n_bins}), " + "target HR=6, n_per_group=25, seed=42" + ) + out.append( + "# Survival times rebinned from methods_comparison_fixed.txt; " + "shared across sibling methods_comparison_* cards" + ) + header_done = True + continue + out.append(line) + # If no stats comment was found, insert after the first # block opener. + if not header_done: + inserted = False + final: list[str] = [] + for line in out: + final.append(line) + if not inserted and line.startswith("#") and "Notebook 07" in line: + final.append( + f"# n=50, n_death_times={n_unique} (quantile-bin medians, K={n_bins}), " + "target HR=6, n_per_group=25, seed=42" + ) + final.append( + "# Survival times rebinned from methods_comparison_fixed.txt; " + "shared across sibling methods_comparison_* cards" + ) + inserted = True + out = final + dest.write_text("\n".join(out) + "\n", encoding="utf-8") + + +def ensure_rebinned(n_bins: int = 4, *, out_dir: pathlib.Path | None = None) -> pathlib.Path: + """Regenerate rebinned cards under ``out_dir`` and return that directory. + + Always rewrites (cheap). Tracked sources in ``CARDS_DIR`` are left unchanged. + """ + if n_bins < 1: + raise ValueError(f"n_bins must be >= 1, got {n_bins}") + dest_dir = out_dir if out_dir is not None else REBINNED_DIR + dest_dir.mkdir(parents=True, exist_ok=True) + + siblings = sorted(CARDS_DIR.glob("methods_comparison_*.txt")) + if not siblings: + raise FileNotFoundError(f"no methods_comparison_*.txt sources in {CARDS_DIR}") + if not BASELINE.is_file(): + raise FileNotFoundError(f"baseline card missing: {BASELINE}") + + baseline_text = BASELINE.read_text(encoding="utf-8") + rows = parse_rows(baseline_text) + raw_times = np.array([float(t) for t in rows["survival_time"]]) + censored = rows["censored"] + binned = rebin_times(raw_times, n_bins) + # Death times among uncensored patients drive the KM grid. + death_mask = [c not in ("1", "True", "true") for c in censored] + n_unique = len({round(t, 10) for t, d in zip(binned, death_mask) if d}) + print( + f"baseline n={len(raw_times)} raw_unique_deaths=" + f"{len({round(t, 10) for t, d in zip(raw_times, death_mask) if d})} " + f"-> K={n_bins} unique_death_times={n_unique}", + flush=True, + ) + + for source in siblings: + dest = dest_dir / source.name + rewrite_card( + source, dest, binned, censored, n_bins=n_bins, n_unique=n_unique, + ) + print(f"wrote {dest.relative_to(CARDS_DIR)}", flush=True) + return dest_dir + + +def main() -> None: + """CLI entry point.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--n-bins", type=int, default=4, + help="quantile bins (default 4; notebook full-comparison mode)", + ) + args = parser.parse_args() + ensure_rebinned(n_bins=args.n_bins) + + +if __name__ == "__main__": + main() diff --git a/kombine/__init__.py b/kombine/__init__.py index 4e7e9e0..eaa4568 100644 --- a/kombine/__init__.py +++ b/kombine/__init__.py @@ -7,7 +7,13 @@ from .kaplan_meier_MINLP import KaplanMeierPatientNLL, MINLPForKM from .kaplan_meier_p_value_MINLP import MINLPforKMPValue from .kaplan_meier_hazard_ratio_MINLP import MINLPforKMHazardRatio -from .yi_correction import YiCorrectionForLogrank, YiCorrectionForCoxPH +from .comparisons import ( + McSimexForCoxPH, + McSimexForKaplanMeier, + McSimexForLogrank, + YiCorrectionForCoxPH, + YiCorrectionForLogrank, +) from .datacard import Datacard __all__ = [ @@ -19,5 +25,8 @@ "MINLPforKMHazardRatio", "YiCorrectionForLogrank", "YiCorrectionForCoxPH", + "McSimexForLogrank", + "McSimexForCoxPH", + "McSimexForKaplanMeier", "Datacard", ] diff --git a/kombine/comparisons/__init__.py b/kombine/comparisons/__init__.py new file mode 100644 index 0000000..a77caa8 --- /dev/null +++ b/kombine/comparisons/__init__.py @@ -0,0 +1,33 @@ +""" +Comparison estimators for published alternatives to KoMbine. + +These methods answer related but distinct questions from KoMbine's profile +likelihood over discrete group assignments. They live here so the KoMbine +core package stays focused on the MINLP / likelihood analysis. +""" + +from .mc_simex import ( + McSimexBase, + McSimexForCoxPH, + McSimexForKaplanMeier, + McSimexForLogrank, +) +from .yi_correction import ( + YiCorrectionBase, + YiCorrectionForCoxPH, + YiCorrectionForKaplanMeier, + YiCorrectionForLogrank, + YiCorrectionWithThreshold, +) + +__all__ = [ + "McSimexBase", + "McSimexForCoxPH", + "McSimexForKaplanMeier", + "McSimexForLogrank", + "YiCorrectionBase", + "YiCorrectionForCoxPH", + "YiCorrectionForKaplanMeier", + "YiCorrectionForLogrank", + "YiCorrectionWithThreshold", +] diff --git a/kombine/comparisons/mc_simex.py b/kombine/comparisons/mc_simex.py new file mode 100644 index 0000000..656711a --- /dev/null +++ b/kombine/comparisons/mc_simex.py @@ -0,0 +1,779 @@ +""" +Küchenhoff MC-SIMEX for a binary high/low survival label. + +This is a comparison estimator, not a KoMbine profile likelihood. Each +simulation-extrapolation step refits a naive hard-label Kaplan-Meier, +logrank, or Cox functional after extra misclassification; the +error-free functional is a quadratic extrapolation to lambda = -1. + +Poisson or discrete-class measurements enter only through +``observable.probability_in_range``, which yields a per-patient flip +rate e_i. The survival model never sees the raw count. POI-SIMEX +(continuous Poisson density as a regression covariate) is intentionally +not implemented. + +The Wald quadratic returned by ``McSimexForCoxPH.compute_2nll_at_hazard_ratio`` +is not a profile likelihood. It is +(log H - log H_hat)^2 / sigma_hat^2, the same information as a Wald CI. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import scipy.optimize +import scipy.stats + +from ..kaplan_meier import KaplanMeierInstance, KaplanMeierPatient +from .yi_correction import YiCorrectionBase, YiCorrectionWithThreshold + +if TYPE_CHECKING: + from ..datacard import Patient + +DEFAULT_LAMBDA_GRID = (0.0, 0.5, 1.0, 1.5, 2.0) +DEFAULT_N_SIMULATIONS = 100 +FLIP_RATE_CEILING = 0.5 - 1e-12 +LOG_HAZARD_RATIO_BOUNDS = (-10.0, 10.0) +WALD_Z_95 = 1.96 +VARIANCE_FLOOR = 1e-12 + + +def flip_probability(misclassification_rate: float, simex_lambda: float) -> float: + """ + Extra-flip probability Pi^lambda for a binary MC-SIMEX step. + + P(flip | lambda, e) = (1 - (1 - 2 e)^lambda) / 2. + lambda = 0 is a no-op. Rates at or above 1/2 are clamped just below 1/2. + """ + if simex_lambda == 0.0: + return 0.0 + rate = min(max(float(misclassification_rate), 0.0), FLIP_RATE_CEILING) + return 0.5 * (1.0 - (1.0 - 2.0 * rate) ** simex_lambda) + + +def extrapolate_quadratic(lambda_grid: np.ndarray, values: np.ndarray) -> float: + """ + Quadratic (or linear if fewer than three finite points) fit in lambda, + evaluated at lambda = -1. + """ + finite = np.isfinite(values) + x_fit = np.asarray(lambda_grid, dtype=float)[finite] + y_fit = np.asarray(values, dtype=float)[finite] + if x_fit.size < 2: + raise ValueError( + "Need at least two finite SIMEX lambda points to extrapolate, " + f"got {x_fit.size}." + ) + degree = 2 if x_fit.size >= 3 else 1 + coefficients = np.polyfit(x_fit, y_fit, deg=degree) + return float(np.polyval(coefficients, -1.0)) + + +def _as_rng(rng: np.random.Generator | int | None) -> np.random.Generator: + if rng is None: + return np.random.default_rng() + if isinstance(rng, np.random.Generator): + return rng + return np.random.default_rng(rng) + + +def _naive_logrank_statistic( # pylint: disable=too-many-locals + times: np.ndarray, + censored: np.ndarray, + is_high: np.ndarray, +) -> float | None: + """ + Hard-label logrank statistic U^2 / V, or None if a group is empty. + """ + if not np.any(is_high) or not np.any(~is_high): + return None + + death_times = np.unique(times[~censored]) + if death_times.size == 0: + return None + + statistic_u = 0.0 + statistic_v = 0.0 + for death_time in death_times: + at_risk = times >= death_time + n_low = int(np.count_nonzero(at_risk & ~is_high)) + n_high = int(np.count_nonzero(at_risk & is_high)) + n_total = n_low + n_high + if n_total == 0: + continue + died = (times == death_time) & ~censored + d_low = int(np.count_nonzero(died & ~is_high)) + d_high = int(np.count_nonzero(died & is_high)) + d_total = d_low + d_high + if d_total == 0: + continue + expected_low = n_low * d_total / n_total + if n_total > 1: + variance_t = ( + n_low * n_high * d_total * (n_total - d_total) + ) / (n_total * n_total * (n_total - 1)) + else: + variance_t = 0.0 + statistic_u += d_low - expected_low + statistic_v += variance_t + + if statistic_v <= 0.0: + return 0.0 + return float(statistic_u * statistic_u / statistic_v) + + +def _risk_sets_at_death( + times: np.ndarray, + censored: np.ndarray, + is_high: np.ndarray, + death_time: float, +) -> tuple[int, int, int, int]: + at_risk = times >= death_time + r_low = int(np.count_nonzero(at_risk & ~is_high)) + r_high = int(np.count_nonzero(at_risk & is_high)) + died = (times == death_time) & ~censored + d_low = int(np.count_nonzero(died & ~is_high)) + d_high = int(np.count_nonzero(died & is_high)) + return r_low, r_high, d_low, d_high + + +def breslow_2nll( + times: np.ndarray, + censored: np.ndarray, + is_high: np.ndarray, + hazard_ratio: float, +) -> float: + """ + Breslow Cox partial 2NLL for a hard high/low label. + """ + if hazard_ratio <= 0.0: + raise ValueError(f"hazard_ratio must be positive, got {hazard_ratio}") + death_times = np.unique(times[~censored]) + log_hazard_ratio = np.log(hazard_ratio) + log_likelihood = 0.0 + for death_time in death_times: + r_low, r_high, d_low, d_high = _risk_sets_at_death( + times, censored, is_high, float(death_time) + ) + d_total = d_low + d_high + if d_total <= 0: + continue + denominator = r_low + hazard_ratio * r_high + if denominator <= 0.0: + denominator = VARIANCE_FLOOR + log_likelihood += d_high * log_hazard_ratio - d_total * np.log(denominator) + return float(-2.0 * log_likelihood) + + +def observed_information_log_hr( + times: np.ndarray, + censored: np.ndarray, + is_high: np.ndarray, + hazard_ratio: float, +) -> float: + """ + Observed information for theta = log H in the Breslow partial likelihood. + """ + death_times = np.unique(times[~censored]) + information = 0.0 + for death_time in death_times: + r_low, r_high, d_low, d_high = _risk_sets_at_death( + times, censored, is_high, float(death_time) + ) + d_total = d_low + d_high + if d_total <= 0 or r_low <= 0 or r_high <= 0: + continue + risk = r_low + hazard_ratio * r_high + information += d_total * hazard_ratio * r_high * r_low / (risk * risk) + return float(information) + + +def mle_log_hazard_ratio( + times: np.ndarray, + censored: np.ndarray, + is_high: np.ndarray, +) -> tuple[float, float] | None: + """ + Bounded MLE of log H and observed-information variance, or None if empty. + """ + if not np.any(is_high) or not np.any(~is_high): + return None + if not np.any(~censored): + return None + + def objective(log_hazard_ratio: float) -> float: + return breslow_2nll( + times, censored, is_high, float(np.exp(log_hazard_ratio)) + ) + + result = scipy.optimize.minimize_scalar( + objective, + bounds=LOG_HAZARD_RATIO_BOUNDS, + method='bounded', + ) + log_hazard_ratio = float(result.x) + hazard_ratio = float(np.exp(log_hazard_ratio)) + information = observed_information_log_hr( + times, censored, is_high, hazard_ratio + ) + if information <= 0.0: + variance = np.nan + else: + variance = 1.0 / information + return log_hazard_ratio, float(variance) + + +class _SimexMonteCarlo: # pylint: disable=too-few-public-methods + """ + Lambda grid, Monte Carlo size, and extra-flip draws shared by MC-SIMEX classes. + """ + + _lambda_grid: tuple[float, ...] + _B: int + _rng: np.random.Generator + _prior_alpha: float + _prior_beta: float + + def _init_simex( # pylint: disable=too-many-arguments,invalid-name + self, + *, + lambda_grid: tuple[float, ...] | list[float] | None, + B: int, + rng: np.random.Generator | int | None, + prior_alpha: float, + prior_beta: float, + ) -> None: + if B < 1: + raise ValueError(f"B must be >= 1, got {B}") + grid = DEFAULT_LAMBDA_GRID if lambda_grid is None else tuple(lambda_grid) + if any(lam < 0.0 for lam in grid): + raise ValueError( + "lambda_grid must be non-negative (do not simulate lambda = -1), " + f"got {grid}" + ) + if 0.0 not in grid: + grid = (0.0,) + grid + self._lambda_grid = tuple(sorted(set(float(lam) for lam in grid))) + self._B = int(B) + self._rng = _as_rng(rng) + self._prior_alpha = float(prior_alpha) + self._prior_beta = float(prior_beta) + + def simulate_labels( + self, + observed_labels: np.ndarray, + flip_rates: np.ndarray, + simex_lambda: float, + ) -> np.ndarray: + """ + Flip observed binary labels with the MC-SIMEX probability at lambda. + """ + if simex_lambda == 0.0: + return observed_labels.copy() + probabilities = np.array( + [flip_probability(rate, simex_lambda) for rate in flip_rates], + dtype=float, + ) + flips = self._rng.random(observed_labels.shape[0]) < probabilities + return observed_labels ^ flips + + +class McSimexBase(YiCorrectionBase, _SimexMonteCarlo): + """ + Shared MC-SIMEX setup: patients, range, lambda grid, Monte Carlo size, RNG. + """ + + def __init__( # pylint: disable=too-many-arguments,invalid-name + self, + patients: list[Patient], + parameter_min: float = -np.inf, + parameter_max: float = np.inf, + *, + lambda_grid: tuple[float, ...] | list[float] | None = None, + B: int = DEFAULT_N_SIMULATIONS, + rng: np.random.Generator | int | None = None, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + ): + YiCorrectionBase.__init__( + self, + patients, + parameter_min=parameter_min, + parameter_max=parameter_max, + ) + self._init_simex( + lambda_grid=lambda_grid, + B=B, + rng=rng, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + + def observed_in_range(self, patient: Patient) -> bool: + """ + Hard observed membership in [parameter_min, parameter_max). + """ + return self._parameter_min <= patient.observed_parameter < self._parameter_max + + def flip_rate_for_range( # pylint: disable=too-many-arguments + self, + patient: Patient, + range_min: float, + range_max: float, + *, + prior_alpha: float, + prior_beta: float, + ) -> float: + """ + e_i = 1 - P(G = G*_i | data) for a single binary range membership. + """ + observed = range_min <= patient.observed_parameter < range_max + probability_observed = self.compute_patient_prob_in_range( + patient, + range_min, + range_max, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + if observed: + return 1.0 - probability_observed + return probability_observed + + +class McSimexWithThreshold(YiCorrectionWithThreshold, _SimexMonteCarlo): + """ + MC-SIMEX for two-group (low / high) hard labels. + """ + + def __init__( # pylint: disable=too-many-arguments,invalid-name + self, + patients: list[Patient], + parameter_min: float = -np.inf, + parameter_max: float = np.inf, + parameter_threshold: float = 0.0, + *, + lambda_grid: tuple[float, ...] | list[float] | None = None, + B: int = DEFAULT_N_SIMULATIONS, + rng: np.random.Generator | int | None = None, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + ): + YiCorrectionWithThreshold.__init__( + self, + patients, + parameter_min=parameter_min, + parameter_max=parameter_max, + parameter_threshold=parameter_threshold, + ) + self._init_simex( + lambda_grid=lambda_grid, + B=B, + rng=rng, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + + def analysis_mask(self) -> np.ndarray: + """ + Patients whose observed parameter lies in [parameter_min, parameter_max). + """ + return np.array( + [ + self._parameter_min <= patient.observed_parameter < self._parameter_max + for patient in self._patients + ], + dtype=bool, + ) + + def observed_is_high(self) -> np.ndarray: + """ + G*_i = 1 iff threshold <= observed_parameter < parameter_max. + """ + return np.array( + [ + self._parameter_threshold <= patient.observed_parameter < self._parameter_max + for patient in self._patients + ], + dtype=bool, + ) + + def flip_rates_high_low( + self, + *, + prior_alpha: float, + prior_beta: float, + ) -> np.ndarray: + """ + Per-patient flip rates for the observed high/low label. + """ + rates = [] + for patient, is_high in zip(self._patients, self.observed_is_high()): + if is_high: + probability_correct = self.compute_patient_prob_in_range( + patient, + self._parameter_threshold, + self._parameter_max, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + else: + probability_correct = self.compute_patient_prob_in_range( + patient, + self._parameter_min, + self._parameter_threshold, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + rates.append(1.0 - probability_correct) + return np.asarray(rates, dtype=float) + + def _survival_arrays(self, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + times = np.array( + [patient.time for patient, keep in zip(self._patients, mask) if keep], + dtype=float, + ) + censored = np.array( + [patient.censored for patient, keep in zip(self._patients, mask) if keep], + dtype=bool, + ) + return times, censored + + +class McSimexForKaplanMeier(McSimexBase): + """ + MC-SIMEX point estimate of a Kaplan-Meier curve (no confidence bands). + """ + + def estimate_survival( # pylint: disable=too-many-locals + self, + times_for_plot: list[float] | None = None, + *, + prior_alpha: float | None = None, + prior_beta: float | None = None, + ) -> dict: + """ + Extrapolate the naive KM of the observed-range group to lambda = -1. + """ + if prior_alpha is None: + prior_alpha = self._prior_alpha + if prior_beta is None: + prior_beta = self._prior_beta + observed_in_range = np.array( + [self.observed_in_range(patient) for patient in self._patients], + dtype=bool, + ) + flip_rates = np.array( + [ + self.flip_rate_for_range( + patient, + self._parameter_min, + self._parameter_max, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + for patient in self._patients + ], + dtype=float, + ) + + group_patients = [ + patient + for patient, in_range in zip(self._patients, observed_in_range) + if in_range + ] + death_times = sorted({ + patient.time for patient in group_patients if not patient.censored + }) + if times_for_plot is None: + if not death_times: + raise ValueError("No death events found in the observed parameter range.") + last_time = max( + [*death_times, *(patient.time for patient in group_patients)] + ) + times_for_plot = [0.0] + death_times + [1.1 * last_time] + + lambda_means = [] + for simex_lambda in self._lambda_grid: + replicates = [] + n_draws = 1 if simex_lambda == 0.0 else self._B + for _ in range(n_draws): + labels = self.simulate_labels( + observed_in_range, flip_rates, simex_lambda + ) + survival = _naive_km_survival( + self._patients, labels, times_for_plot + ) + if survival is not None: + replicates.append(survival) + if not replicates: + lambda_means.append(np.full(len(times_for_plot), np.nan)) + else: + lambda_means.append(np.mean(np.vstack(replicates), axis=0)) + + stacked = np.vstack(lambda_means) + extrapolated = np.array([ + extrapolate_quadratic(np.asarray(self._lambda_grid), stacked[:, i]) + for i in range(stacked.shape[1]) + ]) + extrapolated = np.clip(extrapolated, 0.0, 1.0) + for i, time in enumerate(times_for_plot): + if time == 0.0: + extrapolated[i] = 1.0 + + return { + 'survival_probabilities': extrapolated, + 'times_for_plot': list(times_for_plot), + 'death_times': death_times, + 'method': 'mc_simex', + 'parameter_min': self._parameter_min, + 'parameter_max': self._parameter_max, + 'B': self._B, + 'lambda_grid': self._lambda_grid, + } + + +def _naive_km_survival( + patients: list[Patient], + in_group: np.ndarray, + times_for_plot: list[float], +) -> np.ndarray | None: + if not np.any(in_group): + return None + km_patients = [] + for patient, flag in zip(patients, in_group): + if patient.censored is None: + raise ValueError("Censored status not set") + km_patients.append( + KaplanMeierPatient( + time=patient.time, + censored=patient.censored, + parameter=1.0 if flag else 0.0, + ) + ) + instance = KaplanMeierInstance( + km_patients, + parameter_min=0.5, + parameter_max=np.inf, + ) + if not instance.patients: + return None + return instance.survival_probabilities(times_for_plot=times_for_plot) + + +class McSimexForLogrank(McSimexWithThreshold): + """ + MC-SIMEX for the two-sample logrank statistic, then a chi-square p-value. + """ + + def estimate_pvalue( # pylint: disable=too-many-locals + self, + *, + prior_alpha: float | None = None, + prior_beta: float | None = None, + ) -> dict: + """ + Extrapolate U^2/V to lambda = -1 and convert to a chi-square(1) p-value. + """ + if prior_alpha is None: + prior_alpha = self._prior_alpha + if prior_beta is None: + prior_beta = self._prior_beta + mask = self.analysis_mask() + if not np.any(mask): + raise ValueError("No patients in [parameter_min, parameter_max).") + + observed_high = self.observed_is_high()[mask] + flip_rates = self.flip_rates_high_low( + prior_alpha=prior_alpha, + prior_beta=prior_beta, + )[mask] + times, censored = self._survival_arrays(mask) + + if not np.any(~censored): + raise ValueError("No death events found in patient data.") + + n_low_observed = int(np.count_nonzero(~observed_high)) + n_high_observed = int(np.count_nonzero(observed_high)) + + lambda_means = [] + for simex_lambda in self._lambda_grid: + replicates = [] + n_draws = 1 if simex_lambda == 0.0 else self._B + for _ in range(n_draws): + labels = self.simulate_labels(observed_high, flip_rates, simex_lambda) + statistic = _naive_logrank_statistic(times, censored, labels) + if statistic is not None: + replicates.append(statistic) + if not replicates: + lambda_means.append(np.nan) + else: + lambda_means.append(float(np.mean(replicates))) + + extrapolated = extrapolate_quadratic( + np.asarray(self._lambda_grid), + np.asarray(lambda_means, dtype=float), + ) + logrank_statistic = max(extrapolated, 0.0) + p_value = float(scipy.stats.chi2.sf(logrank_statistic, df=1)) + + return { + 'p_value': p_value, + 'logrank_statistic': logrank_statistic, + 'n_low_observed': n_low_observed, + 'n_high_observed': n_high_observed, + 'method': 'mc_simex', + 'B': self._B, + 'lambda_grid': self._lambda_grid, + } + + +class McSimexForCoxPH(McSimexWithThreshold): + """ + MC-SIMEX for a Cox log hazard ratio, with a Wald interval and Wald Δ2NLL. + """ + + def __init__(self, *args, **kwargs): + """ + Same arguments as McSimexWithThreshold, plus a cache for the Wald fit. + """ + super().__init__(*args, **kwargs) + self._cached_estimate: dict | None = None + self._cached_priors: tuple[float, float] | None = None + + def estimate_hazard_ratio( # pylint: disable=too-many-locals,too-many-branches,too-many-statements + self, + *, + prior_alpha: float | None = None, + prior_beta: float | None = None, + ) -> dict: + """ + Extrapolate log H and its sampling variance to lambda = -1. + + Variance-versus-lambda uses the Breslow observed information of each + successful replicate (Cook-Stefanski / R simex style), then a + quadratic extrapolation. The interval is Wald, not a likelihood-ratio + interval. + """ + if prior_alpha is None: + prior_alpha = self._prior_alpha + if prior_beta is None: + prior_beta = self._prior_beta + priors = (prior_alpha, prior_beta) + if self._cached_estimate is not None and self._cached_priors == priors: + return self._cached_estimate + + mask = self.analysis_mask() + if not np.any(mask): + raise ValueError("No patients in [parameter_min, parameter_max).") + + observed_high = self.observed_is_high()[mask] + flip_rates = self.flip_rates_high_low( + prior_alpha=prior_alpha, + prior_beta=prior_beta, + )[mask] + times, censored = self._survival_arrays(mask) + + lambda_log_hr = [] + lambda_variance = [] + naive_log_hr = None + for simex_lambda in self._lambda_grid: + log_hrs = [] + variances = [] + n_draws = 1 if simex_lambda == 0.0 else self._B + for _ in range(n_draws): + labels = self.simulate_labels(observed_high, flip_rates, simex_lambda) + fit = mle_log_hazard_ratio(times, censored, labels) + if fit is None: + continue + log_hrs.append(fit[0]) + if np.isfinite(fit[1]): + variances.append(fit[1]) + if not log_hrs: + lambda_log_hr.append(np.nan) + lambda_variance.append(np.nan) + continue + mean_log_hr = float(np.mean(log_hrs)) + lambda_log_hr.append(mean_log_hr) + if simex_lambda == 0.0: + naive_log_hr = mean_log_hr + if variances: + lambda_variance.append(float(np.mean(variances))) + else: + lambda_variance.append(np.nan) + + log_hazard_ratio = extrapolate_quadratic( + np.asarray(self._lambda_grid), + np.asarray(lambda_log_hr, dtype=float), + ) + variance = extrapolate_quadratic( + np.asarray(self._lambda_grid), + np.asarray(lambda_variance, dtype=float), + ) + finite_variances = [ + value for value in lambda_variance if np.isfinite(value) + ] + if variance <= 0.0: + if finite_variances: + variance = max(min(finite_variances), VARIANCE_FLOOR) + else: + variance = VARIANCE_FLOOR + se_log_hazard_ratio = float(np.sqrt(variance)) + hazard_ratio = float(np.exp(log_hazard_ratio)) + ci_lower = float(np.exp(log_hazard_ratio - WALD_Z_95 * se_log_hazard_ratio)) + ci_upper = float(np.exp(log_hazard_ratio + WALD_Z_95 * se_log_hazard_ratio)) + naive_hazard_ratio = ( + None if naive_log_hr is None else float(np.exp(naive_log_hr)) + ) + + estimate = { + 'hazard_ratio': hazard_ratio, + 'log_hazard_ratio': log_hazard_ratio, + 'se_log_hazard_ratio': se_log_hazard_ratio, + 'ci_lower': ci_lower, + 'ci_upper': ci_upper, + 'confidence_level': 0.95, + 'naive_hazard_ratio': naive_hazard_ratio, + 'method': 'mc_simex', + 'B': self._B, + 'lambda_grid': self._lambda_grid, + } + self._cached_estimate = estimate + self._cached_priors = priors + return estimate + + def compute_2nll_at_hazard_ratio( + self, + hazard_ratio: float, + *, + prior_alpha: float | None = None, + prior_beta: float | None = None, + ) -> scipy.optimize.OptimizeResult: + """ + Wald quadratic (log H - log H_hat)^2 / sigma_hat^2. + + This is not a profile likelihood and must not be labeled as one. + """ + if hazard_ratio <= 0.0: + raise ValueError(f"hazard_ratio must be positive, got {hazard_ratio}") + estimate = self.estimate_hazard_ratio( + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + delta = ( + (np.log(hazard_ratio) - estimate['log_hazard_ratio']) ** 2 + / estimate['se_log_hazard_ratio'] ** 2 + ) + return scipy.optimize.OptimizeResult( + x=float(delta), + success=True, + hazard_ratio=hazard_ratio, + log_hazard_ratio=float(np.log(hazard_ratio)), + fitted_hazard_ratio=estimate['hazard_ratio'], + se_log_hazard_ratio=estimate['se_log_hazard_ratio'], + ci_lower=estimate['ci_lower'], + ci_upper=estimate['ci_upper'], + cox_2NLL=float(delta), + patient_2NLL=0.0, + method='mc_simex_wald', + ) diff --git a/kombine/yi_correction.py b/kombine/comparisons/yi_correction.py similarity index 88% rename from kombine/yi_correction.py rename to kombine/comparisons/yi_correction.py index 50f5378..5e7941d 100644 --- a/kombine/yi_correction.py +++ b/kombine/comparisons/yi_correction.py @@ -35,8 +35,10 @@ import scipy.optimize import scipy.stats +from ..utilities import brentq_hazard_ratio_ci + if TYPE_CHECKING: - from .datacard import Patient + from ..datacard import Patient class YiCorrectionBase: # pylint: disable=too-few-public-methods @@ -469,6 +471,79 @@ def compute_2nll_at_hazard_ratio( # pylint: disable=too-many-locals return result + def hazard_ratio_confidence_interval( # pylint: disable=too-many-arguments,too-many-locals + self, + *, + confidence_level: float = 0.68, + hazard_ratio_min: float = 0.01, + hazard_ratio_max: float = 100.0, + tolerance: float = 1e-3, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + ) -> tuple[float, float, float, scipy.optimize.OptimizeResult]: + """ + Continuous MLE and likelihood-ratio interval for the Yi-weighted Breslow Cox HR. + + Minimizes 2NLL in log H, then finds where 2NLL crosses + min + chi2.ppf(confidence_level, 1). + """ + if hazard_ratio_min <= 0.0 or hazard_ratio_max <= hazard_ratio_min: + raise ValueError( + "Need 0 < hazard_ratio_min < hazard_ratio_max, got " + f"{hazard_ratio_min}, {hazard_ratio_max}." + ) + if not 0.0 < confidence_level < 1.0: + raise ValueError( + f"confidence_level must be in (0, 1), got {confidence_level}." + ) + + def objective(log_hr: float) -> float: + result = self.compute_2nll_at_hazard_ratio( + float(np.exp(log_hr)), + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + return float(result.x) + + result_opt = scipy.optimize.minimize_scalar( + objective, + bounds=(np.log(hazard_ratio_min), np.log(hazard_ratio_max)), + method='bounded', + options={'xatol': tolerance * 0.1}, + ) + if not result_opt.success: + raise ValueError("Could not find best-fit hazard ratio") + + best_fit_log_hr = float(result_opt.x) # type: ignore[arg-type] + best_fit_hr = float(np.exp(best_fit_log_hr)) + min_twonll = float(result_opt.fun) # type: ignore[arg-type] + best_fit_result = self.compute_2nll_at_hazard_ratio( + best_fit_hr, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + twonll_threshold = min_twonll + float( + scipy.stats.chi2.ppf(confidence_level, df=1) + ) + + def twonll_minus_threshold(log_hr: float) -> float: + result = self.compute_2nll_at_hazard_ratio( + float(np.exp(log_hr)), + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + return float(result.x) - twonll_threshold + + lower_ci, upper_ci = brentq_hazard_ratio_ci( + twonll_minus_threshold, + best_fit_log_hr=best_fit_log_hr, + hazard_ratio_min=hazard_ratio_min, + hazard_ratio_max=hazard_ratio_max, + tolerance=tolerance, + ) + + return best_fit_hr, lower_ci, upper_ci, best_fit_result + class YiCorrectionForKaplanMeier(YiCorrectionBase): """ diff --git a/kombine/datacard.py b/kombine/datacard.py index 492b9af..93617c0 100644 --- a/kombine/datacard.py +++ b/kombine/datacard.py @@ -26,7 +26,14 @@ ) from .kaplan_meier_p_value_MINLP import MINLPforKMPValue from .kaplan_meier_hazard_ratio_MINLP import MINLPforKMHazardRatio -from .yi_correction import YiCorrectionForLogrank, YiCorrectionForCoxPH, YiCorrectionForKaplanMeier +from .comparisons import ( + McSimexForCoxPH, + McSimexForKaplanMeier, + McSimexForLogrank, + YiCorrectionForCoxPH, + YiCorrectionForKaplanMeier, + YiCorrectionForLogrank, +) from .utilities import LOG_ZERO_EPSILON_DEFAULT, prob_poisson_density_in_range, validate_class_probs def _parse_prob_class_label(label: str) -> int | None: @@ -931,7 +938,7 @@ def observed_parameter(self) -> float: raise ValueError("Observable not set") return self.observable.observed_parameter() -class Datacard: +class Datacard: # pylint: disable=too-many-public-methods """ A datacard class to specify the inputs to ROC Picker. Refer to docs/03_examples.md for usage examples. @@ -1714,6 +1721,101 @@ def km_survival_yi( # pylint: disable=too-many-arguments prior_beta=prior_beta, ) + def km_survival_mc_simex( # pylint: disable=too-many-arguments,invalid-name + self, + *, + parameter_min: float = -np.inf, + parameter_max: float = np.inf, + times_for_plot: list[float] | None = None, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + lambda_grid: tuple[float, ...] | list[float] | None = None, + B: int = 100, + rng: np.random.Generator | int | None = None, + ) -> dict: + """ + Kaplan-Meier point estimate using Küchenhoff MC-SIMEX on hard labels. + + Extra flips are simulated at lambda >= 0 and the naive KM curve is + extrapolated quadratically to lambda = -1. There are no KM bands. + """ + calculator = McSimexForKaplanMeier( + patients=self.patients, + parameter_min=parameter_min, + parameter_max=parameter_max, + lambda_grid=lambda_grid, + B=B, + rng=rng, + ) + return calculator.estimate_survival( + times_for_plot=times_for_plot, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + + def km_p_value_logrank_mc_simex( # pylint: disable=too-many-arguments,invalid-name + self, + *, + parameter_threshold: float, + parameter_min: float = -np.inf, + parameter_max: float = np.inf, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + lambda_grid: tuple[float, ...] | list[float] | None = None, + B: int = 100, + rng: np.random.Generator | int | None = None, + ) -> dict: + """ + Logrank p-value using Küchenhoff MC-SIMEX on hard high/low labels. + + The naive logrank statistic is averaged at each lambda and extrapolated + to lambda = -1. The p-value is from a chi-square_1 tail of that number. + """ + calculator = McSimexForLogrank( + patients=self.patients, + parameter_min=parameter_min, + parameter_max=parameter_max, + parameter_threshold=parameter_threshold, + lambda_grid=lambda_grid, + B=B, + rng=rng, + ) + return calculator.estimate_pvalue( + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + + def km_hazard_ratio_mc_simex( # pylint: disable=too-many-arguments,invalid-name + self, + *, + parameter_threshold: float, + parameter_min: float = -np.inf, + parameter_max: float = np.inf, + prior_alpha: float = 0.5, + prior_beta: float = 0.0, + lambda_grid: tuple[float, ...] | list[float] | None = None, + B: int = 100, + rng: np.random.Generator | int | None = None, + ) -> McSimexForCoxPH: + """ + MC-SIMEX Cox calculator: extrapolated log H, Wald CI, and Wald Δ2NLL. + + Call ``estimate_hazard_ratio`` for the point and Wald interval, and + ``compute_2nll_at_hazard_ratio`` to scan the Wald quadratic on the same + H grid used for Yi and KoMbine. + """ + return McSimexForCoxPH( + patients=self.patients, + parameter_min=parameter_min, + parameter_max=parameter_max, + parameter_threshold=parameter_threshold, + lambda_grid=lambda_grid, + B=B, + rng=rng, + prior_alpha=prior_alpha, + prior_beta=prior_beta, + ) + def clear_distributions(self): """ Delete the distributions for all patients. diff --git a/kombine/discrete_optimization.py b/kombine/discrete_optimization.py index 01e28c2..a14bd92 100644 --- a/kombine/discrete_optimization.py +++ b/kombine/discrete_optimization.py @@ -4,6 +4,8 @@ """ import collections.abc +import dataclasses +import time import typing import numpy as np @@ -526,3 +528,230 @@ def remaining_width(index: int) -> float: raise RuntimeError("cached_level_crossings left a level unsolved") crossings.append(x) return crossings + + +SignOracleStatus = typing.Literal["inside", "outside", "unknown"] +SignOracleResult = SignOracleStatus | tuple[SignOracleStatus, float] + + +@dataclasses.dataclass +class OracleCostTracker: # pylint: disable=too-many-instance-attributes + """Online estimate of cost-biased bisection probe fraction ``r``. + + Prior ``r₀=0.5`` (midpoint). Updates from outcome counts and unitless + wall-time ratios; clamps with uncertainty so outliers cannot poison ``r``. + """ + + r_prior: float = 0.5 + kappa: float = 4.0 + z: float = 1.0 + epsilon_floor: float = 0.05 + winsor_cap: float = 4.0 + ema_alpha: float = 0.3 + w_threshold: float = 2.0 + alpha: float = dataclasses.field(init=False) + beta: float = dataclasses.field(init=False) + n_in: int = 0 + n_out: int = 0 + u_in: float = 1.0 + u_out: float = 1.0 + _times: list[float] = dataclasses.field(default_factory=list) + + def __post_init__(self) -> None: + if not 0.0 < self.r_prior < 1.0: + raise ValueError(f"r_prior must be in (0, 1), got {self.r_prior}") + self.alpha = self.kappa * self.r_prior + self.beta = self.kappa * (1.0 - self.r_prior) + + def record(self, outcome: typing.Literal["inside", "outside"], wall_time: float) -> None: + """Record one oracle outcome and its wall time (seconds).""" + t = float(wall_time) + if self._times: + median_all = float(np.median(self._times)) + if median_all > 0.0: + t = min(t, self.winsor_cap * median_all) + self._times.append(t) + t_ref = float(np.median(self._times)) + if t_ref <= 0.0: + t_ref = 1.0 + norm_t = t / t_ref + + if outcome == "inside": + self.n_in += 1 + self.alpha += 1.0 + self.u_in = (1.0 - self.ema_alpha) * self.u_in + self.ema_alpha * norm_t + else: + self.n_out += 1 + self.beta += 1.0 + self.u_out = (1.0 - self.ema_alpha) * self.u_out + self.ema_alpha * norm_t + + def probe_fraction(self) -> float: + """Fraction from ``inside`` toward ``outside`` for the next probe.""" + r_count = self.alpha / (self.alpha + self.beta) + if self.n_in >= 1 and self.n_out >= 1: + r_time = self.u_in / (self.u_in + self.u_out) + w = min(1.0, min(self.n_in, self.n_out) / self.w_threshold) + r_raw = w * r_time + (1.0 - w) * r_count + else: + r_raw = r_count + + n_eff = self.n_in + self.n_out + denom = self.alpha + self.beta + n_eff + se_r = float(np.sqrt(max(r_raw * (1.0 - r_raw) / denom, 0.0))) + epsilon = min(0.5, self.z * se_r + self.epsilon_floor) + return float(np.clip(r_raw, epsilon, 1.0 - epsilon)) + + +def feasibility_assisted_level_crossings( # pylint: disable=too-many-locals, too-many-arguments, too-many-positional-arguments, too-many-branches, too-many-statements + value_func: collections.abc.Callable[[float], float], + sign_oracle: collections.abc.Callable[[float, float], SignOracleResult], + x_outer: float, + x_inner: float, + levels: collections.abc.Sequence[float], + *, + xtol: float = 1e-4, + rtol: float = 1e-4, + cost_biased_bisection: bool = True, + r_prior: float = 0.5, + kappa: float = 4.0, + outer_oracle: list[SignOracleStatus | None] | None = None, +) -> list[float]: + """Locate profile level crossings using a cheap inside/outside oracle. + + ``value_func(p)`` returns nonnegative excess (e.g. ``2NLL(p)-2NLL_min``). + ``sign_oracle(p, level)`` returns whether excess can be ``<= level`` + (``inside``), is proven ``> level`` (``outside``), or ``unknown``. + It may also return ``(status, cost)`` where ``cost`` is a sleep-immune + effort measure (prefer Gurobi ``Work`` over wall-clock). + + Coarse bracketing uses the oracle (falling back to ``value_func`` on + ``unknown``); the final root is polished with ``brentq`` on ``value_func``. + + When ``cost_biased_bisection`` is True, probe points use ``OracleCostTracker`` + (prior ``r_prior=0.5`` = midpoint until data arrives). + """ + if not levels: + return [] + if x_outer == x_inner: + return [x_inner] * len(levels) + + value_cache: dict[float, float] = {} + + def f_cached(x: float) -> float: + if x not in value_cache: + value_cache[x] = float(value_func(x)) + return value_cache[x] + + def _call_oracle( + x: float, level: float, + ) -> tuple[SignOracleStatus, float | None]: + raw = sign_oracle(x, level) + if isinstance(raw, tuple): + status, cost = raw + return status, float(cost) + return raw, None + + def classify(x: float, level: float) -> typing.Literal["inside", "outside"]: + status, _cost = _call_oracle(x, level) + if status == "unknown": + return "inside" if f_cached(x) <= level else "outside" + return status + + def classify_timed( + x: float, + level: float, + tracker: OracleCostTracker | None, + ) -> typing.Literal["inside", "outside"]: + # Prefer oracle-reported cost (Gurobi Work). Fall back to process CPU + # time — not wall clock — so OS sleep cannot poison the tracker. + t0 = time.process_time() + status, cost = _call_oracle(x, level) + if cost is None: + cost = time.process_time() - t0 + if status == "unknown": + return "inside" if f_cached(x) <= level else "outside" + if tracker is not None: + tracker.record(status, cost) + return status + + crossings: list[float] = [] + for level in levels: + tracker = ( + OracleCostTracker(r_prior=r_prior, kappa=kappa) + if cost_biased_bisection + else None + ) + raw_outer, outer_cost = _call_oracle(x_outer, level) + if outer_oracle is not None: + outer_oracle.append(None if raw_outer == "unknown" else raw_outer) + if raw_outer == "unknown": + outer_side = "inside" if f_cached(x_outer) <= level else "outside" + else: + outer_side = raw_outer + if tracker is not None: + cost = outer_cost if outer_cost is not None else 0.0 + tracker.record(outer_side, cost) + if outer_side == "inside": + crossings.append(x_outer) + continue + inner_side = ( + classify_timed(x_inner, level, tracker) + if tracker is not None + else classify(x_inner, level) + ) + if inner_side == "outside": + crossings.append(x_inner) + continue + + inside = x_inner + outside = x_outer + while not _is_close(inside, outside, xtol, rtol): + if tracker is not None: + r = tracker.probe_fraction() + mid = inside + r * (outside - inside) + else: + mid = 0.5 * (inside + outside) + if _is_close(mid, inside, xtol, rtol) or _is_close(mid, outside, xtol, rtol): + break + side = ( + classify_timed(mid, level, tracker) + if tracker is not None + else classify(mid, level) + ) + if side == "inside": + inside = mid + else: + outside = mid + + left, right = (inside, outside) if inside < outside else (outside, inside) + if left == right or _is_close(left, right, xtol, rtol): + closer = left if abs(f_cached(left) - level) <= abs(f_cached(right) - level) else right + crossings.append(closer) + continue + g_left = f_cached(left) - level + g_right = f_cached(right) - level + if g_left == 0.0: + crossings.append(left) + continue + if g_right == 0.0: + crossings.append(right) + continue + if g_left * g_right > 0.0: + # Oracle bracket disagreed with numeric signs; pick closer endpoint. + closer = left if abs(g_left) <= abs(g_right) else right + crossings.append(closer) + continue + crossings.append( + typing.cast( + float, + scipy.optimize.brentq( + lambda x, _level=level: f_cached(x) - _level, + left, + right, + xtol=xtol, + rtol=np.float64(rtol), + full_output=False, + ), + ) + ) + return crossings diff --git a/kombine/kaplan_meier_MINLP.py b/kombine/kaplan_meier_MINLP.py index ef66c87..99cdeb1 100644 --- a/kombine/kaplan_meier_MINLP.py +++ b/kombine/kaplan_meier_MINLP.py @@ -4,11 +4,13 @@ """ import collections.abc +import dataclasses import datetime import functools import itertools import math import os +import typing import gurobipy as gp from gurobipy import GRB @@ -101,6 +103,48 @@ def observed_parameter(self) -> float: """ return self.__observed_parameter + def min_nll_on_interval( + self, + range_min: float, + range_max: float, + ) -> float: + """ + Minimum parameter NLL for a value in [range_min, range_max). + + Endpoints alone are not enough when the NLL is piecewise-constant on + integer class bins [k, k+1): the class boundary is a jump, not a + continuous fence. Candidates are the observed value (if inside), the + inclusive left / exclusive right edges, and class-bin midpoints. + """ + if not range_min < range_max: + return float('inf') + candidates: list[float] = [] + observed = self.observed_parameter + if range_min <= observed < range_max: + candidates.append(float(observed)) + if np.isfinite(range_min): + candidates.append(float(range_min)) + if np.isfinite(range_max): + just_inside = float(range_max) - LOG_ZERO_EPSILON_DEFAULT + if range_min <= just_inside < range_max: + candidates.append(just_inside) + if np.isfinite(range_min): + k_start = max(0, int(math.floor(range_min))) + else: + k_start = 0 + if np.isfinite(range_max): + k_stop = int(math.floor(range_max - 1e-15)) + 1 + else: + k_stop = k_start + 64 + k_stop = min(k_stop, k_start + 64) + for class_index in range(k_start, max(k_start, k_stop)): + midpoint = class_index + 0.5 + if range_min <= midpoint < range_max: + candidates.append(midpoint) + if not candidates: + return float('inf') + return float(min(self.parameter(value) for value in candidates)) + @staticmethod def _solve_0d( full_nll: collections.abc.Callable[[float], float], @@ -501,6 +545,31 @@ def nominal(self) -> KaplanMeierPatient: parameter=self.observed_parameter, ) + +@dataclasses.dataclass +class GurobiWorkStats: + """Accumulated Gurobi Work units for oracle vs full-minimize calls.""" + + oracle_work: float = 0.0 + oracle_calls: int = 0 + oracle_outside_calls: int = 0 + minimize_work: float = 0.0 + minimize_calls: int = 0 + + @property + def total_work(self) -> float: + """Oracle plus minimize Work.""" + return self.oracle_work + self.minimize_work + + def reset(self) -> None: + """Clear accumulated Work counters.""" + self.oracle_work = 0.0 + self.oracle_calls = 0 + self.oracle_outside_calls = 0 + self.minimize_work = 0.0 + self.minimize_calls = 0 + + class MINLPForKM(GurobiOptimizerMixin): # pylint: disable=too-many-public-methods, too-many-instance-attributes """ Mixed Integer Nonlinear Programming for a point on the Kaplan-Meier curve. @@ -543,6 +612,16 @@ def __init__( # pylint: disable=too-many-arguments self.__mip_start_a: dict[int, float] | None = None self.__mip_start_mode: tuple[bool, bool] | None = None self.__mip_start_profile: dict[str, list[float]] | None = None + self.__binom_penalty_var: gp.Var | None = None + self.__binom_piece_vars: list[gp.Var] | None = None + self.__patient_penalty_var: gp.Var | None = None + self.__feasibility_cut_constraints: list = [] + # Sleep-immune cost proxy from the last excess_at_most call (Gurobi Work). + self.last_oracle_work: float | None = None + # Sum of Work across all optimize attempts in the last run_MINLP call. + self.last_minimize_work: float | None = None + self.work_stats = GurobiWorkStats() + self._last_minimize_call_work: float = 0.0 if not np.isfinite(self.__parameter_min and self.__parameter_min != -np.inf): raise ValueError("parameter_min must be finite or -inf") if not np.isfinite(self.__parameter_max and self.__parameter_max != np.inf): @@ -564,6 +643,90 @@ def seed_assignment_starts(self, starts: dict[int, float] | list[int]) -> None: # Mode unknown when seeded externally; allow first solve to use them. self.__mip_start_mode = None + @staticmethod + def _work_from_model(model) -> float | None: + try: + return float(model.Work) + except (AttributeError, gp.GurobiError): + try: + return float(model.Runtime) + except (AttributeError, gp.GurobiError): + return None + + def _record_work( + self, + model, + *, + kind: typing.Literal["oracle", "minimize"], + side: str | None = None, + ) -> float: + work = self._work_from_model(model) or 0.0 + if kind == "oracle": + self.work_stats.oracle_work += work + self.work_stats.oracle_calls += 1 + if side == "outside": + self.work_stats.oracle_outside_calls += 1 + else: + self.work_stats.minimize_work += work + self.work_stats.minimize_calls += 1 + return work + + def _log_run_minlp_completion( + self, + *, + expected_probability: float | None, + status: int, + print_progress: bool, + verbose: bool, + ) -> None: + if not (print_progress or verbose): + return + work = self.last_minimize_work + work_str = f"{work:.3f}" if work is not None else "?" + outcome = "success" if status == GRB.OPTIMAL else f"status={status}" + print( + f"[{datetime.datetime.now()}] run_MINLP -> {outcome} " + f"(status={status}, Work={work_str}, p={expected_probability})", + flush=True, + ) + + def _optimize_with_fallbacks( + self, + model, + initial_params: dict, + fallback_strategies: list[tuple[dict, str]], + verbose: bool, + ): + self._set_gurobi_params(model, initial_params) + call_work = 0.0 + + if verbose: + print("Attempting initial optimization...") + model.optimize() + call_work += self._record_work(model, kind="minimize") + + needs_fallback = model.status in (GRB.SUBOPTIMAL, GRB.TIME_LIMIT) + if needs_fallback: + for i, (fallback_params, description) in enumerate(fallback_strategies): + if verbose: + print( + f"Model returned status {model.status}. " + f"Applying fallback {i + 1}: {description}" + ) + print(f" New parameters: {fallback_params}") + self._set_gurobi_params(model, fallback_params) + model.optimize() + fallback_work = self._record_work(model, kind="minimize") + call_work += fallback_work + if verbose: + print(f" fallback {i + 1} Work={fallback_work:.3f}: {description}") + if model.status == GRB.OPTIMAL: + if verbose: + print(f"Fallback {i + 1} successful. Model is now optimal.") + break + self._last_minimize_call_work = call_work + return model + def _apply_assignment_mip_starts( self, a, @@ -1045,43 +1208,21 @@ def possible_probabilities(self) -> set[float]: @functools.cached_property def nll_penalty_for_patient_in_range(self) -> npt.NDArray[np.float64]: """ - Calculate the negative log-likelihood penalty for each patient - if that patient is within the parameter range. - This is negative if the patient's observed parameter is within the range - and positive if it is outside the range. - """ - sgn_nll_penalty_for_patient_in_range = 2 * self.parameter_in_range - 1 - observed_nll = np.array([ - p.parameter(p.observed_parameter) - for p in self.all_patients - ]) - if np.isfinite(self.parameter_min): - parameter_min_nll = np.array([ - p.parameter(self.parameter_min) - for p in self.all_patients - ]) - else: - parameter_min_nll = np.full(self.n_patients, np.inf) - if np.isfinite(self.parameter_max): - parameter_max_nll = np.array([ - p.parameter(self.parameter_max) - for p in self.all_patients - ]) - else: - parameter_max_nll = np.full(self.n_patients, np.inf) - - range_boundary_nll = np.min( - np.array([parameter_min_nll, parameter_max_nll]), - axis=0 - ) - abs_nll_penalty_for_patient_in_range = observed_nll - range_boundary_nll - - nll_penalty_for_patient_in_range = ( - sgn_nll_penalty_for_patient_in_range - * abs_nll_penalty_for_patient_in_range - ) + Calculate the NLL penalty for assigning each patient to this range. - return nll_penalty_for_patient_in_range + The value is min NLL on [parameter_min, parameter_max) minus min NLL + on the complement. Negative means the range is the better assignment. + """ + return np.array([ + ( + patient.min_nll_on_interval(self.parameter_min, self.parameter_max) + - min( + patient.min_nll_on_interval(-np.inf, self.parameter_min), + patient.min_nll_on_interval(self.parameter_max, np.inf), + ) + ) + for patient in self.all_patients + ], dtype=float) @functools.cached_property def n_choose_d_term_table(self) -> dict[tuple[int, int], float]: @@ -1355,6 +1496,31 @@ def add_kaplan_meier_probability_variables_and_constraints( return km_probability_var + def _attach_binom_piece_vars( + self, + model: gp.Model, + binomial_terms_by_time: list[list], + ) -> list[gp.Var]: + """Per-death-time auxiliaries for finer feasibility cuts.""" + piece_vars: list[gp.Var] = [] + for i in range(self.n_times_to_consider): + piece = model.addVar( + lb=0.0, + vtype=GRB.CONTINUOUS, + name=f"binom_piece[{i}]", + ) + terms_i = binomial_terms_by_time[i] + if terms_i: + model.addConstr( + piece == gp.quicksum(terms_i), + name=f"binom_piece_{i}_definition", + ) + else: + model.addConstr(piece == 0.0, name=f"binom_piece_{i}_zero") + piece_vars.append(piece) + self.__binom_piece_vars = piece_vars + return piece_vars + def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statements, too-many-arguments, too-many-branches self, model: gp.Model, @@ -1363,7 +1529,7 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement d: gp.tupledict[int, gp.Var], sub_d: gp.tupledict[int, gp.Var], s: gp.tupledict[int, gp.Var], - ): + ) -> tuple[gp.Var, gp.Var, gp.Var | None]: """ Add the binomial penalty to the model. This penalty is based on the expected survival probability @@ -1450,28 +1616,31 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement self.__profile_log_p_survived = log_p_survived n_choose_d_table = self.n_choose_d_term_table - binomial_terms = [] + binomial_terms_by_time: list[list] = [ + [] for _ in range(self.n_times_to_consider) + ] if self.__binomial_only: _ = (r, d, sub_d) r_vals, d_vals, _s_vals, sub_d_vals = self._fixed_selected_counts() sub_d_offset = 0 for i, time in enumerate(self.times_to_consider): + terms_i = binomial_terms_by_time[i] r_value = r_vals[i] d_value = d_vals[i] s_value = r_value - d_value penalty = n_choose_d_table[(r_value, d_value)] - binomial_terms.append(-penalty) - binomial_terms.append(-d_value * log_p_died[i]) - binomial_terms.append(-s_value * log_p_survived[i]) + terms_i.append(-penalty) + terms_i.append(-d_value * log_p_died[i]) + terms_i.append(-s_value * log_p_survived[i]) if d_value > 0: - binomial_terms.append( + terms_i.append( -(math.lgamma(d_value + 1) - d_value * np.log(d_value)) ) n_sub = len(self._collapsed_time_groups[time]) for sub_d_value in sub_d_vals[sub_d_offset:sub_d_offset + n_sub]: if sub_d_value > 0: - binomial_terms.append( + terms_i.append( math.lgamma(sub_d_value + 1) - sub_d_value * np.log(sub_d_value) ) sub_d_offset += n_sub @@ -1481,6 +1650,7 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement # are charged on the same indicator. sub_d_counter = -1 for i, time in enumerate(self.times_to_consider): + terms_i = binomial_terms_by_time[i] feasible_indicators = [] for (r_value, d_value), penalty in n_choose_d_table.items(): if not self._rd_pair_is_feasible(i, r_value, d_value): @@ -1507,11 +1677,11 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement name=f"n_choose_d_indicator_d_{i}_{r_value}_{d_value}", ) s_value = r_value - d_value - binomial_terms.append(-penalty * indicator) - binomial_terms.append(-d_value * log_p_died[i] * indicator) - binomial_terms.append(-s_value * log_p_survived[i] * indicator) + terms_i.append(-penalty * indicator) + terms_i.append(-d_value * log_p_died[i] * indicator) + terms_i.append(-s_value * log_p_survived[i] * indicator) if d_value > 0: - binomial_terms.append( + terms_i.append( -indicator * ( math.lgamma(d_value + 1) - d_value * np.log(d_value) ) @@ -1551,7 +1721,7 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement name=f"sub_d_indicator_constr_{i}_{sub_d_counter}_{sub_d_value}", ) if sub_d_value > 0: - binomial_terms.append( + terms_i.append( sub_d_indicator * ( math.lgamma(sub_d_value + 1) - sub_d_value * np.log(sub_d_value) ) @@ -1561,6 +1731,10 @@ def add_binomial_penalty( # pylint: disable=too-many-locals, too-many-statement name=f"one_sub_d_indicator_per_sub_death_time_{i}_{sub_d_counter}", ) + self._attach_binom_piece_vars(model, binomial_terms_by_time) + binomial_terms = [ + term for terms_i in binomial_terms_by_time for term in terms_i + ] binom_penalty_expr = gp.quicksum(binomial_terms) binom_penalty = model.addVar( vtype=GRB.CONTINUOUS, @@ -1675,7 +1849,7 @@ def _make_gurobi_model(self): #pylint: disable=too-many-locals km_probability_var = None expected_probability_var = None use_binomial_penalty_indicator = None - binom_penalty = 0.0 + binom_penalty: gp.Var | None = None if self.__patient_wise_only: km_probability_var = self.add_kaplan_meier_probability_variables_and_constraints( @@ -1719,10 +1893,30 @@ def _make_gurobi_model(self): #pylint: disable=too-many-locals model=model, a=a, ) + # Aux var so feasibility oracles can cut the shifted patient piece alone. + patient_penalty_var = model.addVar( + lb=0.0, + vtype=GRB.CONTINUOUS, + name="patient_penalty", + ) + model.addConstr( + patient_penalty_var == patient_penalty, + name="patient_penalty_definition", + ) + self.__patient_penalty_var = patient_penalty_var + if self.__patient_wise_only: + self.__binom_penalty_var = None + self.__binom_piece_vars = None + objective_binom = 0.0 + else: + if binom_penalty is None: + raise RuntimeError("binom_penalty var missing; rebuild gurobi_model") + self.__binom_penalty_var = binom_penalty + objective_binom = binom_penalty # Objective: minimize total penalty model.setObjective( - 2 * (binom_penalty + patient_penalty), + 2 * (objective_binom + patient_penalty_var), GRB.MINIMIZE, ) model.update() @@ -1791,6 +1985,179 @@ def update_model_with_expected_probability( # pylint: disable=too-many-arguments model.update() + def _clear_feasibility_cuts(self, model: gp.Model) -> None: + """Remove temporary CL-feasibility cuts from a prior excess_at_most call.""" + for constr in self.__feasibility_cut_constraints: + try: + model.remove(constr) + except gp.GurobiError: + pass + self.__feasibility_cut_constraints = [] + model.update() + + def excess_at_most( # pylint: disable=too-many-arguments, too-many-locals, too-many-branches, too-many-statements + self, + expected_probability: float, + *, + twoNLL_min: float, + level: float, + component_cuts: bool = False, + binom_time_cuts: bool = False, + verbose: bool = False, + print_progress: bool = False, + TimeLimit: float | None = None, + Threads: int | None = None, + LogFile: os.PathLike | None = None, + ) -> typing.Literal["inside", "outside", "unknown"]: + """ + Ask whether profile excess 2NLL(p) - twoNLL_min can be <= ``level``. + + Adds ``2*(binom + patient) <= T`` (and optional redundant piece cuts) with + ``T = twoNLL_min + level``, then runs a feasibility-focused solve without + requiring a tight optimality certificate. + + Optional ``component_cuts`` adds aggregate ``2*binom`` and ``2*patient`` + bounds; ``binom_time_cuts`` adds ``2*binom_piece[i]`` per death time. + + Returns + ------- + "inside" + A feasible point under the cuts was found (excess can be <= level). + "outside" + Gurobi proved the cut model infeasible (excess > level). + "unknown" + Status inconclusive (e.g. time limit with no feasible point); caller + should fall back to a full ``run_MINLP`` / ``twoNLL`` evaluation. + """ + if expected_probability <= 0 or expected_probability >= 1: + if not self.__patient_wise_only: + raise ValueError( + f"expected_probability={expected_probability} must be in (0, 1)" + ) + if expected_probability < 0 or expected_probability > 1: + raise ValueError( + f"expected_probability={expected_probability} must be in [0, 1]" + ) + if level < 0: + raise ValueError(f"level must be >= 0, got {level}") + if not np.isfinite(twoNLL_min): + raise ValueError(f"twoNLL_min must be finite, got {twoNLL_min}") + + threshold = float(twoNLL_min + level) + if print_progress or verbose: + print( + f"[{datetime.datetime.now()}] Feasibility excess_at_most " + f"p={expected_probability} level={level} T={threshold} " + f"at time point {self.time_point}", + flush=True, + ) + + ( + model, + a, + km_probability_var, + expected_probability_var, + _use_binomial_penalty_indicator, + ) = self.gurobi_model + self.update_model_with_expected_probability( + model=model, + km_probability_var=km_probability_var, + expected_probability=expected_probability, + expected_probability_var=expected_probability_var, + ) + self._clear_feasibility_cuts(model) + + patient_var = self.__patient_penalty_var + if patient_var is None: + raise RuntimeError("patient_penalty var missing; rebuild gurobi_model") + binom_var = self.__binom_penalty_var + total_2nll = 2 * ( + (binom_var if binom_var is not None else 0.0) + patient_var + ) + self.__feasibility_cut_constraints.append( + model.addConstr(total_2nll <= threshold, name="feas_total_2nll_cut") + ) + if component_cuts: + if binom_var is not None: + self.__feasibility_cut_constraints.append( + model.addConstr(2 * binom_var <= threshold, name="feas_binom_2nll_cut") + ) + self.__feasibility_cut_constraints.append( + model.addConstr( + 2 * patient_var <= threshold, name="feas_patient_2nll_cut" + ) + ) + if binom_time_cuts: + piece_vars = self.__binom_piece_vars + if piece_vars is None: + raise RuntimeError("binom_piece vars missing; rebuild gurobi_model") + for i, piece_var in enumerate(piece_vars): + self.__feasibility_cut_constraints.append( + model.addConstr( + 2 * piece_var <= threshold, + name=f"feas_binom_piece_{i}_cut", + ) + ) + + self._apply_assignment_mip_starts( + a, + binomial_only=self.__binomial_only, + patient_wise_only=self.__patient_wise_only, + ) + self._apply_profile_mip_starts() + + model.setParam("OutputFlag", 1 if verbose else 0) + model.setParam("DisplayInterval", 1) + model.setParam("NonConvex", 2) + model.setParam("Seed", 123456) + model.setParam("MIPFocus", 1) # feasibility + model.setParam("FuncPieces", 1000) + model.setParam("FuncPieceRatio", 0.5) + # Stop at the first cut-feasible incumbent; proving outside still explores. + model.setParam("SolutionLimit", 1) + if TimeLimit is not None: + model.setParam("TimeLimit", TimeLimit) + if Threads is not None: + model.setParam("Threads", Threads) + if LogFile is not None: + model.setParam("LogFile", os.fspath(LogFile)) + + try: + model.optimize() + status = model.status + sol_count = int(model.SolCount) + + if sol_count > 0: + self._store_assignment_mip_starts( + a, + binomial_only=self.__binomial_only, + patient_wise_only=self.__patient_wise_only, + ) + self._store_profile_mip_starts() + result: typing.Literal["inside", "outside", "unknown"] = "inside" + elif status == GRB.INFEASIBLE: + result = "outside" + else: + result = "unknown" + + self.last_oracle_work = self._record_work( + model, kind="oracle", side=result if result != "unknown" else None, + ) + + if print_progress or verbose: + print( + f"[{datetime.datetime.now()}] excess_at_most -> {result} " + f"(status={status}, SolCount={sol_count}" + f"{f', Work={self.last_oracle_work:.3f}' if self.last_oracle_work is not None else ''})", + flush=True, + ) + finally: + # Do not leave SolutionLimit/MIPFocus for subsequent full minimizes. + model.setParam("SolutionLimit", 2000000000) + model.setParam("MIPFocus", 0) + self._clear_feasibility_cuts(model) + return result + def run_MINLP( # pylint: disable=too-many-locals, too-many-statements, too-many-branches, too-many-arguments self, expected_probability: float | None, @@ -1811,8 +2178,9 @@ def run_MINLP( # pylint: disable=too-many-locals, too-many-statements, too-many- """ if print_progress or verbose: print( - "Running MINLP for expected probability ", expected_probability, - " at time point ", self.time_point, " at time ", datetime.datetime.now() + f"[{datetime.datetime.now()}] Running MINLP for expected probability " + f"{expected_probability} at time point {self.time_point}", + flush=True, ) if binomial_only != self.__binomial_only or patient_wise_only != self.__patient_wise_only: raise ValueError( @@ -1840,6 +2208,8 @@ def run_MINLP( # pylint: disable=too-many-locals, too-many-statements, too-many- expected_probability_var, _use_binomial_penalty_indicator, ) = self.gurobi_model + # Drop any leftover feasibility cuts before a full minimize. + self._clear_feasibility_cuts(model) self.update_model_with_expected_probability( model=model, km_probability_var=km_probability_var, @@ -1916,8 +2286,15 @@ def run_MINLP( # pylint: disable=too-many-locals, too-many-statements, too-many- model = self._optimize_with_fallbacks( model, initial_gurobi_params, fallback_strategies, verbose ) + self.last_minimize_work = self._last_minimize_call_work if model.status != GRB.OPTIMAL: + self._log_run_minlp_completion( + expected_probability=expected_probability, + status=model.status, + print_progress=print_progress, + verbose=verbose, + ) if model.status == GRB.INFEASIBLE and patient_wise_only: # If the model is infeasible, it means that no patients can be selected # while satisfying the constraints. This can happen if the expected @@ -1991,6 +2368,13 @@ def run_MINLP( # pylint: disable=too-many-locals, too-many-statements, too-many- print("Patient penalty: ", 2*patient_penalty_val) print("Total penalty: ", model.ObjVal) + self._log_run_minlp_completion( + expected_probability=expected_probability, + status=model.status, + print_progress=print_progress, + verbose=verbose, + ) + return scipy.optimize.OptimizeResult( x=model.ObjVal, success=model.status == GRB.OPTIMAL, diff --git a/kombine/kaplan_meier_hazard_ratio_MINLP.py b/kombine/kaplan_meier_hazard_ratio_MINLP.py index fee80ba..2430b84 100644 --- a/kombine/kaplan_meier_hazard_ratio_MINLP.py +++ b/kombine/kaplan_meier_hazard_ratio_MINLP.py @@ -17,7 +17,7 @@ import scipy.stats from .kaplan_meier_p_value_MINLP import MINLPforKMPValue -from .utilities import LOG_ZERO_EPSILON_DEFAULT +from .utilities import LOG_ZERO_EPSILON_DEFAULT, brentq_hazard_ratio_ci class MINLPforKMHazardRatio(MINLPforKMPValue): @@ -197,7 +197,11 @@ def compute_2nll_at_hazard_ratio( # pylint: disable=too-many-locals, too-many-a The Gurobi model (for advanced users). """ if print_progress or verbose: - print(f"Computing 2NLL at hazard ratio {hazard_ratio} at {datetime.datetime.now()}") + print( + f"[{datetime.datetime.now()}] Computing 2NLL at hazard ratio " + f"{hazard_ratio}", + flush=True, + ) ( # pylint: disable=duplicate-code model, @@ -261,35 +265,8 @@ def compute_2nll_at_hazard_ratio( # pylint: disable=too-many-locals, too-many-a self._store_hazard_ratio_mip_starts(a, cox_only, log_hazard_ratio) - # Extract results - twonll = model.ObjVal - patients_low, patients_high = self._extract_patients_per_curve(a) - patient_penalty = self._compute_patient_wise_penalty_value(a) - cox_penalty = self._compute_cox_penalty(model) - - # Extract curve statistics - (n_total_low, n_alive_low, km_prob_low, - n_total_high, n_alive_high, km_prob_high) = ( - self._extract_curve_statistics(model) - ) - - result = scipy.optimize.OptimizeResult( - x=twonll, - success=True, - patients_low=patients_low, - patients_high=patients_high, - n_total_low=n_total_low, - n_alive_low=n_alive_low, - n_total_high=n_total_high, - n_alive_high=n_alive_high, - km_probability_low=km_prob_low, - km_probability_high=km_prob_high, - cox_2NLL=2 * cox_penalty, - patient_2NLL=2 * patient_penalty, - hazard_ratio=hazard_ratio, - log_hazard_ratio=log_hazard_ratio, - model=model, - ) + result = self._extract_optimize_result(model, a) + result.log_hazard_ratio = log_hazard_ratio # Check if we're at or near the bounds and warn self._check_hazard_ratio_bounds(hazard_ratio, log_hazard_ratio) @@ -468,31 +445,13 @@ def twonll_minus_threshold(log_hr): result = self.compute_2nll_at_hazard_ratio(hr, cox_only=cox_only) return result.x - twonll_threshold - # Lower bound: search between min and best-fit - try: - lower_log_hr = scipy.optimize.brentq( - twonll_minus_threshold, - np.log(hazard_ratio_min), - best_fit_log_hr, - xtol=tolerance - ) - lower_ci = np.exp(float(lower_log_hr)) # type: ignore[arg-type] - except ValueError: - # If no crossing found, the lower bound is at the search limit - lower_ci = hazard_ratio_min - - # Upper bound: search between best-fit and max - try: - upper_log_hr = scipy.optimize.brentq( - twonll_minus_threshold, - best_fit_log_hr, - np.log(hazard_ratio_max), - xtol=tolerance - ) - upper_ci = np.exp(float(upper_log_hr)) # type: ignore[arg-type] - except ValueError: - # If no crossing found, the upper bound is at the search limit - upper_ci = hazard_ratio_max + lower_ci, upper_ci = brentq_hazard_ratio_ci( + twonll_minus_threshold, + best_fit_log_hr=best_fit_log_hr, + hazard_ratio_min=hazard_ratio_min, + hazard_ratio_max=hazard_ratio_max, + tolerance=tolerance, + ) # Check if best fit is at or near the bounds self._check_hazard_ratio_bounds(best_fit_hr, best_fit_result.log_hazard_ratio) diff --git a/kombine/kaplan_meier_likelihood.py b/kombine/kaplan_meier_likelihood.py index 5092ce3..5674b31 100644 --- a/kombine/kaplan_meier_likelihood.py +++ b/kombine/kaplan_meier_likelihood.py @@ -20,12 +20,17 @@ import scipy.optimize import scipy.stats -from .discrete_optimization import binary_search_sign_change, cached_level_crossings +from .discrete_optimization import ( + SignOracleStatus, + binary_search_sign_change, + cached_level_crossings, + feasibility_assisted_level_crossings, +) from .kaplan_meier import ( KaplanMeierBase, KaplanMeierInstance, ) -from .kaplan_meier_MINLP import MINLPForKM, KaplanMeierPatientNLL +from .kaplan_meier_MINLP import GurobiWorkStats, MINLPForKM, KaplanMeierPatientNLL from .utilities import InspectableCache, LOG_ZERO_EPSILON_DEFAULT @dataclasses.dataclass @@ -197,7 +202,7 @@ def __post_init__(self): ) self.CL_hatches = self.CL_hatches[:len(self.CLs)] -class KaplanMeierLikelihood(KaplanMeierBase): +class KaplanMeierLikelihood(KaplanMeierBase): # pylint: disable=too-many-instance-attributes """ Kaplan-Meier curve with error bars calculated using the log-likelihood method. """ @@ -222,6 +227,12 @@ def __init__( # pylint: disable=too-many-arguments self.__log_zero_epsilon = log_zero_epsilon self.__collapse_consecutive_deaths = collapse_consecutive_deaths self.__time_unit = time_unit + self._last_gurobi_work_stats: GurobiWorkStats | None = None + + @property + def last_gurobi_work_stats(self) -> GurobiWorkStats | None: + """Work totals from the most recent ``survival_probabilities_likelihood`` call.""" + return self._last_gurobi_work_stats @property def all_patients(self) -> list[KaplanMeierPatientNLL]: @@ -310,6 +321,7 @@ def get_twoNLL_function( # pylint: disable=too-many-arguments MIPGapAbs=None, rerun_until_convergence=False, assignment_starts=None, + minlp_time_limit: float | None = None, ) -> tuple[ InspectableCache[float | None, scipy.optimize.OptimizeResult], InspectableCache[float | None, float], @@ -349,6 +361,7 @@ def run_MINLP(expected_probability: float | None) -> scipy.optimize.OptimizeResu print_progress=print_progress, MIPGap=MIPGap, MIPGapAbs=MIPGapAbs, + TimeLimit=minlp_time_limit, ) if not rerun_until_convergence or not result.success: @@ -370,6 +383,7 @@ def run_MINLP(expected_probability: float | None) -> scipy.optimize.OptimizeResu print_progress=print_progress, MIPGap=MIPGap, MIPGapAbs=MIPGapAbs, + TimeLimit=minlp_time_limit, ) except Exception: #pylint: disable=broad-exception-caught # If rerun raises an exception, return the last successful result @@ -479,6 +493,14 @@ def survival_probabilities_likelihood( # pylint: disable=too-many-locals, too-ma MIPGap=None, MIPGapAbs=None, rerun_until_convergence=False, + crossing_mode: typing.Literal["full", "feasibility"] = "full", + component_cuts: bool = False, + binom_time_cuts: bool = False, + cost_biased_bisection: bool = True, + r_prior: float = 0.5, + kappa: float = 4.0, + minlp_time_limit: float | None = None, + oracle_time_limit: float | None = None, ) -> tuple[np.ndarray, np.ndarray]: """ Get the survival probabilities for the given quantiles. @@ -505,12 +527,40 @@ def survival_probabilities_likelihood( # pylint: disable=too-many-locals, too-ma Gurobi absolute MIP gap tolerance (used for objective function tolerance) rerun_until_convergence : bool, default False If True, rerun the MINLP until the result converges within tolerances + crossing_mode : {"full", "feasibility"}, default "full" + How to locate profile-likelihood CL endpoints. ``full`` minimizes + 2NLL at every trial ``p`` (legacy). ``feasibility`` uses + ``excess_at_most`` sign tests for bracketing, with full minimizes + only for MLE, oracle fallbacks, and brentq polish. + component_cuts : bool, default False + When ``crossing_mode="feasibility"``, also add redundant + ``2*binom <= T`` and ``2*patient <= T`` cuts. Off by default; + sum-only was faster with matching endpoints in hard-case ablations. + binom_time_cuts : bool, default False + When ``crossing_mode="feasibility"``, also add redundant + ``2*binom_piece[i] <= T`` per death time. + cost_biased_bisection : bool, default True + When ``crossing_mode="feasibility"``, use cost-biased oracle bisection + (prior probe fraction ``r_prior``; adapts from oracle timings). + r_prior : float, default 0.5 + Initial probe fraction between inside and outside endpoints (0.5 = + midpoint bisection). + kappa : float, default 4.0 + Prior strength (pseudo-counts) for the Beta prior on ``r``. + minlp_time_limit : float, optional + Gurobi TimeLimit (seconds) for each full ``run_MINLP`` call. None = no limit. + oracle_time_limit : float, optional + Gurobi TimeLimit (seconds) for each ``excess_at_most`` oracle call. None = no limit. Returns ------- tuple[np.ndarray, np.ndarray] Best probabilities and survival probabilities for each confidence level """ + if crossing_mode not in {"full", "feasibility"}: + raise ValueError( + f"crossing_mode must be 'full' or 'feasibility', got {crossing_mode!r}" + ) # Set default tolerance values if not provided if MIPGap is None: MIPGap = self.__default_MIPGap @@ -522,12 +572,13 @@ def survival_probabilities_likelihood( # pylint: disable=too-many-locals, too-ma for i, t in enumerate(times_for_plot, start=1): if print_progress: print( - f"Calculating survival probabilities for time point {t:.2f} " - f"({i} / {len(times_for_plot)}) at time {datetime.datetime.now()}" + f"[{datetime.datetime.now()}] Calculating survival probabilities " + f"for time point {t:.2f} ({i} / {len(times_for_plot)})", + flush=True, ) survival_probabilities_time_point = [] survival_probabilities.append(survival_probabilities_time_point) - run_MINLP, twoNLL, _minlp = self.get_twoNLL_function( + run_MINLP, twoNLL, minlp = self.get_twoNLL_function( time_point=t, binomial_only=binomial_only, patient_wise_only=patient_wise_only, @@ -536,6 +587,7 @@ def survival_probabilities_likelihood( # pylint: disable=too-many-locals, too-ma MIPGap=MIPGap, MIPGapAbs=MIPGapAbs, rerun_until_convergence=rerun_until_convergence, + minlp_time_limit=minlp_time_limit, ) # Find the expected probability that minimizes the negative log-likelihood # for the given time point @@ -623,7 +675,7 @@ def objective_function( else: levels = [scipy.stats.chi2.ppf(cl, 1).item() for cl in CLs] - def profile_excess( + def profile_excess( # pylint: disable=cell-var-from-loop expected_probability: float, _twoNLL=twoNLL, _twoNLL_min=twoNLL_min, @@ -638,47 +690,114 @@ def profile_excess( eps = self.__endpoint_epsilon one_minus_eps = 1.0 - eps - if best_prob <= eps: - lowers = [0.0] * len(CLs) - else: - lowers = cached_level_crossings( + def resolve_crossings( # pylint: disable=cell-var-from-loop + x_outer: float, + x_inner: float, + ) -> tuple[list[float], list[SignOracleStatus | None]]: + if crossing_mode == "full": + return cached_level_crossings( + profile_excess, + x_outer, + x_inner, + levels, + xtol=brentq_xtol, + rtol=brentq_rtol, + ), [None] * len(levels) + + def sign_oracle( # pylint: disable=cell-var-from-loop + p: float, + level: float, + _minlp=minlp, + _twoNLL_min=twoNLL_min, + ) -> tuple[typing.Literal["inside", "outside", "unknown"], float]: + status = _minlp.excess_at_most( + p, + twoNLL_min=_twoNLL_min, + level=level, + component_cuts=component_cuts, + binom_time_cuts=binom_time_cuts, + verbose=gurobi_verbose, + print_progress=print_progress, + TimeLimit=oracle_time_limit, + ) + # Prefer Gurobi Work (sleep-immune). Fall back to a tiny positive + # cost so the tracker still gets an outcome-count update. + work = _minlp.last_oracle_work + return status, float(work) if work is not None and work > 0.0 else 1e-9 + + outer_oracle: list[SignOracleStatus | None] = [] + crossings = feasibility_assisted_level_crossings( profile_excess, - eps, - best_prob_clipped, + sign_oracle, + x_outer, + x_inner, levels, xtol=brentq_xtol, rtol=brentq_rtol, + cost_biased_bisection=cost_biased_bisection, + r_prior=r_prior, + kappa=kappa, + outer_oracle=outer_oracle, ) - f_eps = profile_excess(eps) - lowers = [ - 0.0 if f_eps <= level else float(x) - for x, level in zip(lowers, levels, strict=True) - ] + return crossings, outer_oracle + + def _clip_lower_endpoint( # pylint: disable=cell-var-from-loop + crossings: list[float], + outer_oracle: list[SignOracleStatus | None], + ) -> list[float]: + clipped: list[float] = [] + for x, level, ostatus in zip(crossings, levels, outer_oracle, strict=True): + if ostatus == "inside": + clipped.append(0.0) + elif ostatus == "outside": + clipped.append(float(x)) + else: + f_outer = profile_excess(eps) + clipped.append(0.0 if f_outer <= level else float(x)) + return clipped + + def _clip_upper_endpoint( # pylint: disable=cell-var-from-loop + crossings: list[float], + outer_oracle: list[SignOracleStatus | None], + ) -> list[float]: + clipped: list[float] = [] + for x, level, ostatus in zip(crossings, levels, outer_oracle, strict=True): + if ostatus == "inside": + clipped.append(1.0) + elif ostatus == "outside": + clipped.append(float(x)) + else: + f_outer = profile_excess(one_minus_eps) + clipped.append(1.0 if f_outer <= level else float(x)) + return clipped + + if best_prob <= eps: + lowers = [0.0] * len(CLs) + else: + lowers, lower_outer = resolve_crossings(eps, best_prob_clipped) + lowers = _clip_lower_endpoint(lowers, lower_outer) if best_prob >= one_minus_eps: uppers = [1.0] * len(CLs) else: - uppers = cached_level_crossings( - profile_excess, - one_minus_eps, - best_prob_clipped, - levels, - xtol=brentq_xtol, - rtol=brentq_rtol, - ) - f_hi = profile_excess(one_minus_eps) - uppers = [ - 1.0 if f_hi <= level else float(x) - for x, level in zip(uppers, levels, strict=True) - ] + uppers, upper_outer = resolve_crossings(one_minus_eps, best_prob_clipped) + uppers = _clip_upper_endpoint(uppers, upper_outer) if print_progress: + ws = minlp.work_stats print( - f" CL crossings unique p evaluations={len(twoNLL.cache)} " - f"for time point {t}" + f"[{datetime.datetime.now()}] CL crossings unique p " + f"evaluations={len(twoNLL.cache)} for time point {t} " + f"(mode={crossing_mode}); " + f"Work oracle={ws.oracle_work:.3f} minimize={ws.minimize_work:.3f} " + f"total={ws.total_work:.3f} " + f"(oracle_calls={ws.oracle_calls}, minimize_calls={ws.minimize_calls}, " + f"oracle_outside={ws.oracle_outside_calls})", + flush=True, ) for lower_bound, upper_bound in zip(lowers, uppers, strict=True): survival_probabilities_time_point.append((lower_bound, upper_bound)) + self._last_gurobi_work_stats = minlp.work_stats return np.array(best_probabilities), np.array(survival_probabilities) def plot(self, config: KaplanMeierPlotConfig | None = None, **kwargs) -> dict: diff --git a/kombine/kaplan_meier_p_value_MINLP.py b/kombine/kaplan_meier_p_value_MINLP.py index 3c8fde3..ea7a2c9 100644 --- a/kombine/kaplan_meier_p_value_MINLP.py +++ b/kombine/kaplan_meier_p_value_MINLP.py @@ -60,6 +60,8 @@ def __init__( # pylint: disable=too-many-arguments self.__null_hypothesis_constraint = None self.__patient_wise_only_constraint = None self.__cox_penalty_constraint = None + self.__risk_set_r_vars = None + self.__risk_set_d_vars = None @property def all_patients(self) -> list[KaplanMeierPatientNLL]: @@ -161,53 +163,49 @@ def parameter_in_range(self) -> npt.NDArray[np.bool_]: & (self.observed_parameters < self.parameter_max) ))).T + @functools.cached_property + def assignment_nlls(self) -> npt.NDArray[np.float64]: + """ + Absolute min NLL for each patient on low, high, and neither. + + Returns an n x 3 array. An infinite bound makes that tail empty + (min NLL +inf). When both bounds are infinite, column 2 is +inf + and every patient must be assigned to a group. + """ + nlls = np.empty((self.n_patients, 3), dtype=float) + for i, patient in enumerate(self.all_patients): + nlls[i, 0] = patient.min_nll_on_interval( + self.parameter_min, self.parameter_threshold, + ) + nlls[i, 1] = patient.min_nll_on_interval( + self.parameter_threshold, self.parameter_max, + ) + nlls[i, 2] = min( + patient.min_nll_on_interval(-np.inf, self.parameter_min), + patient.min_nll_on_interval(self.parameter_max, np.inf), + ) + return nlls + @functools.cached_property def nll_penalty_for_patient_in_range(self) -> npt.NDArray[np.float64]: """ - Calculate the negative log-likelihood penalty for each patient - if that patient is within the parameter range. - This is negative if the patient's observed parameter is within the range - and positive if it is outside the range. - Returns an n x 2 array: for each patient, the penalty to be included - in the low and high curves. - """ - sgn_nll_penalty_for_patient_in_range = 2 * self.parameter_in_range - 1 - observed_nll = np.array([ - p.parameter(p.observed_parameter) - for p in self.all_patients - ]) - parameter_min_nll: npt.NDArray[np.float64] = np.array([ - p.parameter(self.parameter_min) if np.isfinite(self.parameter_min) else np.inf - for p in self.all_patients - ]) - parameter_threshold_nll: npt.NDArray[np.float64] = np.array([ - p.parameter(self.parameter_threshold) #parameter threshold must be finite - for p in self.all_patients - ]) - parameter_max_nll: npt.NDArray[np.float64] = np.array([ - p.parameter(self.parameter_max) if np.isfinite(self.parameter_max) else np.inf - for p in self.all_patients - ]) - - range_boundary_nll_low = np.min( - np.array([parameter_min_nll, parameter_threshold_nll]), - axis=0 - ) - range_boundary_nll_high = np.min( - np.array([parameter_threshold_nll, parameter_max_nll]), - axis=0 - ) + Relative NLL for assigning each patient to the low and high ranges. - range_boundary_nll: npt.NDArray[np.float64] = \ - np.array([range_boundary_nll_low, range_boundary_nll_high]).T - abs_nll_penalty_for_patient_in_range = observed_nll - range_boundary_nll.T + Each column is min NLL on that range minus the best of low, high, + and neither. Returns an n x 2 array of non-negative values. + """ + best = np.min(self.assignment_nlls, axis=1, keepdims=True) + return self.assignment_nlls[:, :2] - best - nll_penalty_for_patient_in_range = ( - sgn_nll_penalty_for_patient_in_range - * abs_nll_penalty_for_patient_in_range.T - ) + @functools.cached_property + def nll_penalty_for_unassigned(self) -> npt.NDArray[np.float64]: + """ + Relative NLL for leaving each patient in neither group. - return nll_penalty_for_patient_in_range + This is min NLL on the tails minus the best of low, high, and neither. + """ + best = np.min(self.assignment_nlls, axis=1) + return self.assignment_nlls[:, 2] - best def add_counter_variables_and_constraints( self, @@ -271,6 +269,43 @@ def add_counter_variables_and_constraints( return r, d, n_survived + def _refresh_death_incidence_constraints( + self, + model: gp.Model, + a: gp.tupledict[tuple[int, ...], gp.Var], + r: gp.tupledict[tuple[int, ...], gp.Var], + d: gp.tupledict[tuple[int, ...], gp.Var], + ) -> None: + """ + Rebuild r/d incidence constraints after (time, censored) change. + + Death-time grid size is invariant under outcome permutation; only which + patients contribute to each death/risk-set row changes. + """ + death_times = self.all_death_times + for k, t in enumerate(death_times): + at_risk = self.patient_still_at_risk(t) + for j in range(2): + for name in (f"r_{k}_{j}", f"d_{k}_{j}"): + constr = model.getConstrByName(name) + if constr is not None: + model.remove(constr) + model.addConstr( + r[k, j] == gp.quicksum( + a[i, j] for i in range(self.n_patients) if at_risk[i] + ), + name=f"r_{k}_{j}", + ) + model.addConstr( + d[k, j] == gp.quicksum( + a[i, j] for i in range(self.n_patients) + if self.all_patients[i].time == t + and not self.all_patients[i].censored + ), + name=f"d_{k}_{j}", + ) + model.update() + def add_kaplan_meier_probability_variables_and_constraints( #pylint: disable=too-many-locals self, model: gp.Model, @@ -612,35 +647,42 @@ def add_patient_wise_penalty( a: gp.tupledict[tuple[int, ...], gp.Var], ): """ - Add the patient-wise penalty to the model. - This penalty is based on the negative log-likelihood of the patient's observed parameter - being within the specified range. + Add the patient-wise measurement NLL to the model. + + Each patient pays the relative min NLL of the chosen bin (low, high, + or neither). The best bin is 0. A group-to-group flip costs one + measurement NLL difference, not two. Exclusion uses the tail NLL, + not the cost of flipping into the other group. """ patient_penalties = [] + rel_range = self.nll_penalty_for_patient_in_range + rel_neither = self.nll_penalty_for_unassigned for i in range(self.n_patients): for j in range(2): - if np.isfinite(self.nll_penalty_for_patient_in_range[i, j]): - penalty = self.nll_penalty_for_patient_in_range[i, j] * a[i, j] - if self.nll_penalty_for_patient_in_range[i, j] < 0: - # If the penalty is negative, it means the patient is nominally in the range - # We want the penalty to be 0 when all patients are at their nominal values - penalty -= self.nll_penalty_for_patient_in_range[i, j] - patient_penalties.append(penalty) - elif np.isneginf(self.nll_penalty_for_patient_in_range[i, j]): - #the patient must be selected, so we add a constraint - model.addConstr( - a[i, j] == 1, name=f"patient_{i}_must_be_in_curve_{j}", - ) - elif np.isposinf(self.nll_penalty_for_patient_in_range[i, j]): - # The patient must not be selected, so we add a constraint + penalty_ij = rel_range[i, j] + if np.isfinite(penalty_ij): + patient_penalties.append(penalty_ij * a[i, j]) + elif np.isposinf(penalty_ij): model.addConstr( a[i, j] == 0, name=f"patient_{i}_must_not_be_in_curve_{j}", ) else: raise ValueError( f"Invalid negative log-likelihood penalty value for patient {i}, curve {j}:" - f"{self.nll_penalty_for_patient_in_range[i, j]}" + f"{penalty_ij}" ) + neither_i = rel_neither[i] + if np.isfinite(neither_i): + patient_penalties.append(neither_i * (1 - a[i, 0] - a[i, 1])) + elif np.isposinf(neither_i): + model.addConstr( + a[i, 0] + a[i, 1] == 1, + name=f"patient_{i}_must_be_assigned", + ) + else: + raise ValueError( + f"Invalid unassigned NLL penalty for patient {i}: {neither_i}" + ) patient_penalty = gp.quicksum(patient_penalties) return patient_penalty @@ -722,13 +764,17 @@ def _compute_patient_wise_penalty_value(self, a: gp.tupledict[tuple[int, ...], g float: The patient-wise penalty value (not multiplied by 2) """ penalty = 0.0 + rel_range = self.nll_penalty_for_patient_in_range + rel_neither = self.nll_penalty_for_unassigned for i in range(self.n_patients): - for j in range(2): - if np.isfinite(self.nll_penalty_for_patient_in_range[i, j]): - contribution = self.nll_penalty_for_patient_in_range[i, j] * ( - a[i, j].X - (1 if self.nll_penalty_for_patient_in_range[i, j] < 0 else 0) - ) - penalty += contribution + assigned_low = a[i, 0].X + assigned_high = a[i, 1].X + if np.isfinite(rel_range[i, 0]): + penalty += rel_range[i, 0] * assigned_low + if np.isfinite(rel_range[i, 1]): + penalty += rel_range[i, 1] * assigned_high + if np.isfinite(rel_neither[i]): + penalty += rel_neither[i] * (1.0 - assigned_low - assigned_high) return penalty def _compute_cox_penalty(self, model: gp.Model): @@ -774,6 +820,8 @@ def _make_gurobi_model(self): GRB.MINIMIZE, ) model.update() + self.__risk_set_r_vars = r + self.__risk_set_d_vars = d return ( # pylint: disable=duplicate-code model, @@ -927,11 +975,209 @@ def update_model_with_patient_wise_only_constraint( #pylint: disable=too-many-ar model.update() - def solve_and_pvalue( # pylint: disable=too-many-locals, too-many-arguments, too-many-branches, too-many-statements + def _get_risk_set_vars(self): + """Return cached r/d tupledicts from the current Gurobi model, if any.""" + return self.__risk_set_r_vars, self.__risk_set_d_vars + + def _dispose_gurobi_model(self) -> None: + """Dispose a cached Gurobi model and drop outcome-tied constraint handles.""" + cached = self.__dict__.get("gurobi_model") + if cached is not None: + try: + cached[0].dispose() + except gp.GurobiError: + pass + self.__dict__.pop("gurobi_model", None) + self.__null_hypothesis_constraint = None + self.__patient_wise_only_constraint = None + self.__cox_penalty_constraint = None + self.__risk_set_r_vars = None + self.__risk_set_d_vars = None + + def _invalidate_outcome_dependent_state(self, *, keep_model: bool = False) -> None: + """ + Clear caches that depend on (time, censored). + + Measurement / assignment-NLL caches are kept: permutations shuffle + outcomes only. When keep_model is True, the Gurobi model is retained so + death-incidence constraints can be refreshed in place. + """ + if keep_model: + for name in ( + "patient_times", + "all_death_times", + "patient_censored", + "nominal_hazard_ratio", + ): + self.__dict__.pop(name, None) + return + self._dispose_gurobi_model() + for name in ( + "patient_times", + "all_death_times", + "patient_censored", + "nominal_hazard_ratio", + ): + self.__dict__.pop(name, None) + + def _set_patient_outcomes( + self, + times: list[float], + censored: list[bool], + *, + keep_model: bool = False, + ) -> None: + """ + Replace each patient's (time, censored), keeping measurement NLLs fixed. + """ + if len(times) != self.n_patients or len(censored) != self.n_patients: + raise ValueError( + "times and censored must have length n_patients " + f"({self.n_patients}), got {len(times)} and {len(censored)}" + ) + self.__all_patients = [ + KaplanMeierPatientNLL( + time=float(times[i]), + censored=bool(censored[i]), + parameter_nll=patient.parameter, + observed_parameter=patient.observed_parameter, + ) + for i, patient in enumerate(self.all_patients) + ] + self._invalidate_outcome_dependent_state(keep_model=keep_model) + + def _extract_optimize_result( + self, + model: gp.Model, + a: gp.tupledict[tuple[int, ...], gp.Var], + ) -> scipy.optimize.OptimizeResult: + """ + Pack the current Gurobi solution into an OptimizeResult. + """ + patients_low, patients_high = self._extract_patients_per_curve(a) + ( + n_total_low, n_alive_low, km_prob_low, + n_total_high, n_alive_high, km_prob_high, + ) = self._extract_curve_statistics(model) + log_hazard_ratio_var = model.getVarByName("log_hazard_ratio") + assert log_hazard_ratio_var is not None + return scipy.optimize.OptimizeResult( + x=model.ObjVal, + success=model.status == GRB.OPTIMAL, + patients_low=patients_low, + patients_high=patients_high, + n_total_low=n_total_low, + n_alive_low=n_alive_low, + n_total_high=n_total_high, + n_alive_high=n_alive_high, + km_probability_low=km_prob_low, + km_probability_high=km_prob_high, + cox_2NLL=2 * self._compute_cox_penalty(model), + patient_2NLL=2 * self._compute_patient_wise_penalty_value(a), + patient_penalties=self.nll_penalty_for_patient_in_range, + hazard_ratio=np.exp(log_hazard_ratio_var.X), + model=model, + ) + + def _fit_null_and_alternative( # pylint: disable=too-many-arguments, too-many-locals + self, + *, + cox_only: bool, + patient_wise_only: bool, + verbose: bool, + print_progress: bool, + solve_null: bool, + MIPGap: float, + MIPGapAbs: float, + TimeLimit: float | None, + Threads: int | None, + MIPFocus: int | None, + LogFile: os.PathLike | None, + ) -> tuple[ + scipy.optimize.OptimizeResult | None, + scipy.optimize.OptimizeResult, + ]: + """ + Fit H0 and/or H1 on this calculator's observed (or permuted) data. + """ + ( + model, + null_hypothesis_indicator, + a, + beta, + use_cox_penalty_indicator, + ) = self.gurobi_model + + self.update_model_with_cox_only_constraints(model, a, cox_only) + self.update_model_with_patient_wise_only_constraint( + model, + beta=beta, + null_hypothesis_indicator=null_hypothesis_indicator, + patient_wise_only=patient_wise_only, + use_cox_penalty_indicator=use_cox_penalty_indicator, + ) + + solver_kwargs = { + "verbose": verbose, + "MIPGap": MIPGap, + "MIPGapAbs": MIPGapAbs, + "TimeLimit": TimeLimit, + "Threads": Threads, + "MIPFocus": MIPFocus, + "LogFile": LogFile, + } + + result_null = None + start_a = None + start_beta = None + if solve_null: + if print_progress or verbose: + print( + f"[{datetime.datetime.now()}] Solving for null hypothesis...", + flush=True, + ) + self.update_model_for_null_hypothesis_or_not( + model, null_hypothesis_indicator, True, + ) + model = self._setup_and_optimize(model, **solver_kwargs) + if model.status != GRB.OPTIMAL: + raise ValueError(f"Null model failed with status {model.status}") + result_null = self._extract_optimize_result(model, a) + if not cox_only: + start_a = { + (i, j): float(a[i, j].X) + for i in range(self.n_patients) + for j in range(2) + } + if beta is not None: + start_beta = float(beta.X) + + if print_progress or verbose: + print( + f"[{datetime.datetime.now()}] Solving for alternative hypothesis...", + flush=True, + ) + self.update_model_for_null_hypothesis_or_not( + model, null_hypothesis_indicator, False, + ) + if start_a is not None: + for (i, j), value in start_a.items(): + a[i, j].Start = value + if start_beta is not None: + beta.Start = start_beta + model = self._setup_and_optimize(model, **solver_kwargs) + if model.status != GRB.OPTIMAL: + raise ValueError(f"Alternative model failed with status {model.status}") + result_alt = self._extract_optimize_result(model, a) + return result_null, result_alt + + def solve_and_pvalue( # pylint: disable=too-many-locals, too-many-arguments, too-many-branches self, *, cox_only: bool = False, patient_wise_only: bool = False, + n_permutations: int = 199, + rng: int | np.random.Generator | None = None, verbose: bool = False, print_progress: bool = False, MIPGap: float | None = None, @@ -944,6 +1190,12 @@ def solve_and_pvalue( # pylint: disable=too-many-locals, too-many-arguments, too """ Solve the MINLP and return the p value. + When cox_only is False, the p-value is a permutation LRT: (time, + censored) are shuffled across patients while measurements stay fixed, + so the null has the same assignment freedom as the alternative. + When cox_only is True, assignments are locked and the reference is + chi-squared with 1 degree of freedom. + Parameters ---------- cox_only : bool, optional @@ -953,13 +1205,25 @@ def solve_and_pvalue( # pylint: disable=too-many-locals, too-many-arguments, too If True, only consider patient-wise errors and constrain the curves to be flipped relative to nominal at each death time point under the null hypothesis. Default is False. + n_permutations : int, optional + Number of outcome permutations when cox_only is False. Ignored + when cox_only is True. Default is 199. + rng : int or numpy.random.Generator or None, optional + Seed or generator for the permutations. Default is None. verbose : bool, optional If True, enable verbose output from Gurobi solver. Default is False. """ if print_progress or verbose: - print(f"Running p-value MINLP at {datetime.datetime.now()}") + print( + f"[{datetime.datetime.now()}] Running p-value MINLP", + flush=True, + ) if cox_only and patient_wise_only: raise ValueError("cox_only and patient_wise_only cannot both be True") + if n_permutations < 0: + raise ValueError( + f"n_permutations must be >= 0, got {n_permutations}" + ) if patient_wise_only: #make sure the nominal hazard ratio is cached before doing anything with the Gurobi model @@ -971,156 +1235,101 @@ def solve_and_pvalue( # pylint: disable=too-many-locals, too-many-arguments, too if MIPGapAbs is None: MIPGapAbs = self.__default_MIPGapAbs - ( # pylint: disable=duplicate-code - model, - null_hypothesis_indicator, - a, - beta, - use_cox_penalty_indicator, - ) = self.gurobi_model - - # Apply cox_only constraints if specified - self.update_model_with_cox_only_constraints(model, a, cox_only) - - # Apply patient_wise_only constraints if specified - self.update_model_with_patient_wise_only_constraint( - model, - beta=beta, - null_hypothesis_indicator=null_hypothesis_indicator, - patient_wise_only=patient_wise_only, - use_cox_penalty_indicator=use_cox_penalty_indicator, - ) - - # Setup and optimize with standard parameters and fallback strategies - if print_progress or verbose: - print("Solving for null hypothesis...") - self.update_model_for_null_hypothesis_or_not(model, null_hypothesis_indicator, True) - # pylint: disable=duplicate-code - model = self._setup_and_optimize( - model, - verbose=verbose, - MIPGap=MIPGap, - MIPGapAbs=MIPGapAbs, - TimeLimit=TimeLimit, - Threads=Threads, - MIPFocus=MIPFocus, - LogFile=LogFile, - ) - if model.status != GRB.OPTIMAL: - raise ValueError(f"Null model failed with status {model.status}") - # pylint: enable=duplicate-code - twonll_null = model.ObjVal - - # Extract detailed information for null hypothesis result - patients_low_null, patients_high_null = self._extract_patients_per_curve(a) - patient_penalty_null = self._compute_patient_wise_penalty_value(a) - cox_penalty_null = self._compute_cox_penalty(model) - - # Extract curve statistics for null hypothesis - (n_total_low_null, n_alive_low_null, km_prob_low_null, - n_total_high_null, n_alive_high_null, km_prob_high_null) = ( - self._extract_curve_statistics(model) + fit_kwargs = { + "cox_only": cox_only, + "patient_wise_only": patient_wise_only, + "verbose": verbose, + "print_progress": print_progress, + "MIPGap": MIPGap, + "MIPGapAbs": MIPGapAbs, + "TimeLimit": TimeLimit, + "Threads": Threads, + "MIPFocus": MIPFocus, + "LogFile": LogFile, + } + result_null, result_alt = self._fit_null_and_alternative( + solve_null=True, + **fit_kwargs, ) - - # Extract hazard ratio for null hypothesis (should be 1.0) - log_hazard_ratio_var = model.getVarByName("log_hazard_ratio") - assert log_hazard_ratio_var is not None - hazard_ratio_null = np.exp(log_hazard_ratio_var.X) - - result_null = scipy.optimize.OptimizeResult( - x=model.ObjVal, - success=model.status == GRB.OPTIMAL, - patients_low=patients_low_null, - patients_high=patients_high_null, - n_total_low=n_total_low_null, - n_alive_low=n_alive_low_null, - n_total_high=n_total_high_null, - n_alive_high=n_alive_high_null, - km_probability_low=km_prob_low_null, - km_probability_high=km_prob_high_null, - cox_2NLL=2*cox_penalty_null, - patient_2NLL=2*patient_penalty_null, - patient_penalties=self.nll_penalty_for_patient_in_range, - hazard_ratio=hazard_ratio_null, - model=model, - ) - - start_a = None - start_beta = None - if not cox_only: - start_a = { - (i, j): float(a[i, j].X) - for i in range(self.n_patients) - for j in range(2) - } - if beta is not None: - start_beta = float(beta.X) - - if print_progress or verbose: - print("Solving for alternative hypothesis...") - self.update_model_for_null_hypothesis_or_not(model, null_hypothesis_indicator, False) - if start_a is not None: - for (i, j), value in start_a.items(): - a[i, j].Start = value - if start_beta is not None: - beta.Start = start_beta - # pylint: disable=duplicate-code - model = self._setup_and_optimize( - model, - verbose=verbose, - MIPGap=MIPGap, - MIPGapAbs=MIPGapAbs, - TimeLimit=TimeLimit, - Threads=Threads, - MIPFocus=MIPFocus, - LogFile=LogFile, - ) - if model.status != GRB.OPTIMAL: - raise ValueError(f"Alternative model failed with status {model.status}") - # pylint: enable=duplicate-code - twonll_alt = model.ObjVal - - # Extract detailed information for alternative hypothesis result - patients_low_alt, patients_high_alt = self._extract_patients_per_curve(a) - patient_penalty_alt = self._compute_patient_wise_penalty_value(a) - cox_penalty_alt = self._compute_cox_penalty(model) - - # Extract curve statistics for alternative hypothesis - (n_total_low_alt, n_alive_low_alt, km_prob_low_alt, - n_total_high_alt, n_alive_high_alt, km_prob_high_alt) = ( - self._extract_curve_statistics(model) - ) - - # Extract hazard ratio for alternative hypothesis (can be any value) - hazard_ratio_alt = np.exp(log_hazard_ratio_var.X) - - result_alt = scipy.optimize.OptimizeResult( - x=model.ObjVal, - success=model.status == GRB.OPTIMAL, - patients_low=patients_low_alt, - patients_high=patients_high_alt, - n_total_low=n_total_low_alt, - n_alive_low=n_alive_low_alt, - n_total_high=n_total_high_alt, - n_alive_high=n_alive_high_alt, - km_probability_low=km_prob_low_alt, - km_probability_high=km_prob_high_alt, - cox_2NLL=2*cox_penalty_alt, - patient_2NLL=2*patient_penalty_alt, - patient_penalties=self.nll_penalty_for_patient_in_range, - hazard_ratio=hazard_ratio_alt, - model=model, + assert result_null is not None + + if cox_only: + lr_stat = result_null.x - result_alt.x + p_value = scipy.stats.chi2.sf(lr_stat, 1) + result_alt.n_permutations = None + result_alt.n_extreme = None + return p_value, result_null, result_alt + + base_times = [patient.time for patient in self.all_patients] + base_censored = [patient.censored for patient in self.all_patients] + perm_calc = type(self)( + list(self.all_patients), + parameter_min=self.parameter_min, + parameter_threshold=self.parameter_threshold, + parameter_max=self.parameter_max, + log_zero_epsilon=self.log_zero_epsilon, + tie_handling=self.tie_handling, + log_hazard_ratio_bounds=self.log_hazard_ratio_bounds, ) - lr_stat = twonll_null - twonll_alt - - # The degrees of freedom is 1: the only difference between null and alternative - # is whether the log hazard ratio is constrained to 0 (null) or free to float (alternative) - df = 1 - - p_value = scipy.stats.chi2.sf(lr_stat, df) + generator = np.random.default_rng(rng) + n_extreme = 0 + try: + for i_perm in range(n_permutations): + if print_progress or verbose: + print( + f"[{datetime.datetime.now()}] Permutation " + f"{i_perm + 1}/{n_permutations}...", + flush=True, + ) + order = generator.permutation(self.n_patients) + reshuffled_times = [base_times[order[i]] for i in range(self.n_patients)] + reshuffled_censored = [ + bool(base_censored[order[i]]) for i in range(self.n_patients) + ] + if i_perm == 0: + perm_calc._set_patient_outcomes( # pylint: disable=protected-access + reshuffled_times, + reshuffled_censored, + keep_model=False, + ) + else: + perm_calc._set_patient_outcomes( # pylint: disable=protected-access + reshuffled_times, + reshuffled_censored, + keep_model=True, + ) + ( + model, + _null_hypothesis_indicator, + a, + _beta, + _use_cox_penalty_indicator, + ) = perm_calc.gurobi_model + r_vars, d_vars = perm_calc._get_risk_set_vars() # pylint: disable=protected-access + assert r_vars is not None and d_vars is not None + perm_calc._refresh_death_incidence_constraints( # pylint: disable=protected-access + model, + a, + r_vars, + d_vars, + ) + _, perm_alt = perm_calc._fit_null_and_alternative( # pylint: disable=protected-access + solve_null=False, + **fit_kwargs, + ) + # 2NLL_0 is invariant to pairing, so T_perm >= T_obs iff alt is as good. + if perm_alt.x <= result_alt.x: + n_extreme += 1 + finally: + perm_calc._dispose_gurobi_model() # pylint: disable=protected-access + + p_value = (1.0 + n_extreme) / (1.0 + n_permutations) + result_alt.n_permutations = n_permutations + result_alt.n_extreme = n_extreme return p_value, result_null, result_alt + def survival_curves_pvalue_logrank( #pylint: disable=too-many-locals, too-many-branches self, *, diff --git a/kombine/utilities.py b/kombine/utilities.py index 9d65bb0..febf657 100644 --- a/kombine/utilities.py +++ b/kombine/utilities.py @@ -5,6 +5,7 @@ import typing import numpy as np +import scipy.optimize import scipy.stats # Default log zero epsilon value used across Kaplan-Meier likelihood methods @@ -334,3 +335,44 @@ def _optimize_with_fallbacks( print(f"Fallback {i+1} successful. Model is now optimal.") break return model + + +def brentq_hazard_ratio_ci( # pylint: disable=too-many-arguments + twonll_minus_threshold: typing.Callable[[float], float], + *, + best_fit_log_hr: float, + hazard_ratio_min: float, + hazard_ratio_max: float, + tolerance: float, +) -> tuple[float, float]: + """ + Find the profile-likelihood hazard-ratio CI by brentq on each side of the MLE. + + ``twonll_minus_threshold`` should return 2NLL(log H) minus the chi-square + threshold, so the roots are the CI endpoints in log-H space. If a side has + no sign change inside the search window, that endpoint is left at the + corresponding search limit. + """ + try: + lower_log_hr = scipy.optimize.brentq( + twonll_minus_threshold, + np.log(hazard_ratio_min), + best_fit_log_hr, + xtol=tolerance, + ) + lower_ci = float(np.exp(float(lower_log_hr))) # type: ignore[arg-type] + except ValueError: + lower_ci = hazard_ratio_min + + try: + upper_log_hr = scipy.optimize.brentq( + twonll_minus_threshold, + best_fit_log_hr, + np.log(hazard_ratio_max), + xtol=tolerance, + ) + upper_ci = float(np.exp(float(upper_log_hr))) # type: ignore[arg-type] + except ValueError: + upper_ci = hazard_ratio_max + + return lower_ci, upper_ci diff --git a/test/kombine/datacards/simple_examples/discrete_classes_hr_example_large.txt b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_large.txt index 027e7c8..c934321 100644 --- a/test/kombine/datacards/simple_examples/discrete_classes_hr_example_large.txt +++ b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_large.txt @@ -1,11 +1,11 @@ observable_type discrete_classes ------------ # Example datacard for hazard ratio calculation -# Binary misclassification matrix with large error (e=0.30) -# Group 0 (below threshold 0.5) has better survival than Group 1 (above 0.5) +# Binary misclassification matrix with large error (e=0.25) +# Class 0 has better survival than class 1 # Probabilities are homogeneous across patients within each true group ------------ survival_time 10 12 15 18 20 24 28 30 35 40 3 5 7 9 12 15 18 20 25 30 censored 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 1 1 1 1 1 -prob0 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 -prob1 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.30 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 0.70 +prob0 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 +prob1 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 diff --git a/test/kombine/datacards/simple_examples/discrete_classes_hr_example_moderate.txt b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_moderate.txt index 556b5ce..b36ddc1 100644 --- a/test/kombine/datacards/simple_examples/discrete_classes_hr_example_moderate.txt +++ b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_moderate.txt @@ -1,11 +1,11 @@ observable_type discrete_classes ------------ # Example datacard for hazard ratio calculation -# Binary misclassification matrix with moderate error (e=0.15) -# Group 0 (below threshold 0.5) has better survival than Group 1 (above 0.5) +# Binary misclassification matrix with moderate error (e=0.20) +# Class 0 has better survival than class 1 # Probabilities are homogeneous across patients within each true group ------------ survival_time 10 12 15 18 20 24 28 30 35 40 3 5 7 9 12 15 18 20 25 30 censored 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 1 1 1 1 1 -prob0 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 -prob1 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.15 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 0.85 +prob0 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 +prob1 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.20 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 0.80 diff --git a/test/kombine/datacards/simple_examples/discrete_classes_hr_example_very_large.txt b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_very_large.txt new file mode 100644 index 0000000..6e579c7 --- /dev/null +++ b/test/kombine/datacards/simple_examples/discrete_classes_hr_example_very_large.txt @@ -0,0 +1,11 @@ +observable_type discrete_classes +------------ +# Example datacard for hazard ratio calculation +# Binary misclassification matrix with very large error (e=0.40) +# Class 0 has better survival than class 1 +# Probabilities are homogeneous across patients within each true group +------------ +survival_time 10 12 15 18 20 24 28 30 35 40 3 5 7 9 12 15 18 20 25 30 +censored 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 1 1 1 1 1 +prob0 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 +prob1 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.40 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 0.60 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_discrete_e05.txt b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e05.txt new file mode 100644 index 0000000..5073056 --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e05.txt @@ -0,0 +1,9 @@ +observable_type discrete_classes +------------ +# Notebook 07 methods-comparison discrete classes, e=0.05 +# Same survival_time/censored as methods_comparison_fixed.txt +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +prob0 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 +prob1 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.05 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 0.95 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_discrete_e10.txt b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e10.txt new file mode 100644 index 0000000..9e95f3c --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e10.txt @@ -0,0 +1,9 @@ +observable_type discrete_classes +------------ +# Notebook 07 methods-comparison discrete classes, e=0.10 +# Same survival_time/censored as methods_comparison_fixed.txt +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +prob0 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 +prob1 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.10 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 0.90 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_discrete_e25.txt b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e25.txt new file mode 100644 index 0000000..27c8e03 --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_discrete_e25.txt @@ -0,0 +1,9 @@ +observable_type discrete_classes +------------ +# Notebook 07 methods-comparison discrete classes, e=0.25 +# Same survival_time/censored as methods_comparison_fixed.txt +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +prob0 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 +prob1 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.25 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_fixed.txt b/test/kombine/datacards/simple_examples/methods_comparison_fixed.txt new file mode 100644 index 0000000..31f903e --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_fixed.txt @@ -0,0 +1,9 @@ +observable_type fixed +------------ +# Notebook 07 methods-comparison baseline (CI-friendly stronger signal) +# n=50, deaths=33, target HR=6, n_per_group=25, seed=42 +# Threshold 0.5001; shared survival with sibling methods_comparison_* cards +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +observable 0.3869 0.1058 0.2568 0.2944 0.0323 0.3016 0.0919 0.0412 0.4655 0.4735 0.3980 0.1562 0.0569 0.3384 0.2213 0.0686 0.2477 0.0265 0.4465 0.1342 0.3280 0.1596 0.2596 0.2724 0.0987 0.8599 0.8802 0.5455 0.6821 0.5656 0.9243 0.8092 0.6688 0.5405 0.6593 0.6661 0.8602 0.8160 0.9359 0.7367 0.5674 0.8524 0.8752 0.7794 0.8801 0.7470 0.7609 0.7152 0.5222 0.5618 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_poisson_large.txt b/test/kombine/datacards/simple_examples/methods_comparison_poisson_large.txt new file mode 100644 index 0000000..29d9fa1 --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_poisson_large.txt @@ -0,0 +1,11 @@ +observable_type poisson_density +------------ +# Notebook 07 methods-comparison Poisson density (large) +# Large counts; relative Poisson error ~2-5% +# Same survival_time/censored as methods_comparison_fixed.txt +# Observed density ≈ fixed observable; hard split at 0.5001 +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +num 387 106 257 294 32 302 92 41 466 474 398 156 57 338 221 69 248 26 446 134 328 160 260 272 99 860 880 546 682 566 924 809 669 540 659 666 860 816 936 737 567 852 875 779 880 747 761 715 522 562 +area 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 1000 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_poisson_moderate.txt b/test/kombine/datacards/simple_examples/methods_comparison_poisson_moderate.txt new file mode 100644 index 0000000..b1ec69d --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_poisson_moderate.txt @@ -0,0 +1,11 @@ +observable_type poisson_density +------------ +# Notebook 07 methods-comparison Poisson density (moderate) +# Moderate counts; relative Poisson error ~10-20% +# Same survival_time/censored as methods_comparison_fixed.txt +# Observed density ≈ fixed observable; hard split at 0.5001 +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +num 39 11 26 29 3 30 9 4 47 47 40 16 6 34 22 7 25 3 45 13 33 16 26 27 10 86 88 55 68 57 92 81 67 54 66 67 86 82 94 74 57 85 88 78 88 75 76 72 52 56 +area 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 100 diff --git a/test/kombine/datacards/simple_examples/methods_comparison_poisson_small.txt b/test/kombine/datacards/simple_examples/methods_comparison_poisson_small.txt new file mode 100644 index 0000000..94c3241 --- /dev/null +++ b/test/kombine/datacards/simple_examples/methods_comparison_poisson_small.txt @@ -0,0 +1,11 @@ +observable_type poisson_density +------------ +# Notebook 07 methods-comparison Poisson density (small) +# Small counts; relative Poisson error ~25-70% +# Same survival_time/censored as methods_comparison_fixed.txt +# Observed density ≈ fixed observable; hard split at 0.5001 +------------ +survival_time 4.6927 30.1012 13.1675 9.1294 1.6962 1.6960 0.5984 20.1123 9.1908 12.3125 0.2080 35.0356 17.8643 2.3869 2.0068 2.0261 3.6275 7.4393 5.6554 3.4422 9.4637 1.5023 3.4552 4.5628 6.0893 5.8213 2.4871 4.6752 3.7536 1.5184 4.2491 0.1544 0.3636 0.0771 0.6559 0.8202 0.5276 2.9409 0.7354 0.5497 1.3040 0.2532 2.7008 0.1291 7.2236 2.4658 0.3692 0.0092 2.8165 2.0452 +censored 1 0 0 0 0 1 0 0 1 1 1 1 0 0 0 0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 0 0 0 0 0 0 0 1 0 0 1 1 0 0 1 1 +num 8 2 5 6 1 6 2 1 9 9 8 3 1 7 4 1 5 1 9 3 7 3 5 5 2 17 18 11 14 11 18 16 13 11 13 13 17 16 19 15 11 17 18 16 18 15 15 14 11 11 +area 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 20 diff --git a/test/kombine/reference/km_likelihood.json b/test/kombine/reference/km_likelihood.json index e32acad..af9ff89 100644 --- a/test/kombine/reference/km_likelihood.json +++ b/test/kombine/reference/km_likelihood.json @@ -437,7 +437,7 @@ 0.0 ], "p_value_breslow": [ - 0.0199 + 0.25 ], "p_value_cox_breslow": [ 0.175 diff --git a/test/kombine/reference/km_likelihood_with_censoring.json b/test/kombine/reference/km_likelihood_with_censoring.json index 4022312..fec9774 100644 --- a/test/kombine/reference/km_likelihood_with_censoring.json +++ b/test/kombine/reference/km_likelihood_with_censoring.json @@ -901,7 +901,7 @@ 0.0 ], "p_value_breslow": [ - 0.02 + 0.25 ], "p_value_cox_breslow": [ 0.175 diff --git a/test/kombine/test_discrete_class_observable.py b/test/kombine/test_discrete_class_observable.py index 994e045..1caaa4d 100644 --- a/test/kombine/test_discrete_class_observable.py +++ b/test/kombine/test_discrete_class_observable.py @@ -7,9 +7,11 @@ import tempfile import numpy as np +import scipy.stats import kombine.datacard from kombine.datacard import DiscreteClassObservable +from kombine.kaplan_meier_p_value_MINLP import MINLPforKMPValue def test_discrete_class_probability_in_range(): @@ -65,3 +67,184 @@ def test_discrete_class_datacard_parsing(): assert np.isclose(patients[1].observable.probability_in_range(0.0, 1.0), 0.7) assert isinstance(patients[2].observable, DiscreteClassObservable) assert np.isclose(patients[2].observable.probability_in_range(1.0, 2.0), 0.5) + + +def test_discrete_class_two_group_penalty_uses_class_bins(): + """ + Two-group NLL penalties must compare class bins, not the integer cut. + + Evaluating NLL only at threshold=1 would give class-1 patients a zero + penalty for both groups, so KoMbine could reassign them for free. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + pathlib.Path(__file__).parent / "datacards" / "simple_examples" + / "discrete_classes_hr_example_small.txt" + ) + hr_calc = datacard.km_hazard_ratio( + parameter_threshold=1.0, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + penalties = hr_calc.nll_penalty_for_patient_in_range + flip_cost = -math.log(0.05) + math.log(0.95) + for i, patient in enumerate(hr_calc.all_patients): + observed_high = patient.observed_parameter >= 1.0 + if observed_high: + np.testing.assert_allclose(penalties[i, 0], flip_cost, rtol=1e-10) + np.testing.assert_allclose(penalties[i, 1], 0.0, atol=1e-12) + else: + np.testing.assert_allclose(penalties[i, 0], 0.0, atol=1e-12) + np.testing.assert_allclose(penalties[i, 1], flip_cost, rtol=1e-10) + assert np.all(np.isposinf(hr_calc.nll_penalty_for_unassigned)) + + +def test_two_group_neither_bin_uses_tail_nll(): + """ + Exclusion must cost the tail NLL, not a group-to-group flip. + + With three classes and finite bounds, neither is class 2. That cost + equals m_neither - m_best, which need not be flip_cost or 2*flip_cost. + """ + obs_low = DiscreteClassObservable(class_probs=[0.90, 0.05, 0.05]) + patient = obs_low.patient_nll(time=1.0, censored=False, systematics=None) + calc = MINLPforKMPValue( + [patient], + parameter_min=0.0, + parameter_threshold=1.0, + parameter_max=2.0, + ) + m_low, m_high, m_neither = calc.assignment_nlls[0] + np.testing.assert_allclose(m_low, -math.log(0.90), rtol=1e-10) + np.testing.assert_allclose(m_high, -math.log(0.05), rtol=1e-10) + np.testing.assert_allclose(m_neither, -math.log(0.05), rtol=1e-10) + np.testing.assert_allclose(calc.nll_penalty_for_patient_in_range[0, 0], 0.0, atol=1e-12) + np.testing.assert_allclose( + calc.nll_penalty_for_patient_in_range[0, 1], + m_high - m_low, + rtol=1e-10, + ) + np.testing.assert_allclose( + calc.nll_penalty_for_unassigned[0], + m_neither - m_low, + rtol=1e-10, + ) + flip_to_high = m_high - m_low + assert not np.isclose(calc.nll_penalty_for_unassigned[0], 2.0 * flip_to_high) + + +def test_discrete_class_hazard_ratio_stays_near_hard_labels(): + """ + At e=0.05, KoMbine's HR should stay near the hard-label Cox MLE. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + pathlib.Path(__file__).parent / "datacards" / "simple_examples" + / "discrete_classes_hr_example_small.txt" + ) + hr_calc = datacard.km_hazard_ratio( + parameter_threshold=1.0, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + hard_hr, _, _, _ = hr_calc.hazard_ratio_confidence_interval( + cox_only=True, + confidence_level=0.95, + hazard_ratio_min=0.01, + hazard_ratio_max=100.0, + ) + kombine_hr, lower, upper, _ = hr_calc.hazard_ratio_confidence_interval( + cox_only=False, + confidence_level=0.95, + hazard_ratio_min=0.01, + hazard_ratio_max=100.0, + ) + assert 0.5 < kombine_hr < 10.0, ( + f"Expected a finite HR near the hard-label fit, got {kombine_hr}" + ) + assert abs(np.log(kombine_hr) - np.log(hard_hr)) < 0.5, ( + f"KoMbine HR {kombine_hr:.3f} drifted from hard-label {hard_hr:.3f}" + ) + assert lower < kombine_hr < upper + + +def test_permutation_pvalue_agrees_with_chi2_at_small_e(): + """ + At e=0.05 assignments do not move, so permutation p stays large. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + pathlib.Path(__file__).parent / "datacards" / "simple_examples" + / "discrete_classes_hr_example_small.txt" + ) + calc = datacard.km_p_value( + parameter_threshold=1.0, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + p_perm, result_null, result_alt = calc.solve_and_pvalue( + cox_only=False, + n_permutations=19, + rng=0, + ) + chi2_p = scipy.stats.chi2.sf(result_null.x - result_alt.x, 1) + assert chi2_p > 0.1, f"Expected a large chi2 p at e=0.05, got {chi2_p}" + assert p_perm > 0.1, f"Expected a large permutation p at e=0.05, got {p_perm}" + + +def test_permutation_pvalue_not_tiny_when_chi2_is(): + """ + At large e the profile LRT chi2 p-value is tiny, but the permutation + p-value must not treat reassignment search as extra evidence. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + pathlib.Path(__file__).parent / "datacards" / "simple_examples" + / "discrete_classes_hr_example_very_large.txt" + ) + calc = datacard.km_p_value( + parameter_threshold=1.0, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + p_perm, result_null, result_alt = calc.solve_and_pvalue( + cox_only=False, + n_permutations=39, + rng=0, + ) + chi2_p = scipy.stats.chi2.sf(result_null.x - result_alt.x, 1) + assert chi2_p < 0.01, f"Expected a tiny chi2 p at e=0.40, got {chi2_p}" + assert p_perm > 0.05, f"Permutation p should not be tiny, got {p_perm}" + p_cox, _, _ = calc.solve_and_pvalue(cox_only=True) + assert p_cox > 0.1, f"Cox-only chi2 p should stay near the hard-label LRT, got {p_cox}" + + +def test_solve_and_pvalue_rejects_negative_permutations(): + """ + n_permutations must be non-negative. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + pathlib.Path(__file__).parent / "datacards" / "simple_examples" + / "discrete_classes_hr_example_small.txt" + ) + calc = datacard.km_p_value( + parameter_threshold=1.0, + parameter_min=-np.inf, + parameter_max=np.inf, + ) + try: + calc.solve_and_pvalue(n_permutations=-1) + except ValueError as exc: + assert "n_permutations" in str(exc) + else: + raise AssertionError("Expected ValueError for n_permutations=-1") + + +if __name__ == "__main__": + test_discrete_class_probability_in_range() + test_discrete_class_observed_parameter_tie_break() + test_discrete_class_patient_nll() + test_discrete_class_datacard_parsing() + test_discrete_class_two_group_penalty_uses_class_bins() + test_two_group_neither_bin_uses_tail_nll() + test_discrete_class_hazard_ratio_stays_near_hard_labels() + test_permutation_pvalue_agrees_with_chi2_at_small_e() + test_permutation_pvalue_not_tiny_when_chi2_is() + test_solve_and_pvalue_rejects_negative_permutations() + print("[SUCCESS] Discrete class observable tests passed") diff --git a/test/kombine/test_discrete_optimization.py b/test/kombine/test_discrete_optimization.py index af24444..f72b238 100644 --- a/test/kombine/test_discrete_optimization.py +++ b/test/kombine/test_discrete_optimization.py @@ -12,7 +12,10 @@ minimize_discrete_single_minimum, binary_search_sign_change, cached_level_crossings, + feasibility_assisted_level_crossings, + OracleCostTracker, ) +from kombine.kaplan_meier_MINLP import GurobiWorkStats, MINLPForKM class TestMinimizeDiscreteSingleMinimum(unittest.TestCase): """ @@ -539,5 +542,232 @@ def f_loose(x): self.assertLess(len(loose), len(tight)) +class TestFeasibilityAssistedLevelCrossings(unittest.TestCase): + """Tests for feasibility_assisted_level_crossings.""" + + @staticmethod + def _quadratic(x: float) -> float: + return (x - 0.5) ** 2 + + def test_oracle_matches_cached_crossings(self): + """Boolean oracle + polish matches numeric cached_level_crossings.""" + levels = [0.01, 0.04] + + def oracle(p: float, level: float): + return "inside" if self._quadratic(p) <= level else "outside" + + for outer, inner in ((0.0, 0.5), (1.0, 0.5)): + expected = cached_level_crossings( + self._quadratic, outer, inner, levels, xtol=1e-8, rtol=1e-8 + ) + got = feasibility_assisted_level_crossings( + self._quadratic, oracle, outer, inner, levels, xtol=1e-8, rtol=1e-8, + cost_biased_bisection=False, + ) + np.testing.assert_allclose(got, expected, atol=1e-6) + + def test_unknown_falls_back_to_value_func(self): + """unknown status uses value_func for the sign test.""" + level = 0.04 + + def always_unknown(_p: float, _level: float): + return "unknown" + + got = feasibility_assisted_level_crossings( + self._quadratic, always_unknown, 0.0, 0.5, [level], xtol=1e-8, rtol=1e-8 + ) + expected = cached_level_crossings( + self._quadratic, 0.0, 0.5, [level], xtol=1e-8, rtol=1e-8 + ) + np.testing.assert_allclose(got, expected, atol=1e-6) + + def test_outer_inside_returns_boundary(self): + """Outer endpoint already inside the cut returns that endpoint.""" + def oracle(p: float, level: float): + return "inside" if self._quadratic(p) <= level else "outside" + + result = feasibility_assisted_level_crossings( + self._quadratic, oracle, 0.0, 0.5, [0.4], xtol=1e-6, rtol=1e-6 + ) + self.assertEqual(result, [0.0]) + + +class TestOracleCostTracker(unittest.TestCase): + """Tests for cost-biased bisection probe fraction estimation.""" + + def test_initial_probe_fraction_is_prior(self): + """Unupdated tracker returns the prior probe fraction.""" + tracker = OracleCostTracker(r_prior=0.5, kappa=4.0, epsilon_floor=0.0) + self.assertAlmostEqual(tracker.probe_fraction(), 0.5) + + def test_inside_cheap_drives_r_below_half(self): + """Cheap inside / expensive outside pulls r below 1/2.""" + tracker = OracleCostTracker(r_prior=0.5, kappa=4.0, epsilon_floor=0.01) + for _ in range(5): + tracker.record("inside", 1.0) + for _ in range(5): + tracker.record("outside", 10.0) + self.assertLess(tracker.probe_fraction(), 0.5) + + def test_reversed_costs_drive_r_above_half(self): + """Expensive inside / cheap outside pulls r above 1/2.""" + tracker = OracleCostTracker(r_prior=0.5, kappa=4.0, epsilon_floor=0.01) + for _ in range(5): + tracker.record("inside", 10.0) + for _ in range(5): + tracker.record("outside", 1.0) + self.assertGreater(tracker.probe_fraction(), 0.5) + + def test_stall_outlier_does_not_poison_r(self): + """A single huge inside time is winsorized and does not pin r.""" + tracker = OracleCostTracker(r_prior=0.5, kappa=4.0, epsilon_floor=0.05) + for _ in range(4): + tracker.record("inside", 1.0) + tracker.record("inside", 1_000.0) + for _ in range(4): + tracker.record("outside", 10.0) + r = tracker.probe_fraction() + self.assertGreaterEqual(r, 0.05) + self.assertLessEqual(r, 0.95) + self.assertLess(r, 0.5) + + def test_only_inside_early_stays_near_prior(self): + """Few inside-only samples keep r in a wide band around the prior.""" + tracker = OracleCostTracker(r_prior=0.5, kappa=4.0, epsilon_floor=0.05) + for _ in range(3): + tracker.record("inside", 1.0) + r = tracker.probe_fraction() + self.assertGreater(r, 0.2) + self.assertLess(r, 0.8) + + +class TestCostBiasedFeasibilityCrossings(unittest.TestCase): + """Integration tests for timed oracle bisection.""" + + @staticmethod + def _quadratic(x: float) -> float: + return (x - 0.5) ** 2 + + def test_cost_biased_matches_midpoint_on_symmetric_oracle(self): + """Instant symmetric oracle: biased bisection matches midpoint.""" + level = 0.04 + + def oracle(p: float, lvl: float): + return "inside" if self._quadratic(p) <= lvl else "outside" + + kwargs = { + "value_func": self._quadratic, + "sign_oracle": oracle, + "x_outer": 0.0, + "x_inner": 0.5, + "levels": [level], + "xtol": 1e-8, + "rtol": 1e-8, + } + midpoint = feasibility_assisted_level_crossings( + **kwargs, cost_biased_bisection=False, + ) + biased = feasibility_assisted_level_crossings( + **kwargs, cost_biased_bisection=True, + ) + np.testing.assert_allclose(biased, midpoint, atol=1e-6) + + def test_cost_biased_fewer_outside_on_asymmetric_timed_oracle(self): + """Oracle with costlier outside proofs; biased bisection should call outside less.""" + level = 0.04 + outside_calls = {"midpoint": 0, "biased": 0} + + def make_oracle(bucket: str): + def oracle(p: float, lvl: float): + inside = self._quadratic(p) <= lvl + # Explicit sleep-immune costs (not wall-clock). + cost = 1.0 if inside else 10.0 + if not inside: + outside_calls[bucket] += 1 + return ("inside" if inside else "outside"), cost + return oracle + + kwargs = { + "value_func": self._quadratic, + "x_outer": 0.0, + "x_inner": 0.5, + "levels": [level], + "xtol": 1e-6, + "rtol": 1e-6, + } + feasibility_assisted_level_crossings( + sign_oracle=make_oracle("midpoint"), + cost_biased_bisection=False, + **kwargs, + ) + feasibility_assisted_level_crossings( + sign_oracle=make_oracle("biased"), + cost_biased_bisection=True, + **kwargs, + ) + self.assertLessEqual(outside_calls["biased"], outside_calls["midpoint"]) + + def test_outer_oracle_outside_recorded_at_endpoint(self): + """Outer endpoint marked outside is stored in outer_oracle.""" + level = 0.04 + outer_oracle: list = [] + feasibility_assisted_level_crossings( + value_func=self._quadratic, + sign_oracle=lambda p, lvl: ( + "outside" if abs(p) < 1e-5 else ( + "inside" if self._quadratic(p) <= lvl else "outside" + ) + ), + x_outer=0.0, + x_inner=0.5, + levels=[level], + xtol=1e-8, + rtol=1e-8, + cost_biased_bisection=False, + outer_oracle=outer_oracle, + ) + self.assertEqual(outer_oracle, ["outside"]) + + +class TestGurobiWorkStats(unittest.TestCase): + """Tests for Gurobi Work accumulation helpers (no Gurobi license needed).""" + + def test_work_stats_reset_and_total(self): + """total_work sums both sides; reset clears counters.""" + stats = GurobiWorkStats( + oracle_work=3.0, + oracle_calls=2, + oracle_outside_calls=1, + minimize_work=7.0, + minimize_calls=4, + ) + self.assertAlmostEqual(stats.total_work, 10.0) + stats.reset() + self.assertEqual(stats.total_work, 0.0) + self.assertEqual(stats.oracle_calls, 0) + self.assertEqual(stats.minimize_calls, 0) + + def test_record_work_accumulates_oracle_and_minimize(self): + """Oracle and minimize Work accumulate separately, including outside oracles.""" + # pylint: disable=protected-access + class _StubModel: # pylint: disable=too-few-public-methods + def __init__(self, work: float): + self.Work = work + + minlp = MINLPForKM.__new__(MINLPForKM) + minlp.work_stats = GurobiWorkStats() + + minlp._record_work(_StubModel(2.5), kind="oracle", side="inside") + minlp._record_work(_StubModel(18.0), kind="oracle", side="outside") + minlp._record_work(_StubModel(12.0), kind="minimize") + + self.assertAlmostEqual(minlp.work_stats.oracle_work, 20.5) + self.assertEqual(minlp.work_stats.oracle_calls, 2) + self.assertEqual(minlp.work_stats.oracle_outside_calls, 1) + self.assertAlmostEqual(minlp.work_stats.minimize_work, 12.0) + self.assertEqual(minlp.work_stats.minimize_calls, 1) + self.assertAlmostEqual(minlp.work_stats.total_work, 32.5) + + if __name__ == "__main__": unittest.main() diff --git a/test/kombine/test_hazard_ratio.py b/test/kombine/test_hazard_ratio.py index 5959d85..328b41d 100644 --- a/test/kombine/test_hazard_ratio.py +++ b/test/kombine/test_hazard_ratio.py @@ -277,7 +277,10 @@ def test_consistency_with_p_value(): ) # Compute p-value - p_value, _, result_alt = pval_calc.solve_and_pvalue(cox_only=False) + p_value, _, result_alt = pval_calc.solve_and_pvalue( + cox_only=False, + n_permutations=0, + ) # The hazard ratio from p-value calculation should match # the hazard ratio we get from the hazard ratio calculator at the same point @@ -326,7 +329,10 @@ def test_hazard_ratio_at_null(): result_hr_null = hr_calc.compute_2nll_at_hazard_ratio(1.0, cox_only=False) # Get the null hypothesis result from p-value calculator - _, result_null, _ = pval_calc.solve_and_pvalue(cox_only=False) + _, result_null, _ = pval_calc.solve_and_pvalue( + cox_only=False, + n_permutations=0, + ) # The 2NLL values should match assert abs(result_hr_null.x - result_null.x) < 1e-2, \ diff --git a/test/kombine/test_km_likelihood.py b/test/kombine/test_km_likelihood.py index 3ef18c0..001192d 100644 --- a/test/kombine/test_km_likelihood.py +++ b/test/kombine/test_km_likelihood.py @@ -252,7 +252,10 @@ def runtest( ) nominal_hazard_ratio_breslow = km_p_value_minlp_breslow.nominal_hazard_ratio - p_value_breslow, _, _ = km_p_value_minlp_breslow.solve_and_pvalue() + p_value_breslow, _, _ = km_p_value_minlp_breslow.solve_and_pvalue( + n_permutations=19, + rng=0, + ) p_value_cox_breslow, _, _ = km_p_value_minlp_breslow.solve_and_pvalue(cox_only=True) try: diff --git a/test/kombine/test_mc_simex.py b/test/kombine/test_mc_simex.py new file mode 100644 index 0000000..1df4223 --- /dev/null +++ b/test/kombine/test_mc_simex.py @@ -0,0 +1,410 @@ +# pylint: disable=duplicate-code +""" +Test Küchenhoff MC-SIMEX comparison estimators. +""" + +import pathlib +import warnings + +import numpy as np + +import kombine.datacard +from kombine.datacard import ( + DiscreteClassObservable, + FixedObservable, + Patient, + PoissonRatioObservable, +) +from kombine.comparisons import McSimexBase +from .test_yi_correction import generate_synthetic_datacard_with_perfect_classification + +warnings.simplefilter("error") + +here = pathlib.Path(__file__).parent +datacards = here / "datacards" / "simple_examples" + + +def test_mc_simex_poisson_ratio_not_implemented(): + """ + MC-SIMEX should raise for Poisson ratio observables, same as Yi. + """ + obs = PoissonRatioObservable( + numerator=10, + denominator=5, + unique_id_numerator=1, + unique_id_denominator=2, + ) + patient = Patient( + survival_time=1.0, + censored=False, + observable=obs, + ) + simex = McSimexBase([patient]) + + try: + simex.compute_patient_prob_in_range(patient, 0.0, 1.0) + except NotImplementedError: + return + raise AssertionError("Expected NotImplementedError for PoissonRatioObservable") + + +def test_mc_simex_km_matches_nominal_fixed_observable(): + """ + With e_i = 0, MC-SIMEX KM matches nominal KM and Yi. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + datacards / "fixed_hr_example.txt" + ) + threshold = 0.5 + + simex_low = datacard.km_survival_mc_simex( + parameter_min=-np.inf, + parameter_max=threshold, + B=1, + rng=0, + ) + simex_high = datacard.km_survival_mc_simex( + parameter_min=threshold, + parameter_max=np.inf, + B=1, + rng=0, + ) + yi_low = datacard.km_survival_yi( + parameter_min=-np.inf, + parameter_max=threshold, + times_for_plot=simex_low['times_for_plot'], + ) + yi_high = datacard.km_survival_yi( + parameter_min=threshold, + parameter_max=np.inf, + times_for_plot=simex_high['times_for_plot'], + ) + km_low = datacard.km_likelihood( + parameter_min=-np.inf, + parameter_max=threshold, + ) + km_high = datacard.km_likelihood( + parameter_min=threshold, + parameter_max=np.inf, + ) + low_nominal = km_low.nominalkm.survival_probabilities( + times_for_plot=simex_low['times_for_plot'] + ) + high_nominal = km_high.nominalkm.survival_probabilities( + times_for_plot=simex_high['times_for_plot'] + ) + + np.testing.assert_allclose( + simex_low['survival_probabilities'], + low_nominal, + atol=1e-10, + rtol=1e-10, + ) + np.testing.assert_allclose( + simex_high['survival_probabilities'], + high_nominal, + atol=1e-10, + rtol=1e-10, + ) + np.testing.assert_allclose( + simex_low['survival_probabilities'], + yi_low['survival_probabilities'], + atol=1e-10, + rtol=1e-10, + ) + np.testing.assert_allclose( + simex_high['survival_probabilities'], + yi_high['survival_probabilities'], + atol=1e-10, + rtol=1e-10, + ) + assert simex_low['method'] == 'mc_simex' + + +def test_mc_simex_logrank_and_hr_match_hard_labels(): + """ + With e_i = 0, logrank and HR match cox_only / hard-label estimators. + """ + datacard = generate_synthetic_datacard_with_perfect_classification( + target_hr=2.5, + n_patients_per_group=40, + random_seed=123, + ) + threshold = 100.0 + + p_standard = datacard.km_p_value_logrank( + parameter_threshold=threshold, + cox_only=True, + ) + p_simex = datacard.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + B=1, + rng=0, + ) + rel_diff = abs(p_standard - p_simex['p_value']) / max(p_standard, 1e-300) + assert rel_diff < 0.1, ( + f"MC-SIMEX logrank differs from cox_only under perfect classification: " + f"standard={p_standard:.4f}, SIMEX={p_simex['p_value']:.4f}" + ) + + simex = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=1, + rng=0, + ) + estimate = simex.estimate_hazard_ratio() + yi_hrs = np.logspace(-1, 1, 41) + yi_2nlls = [ + datacard.km_hazard_ratio_yi( + parameter_threshold=threshold, + hazard_ratio=float(hr), + ).x + for hr in yi_hrs + ] + yi_best = float(yi_hrs[int(np.argmin(yi_2nlls))]) + assert np.isclose(estimate['hazard_ratio'], yi_best, rtol=0.15), ( + f"MC-SIMEX HR {estimate['hazard_ratio']:.3f} should match Yi/naive " + f"{yi_best:.3f} when e_i = 0" + ) + + +def test_mc_simex_discrete_small_error_moves_away_from_one(): + """ + At e = 0.05, SIMEX HR moves away from 1 relative to naive hard labels. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + datacards / "discrete_classes_hr_example_small.txt" + ) + threshold = 0.5001 + hard_label_card = kombine.datacard.Datacard([ + Patient( + survival_time=patient.time, + censored=patient.censored, + observable=FixedObservable(patient.observed_parameter), + ) + for patient in datacard.patients + ]) + naive_hr = hard_label_card.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=1, + rng=0, + ).estimate_hazard_ratio()['hazard_ratio'] + + simex = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=200, + rng=0, + ) + simex_estimate = simex.estimate_hazard_ratio() + + yi_hrs = np.logspace(-1, 1, 41) + yi_2nlls = [ + datacard.km_hazard_ratio_yi( + parameter_threshold=threshold, + hazard_ratio=float(hr), + ).x + for hr in yi_hrs + ] + yi_best = float(yi_hrs[int(np.argmin(yi_2nlls))]) + + naive_log = abs(np.log(naive_hr)) + simex_log = abs(np.log(simex_estimate['hazard_ratio'])) + assert simex_log > naive_log, ( + f"SIMEX should move away from 1 relative to naive: " + f"naive={naive_hr:.3f}, " + f"SIMEX={simex_estimate['hazard_ratio']:.3f}" + ) + assert not np.isclose(simex_estimate['hazard_ratio'], yi_best, rtol=1e-3), ( + "MC-SIMEX point HR should not be identical to Yi weights" + ) + + +def _high_noise_discrete_datacard() -> kombine.datacard.Datacard: + baseline = kombine.datacard.Datacard.parse_datacard( + datacards / "discrete_classes_hr_example_small.txt" + ) + # Draw observed labels from the e ~ 0.5 misclassification process so the + # naive (and therefore SIMEX) hard-label HR is already near 1. + label_rng = np.random.default_rng(0) + patients = [] + for patient in baseline.patients: + observed_high = bool(label_rng.random() < 0.5) + if observed_high: + probs = [0.49, 0.51] + else: + probs = [0.51, 0.49] + patients.append(Patient( + survival_time=patient.time, + censored=patient.censored, + observable=DiscreteClassObservable(class_probs=probs), + )) + return kombine.datacard.Datacard(patients) + + +def test_mc_simex_high_noise_hr_near_one_finite_wald(): + """ + When e_i -> 0.5, the point HR is near 1 and the Wald CI stays finite. + """ + datacard = _high_noise_discrete_datacard() + simex = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=0.5001, + B=80, + rng=0, + ) + estimate = simex.estimate_hazard_ratio() + assert 0.25 < estimate['hazard_ratio'] < 4.0, ( + f"Expected HR near 1 under e~0.5, got {estimate['hazard_ratio']:.3f}" + ) + assert np.isfinite(estimate['ci_lower']) + assert np.isfinite(estimate['ci_upper']) + assert estimate['ci_lower'] > 0.0 + assert estimate['ci_upper'] < np.inf + assert estimate['ci_lower'] < estimate['ci_upper'] + + +def test_mc_simex_rng_reproducible(): + """ + The same rng seed should reproduce logrank and HR results. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + datacards / "discrete_classes_hr_example_small.txt" + ) + threshold = 0.5001 + first = datacard.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + B=40, + rng=7, + ) + second = datacard.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + B=40, + rng=7, + ) + third = datacard.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + B=40, + rng=8, + ) + assert first['p_value'] == second['p_value'] + assert first['logrank_statistic'] == second['logrank_statistic'] + assert first['p_value'] != third['p_value'] + + hr_a = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=40, + rng=7, + ).estimate_hazard_ratio() + hr_b = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=40, + rng=7, + ).estimate_hazard_ratio() + assert hr_a['hazard_ratio'] == hr_b['hazard_ratio'] + assert hr_a['se_log_hazard_ratio'] == hr_b['se_log_hazard_ratio'] + + +def test_mc_simex_return_structure(): + """ + Return dicts / Wald 2NLL objects should have the documented keys. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + datacards / "fixed_hr_example.txt" + ) + threshold = 0.5 + + km = datacard.km_survival_mc_simex( + parameter_min=-np.inf, + parameter_max=threshold, + B=1, + rng=0, + ) + for key in ( + 'survival_probabilities', + 'times_for_plot', + 'death_times', + 'method', + 'parameter_min', + 'parameter_max', + 'lambda_grid', + 'B', + ): + assert key in km, f"Missing KM key {key}" + assert km['method'] == 'mc_simex' + assert km['survival_probabilities'][0] == 1.0 + + logrank = datacard.km_p_value_logrank_mc_simex( + parameter_threshold=threshold, + B=1, + rng=0, + ) + for key in ( + 'p_value', + 'logrank_statistic', + 'n_low_observed', + 'n_high_observed', + 'method', + 'lambda_grid', + 'B', + ): + assert key in logrank, f"Missing logrank key {key}" + assert 0.0 <= logrank['p_value'] <= 1.0 + assert logrank['logrank_statistic'] >= 0.0 + + simex = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + B=1, + rng=0, + ) + estimate = simex.estimate_hazard_ratio() + for key in ( + 'hazard_ratio', + 'log_hazard_ratio', + 'se_log_hazard_ratio', + 'ci_lower', + 'ci_upper', + 'method', + ): + assert key in estimate, f"Missing HR key {key}" + assert estimate['method'] == 'mc_simex' + assert estimate['ci_lower'] < estimate['hazard_ratio'] < estimate['ci_upper'] + + wald = simex.compute_2nll_at_hazard_ratio(estimate['hazard_ratio']) + assert wald.method == 'mc_simex_wald' + assert wald.x == 0.0 or abs(wald.x) < 1e-12 + wald_edge = estimate['hazard_ratio'] * np.exp(1.96 * estimate['se_log_hazard_ratio']) + off_mle = simex.compute_2nll_at_hazard_ratio(wald_edge) + assert np.isclose(off_mle.x, 1.96 ** 2, rtol=1e-6) + + +if __name__ == "__main__": + print("Running MC-SIMEX tests...") + + print("Test 1: Poisson ratio not implemented...") + test_mc_simex_poisson_ratio_not_implemented() + print("[PASS] Poisson ratio not implemented") + + print("Test 2: KM matches nominal / Yi when e_i = 0...") + test_mc_simex_km_matches_nominal_fixed_observable() + print("[PASS] KM matches nominal / Yi when e_i = 0") + + print("Test 3: logrank and HR match hard labels...") + test_mc_simex_logrank_and_hr_match_hard_labels() + print("[PASS] logrank and HR match hard labels") + + print("Test 4: discrete e=0.05 moves away from 1...") + test_mc_simex_discrete_small_error_moves_away_from_one() + print("[PASS] discrete e=0.05 moves away from 1") + + print("Test 5: high-noise HR near 1 with finite Wald CI...") + test_mc_simex_high_noise_hr_near_one_finite_wald() + print("[PASS] high-noise HR near 1 with finite Wald CI") + + print("Test 6: rng reproducibility...") + test_mc_simex_rng_reproducible() + print("[PASS] rng reproducibility") + + print("Test 7: return structure...") + test_mc_simex_return_structure() + print("[PASS] return structure") + + print("\n[SUCCESS] All tests passed!") diff --git a/test/kombine/test_yi_correction.py b/test/kombine/test_yi_correction.py index e807ce9..fdfef81 100644 --- a/test/kombine/test_yi_correction.py +++ b/test/kombine/test_yi_correction.py @@ -1,3 +1,4 @@ +# pylint: disable=duplicate-code """ Test Yi's misclassification correction methods. @@ -10,6 +11,7 @@ import warnings import numpy as np +import scipy.stats import kombine.datacard from kombine.datacard import ( @@ -19,7 +21,7 @@ PoissonRatioObservable, Patient, ) -from kombine.yi_correction import YiCorrectionBase +from kombine.comparisons import YiCorrectionBase, YiCorrectionForCoxPH from kombine.utilities import prob_poisson_density_in_range from ..utility_testing_functions import generate_two_group_datacard_from_hr @@ -479,6 +481,46 @@ def test_yi_kaplan_meier_survival(): print(f" Final survival (high): {result_yi_high['survival_probabilities'][-1]:.4f}") +def test_yi_continuous_hr_matches_naive_breslow_fixed_card(): + """ + On a fixed observable, Yi's continuous Breslow MLE matches MC-SIMEX. + """ + datacard = kombine.datacard.Datacard.parse_datacard( + datacards / "fixed_hr_example.txt" + ) + threshold = 0.5001 + yi_calc = YiCorrectionForCoxPH( + patients=datacard.patients, + parameter_min=-np.inf, + parameter_max=np.inf, + parameter_threshold=threshold, + ) + best_hr_yi, lower_ci, upper_ci, best_fit = yi_calc.hazard_ratio_confidence_interval( + confidence_level=0.95, + hazard_ratio_min=0.01, + hazard_ratio_max=100.0, + ) + simex = datacard.km_hazard_ratio_mc_simex( + parameter_threshold=threshold, + parameter_min=-np.inf, + parameter_max=np.inf, + B=1, + rng=0, + ) + estimate = simex.estimate_hazard_ratio() + np.testing.assert_allclose( + best_hr_yi, + estimate['hazard_ratio'], + rtol=1e-4, + atol=0.0, + ) + chi2_95 = float(scipy.stats.chi2.ppf(0.95, df=1)) + delta_lower = yi_calc.compute_2nll_at_hazard_ratio(lower_ci).x - best_fit.x + delta_upper = yi_calc.compute_2nll_at_hazard_ratio(upper_ci).x - best_fit.x + np.testing.assert_allclose(delta_lower, chi2_95, rtol=1e-3, atol=1e-2) + np.testing.assert_allclose(delta_upper, chi2_95, rtol=1e-3, atol=1e-2) + assert lower_ci < best_hr_yi < upper_ci + if __name__ == "__main__": # Run tests @@ -512,4 +554,8 @@ def test_yi_kaplan_meier_survival(): test_yi_kaplan_meier_survival() print("[PASS] Yi's Kaplan-Meier survival probabilities") + print("Test 8: Continuous Yi HR matches naive Breslow on a fixed card...") + test_yi_continuous_hr_matches_naive_breslow_fixed_card() + print("[PASS] Continuous Yi HR matches naive Breslow on a fixed card") + print("\n[SUCCESS] All tests passed!")