Adding the implementation of the FlavorScheme and FlavorMatrix

[Sheshuk]

This is a suggestion for #309

• Implement FlavorScheme base class and building TwoFlavor, ThreeFlavor, FourFlavor enums as it's children
• Implement FlavorMatrix class, which contains:
  • flavors1, flavors2: children of FlavorSchemeflavor schemes - it means that matrix transforms from flavor1 space to flavors2
  • array: np.ndarray the matrix itself
• Implement conversion matrices between various FlavorShemes, and make them easily accessible
• Keep the FlavorScheme inside the flux.Container class (before it stored an array of flavor values, but we need to keep the full scheme, to be able to operate on it)
• Implement matrix multiplication FlavorMatrix @ FlavorMatrix -> FlavorMatrix
• Implement matrix multiplication FlavorMatrix @ flux.Container -> flux.Container

See the description below

[Sheshuk]

Edit: Updated the conversion matrix

User interface

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.

Flavor schemes:

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

Conversion matrices are defined using operators >> and << on the FlavorSchemes. 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. ]])

Flux Containers

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)

Containers can be multiplied by the flavor matrices, in this manner:

flux_oscillated = M @ flux 

Containers can also be converted to the different flavor schemes:

#no need to know what was the flavor scheme of the flux
flux_3f = flux>>ThreeFlavor

Accessing and slicing of the flux containers

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]

[Sheshuk]

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']

[jpkneller]

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.

[Sheshuk]

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

[JostMigenda]

The implementation is looking quite good; I have a couple of mostly minor comments.

Regarding the high-level API design questions, let’s continue in #309.

[JostMigenda]

As agreed during yesterday’s SNEWS telecon, this + #335 looks good and is ready to be merged into the release_v2.0 branch. Thanks for the nice work!

Like we discussed, we’ll keep the TwoFlavor scheme so we can use the TwoFlavor>>ThreeFlavor matrices in the FlavorTransformations; once those are ready, we’ll see if there’s a bit more cleanup we can do.