Sheshuk/flavor transformations with flavor matrices

[Sheshuk]

To be merged after #350

In this PR I do the following:

• Take the FlavorTransformation classes, implemented by @jpkneller in Jpkneller version 2.0 #308
• Switch to using the FlavorMatrix class everywhere in get_probabilities
• Move the Three/FourFlavorTransformation.Pmf_HighDensityLimit methods to Three/FourFlavorMixingParameters.Pmf_HighDensityLimit - I think they belong next there, next to VacuumMixingMatrix
• refactor the class hierarchy
• Finalize EarthMatter transform (based on BEMEWS)
• Finalize MSWEffect transform (based on SNOSHEWS)
• Update tests

[sybenzvi]

Comment from @Sheshuk (2024-08-13): some classes are implementing all flavor transformations. Others are handling transformations only in certain environments, e.g., in Earth or in vacuum. As a result, the matrix outputs are not consistent or equivalent across all classes (e.g., flavor-to-flavor, flavor-to-mass, mass-to-mass, etc.).

This needs to be checked carefully in the code to ensure that mixing operations between the classes is not allowed if physically the transformations should not be mixed. All classes inherit from FlavorTransformation and the class hierarchy needs to be refactored to prevent unphysical combinations.

[Sheshuk]

Finally I think this is ready for review.
Here is a summary so we can discuss it.

What I did

1. Separated "concrete" flavor transformation classes into a tree, so that SN, Vacuum and Earth transformations are different classes, and distinct from the "full" transformations (children of FlavorTransformation class)
   
2. Made flavor_transformation a nested package instead of single file (so in_sn,in_vacuum and in_earth contain corresponding transforms)
3. Made each transformation produce a FlavorMatrix object, which is very convenient for accessing, multiplication etc. instead of a plain np.array
4. Updated the tests: convert the transformation matrix TwoFlavor>>P_ff>>TwoFlavor and use the same 2Flavor formulas in the tests as before (they still need update in the future)
5. Added new tests: using the FlavorMatrix we can more easily check the matrix form against the equations in snewpy v2 draft, for example:

   snewpy/python/snewpy/test/test_04_xforms.py

   Lines 312 to 319 in d8664a6

   	P_SN = xform.in_sn.P_mf(t,E)
   	#check with eq (A54) from snewpy_v2.0 paper
   	assert np.allclose(P_SN['NU','NU'],
   	[[0, U2['mu','1'], U2['tau','1'],0],
   	[1, 0, 0, 0],
   	[0, U2['mu','3'], U2['tau','3'],0],
   	[0, 0, 0, 1]]
   	)

   
   
6. User interface: since now we have all the concrete classes like AdiabaticMSW as building blocks of the TransformationChain the user would have to build the desired chain from scratch... So for backward compatibility I declared the pre-defined TransformationChains with the same names, as their main components:
   

   snewpy/python/snewpy/flavor_transformation/__init__.py

   Lines 44 to 51 in d8664a6

   	#define default values for backward compatibility
   	AdiabaticMSW = TransformationChain(in_sn.AdiabaticMSW())
   	NonAdiabaticMSWH = TransformationChain(in_sn.NonAdiabaticMSWH())
   	AdiabaticMSWes = TransformationChain(in_sn.AdiabaticMSWes())
   	NonAdiabaticMSWes = TransformationChain(in_sn.NonAdiabaticMSWes())
   	TwoFlavorDecoherence = TransformationChain(in_sn.TwoFlavorDecoherence())
   	NeutrinoDecay = TransformationChain(in_sn.AdiabaticMSW(), in_vacuum.NeutrinoDecay())
   	QuantumDecoherence = TransformationChain(in_sn.AdiabaticMSW(), in_vacuum.QuantumDecoherence())

   
   I'm not sure this is the best way to go. We will probably need to revisit this.
7. Added some basic documentation, we'll need to revisit it as well, once the interface is finalized.

Next steps (to be done in separate PRs):

• Update tests with 3flavor formulas
• Finalize the user interface
• Update the documentation
• SNOSHEWS support needs to be tested separately
• Implement 3flavor rate calculation

[Sheshuk]

