Update snowglobes.py - Experimental

python/snewpy/flavor.py

@@ -1,7 +1,7 @@
 import enum
 import numpy as np
 import typing
-import snewpy.utils
+from snewpy.utils import expand_dimensions_to
 
 class FlavorEnumMeta(enum.EnumMeta):
     """Meta class for setting flavor enums.
@@ -198,7 +198,7 @@ def __matmul__(self, other):
             try:
                 m0, m1 = self.array, other.array
                 ndims = max(m0.ndim, m1.ndim)
-                m0,m1 = [snewpy.utils.expand_dimensions_to(m, ndim=ndims) for m in [m0,m1]]
+                m0,m1 = [expand_dimensions_to(m, ndim=ndims) for m in [m0,m1]]
                 array = np.einsum('ij...,jk...->ik...',m0,m1)
                 return FlavorMatrix(array, self.flavor_out, from_flavor = other.flavor_in)
             except Exception as e:

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/base.py

@@ -47,6 +47,8 @@ def __init__(self, time, metadata):
         """
         self.time = time
         self.metadata = metadata
+        for entry in metadata:
+            self.name += f"-{metadata[entry]}"
 
     def __repr__(self):
         """Default representation of the model.
@@ -132,10 +134,10 @@ def get_initial_spectra(self, t, E):
             A container with the information about the initial neutrino spectra
         """
         spectra_dict = self._get_initial_spectra_dict(t, E, flavors=ThreeFlavor)
-        initial_spectra =  flux.Container['1/(MeV*s)'].from_dict(spectra_dict, 
-                                                                time=t,
-                                                                energy=E,
-                                                                flavor_scheme=ThreeFlavor)
+        initial_spectra =  flux.Spectrum.from_dict(spectra_dict, 
+                                                   time=t,
+                                                   energy=E,
+                                                   flavor_scheme=ThreeFlavor)
         return initial_spectra
 
     def get_transformed_spectra(self, t, E, flavor_xform):

python/snewpy/models/ccsn.py

@@ -50,6 +50,7 @@ def __init__(self):
         self.metadata["EOS"] = "HS(DD2)"
         self.metadata["Progenitor mass"] = 18 * u.Msun
         filename='Fischer_2020.tar.gz'
