Merge pull request #62 from HEPData/root-support-inf

GraemeWatt · web-flow · commit 8af9bc87aae9 · 2025-04-10T23:57:46.000+01:00
Add support for underflow/overflow bins in ROOT export
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -16,7 +16,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python 3.12
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.12'
     - name: Install pip dependencies
diff --git a/README.rst b/README.rst
@@ -37,7 +37,7 @@ This Python 3 library provides support for converting:
 * Old HepData input format (`sample <https://github.com/HEPData/hepdata-submission/blob/main/examples/oldhepdata/sample.oldhepdata>`_) to `YAML <https://github.com/HEPData/hepdata-submission>`_
 * `YAML <https://github.com/HEPData/hepdata-submission>`_ to:
 
-  * `ROOT <https://root.cern.ch>`_ (tested with v6.32/04)
+  * `ROOT <https://root.cern.ch>`_ (tested with v6.32.04)
   * `YODA <https://yoda.hepforge.org>`_ (tested with v2.1.0)
   * `CSV <https://en.wikipedia.org/wiki/Comma-separated_values>`_
 
diff --git a/hepdata_converter/testsuite/test_rootwriter.py b/hepdata_converter/testsuite/test_rootwriter.py
@@ -78,8 +78,7 @@ def test_simple_parse(self, yaml_full_path, full_root_path):
         f_orig.Close()
 
     @insert_path('yaml_full')
-    @insert_path('root/full.root')
-    def test_th1_parse(self, yaml_full_path, full_root_path):
+    def test_th1_parse(self, yaml_full_path):
         output_file_path = os.path.join(self.current_tmp, 'datafile.root')
         hepdata_converter.convert(yaml_full_path, output_file_path,
                                   options={'output_format': 'root',
@@ -97,3 +96,10 @@ def test_parse_all(self, test_submissions):
                                                'hepdata_doi': '10.17182/hepdata.12345.v1'})
 
             self.assertNotEqual(os.stat(output_file_path).st_size, 0, 'output root file is empty')
+
+    @insert_path('yaml_inf')
+    def test_parse_with_overflows(self, yaml_inf_path):
+        output_file_path = os.path.join(self.current_tmp, 'data-inf.root')
+        hepdata_converter.convert(yaml_inf_path, output_file_path,
+                                  options={'output_format': 'root'})
+        pass
diff --git a/hepdata_converter/version.py b/hepdata_converter/version.py
@@ -1,3 +1,3 @@
 # this file ideally should only contain __version__ declaration, as anything else
 # may break setup.py and PyPI uploads
-__version__ = '0.3.0'
+__version__ = '0.3.1'
diff --git a/hepdata_converter/writers/array_writer.py b/hepdata_converter/writers/array_writer.py
@@ -83,7 +83,7 @@ def match_and_create(cls, independent_variables_map, dependent_variable, depende
             return cls(independent_variables_map, dependent_variable, dependent_variable_index).create_objects()
         return []
 
-    def calculate_total_errors(self):
+    def calculate_total_errors(self, for_tgraph=False):
         is_number_list = self.is_number_var(self.dependent_variable)
         for independent_variable in self.independent_variable_map:
             xerr_minus = []
@@ -93,7 +93,7 @@ def calculate_total_errors(self):
             xval = []
             self.xval.append(xval)
             ArrayWriter.calculate_total_errors(independent_variable, is_number_list,
-                                               xerr_minus, xerr_plus, xval)
+                                               xerr_minus, xerr_plus, xval, {}, for_tgraph)
         ArrayWriter.calculate_total_errors(self.dependent_variable, is_number_list,
                                            self.yerr_minus, self.yerr_plus, self.yval, self.err_breakdown)
 
@@ -158,7 +158,7 @@ def process_error_labels(value):
                         break
 
     @staticmethod
-    def calculate_total_errors(variable, is_number_list, min_errs, max_errs, values, err_breakdown={}):
+    def calculate_total_errors(variable, is_number_list, min_errs, max_errs, values, err_breakdown={}, for_tgraph=False):
         i_numeric = -1  # bin number excluding non-numeric y values
         for i, entry in enumerate(variable['values']):
             if not is_number_list[i]:
@@ -209,10 +209,29 @@ def calculate_total_errors(variable, is_number_list, min_errs, max_errs, values,
                     min_errs.append(0.0)
                     max_errs.append(0.0)
             else:
-                middle_val = (entry['high'] - entry['low']) * 0.5 + entry['low']
-                values.append(middle_val)
-                min_errs.append(middle_val - entry['low'])
-                max_errs.append(entry['high'] - middle_val)
+                if entry['low'] == float('-inf'):  # underflow bin
+                    if for_tgraph:  # zero-width bin centred on upper limit
+                        values.append(entry['high'])
+                        min_errs.append(0.0)
+                        max_errs.append(0.0)
+                    else:  # infinite-width bin centred on -inf
+                        values.append(entry['low'])
+                        min_errs.append(float('inf'))
+                        max_errs.append(float('inf'))
+                elif entry['high'] == float('inf'):  # overflow bin
+                    if for_tgraph:  # zero-width bin centred on lower limit
+                        values.append(entry['low'])
+                        min_errs.append(0.0)
+                        max_errs.append(0.0)
+                    else:  # infinite-width bin centred on +inf
+                        values.append(entry['high'])
+                        min_errs.append(float('inf'))
+                        max_errs.append(float('inf'))
+                else:
+                    middle_val = (entry['high'] - entry['low']) * 0.5 + entry['low']
+                    values.append(middle_val)
+                    min_errs.append(middle_val - entry['low'])
+                    max_errs.append(entry['high'] - middle_val)
 
 
     @classmethod
diff --git a/hepdata_converter/writers/root_writer.py b/hepdata_converter/writers/root_writer.py
@@ -5,6 +5,7 @@
 import array
 import tempfile
 import os
+import math
 from ctypes import c_char_p
 from hepdata_converter.writers.utils import error_value_processor
 
@@ -55,31 +56,36 @@ def _create_empty_hist(self, dependent_var_title, index, yval):
 
         name = "Hist%sD_y%s_e%s" % (self.dim, self.dependent_variable_index + 1, index)
 
-        # order bin values of independent variables
+        # order bin values of independent variables and check if underflow/overflow bins present
         xval_ordered = []
+        underflow = []
+        overflow = []
         for i in range(self.dim):
             xval_ordered.append([])
             xval_ordered[i] = sorted(xval[i])
+            nbinsi = len(xval_ordered[i]) - 1
+            underflow.append(1 if not math.isfinite(xval_ordered[i][0]) else 0)
+            overflow.append(1 if not math.isfinite(xval_ordered[i][nbinsi]) else 0)
 
         if 1 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx)
 
         if 2 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
