fix: cost function bugs (masked BinnedNLL gradient, pulls masks, visualize)

- BinnedNLL._grad: per-parameter renormalization correction for masked
  bins with >=2 parameters (was a scalar collapsing all parameters)
- BinnedCost._n_err / LeastSquares._pulls: handle index masks correctly
  (~index_array is bitwise NOT) and OR the zero-error mask with the user
  mask instead of replacing it, so zero-error bins yield NaN pulls
- CostSum.visualize: squeeze=False so a single visualizable component
  does not crash; drop the stale **kwargs docstring paragraph
- UnbinnedNLL.scaled_pdf: scale by number of data points (data.shape[-1])
  instead of D*N for multivariate data
- _normalize_output: do not bypass shape validation for non-float outputs
- BinnedNLL: cache total counts in _update_cache (avoid per-call np.sum);
  Template._pred: avoid per-component full-array temporaries for the
  variance floor
- replace abc.abstractproperty with property over abstractmethod

Assisted-by: ClaudeCode:claude-opus-4-8

Author: henryiii

src/iminuit/cost.py

@@ -179,6 +179,16 @@ def _replace_none(x, replacement):
     return x
 
 
+def _exclusion_from_mask(mask: NDArray, n: int) -> NDArray:
+    # Return a boolean array of length n that is True for excluded entries.
+    # mask may be a boolean array or an array of indices of selected entries.
+    if mask.dtype == bool:
+        return ~mask
+    excluded = np.ones(n, dtype=bool)
+    excluded[mask] = False
+    return excluded
+
+
 def chi2(y: ArrayLike, ye: ArrayLike, ym: ArrayLike) -> float:
     """
     Compute (potentially) chi2-distributed cost.
@@ -758,16 +768,14 @@ def visualize(
             Dict that maps an index to dict of keyword arguments. This can be
             used to pass keyword arguments to a visualize method of a component with
             that index.
-        **kwargs :
-            Other keyword arguments are forwarded to all components.
         """
         from matplotlib import pyplot as plt
 
         n = sum(hasattr(comp, "visualize") for comp in self)
 
         fig = plt.gcf()
         fig.set_figwidth(n * fig.get_figwidth() / 1.5)
-        _, ax = plt.subplots(1, n, num=fig.number)
+        _, ax = plt.subplots(1, n, num=fig.number, squeeze=False)
 
         if component_kwargs is None:
             component_kwargs = {}
