Add profile-curve scan helpers and vectorize Neyman belts

src/npkit/__init__.py

@@ -9,8 +9,20 @@
 from .observables import Observable, ObservableSet, Params
 from .measurements import Combination
 from .likelihood import GaussianModel, GaussianLikelihood
-from .stats import fit_mle, q_profile
-from .neyman import Belt, build_belt, invert_belt, check_coverage
+from .stats import (
+    fit_mle,
+    profile_curve_from_grid,
+    profile_curve_from_likelihood,
+    q_profile,
+)
+from .neyman import (
+    Belt,
+    build_belt,
+    build_belts_from_grid,
+    check_coverage,
+    invert_belt,
+    invert_belt_from_curve,
+)
 
 __all__ = [
     "Observable",
@@ -20,9 +32,13 @@
     "GaussianModel",
     "GaussianLikelihood",
     "fit_mle",
+    "profile_curve_from_grid",
+    "profile_curve_from_likelihood",
     "q_profile",
     "Belt",
     "build_belt",
+    "build_belts_from_grid",
+    "invert_belt_from_curve",
     "invert_belt",
     "check_coverage",
 ]

src/npkit/likelihood.py

@@ -78,13 +78,30 @@ def __post_init__(self) -> None:
             raise ValueError("covariance must be positive-definite")
         self._logdet = float(logdet)
 
-    def simulate(self, params: Params, rng: np.random.Generator) -> NDArray[np.float64]:
+    @property
+    def covariance_matrix(self) -> NDArray[np.float64]:
+        assert self._cov is not None
+        return self._cov
+
+    @property
+    def inverse_covariance(self) -> NDArray[np.float64]:
+        return self._cov_inv
+
+    def simulate(
+        self,
+        params: Params,
+        rng: np.random.Generator,
+        size: int | tuple[int, ...] | None = None,
+    ) -> NDArray[np.float64]:
         """
-        Draw one pseudo-experiment vector y ~ N(μ(params), V).
+        Draw pseudo-experiments y ~ N(μ(params), V).
+
+        If `size` is None, return one vector with shape (n,).
+        If `size` is an int or tuple, return an array with shape `size + (n,)`.
         """
         mean = self.obs.predict_vector(params)
         assert self._cov is not None  # for type-checkers
-        y = rng.multivariate_normal(mean=mean, cov=self._cov)
+        y = rng.multivariate_normal(mean=mean, cov=self._cov, size=size)
         return cast(NDArray[np.float64], np.asarray(y, dtype=float))
 
     def likelihood(self, data: Combination) -> "GaussianLikelihood":

src/npkit/neyman.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Callable
 
@@ -36,33 +37,116 @@ def build_belt(
 
     For each grid point C:
       1) Generate y_toy ~ N(μ(C), V)
-      2) Compute q(value=C) on y_toy via profile-likelihood
+      2) Compute q(value=C) on y_toy by scanning the 1D NLL curve
       3) Take the (1 - alpha)-quantile over toys as qcrit[C]
     """
-    m = grid.size
-    qcrit = np.empty(m, dtype=float)
+    _ = like_builder  # kept for API compatibility; the batched path does not need it.
 
-    for i, c in enumerate(grid):
-        q_vals = np.empty(n_toys, dtype=float)
-        for t in range(n_toys):
-            y = model.simulate({**start, param: float(c)}, rng=rng)
+    grid_arr = np.asarray(grid, dtype=float)
+    means_grid = np.asarray(
+        [model.obs.predict_vector({**start, param: float(c)}) for c in grid_arr],
+        dtype=float,
+    )
+    if means_grid.ndim == 1:
+        means_grid = means_grid[:, None]
+
+    cov_inv = model.inverse_covariance
+    qcrit = np.empty(grid_arr.size, dtype=float)
+
+    for i, c in enumerate(grid_arr):
+        y = np.asarray(
+            model.simulate({**start, param: float(c)}, rng=rng, size=n_toys),
+            dtype=float,
+        )
+        if y.ndim == 1:
+            y = y[:, None]
+
+        resid = y[:, None, :] - means_grid[None, :, :]
+        q_curve = np.sum((resid @ cov_inv) * resid, axis=-1)
+        q_vals = q_curve[:, i] - np.min(q_curve, axis=1)
+        qcrit[i] = float(np.quantile(q_vals, 1.0 - alpha))
+
+    return Belt(
+        param=param, grid=grid_arr, qcrit=qcrit, alpha=alpha
+    )
+
+
+def build_belts_from_grid(
+    param: str,
+    model: GaussianModel,
+    grid: np.ndarray,
+    n_toys: int,
+    alphas: Sequence[float],
+    rng: np.random.Generator,
+    start: Params,
+) -> tuple[Belt, ...]:
+    """
+    Build several Neyman belts from the same toy ensemble without fitting.
 