-            nbinsy = len(xval_ordered[1]) - 1
-            binsy = array.array('d', xval_ordered[1])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
+            nbinsy = len(xval_ordered[1]) - 1 - underflow[1] - overflow[1]
+            binsy = array.array('d', xval_ordered[1][underflow[1]:nbinsy+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx, nbinsy, binsy)
 
         if 3 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
-            nbinsy = len(xval_ordered[1]) - 1
-            binsy = array.array('d', xval_ordered[1])
-            nbinsz = len(xval_ordered[2]) - 1
-            binsz = array.array('d', xval_ordered[2])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
+            nbinsy = len(xval_ordered[1]) - 1 - underflow[1] - overflow[1]
+            binsy = array.array('d', xval_ordered[1][underflow[1]:nbinsy+2])
+            nbinsz = len(xval_ordered[2]) - 1 - underflow[2] - overflow[2]
+            binsz = array.array('d', xval_ordered[2][underflow[2]:nbinsz+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx, nbinsy, binsy, nbinsz, binsz)
 
         for i in range(self.dim):
@@ -105,31 +111,36 @@ def _create_hist(self, xval):
         name = "Hist%sD_y%s" % (self.dim, self.dependent_variable_index + 1)
         args = []
 
-        # order bin values of independent variables
+        # order bin values of independent variables and check if underflow/overflow bins present
         xval_ordered = []
+        underflow = []
+        overflow = []
         for i in range(self.dim):
             xval_ordered.append([])
             xval_ordered[i] = sorted(xval[i])
+            nbinsi = len(xval_ordered[i]) - 1
+            underflow.append(1 if not math.isfinite(xval_ordered[i][0]) else 0)
+            overflow.append(1 if not math.isfinite(xval_ordered[i][nbinsi]) else 0)
 
         if 1 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx)
 
         if 2 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
-            nbinsy = len(xval_ordered[1]) - 1
-            binsy = array.array('d', xval_ordered[1])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
+            nbinsy = len(xval_ordered[1]) - 1 - underflow[1] - overflow[1]
+            binsy = array.array('d', xval_ordered[1][underflow[1]:nbinsy+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx, nbinsy, binsy)
 
         if 3 == self.dim:
-            nbinsx = len(xval_ordered[0]) - 1
-            binsx = array.array('d', xval_ordered[0])
-            nbinsy = len(xval_ordered[1]) - 1
-            binsy = array.array('d', xval_ordered[1])
-            nbinsz = len(xval_ordered[2]) - 1
-            binsz = array.array('d', xval_ordered[2])
+            nbinsx = len(xval_ordered[0]) - 1 - underflow[0] - overflow[0]
+            binsx = array.array('d', xval_ordered[0][underflow[0]:nbinsx+2])
+            nbinsy = len(xval_ordered[1]) - 1 - underflow[1] - overflow[1]
+            binsy = array.array('d', xval_ordered[1][underflow[1]:nbinsy+2])
+            nbinsz = len(xval_ordered[2]) - 1 - underflow[2] - overflow[2]
+            binsz = array.array('d', xval_ordered[2][underflow[2]:nbinsz+2])
             hist = self.get_hist_classes()[self.dim - 1](self.sanitize_name(name), '', nbinsx, binsx, nbinsy, binsy, nbinsz, binsz)
 
         for i in range(self.dim):
@@ -246,6 +257,9 @@ def create_objects(self):
                 if x['high'] not in xval[i]:
                     xval[i].append(x['high'])
 
+        if not any(xval):
+            return []
+
         try:
             hist = self._create_hist(xval)
         except:
@@ -280,7 +294,7 @@ def match(cls, independent_variables_map, dependent_variable):
         return False
 
     def create_objects(self):
-        self.calculate_total_errors()
+        self.calculate_total_errors(for_tgraph=True)
 
         # check that errors are symmetric (within a tolerance to allow for numerical rounding)
         tol = 1e-15
@@ -335,7 +349,7 @@ def match(cls, independent_variables_map, dependent_variable):
         return False
 
     def create_objects(self):
-        self.calculate_total_errors()
+        self.calculate_total_errors(for_tgraph=True)
 
         if len(self.xval[0]):
             graph = ROOTModule.TGraphAsymmErrors(len(self.xval[0]),
diff --git a/setup.py b/setup.py
@@ -33,7 +33,7 @@ def get_version():
 setup(
     name='hepdata-converter',
     version=get_version(),
-    install_requires=['hepdata-validator>=0.3.5'],
+    install_requires=['hepdata-validator>=0.3.6'],
     entry_points={
         'console_scripts': [
             'hepdata-converter = hepdata_converter:main',
