Alternative variable binning approach

pisa/analysis/analysis.py

@@ -486,16 +486,22 @@ def __init__(
             assert hypo_maker is not None, msg
             assert data_dist is not None, msg
             assert metric is not None, msg
+            # this passes through the setter method, but it should just pass through
+            # without actually doing anything
             if hypo_maker.__class__.__name__ == "Detectors":
-                # this passes through the setter method, but it should just pass through
-                # without actually doing anything
                 self.detailed_metric_info = [self.get_detailed_metric_info(
                     data_dist=data_dist[i], hypo_asimov_dist=self.hypo_asimov_dist[i],
                     params=hypo_maker.distribution_makers[i].params, metric=metric[i],
-                    other_metrics=other_metrics, detector_name=hypo_maker.det_names[i], hypo_maker=hypo_maker
+                    other_metrics=other_metrics, detector_name=hypo_maker.det_names[i], hypo_maker=hypo_maker,
                 ) for i in range(len(data_dist))]
-            else: # DistributionMaker object
-                if 'generalized_poisson_llh' == metric[0]:
+            elif isinstance(data_dist, list): # DistributionMaker object with variable binning
+                self.detailed_metric_info = [self.get_detailed_metric_info(
+                    data_dist=data_dist[i], hypo_asimov_dist=self.hypo_asimov_dist[i],
+                    params=hypo_maker.params, metric=metric[0],
+                    other_metrics=other_metrics, detector_name=hypo_maker.detector_name, hypo_maker=hypo_maker,
+                ) for i in range(len(data_dist))]
+            else: # DistributionMaker object with regular binning
+                if metric[0] == 'generalized_poisson_llh':
                     generalized_poisson_dist = hypo_maker.get_outputs(return_sum=False, force_standard_output=False)
                     generalized_poisson_dist = merge_mapsets_together(mapset_list=generalized_poisson_dist)
                 else:
@@ -504,7 +510,7 @@ def __init__(
                 self.detailed_metric_info = self.get_detailed_metric_info(
                     data_dist=data_dist, hypo_asimov_dist=self.hypo_asimov_dist, generalized_poisson_hypo=generalized_poisson_dist,
                     params=hypo_maker.params, metric=metric[0], other_metrics=other_metrics,
-                    detector_name=hypo_maker.detector_name, hypo_maker=hypo_maker
+                    detector_name=hypo_maker.detector_name, hypo_maker=hypo_maker,
                 )
 
     def __getitem__(self, i):
@@ -1006,17 +1012,6 @@ def fit_recursively(
 
         """
 
-        if isinstance(metric, str):
-            metric = [metric]
-
-        # Before starting any fit, check if we already have a perfect match between data and template
-        # This can happen if using pseudodata that was generated with the nominal values for parameters
-        # (which will also be the initial values in the fit) and blah...
-        # If this is the case, don't both to fit and return results right away.
-
-        if isinstance(metric, str):
-            metric = [metric]
-
         # Grab the hypo map
         hypo_asimov_dist = hypo_maker.get_outputs(return_sum=True)
 
@@ -1027,12 +1022,18 @@ def fit_recursively(
             if len(metric) == 1: # One metric for all detectors
                 metric = list(metric) * len(hypo_maker.distribution_makers)
             elif len(metric) != len(hypo_maker.distribution_makers):
-                raise IndexError('Number of defined metrics does not match with number of detectors.')
-        else: # DistributionMaker object
-            assert len(metric) == 1
+                raise IndexError(f"Number of defined metrics does not match with number of detectors.")
+        elif isinstance(hypo_asimov_dist, MapSet) or isinstance(hypo_asimov_dist, list):
+            # DistributionMaker object (list means variable binning)
+            assert len(metric) == 1, f"Only one metric allowed for DistributionMaker"
+        else: 
+            raise NotImplementedError(f"hypo_maker returned output of type {type(hypo_asimov_dist)}")
 
-        # Check if the hypo matches data
-        if hypo_maker.__class__.__name__ != "Detectors" and data_dist.allclose(hypo_asimov_dist) :
+        # Before starting any fit, check if we already have a perfect match between data and template
+        # This can happen if using pseudodata that was generated with the nominal values for parameters
+        # (which will also be the initial values in the fit)
+        # If this is the case, don't both to fit and return results right away.
+        if isinstance(data_dist, MapSet) and data_dist.allclose(hypo_asimov_dist):
 
             msg = 'Initial hypo matches data, no need for fit'
             logging.info(msg)
@@ -2602,11 +2603,17 @@ def _minimizer_callable(self, scaled_param_vals, hypo_maker, data_dist,
                 metric_val = 0
                 for i in range(len(hypo_maker.distribution_makers)):
                     data = data_dist[i].metric_total(expected_values=hypo_asimov_dist[i],
-                                                  metric=metric[i], metric_kwargs=metric_kwargs)
+                                                     metric=metric[i], metric_kwargs=metric_kwargs)
                     metric_val += data
                 priors = hypo_maker.params.priors_penalty(metric=metric[0]) # uses just the "first" metric for prior
                 metric_val += priors
-            else: # DistributionMaker object
+            elif isinstance(hypo_asimov_dist, list): # DistributionMaker object with variable binning
+                metric_val = 0
+                for i in range(len(hypo_asimov_dist)):
+                    metric_val += data_dist[i].metric_total(expected_values=hypo_asimov_dist[i],
+                                                            metric=metric[0], metric_kwargs=metric_kwargs)
+                metric_val += hypo_maker.params.priors_penalty(metric=metric[0])
+            else: # DistributionMaker object with regular binning
                 if metric[0] == 'weighted_chi2':
                     actual_values = data_dist.hist['total']
                     expected_values = hypo_asimov_dist.hist['total']
@@ -2938,9 +2945,9 @@ def fit_hypo(self, data_dist, hypo_maker, metric, minimizer_settings,
             if len(metric) == 1: # One metric for all detectors
                 metric = list(metric) * len(hypo_maker.distribution_makers)
             elif len(metric) != len(hypo_maker.distribution_makers):
-                raise IndexError('Number of defined metrics does not match with number of detectors.')
-        else: # DistributionMaker object
-            assert len(metric) == 1
+                raise IndexError(f"Number of defined metrics does not match with number of detectors.")
+        else:
+            assert len(metric) == 1, f"Only one metric allowed for DistributionMaker"
 
         if check_ordering:
             if 'nh' in hypo_param_selections or 'ih' in hypo_param_selections:
@@ -3026,9 +3033,12 @@ def nofit_hypo(self, data_dist, hypo_maker, hypo_param_selections,
             if len(metric) == 1: # One metric for all detectors
                 metric = list(metric) * len(hypo_maker.distribution_makers)
             elif len(metric) != len(hypo_maker.distribution_makers):
-                raise IndexError('Number of defined metrics does not match with number of detectors.')
-        else: # DistributionMaker object
-            assert len(metric) == 1
+                raise IndexError(f"Number of defined metrics does not match with number of detectors.")
+        elif isinstance(hypo_asimov_dist, MapSet) or isinstance(hypo_asimov_dist, list):
+            # DistributionMaker object (list means variable binning)
+            assert len(metric) == 1, f"Only one metric allowed for DistributionMaker"
+        else: 
+            raise NotImplementedError(f"hypo_maker returned output of type {type(hypo_asimov_dist)}")
         fit_info.metric = metric
 
         # Assess the fit: whether the data came from the hypo_asimov_dist
@@ -3040,6 +3050,8 @@ def nofit_hypo(self, data_dist, hypo_maker, hypo_param_selections,
                     metric_val += data
                 priors = hypo_maker.params.priors_penalty(metric=metric[0]) # uses just the "first" metric for prior
                 metric_val += priors
+            elif isinstance(data_dist, list):
+                raise NotImplementedError()
             else: # DistributionMaker object
 
                 if 'generalized_poisson_llh' == metric[0]:
@@ -3088,6 +3100,8 @@ def nofit_hypo(self, data_dist, hypo_maker, hypo_param_selections,
                 params=hypo_maker.distribution_makers[i].params, metric=metric[i],
                 other_metrics=other_metrics, detector_name=hypo_maker.det_names[i]
             ) for i in range(len(data_dist))]
+        elif isinstance(data_dist, list):
+                raise NotImplementedError()
         else: # DistributionMaker object
 
             if 'generalized_poisson_llh' == metric[0]:

pisa/core/binning.py

@@ -1590,7 +1590,6 @@ def __init__(self, dimensions, name=None, mask=None):
         # Handle masking
         self._init_mask(mask) 
 
-
     def _init_mask(self, mask) :
         '''
         Initialize the bin mask. This can either be specified as:
@@ -1639,13 +1638,11 @@ def _init_mask(self, mask) :
         # Done, store the mask
         self._mask = mask
 
-
     @property
     def name(self):
         """Name of the dimension"""
         return self._name
 
-
     def __repr__(self):
         previous_precision = np.get_printoptions()['precision']
         np.set_printoptions(precision=18)
@@ -3043,6 +3040,88 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
 
+class VarBinning(object):
+    # pylint: disable=line-too-long
+    r"""
+    Variable binning.
+
+    """
+    # pylint: enable=line-too-long
+    def __init__(self, binnings, selections, name=None):
+
+        assert (isinstance(selections, OneDimBinning) or
+                isinstance(selections, list))
+        assert isinstance(binnings, list) and len(binnings) == len(selections)
+        all_equal = True
+        for b in binnings:
+            assert isinstance(b, MultiDimBinning)
+            if isinstance(selections, OneDimBinning):
+                assert selections.name not in b.names
+            else:
+                assert binnings.count(b) == 1, 'Binning used more than once, modify your selection'
+            all_equal = all_equal and b == binnings[0]
+        assert not all_equal, 'No need for VarBinning'
+
+        self._binnings = binnings
+        self._selections = selections
+        self._nselections = len(selections)
+        self._name = name
+        self._names = None
+
+    @property
+    def binnings(self):
+        """list of MultiDimBinning : each binning in a list"""
+        return self._binnings
+
+    @property
+    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):
+        """Name of the dimension"""
+        return self._name
+
+    @property
+    def names(self):
+        """list of strings : names of each dimension contained plus cut var"""
+        if self._names is None:
+            self._names = []
+            for b in self.binnings:
+                self._names.append(b.names)
+        return self._names
+
+    def __pretty__(self, p, cycle):
+        """Method used by the `pretty` library for formatting"""
+        if cycle:
+            p.text('%s(...)' % self.__class__.__name__)
+        else:
+            p.begin_group(4, '%s([' % self.__class__.__name__)
+            for n, dim in enumerate(self):
+                p.breakable()
+                p.pretty(dim)
+                if n < len(self)-1:
+                    p.text(',')
+            p.end_group(4, '])')
+
+    def _repr_pretty_(self, p, cycle):
+        """Method used by e.g. ipython/Jupyter for formatting"""
+        return self.__pretty__(p, cycle)
+
+    def __iter__(self):
+        """Iterate over dimensions. Use `iterbins` to iterate over bins."""
+        return iter(self._binnings)
+
+    def __len__(self):
+        return self._nselections
+
+
 def test_OneDimBinning():
     """Unit tests for OneDimBinning class"""
     # pylint: disable=line-too-long, wrong-import-position
@@ -3499,6 +3578,28 @@ def test_MultiDimBinning():
     logging.info('<< PASS : test_MultiDimBinning >>')
 
 
+def test_VarBinning():
+    """Unit tests for VarBinning class"""
+    # pylint: disable=wrong-import-position
+
+    b1 = OneDimBinning(name='energy', num_bins=40, is_log=True,
+                       domain=[1, 80]*ureg.GeV)
+    b2 = OneDimBinning(name='coszen', num_bins=40, is_lin=True,
+                       domain=[-1, 1])
+    mdb1 = MultiDimBinning([b1, b2])
+    b3 = OneDimBinning(name='energy', num_bins=20, is_log=True,
+                       domain=[1, 80]*ureg.GeV)
+    mdb2 = MultiDimBinning([b3, b2])
+
+    selection = OneDimBinning(name='pid', is_lin=True, num_bins=2, domain=[0, 1])
+    varbin = VarBinning([mdb1, mdb2], selection)
+
+    selection = ['pid < 0.5', 'pid >= 0.5']
+    varbin = VarBinning([mdb1, mdb2], selection)
+
+    logging.info('<< PASS : test_VarBinning >>')
+
+
 if __name__ == "__main__":
     set_verbosity(1)
     test_OneDimBinning()

pisa/core/container.py

@@ -11,6 +11,7 @@
 from collections import defaultdict
 
 import numpy as np
+import re
 
 from pisa import FTYPE
 from pisa.core.binning import OneDimBinning, MultiDimBinning
@@ -276,7 +277,7 @@ class Container():
     """
 
     default_translation_mode = "average"
-    translation_modes = ("average", "sum", None)
+    translation_modes = ("average", "sum")
     array_representations = ("events", "log_events")
 
 
@@ -590,6 +591,18 @@ def translate(self, key, src_representation):
             else:
                 raise NotImplementedError(f"translating {src_representation} to {dest_representation}")
 
+        elif self.tranlation_modes[key] == 'sum':
+            if from_map and to_map:
+                raise NotImplementedError()
+
+            elif to_map:
+                out = self.array_to_binned(key, src_representation, dest_representation, averaged=False)
+                self.representation = dest_representation
+                self[key] = out
+
+            else:
+                raise NotImplementedError()
+
         else:
             raise NotImplementedError()
 
@@ -611,13 +624,13 @@ def find_valid_representation(self, key):
 
         precedence = np.inf
         representation = None
-        
+
         for h in validity.keys():
             if validity[h]:
                 if self.precedence[h] < precedence:
                     precedence = self.precedence[h]
                     representation = self._representations[h]
-                    
+
         return representation
 
     def resample(self, key, src_representation, dest_representation):
@@ -724,6 +737,19 @@ def binned_to_array(self, key, src_representation, dest_representation):
 
         return lookup(sample, weights, hist_binning)
 
+    def get_keep_mask(self, keep_criteria):
+        """Returns a mask that only keeps the events that survive the given cut(s)."""
+        assert isinstance(keep_criteria, str)
+
+        for var in self.keys:
+            # using word boundary \b to replace whole words only
+            keep_criteria = re.sub(
+                r'\b{}\b'.format(var),
+                'self["%s"]'%(var),
+                keep_criteria
+            )
+        return eval(keep_criteria)  # pylint: disable=eval-used
+
 
 
 def test_container():

pisa/core/distribution_maker.py

@@ -281,16 +281,15 @@ def get_outputs(self, return_sum=False, sum_map_name='total',
                 outputs.tex = sum_map_tex_name
                 outputs = MapSet(outputs) # final output must be a MapSet
 
-            # Case where the output of a pipeline is a dict of different MapSets
-            elif isinstance(outputs[0], OrderedDict):
-                output_dict = OrderedDict()
-                for key in outputs[0].keys():
-                    output_dict[key] = sum([sum(A[key]) for A in outputs]) # This produces a Map objects
-                    output_dict[key].name = sum_map_name
-                    output_dict[key].tex = sum_map_tex_name
-                    output_dict[key] = MapSet(output_dict[key])
-
-                outputs = output_dict
+            # Case where the output of a pipeline is a list of different MapSets
+            elif isinstance(outputs[0], list):
+                outs = []
+                for i in range(len(outputs[0])):
+                    o = sum([sum(x) for x in np.array(outputs)[:, i]])
+                    o.name = sum_map_name
+                    o.tex = sum_map_tex_name
+                    outs.append(MapSet(o))
+                outputs = outs
 
         return outputs