Speed KoMbine KM/HR profiles and cut plot compile ~2.5×

docs/kombine/06_compare_thomas_grunkemeier.md

@@ -249,7 +249,7 @@ def compare_tg_vs_kombine(  # pylint: disable=too-many-arguments,too-many-positi
   datacard_path,
   conf_level=0.95,
   atol_surv=1e-6,
-  atol_ci=3e-5,
+  atol_ci=2e-4,  # brentq on MINLP 2NLL uses xtol/rtol 1e-4
 ):
   """Compare TG (R) and KoMbine (Python), return merged table and summary."""
   kombine_df, datacard = kombine_binomial_only(datacard_path, conf_level=conf_level)

docs/kombine/compile_km_plots.py

@@ -23,6 +23,7 @@
 import re
 import subprocess
 import sys
+import time
 import warnings
 
 import matplotlib.lines
@@ -838,14 +839,22 @@ def main():
     print("TESTING MODE: Using smaller datasets")
     print("=" * 60)
 
+  t_all = time.perf_counter()
+
   if generate_all or args.km_example:
+    t0 = time.perf_counter()
     plot_km_example(testing=args.testing)
+    print(f"  [km_example] elapsed {time.perf_counter() - t0:.1f}s", flush=True)
 
   if generate_all or args.greenwood:
+    t0 = time.perf_counter()
     plot_compare_to_greenwood(testing=args.testing)
+    print(f"  [greenwood] elapsed {time.perf_counter() - t0:.1f}s", flush=True)
 
   if generate_all or args.p_value:
+    t0 = time.perf_counter()
     plot_compare_p_value(testing=args.testing)
+    print(f"  [p_value] elapsed {time.perf_counter() - t0:.1f}s", flush=True)
 
   if generate_all or args.lung:
     # Handle lung plot testing options
@@ -867,10 +876,15 @@ def main():
             f"Testing panels: {', '.join(p for p, t in zip(panel_names, lung_testing) if t)}")
       print("=" * 60)
 
+    t0 = time.perf_counter()
     plot_lung_dataset(testing=lung_testing, survival_type=args.lung_survival_type)
+    print(f"  [lung] elapsed {time.perf_counter() - t0:.1f}s", flush=True)
 
   print("=" * 60)
-  print("Plot generation complete!")
+  print(
+    f"Plot generation complete! total_elapsed={time.perf_counter() - t_all:.1f}s",
+    flush=True,
+  )
 
 
 if __name__ == "__main__":

kombine/discrete_optimization.py

@@ -7,6 +7,7 @@
 import typing
 
 import numpy as np
+import scipy.optimize
 
 from .utilities import InspectableCache
 