One more thing: I had to swap the formulas for NMO and IMO for NeutrinoDecay in the tests: 09735b5 I'm not sure this is correct. we need a cross-check

[JostMigenda]

Thanks a lot for this work @Sheshuk! Sorry it’s taking me a while to review this; I’m still looking through the code but wanted to get at least some high level comments to you now:

1. Separated "concrete" flavor transformation classes into a tree, so that SN, Vacuum and Earth transformations are different classes, and distinct from the "full" transformations (children of FlavorTransformation class)

This class structure looks really good to me. Just two comments:

• I can see in_earth.NoEarthMatter and in_vacuum.NoVacuumTransformation, but no equivalent in_sn.NoSNTransformation? And that’s presumably why in_sn is a required argument to a TransformationChain, while in_vacuum and in_earth are optional? For cases like this, where the symmetry of the three transformation types is broken, we should take care to explicitly point them out in the documentation.
• I’d like to rename MSWEffect, since that name sounds so generic that novice users would likely expect that to be the basic MSW case. Something like CustomMSW or AdvancedMSW would be clearer (and more consistent with AdiabaticMSW, NonAdiabaticMSW and their four-flavour equivalents, which all drop the “Effect” from the class name).

5. Added new tests: using the FlavorMatrix we can more easily check the matrix form against the equations in snewpy v2 draft, for example:

   snewpy/python/snewpy/test/test_04_xforms.py

   Lines 312 to 319 in d8664a6

   	P_SN = xform.in_sn.P_mf(t,E)
   	#check with eq (A54) from snewpy_v2.0 paper
   	assert np.allclose(P_SN['NU','NU'],
   	[[0, U2['mu','1'], U2['tau','1'],0],
   	[1, 0, 0, 0],
   	[0, U2['mu','3'], U2['tau','3'],0],
   	[0, 0, 0, 1]]
   	)

   

This is really nice for readability!

6. User interface: since now we have all the concrete classes like AdiabaticMSW as building blocks of the TransformationChain the user would have to build the desired chain from scratch... So for backward compatibility I declared the pre-defined TransformationChains with the same names, as their main components:
   

   snewpy/python/snewpy/flavor_transformation/__init__.py

   Lines 44 to 51 in d8664a6

   	#define default values for backward compatibility
   	AdiabaticMSW = TransformationChain(in_sn.AdiabaticMSW())
   	NonAdiabaticMSWH = TransformationChain(in_sn.NonAdiabaticMSWH())
   	AdiabaticMSWes = TransformationChain(in_sn.AdiabaticMSWes())
   	NonAdiabaticMSWes = TransformationChain(in_sn.NonAdiabaticMSWes())
   	TwoFlavorDecoherence = TransformationChain(in_sn.TwoFlavorDecoherence())
   	NeutrinoDecay = TransformationChain(in_sn.AdiabaticMSW(), in_vacuum.NeutrinoDecay())
   	QuantumDecoherence = TransformationChain(in_sn.AdiabaticMSW(), in_vacuum.QuantumDecoherence())

   I'm not sure this is the best way to go. We will probably need to revisit this.

At first glance, that makes sense to me; but happy to revisit it later.

7. Added some basic documentation, we'll need to revisit it as well, once the interface is finalized.

That makes sense; I’ll largely skip that during my review.

---

And a first comment on the code:

snewpy/python/snewpy/flavor_transformation/transforms.py

Lines 116 to 119 in d8664a6

	self.in_sn = in_sn
	self.in_vacuum = in_vacuum
	self.in_earth = in_earth
	self.transforms = [in_sn, in_vacuum, in_earth]

The TransformationChain instances are currently saving the transforms in two different places; can we simplify that?

[JostMigenda]

Individual comments are at the appropriate line in the code.

Regarding the code arrangement in flavor_transformation/:
In addition to the in_* submodules (which makes sense), there are three fairly generic submodules (__init__.py, base.py and transforms.py) and it’s not obvious to me what code goes where. My intuition would probably be to drop transforms.py by putting NoTransformation/CompleteExchange/ThreeFlavorDecoherence into __init__.py (alongside the backwards-compatible transformation chains) and moving FlavorTransformation and TransformationChain into base.py?

