Jpkneller add collate and aggregate functions for rate calculator tables

python/snewpy/flux.py

@@ -402,6 +402,25 @@ def __eq__(self, other:'Container')->bool:
                  all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]])
         return result
 
+    def __add__(self,other:'Container'):
+        # Overload the + operator. 
+        # Don't compare the flavors, only that they have the same number
+        if self.__class__==other.__class__ and \
+            self.unit == other.unit and \
+            self.flavor_scheme==other.flavor_scheme and \
+            len(self.flavor)==len(other.flavor) and \
+            all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]):
+                array = self.array+other.array
+                axes = list(self.axes)
+                return Container(array,*axes)
+        else:
+            return NotImplemented
+
+    def __radd__(self,other):
+        if other == 0:
+            return self
+        return self.__add__(other)    
+
     def _is_full_flavor(self):
         return all(self.flavor==list(self.flavor_scheme))
 

python/snewpy/models/model_files.yml

@@ -6,7 +6,7 @@ config:
     - &snewpy "https://github.com/SNEWS2/snewpy/raw/v{snewpy_version}/models/{model}/{filename}"
     - &ccsn_repository "https://github.com/SNEWS2/snewpy-models-ccsn/raw/v0.4/models/{model}/{filename}"
     - &presn_repository "https://github.com/SNEWS2/snewpy-models-presn/raw/v0.2/models/{model}/{filename}"
-    - &presn_repository_main "https://github.com/SNEWS2/snewpy-models-presn/raw/master/models/{model}/{filename}"
+    - &presn_repository_main "https://github.com/SNEWS2/snewpy-models-presn/raw/main/models/{model}/{filename}"
 
 models:
     ccsn:

python/snewpy/rate_calculator.py

@@ -425,4 +425,45 @@ def run(self, flux:Container, detector:str, material:str=None, detector_effects:
             dict[str, Container]
                 A dictionary with interaction rates (as instances of :class:`snewpy.flux.Container`) for each channel.
         """
-        return self.read_detector(detector,material).run(flux, detector_effects=detector_effects)
+        return self.read_detector(detector,material).run(flux, detector_effects=detector_effects)
+
+def collate(rates):
+    """Collates the event rates / numbers table returned by RateCalculator.run 
+       into distinct channels e.g. add all electron elastic scattering and NC channels
+
+    Parameters
+    ----------
+    dict[str, Container]
+                A dictionary with interaction rates (as instances of :class:`snewpy.flux.Container`) for each channel.
+
+    Returns
+    -------
+    dict[str, Container]
+                A dictionary with interaction rates (as instances of :class:`snewpy.flux.Container`) for the collated channels.
+    """
+
+    def aggregate_channels(rates,patterns):
+        for name, pattern in patterns.items():
+            #get channels in rates with names that contain the pattern
+            matches = [channel for channel in rates.keys() if re.search(pattern,channel)]
+            #sum over the matches
+            rates_agg = sum(rates[channel] for channel in matches)
+            #remove matching channels from rates
+            for channel in matches:
+                del rates[channel]
+            #make a new entry with the aggregate 
+            if len(matches) > 0:
+                rates[name] = rates_agg
+        return rates
+
+    # make collated rate table
+    collated_rates = {}
+    patterns = {'nc':'nc_',
+                'eES':'_e', 
+                'coh_helm_Ar':r'coh_helm.*_Ar', 'coh_helm_Ge':r'coh_helm.*_Ge', 'coh_helm_Xe':r'coh_helm.*_Xe',
+                'coh_klein-nystrand_Ar':r'coh_klein.*_Ar', 'coh_klein-nystrand_Ge':r'coh_klein.*_Ge', 'coh_klein-nystrand_Xe':r'coh_kelin.*_Xe'                
+               }
+    for detector in rates:
+        collated_rates[detector] = aggregate_channels(rates[detector],patterns)
+
+    return collated_rates