@@ -351,3 +352,177 @@ def minimize_discrete_single_minimum( #pylint: disable=too-many-locals, too-many
     print(f"{i_min + left:3d} {candidates[i_min]:6.3f} {values[i_min]:9.5g}")
 
   return candidates[i_min], values[i_min]
+
+
+def _innermost_bracket(
+  cache: dict[float, float],
+  level: float,
+  x_outer: float,
+  x_inner: float,
+) -> tuple[float, float]:
+  """Return ``(last_inside, first_outside)`` for the MLE-connected crossing.
+
+  ``first_outside`` is the cached outside point closest to ``x_inner``.
+  ``last_inside`` is ``x_inner`` (or an inside point still closer to the MLE
+  than that outside point).  Do not walk through other inside points: a
+  MINLP profile can dip back below the cut in a disconnected valley.
+  Endpoints must be in ``cache`` with ``f(x_outer) >= level >= f(x_inner)``.
+  """
+  lo = min(x_outer, x_inner)
+  hi = max(x_outer, x_inner)
+  outside = [x for x, fx in cache.items() if lo <= x <= hi and fx >= level]
+  inside = [x for x, fx in cache.items() if lo <= x <= hi and fx < level]
+  if not outside or not inside:
+    raise ValueError(
+      f"No cached sign change for level {level} on [{x_outer}, {x_inner}]"
+    )
+  if x_outer < x_inner:
+    first_outside = max(outside)
+    inner_side = [x for x in inside if x > first_outside]
+    last_inside = max(inner_side) if inner_side else x_inner
+  else:
+    first_outside = min(outside)
+    inner_side = [x for x in inside if x < first_outside]
+    last_inside = min(inner_side) if inner_side else x_inner
+  if last_inside == first_outside:
+    raise ValueError(
+      f"Cached bracket collapsed for level {level}: {last_inside}"
+    )
+  return last_inside, first_outside
+
+
+def _adjacent_outside_bracket(  # pylint: disable=too-many-arguments, too-many-positional-arguments
+  f_cached: collections.abc.Callable[[float], float],
+  level: float,
+  last_inside: float,
+  first_outside: float,
+  xtol: float,
+  rtol: float,
+) -> tuple[float, float]:
+  """Double from ``last_inside`` toward ``first_outside`` until ``f >= level``.
+
+  Returns a left-to-right pair for ``brentq``.
+  """
+  if _is_close(last_inside, first_outside, xtol, rtol):
+    return min(last_inside, first_outside), max(last_inside, first_outside)
+  direction = 1.0 if first_outside > last_inside else -1.0
+  remaining = abs(first_outside - last_inside)
+  step = min(
+    max(
+      xtol,
+      rtol * max(abs(last_inside), abs(first_outside), 1e-12),
+      0.05 * remaining,
+    ),
+    remaining,
+  )
+  x = last_inside
+  while True:
+    x_next = x + direction * step
+    if (x_next - first_outside) * direction >= 0.0:
+      x_next = first_outside
+    if f_cached(x_next) >= level:
+      return (x, x_next) if x < x_next else (x_next, x)
+    if x_next == first_outside:
+      return (x, x_next) if x < x_next else (x_next, x)
+    x = x_next
+    step = min(step * 2.0, abs(first_outside - x))
+    if step <= 0.0:
+      return (x, first_outside) if x < first_outside else (first_outside, x)
+
+
+def cached_level_crossings(  # pylint: disable=too-many-locals, too-many-arguments, too-many-positional-arguments
+  func: collections.abc.Callable[[float], float],
+  x_outer: float,
+  x_inner: float,
+  levels: collections.abc.Sequence[float],
+  *,
+  xtol: float = 1e-4,
+  rtol: float = 1e-4,
+) -> list[float]:
+  """Find where a 1D profile crosses several levels, sharing evaluations.
+
+  ``func`` should be nonnegative and smaller toward ``x_inner`` (e.g.
+  ``2NLL(p) - 2NLL_min`` on one side of the MLE).  ``levels`` are the
+  target heights (e.g. χ² cuts).  After every new evaluation, each
+  remaining cut walks cached points from ``x_inner`` outward and takes
+  the first exit (closest outside point, last inside point), then doubles
+  from that inside point until the cut is exceeded so a thin near-MLE
+  crossing is not skipped.
+
+  If ``func(x_outer) <= level``, the whole interval is inside that cut and
+  the crossing is ``x_outer``.  Remaining cuts are solved widest-bracket
+  first (typically the outermost level).
+  """
+  if not levels:
+    return []
+  if x_outer == x_inner:
+    return [x_inner] * len(levels)
+
+  cache: dict[float, float] = {}
+
+  def f_cached(x: float) -> float:
+    if x not in cache:
+      cache[x] = float(func(x))
+    return cache[x]
+
+  f_cached(x_outer)
+  f_cached(x_inner)
+
+  results: list[float | None] = [None] * len(levels)
+  pending = list(range(len(levels)))
+
+  def remaining_width(index: int) -> float:
+    level = levels[index]
+    if f_cached(x_outer) <= level:
+      return -1.0
+    if f_cached(x_inner) >= level:
+      return -1.0
+    left, right = _innermost_bracket(cache, level, x_outer, x_inner)
+    return abs(right - left)
+
+  while pending:
+    pending.sort(key=lambda i: (remaining_width(i), levels[i]), reverse=True)
+    index = pending.pop(0)
+    level = levels[index]
+    f_outer = f_cached(x_outer)
+    f_inner = f_cached(x_inner)
+    if f_outer <= level:
+      results[index] = x_outer
+      continue
+    if f_inner >= level:
+      results[index] = x_inner
+      continue
+    last_inside, first_outside = _innermost_bracket(cache, level, x_outer, x_inner)
+    left, right = _adjacent_outside_bracket(
+      f_cached, level, last_inside, first_outside, xtol, rtol,
+    )
+    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
+      results[index] = closer
+      continue
+    g_left = f_cached(left) - level
+    g_right = f_cached(right) - level
+    if g_left == 0.0:
+      results[index] = left
+      continue
+    if g_right == 0.0:
+      results[index] = right
+      continue
+    results[index] = 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,
+      ),
+    )
+
+  crossings: list[float] = []
+  for x in results:
+    if x is None:
+      raise RuntimeError("cached_level_crossings left a level unsolved")
+    crossings.append(x)
+  return crossings