-            def fresh_like() -> GaussianLikelihood:
-                # Bind a fresh likelihood to this toy data
-                return GaussianLikelihood(model.obs, y, model.covariance)
+    For each grid point C:
+      1) generate toys at C,
+      2) scan the full 1D NLL curve on the supplied grid,
+      3) infer q(C) from the minimum of that curve,
+      4) reuse the same q samples to extract multiple critical values.
+    """
+    alpha_list = [float(alpha) for alpha in alphas]
+    if not alpha_list:
+        raise ValueError("alphas must contain at least one confidence level.")
 
-            q_vals[t] = q_profile(
+    grid_arr = np.asarray(grid, dtype=float)
+    means_grid = np.asarray(
+        [model.obs.predict_vector({**start, param: float(c)}) for c in grid_arr],
+        dtype=float,
+    )
+    if means_grid.ndim == 1:
+        means_grid = means_grid[:, None]
+
+    cov_inv = model.inverse_covariance
+    q_samples = np.empty((grid_arr.size, n_toys), dtype=float)
+
+    for i, c in enumerate(grid_arr):
+        print(f"Generating toys for {param}={c:.3f} ({i+1}/{grid_arr.size})")
+        y = np.asarray(
+            model.simulate({**start, param: float(c)}, rng=rng, size=n_toys),
+            dtype=float,
+        )
+        if y.ndim == 1:
+            y = y[:, None]
+        resid = y[:, None, :] - means_grid[None, :, :]
+        q_curve = np.sum((resid @ cov_inv) * resid, axis=-1)
+        q_samples[i, :] = q_curve[:, i] - np.min(q_curve, axis=1)
+
+    belts = []
+    for alpha in alpha_list:
+        qcrit = np.quantile(q_samples, 1.0 - alpha, axis=1)
+        belts.append(
+            Belt(
                 param=param,
-                value=float(c),
-                like_builder=fresh_like,
-                start=start,
-                bounds=bounds,
+                grid=grid_arr,
+                qcrit=np.asarray(qcrit, dtype=float),
+                alpha=alpha,
             )
-        qcrit[i] = float(np.quantile(q_vals, 1.0 - alpha))
+        )
+    return tuple(belts)
 
-    return Belt(
-        param=param, grid=np.asarray(grid, dtype=float), qcrit=qcrit, alpha=alpha
-    )
+
+def invert_belt_from_curve(belt: Belt, q_obs: np.ndarray) -> tuple[float, float]:
+    """
+    Invert a belt using a precomputed observed q curve.
+
+    This avoids any likelihood refits and is the right companion to
+    `profile_curve_from_grid`.
+    """
+    q_obs_arr = np.asarray(q_obs, dtype=float)
+    if q_obs_arr.shape != belt.grid.shape:
+        raise ValueError("q_obs must have the same shape as belt.grid.")
+
+    mask = q_obs_arr <= belt.qcrit
+    if not mask.any():
+        return (np.nan, np.nan)
+
+    idx = np.where(mask)[0]
+    return (float(belt.grid[idx[0]]), float(belt.grid[idx[-1]]))
 
 
 def invert_belt(

src/npkit/stats.py

@@ -138,3 +138,47 @@ def constrained_nll(x: np.ndarray) -> float:
 
     q = float(res.fun - nll_min)
     return max(0.0, q)
+
+
+def profile_curve_from_likelihood(
+    param: str,
+    grid: np.ndarray,
+    like: GaussianLikelihood,
+    start: Params,
+) -> tuple[float, float, np.ndarray]:
+    """
+    Evaluate the NLL on a 1D parameter grid without any minimization.
+
+    This is the no-nuisance scan path:
+      - compute NLL at every grid point,
+      - infer the best fit from the minimum of that curve,
+      - shift the curve so q_min = 0.
+    """
+    grid_arr = np.asarray(grid, dtype=float)
+    curve = np.asarray(
+        [like.nll({**start, param: float(value)}) for value in grid_arr],
+        dtype=float,
+    )
+    best_idx = int(np.argmin(curve))
+    nll_min = float(curve[best_idx])
+    q_curve = np.maximum(0.0, curve - nll_min)
+    return float(grid_arr[best_idx]), nll_min, q_curve
+
+
+def profile_curve_from_grid(
+    param: str,
+    grid: np.ndarray,
+    like_builder: Callable[[], GaussianLikelihood],
+    start: Params,
+) -> tuple[float, float, np.ndarray]:
+    """
+    Convenience wrapper around `profile_curve_from_likelihood`.
+
+    A fresh likelihood is built once, then scanned on the supplied grid.
+    """
+    return profile_curve_from_likelihood(
+        param=param,
+        grid=grid,
+        like=like_builder(),
+        start=start,
+    )