Jpkneller add collate and aggregate functions for rate calculator tables

doc/source/nb/README.md

@@ -6,7 +6,7 @@ The Jupyter notebooks in this directory contain different examples for how to us
 
 These directories contain notebooks demonstrating how to use the core-collapse and pre-supernova models available through SNEWPY.
 
-## AnalyticFluence
+## Analytic3Species
 
 This notebook demonstrates how to use the `Analytic3Species` class from `snewpy.models` to create an analytic supernova model by specifying the luminosity, mean energy and mean squared energy for three neutrino flavors.
 
@@ -25,4 +25,4 @@ This notebook demonstrates how to use SNEWPY’s `snewpy.snowglobes` module to i
 ## `dev` Directory
 
 This directory contains notebooks which may be under development or illustrate usage of internal/undocumented APIs.
-They are not recommended for general users.
+They are not recommended for general users.

python/snewpy/flavor.py

@@ -46,7 +46,7 @@ def __getitem__(cls, key):
         return np.array(list(cls.__members__.values()),dtype=object)[key]
 
 class FlavorScheme(enum.IntEnum, metaclass=FlavorEnumMeta):
-    """Configurable enumeration for different flavor schems (2, 3, 4, ... flavors).
+    """Configurable enumeration for different flavor schems (1, 2, 3, 4, ... flavors).
     """
 
     def to_tex(self):
@@ -89,14 +89,18 @@ def lepton(self):
 
     @classmethod
     def from_lepton_names(cls, name:str, leptons:list):
-        enum_class =  cls(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']])
+        if leptons != None:
+             enum_class =  cls(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']])
+        else:
+             enum_class =  cls(name, start=0, names = [f'NU_{BAR}' for BAR in ['','_BAR']])
         return enum_class
 
     @classmethod
     def take(cls, index):
         return cls[index]
 
-#- Define 2, 3, and 4-flavor schemes for the module.
+#- Define 1, 2, 3, and 4-flavor schemes for the module.
+OneFlavor = FlavorScheme.from_lepton_names('OneFlavor',None)
 TwoFlavor = FlavorScheme.from_lepton_names('TwoFlavor',['E','X'])
 ThreeFlavor = FlavorScheme.from_lepton_names('ThreeFlavor',['E','MU','TAU'])
 FourFlavor = FlavorScheme.from_lepton_names('FourFlavor',['E','MU','TAU','S'])

python/snewpy/flux.py

@@ -108,7 +108,7 @@ def __init__(self,
                  *,
                  integrable_axes: set[Axes] | None = None,
                  flavor_scheme: FlavorScheme | None = None
-    ):
+                ):
         """A container class storing the physical quantity (flux, fluence, rate...), which depends on flavor, time and energy.
 
         Parameters
