Closed Sheshuk closed 3 months ago
Edit: Updated the conversion matrix
Here is the result of the interface I ended up with. Probably it's not minimal, but should be quite convenient when we implement the flavor transformations - the flavor conversions are done easily.
Working with the flavor schemes is simple: each one is an IntEnum in snewpy.flavor
, being a subclass of FlavorScheme
.
It's easy to define a new flavor scheme using this function:
from snewpy.flavor import FlavorScheme #base class for all flavors
ExtraSterileFavors = FlavorScheme.from_lepton_names(name='ExtraSteriles', leptons=['E','MU','TAU','S1','S2','S3'])
print(list(ExtraSterileFavors))
prints:
[<ExtraSterile.NU_E: 0>,
<ExtraSterile.NU_E_BAR: 1>,
<ExtraSterile.NU_MU: 2>,
<ExtraSterile.NU_MU_BAR: 3>,
<ExtraSterile.NU_TAU: 4>,
<ExtraSterile.NU_TAU_BAR: 5>,
<ExtraSterile.NU_S1: 6>,
<ExtraSterile.NU_S1_BAR: 7>,
<ExtraSterile.NU_S2: 8>,
<ExtraSterile.NU_S2_BAR: 9>,
<ExtraSterile.NU_S3: 10>,
<ExtraSterile.NU_S3_BAR: 11>]
I made three schemes:
from snewpy.flavor import TwoFlavor, ThreeFlavor, FourFlavor
Conversion matrices are defined using operators >>
and <<
on the FlavorScheme
s. For example:
M2_to_3 = TwoFlavor>>ThreeFlavor
M4_to_2 = TwoFlavor<<FourFlavor
print(M2_to_3)
prints:
FlavorMatrix:<TwoFlavor->ThreeFlavor> shape=(6, 4)
array([[1., 0., 0., 0.],
[0., 1., 0., 0.],
[0., 0., 1., 0.],
[0., 0., 0., 1.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]])
and they can be multiplied by each other:
(ThreeFlavor<<TwoFlavor) @ (TwoFlavor<<FourFlavor)
gives a matrix (obviously,we're losing the information about nu_mu/nu_tau):
FlavorMatrix:<FourFlavor->ThreeFlavor> shape=(6, 8)
array([[1. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ],
[0. , 1. , 0. , 0. , 0. , 0. , 0. , 0. ],
[0. , 0. , 0.5, 0. , 0.5, 0. , 0. , 0. ],
[0. , 0. , 0. , 0.5, 0. , 0.5, 0. , 0. ],
[0. , 0. , 0.5, 0. , 0.5, 0. , 0. , 0. ],
[0. , 0. , 0. , 0.5, 0. , 0.5, 0. , 0. ]])
Container classes from snewpy.flux
now also keep track of flavor_scheme
(alongside with the flavor
attribute, which keeps the actual list of the flavors - which is used for slicing)
flux_oscillated = M @ flux
#no need to know what was the flavor scheme of the flux
flux_3f = flux>>ThreeFlavor
User can directly access the desired components of the flux:
#this works in all schemes,where we have NU_E
flux_nue = flux["NU_E"]
#this will fail on all schemes without NU_X
flux_nux = flux["NU_X"]
#this works only if the flux.flavor_scheme==TwoFlavor, otherwise raises
flux_nue = flux[TwoFlavor.NU_E]
A more real example of making a plot with the flux from the model:
import matplotlib.pyplot as plt
from snewpy.models import ccsn
import astropy.units as u
import numpy as np
#get the model
m = ccsn.Bollig_2016(progenitor_mass=27*u.Msun)
T = np.linspace(0,1, 201)<<u.s
E = np.linspace(0,100, 101)<<u.MeV
flux = m.get_flux(T,E, distance=10*u.kpc)
#integrate over energy
flux_vs_time_2f = flux.integrate('energy')
flux_vs_time_3f = flux_vs_time_2f>>ThreeFlavor
#function for plotting
def plot_time(flux, label='', **kwargs):
for flv in flux.flavor:
plt.plot(flux.time, flux.array[flv].squeeze(), label=label+flv.to_tex(), **kwargs)
plt.plot(flux.time, flux.array.sum(axis=0).squeeze(),label=label+'total', color='k', **kwargs)
plot_time(flux_vs_time_2f, lw=1, label='TwoFlavor ')
plot_time(flux_vs_time_3f, ls=':', lw=2, label='ThreeFlavor ')
plt.legend(ncols=2)
plt.show()
Edit: updated the image
Now the components are the same F['NU_X']==F['NU_MU']==F['NU_TAU']
Thanks Andrey, I'll check it out. It seems that NU_X means MU and TAU. Is that correct? It needs to mean MU or TAU.
It needs to mean MU or TAU.
Sorry, I'm mixing our definition all the time. So it should be F[NU_X] = 0.5*(F[NU_MU]+F[NU_TAU])
? I'll fix that
This is a suggestion for #309
FlavorScheme
base class and buildingTwoFlavor, ThreeFlavor, FourFlavor
enums as it's childrenFlavorMatrix
class, which contains:flavors1
,flavors2
: children ofFlavorScheme
flavor schemes - it means that matrix transforms fromflavor1
space toflavors2
array: np.ndarray
the matrix itselfFlavorShemes
, and make them easily accessibleFlavorScheme
inside theflux.Container
class (before it stored an array offlavor
values, but we need to keep the full scheme, to be able to operate on it)FlavorMatrix @ FlavorMatrix -> FlavorMatrix
FlavorMatrix @ flux.Container -> flux.Container
See the description below