Finally, I added a commit cleaning up unused imports.

[Sheshuk]

@JostMigenda thanks a lot for your review!

• I can see in_earth.NoEarthMatter and in_vacuum.NoVacuumTransformation, but no equivalent in_sn.NoSNTransformation? And that’s presumably why in_sn is a required argument to a TransformationChain, while in_vacuum and in_earth are optional?

Yes, it is correct. I also don't like it being asymmetric, so if you have any idea of making it more clear for the user, please tell.
Another misleading thing is that we can't reproduce NoTransformation in TransformationChain - even if we completely ignore the SN matter and other effects, the flavor states are transformed to mass ones $P_{\alpha\to i}$ and then back $P_{i\to\beta}$, but because of the decoherence of mass states, these transformation matrices don't produce a unit matrix $I_{\alpha\beta}$.

• I’d like to rename MSWEffect, since that name sounds so generic that novice users would likely expect that to be the basic MSW case. Something like CustomMSW or AdvancedMSW would be clearer (and more consistent with AdiabaticMSW, NonAdiabaticMSW and their four-flavour equivalents, which all drop the “Effect” from the class name).

I agree, CustomMSW would be better here.
I kept all the class names initially defined by Jim in his branch, and didn't change anything because I'm not sure I have enough understanding of underlying physics 😃
But it will be good if we can come with a more clear naming convention in a discussion.

But probably we should do the renaming in a separate PR, it will be rather straightforward.

The TransformationChain instances are currently saving the transforms in two different places; can we simplify that?

Sure. I can make in_sn, in_vacuum, in_earth properties, accessing the elements of self.transforms. Or vice versa. Would that be better?

In addition to the in_* submodules (which makes sense), there are three fairly generic submodules (__init__.py, base.py and transforms.py) and it’s not obvious to me what code goes where. My intuition would probably be to drop transforms.py by putting NoTransformation/CompleteExchange/ThreeFlavorDecoherence into __init__.py (alongside the backwards-compatible transformation chains) and moving FlavorTransformation and TransformationChain into base.py?

That makes sense, thank you, I'll do that!

The only classes standing out here are Three/FourFlavorTransformation in base.py, which are not base classes per se, but a small functionality, providing a way to say "this class should have Three(Four)FlavorMixingParameters,no matter what. If you try to change the parameters, just change the values and don't change their type".
I'm not sure this is the most elegant way to implement it. If you have any ideas, please share

[Sheshuk]

In addition to the in_* submodules (which makes sense), there are three fairly generic submodules (__init__.py, base.py and transforms.py) and it’s not obvious to me what code goes where. My intuition would probably be to drop transforms.py by putting NoTransformation/CompleteExchange/ThreeFlavorDecoherence into __init__.py (alongside the backwards-compatible transformation chains) and moving FlavorTransformation and TransformationChain into base.py?

On second thought, this makes a circular import:
TransformationChain needs NoVacuumTransformation,NoEarthMatter for the default values (and SNTransformation,VacuumTransformation and EarthTransformation for type hints), which depend on FlavorTransformation base class. So we can't put FlavorTransformation in base.py.
We can put it together with the rest transorms in __init__.py, but that will bascially mean that we migrated our transforms.py into __init__.py. I personally don't like when there is a lot of code in __init__.py except for some small definitions of the objects to be imported etc.

[JostMigenda]

• I can see in_earth.NoEarthMatter and in_vacuum.NoVacuumTransformation, but no equivalent in_sn.NoSNTransformation? And that’s presumably why in_sn is a required argument to a TransformationChain, while in_vacuum and in_earth are optional?

Yes, it is correct. I also don't like it being asymmetric, so if you have any idea of making it more clear for the user, please tell. Another misleading thing is that we can't reproduce NoTransformation in TransformationChain - even if we completely ignore the SN matter and other effects, the flavor states are transformed to mass ones P α → i and then back P i → β , but because of the decoherence of mass states, these transformation matrices don't produce a unit matrix I α β .

Hm … good point; I can see it being very confusing if TransformationChain(NoSNTransformation, NoVacuumTransformation, NoEarthMatter) does not match NoTransformation. @jpkneller, do you have any idea how to make it clearer? Or should we just live with that asymmetry and document it well?

