generalise pipeline config parser and varbinning to accept selection string also

Author: thehrh

pisa/core/binning.py

@@ -3047,16 +3047,19 @@ class VarBinning(object):
 
     """
     # pylint: enable=line-too-long
-    def __init__(self, binnings, cut_var, name=None, mask=None):
-        
-        assert isinstance(cut_var, OneDimBinning)
-        assert isinstance(binnings, list) and len(binnings) == cut_var.size
+    def __init__(self, binnings, selections, name=None, mask=None):
+
+        assert (isinstance(selections, OneDimBinning) or
+                isinstance(selections, list))
+        assert isinstance(binnings, list) and len(binnings) == len(selections)
         for b in binnings:
             assert isinstance(b, MultiDimBinning)
-            assert cut_var.name not in b.names
+            if isinstance(selections, OneDimBinning):
+                assert selections.name not in b.names
 
         self._binnings = binnings
-        self._cut_var = cut_var
+        self._selections = selections
+        self._nselections = len(selections)
         self._name = name
         self._names = None
 
@@ -3066,9 +3069,14 @@ def binnings(self):
         return self._binnings
 
     @property
-    def cut_var(self):
-        """OneDimBinning : variable for which to use different binnings"""
-        return self._cut_var
+    def selections(self):
+        """list of strs or OneDimBinning : selections for which to use different binnings"""
+        return self._selections
+
+    @property
+    def nselections(self):
+        """int : number of selections with possibly different binnings"""
+        return self._nselections
 
     @property
     def name(self):
@@ -3079,7 +3087,7 @@ def name(self):
     def names(self):
         """list of strings : names of each dimension contained plus cut var"""
         if self._names is None:
-            self._names = [self.cut_var.name]
+            self._names = [self.cut_var.name] #FIXME
             for b in self.binnings:
                 self._names.extend([n for n in b.names if n not in self._names])
         return self._names
@@ -3106,7 +3114,7 @@ def __iter__(self):
         return iter(self._binnings)
 
     def __len__(self):
-        return len(self._binnings)
+        return self._nselections
 
 
 def test_OneDimBinning():

pisa/utils/config_parser.py

@@ -239,6 +239,7 @@
 
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from collections import OrderedDict
+from collections.abc import Mapping
 from io import StringIO
 from os.path import abspath, expanduser, expandvars, isfile, join
 import re
@@ -620,38 +621,68 @@ def parse_pipeline_config(config):
             binning, _ = split(name, sep='.')
 
             bin_split = config['binning'].get(binning + '.split', None)
-            if bin_split is not None: # Use multiple MultiDimBinning(s)
-                bin_split = OneDimBinning(**eval(config.get('binning', binning + '.split')))
-
-                bins = [[] for i in range(bin_split.size)]
+            if bin_split is not None:
+                # User requested split into several event samples with their
+                # own MultiDimBinning definitions.
+                try:
+                    bin_split = eval(bin_split) # pylint: disable=eval-used
+                except:
+                    assert isinstance(bin_split, str)
+                    # Just split original str into individual selection strs
+                    bin_split = split(bin_split)
+                else:
+                    assert isinstance(bin_split, Mapping)
+                    # If input can be parsed as dict, split events according to
+                    # the presumably contained OneDimBinning definition
+                    bin_split = OneDimBinning(**bin_split)
+
+                nselections = len(bin_split)
+                # instantiate the OneDimBinnings corresponding to each selection
+                bins = [[] for i in range(nselections)]
                 for bin_name in order:
                     def_raw = config.get('binning', binning + '.' + bin_name)
-                    kwargs = eval(def_raw)
+                    kwargs = eval(def_raw) # pylint: disable=eval-used
                     if isinstance(kwargs, list):
-                        assert len(kwargs) == bin_split.size
+                        # Dedicated OneDimBinning kwargs for each selection
+                        assert len(kwargs) == nselections
                     else:
-                        kwargs = [kwargs] * bin_split.size
-                    for i in range(bin_split.size):
-                        kw = kwargs[i]
-                        bins[i].append(OneDimBinning(bin_name, **kw))
+                        # Broadcast the universal OneDimBinning kwargs across
+                        # all selections
+                        kwargs = [kwargs] * nselections
+                    for i, kw in enumerate(kwargs):
+                        bins[i].append(OneDimBinning(name=bin_name, **kw))
 
                 mask = config['binning'].get(binning + '.mask', None)
                 if mask is not None:
-                    mask = eval(mask)
+                    mask = eval(mask) # pylint: disable=eval-used
                     if isinstance(mask[0], list):
-                        assert len(mask) == bin_split.size
+                        # Dedicated mask for each selection
+                        assert len(mask) == nselections
                     else:
-                        mask = [mask] * bin_split.size
+                        # Broadcast the universal mask across all selections
+                        mask = [mask] * nselections
                 else:
-                    mask = [mask] * bin_split.size
+                    # No mask for any selection
+                    mask = [None] * nselections
 
                 multibins = []
-                for i in range(bin_split.size):
-                    multibins.append(MultiDimBinning(bins[i], name=binning+f"_{i}", mask=mask[i]))
+                for i in range(nselections):
+                    mb = MultiDimBinning(
+                            dimensions=bins[i],
+                            name=binning+f"_{i}",
+                            mask=mask[i]
+                        )
+                    multibins.append(mb)
 
-                binning_dict[binning] = VarBinning(multibins, bin_split, name=binning)
+                binning_dict[binning] = VarBinning(
+                    binnings=multibins,
+                    selections=bin_split,
+                    name=binning,
+                    mask=mask
+                )
 
-            else: # Use only one MultiDimBinning
+            else:
+                # Requested only one single MultiDimBinning for all events
                 bins = []
                 for bin_name in order:
                     try:
@@ -680,7 +711,7 @@ def parse_pipeline_config(config):
                         )
                         raise
                     try:
-                        bins.append(OneDimBinning(bin_name, **kwargs))
+                        bins.append(OneDimBinning(name=bin_name, **kwargs))
                     except:
                         logging.error(
                             "Failed to instantiate new `OneDimBinning` from '%s'"
@@ -691,9 +722,11 @@ def parse_pipeline_config(config):
                 # Get the bin mask, if there is one
                 mask = config['binning'].get(binning + '.mask', None)
                 if mask is not None :
-                    mask = eval(mask)
+                    mask = eval(mask) # pylint: disable=eval-used
                 # Create the binning object
-                binning_dict[binning] = MultiDimBinning(bins, name=binning, mask=mask)
+                binning_dict[binning] = MultiDimBinning(
+                    dimensions=bins, name=binning, mask=mask
+                )
 
 
     stage_dicts = OrderedDict()

pisa_examples/resources/settings/binning/example.cfg

@@ -75,3 +75,8 @@ reco_var_binning.order = reco_energy, reco_coszen
 reco_var_binning.split = {'name':'pid', 'bin_edges': [-3.,0.,1000.], 'tex':r'{\rm PID}'}
 reco_var_binning.reco_energy = {'num_bins':10, 'is_log':True, 'domain':[5.,100.], 'tex':r'E_{\rm reco}'}
 reco_var_binning.reco_coszen = [{'num_bins':10, 'is_lin':True, 'domain':[-1,1], 'tex':r'\cos{\theta}_{\rm reco}'}, {'num_bins':20, 'is_lin':True, 'domain':[-1,1], 'tex':r'\cos{\theta}_{\rm reco}'}]
+
+reco_var_binning_2.order = reco_energy, reco_coszen
+reco_var_binning_2.split = (true_energy > 10) & (true_coszen > 0), (true_coszen <= 0)
+reco_var_binning_2.reco_energy = {'num_bins':10, 'is_log':True, 'domain':[5.,100.], 'tex':r'E_{\rm reco}'}
+reco_var_binning_2.reco_coszen = [{'num_bins':10, 'is_lin':True, 'domain':[-1,1], 'tex':r'\cos{\theta}_{\rm reco}'}, {'num_bins':20, 'is_lin':True, 'domain':[-1,1], 'tex':r'\cos{\theta}_{\rm reco}'}]