@@ -777,7 +785,7 @@ def visualize(
             if not hasattr(comp, "visualize"):
                 continue
             kwargs = component_kwargs.get(k, {})
-            plt.sca(ax[i])
+            plt.sca(ax[0, i])
             comp.visualize(cargs, **kwargs)
             i += 1
 
@@ -907,12 +915,14 @@ def __init__(
         self._model_grad = grad
         super().__init__(_model_parameters(model, name), _norm(data), verbose)
 
-    @abc.abstractproperty
+    @property
+    @abc.abstractmethod
     def pdf(self):
         """Get probability density model."""
         ...  # pragma: no cover
 
-    @abc.abstractproperty
+    @property
+    @abc.abstractmethod
     def scaled_pdf(self):
         """Get number density model."""
         ...  # pragma: no cover
@@ -1042,7 +1052,9 @@ def pdf(self):
     @property
     def scaled_pdf(self):
         """Get number density model."""
-        scale = np.prod(self.data.shape)
+        # number of data points: for multivariate data of shape (D, N) this
+        # is N, the last axis; for 1D data of shape (N,) it is also N
+        scale = self.data.shape[-1]
         if self._log:
             return lambda *args: scale * np.exp(self._model(*args))
         return lambda *args: scale * self._model(*args)
@@ -1305,12 +1317,13 @@ class BinnedCost(MaskedCostWithPulls):
     :meta private:
     """
 
-    __slots__ = "_xe", "_ndim", "_bohm_zech_n", "_bohm_zech_s"
+    __slots__ = "_xe", "_ndim", "_bohm_zech_n", "_bohm_zech_s", "_counts_total"
 
     _xe: Union[NDArray, Tuple[NDArray, ...]]
     _ndim: int
     _bohm_zech_n: NDArray
     _bohm_zech_s: Optional[NDArray]
+    _counts_total: float
 
     n = MaskedCost.data
 
@@ -1434,7 +1447,11 @@ def _n_err(self) -> Tuple[NDArray, NDArray]:
         # mask values where error is zero
         ma = err == 0
         if self.mask is not None:
-            ma = ~self.mask
+            # the mask acts on the first dimension; broadcast the per-bin
+            # exclusion across the remaining dimensions before combining
+            excluded = _exclusion_from_mask(self.mask, n.shape[0])
+            excluded = excluded.reshape((-1,) + (1,) * (n.ndim - 1))
+            ma = ma | excluded
         n[ma] = np.nan
         err[ma] = np.nan
         return n, err
@@ -1460,6 +1477,8 @@ def _update_cache(self):
             self._bohm_zech_n = val * s
         else:
             self._bohm_zech_n = n
+        # cache the total number of entries in the unmasked bins
+        self._counts_total = np.sum(self._counts())
 
     def _transformed(self, val: NDArray) -> Tuple[NDArray, NDArray]:
         s = self._bohm_zech_s
@@ -1793,8 +1812,8 @@ def __init__(
         self._model_len = np.prod(self._xe_shape)
 
     def _pred(self, args: Sequence[float]) -> Tuple[NDArray, NDArray]:
-        mu: NDArray = 0  # type:ignore
-        mu_var: NDArray = 0  # type:ignore
+        mu = np.zeros(self._data.shape[: self._ndim])
+        mu_var = np.zeros_like(mu)
         i = 0
         for t1, t2 in self._model_data:
             if isinstance(t1, np.ndarray) and isinstance(t2, np.ndarray):
@@ -1813,7 +1832,9 @@ def _pred(self, args: Sequence[float]) -> Tuple[NDArray, NDArray]:
                 # subtraction, we set negative values to zero
                 d[d < 0] = 0
                 mu += d
-                mu_var += np.ones_like(mu) * 1e-300
+                # add a tiny floor to the variance to avoid exactly zero values;
+                # a scalar add avoids allocating full temporaries
+                mu_var += 1e-300
                 i += t2
             else:  # never arrive here
                 assert False  # pragma: no cover
@@ -1986,7 +2007,7 @@ def _pred(self, args: Sequence[float]) -> NDArray:
         if ma is not None:
             p /= np.sum(p[ma])
         # scale probabilities with total number of entries of unmasked bins in histogram
-        return p * np.sum(self._counts())
+        return p * self._counts_total
 
     def _value(self, args: Sequence[float]) -> float:
         mu = self._pred(args)
@@ -2001,11 +2022,11 @@ def _grad(self, args: Sequence[float]) -> NDArray:
         # normalise probability of remaining bins
         if ma is not None:
             psum = np.sum(p[ma])
-            pg = pg / psum - p * np.sum(pg[:, ma]) / psum**2
+            pg = pg / psum - p * np.sum(pg[:, ma], axis=1)[:, np.newaxis] / psum**2
             p /= psum
         # scale probabilities with total number of entries of unmasked bins in histogram
         n = self._counts()
-        ntot = np.sum(n)
+        ntot = self._counts_total
         mu = p * ntot
         gmu = pg * ntot
         ma = self.mask
@@ -2345,10 +2366,11 @@ def _pulls(self, args: Sequence[float]) -> NDArray:
         ye = self.yerror.copy()
         ym = self.prediction(args)
 
+        ma = ye == 0
         if self.mask is not None:
-            ma = ~self.mask
-            y[ma] = np.nan
-            ye[ma] = np.nan
+            ma = ma | _exclusion_from_mask(self.mask, y.shape[0])
+        y[ma] = np.nan
+        ye[ma] = np.nan
         return (y - ym) / ye
 
     def _pred(self, args: Sequence[float]) -> NDArray:
@@ -2551,7 +2573,7 @@ def _normalize_output(x, kind, *shape, msg=None):
         warnings.warn(msg, PerformanceWarning)
         x = np.array(x)
         if x.dtype.kind != "f":
-            return x.astype(float)
+            x = x.astype(float)
     if x.ndim < len(shape):
         return x.reshape(*shape)
     elif x.shape != shape:

tests/test_cost.py

@@ -384,6 +384,21 @@ def test_UnbinnedNLL_properties(log):
     assert c.verbose == 1
 
 
+def test_UnbinnedNLL_scaled_pdf_2D():
+    # for multivariate data of shape (D, N) the scale must be the number of
+    # data points N, not D * N
+    def model(x_y, mux, muy, sx, sy):
+        return mvnorm(mux, muy, sx, sy).pdf(x_y.T)
+
+    truth = 0.1, 0.2, 0.3, 0.4
+    x, y = mvnorm(*truth).rvs(size=15, random_state=1).T
+    c = UnbinnedNLL((x, y), model)
+
+    assert c.data.shape == (2, 15)
+    expected = 15 * model(c.data, *truth)
+    assert_allclose(c.scaled_pdf(c.data, *truth), expected)
+
+
 @pytest.mark.parametrize("log", (False, True))
 def test_UnbinnedNLL_visualize(log):
     pytest.importorskip("matplotlib")
@@ -643,6 +658,30 @@ def test_BinnedNLL_pulls(binned):
     assert np.nanvar(pulls) == pytest.approx(1, abs=0.2)
 
 
+def test_BinnedNLL_pulls_mask():
+    # masked bins must produce NaN pulls for both boolean and index masks
+    n = np.array([5, 1000, 50, 1])
+    xe = [0, 1, 2, 3, 4]
+    c = BinnedNLL(n, xe, expon_cdf)
+    args = (1.0,)
+
+    full = c.pulls(args)
+    assert not np.any(np.isnan(full))
+
+    # boolean mask: masked bin is NaN, the rest finite
+    c.mask = np.array([True, True, False, True])
+    p = c.pulls(args)
+    assert np.isnan(p[2])
+    assert not np.any(np.isnan(p[[0, 1, 3]]))
+
+    # index mask selecting the same bins must give identical pulls;
+    # bitwise NOT of an index array would poison the wrong bins
+    c.mask = np.array([0, 1, 3])
+    p2 = c.pulls(args)
+    assert np.isnan(p2[2])
+    assert_allclose(p2[[0, 1, 3]], p[[0, 1, 3]])
+
+
 @pytest.mark.parametrize("use_grad", (False, True))
 def test_BinnedNLL_weighted(use_grad):
     xe = np.array([0, 0.2, 0.4, 0.8, 1.5, 10])
@@ -840,6 +879,23 @@ def test_BinnedNLL_mask():
     assert_allclose(c.grad(2), ref(2))
 
 
+def test_BinnedNLL_mask_grad_multipar():
+    # regression test: the masked gradient renormalization correction must be
+    # per-parameter; with >=2 parameters a scalar correction is wrong
+    pytest.importorskip("jacobi")
+    xe = np.linspace(-2, 2, 6)
+    n = np.diff(norm_cdf(xe, 0.1, 1.2)) * 1000
+    c = BinnedNLL(n, xe, norm_cdf, grad=numerical_model_gradient(norm_cdf))
+    c.mask = np.arange(len(n)) != 2
+
+    ref = numerical_cost_gradient(c)
+    # evaluate away from the truth so the gradient is clearly non-zero
+    for args in [(0.0, 1.0), (0.3, 0.8), (-0.2, 1.5)]:
+        g = c.grad(*args)
+        assert np.linalg.norm(g) > 1.0
+        assert_allclose(g, ref(*args), rtol=1e-3)
+
+
 def test_BinnedNLL_properties():
     def cdf(x, a, b):
         return 0
@@ -1457,6 +1513,25 @@ def test_LeastSquares_pulls():
     assert_equal(c.pulls((0, 1)), [10, np.nan])
 
 
+def test_LeastSquares_pulls_mask_index():
+    # an index mask must mask the same bins as the equivalent boolean mask;
+    # bitwise NOT of an index array would poison the wrong bins
+    c = LeastSquares([1, 2, 3], [2, 3, 4], 0.1, line)
+    c.mask = [0, 2]
+    assert_equal(c.pulls((0, 1)), [10, np.nan, 10])
+
+
+def test_LeastSquares_pulls_zero_error():
+    # bins with zero error must yield NaN pulls (documented), not +-inf,
+    # also when inside the user mask
+    c = LeastSquares([1, 2, 3], [2, 3, 4], [0.1, 0.0, 0.1], line)
+    p = c.pulls((0, 1))
+    assert_equal(p, [10, np.nan, 10])
+    c.mask = [True, True, True]
+    p = c.pulls((0, 1))
+    assert_equal(p, [10, np.nan, 10])
+
+
 @pytest.mark.parametrize("use_grad", (False, True))
 def test_CostSum_1(use_grad):
     def model1(x, a):
@@ -1632,6 +1707,18 @@ def test_CostSum_visualize():
     c.visualize((1, 2))
 
 
+def test_CostSum_visualize_single_component():
+    # regression test: a CostSum with exactly one visualizable component must
+    # not crash; subplots(1, 1) returns a bare Axes unless squeeze=False
+    pytest.importorskip("matplotlib")
+    from matplotlib import pyplot as plt
+
+    c = UnbinnedNLL([1.0, 2.0, 3.0], norm_pdf) + 10.0
+    plt.figure()
+    c.visualize((0.0, 1.0))
+    plt.close("all")
+
+
 def test_NormalConstraint_1():
     c1 = NormalConstraint("a", 1, 1.5)
     c2 = NormalConstraint(("a", "b"), (1, 2), (3, 4))
@@ -1791,6 +1878,20 @@ def test_NormalConstraint_pickle():
     assert_equal(c.covariance, c2.covariance)
 
 
+def test_normalize_output_wrong_shape_int():
+    # an integer array with the wrong shape must still raise the descriptive
+    # shape error instead of silently slipping through the float early-return
+    from iminuit.cost import _normalize_output
+
+    x = np.array([1, 2], dtype=int)
+    with pytest.warns(PerformanceWarning):
+        with pytest.raises(
+            ValueError,
+            match=r"output of model has shape \(2,\), but \(3,\) is required",
+        ):
+            _normalize_output(list(x), "model", 3)
+
+
 def test_NormalConstraint_bad_input_1():
     with pytest.raises(ValueError, match="scalar or one-dimensional"):
         NormalConstraint("par", [[[1, 2]]], np.eye(2))