• I’d like to rename MSWEffect, since that name sounds so generic that novice users would likely expect that to be the basic MSW case. Something like CustomMSW or AdvancedMSW would be clearer (and more consistent with AdiabaticMSW, NonAdiabaticMSW and their four-flavour equivalents, which all drop the “Effect” from the class name).

I agree, CustomMSW would be better here. I kept all the class names initially defined by Jim in his branch, and didn't change anything because I'm not sure I have enough understanding of underlying physics 😃 But it will be good if we can come with a more clear naming convention in a discussion.

But probably we should do the renaming in a separate PR, it will be rather straightforward.

Yep, let’s do that! I’ve created #366 to keep track of it.

The TransformationChain instances are currently saving the transforms in two different places; can we simplify that?

Sure. I can make in_sn, in_vacuum, in_earth properties, accessing the elements of self.transforms. Or vice versa. Would that be better?

It would certainly eliminate any risk of bugs where we change e.g. self.in_sn but forget to update self.transforms, yes.
After thinking about it a bit longer yesterday, however, I think making self.transforms a namedtuple would be even better?

>>> from collections import namedtuple
>>> Transformations = namedtuple('Transformations', 'in_sn, in_vacuum, in_earth')
>>> transforms = Transformations('AdiabaticMSW', 'NoVacuumTransformation', 'NoEarthMatter')
>>> for t in transforms:
...     print(t)
... 
AdiabaticMSW
NoVacuumTransformation
NoEarthMatter
>>> transforms.in_vacuum
'NoVacuumTransformation'

In particular, the current design doesn’t make it clear that iterating over self.transforms and accessing self.in_* is equivalent; whereas the namedtuple makes that obvious.
(Personally, I also think that self.transforms.in_vacuum, while a bit more verbose, is more readable than self.in_vacuum; but that’s more subjective.)

The only classes standing out here are Three/FourFlavorTransformation in base.py, which are not base classes per se, but a small functionality, providing a way to say "this class should have Three(Four)FlavorMixingParameters,no matter what. If you try to change the parameters, just change the values and don't change their type". I'm not sure this is the most elegant way to implement it. If you have any ideas, please share

I guess they are more of a mix-in than a base class; but I’m fine with that.

In addition to the in_* submodules (which makes sense), there are three fairly generic submodules (__init__.py, base.py and transforms.py) and it’s not obvious to me what code goes where. My intuition would probably be to drop transforms.py by putting NoTransformation/CompleteExchange/ThreeFlavorDecoherence into __init__.py (alongside the backwards-compatible transformation chains) and moving FlavorTransformation and TransformationChain into base.py?

On second thought, this makes a circular import:

Oh, good point!

After thinking about it for a while, I think the main issue is the TransformationChain—it’s a high-level class which depends on the SN/Vacuum/EarthTransformation classes, so it shouldn’t be lumped in with the base classes. With that in mind, I would suggest restructuring as follows:

• One submodule for all concrete transformation presets (both phenomenological ones (NoTransformation/CompleteExchange/ThreeFlavorDecoherence) and the backwards-compatible TransformationChains)
• One submodule for base classes and mix-ins (i.e. FlavorTransformation and Three/FourFlavorTransformation)
• A separate submodule for the TransformationChain class

The first of these should probably be __init__.py (where the backwards-compatible TransformationChains already live); the second should be base.py; and I’d probably rename the last one from transforms.py to TransformationChain.py for clarity.
Since I already went through the whole process, I’ll post a separate PR in a minute with these changes for you to look at (and merge into this PR branch, if you’re happy with it).

[Sheshuk]

After thinking about it a bit longer yesterday, however, I think making self.transforms a namedtuple would be even better?

Good idea, that's perfect use case for that. I also like self.transforms.in_vacuum - it's always nice when the code reads as a natural language.

I'll do it now.

[jpkneller]

One more thing: I had to swap the formulas for NMO and IMO for NeutrinoDecay in the tests: 09735b5 I'm not sure this is correct. we need a cross-check

I'm glad you caught this. I checked the formulas, they are now correct. I must have mixed up NMO and IMO.