@@ -149,7 +149,8 @@ def __init__(self,
 
         Nf,Nt,Ne = len(self.flavor), len(self.time), len(self.energy)
         #list all valid shapes of the input array
-        expected_shapes=[(nf,nt,ne) for nf in (Nf,Nf-1) for nt in (Nt,Nt-1) for ne in (Ne,Ne-1)]
+        expected_shapes=[(nf,nt,ne) for nf in (Nf-1,Nf) for nt in (Nt-1,Nt) for ne in (Ne-1,Ne)]
+
         #treat special case if data is 1d array
         if self.array.ndim==1:
             #try to reshape the array to expected shape
@@ -335,6 +336,7 @@ def integrate_or_sum(self, axis: Axes | str)->'Container':
     def can_integrate(self, axis):
         "return true if can be integrated along given axis"
         return Axes.get(axis) in self._integrable_axes
+
     def can_sum(self, axis):
         "return true if can be summed along given axis"
         return Axes.get(axis) not in self._integrable_axes
@@ -484,8 +486,7 @@ def project_to(self, axis='energy', squeeze=False):
         if squeeze:
             return x, fP.array.squeeze().T
         else:
-            return x, fP
-
+            return x, fP        
 
     def plot(flux, projection='energy', styles=None, **kwargs):
         x, fP = flux.project_to(projection, squeeze=False)
@@ -511,6 +512,28 @@ def plot(flux, projection='energy', styles=None, **kwargs):
         plt.ylabel(f'{fP.__class__.__name__}, {x.unit._repr_latex_()}')
         return lines
 
+    @staticmethod
+    def _reconstruct(array, flavor, time, energy, integrable_axes, flavor_scheme):
+        return Container(array,
+                         flavor,
+                         time,
+                         energy,
+                         integrable_axes=integrable_axes,
+                         flavor_scheme=flavor_scheme,
+                         )
+
+    def __reduce__(self):
+        return ( Container._reconstruct,
+                      ( self.array,
+                        self.flavor,
+                        self.time,
+                        self.energy,
+                        self._integrable_axes,
+                        self.flavor_scheme,
+                       ),
+               )    
+
+
 #some standard container classes that can be used for 
 Flux = Container['1/(MeV*s*m**2)', "d2FdEdT"]
 Fluence = Container[Flux.unit*u.s, "dFdE"]

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

@@ -6,10 +6,14 @@
 .. autoclass:: RateCalculator
     :members: run
 """
+import re
+
 import numpy as np
+
 from snewpy.snowglobes_interface import SnowglobesData, guess_material
-from snewpy.neutrino import Flavor
+from snewpy.flavor import ThreeFlavor
 from snewpy.flux import Container
+
 from astropy import units as u
 from warnings import warn
 from typing import Callable
@@ -131,11 +135,13 @@ def __init__(self, callable):
             A function of one parameter (energy). This can be an analitical function, or an interpolation of a (E,value) dataset
         """
         self.value = callable
+
     def __mul__(self, f:Container)->Container:
         e = f.energy #Define sample points
         if not f.can_integrate('energy'): #we have bins, let's use central values for sampling
             e = center(f.energy)
         return f*self.value(e)
+
     def __rmul__(self, f:Container)->Container:
         #same as multiplication from the left
         return self.__mul__(f)
@@ -150,14 +156,14 @@ def from_threshold(cls, e_min=1<<u.MeV):
 class DetectionChannel:
     """Description of a single detection channel in the detector"""
 
-    def __init__(self, flavor:Flavor, xsec:callable, smearing:SmearingMatrix=None,
+    def __init__(self, flavor:ThreeFlavor, xsec:callable, smearing:SmearingMatrix=None,
                  efficiency:FunctionOfEnergy=1., weight:float=1.):
         """
         Parameters
         ----------
         name:str
             channel name
-        flavor:Flavor
+        flavor:ThreeFlavor
             flavor of the interacting neutrino
         xsec:callable or FunctionOfEnergy
             crossection as a function of energy
@@ -184,7 +190,7 @@ def flavor(self):
         return self._flavor
     @flavor.setter
     def flavor(self, flavors):
-        if isinstance(flavors, Flavor):
+        if isinstance(flavors, ThreeFlavor):
             flavors = [flavors]
         self._flavor = flavors
     @property
@@ -208,6 +214,7 @@ def xsec(self, xsec):
 
     def __repr__(self):
         return f'{self.__class__.__name__} (flavor={",".join([f.name for f in self.flavor])}, smearing={self.smearing is not None}, weight={self.weight})'
+
     def calc_rate(self, flux:Container, apply_smearing=True, apply_efficiency=True)->Container:
         """Calculate the event rate in this channel
 
@@ -232,14 +239,15 @@ def _calc_interaction_rate(self, flux):
         """calculate interaction rate for given channel"""
         tgt_mass = 1<<u.kt
         Ntargets = tgt_mass.to_value(u.Dalton)
-        #sum flux over flavors
-        array_total = sum([flux[flv].array for flv in self.flavor])
-        #create a summary flux container
-        flux_total = Container(array_total, flavor=self.flavor, 
-                               time=flux.time, energy=flux.energy,
-                               integrable_axes=flux._integrable_axes)
-        rate = self.xsec*flux_total*self.weight*Ntargets
-        return rate
+        if len(self.flavor) > 1:
+           #sum flux over flavors
+           sumfluxarray = sum([flux[flv].array for flv in self.flavor])
+           #create a summary flux container
+           sumflux = Container(sumfluxarray, flavor=[self.flavor[0],self.flavor[-1]], 
+                               time=flux.time, energy=flux.energy)
+           return self.xsec*sumflux*self.weight*Ntargets
+        else:
+             return self.xsec*flux[self.flavor]*self.weight*Ntargets
 
 class Detector:
     """A detector configuration for the rate calculation. """
@@ -288,12 +296,12 @@ def run(self, flux:Container, detector_effects:bool=True)->dict[str, Container]:
         return result
 
 def _get_flavor_index(channel):
-    _map = {'+e':Flavor.NU_E,
-            '-e':Flavor.NU_E_BAR,
-            '+m':Flavor.NU_MU,
-            '-m':Flavor.NU_MU_BAR,
-            '+t':Flavor.NU_TAU,
-            '-t':Flavor.NU_TAU_BAR
+    _map = {'+e':ThreeFlavor.NU_E,
+            '-e':ThreeFlavor.NU_E_BAR,
+            '+m':ThreeFlavor.NU_MU,
+            '-m':ThreeFlavor.NU_MU_BAR,
+            '+t':ThreeFlavor.NU_TAU,
+            '-t':ThreeFlavor.NU_TAU_BAR
             }
     return _map[channel.parity+channel.flavor]
 
@@ -303,7 +311,9 @@ def _bin_edges_from_centers(centers:np.ndarray)->np.ndarray:
     edges = centers-0.5*np.pad(binw,(0,1),mode='edge') #get lower edges
     edges = np.append(edges,edges[-1]+binw[-1])
     return edges
+
 #--------------------------------------
+
 class RateCalculator(SnowglobesData):
     r"""Simple rate calculation interface.
         Computes expected rate for a detector using SNOwGLoBES data. 
@@ -340,13 +350,13 @@ def __init__(self, base_dir=''):
         """
         super().__init__(base_dir=base_dir)
 
-    def load_xsec(self, channel_name:str, flavor:Flavor)->FunctionOfEnergy:
+    def load_xsec(self, channel_name:str, flavor:ThreeFlavor)->FunctionOfEnergy:
         """Load cross-section for a given channel, interpolated in the energies"""
         xsec = np.loadtxt(self.base_dir/f"xscns/xs_{channel_name}.dat")
         # Cross-section in 10^-38 cm^2
         xp = xsec[:,0]
         #get the column to read from the file
-        column = {Flavor.NU_E:1, Flavor.NU_MU:2, Flavor.NU_TAU:3, Flavor.NU_E_BAR:4, Flavor.NU_MU_BAR:5, Flavor.NU_TAU_BAR:6}[flavor]
+        column = {ThreeFlavor.NU_E:1, ThreeFlavor.NU_MU:2, ThreeFlavor.NU_TAU:3, ThreeFlavor.NU_E_BAR:4, ThreeFlavor.NU_MU_BAR:5, ThreeFlavor.NU_TAU_BAR:6}[flavor]
         yp = xsec[:, column]
         def xsec(energies):
             E = energies.to_value('GeV')
@@ -425,4 +435,64 @@ 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 nested dictionary with interaction rates (as instances of :class:`snewpy.flux.Container`) for each channel.
+
+    Returns
+    -------
+    dict[str, Container]
+                A nested dictionary with interaction rates (as instances of :class:`snewpy.flux.Container`) for the collated channels.
+    """
+
+    def aggregate_channels(rates,patterns):
+        for aggname, 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
+            sumrates = sum([rates[channel].array for channel in matches])
+            #make a new entry with the aggregate 
+            if len(matches) > 0:
+                rates[aggname] = Container(sumrates,ThreeFlavor.take([0,-1]), rates[matches[0]].time, rates[matches[0]].energy)
+            #remove matching channels from rates
+            for channel in matches:
+                del rates[channel]
+        return rates
+
+    # make collated rate table
+    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_klein.*_Xe'                
+               }
+
+    collated_rates = aggregate_channels(rates,patterns)
+
+    return collated_rates  
+
+
+def aggregate(rates):
+    """Sum all the channels in the event rates / numbers table returned by RateCalculator.run 
+
+    Parameters
+    ----------
+    dict[str, Container]
+                A nested 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`) summed over all channels for a given detector.
+    """    
+    sumrates = sum([rates[channel].array for channel in rates])
+    return Container(sumrates,ThreeFlavor.take([0,-1]), rates[0].time, rates[0].energy)
+
+
+