From 693d5eb09b62f70d11dd61f721353d11cacf9367 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 16:23:03 +0300 Subject: [PATCH 01/31] Adding the implementation of the FlavorScheme and FlavorMatrix --- python/snewpy/flavor.py | 123 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 python/snewpy/flavor.py diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py new file mode 100644 index 000000000..f5a68fb00 --- /dev/null +++ b/python/snewpy/flavor.py @@ -0,0 +1,123 @@ +def MakeFlavorScheme(name:str, leptons:list[str]): + return enum.Enum(name, start=0, type=FlavorScheme, + names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] + ) + +TwoFlavor = MakeFlavorScheme('TwoFlavor',['E','X']) +ThreeFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU']) +FourFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) + +class FlavorScheme: + def to_tex(self): + """LaTeX-compatible string representations of flavor.""" + if self.is_antineutrino: + return r'$\overline{{\nu}}_{0}$'.format(self.name[3].lower()) + return r'$\{0}$'.format(self.name.lower()) + + @classmethod + def to_value(cls,key)->int: + if isinstance(key, str): + try: + return cls[key].value + except KeyError as e: + raise KeyError( + f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' + ) + elif isinstance(key, FlavorScheme): + if not isinstance(key, cls): + raise TypeError(f'Value {repr(key)} is not from {cls.__name__} sheme!') + return key.value + return key + + @property + def is_neutrino(self): + return not self.is_antineutrino + + @property + def is_antineutrino(self): + return '_BAR' in self.name + + @property + def is_electron(self): + """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" + return self.lepton=='E' + + @property + def is_muon(self): + """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" + return self.lepton=='MU' + + @property + def is_tauon(self): + """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" + return self.lepton=='TAU' + + @property + def is_sterile(self): + return self.lepton=='S' + + @property + def lepton(self): + return self.name.split('_')[1] + + +def _makeFlavorScheme(name:str, leptons:list[str]): + return enum.Enum(name, start=0, type=FlavorScheme, + names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] + ) + +TwoFlavor = MakeFlavorScheme('TwoFlavor',['E','X']) +ThreeFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU']) +FourFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) + + +class FlavorMatrix: + def __init__(self, + array:np.ndarray, + flavors:FlavorScheme, + flavors2:FlavorScheme|None = None + ): + self.array = np.asarray(array) + self.flavors1 = flavors + self.flavors2 = flavors2 or flavors + assert self.array.shape == (len(self.flavors2), len(self.flavors1)) + + def _convert_index(self, index): + if isinstance(index, str) or (not isinstance(index,typing.Iterable)): + index = [index] + new_idx = [flavors.to_value(idx) for idx,flavors in zip(index, [self.flavors2,self.flavors1])] + print(new_idx) + return tuple(new_idx) + + def __getitem__(self, index): + return self.array[self._convert_index(index)] + + def __setitem__(self, index, value): + self.array[self._convert_index(index)]=value + + def __repr__(self): + s = f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}>:' + s+='\n'+repr(self.array) + return s + #methods for creating matrix + @classmethod + def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + flavors2 = flavors2 or flavors1 + shape = (len(flavors2), len(flavors1)) + data = np.zeros(shape) + return cls(data, flavors1, flavors2) + + @classmethod + def eye(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + flavors2 = flavors2 or flavors1 + shape = (len(flavors2), len(flavors1)) + data = np.eye(*shape) + return cls(data, flavors1, flavors2) + + @classmethod + def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + flavors2 = flavors2 or flavors1 + data = [[f1.name==f2.name + for f1 in flavors1] + for f2 in flavors2] + return cls(np.array(data,dtype=float), flavors1, flavors2) \ No newline at end of file From 698b0776cf2fda43474c8258928384ec0d4d1fa0 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 16:27:40 +0300 Subject: [PATCH 02/31] Fix --- python/snewpy/flavor.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index f5a68fb00..2c83c022e 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -1,11 +1,6 @@ -def MakeFlavorScheme(name:str, leptons:list[str]): - return enum.Enum(name, start=0, type=FlavorScheme, - names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] - ) - -TwoFlavor = MakeFlavorScheme('TwoFlavor',['E','X']) -ThreeFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU']) -FourFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) +import enum +import numpy as np +import typing class FlavorScheme: def to_tex(self): @@ -66,9 +61,9 @@ def _makeFlavorScheme(name:str, leptons:list[str]): names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] ) -TwoFlavor = MakeFlavorScheme('TwoFlavor',['E','X']) -ThreeFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU']) -FourFlavor = MakeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) +TwoFlavor = _makeFlavorScheme('TwoFlavor',['E','X']) +ThreeFlavor = _makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) +FourFlavor = _makeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) class FlavorMatrix: From ae917a4b198ac0e4277143d09efb7fc29e8352dc Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 20:48:34 +0300 Subject: [PATCH 03/31] Adding 'FlavorMatrix.from_function' and '__matmul__' methods --- python/snewpy/flavor.py | 47 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 42 insertions(+), 5 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 2c83c022e..86e450e61 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -75,7 +75,9 @@ def __init__(self, self.array = np.asarray(array) self.flavors1 = flavors self.flavors2 = flavors2 or flavors - assert self.array.shape == (len(self.flavors2), len(self.flavors1)) + expected_shape = (len(self.flavors2), len(self.flavors1)) + assert self.array.shape == expected_shape, \ + f"Array shape {self.array.shape} mismatch expected {expected_shape}" def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): @@ -94,6 +96,13 @@ def __repr__(self): s = f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}>:' s+='\n'+repr(self.array) return s + + def __matmul__(self, other): + if isinstance(other, FlavorMatrix): + data = np.tensordot(self.array, other.array, axes=[0,1]) + return FlavorMatrix(data, self.flavors1, other.flavors2) + raise TypeError(f"Cannot multiply object of {self.__class__} by {other.__class__}") + #methods for creating matrix @classmethod def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): @@ -108,11 +117,39 @@ def eye(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): shape = (len(flavors2), len(flavors1)) data = np.eye(*shape) return cls(data, flavors1, flavors2) - + @classmethod - def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + def from_function(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None, *, function): flavors2 = flavors2 or flavors1 - data = [[f1.name==f2.name + data = [[function(f1,f2) for f1 in flavors1] for f2 in flavors2] - return cls(np.array(data,dtype=float), flavors1, flavors2) \ No newline at end of file + return cls(np.array(data,dtype=float), flavors1, flavors2) + @classmethod + def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + return cls.from_function(flavors1, flavors2, lambda f1,f2: 1.*(f1.name==f2.name)) + + +def flavor_matrix(flavor1:FlavorScheme, flavor2:FlavorScheme=None): + """A decorator for creating the flavor matrix from the given function""" + flavor2 = flavor2 or flavor1 + def _decorator(func): + return FlavorMatrix.from_function(flavor1, flavor2, function=func) + return _decorator + +#define some conversion matrices +@flavor_matrix(ThreeFlavor,TwoFlavor) +def convert_3to2(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): + return 0.5 + return 0 + +@flavor_matrix(TwoFlavor,ThreeFlavor) +def convert_2to3(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): + return 1 + return 0 \ No newline at end of file From 024ec358368f2df2c7695ef18487c5f8af4c562d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 21:30:37 +0300 Subject: [PATCH 04/31] Using TwoFlavor as neutrino.Flavor --- python/snewpy/flavor.py | 9 ++++++--- python/snewpy/neutrino.py | 27 +-------------------------- 2 files changed, 7 insertions(+), 29 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 86e450e61..27d58c109 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -10,7 +10,7 @@ def to_tex(self): return r'$\{0}$'.format(self.name.lower()) @classmethod - def to_value(cls,key)->int: + def _to_value(cls,key)->int: if isinstance(key, str): try: return cls[key].value @@ -22,6 +22,8 @@ def to_value(cls,key)->int: if not isinstance(key, cls): raise TypeError(f'Value {repr(key)} is not from {cls.__name__} sheme!') return key.value + elif isinstance(key, typing.Iterable): + return tuple(map(cls._to_value, key)) return key @property @@ -57,9 +59,10 @@ def lepton(self): def _makeFlavorScheme(name:str, leptons:list[str]): - return enum.Enum(name, start=0, type=FlavorScheme, + enum_class = enum.IntEnum(name, start=0, type=FlavorScheme, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] ) + return enum_class TwoFlavor = _makeFlavorScheme('TwoFlavor',['E','X']) ThreeFlavor = _makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) @@ -82,7 +85,7 @@ def __init__(self, def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): index = [index] - new_idx = [flavors.to_value(idx) for idx,flavors in zip(index, [self.flavors2,self.flavors1])] + new_idx = [flavors._to_value(idx) for idx,flavors in zip(index, [self.flavors2,self.flavors1])] print(new_idx) return tuple(new_idx) diff --git a/python/snewpy/neutrino.py b/python/snewpy/neutrino.py index bcbb53e12..e1612360f 100644 --- a/python/snewpy/neutrino.py +++ b/python/snewpy/neutrino.py @@ -7,6 +7,7 @@ from typing import Optional import numpy as np from collections.abc import Mapping +from .flavor import TwoFlavor as Flavor class MassHierarchy(IntEnum): """Neutrino mass ordering: ``NORMAL`` or ``INVERTED``.""" @@ -23,32 +24,6 @@ def derive_from_dm2(cls, dm12_2, dm32_2, dm31_2): else: return MassHierarchy.INVERTED -class Flavor(IntEnum): - """Enumeration of CCSN Neutrino flavors.""" - NU_E = 0 - NU_X = 1 - NU_E_BAR = 2 - NU_X_BAR = 3 - - def to_tex(self): - """LaTeX-compatible string representations of flavor.""" - if '_BAR' in self.name: - return r'$\overline{{\nu}}_{0}$'.format(self.name[3].lower()) - return r'$\{0}$'.format(self.name.lower()) - - @property - def is_electron(self): - """Return ``True`` for ``Flavor.NU_E`` and ``Flavor.NU_E_BAR``.""" - return self.value in (Flavor.NU_E.value, Flavor.NU_E_BAR.value) - - @property - def is_neutrino(self): - """Return ``True`` for ``Flavor.NU_E`` and ``Flavor.NU_X``.""" - return self.value in (Flavor.NU_E.value, Flavor.NU_X.value) - - @property - def is_antineutrino(self): - return self.value in (Flavor.NU_E_BAR.value, Flavor.NU_X_BAR.value) @dataclass class MixingParameters3Flavor(Mapping): From 94c7b98ab1f159aafce2188818faa730b807b18b Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 21:44:16 +0300 Subject: [PATCH 05/31] Fixing the to_tex --- python/snewpy/flavor.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 27d58c109..383c92150 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -5,9 +5,13 @@ class FlavorScheme: def to_tex(self): """LaTeX-compatible string representations of flavor.""" + base = r'\nu' if self.is_antineutrino: - return r'$\overline{{\nu}}_{0}$'.format(self.name[3].lower()) - return r'$\{0}$'.format(self.name.lower()) + base = r'\overline{\nu}' + lepton = self.lepton.lower() + if self.is_muon or self.is_tauon: + lepton = f"\{lepton}" + return f"${base}_{{{lepton}}}$" @classmethod def _to_value(cls,key)->int: @@ -93,7 +97,7 @@ def __getitem__(self, index): return self.array[self._convert_index(index)] def __setitem__(self, index, value): - self.array[self._convert_index(index)]=value + self.array[self._convert_index(index)] = value def __repr__(self): s = f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}>:' From 38a6879ab4388c9e35e1aa13be2c9d5ef4d9e80d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 21:49:47 +0300 Subject: [PATCH 06/31] fix typo --- python/snewpy/flavor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 383c92150..022d70b92 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -10,7 +10,7 @@ def to_tex(self): base = r'\overline{\nu}' lepton = self.lepton.lower() if self.is_muon or self.is_tauon: - lepton = f"\{lepton}" + lepton = '\\'+lepton return f"${base}_{{{lepton}}}$" @classmethod From 47ef94f8371fb20000ccea32e2cf57a16eb7af12 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 30 Apr 2024 21:51:50 +0300 Subject: [PATCH 07/31] remove type annotation breaking python 3.8 --- python/snewpy/flavor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 022d70b92..9e8c17438 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -62,7 +62,7 @@ def lepton(self): return self.name.split('_')[1] -def _makeFlavorScheme(name:str, leptons:list[str]): +def _makeFlavorScheme(name:str, leptons:list): enum_class = enum.IntEnum(name, start=0, type=FlavorScheme, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] ) From e80bc239d8f4e4d9cd9b02666224125ef2e0615d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 3 May 2024 12:01:52 +0300 Subject: [PATCH 08/31] Update type annotations --- python/snewpy/flavor.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 9e8c17438..85e88cbd0 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -77,7 +77,7 @@ class FlavorMatrix: def __init__(self, array:np.ndarray, flavors:FlavorScheme, - flavors2:FlavorScheme|None = None + flavors2:FlavorScheme = None ): self.array = np.asarray(array) self.flavors1 = flavors @@ -112,28 +112,28 @@ def __matmul__(self, other): #methods for creating matrix @classmethod - def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): flavors2 = flavors2 or flavors1 shape = (len(flavors2), len(flavors1)) data = np.zeros(shape) return cls(data, flavors1, flavors2) @classmethod - def eye(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + def eye(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): flavors2 = flavors2 or flavors1 shape = (len(flavors2), len(flavors1)) data = np.eye(*shape) return cls(data, flavors1, flavors2) @classmethod - def from_function(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None, *, function): + def from_function(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None, *, function): flavors2 = flavors2 or flavors1 data = [[function(f1,f2) for f1 in flavors1] for f2 in flavors2] return cls(np.array(data,dtype=float), flavors1, flavors2) @classmethod - def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme|None = None): + def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): return cls.from_function(flavors1, flavors2, lambda f1,f2: 1.*(f1.name==f2.name)) @@ -159,4 +159,5 @@ def convert_2to3(f1,f2): return 1. if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): return 1 - return 0 \ No newline at end of file + return 0 + From cd6e7fe6be063d088783eec15ab89432e156ce75 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 3 May 2024 15:43:39 +0300 Subject: [PATCH 09/31] Using custom metaclass with __getitem__ method --- python/snewpy/flavor.py | 45 ++++++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 85e88cbd0..a7501c1a3 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -2,7 +2,28 @@ import numpy as np import typing -class FlavorScheme: +class EnumMeta(enum.EnumMeta): + def __getitem__(cls, key): + print(cls) + #if this is from a flavor scheme: check that it's ours + if isinstance(key, FlavorScheme): + if not isinstance(key, cls): + raise TypeError(f'Value {repr(key)} is not from {cls.__name__} sheme!') + return key + #if this is a string find it by name + if isinstance(key, str): + try: + return super().__getitem__(key) + except KeyError as e: + raise KeyError( + f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' + ) + #if this is an iterable: apply to each value, and construct a tuple + if isinstance(key, typing.Iterable): + return tuple(map(cls.__getitem__, key)) + #if this is an int value: find a matching + +class FlavorScheme(enum.IntEnum, metaclass=EnumMeta): def to_tex(self): """LaTeX-compatible string representations of flavor.""" base = r'\nu' @@ -12,23 +33,6 @@ def to_tex(self): if self.is_muon or self.is_tauon: lepton = '\\'+lepton return f"${base}_{{{lepton}}}$" - - @classmethod - def _to_value(cls,key)->int: - if isinstance(key, str): - try: - return cls[key].value - except KeyError as e: - raise KeyError( - f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' - ) - elif isinstance(key, FlavorScheme): - if not isinstance(key, cls): - raise TypeError(f'Value {repr(key)} is not from {cls.__name__} sheme!') - return key.value - elif isinstance(key, typing.Iterable): - return tuple(map(cls._to_value, key)) - return key @property def is_neutrino(self): @@ -63,7 +67,7 @@ def lepton(self): def _makeFlavorScheme(name:str, leptons:list): - enum_class = enum.IntEnum(name, start=0, type=FlavorScheme, + enum_class = FlavorScheme(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] ) return enum_class @@ -89,8 +93,7 @@ def __init__(self, def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): index = [index] - new_idx = [flavors._to_value(idx) for idx,flavors in zip(index, [self.flavors2,self.flavors1])] - print(new_idx) + new_idx = [flavors[idx] for idx,flavors in zip(index, [self.flavors2,self.flavors1])] return tuple(new_idx) def __getitem__(self, index): From b2e18d74ec7d20e2892e827f47e1cd7d3f3db8fd Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Sat, 4 May 2024 15:19:08 +0300 Subject: [PATCH 10/31] Adding tests for FlavorMatrix --- python/snewpy/test/test_flavors.py | 97 ++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 python/snewpy/test/test_flavors.py diff --git a/python/snewpy/test/test_flavors.py b/python/snewpy/test/test_flavors.py new file mode 100644 index 000000000..87a1dc94f --- /dev/null +++ b/python/snewpy/test/test_flavors.py @@ -0,0 +1,97 @@ +import pytest +import numpy as np +from snewpy import flavor +from snewpy.flavor import TwoFlavor,ThreeFlavor,FourFlavor + +class TestFlavorScheme: + @staticmethod + def test_flavor_scheme_lengths(): + assert len(TwoFlavor)==4 + assert len(ThreeFlavor)==6 + assert len(FourFlavor)==8 + + @staticmethod + def test_getitem_string(): + assert TwoFlavor['NU_E'] == TwoFlavor.NU_E + assert TwoFlavor['NU_X'] == TwoFlavor.NU_X + with pytest.raises(KeyError): + TwoFlavor['NU_MU'] + + @staticmethod + def test_getitem_enum(): + assert TwoFlavor[TwoFlavor.NU_E] == TwoFlavor.NU_E + assert TwoFlavor[TwoFlavor.NU_X] == TwoFlavor.NU_X + with pytest.raises(TypeError): + TwoFlavor[ThreeFlavor.NU_E] + + @staticmethod + def test_values_from_different_enums(): + assert TwoFlavor.NU_E==ThreeFlavor.NU_E + assert TwoFlavor.NU_E_BAR==ThreeFlavor.NU_E_BAR + + @staticmethod + def test_makeFlavorScheme(): + TestFlavor = flavor.makeFlavorScheme('TestFlavor',['A','B','C']) + assert len(TestFlavor)==6 + assert [f.name for f in TestFlavor]==['NU_A','NU_A_BAR','NU_B','NU_B_BAR','NU_C','NU_C_BAR'] + + @staticmethod + def test_flavor_properties(): + f = ThreeFlavor.NU_E + assert f.is_neutrino + assert f.is_electron + assert not f.is_muon + assert not f.is_tauon + assert f.lepton=='E' + + f = ThreeFlavor.NU_MU + assert f.is_neutrino + assert not f.is_electron + assert f.is_muon + assert not f.is_tauon + assert f.lepton=='MU' + + f = ThreeFlavor.NU_E_BAR + assert not f.is_neutrino + assert f.is_electron + assert not f.is_muon + assert not f.is_tauon + assert f.lepton=='E' + + f = ThreeFlavor.NU_MU_BAR + assert not f.is_neutrino + assert not f.is_electron + assert f.is_muon + assert not f.is_tauon + assert f.lepton=='MU' + + f = ThreeFlavor.NU_TAU + assert f.is_neutrino + assert not f.is_electron + assert not f.is_muon + assert f.is_tauon + assert f.lepton=='TAU' + + f = ThreeFlavor.NU_TAU_BAR + assert not f.is_neutrino + assert not f.is_electron + assert not f.is_muon + assert f.is_tauon + assert f.lepton=='TAU' + +class TestFlavorMatrix: + @staticmethod + def test_init_square_matrix(): + m = flavor.FlavorMatrix(array=np.ones(shape=(4,4)), flavors=TwoFlavor) + assert m.shape == (4,4) + assert m.flavors_from == TwoFlavor + assert m.flavors_to == TwoFlavor + + @staticmethod + def test_init_square_matrix_with_wrong_shape_raises_ValueError(): + with pytest.raises(ValueError): + m = flavor.FlavorMatrix(array=np.ones(shape=(4,5)), flavors=TwoFlavor) + with pytest.raises(ValueError): + m = flavor.FlavorMatrix(array=np.ones(shape=(5,5)), flavors=TwoFlavor) + with pytest.raises(ValueError): + m = flavor.FlavorMatrix(array=np.ones(shape=(5,4)), flavors=TwoFlavor) \ No newline at end of file From f12c747ed8837e7ace8d26b34d7e4c886cc02965 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Sat, 4 May 2024 15:20:14 +0300 Subject: [PATCH 11/31] Add convenience properties for the matrix --- python/snewpy/flavor.py | 49 +++++++++++++++++++++++++++-------------- 1 file changed, 32 insertions(+), 17 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index a7501c1a3..8981d0e59 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -4,7 +4,9 @@ class EnumMeta(enum.EnumMeta): def __getitem__(cls, key): - print(cls) + #if this is an iterable: apply to each value, and construct a tuple + if isinstance(key, typing.Iterable) and not isinstance(key, str): + return tuple(map(cls.__getitem__, key)) #if this is from a flavor scheme: check that it's ours if isinstance(key, FlavorScheme): if not isinstance(key, cls): @@ -18,9 +20,7 @@ def __getitem__(cls, key): raise KeyError( f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' ) - #if this is an iterable: apply to each value, and construct a tuple - if isinstance(key, typing.Iterable): - return tuple(map(cls.__getitem__, key)) + #if this is an int value: find a matching class FlavorScheme(enum.IntEnum, metaclass=EnumMeta): @@ -44,17 +44,14 @@ def is_antineutrino(self): @property def is_electron(self): - """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" return self.lepton=='E' @property def is_muon(self): - """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" return self.lepton=='MU' @property def is_tauon(self): - """Return ``True`` for ``TwoFlavor.NU_E`` and ``TwoFlavor.NU_E_BAR``.""" return self.lepton=='TAU' @property @@ -66,15 +63,15 @@ def lepton(self): return self.name.split('_')[1] -def _makeFlavorScheme(name:str, leptons:list): +def makeFlavorScheme(name:str, leptons:list): enum_class = FlavorScheme(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] ) return enum_class -TwoFlavor = _makeFlavorScheme('TwoFlavor',['E','X']) -ThreeFlavor = _makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) -FourFlavor = _makeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) +TwoFlavor = makeFlavorScheme('TwoFlavor',['E','X']) +ThreeFlavor = makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) +FourFlavor = makeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) class FlavorMatrix: @@ -87,8 +84,8 @@ def __init__(self, self.flavors1 = flavors self.flavors2 = flavors2 or flavors expected_shape = (len(self.flavors2), len(self.flavors1)) - assert self.array.shape == expected_shape, \ - f"Array shape {self.array.shape} mismatch expected {expected_shape}" + if(self.array.shape != expected_shape): + raise ValueError(f"FlavorMatrix array shape {self.array.shape} mismatch expected {expected_shape}") def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): @@ -109,10 +106,25 @@ def __repr__(self): def __matmul__(self, other): if isinstance(other, FlavorMatrix): - data = np.tensordot(self.array, other.array, axes=[0,1]) - return FlavorMatrix(data, self.flavors1, other.flavors2) + data = np.tensordot(self.array, other.array, axes=[1,0]) + return FlavorMatrix(data, other.flavors1, self.flavors2) raise TypeError(f"Cannot multiply object of {self.__class__} by {other.__class__}") - + #properties + @property + def shape(self): + return self.array.shape + @property + def flavors_right(self): + return self.flavors1 + @property + def flavors_left(self): + return self.flavors2 + @property + def flavors_from(self): + return self.flavors_right + @property + def flavors_to(self): + return self.flavors_left #methods for creating matrix @classmethod def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): @@ -147,7 +159,9 @@ def _decorator(func): return FlavorMatrix.from_function(flavor1, flavor2, function=func) return _decorator -#define some conversion matrices +#define the conversion matrices +M_convert = {} + @flavor_matrix(ThreeFlavor,TwoFlavor) def convert_3to2(f1,f2): if (f1.name==f2.name): @@ -164,3 +178,4 @@ def convert_2to3(f1,f2): return 1 return 0 +#other conversion matrices can be defined \ No newline at end of file From c3130f98255ec7fbefdb9ec4047e3785f72d66ba Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Sat, 4 May 2024 15:51:37 +0300 Subject: [PATCH 12/31] Make conversion matrix dict --- python/snewpy/flavor.py | 68 ++++++++++++++++++++++++++--------------- 1 file changed, 43 insertions(+), 25 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 8981d0e59..2978ff04a 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -71,7 +71,7 @@ def makeFlavorScheme(name:str, leptons:list): TwoFlavor = makeFlavorScheme('TwoFlavor',['E','X']) ThreeFlavor = makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) -FourFlavor = makeFlavorScheme('ThreeFlavor',['E','MU','TAU','S']) +FourFlavor = makeFlavorScheme('FourFlavor',['E','MU','TAU','S']) class FlavorMatrix: @@ -98,16 +98,20 @@ def __getitem__(self, index): def __setitem__(self, index, value): self.array[self._convert_index(index)] = value - + + def _repr_short(self): + return f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}> shape={self.shape}' def __repr__(self): - s = f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}>:' - s+='\n'+repr(self.array) + s=self._repr_short()+'\n'+repr(self.array) return s def __matmul__(self, other): if isinstance(other, FlavorMatrix): - data = np.tensordot(self.array, other.array, axes=[1,0]) - return FlavorMatrix(data, other.flavors1, self.flavors2) + try: + data = np.tensordot(self.array, other.array, axes=[1,0]) + return FlavorMatrix(data, other.flavors1, self.flavors2) + except: + raise ValueError(f"Cannot multiply {self._repr_short()} by {other._repr_short()}") raise TypeError(f"Cannot multiply object of {self.__class__} by {other.__class__}") #properties @property @@ -160,22 +164,36 @@ def _decorator(func): return _decorator #define the conversion matrices -M_convert = {} - -@flavor_matrix(ThreeFlavor,TwoFlavor) -def convert_3to2(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): - return 0.5 - return 0 - -@flavor_matrix(TwoFlavor,ThreeFlavor) -def convert_2to3(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): - return 1 - return 0 - -#other conversion matrices can be defined \ No newline at end of file +def _define_the_conversion_matrices(): + flavor_schemes = (TwoFlavor, ThreeFlavor, FourFlavor) + #define the basic matrices, where we can just make 1 if the names are the same + M = {fs1: + {fs2: flavor_matrix(fs1,fs2)(lambda f1,f2:1.0*(f1.name==f2.name)) + for fs2 in flavor_schemes[1:]} + for fs1 in flavor_schemes[1:] + } + #define special cases + @flavor_matrix(ThreeFlavor,TwoFlavor) + def M_3to2(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): + return 0.5 + return 0 + + @flavor_matrix(TwoFlavor,ThreeFlavor) + def M_2to3(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): + return 1 + return 0 + #override the matrix for these special cases + M[TwoFlavor] = {} + for fs2 in flavor_schemes[1:]: + M[TwoFlavor][fs2] = M[ThreeFlavor][fs2] @ M_2to3 + for fs1 in flavor_schemes[1:]: + M[fs1][TwoFlavor] = M_3to2 @ M[fs1][ThreeFlavor] + return M + +M_convert = _define_the_conversion_matrices() \ No newline at end of file From 105d466d768756bae644bae90ca98ef5b89c285e Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 10:35:30 +0300 Subject: [PATCH 13/31] Cleanup the interface: change method and argument names etc --- python/snewpy/flavor.py | 154 ++++++++++++++++------------------------ 1 file changed, 63 insertions(+), 91 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 2978ff04a..a51346628 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -61,36 +61,34 @@ def is_sterile(self): @property def lepton(self): return self.name.split('_')[1] - -def makeFlavorScheme(name:str, leptons:list): - enum_class = FlavorScheme(name, start=0, - names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']] - ) - return enum_class + @classmethod + def from_lepton_names(cls, name:str, leptons:list): + enum_class = cls(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']]) + return enum_class -TwoFlavor = makeFlavorScheme('TwoFlavor',['E','X']) -ThreeFlavor = makeFlavorScheme('ThreeFlavor',['E','MU','TAU']) -FourFlavor = makeFlavorScheme('FourFlavor',['E','MU','TAU','S']) +TwoFlavor = FlavorScheme.from_lepton_names('TwoFlavor',['E','X']) +ThreeFlavor = FlavorScheme.from_lepton_names('ThreeFlavor',['E','MU','TAU']) +FourFlavor = FlavorScheme.from_lepton_names('FourFlavor',['E','MU','TAU','S']) class FlavorMatrix: def __init__(self, array:np.ndarray, - flavors:FlavorScheme, - flavors2:FlavorScheme = None + flavor:FlavorScheme, + from_flavor:FlavorScheme = None ): self.array = np.asarray(array) - self.flavors1 = flavors - self.flavors2 = flavors2 or flavors - expected_shape = (len(self.flavors2), len(self.flavors1)) + self.flavor_out = flavor + self.flavor_in = from_flavor or flavor + expected_shape = (len(self.flavor_out), len(self.flavor_in)) if(self.array.shape != expected_shape): raise ValueError(f"FlavorMatrix array shape {self.array.shape} mismatch expected {expected_shape}") - + def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): index = [index] - new_idx = [flavors[idx] for idx,flavors in zip(index, [self.flavors2,self.flavors1])] + new_idx = [flavors[idx] for idx,flavors in zip(index, self.flavors)] return tuple(new_idx) def __getitem__(self, index): @@ -100,7 +98,8 @@ def __setitem__(self, index, value): self.array[self._convert_index(index)] = value def _repr_short(self): - return f'{self.__class__.__name__}:<{self.flavors1.__name__}->{self.flavors2.__name__}> shape={self.shape}' + return f'{self.__class__.__name__}:<{self.flavor_in.__name__}->{self.flavor_out.__name__}> shape={self.shape}' + def __repr__(self): s=self._repr_short()+'\n'+repr(self.array) return s @@ -109,91 +108,64 @@ def __matmul__(self, other): if isinstance(other, FlavorMatrix): try: data = np.tensordot(self.array, other.array, axes=[1,0]) - return FlavorMatrix(data, other.flavors1, self.flavors2) - except: - raise ValueError(f"Cannot multiply {self._repr_short()} by {other._repr_short()}") + return FlavorMatrix(data, self.flavor_out, from_flavor = other.flavor_in) + except Exception as e: + raise ValueError(f"Cannot multiply {self._repr_short()} by {other._repr_short()}") from e + raise TypeError(f"Cannot multiply object of {self.__class__} by {other.__class__}") #properties @property def shape(self): return self.array.shape @property - def flavors_right(self): - return self.flavors1 - @property - def flavors_left(self): - return self.flavors2 - @property - def flavors_from(self): - return self.flavors_right - @property - def flavors_to(self): - return self.flavors_left - #methods for creating matrix + def flavor(self): + return self.flavor_out + @classmethod - def zeros(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): - flavors2 = flavors2 or flavors1 - shape = (len(flavors2), len(flavors1)) + def zeros(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): + from_flavor = from_flavor or flavor + shape = (len(from_flavor), len(flavor)) data = np.zeros(shape) - return cls(data, flavors1, flavors2) + return cls(data, flavor, from_flavor) @classmethod - def eye(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): - flavors2 = flavors2 or flavors1 - shape = (len(flavors2), len(flavors1)) + def eye(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): + from_flavor = from_flavor or flavor + shape = (len(from_flavor), len(flavor)) data = np.eye(*shape) - return cls(data, flavors1, flavors2) + return cls(data, flavor, from_flavor) @classmethod - def from_function(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None, *, function): - flavors2 = flavors2 or flavors1 - data = [[function(f1,f2) - for f1 in flavors1] - for f2 in flavors2] - return cls(np.array(data,dtype=float), flavors1, flavors2) - @classmethod - def identity(cls, flavors1:FlavorScheme, flavors2:FlavorScheme = None): - return cls.from_function(flavors1, flavors2, lambda f1,f2: 1.*(f1.name==f2.name)) - - -def flavor_matrix(flavor1:FlavorScheme, flavor2:FlavorScheme=None): - """A decorator for creating the flavor matrix from the given function""" - flavor2 = flavor2 or flavor1 - def _decorator(func): - return FlavorMatrix.from_function(flavor1, flavor2, function=func) - return _decorator + def from_function(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): + """A decorator for creating the flavor matrix from the given function""" + from_flavor = from_flavor or flavor + def _decorator(function): + data = [[function(f1,f2) + for f2 in from_flavor] + for f1 in flavor] + + return cls(np.array(data,dtype=float), flavor, from_flavor) + return _decorator -#define the conversion matrices -def _define_the_conversion_matrices(): - flavor_schemes = (TwoFlavor, ThreeFlavor, FourFlavor) - #define the basic matrices, where we can just make 1 if the names are the same - M = {fs1: - {fs2: flavor_matrix(fs1,fs2)(lambda f1,f2:1.0*(f1.name==f2.name)) - for fs2 in flavor_schemes[1:]} - for fs1 in flavor_schemes[1:] - } - #define special cases - @flavor_matrix(ThreeFlavor,TwoFlavor) - def M_3to2(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): - return 0.5 - return 0 - - @flavor_matrix(TwoFlavor,ThreeFlavor) - def M_2to3(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): - return 1 - return 0 - #override the matrix for these special cases - M[TwoFlavor] = {} - for fs2 in flavor_schemes[1:]: - M[TwoFlavor][fs2] = M[ThreeFlavor][fs2] @ M_2to3 - for fs1 in flavor_schemes[1:]: - M[fs1][TwoFlavor] = M_3to2 @ M[fs1][ThreeFlavor] - return M - -M_convert = _define_the_conversion_matrices() \ No newline at end of file + @classmethod + def conversion_matrix(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): + from_flavor = from_flavor or flavor + if(flavor==TwoFlavor): + #define special cases + @FlavorMatrix.from_function(flavor, from_flavor) + def convert_Nto2(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): + return 0.5 + return 0 + return convert_Nto2 + else: + @FlavorMatrix.from_function(flavor, from_flavor) + def convert_2toN(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): + return 1 + return 0 + return convert_2toN \ No newline at end of file From 0efaf6f3e0d78eee6d84d88afdeb1f358bbc20bb Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 10:36:15 +0300 Subject: [PATCH 14/31] Update tests --- python/snewpy/test/test_flavors.py | 36 ++++++++++++++++++++++-------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/python/snewpy/test/test_flavors.py b/python/snewpy/test/test_flavors.py index 87a1dc94f..9e0ecdd87 100644 --- a/python/snewpy/test/test_flavors.py +++ b/python/snewpy/test/test_flavors.py @@ -1,7 +1,9 @@ import pytest import numpy as np -from snewpy import flavor -from snewpy.flavor import TwoFlavor,ThreeFlavor,FourFlavor +import snewpy.flavor +from snewpy.flavor import TwoFlavor,ThreeFlavor,FourFlavor, FlavorMatrix, FlavorScheme + +flavor_schemes = TwoFlavor,ThreeFlavor,FourFlavor class TestFlavorScheme: @staticmethod @@ -11,6 +13,7 @@ def test_flavor_scheme_lengths(): assert len(FourFlavor)==8 @staticmethod + def test_getitem_string(): assert TwoFlavor['NU_E'] == TwoFlavor.NU_E assert TwoFlavor['NU_X'] == TwoFlavor.NU_X @@ -31,7 +34,7 @@ def test_values_from_different_enums(): @staticmethod def test_makeFlavorScheme(): - TestFlavor = flavor.makeFlavorScheme('TestFlavor',['A','B','C']) + TestFlavor = FlavorScheme.from_lepton_names('TestFlavor',leptons=['A','B','C']) assert len(TestFlavor)==6 assert [f.name for f in TestFlavor]==['NU_A','NU_A_BAR','NU_B','NU_B_BAR','NU_C','NU_C_BAR'] @@ -82,16 +85,31 @@ def test_flavor_properties(): class TestFlavorMatrix: @staticmethod def test_init_square_matrix(): - m = flavor.FlavorMatrix(array=np.ones(shape=(4,4)), flavors=TwoFlavor) + m = FlavorMatrix(array=np.ones(shape=(4,4)), flavor=TwoFlavor) assert m.shape == (4,4) - assert m.flavors_from == TwoFlavor - assert m.flavors_to == TwoFlavor + assert m.flavor_in == TwoFlavor + assert m.flavor_out == TwoFlavor @staticmethod def test_init_square_matrix_with_wrong_shape_raises_ValueError(): with pytest.raises(ValueError): - m = flavor.FlavorMatrix(array=np.ones(shape=(4,5)), flavors=TwoFlavor) + m = FlavorMatrix(array=np.ones(shape=(4,5)), flavor=TwoFlavor) with pytest.raises(ValueError): - m = flavor.FlavorMatrix(array=np.ones(shape=(5,5)), flavors=TwoFlavor) + m = FlavorMatrix(array=np.ones(shape=(5,5)), flavor=TwoFlavor) with pytest.raises(ValueError): - m = flavor.FlavorMatrix(array=np.ones(shape=(5,4)), flavors=TwoFlavor) \ No newline at end of file + m = FlavorMatrix(array=np.ones(shape=(5,4)), flavor=TwoFlavor) + + @staticmethod + def test_conversion_matrices_for_same_flavor_are_unity(): + for flavor in [TwoFlavor,ThreeFlavor,FourFlavor]: + matrix = FlavorMatrix.conversion_matrix(flavor,flavor) + print(matrix) + assert np.allclose(matrix.array, np.eye(len(flavor))) + + @staticmethod + @pytest.mark.parametrize('flavor_in',flavor_schemes) + @pytest.mark.parametrize('flavor_out',flavor_schemes) + def test_conversion_matrices(flavor_in, flavor_out): + M = FlavorMatrix.conversion_matrix(flavor_out,flavor_in) + assert M.flavor_in == flavor_in + assert M.flavor_out == flavor_out \ No newline at end of file From 5f200eeaddda24299df5c808a1046cc2308defac Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 11:10:17 +0300 Subject: [PATCH 15/31] Implement slicing for the flavors --- python/snewpy/flavor.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index a51346628..196072748 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -21,7 +21,8 @@ def __getitem__(cls, key): f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' ) - #if this is an int value: find a matching + #if this is anything else - treat it as a slice + return np.array(list(cls.__members__.values()),dtype=object)[key] class FlavorScheme(enum.IntEnum, metaclass=EnumMeta): def to_tex(self): @@ -66,6 +67,9 @@ def lepton(self): def from_lepton_names(cls, name:str, leptons:list): enum_class = cls(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']]) return enum_class + @classmethod + def take(cls, index): + return cls[index] TwoFlavor = FlavorScheme.from_lepton_names('TwoFlavor',['E','X']) ThreeFlavor = FlavorScheme.from_lepton_names('ThreeFlavor',['E','MU','TAU']) From 3a7e2f50a9caa09822e2058bc39684a41f0bc260 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 11:29:47 +0300 Subject: [PATCH 16/31] Using FlavorScheme in the flux.Container --- python/snewpy/flux.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 0fb330374..4a170589b 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -53,6 +53,7 @@ """ from typing import Union, Optional, Set, List from snewpy.neutrino import Flavor +from snewpy.flavor import FlavorScheme from astropy import units as u import numpy as np @@ -121,8 +122,15 @@ def __init__(self, self.array = u.Quantity(data) self.time = u.Quantity(time, ndmin=1) self.energy = u.Quantity(energy, ndmin=1) - self.flavor = np.sort(np.array(flavor, ndmin=1)) + if isinstance(flavor, type): + if issubclass(flavor, FlavorScheme): + self.flavor = flavor + else: + raise TypeError(f"Wrong type of flavor={flavor}, should be a FlavorScheme") + else: + self.flavor = np.sort(np.array(flavor, ndmin=1, dtype=object)) + Nf,Nt,Ne = len(self.flavor), len(self.time), len(self.energy) #list all valid shapes of the input array expected_shapes=[(nf,nt,ne) for nf in (Nf,Nf-1) for nt in (Nt,Nt-1) for ne in (Ne,Ne-1)] @@ -162,11 +170,14 @@ def shape(self): def __getitem__(self, args)->'Container': """Slice the flux array and produce a new Flux object""" + if not isinstance(args, tuple): + args = tuple([args],) + print(args) try: iter(args) except TypeError: args = [args] - args = [a if isinstance(a, slice) else slice(a, a + 1) for a in args] + #args = [a if isinstance(a, slice) else slice(a, a + 1) for a in args] #expand args to match axes args+=[slice(None)]*(len(Axes)-len(args)) array = self.array.__getitem__(tuple(args)) @@ -175,7 +186,8 @@ def __getitem__(self, args)->'Container': def __repr__(self) -> str: """print information about the container""" - s = [f"{len(values)} {label.name}({values.min()};{values.max()})" + s = [f"{len(values)} {label.name}({values.min()};{values.max()})" if isinstance(values, np.ndarray) + else f"{len(values)} {label.name}({repr(values)})" for label, values in zip(Axes,self.axes)] return f"{self.__class__.__name__} {self.array.shape} [{self.array.unit}]: <{' x '.join(s)}>" @@ -329,8 +341,9 @@ def _save_quantity(name): except: return {name:values} data_dict = {} - for name in ['array','time','energy','flavor']: + for name in ['array','time','energy']: data_dict.update(_save_quantity(name)) + data_dict['flavor'] = np.array(self.flavor, dtype=object) np.savez(fname, _class_name=self.__class__.__name__, **data_dict, @@ -340,7 +353,7 @@ def _save_quantity(name): @classmethod def load(cls, fname:str)->'Container': """Load container from a given file""" - with np.load(fname) as f: + with np.load(fname, allow_pickle=True) as f: def _load_quantity(name): array = f[name] try: @@ -361,7 +374,8 @@ def __eq__(self, other:'Container')->bool: result = self.__class__==other.__class__ and \ self.unit == other.unit and \ np.allclose(self.array, other.array) and \ - all([np.allclose(self.axes[ax], other.axes[ax]) for ax in Axes]) + all(self.flavor==other.flavor) and \ + all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]) return result class Container(_ContainerBase): From e44422c3f3d7ee119cd028142d072f189ccb9315 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 11:32:24 +0300 Subject: [PATCH 17/31] Remove trailing character --- python/snewpy/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 4a170589b..36b235646 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -374,7 +374,7 @@ def __eq__(self, other:'Container')->bool: result = self.__class__==other.__class__ and \ self.unit == other.unit and \ np.allclose(self.array, other.array) and \ - all(self.flavor==other.flavor) and \ + all(self.flavor==other.flavor) and \ all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]) return result From 017faeb4c4060dbcef007275edd4e697005074a7 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 11:47:39 +0300 Subject: [PATCH 18/31] Update the flavor checks --- python/snewpy/flux.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 36b235646..3c60ff27e 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -123,13 +123,27 @@ def __init__(self, self.time = u.Quantity(time, ndmin=1) self.energy = u.Quantity(energy, ndmin=1) - if isinstance(flavor, type): + if not isinstance(flavor, type): + #convert to array + flavor_array = np.sort(np.array(flavor, ndmin=1, dtype=object)) + #try to get the flavor schemes + flavor_schemes = set(f.__class__ for f in flavor_array) + if len(flavor_schemes)!=1: + raise ValueError(f"Flavors {flavor} must be from a single flavor scheme, but are from {flavor_schemes}") + flavor_scheme = list(flavor_schemes)[0] + #check if we have the full set of flavors from this scheme + if(len(flavor_array)==len(flavor_scheme))and all(flavor_array==flavor_scheme[:]): + self.flavor = flavor_scheme + else: + self.flavor = flavor_array + + else: if issubclass(flavor, FlavorScheme): self.flavor = flavor else: raise TypeError(f"Wrong type of flavor={flavor}, should be a FlavorScheme") - else: - self.flavor = np.sort(np.array(flavor, ndmin=1, dtype=object)) + + Nf,Nt,Ne = len(self.flavor), len(self.time), len(self.energy) #list all valid shapes of the input array @@ -374,7 +388,8 @@ def __eq__(self, other:'Container')->bool: result = self.__class__==other.__class__ and \ self.unit == other.unit and \ np.allclose(self.array, other.array) and \ - all(self.flavor==other.flavor) and \ + type(self.flavor[0])==type(other.flavor[0]) and \ + all(self.flavor[:]==other.flavor) and \ all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]) return result From 6f01ceffc9c2e77ae5cec0ab36ddb876cf73277d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 11:48:44 +0300 Subject: [PATCH 19/31] Update the flavor checks --- python/snewpy/flux.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 3c60ff27e..7e47463bc 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -52,8 +52,8 @@ """ from typing import Union, Optional, Set, List -from snewpy.neutrino import Flavor -from snewpy.flavor import FlavorScheme +#from snewpy.neutrino import Flavor +from snewpy.flavor import FlavorScheme, FlavorMatrix from astropy import units as u import numpy as np @@ -344,6 +344,9 @@ def __mul__(self, factor) -> 'Container': axes = list(self.axes) return Container(array, *axes) + def __rmatmul__(self, matrix:FlavorMatrix) -> 'Container': + + def save(self, fname:str)->None: """Save container data to a given file (using `numpy.savez`)""" def _save_quantity(name): From d8cf66771d0915a8829e33476229a42ebf3f756c Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 12:29:23 +0300 Subject: [PATCH 20/31] Using >> and << operators for generating conversion matrices --- python/snewpy/flavor.py | 54 +++++++++++++++++------------- python/snewpy/test/test_flavors.py | 8 +++-- 2 files changed, 35 insertions(+), 27 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 196072748..cb0c303b5 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -67,6 +67,7 @@ def lepton(self): def from_lepton_names(cls, name:str, leptons:list): enum_class = cls(name, start=0, names = [f'NU_{L}{BAR}' for L in leptons for BAR in ['','_BAR']]) return enum_class + @classmethod def take(cls, index): return cls[index] @@ -75,7 +76,6 @@ def take(cls, index): ThreeFlavor = FlavorScheme.from_lepton_names('ThreeFlavor',['E','MU','TAU']) FourFlavor = FlavorScheme.from_lepton_names('FourFlavor',['E','MU','TAU','S']) - class FlavorMatrix: def __init__(self, array:np.ndarray, @@ -107,7 +107,9 @@ def _repr_short(self): def __repr__(self): s=self._repr_short()+'\n'+repr(self.array) return s - + def __eq__(self,other): + return self.flavor_in==other.flavor_in and self.flavor_out==other.flavor_out and np.allclose(self.array,other.array) + def __matmul__(self, other): if isinstance(other, FlavorMatrix): try: @@ -150,26 +152,30 @@ def _decorator(function): return cls(np.array(data,dtype=float), flavor, from_flavor) return _decorator + #flavor conversion utils + +def conversion_matrix(from_flavor:FlavorScheme, to_flavor:FlavorScheme): + print(from_flavor, to_flavor) + if(from_flavor==TwoFlavor): + #define special cases + @FlavorMatrix.from_function(to_flavor, from_flavor) + def convert_2toN(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): + return 0.5 + return 0 + return convert_2toN + else: + @FlavorMatrix.from_function(to_flavor, from_flavor) + def convert_Nto2(f1,f2): + if (f1.name==f2.name): + return 1. + if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): + return 1. + return 0. + return convert_Nto2 - @classmethod - def conversion_matrix(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): - from_flavor = from_flavor or flavor - if(flavor==TwoFlavor): - #define special cases - @FlavorMatrix.from_function(flavor, from_flavor) - def convert_Nto2(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): - return 0.5 - return 0 - return convert_Nto2 - else: - @FlavorMatrix.from_function(flavor, from_flavor) - def convert_2toN(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): - return 1 - return 0 - return convert_2toN \ No newline at end of file +FlavorScheme.conversion_matrix = classmethod(conversion_matrix) +EnumMeta.__rshift__ = conversion_matrix +EnumMeta.__lshift__ = lambda f1,f2:conversion_matrix(f2,f1) diff --git a/python/snewpy/test/test_flavors.py b/python/snewpy/test/test_flavors.py index 9e0ecdd87..ce7bd0ba5 100644 --- a/python/snewpy/test/test_flavors.py +++ b/python/snewpy/test/test_flavors.py @@ -102,14 +102,16 @@ def test_init_square_matrix_with_wrong_shape_raises_ValueError(): @staticmethod def test_conversion_matrices_for_same_flavor_are_unity(): for flavor in [TwoFlavor,ThreeFlavor,FourFlavor]: - matrix = FlavorMatrix.conversion_matrix(flavor,flavor) - print(matrix) + matrix = flavor>>flavor + assert isinstance(matrix, FlavorMatrix) assert np.allclose(matrix.array, np.eye(len(flavor))) @staticmethod @pytest.mark.parametrize('flavor_in',flavor_schemes) @pytest.mark.parametrize('flavor_out',flavor_schemes) def test_conversion_matrices(flavor_in, flavor_out): - M = FlavorMatrix.conversion_matrix(flavor_out,flavor_in) + M = flavor_in>>flavor_out + assert M==flavor_out< Date: Fri, 10 May 2024 15:08:38 +0300 Subject: [PATCH 21/31] Deriving the flavor_scheme --- python/snewpy/flux.py | 81 +++++++++++++++++++++---------------------- 1 file changed, 39 insertions(+), 42 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 7e47463bc..e6d014589 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -87,11 +87,12 @@ class _ContainerBase: unit = None def __init__(self, data: u.Quantity, - flavor: List[Flavor], + flavor: List[FlavorScheme], time: u.Quantity[u.s], energy: u.Quantity[u.MeV], *, - integrable_axes: Optional[Set[Axes]] = None + integrable_axes: Optional[Set[Axes]] = None, + flavor_scheme:Optional[FlavorScheme] = None ): """A container class storing the physical quantity (flux, fluence, rate...), which depends on flavor, time and energy. @@ -113,7 +114,11 @@ def __init__(self, integrable_axes: set of :class:`Axes` or None List of axes which can be integrated. - If None (default) this set will be derived from the axes shapes + If None (default) this set will be derived from the axes shapes + + flavor_scheme: a subclass of :class:`snewpy.flavor.FlavorSchemes` or None + A class which lists all the allowed flavors. + If None (default) this value will be retrieved from the ``flavor`` arguemnt. """ if self.unit is not None: #try to convert to the unit @@ -122,28 +127,15 @@ def __init__(self, self.array = u.Quantity(data) self.time = u.Quantity(time, ndmin=1) self.energy = u.Quantity(energy, ndmin=1) - - if not isinstance(flavor, type): - #convert to array - flavor_array = np.sort(np.array(flavor, ndmin=1, dtype=object)) - #try to get the flavor schemes - flavor_schemes = set(f.__class__ for f in flavor_array) - if len(flavor_schemes)!=1: + self.flavor = np.array(flavor,ndmin=1, dtype=object) + self.flavor_scheme = flavor_scheme + #define the flavor scheme + if not flavor_scheme: + flavor_schemes = set(f.__class__ for f in self.flavor) + if len(flavor_schemes)>1: raise ValueError(f"Flavors {flavor} must be from a single flavor scheme, but are from {flavor_schemes}") - flavor_scheme = list(flavor_schemes)[0] - #check if we have the full set of flavors from this scheme - if(len(flavor_array)==len(flavor_scheme))and all(flavor_array==flavor_scheme[:]): - self.flavor = flavor_scheme - else: - self.flavor = flavor_array - - else: - if issubclass(flavor, FlavorScheme): - self.flavor = flavor - else: - raise TypeError(f"Wrong type of flavor={flavor}, should be a FlavorScheme") - - + elif len(flavor_schemes)==1: + self.flavor_scheme = flavor_schemes.pop() Nf,Nt,Ne = len(self.flavor), len(self.time), len(self.energy) #list all valid shapes of the input array @@ -185,24 +177,27 @@ def shape(self): def __getitem__(self, args)->'Container': """Slice the flux array and produce a new Flux object""" if not isinstance(args, tuple): - args = tuple([args],) - print(args) - try: - iter(args) - except TypeError: - args = [args] - #args = [a if isinstance(a, slice) else slice(a, a + 1) for a in args] - #expand args to match axes - args+=[slice(None)]*(len(Axes)-len(args)) - array = self.array.__getitem__(tuple(args)) - newaxes = [ax.__getitem__(arg) for arg, ax in zip(args, self.axes)] + args = list([args],) + + arg_slices = [slice(None)]*len(Axes) + if isinstance(args[0],str) or isinstance(args[0],FlavorScheme): + args[0] = self.flavor_scheme[args[0]] + for n,arg in enumerate(args): + if not isinstance(arg, slice): + arg = slice(arg, arg + 1) + arg_slices[n] = arg + + array = self.array.__getitem__(tuple(arg_slices)) + newaxes = [ax.__getitem__(arg) for arg, ax in zip(arg_slices, self.axes)] return self.__class__(array, *newaxes) def __repr__(self) -> str: """print information about the container""" - s = [f"{len(values)} {label.name}({values.min()};{values.max()})" if isinstance(values, np.ndarray) - else f"{len(values)} {label.name}({repr(values)})" - for label, values in zip(Axes,self.axes)] + s = [f"{len(values)} {label.name}({values.min()};{values.max()})" + if label!=Axes.flavor + else f"{len(values)} {label.name}[{self.flavor_scheme.__name__}]({values.min()};{values.max()})" + for label, values in zip(Axes,self.axes) + ] return f"{self.__class__.__name__} {self.array.shape} [{self.array.unit}]: <{' x '.join(s)}>" def sum(self, axis: Union[Axes,str])->'Container': @@ -344,8 +339,9 @@ def __mul__(self, factor) -> 'Container': axes = list(self.axes) return Container(array, *axes) - def __rmatmul__(self, matrix:FlavorMatrix) -> 'Container': - + def convert_to_flavor_scheme(new_flavor:FlavorScheme): + M = self.flavor.conversion_matrix(new_flavor) + return M @ self def save(self, fname:str)->None: """Save container data to a given file (using `numpy.savez`)""" @@ -391,8 +387,9 @@ def __eq__(self, other:'Container')->bool: result = self.__class__==other.__class__ and \ self.unit == other.unit and \ np.allclose(self.array, other.array) and \ - type(self.flavor[0])==type(other.flavor[0]) and \ - all(self.flavor[:]==other.flavor) and \ + self.flavor_scheme==other.flavor_scheme and \ + len(self.flavor)==len(other.flavor) and \ + all(self.flavor==other.flavor) and \ all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]) return result From 168b22c8879edc7ef10c890a5478a58cdbabaf84 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 15:33:59 +0300 Subject: [PATCH 22/31] Fixing the chechs in getitem --- python/snewpy/flux.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index e6d014589..8ca71b7ef 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -129,14 +129,17 @@ def __init__(self, self.energy = u.Quantity(energy, ndmin=1) self.flavor = np.array(flavor,ndmin=1, dtype=object) self.flavor_scheme = flavor_scheme - #define the flavor scheme if not flavor_scheme: - flavor_schemes = set(f.__class__ for f in self.flavor) - if len(flavor_schemes)>1: - raise ValueError(f"Flavors {flavor} must be from a single flavor scheme, but are from {flavor_schemes}") - elif len(flavor_schemes)==1: - self.flavor_scheme = flavor_schemes.pop() - + #guess the flavor scheme + if isinstance(flavor, type) and issubclass(flavor, FlavorScheme): + self.flavor_scheme = flavor + else: + #get schemes from the data + flavor_schemes = set(f.__class__ for f in self.flavor) + if len(flavor_schemes)!=1: + raise ValueError(f"Flavors {flavor} must be from a single flavor scheme, but are from {flavor_schemes}") + else: + self.flavor_scheme = flavor_schemes.pop() Nf,Nt,Ne = len(self.flavor), len(self.time), len(self.energy) #list all valid shapes of the input array expected_shapes=[(nf,nt,ne) for nf in (Nf,Nf-1) for nt in (Nt,Nt-1) for ne in (Ne,Ne-1)] @@ -176,9 +179,9 @@ def shape(self): def __getitem__(self, args)->'Container': """Slice the flux array and produce a new Flux object""" - if not isinstance(args, tuple): - args = list([args],) - + if not isinstance(args,tuple): + args = [args] + args = list(args) arg_slices = [slice(None)]*len(Axes) if isinstance(args[0],str) or isinstance(args[0],FlavorScheme): args[0] = self.flavor_scheme[args[0]] @@ -189,13 +192,13 @@ def __getitem__(self, args)->'Container': array = self.array.__getitem__(tuple(arg_slices)) newaxes = [ax.__getitem__(arg) for arg, ax in zip(arg_slices, self.axes)] - return self.__class__(array, *newaxes) + return self.__class__(array, *newaxes, flavor_scheme=self.flavor_scheme) def __repr__(self) -> str: """print information about the container""" s = [f"{len(values)} {label.name}({values.min()};{values.max()})" if label!=Axes.flavor - else f"{len(values)} {label.name}[{self.flavor_scheme.__name__}]({values.min()};{values.max()})" + else f"{len(values)} {label.name}[{self.flavor_scheme}]({values.min()};{values.max()})" for label, values in zip(Axes,self.axes) ] return f"{self.__class__.__name__} {self.array.shape} [{self.array.unit}]: <{' x '.join(s)}>" From ffaef43ad71b9e1f6c31a9e4bd6b8a9cca657321 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 15:34:26 +0300 Subject: [PATCH 23/31] Adding tests for access by value and by string --- python/snewpy/test/test_flux_container.py | 31 +++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/python/snewpy/test/test_flux_container.py b/python/snewpy/test/test_flux_container.py index 35f1c6e14..5c2337c2c 100644 --- a/python/snewpy/test/test_flux_container.py +++ b/python/snewpy/test/test_flux_container.py @@ -4,6 +4,7 @@ from snewpy.flux import Container,Axes from snewpy.neutrino import Flavor +from snewpy.flavor import TwoFlavor,ThreeFlavor,FourFlavor import numpy as np import astropy.units as u import pytest @@ -29,13 +30,24 @@ def quantities(draw, values, unit): return draw(values)< Date: Fri, 10 May 2024 15:54:08 +0300 Subject: [PATCH 24/31] Matrix mutiplication for the flux --- python/snewpy/flavor.py | 3 ++- python/snewpy/flux.py | 21 +++++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index cb0c303b5..b25f52242 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -117,7 +117,8 @@ def __matmul__(self, other): return FlavorMatrix(data, self.flavor_out, from_flavor = other.flavor_in) except Exception as e: raise ValueError(f"Cannot multiply {self._repr_short()} by {other._repr_short()}") from e - + elif hasattr(other, '__rmatmul__'): + return other.__rmatmul__(self) raise TypeError(f"Cannot multiply object of {self.__class__} by {other.__class__}") #properties @property diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 8ca71b7ef..66eb4bb1a 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -341,10 +341,6 @@ def __mul__(self, factor) -> 'Container': array = self.array*factor axes = list(self.axes) return Container(array, *axes) - - def convert_to_flavor_scheme(new_flavor:FlavorScheme): - M = self.flavor.conversion_matrix(new_flavor) - return M @ self def save(self, fname:str)->None: """Save container data to a given file (using `numpy.savez`)""" @@ -396,6 +392,23 @@ def __eq__(self, other:'Container')->bool: all([np.allclose(self.axes[ax], other.axes[ax]) for ax in list(Axes)[1:]]) return result + def _is_full_flavor(self): + return all(self.flavor==list(self.flavor_scheme)) + + def convert_to_flavor(self, flavor:FlavorScheme): + if(self.flavor_scheme==flavor): + return self + return (self.flavor_scheme>>flavor)@self + def __rshift__(self, flavor:FlavorScheme): + return self.convert_to_flavor(flavor) + + def __rmatmul__(self, matrix:FlavorMatrix): + if not self._is_full_flavor(): + raise RuntimeError(f"Cannot multiply flavor matrix object {self}, expected {len(self.flavor_scheme)} flavors") + if matrix.flavor_in!=self.flavor_scheme: + raise ValueError(f"Cannot multiply flavor matrix {matrix} by {self} - flavor scheme mismatch!") + array = np.tensordot(matrix.array,self.array, axes=[1,0]) + return Container(array, flavor=matrix.flavor_out, time=self.time, energy=self.energy) class Container(_ContainerBase): #a dictionary holding classes for each unit _unit_classes = {} From f9a2ff8cf876b4bd36dd87b5687e811102c9a03d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 16:05:11 +0300 Subject: [PATCH 25/31] Remove debug prints --- python/snewpy/flavor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index b25f52242..49bd8bc71 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -156,7 +156,6 @@ def _decorator(function): #flavor conversion utils def conversion_matrix(from_flavor:FlavorScheme, to_flavor:FlavorScheme): - print(from_flavor, to_flavor) if(from_flavor==TwoFlavor): #define special cases @FlavorMatrix.from_function(to_flavor, from_flavor) From b52f8ee26c487ed9a0de3b4305fc05e346e95cb3 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 16:49:14 +0300 Subject: [PATCH 26/31] initialize flux with the flavor enum, not its sorted version --- python/snewpy/models/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/snewpy/models/base.py b/python/snewpy/models/base.py index 0d14e30f4..0d81efc0c 100644 --- a/python/snewpy/models/base.py +++ b/python/snewpy/models/base.py @@ -190,8 +190,8 @@ def get_flux (self, t, E, distance, flavor_xform=NoTransformation()): factor = 1/(4*np.pi*(distance.to('cm'))**2) f = self.get_transformed_spectra(t, E, flavor_xform) - array = np.stack([f[flv] for flv in sorted(Flavor)]) - return Flux(data=array*factor, flavor=np.sort(Flavor), time=t, energy=E) + array = np.stack([f[flv] for flv in Flavor]) + return Flux(data=array*factor, flavor=Flavor, time=t, energy=E) From 7b300547160c91a926a5d0984e717f66bd57b035 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Fri, 10 May 2024 18:24:54 +0300 Subject: [PATCH 27/31] Correcting the Three->Two flavor conversion --- python/snewpy/flavor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 49bd8bc71..0f760c4a9 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -163,7 +163,7 @@ def convert_2toN(f1,f2): if (f1.name==f2.name): return 1. if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): - return 0.5 + return 1. return 0 return convert_2toN else: @@ -172,7 +172,7 @@ def convert_Nto2(f1,f2): if (f1.name==f2.name): return 1. if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): - return 1. + return 0.5 return 0. return convert_Nto2 From 00aa2ea259922e644b96b84100acd5ef955f9183 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 14 May 2024 13:37:13 +0300 Subject: [PATCH 28/31] Update python/snewpy/flavor.py More compact flavor matrix definition Co-authored-by: Jost Migenda --- python/snewpy/flavor.py | 31 ++++++++++++------------------- 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 0f760c4a9..4aca540e7 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -156,25 +156,18 @@ def _decorator(function): #flavor conversion utils def conversion_matrix(from_flavor:FlavorScheme, to_flavor:FlavorScheme): - if(from_flavor==TwoFlavor): - #define special cases - @FlavorMatrix.from_function(to_flavor, from_flavor) - def convert_2toN(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f2.lepton=='X' and f1.lepton in ['MU','TAU']): - return 1. - return 0 - return convert_2toN - else: - @FlavorMatrix.from_function(to_flavor, from_flavor) - def convert_Nto2(f1,f2): - if (f1.name==f2.name): - return 1. - if (f1.is_neutrino==f2.is_neutrino)and(f1.lepton=='X' and f2.lepton in ['MU','TAU']): - return 0.5 - return 0. - return convert_Nto2 + @FlavorMatrix.from_function(to_flavor, from_flavor) + def convert(f1, f2): + if f1.name == f2.name: + return 1. + if (f1.is_neutrino == f2.is_neutrino) and (f2.lepton == 'X' and f1.lepton in ['MU', 'TAU']): + # convert from TwoFlavor to more flavors + return 1. + if (f1.is_neutrino == f2.is_neutrino) and (f1.lepton == 'X' and f2.lepton in ['MU', 'TAU']): + # convert from more flavors to TwoFlavor + return 0.5 + return 0. + return convert FlavorScheme.conversion_matrix = classmethod(conversion_matrix) EnumMeta.__rshift__ = conversion_matrix From 76f150ad8dc0d65fda84584857ff52c1a83e0d6d Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 14 May 2024 14:58:58 +0300 Subject: [PATCH 29/31] Add tests for the FlavorMatrix.__getitem__ --- python/snewpy/test/test_flavors.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/snewpy/test/test_flavors.py b/python/snewpy/test/test_flavors.py index ce7bd0ba5..ead9e0dc6 100644 --- a/python/snewpy/test/test_flavors.py +++ b/python/snewpy/test/test_flavors.py @@ -89,6 +89,28 @@ def test_init_square_matrix(): assert m.shape == (4,4) assert m.flavor_in == TwoFlavor assert m.flavor_out == TwoFlavor + + @staticmethod + def test_getitem(): + m = FlavorMatrix.eye(TwoFlavor,TwoFlavor) + assert m[TwoFlavor.NU_E, TwoFlavor.NU_E]==1 + assert m['NU_E','NU_E']==1 + assert m['NU_E','NU_X']==0 + assert np.allclose(m['NU_E'], [1,0,0,0]) + assert np.allclose(m['NU_E'], m['NU_E',:]) + assert np.allclose(m[:,:], m.array) + + @staticmethod + def test_setitem(): + m = FlavorMatrix.eye(TwoFlavor,TwoFlavor) + m['NU_E']=[2,3,4,5] + assert m['NU_E','NU_E']==2 + assert m['NU_E','NU_X']==4 + #check that nothing changed in other parts + assert m['NU_X','NU_E']==0 + assert m['NU_X','NU_X']==1 + m['NU_E','NU_E_BAR']=123 + assert m['NU_E','NU_E_BAR']==123 @staticmethod def test_init_square_matrix_with_wrong_shape_raises_ValueError(): From 07a090d046c1fb0a1021338e6c3f7e5eec62c2b6 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 14 May 2024 14:59:44 +0300 Subject: [PATCH 30/31] Fixing the _convert_index and __getitem__ --- python/snewpy/flavor.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/python/snewpy/flavor.py b/python/snewpy/flavor.py index 4aca540e7..9ed92eaaf 100644 --- a/python/snewpy/flavor.py +++ b/python/snewpy/flavor.py @@ -5,6 +5,8 @@ class EnumMeta(enum.EnumMeta): def __getitem__(cls, key): #if this is an iterable: apply to each value, and construct a tuple + if isinstance(key, slice): + return slice(cls[key.start],cls[key.stop],key.step) if isinstance(key, typing.Iterable) and not isinstance(key, str): return tuple(map(cls.__getitem__, key)) #if this is from a flavor scheme: check that it's ours @@ -21,6 +23,8 @@ def __getitem__(cls, key): f'Cannot find key "{key}" in {cls.__name__} sheme! Valid options are {list(cls)}' ) + if key is None: + return None #if this is anything else - treat it as a slice return np.array(list(cls.__members__.values()),dtype=object)[key] @@ -92,7 +96,7 @@ def __init__(self, def _convert_index(self, index): if isinstance(index, str) or (not isinstance(index,typing.Iterable)): index = [index] - new_idx = [flavors[idx] for idx,flavors in zip(index, self.flavors)] + new_idx = [flavors[idx] for idx,flavors in zip(index, [self.flavor_out, self.flavor_in])] return tuple(new_idx) def __getitem__(self, index): @@ -128,13 +132,6 @@ def shape(self): def flavor(self): return self.flavor_out - @classmethod - def zeros(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): - from_flavor = from_flavor or flavor - shape = (len(from_flavor), len(flavor)) - data = np.zeros(shape) - return cls(data, flavor, from_flavor) - @classmethod def eye(cls, flavor:FlavorScheme, from_flavor:FlavorScheme = None): from_flavor = from_flavor or flavor From 5b2fb84adc9e5b4a371053b5725cfcfcd7f2cf26 Mon Sep 17 00:00:00 2001 From: Andrey Sheshukov Date: Tue, 14 May 2024 15:08:24 +0300 Subject: [PATCH 31/31] Adding comment for the isinstance and issubclass checks --- python/snewpy/flux.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/snewpy/flux.py b/python/snewpy/flux.py index 66eb4bb1a..fc9469153 100644 --- a/python/snewpy/flux.py +++ b/python/snewpy/flux.py @@ -131,8 +131,9 @@ def __init__(self, self.flavor_scheme = flavor_scheme if not flavor_scheme: #guess the flavor scheme - if isinstance(flavor, type) and issubclass(flavor, FlavorScheme): - self.flavor_scheme = flavor + if isinstance(flavor, type):#issubclass without isinstance(type) check raises TypeError + if issubclass(flavor, FlavorScheme): + self.flavor_scheme = flavor else: #get schemes from the data flavor_schemes = set(f.__class__ for f in self.flavor)