WIP: Trying to enable this with binning

Author: acampove

src/dmu/stats/gof_calculator.py

@@ -2,14 +2,17 @@
 Module holding GofCalculator class
 '''
 from functools import lru_cache
+from typing    import cast
 
 import zfit
 import numpy
 import pandas as pnd
 
 from scipy                  import stats
-from zfit.core.basepdf      import BasePDF   as zpdf
-from zfit.core.parameter    import Parameter as zpar
+from zfit.interface         import ZfitBinnedData
+from zfit.interface         import ZfitData      as zdat
+from zfit.interface         import ZfitPDF       as zpdf
+from zfit.interface         import ZfitParameter as zpar
 from dmu.logging.log_store  import LogStore
 
 log = LogStore.add_logger('dmu:stats:gofcalculator')
@@ -26,14 +29,35 @@ def __init__(self, nll, ndof : int = 10):
         self._pdf     = self._pdf_from_nll()
         self._data_in = self._data_from_nll()
         self._data_np = self._data_np_from_data(self._data_in)
-        self._data_zf = zfit.Data.from_numpy(obs=self._pdf.space, array=self._data_np)
+        self._data_bn = self._data_bn_from_data(self._data_in)
+        self._data_zf = self._get_zdata()
+    # ----------------------
+    def _get_zdata(self) -> zdat|ZfitBinnedData:
+        '''
+        Returns
+        -------------
+        Data usable by zfit to do fits
+        '''
+        if self._data_np is not None:
+            return zfit.Data.from_numpy(obs=self._pdf.space, array=self._data_np)
+
+        if isinstance(self._data_in, ZfitBinnedData):
+            return self._data_in
+
+        raise ValueError('Data is neither binned, nor convertible to numpy array')
     # ---------------------
-    def _data_np_from_data(self, dat) -> numpy.ndarray:
+    def _data_np_from_data(self, dat) -> numpy.ndarray|None:
         if isinstance(dat, numpy.ndarray):
             return dat
 
+        if isinstance(dat, zfit.data.BinnedData):
+            log.info('Input data is binned, cannot convert to numpy')
+            return None
+
         if isinstance(dat, zfit.Data):
-            return zfit.run(zfit.z.unstack_x(dat))
+            arr_val = zfit.run(zfit.z.unstack_x(dat))
+            arr_val = cast(numpy.ndarray, arr_val)
+            return arr_val
 
         if isinstance(dat, pnd.DataFrame):
             return dat.to_numpy()
@@ -44,6 +68,21 @@ def _data_np_from_data(self, dat) -> numpy.ndarray:
 
         data_type = str(type(dat))
         raise ValueError(f'Data is not a numpy array, zfit.Data or pandas.DataFrame, but {data_type}')
+    # ----------------------
+    def _data_bn_from_data(self, dat) -> ZfitBinnedData|None:
+        '''
+        Parameters
+        -------------
+        dat: Zfit data
+
+        Returns
+        -------------
+        ZfitBinned data or None if input is not binned
+        '''
+        if isinstance(dat, ZfitBinnedData):
+            return dat
+
+        return None
     # ---------------------
     def _pdf_from_nll(self) -> zpdf:
         l_model = self._nll.model
@@ -52,7 +91,7 @@ def _pdf_from_nll(self) -> zpdf:
 
         return l_model[0]
     # ---------------------
-    def _data_from_nll(self) -> zpdf:
+    def _data_from_nll(self) -> zdat:
         l_data = self._nll.data
         if len(l_data) != 1:
             raise ValueError('Not found one and only one dataset')
@@ -81,7 +120,9 @@ def _get_binning(self) -> tuple[int, float, float]:
     # ---------------------
     def _get_pdf_bin_contents(self) -> numpy.ndarray:
         nbins, min_x, max_x  = self._get_binning()
-        _, arr_edg = numpy.histogram(self._data_np, bins = nbins, range=(min_x, max_x))
+
+        if self._data_np is not None:
+            _, arr_edg = numpy.histogram(self._data_np, bins = nbins, range=(min_x, max_x))
 
         size = arr_edg.size
 
@@ -139,7 +180,7 @@ def get_gof(self, kind : str) -> float:
         -----------------
         kind: Type of goodness of fit: pvalue, chi2, chi2/ndof
 
-        Returns 
+        Returns
         -----------------
         Goodness of fit of a given kind
         '''

tests/stats/test_gofcalculator.py

@@ -8,6 +8,7 @@
 import numpy
 import pytest
 
+from zfit                     import Data       as zdata
 from dmu.stats.gof_calculator import GofCalculator
 from dmu.logging.log_store    import LogStore
 
@@ -20,6 +21,7 @@ class Data:
     '''
     minimizer = zfit.minimize.Minuit()
     obs       = zfit.Space('x', limits=(-10, 10))
+    obs_bin   = zfit.Space('x', limits=(-10, 10), binning=50)
 #---------------------------------------------
 @pytest.fixture(scope='session', autouse=True)
 def _initialize():
@@ -33,24 +35,43 @@ def _get_model():
 
     return pdf
 # -------------------------------------------
-def _get_data():
+def _get_data() -> zdata:
     data_np = numpy.random.normal(0, 1, size=10000)
     data_zf = zfit.Data.from_numpy(obs=Data.obs, array=data_np)
 
     return data_zf
 # -------------------------------------------
-def _get_nll():
+def _get_nll(binned : bool):
     pdf = _get_model()
     dat = _get_data()
-    nll = zfit.loss.UnbinnedNLL(model=pdf, data=dat)
+
+    if binned:
+        dat = dat.to_binned(space=Data.obs_bin)
+        pdf = pdf.to_binned(space=Data.obs_bin)
+        nll = zfit.loss.BinnedNLL(model=pdf, data=dat)
+    else:
+        nll = zfit.loss.UnbinnedNLL(model=pdf, data=dat)
 
     return nll
 # -------------------------------------------
-def test_simple():
+def test_unbinned():
+    '''
+    Test GofCalculator with unbinned data
+    '''
+    nll = _get_nll(binned=False)
+    res = Data.minimizer.minimize(nll)
+    print(res)
+
+    gcl = GofCalculator(nll, ndof=10)
+    gof = gcl.get_gof(kind='pvalue')
+
+    assert math.isclose(gof, 0.965, abs_tol=0.01)
+# -------------------------------------------
+def test_binned():
     '''
-    Simplest test of GofCalculator
+    Test GofCalculator with binned data
     '''
-    nll = _get_nll()
+    nll = _get_nll(binned=True)
     res = Data.minimizer.minimize(nll)
     print(res)