From 7d4f544ad54b220b37fe92637c430492b955a5f1 Mon Sep 17 00:00:00 2001
From: Hans Dembinski <HDembinski@users.noreply.github.com>
Date: Fri, 10 Feb 2023 15:54:41 +0100
Subject: [PATCH] prediction for all binned likelihoods (#831)

---
 src/iminuit/cost.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/iminuit/cost.py b/src/iminuit/cost.py
index 36efbd73..93b19aa6 100644
--- a/src/iminuit/cost.py
+++ b/src/iminuit/cost.py
@@ -1002,12 +1002,32 @@ def __init__(self, args, n, xe, verbose, *updater):
     def _ndata(self):
         return np.prod(self._masked.shape[: self._ndim])
 
+    @abc.abstractmethod
+    def _pred(self, args: _tp.Sequence[float]) -> np.ndarray:
+        ...  # pragma: no cover
+
     def _update_cache(self):
         super()._update_cache()
         if self._bztrafo:
             ma = _replace_none(self._mask, ...)
             self._bztrafo = BohmZechTransform(self._data[ma, 0], self._data[ma, 1])
 
+    def prediction(self, args: _tp.Sequence[float]) -> np.ndarray:
+        """
+        Return the bin expectation for the fitted model.
+
+        Parameters
+        ----------
+        args : array-like
+            Parameter values.
+
+        Returns
+        -------
+        np.ndarray
+            Model prediction for each bin.
+        """
+        return self._pred(args)
+
     def visualize(self, args: _tp.Sequence[float]):
         """
         Visualize data and model agreement (requires matplotlib).
@@ -1031,13 +1051,9 @@ def visualize(self, args: _tp.Sequence[float]):
         if self.mask is not None:
             cx = cx[self.mask]
         plt.errorbar(cx, n, ne, fmt="ok")
-        mu = self._pred(args)
+        mu = self.prediction(args)
         plt.stairs(mu, xe, fill=True, color="C0")
 
-    @abc.abstractmethod
-    def _pred(self, args: _tp.Sequence[float]) -> np.ndarray:
-        ...  # pragma: no cover
-
 
 class BinnedCostWithModel(BinnedCost):
     """