+
         return super().__init__(filename, metadata=self.metadata)
 
 @RegistryModel(
@@ -78,6 +79,7 @@ def __init__(self, progenitor_mass:u.Quantity, revival_time:u.Quantity, metallic
             filename = f"nakazato-{eos}-BH-z{metallicity}-s{progenitor_mass:3.1f}.fits"
         #modify metadata if needed...
         #self.metadata['name']=value
+
         return super().__init__(filename, self.metadata)
 
 
@@ -93,6 +95,7 @@ def __init__(self, progenitor_mass:u.Quantity, eos:str):
             filename = f'sukhbold-{eos}-z{progenitor_mass.value:3.1f}.fits'
         else:
             filename = f'sukhbold-{eos}-s{progenitor_mass.value:3.1f}.fits'
+
         return super().__init__(filename, self.metadata)
 
 
@@ -111,6 +114,7 @@ class Tamborra_2014(loaders.Tamborra_2014):
     def __init__(self, *, progenitor_mass:u.Quantity, direction:int):
         filename = f's{progenitor_mass.value:3.1f}c_3D_dir{direction}'
         # Metadata is handled by __init__ in _GarchingArchiveModel
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 @RegistryModel(
@@ -123,6 +127,7 @@ class Bollig_2016(loaders.Bollig_2016):
     """
     def __init__(self, progenitor_mass:u.Quantity, eos:str='LS220'):
         filename = f's{progenitor_mass.value:3.1f}c'
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 @RegistryModel(
@@ -139,6 +144,7 @@ class Walk_2018(loaders.Walk_2018):
     def __init__(self, *, progenitor_mass:u.Quantity, rotation:str, direction:int):
 
         filename = f's{progenitor_mass.value:3.1f}c_3D_{rotation}rot_dir{direction}'
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 
@@ -156,6 +162,7 @@ class Walk_2019(loaders.Walk_2019):
     """
     def __init__(self, * ,progenitor_mass:u.Quantity, direction:int):
         filename = f's{progenitor_mass.value:3.1f}c_3DBH_dir{direction}'
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 
@@ -173,6 +180,7 @@ class OConnor_2013(loaders.OConnor_2013):
     def __init__(self, eos:str, progenitor_mass:u.Quantity):
         # Load from Parameters
         filename = f'{eos}_timeseries.tar.gz'
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 @RegistryModel(
@@ -185,6 +193,7 @@ def __init__(self, progenitor_mass:u.Quantity):
         self.metadata["EOS"] = "LS220"
         # Filename is currently the same regardless of parameters
         filename = 'M1_neutrinos.dat'
+
         return super().__init__(filename, self.metadata)
 
 @RegistryModel(
@@ -198,6 +207,7 @@ class Zha_2021(loaders.Zha_2021):
     def __init__(self, *, progenitor_mass:u.Quantity):
         self.metadata["EOS"] = "STOS_B145"
         filename = f's{progenitor_mass.value:g}.dat'
+
         return super().__init__(filename, self.metadata)
 
 @RegistryModel(
@@ -225,6 +235,7 @@ def __init__(self, *, progenitor_mass, turbmixing_param):
             fname = f'stir_a{turbmixing_param:3.2f}/stir_multimessenger_a{turbmixing_param:3.2f}_m{progenitor_mass.value:.1f}.h5'
         else:
             fname = f'stir_a{turbmixing_param:3.2f}/stir_multimessenger_a{turbmixing_param:3.2f}_m{progenitor_mass.value:g}.h5'
+
         return super().__init__(fname, self.metadata)
 
 @RegistryModel(
@@ -241,6 +252,7 @@ class Kuroda_2020(loaders.Kuroda_2020):
     def __init__(self, *, rotational_velocity, magnetic_field_exponent):
         self.metadata["EOS"] = "LS220"
         filename = f'LnuR{int(rotational_velocity.value):1d}0B{int(magnetic_field_exponent):02d}.dat'
+
         return super().__init__(filename, self.metadata)
 
 @RegistryModel(
@@ -261,6 +273,7 @@ def __init__(self, cache_flux=False, *, progenitor_mass):
             filename = f'lum_spec_{int(progenitor_mass.value):d}M_r250.h5'
         else:
             filename = f'lum_spec_{int(progenitor_mass.value):d}M.h5'
+
         return super().__init__(filename, self.metadata, cache_flux=cache_flux)
 
 @RegistryModel(
@@ -277,6 +290,7 @@ def __init__(self, progenitor_mass:u.Quantity):
             filename = f'lum_spec_{int(progenitor_mass.value):2d}M_r10000_dat.h5'
         else:
             filename = f'lum_spec_{progenitor_mass.value:.2f}M_r10000_dat.h5'
+
         return super().__init__(filename, self.metadata)
 
 
@@ -316,6 +330,7 @@ def __init__(self, progenitor_mass:u.Quantity):
         progenitor = self._mass_to_progenitor[progenitor_mass]
         self.metadata['Black hole'] = progenitor.endswith('.bh')
         filename = f'lum_spec_{progenitor}_dat.h5'
+
         return super().__init__(filename, self.metadata)
 
 
@@ -341,6 +356,7 @@ def __init__(self, progenitor_mass:u.Quantity):
         progenitor = self._mass_to_progenitor[progenitor_mass]
         self.metadata['Black hole'] = progenitor == '12.25' or progenitor == '14' or progenitor == '19.56' or progenitor == '40'
         filename = f'lum_spec_{progenitor}_dat.h5'
+
         return super().__init__(filename, self.metadata)
 
 
@@ -393,6 +409,7 @@ def __init__(self, axion_mass:u.Quantity, axion_coupling:u.Quantity):
 
         # Set the metadata.
         self.metadata['PNS mass'] = pns_mass*u.Msun
+
         return super().__init__(filename, self.metadata)
 
 
@@ -467,6 +484,7 @@ def __init__(self, *, Bfield:str, direction:str, rotation:int=None, grav:str=Non
             filename = f'{Bfield}_b12_dipdecay_3d_grav{grav}_snewpy_{direction}.dat'
         if Bfield=='L1':
             filename = f'{Bfield}_b12_3d_{rotation}deg_snewpy_{direction}.dat'
+
         return super().__init__(filename=filename, metadata=self.metadata)
 
 class SNOwGLoBES:
@@ -479,6 +497,8 @@ def __init__(self, tarfilename):
         tarfilename: str
             Absolute or relative path to tar archive with SNOwGLoBES files.
         """
+        self.name = self.__class__.__name__     
+
         self.tfname = tarfilename
         tf = tarfile.open(self.tfname)
 
@@ -547,25 +567,4 @@ def get_fluence(self, t):
 
         return fluence
 
-class Analytic3Species(PinchedModel):
-    """An analytical model calculating spectra given total luminosity,
-    average energy, and rms or pinch, for each species.
-    """
-
-    param = "There are no input files available for this class. Use `doc/scripts/Analytic.py` in the SNEWPY GitHub repo to create a custom input file."
-
-    def get_param_combinations(cls):
-        print(cls.param)
-        return []
-
-    def __init__(self, filename):
-        """
-        Parameters
-        ----------
-        filename : str
-            Absolute or relative path to file with model data.
-        """
 
-        simtab = Table.read(filename,format='ascii')
-        self.filename = filename
-        super().__init__(simtab, metadata={})