Closes #342: adding the function

python/snewpy/flux.py

@@ -66,6 +66,8 @@
 #list of units which will be used as units for decomposition inside the Container
 snewpy_unit_bases = [u.MeV, u.m, u.s, u.kg]
 
+import matplotlib.pyplot as plt
+
 class Axes(IntEnum):
     """Enum to keep the number number of the array dimension for each axis""" 
     flavor=0, #Flavor dimension
@@ -442,7 +444,27 @@ def __class_getitem__(cls, args):
             name = name or f'{cls.__name__}[{unit}]'
             cls._unit_classes[unit] = type(name,(cls,),{'unit':unit})
         return cls._unit_classes[unit]
-
+
+    def plot(flux, projection='energy', styles=None, **kwargs):
+        if projection=='energy':
+            fP = flux.integrate_or_sum('time')
+        else:
+            fP = flux.integrate_or_sum('energy')
+        x = fP.__dict__[projection]
+        if isinstance(styles, dict):
+            styles = styles.get
+        elif styles==None:
+            styles = lambda flv: {'ls':'-' if flv.is_neutrino else ':',
+                                  'color':f'C{flv//2:d}'}
+        for flv in fP.flavor:
+            style = styles(flv)
+            style.update(kwargs)
+            y = fP[flv].array.squeeze()
+            plt.plot(x,y,label=flv.to_tex(), **style)
+        plt.legend(ncols=2)
+        plt.xlabel(f'{projection}, {x.unit._repr_latex_()}')
+        plt.ylabel(f'{fP.__class__.__name__}, {x.unit._repr_latex_()}')
+        return
 
 #some standard container classes that can be used for 
 Flux = Container['1/(MeV*s*m**2)', "d2FdEdT"]