Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding the implementation of the FlavorScheme and FlavorMatrix #324

Merged
merged 31 commits into from
Jun 7, 2024

Conversation

Sheshuk
Copy link
Contributor

@Sheshuk Sheshuk commented Apr 30, 2024

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
Copy link
Contributor Author

Sheshuk commented May 10, 2024

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
Copy link
Contributor Author

Sheshuk commented May 10, 2024

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()

image

Edit: updated the image
Now the components are the same F['NU_X']==F['NU_MU']==F['NU_TAU']

@Sheshuk Sheshuk added the enhancement New feature or request label May 10, 2024
@jpkneller
Copy link
Contributor

jpkneller commented May 10, 2024

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
Copy link
Contributor Author

Sheshuk commented May 10, 2024

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

@jpkneller jpkneller marked this pull request as ready for review May 13, 2024 12:28
@JostMigenda JostMigenda mentioned this pull request May 13, 2024
Copy link
Member

@JostMigenda JostMigenda left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

python/snewpy/flavor.py Outdated Show resolved Hide resolved
python/snewpy/flavor.py Outdated Show resolved Hide resolved
python/snewpy/flavor.py Outdated Show resolved Hide resolved
python/snewpy/flavor.py Outdated Show resolved Hide resolved
python/snewpy/flux.py Outdated Show resolved Hide resolved
python/snewpy/flux.py Show resolved Hide resolved
Copy link
Member

@JostMigenda JostMigenda left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Sheshuk Sheshuk merged commit 65bc4d5 into release_v2.0 Jun 7, 2024
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants