From c9ddfc680a6b1ffb98196b8794de8208f454669f Mon Sep 17 00:00:00 2001 From: Gopal Garg Date: Tue, 23 Jan 2024 02:04:58 -0600 Subject: [PATCH 1/2] Updated PISA with LRI potential --- pisa/stages/osc/lri_params.py | 121 +++ pisa/stages/osc/prob3.py | 43 +- pisa/stages/osc/prob3_bk.py | 407 ++++++++ .../osc/prob3numba/numba_osc_hostfuncs.py | 17 +- .../osc/prob3numba/numba_osc_hostfuncs_bk.py | 218 +++++ .../osc/prob3numba/numba_osc_kernel_bk.py | 879 ++++++++++++++++++ .../osc/prob3numba/numba_osc_kernels.py | 25 +- pisa/stages/osc/prob3numba/numba_osc_tests.py | 15 + .../osc/prob3numba/numba_osc_tests_bk.py | 791 ++++++++++++++++ ...tate_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 392 bytes ...tate_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 493 bytes ...ecay_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 590 bytes ...ecay_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 843 bytes ..._mat_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 497 bytes ..._mat_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 650 bytes ..._vac_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 589 bytes ..._vac_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 842 bytes ..._dms_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 631 bytes ..._dms_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 888 bytes ...sition_matrix_hostfunc__nufit32_io__f4.pkl | Bin 1107 -> 1180 bytes ...sition_matrix_hostfunc__nufit32_io__f8.pkl | Bin 1596 -> 1705 bytes ...trix_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 1180 bytes ...trix_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 1705 bytes ..._matrix_hostfunc__nufit32_no_E1TeV__f4.pkl | Bin 1107 -> 1180 bytes ..._matrix_hostfunc__nufit32_no_E1TeV__f8.pkl | Bin 1596 -> 1705 bytes ...sition_matrix_hostfunc__nufit32_no__f4.pkl | Bin 1107 -> 1180 bytes ...sition_matrix_hostfunc__nufit32_no__f8.pkl | Bin 1596 -> 1705 bytes ...atrix_hostfunc__nufit32_no_blearth__f4.pkl | Bin 1107 -> 1180 bytes ...atrix_hostfunc__nufit32_no_blearth__f8.pkl | Bin 1596 -> 1705 bytes ..._matrix_hostfunc__nufit32_no_nubar__f4.pkl | Bin 1107 -> 1180 bytes ..._matrix_hostfunc__nufit32_no_nubar__f8.pkl | Bin 1596 -> 1705 bytes ...matrix_hostfunc__nufit32_std_decay__f4.pkl | Bin 1107 -> 1180 bytes ...matrix_hostfunc__nufit32_std_decay__f8.pkl | Bin 1596 -> 1705 bytes ...rix_hostfunc__nufit32_std_decay_no__f4.pkl | Bin 1107 -> 1180 bytes ...rix_hostfunc__nufit32_std_decay_no__f8.pkl | Bin 1596 -> 1705 bytes ...atrix_hostfunc__nufit32_std_nsi_no__f4.pkl | Bin 1107 -> 1180 bytes ...atrix_hostfunc__nufit32_std_nsi_no__f8.pkl | Bin 1596 -> 1705 bytes ...atrix_hostfunc__nufit32_vac_nsi_no__f4.pkl | Bin 1107 -> 1180 bytes ...atrix_hostfunc__nufit32_vac_nsi_no__f8.pkl | Bin 1596 -> 1705 bytes ...asis_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 722 bytes ...asis_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 1019 bytes ...duct_hostfunc__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 832 bytes ...duct_hostfunc__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 1269 bytes .../propagate_scalar__nufit32_io__f4.pkl | Bin 936 -> 1009 bytes .../propagate_scalar__nufit32_io__f8.pkl | Bin 1317 -> 1426 bytes ...pagate_scalar__nufit32_lri_std_mat__f4.pkl | Bin 0 -> 1009 bytes ...pagate_scalar__nufit32_lri_std_mat__f8.pkl | Bin 0 -> 1426 bytes ...propagate_scalar__nufit32_no_E1TeV__f4.pkl | Bin 936 -> 1009 bytes ...propagate_scalar__nufit32_no_E1TeV__f8.pkl | Bin 1317 -> 1426 bytes .../propagate_scalar__nufit32_no__f4.pkl | Bin 936 -> 1009 bytes .../propagate_scalar__nufit32_no__f8.pkl | Bin 1317 -> 1426 bytes ...opagate_scalar__nufit32_no_blearth__f4.pkl | Bin 936 -> 1009 bytes ...opagate_scalar__nufit32_no_blearth__f8.pkl | Bin 1317 -> 1426 bytes ...propagate_scalar__nufit32_no_nubar__f4.pkl | Bin 936 -> 1009 bytes ...propagate_scalar__nufit32_no_nubar__f8.pkl | Bin 1317 -> 1426 bytes ...ropagate_scalar__nufit32_std_decay__f4.pkl | Bin 936 -> 1009 bytes ...ropagate_scalar__nufit32_std_decay__f8.pkl | Bin 1317 -> 1426 bytes ...agate_scalar__nufit32_std_decay_no__f4.pkl | Bin 936 -> 1009 bytes ...agate_scalar__nufit32_std_decay_no__f8.pkl | Bin 1317 -> 1426 bytes ...opagate_scalar__nufit32_std_nsi_no__f4.pkl | Bin 936 -> 1009 bytes ...opagate_scalar__nufit32_std_nsi_no__f8.pkl | Bin 1317 -> 1426 bytes ...opagate_scalar__nufit32_vac_nsi_no__f4.pkl | Bin 936 -> 1009 bytes ...opagate_scalar__nufit32_vac_nsi_no__f8.pkl | Bin 1317 -> 1426 bytes 63 files changed, 2507 insertions(+), 9 deletions(-) create mode 100644 pisa/stages/osc/lri_params.py create mode 100644 pisa/stages/osc/prob3_bk.py create mode 100644 pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py create mode 100644 pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py create mode 100755 pisa/stages/osc/prob3numba/numba_osc_tests_bk.py create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/convert_from_mass_eigenstate_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/convert_from_mass_eigenstate_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_vac_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_H_vac_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_dms_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_dms_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_massbasis_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_massbasis_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/product_hostfunc__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/product_hostfunc__nufit32_lri_std_mat__f8.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f4.pkl create mode 100644 pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f8.pkl diff --git a/pisa/stages/osc/lri_params.py b/pisa/stages/osc/lri_params.py new file mode 100644 index 000000000..52d3a0c67 --- /dev/null +++ b/pisa/stages/osc/lri_params.py @@ -0,0 +1,121 @@ +""" +LRI Params: Charecterize Long Range Interaction mediator + +Developed by following osc_params.py and nsi_params.py +""" + +from __future__ import absolute_import, division + +import numpy as np +from pisa import CTYPE, FTYPE +from pisa.utils.comparisons import ALLCLOSE_KW, isscalar, recursiveEquality +from pisa.utils.log import Levels, set_verbosity, logging + +__all__ = ['LRIParams'] + +class LRIParams(object): + """ + Holds the mediator information of long range interaction:z' + Assumed three anamoly free symmetries, Le_mu, Le_tau, Lmu_tau(by mixing z and z')). + + Attributes + ---------- + potential matrix : Three 2d float array of shape (3,3), one for each symmetry + + Potential matrix holding the potential term of three different symmetris, which is a + function of mediator mass, and the coupling constant of the interaction. + + + """ + + def __init__(self): + + self._v_lri = 0. + self._potential_matrix_emu = np.zeros((3, 3), dtype=FTYPE) + self._potential_matrix_etau = np.zeros((3, 3), dtype=FTYPE) + self._potential_matrix_mutau = np.zeros((3, 3), dtype=FTYPE) + + # --- LRI potential --- + + @property + def v_lri(self): + """Potential term of symmetry e mu""" + return self._v_lri + + @v_lri.setter + def v_lri(self, value): + assert value <1. + self._v_lri = value + + @property + def potential_matrix_emu(self): + """LRI matter interaction potential matrix e mu symmetry""" + + v_matrix = np.zeros((3, 3), dtype=FTYPE) + + v_matrix[0, 0] = self.v_lri + v_matrix[0, 1] = 0. + v_matrix[0, 2] = 0. + v_matrix[1, 0] = 0. + v_matrix[1, 1] = - self.v_lri + v_matrix[1, 2] = 0. + v_matrix[2, 0] = 0. + v_matrix[2, 1] = 0. + v_matrix[2, 2] = 0. + + assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW) + + return v_matrix + + @property + def potential_matrix_etau(self): + """LRI matter interaction potential matrix e tau symmetry""" + + v_matrix = np.zeros((3, 3), dtype=FTYPE) + + v_matrix[0, 0] = self.v_lri + v_matrix[0, 1] = 0. + v_matrix[0, 2] = 0. + v_matrix[1, 0] = 0. + v_matrix[1, 1] = 0. + v_matrix[1, 2] = 0. + v_matrix[2, 0] = 0. + v_matrix[2, 1] = 0. + v_matrix[2, 2] = - self.v_lri + + assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW) + + return v_matrix + + @property + def potential_matrix_mutau(self): + """LRI matter interaction potential matrix mu tau symmetry""" + + v_matrix = np.zeros((3, 3), dtype=FTYPE) + + v_matrix[0, 0] = 0. + v_matrix[0, 1] = 0. + v_matrix[0, 2] = 0. + v_matrix[1, 0] = 0. + v_matrix[1, 1] = self.v_lri + v_matrix[1, 2] = 0. + v_matrix[2, 0] = 0. + v_matrix[2, 1] = 0. + v_matrix[2, 2] = - self.v_lri + + assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW) + + return v_matrix + + +def test_lri_params(): + """ + # TODO: implement me! + """ + pass + +if __name__=='__main__': + from pisa import TARGET + from pisa.utils.log import set_verbosity, logging + set_verbosity(1) + test_lri_params() diff --git a/pisa/stages/osc/prob3.py b/pisa/stages/osc/prob3.py index 391bb4d43..d4a86a9c3 100644 --- a/pisa/stages/osc/prob3.py +++ b/pisa/stages/osc/prob3.py @@ -18,6 +18,7 @@ from pisa.stages.osc.nsi_params import StdNSIParams, VacuumLikeNSIParams from pisa.stages.osc.osc_params import OscParams from pisa.stages.osc.decay_params import DecayParams +from pisa.stages.osc.lri_params import LRIParams from pisa.stages.osc.layers import Layers from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import propagate_array, fill_probs from pisa.utils.numba_tools import WHERE @@ -64,6 +65,8 @@ class prob3(Stage): eps_mutau_phase : quantity (angle) eps_tautau : quantity (dimensionless) decay_alpha3 : quantity (energy^2) + v_lri : quantity (eV) + **kwargs Other kwargs are handled by Stage @@ -77,6 +80,7 @@ def __init__( reparam_mix_matrix=False, neutrino_decay=False, scale_density=False, + lri_type=None, **std_kwargs, ): @@ -155,7 +159,22 @@ def __init__( else: decay_params = () - expected_params = expected_params + nsi_params + decay_params + if lri_type is not None: + choices = ['emu-symmetry', 'etau-symmetry', 'mutau-symmetry'] + lri_type = lri_type.strip().lower() + if not lri_type in choices: + raise ValueError( + 'Chosen LRI symmetry type "%s" not available! Choose one of %s.' + % (lri_type, choices) + ) + self.lri_type = lri_type + + if self.lri_type is None: + lri_params = () + else: + lri_params = ('v_lri',) + + expected_params = expected_params + nsi_params + decay_params + lri_params # init base class super().__init__( @@ -169,6 +188,8 @@ def __init__( self.nsi_params = None self.decay_params = None self.decay_matrix = None + self.lri_params = None + self.lri_pot = None # Note that the interaction potential (Hamiltonian) just scales with the # electron density N_e for propagation through the Earth, # even(to very good approx.) in the presence of generalised interactions @@ -201,6 +222,10 @@ def setup_function(self): if self.neutrino_decay: logging.debug('Working with neutrino decay') self.decay_params = DecayParams() + + if self.lri_type is not None: + logging.debug('Working with LRI') + self.lri_params = LRIParams() # setup the layers #if self.params.earth_model.value is not None: @@ -262,6 +287,7 @@ def calc_probs(self, nubar, e_array, rho_array, len_array, out): self.gen_mat_pot_matrix_complex, self.decay_flag, self.decay_matix, + self.lri_pot, nubar, e_array, rho_array, @@ -337,6 +363,9 @@ def compute_function(self): if self.neutrino_decay: self.decay_params.decay_alpha3 = self.params.decay_alpha3.value.m_as('eV**2') + + if self.lri_type is not None: + self.lri_params.v_lri = self.params.v_lri.value.m_as('eV') # now we can proceed to calculate the generalised matter potential matrix std_mat_pot_matrix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) @@ -361,6 +390,18 @@ def compute_function(self): else : self.decay_matix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) + self.lri_pot = np.zeros((3, 3), dtype=FTYPE) + types_lri = ['emu-symmetry', 'etau-symmetry', 'etau-symmetry'] + if self.lri_type is not None: + if self.lri_type == 'emu-symmetry': + self.lri_pot = self.lri_params.potential_matrix_emu + elif self.lri_type == 'etau-symmetry': + self.lri_pot = self.lri_params.potential_matrix_etau + elif self.lri_type == 'mutau-symmetry': + self.lri_pot = self.lri_params.potential_matrix_mutau + else: + raise Exception("Implemented symmetries are" % types_lri) + for container in self.data: self.calc_probs(container['nubar'], diff --git a/pisa/stages/osc/prob3_bk.py b/pisa/stages/osc/prob3_bk.py new file mode 100644 index 000000000..391bb4d43 --- /dev/null +++ b/pisa/stages/osc/prob3_bk.py @@ -0,0 +1,407 @@ +""" +PISA pi stage for the calculation of earth layers and osc. probabilities + +Maybe it would amke sense to split this up into a separate earth layer stage +and an osc. stage....todo + +""" + +from __future__ import absolute_import, print_function, division + +import numpy as np +from numba import guvectorize + +from pisa import FTYPE, CTYPE, TARGET, ureg +from pisa.core.stage import Stage +from pisa.utils.log import logging +from pisa.utils.profiler import profile +from pisa.stages.osc.nsi_params import StdNSIParams, VacuumLikeNSIParams +from pisa.stages.osc.osc_params import OscParams +from pisa.stages.osc.decay_params import DecayParams +from pisa.stages.osc.layers import Layers +from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import propagate_array, fill_probs +from pisa.utils.numba_tools import WHERE +from pisa.utils.resources import find_resource + + +class prob3(Stage): + """ + Prob3-like oscillation PISA Pi class + + Parameters + ---------- + params + Expected params .. :: + + detector_depth : float + earth_model : PREM file path + prop_height : quantity (dimensionless) + YeI : quantity (dimensionless) + YeO : quantity (dimensionless) + YeM : quantity (dimensionless) + density_scale : quantity (dimensionless) + theta12 : quantity (angle) + theta13 : quantity (angle) + theta23 : quantity (angle) + deltam21 : quantity (mass^2) + deltam31 : quantity (mass^2) + deltacp : quantity (angle) + eps_scale : quantity(dimensionless) + eps_prime : quantity(dimensionless) + phi12 : quantity(angle) + phi13 : quantity(angle) + phi23 : quantity(angle) + alpha1 : quantity(angle) + alpha2 : quantity(angle) + deltansi : quantity(angle) + eps_ee : quantity (dimensionless) + eps_emu_magn : quantity (dimensionless) + eps_emu_phase : quantity (angle) + eps_etau_magn : quantity (dimensionless) + eps_etau_phase : quantity (angle) + eps_mumu : quantity(dimensionless) + eps_mutau_magn : quantity (dimensionless) + eps_mutau_phase : quantity (angle) + eps_tautau : quantity (dimensionless) + decay_alpha3 : quantity (energy^2) + + **kwargs + Other kwargs are handled by Stage + ----- + + """ + + def __init__( + self, + nsi_type=None, + reparam_mix_matrix=False, + neutrino_decay=False, + scale_density=False, + **std_kwargs, + ): + + expected_params = ( + 'detector_depth', + 'earth_model', + 'prop_height', + 'YeI', + 'YeO', + 'YeM', + 'theta12', + 'theta13', + 'theta23', + 'deltam21', + 'deltam31', + 'deltacp' + ) + + # Add (matter) density scale as free parameter? + if scale_density: + expected_params = expected_params + ('density_scale',) + + # Check whether and if so with which NSI parameters we are to work. + if nsi_type is not None: + choices = ['standard', 'vacuum-like'] + nsi_type = nsi_type.strip().lower() + if not nsi_type in choices: + raise ValueError( + 'Chosen NSI type "%s" not available! Choose one of %s.' + % (nsi_type, choices) + ) + self.nsi_type = nsi_type + """Type of NSI to assume.""" + + self.reparam_mix_matrix = reparam_mix_matrix + """Use a PMNS mixing matrix parameterisation that differs from + the standard one by an overall phase matrix + diag(e^(i*delta_CP), 1, 1). This has no impact on + oscillation probabilities in the *absence* of NSI.""" + + self.neutrino_decay = neutrino_decay + + if neutrino_decay: + self.decay_flag = 1 + else : + self.decay_flag = -1 + + """Invoke neutrino decay with neutrino oscillation.""" + + if self.nsi_type is None: + nsi_params = () + elif self.nsi_type == 'vacuum-like': + nsi_params = ('eps_scale', + 'eps_prime', + 'phi12', + 'phi13', + 'phi23', + 'alpha1', + 'alpha2', + 'deltansi' + ) + elif self.nsi_type == 'standard': + nsi_params = ('eps_ee', + 'eps_emu_magn', + 'eps_emu_phase', + 'eps_etau_magn', + 'eps_etau_phase', + 'eps_mumu', + 'eps_mutau_magn', + 'eps_mutau_phase', + 'eps_tautau' + ) + + if self.neutrino_decay : + decay_params = ('decay_alpha3',) + else: + decay_params = () + + expected_params = expected_params + nsi_params + decay_params + + # init base class + super().__init__( + expected_params=expected_params, + **std_kwargs, + ) + + + self.layers = None + self.osc_params = None + self.nsi_params = None + self.decay_params = None + self.decay_matrix = None + # Note that the interaction potential (Hamiltonian) just scales with the + # electron density N_e for propagation through the Earth, + # even(to very good approx.) in the presence of generalised interactions + # (NSI), which is why we can simply treat it as a constant here. + self.gen_mat_pot_matrix_complex = None + """Interaction Hamiltonian without the factor sqrt(2)*G_F*N_e.""" + self.YeI = None + self.YeO = None + self.YeM = None + + def setup_function(self): + + # object for oscillation parameters + self.osc_params = OscParams() + if self.reparam_mix_matrix: + logging.debug( + 'Working with reparameterizated version of mixing matrix.' + ) + else: + logging.debug( + 'Working with standard parameterization of mixing matrix.' + ) + if self.nsi_type == 'vacuum-like': + logging.debug('Working in vacuum-like NSI parameterization.') + self.nsi_params = VacuumLikeNSIParams() + elif self.nsi_type == 'standard': + logging.debug('Working in standard NSI parameterization.') + self.nsi_params = StdNSIParams() + + if self.neutrino_decay: + logging.debug('Working with neutrino decay') + self.decay_params = DecayParams() + + # setup the layers + #if self.params.earth_model.value is not None: + earth_model = find_resource(self.params.earth_model.value) + self.YeI = self.params.YeI.value.m_as('dimensionless') + self.YeO = self.params.YeO.value.m_as('dimensionless') + self.YeM = self.params.YeM.value.m_as('dimensionless') + prop_height = self.params.prop_height.value.m_as('km') + detector_depth = self.params.detector_depth.value.m_as('km') + self.layers = Layers(earth_model, detector_depth, prop_height) + self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) + + # --- calculate the layers --- + if self.is_map: + # speed up calculation by adding links + # as layers don't care about flavour + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc', + 'nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + + for container in self.data: + self.layers.calcLayers(container['true_coszen']) + container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) + container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) + + # don't forget to un-link everything again + self.data.unlink_containers() + + # --- setup empty arrays --- + if self.is_map: + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc']) + self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + for container in self.data: + container['probability'] = np.empty((container.size, 3, 3), dtype=FTYPE) + self.data.unlink_containers() + + # setup more empty arrays + for container in self.data: + container['prob_e'] = np.empty((container.size), dtype=FTYPE) + container['prob_mu'] = np.empty((container.size), dtype=FTYPE) + + def calc_probs(self, nubar, e_array, rho_array, len_array, out): + ''' wrapper to execute osc. calc ''' + if self.reparam_mix_matrix: + mix_matrix = self.osc_params.mix_matrix_reparam_complex + else: + mix_matrix = self.osc_params.mix_matrix_complex + + logging.debug('mat pot:\n%s' + % self.gen_mat_pot_matrix_complex) + logging.debug('decay mat:\n%s' + % self.decay_matix) + + propagate_array(self.osc_params.dm_matrix, # pylint: disable = unexpected-keyword-arg, no-value-for-parameter + mix_matrix, + self.gen_mat_pot_matrix_complex, + self.decay_flag, + self.decay_matix, + nubar, + e_array, + rho_array, + len_array, + out=out + ) + + def compute_function(self): + + if self.is_map: + # speed up calculation by adding links + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc']) + self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + + # this can be done in a more clever way (don't have to recalculate all paths) + YeI = self.params.YeI.value.m_as('dimensionless') + YeO = self.params.YeO.value.m_as('dimensionless') + YeM = self.params.YeM.value.m_as('dimensionless') + if YeI != self.YeI or YeO != self.YeO or YeM != self.YeM: + self.YeI = YeI; self.YeO = YeO; self.YeM = YeM + self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) + for container in self.data: + self.layers.calcLayers(container['true_coszen']) + container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) + container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) + + # Matter density scale is a free parameter? + if 'density_scale' in self.params.names: + density_scale = self.params.density_scale.value.m_as('dimensionless') + else: + density_scale = 1. + + # some safety checks on units + # trying to avoid issue of angles with no dimension being assumed to be radians + # here we enforce the user must speficy a valid angle unit + for angle_param in [self.params.theta12, self.params.theta13, self.params.theta23, self.params.deltacp] : + assert angle_param.value.units != ureg.dimensionless, "Param %s is dimensionless, but should have angle units [rad, degree]" % angle_param.name + + # --- update mixing params --- + self.osc_params.theta12 = self.params.theta12.value.m_as('rad') + self.osc_params.theta13 = self.params.theta13.value.m_as('rad') + self.osc_params.theta23 = self.params.theta23.value.m_as('rad') + self.osc_params.dm21 = self.params.deltam21.value.m_as('eV**2') + self.osc_params.dm31 = self.params.deltam31.value.m_as('eV**2') + self.osc_params.deltacp = self.params.deltacp.value.m_as('rad') + if self.nsi_type == 'vacuum-like': + self.nsi_params.eps_scale = self.params.eps_scale.value.m_as('dimensionless') + self.nsi_params.eps_prime = self.params.eps_prime.value.m_as('dimensionless') + self.nsi_params.phi12 = self.params.phi12.value.m_as('rad') + self.nsi_params.phi13 = self.params.phi13.value.m_as('rad') + self.nsi_params.phi23 = self.params.phi23.value.m_as('rad') + self.nsi_params.alpha1 = self.params.alpha1.value.m_as('rad') + self.nsi_params.alpha2 = self.params.alpha2.value.m_as('rad') + self.nsi_params.deltansi = self.params.deltansi.value.m_as('rad') + elif self.nsi_type == 'standard': + self.nsi_params.eps_ee = self.params.eps_ee.value.m_as('dimensionless') + self.nsi_params.eps_emu = ( + (self.params.eps_emu_magn.value.m_as('dimensionless'), + self.params.eps_emu_phase.value.m_as('rad')) + ) + self.nsi_params.eps_etau = ( + (self.params.eps_etau_magn.value.m_as('dimensionless'), + self.params.eps_etau_phase.value.m_as('rad')) + ) + self.nsi_params.eps_mumu = self.params.eps_mumu.value.m_as('dimensionless') + self.nsi_params.eps_mutau = ( + (self.params.eps_mutau_magn.value.m_as('dimensionless'), + self.params.eps_mutau_phase.value.m_as('rad')) + ) + self.nsi_params.eps_tautau = self.params.eps_tautau.value.m_as('dimensionless') + + if self.neutrino_decay: + self.decay_params.decay_alpha3 = self.params.decay_alpha3.value.m_as('eV**2') + + # now we can proceed to calculate the generalised matter potential matrix + std_mat_pot_matrix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) + std_mat_pot_matrix[0, 0] += 1.0 + + # add effective nsi coupling matrix + if self.nsi_type is not None: + logging.debug('NSI matrix:\n%s' % self.nsi_params.eps_matrix) + self.gen_mat_pot_matrix_complex = ( + std_mat_pot_matrix + self.nsi_params.eps_matrix + ) + logging.debug('Using generalised matter potential:\n%s' + % self.gen_mat_pot_matrix_complex) + else: + self.gen_mat_pot_matrix_complex = std_mat_pot_matrix + logging.debug('Using standard matter potential:\n%s' + % self.gen_mat_pot_matrix_complex) + + if self.neutrino_decay: + self.decay_matix = self.decay_params.decay_matrix + logging.debug('Decay matrix:\n%s' % self.decay_params.decay_matrix) + else : + self.decay_matix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) + + + for container in self.data: + self.calc_probs(container['nubar'], + container['true_energy'], + container['densities']*density_scale, + container['distances'], + out=container['probability'], + ) + container.mark_changed('probability') + + # the following is flavour specific, hence unlink + self.data.unlink_containers() + + for container in self.data: + # initial electrons (0) + fill_probs(container['probability'], + 0, + container['flav'], + out=container['prob_e'], + ) + # initial muons (1) + fill_probs(container['probability'], + 1, + container['flav'], + out=container['prob_mu'], + ) + + container.mark_changed('prob_e') + container.mark_changed('prob_mu') + + + def apply_function(self): + + # maybe speed up like this? + #self.data.representation = self.calc_mode + #for container in self.data: + # container['oscillated_flux'] = (container['nu_flux'][:,0] * container['prob_e']) + (container['nu_flux'][:,1] * container['prob_mu']) + + #self.data.representation = self.apply_mode + + # update the outputted weights + for container in self.data: + container['weights'] *= (container['nu_flux'][:,0] * container['prob_e']) + (container['nu_flux'][:,1] * container['prob_mu']) + diff --git a/pisa/stages/osc/prob3numba/numba_osc_hostfuncs.py b/pisa/stages/osc/prob3numba/numba_osc_hostfuncs.py index 2a988f6bb..652784ca1 100644 --- a/pisa/stages/osc/prob3numba/numba_osc_hostfuncs.py +++ b/pisa/stages/osc/prob3numba/numba_osc_hostfuncs.py @@ -58,29 +58,29 @@ @guvectorize( - [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], - "(a,a), (a,a), (b,c), (), (b,c), (), (), (i), (i) -> (a,a)", + [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {FX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], + "(a,a), (a,a), (b,c), (), (b,c), (b,c), (), (), (i), (i) -> (a,a)", target=TARGET, ) -def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability): +def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability): """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET is "cuda" or "host")""" osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability + dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability ) @njit( - [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], + [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {FX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], parallel=TARGET == "parallel" ) def propagate_scalar( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability + dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability ): """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET is "cuda" or "host")""" osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability + dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability ) @@ -97,6 +97,7 @@ def propagate_scalar( f"{CX}[:,:], " # H_vac f"{IX}, " # neutrino decay flag f"{CX}[:,:], " # H_decay + f"{FX}[:,:], " # lri_pot f"{FX}[:,:], " # dm f"{CX}[:,:], " # transition_matrix ")" @@ -114,6 +115,7 @@ def get_transition_matrix_hostfunc( H_vac, decay_flag, H_decay, + lri_pot, dm, transition_matrix, ): @@ -130,6 +132,7 @@ def get_transition_matrix_hostfunc( H_vac, decay_flag, H_decay, + lri_pot, dm, transition_matrix, ) diff --git a/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py b/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py new file mode 100644 index 000000000..2a988f6bb --- /dev/null +++ b/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py @@ -0,0 +1,218 @@ +# pylint: disable = invalid-name + +""" +Host function wrappers for numba_osc_kernels. +""" + +from __future__ import absolute_import, print_function, division + +__all__ = ["FX", "CX", "IX", "propagate_array", "fill_probs"] + +import numpy as np +from numba import guvectorize, njit + +from pisa import FTYPE, ITYPE, TARGET +from pisa.stages.osc.prob3numba.numba_osc_kernels import ( + # osc_probs_vacuum_kernel, + osc_probs_layers_kernel, + get_transition_matrix, + get_transition_matrix_massbasis, + get_H_vac, + get_H_decay, + get_H_mat, + get_dms, + get_dms_numerical, + get_product, + convert_from_mass_eigenstate, +) + + +assert FTYPE in [np.float32, np.float64], str(FTYPE) + +FX = "f4" if FTYPE == np.float32 else "f8" +"""Float string code to use, understood by both Numba and Numpy""" + +CX = "c8" if FTYPE == np.float32 else "c16" +"""Complex string code to use, understood by both Numba and Numpy""" + +IX = "i4" if ITYPE == np.int32 else "i8" +"""Signed integer string code to use, understood by both Numba and Numpy""" + + +# @guvectorize( +# [f"({FX}[:,:], {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:,:])"], +# "(a,a), (a,a), (), (), (i) -> (a,a)", +# target=TARGET, +# ) +# def propagate_array_vacuum(dm, mix, nubar, energy, distances, probability): +# """wrapper to run `osc_probs_vacuum_kernel` from host (whether TARGET is +# "cuda" or "host")""" +# osc_probs_vacuum_kernel(dm, mix, nubar, energy, distances, probability) +# +# +# @njit([f"({FX}[:,:], {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:,:])"], target=TARGET) +# def propagate_scalar_vacuum(dm, mix, nubar, energy, distances, probability): +# """wrapper to run `osc_probs_vacuum_kernel` from host (whether TARGET is +# "cuda" or "host")""" +# osc_probs_vacuum_kernel(dm, mix, nubar, energy, distances, probability) + + +@guvectorize( + [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], + "(a,a), (a,a), (b,c), (), (b,c), (), (), (i), (i) -> (a,a)", + target=TARGET, +) +def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability): + """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET + is "cuda" or "host")""" + osc_probs_layers_kernel( + dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability + ) + + +@njit( + [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], + parallel=TARGET == "parallel" +) +def propagate_scalar( + dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability +): + """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET + is "cuda" or "host")""" + osc_probs_layers_kernel( + dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability + ) + + +@njit( + [ + "(" + f"{IX}, " # nubar + f"{FX}, " # energy + f"{FX}, " # rho + f"{FX}, " # baseline + f"{CX}[:,:], " # mix_nubar + f"{CX}[:,:], " # mix_nubar_conj_transp + f"{CX}[:,:], " # mat_pot + f"{CX}[:,:], " # H_vac + f"{IX}, " # neutrino decay flag + f"{CX}[:,:], " # H_decay + f"{FX}[:,:], " # dm + f"{CX}[:,:], " # transition_matrix + ")" + ], + parallel=TARGET == "parallel" +) +def get_transition_matrix_hostfunc( + nubar, + energy, + rho, + baseline, + mix_nubar, + mix_nubar_conj_transp, + mat_pot, + H_vac, + decay_flag, + H_decay, + dm, + transition_matrix, +): + """wrapper to run `get_transition_matrix` from host (whether TARGET is + "cuda" or "host")""" + get_transition_matrix( + nubar, + energy, + rho, + baseline, + mix_nubar, + mix_nubar_conj_transp, + mat_pot, + H_vac, + decay_flag, + H_decay, + dm, + transition_matrix, + ) + + +@njit([f"({FX}, {FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") +def get_transition_matrix_massbasis_hostfunc( + baseline, + energy, + dm_mat, + dm_mat_mat, + H_full_mass_eigenstate_basis, + transition_matrix, +): + """wrapper to run `get_transition_matrix_massbasis` from host (whether + TARGET is "cuda" or "host")""" + get_transition_matrix_massbasis( + baseline, + energy, + dm_mat, + dm_mat_mat, + H_full_mass_eigenstate_basis, + transition_matrix, + ) + + +@njit([f"({CX}[:,:], {CX}[:,:], {FX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") +def get_H_vac_hostfunc(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac): + """wrapper to run `get_H_vac` from host (whether TARGET is "cuda" or "host")""" + get_H_vac(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac) + +@njit([f"({CX}[:,:], {CX}[:,:], {FX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") +def get_H_decay_hostfunc(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay): + """wrapper to run `get_H_decay` from host (whether TARGET is "cuda" or "host")""" + get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay) + +# @guvectorize( +# [f"({FX}, {CX}[:,:], {IX}, {CX}[:,:])"], "(), (m, m), () -> (m, m)" +# ) +@njit([f"({FX}, {CX}[:,:], {IX}, {CX}[:,:])"], parallel=TARGET == "parallel") +def get_H_mat_hostfunc(rho, mat_pot, nubar, H_mat): + """wrapper to run `get_H_mat` from host (whether TARGET is "cuda" or "host")""" + get_H_mat(rho, mat_pot, nubar, H_mat) + + +@njit([f"({FX}, {CX}[:,:], {FX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") +def get_dms_hostfunc(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat): + """wrapper to run `get_dms` from host (whether TARGET is "cuda" or "host")""" + get_dms(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat) + +@njit([f"({FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") +def get_dms_numerical_hostfunc(energy, H_full, dm_mat_mat, dm_mat): + """wrapper to run `get_dms` from host (whether TARGET is "cuda" or "host")""" + get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat) + +@njit([f"({FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:], {CX}[:,:,:])"], parallel=TARGET == "parallel") +def get_product_hostfunc( + energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product +): + """wrapper to run `get_product` from host (whether TARGET is "cuda" or "host")""" + get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product) + + +@njit([f"({IX}, {CX}[:,:], {CX}[:])"], parallel=TARGET == "parallel") +def convert_from_mass_eigenstate_hostfunc(state, mix_nubar, psi): + """wrapper to run `convert_from_mass_eigenstate` from host (whether TARGET + is "cuda" or "host")""" + convert_from_mass_eigenstate(state, mix_nubar, psi) + + +@guvectorize( + [f"({FX}[:,:], {IX}, {IX}, {FX}[:])"], "(a,b), (), () -> ()" +) +def fill_probs(probability, initial_flav, flav, out): + """Fill `out` with transition probabilities to go from `initial_flav` to + `flav`, from values in `probaility` matrix. + + Parameters + ---------- + probability : real 2d array + initial_flav : signed(?) int (int4 or int8) + flav : signed(?) int (int4 or int8) + out : real 1-d array + + """ + out[0] = probability[initial_flav, flav] diff --git a/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py b/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py new file mode 100644 index 000000000..76d1f9885 --- /dev/null +++ b/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py @@ -0,0 +1,879 @@ +# pylint: disable = not-callable, invalid-name, too-many-nested-blocks + + +""" +Neutrino flavour oscillation in matter calculation +Based on the original prob3++ implementation of Roger Wendell +http://www.phy.duke.edu/~raw22/public/Prob3++/ (2012) + +See `numba_osc_tests.py` for unit tests of functions in this module. +""" + + +from __future__ import absolute_import, print_function, division + +__all__ = [ + # "osc_probs_vacuum_kernel", + "osc_probs_layers_kernel", + "get_transition_matrix", +] + +__version__ = "0.2" + +import cmath +import math +import numpy as np + +from pisa.utils.numba_tools import ( + myjit, + conjugate_transpose, + conjugate, + matrix_dot_matrix, + matrix_dot_vector, + clear_matrix, + copy_matrix, + cuda, + ctype, + ftype, +) + + +# TODO/FIXME: osc_probs_vacuum_kernel produces non-unitary results. No one +# should use this until the issue is resolved. + +# @myjit +# def osc_probs_vacuum_kernel(dm, mix, nubar, energy, distance_in_layer, osc_probs): +# """ Calculate vacumm mixing probabilities +# +# Parameters +# ---------- +# dm : real 2d array +# Mass splitting matrix, eV^2 +# +# mix : complex 2d array +# PMNS mixing matrix +# +# nubar : int +# +1 for neutrinos, -1 for antineutrinos +# +# energy : float +# Neutrino energy, GeV +# +# distance_in_layer : real 1d-array +# Baselines (will be summed up), km +# +# osc_probs : real 2d array (empty) +# Returned oscillation probabilities in the form: +# osc_prob[i,j] = probability of flavor i to oscillate into flavor j +# with 0 = electron, 1 = muon, 3 = tau +# +# +# Notes +# ----- +# This is largely unvalidated so far +# +# """ +# +# # no need to conjugate mix matrix, as we anyway only need real part +# # can this be right? +# +# clear_matrix(osc_probs) +# osc_probs_local = cuda.local.array(shape=(3, 3), dtype=ftype) +# +# # sum up length from all layers +# baseline = 0.0 +# for i in range(distance_in_layer.shape[0]): +# baseline += distance_in_layer[i] +# +# # make more precise 20081003 rvw +# l_over_e = 1.26693281 * baseline / energy +# s21 = math.sin(dm[1, 0] * l_over_e) +# s32 = math.sin(dm[2, 0] * l_over_e) +# s31 = math.sin((dm[2, 1] + dm[3, 2]) * l_over_e) +# +# # does anybody understand this loop? +# # ista = abs(*nutype) - 1 +# for ista in range(3): +# for iend in range(2): +# osc_probs_local[ista, iend] = ( +# (mix[ista, 0].real * mix[ista, 1].real * s21) ** 2 +# + (mix[ista, 1].real * mix[ista, 2].real * s32) ** 2 +# + (mix[ista, 2].real * mix[ista, 0].real * s31) ** 2 +# ) +# if iend == ista: +# osc_probs_local[ista, iend] = 1.0 - 4.0 * osc_probs_local[ista, iend] +# else: +# osc_probs_local[ista, iend] = -4.0 * osc_probs_local[ista, iend] +# +# osc_probs_local[ista, 2] = ( +# 1.0 - osc_probs_local[ista, 0] - osc_probs_local[ista, 1] +# ) +# +# # is this necessary? +# if nubar > 0: +# copy_matrix(osc_probs_local, osc_probs) +# else: +# for i in range(3): +# for j in range(3): +# osc_probs[i, j] = osc_probs_local[j, i] + + +@myjit +def osc_probs_layers_kernel( + dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, density_in_layer, distance_in_layer, osc_probs +): + """ Calculate oscillation probabilities + + given layers of length and density + + Parameters + ---------- + dm : real 2d array + Mass splitting matrix, eV^2 + + mix : complex 2d array + PMNS mixing matrix + + mat_pot : complex 2d array + Generalised matter potential matrix without "a" factor (will be + multiplied with "a" factor); set to diag([1, 0, 0]) for only standard + oscillations + + decay_flag: int + +1 forstandard oscillations + decay, -1 for standard oscillations + + mat_decay : complex 2d array + decay matrix with -j*alpha3 = [2,2] element + + nubar : real int, scalar or Nd array (broadcast dim) + 1 for neutrinos, -1 for antineutrinos + + energy : real float, scalar or Nd array (broadcast dim) + Neutrino energy, GeV + + density_in_layer : real 1d array + Density of each layer, moles of electrons / cm^2 + + distance_in_layer : real 1d array + Distance of each layer traversed, km + + osc_probs : real (N+2)-d array (empty) + Returned oscillation probabilities in the form: + osc_prob[i,j] = probability of flavor i to oscillate into flavor j + with 0 = electron, 1 = muon, 3 = tau + + + Notes + ----- + !!! Right now, because of CUDA, the maximum number of layers + is hard coded and set to 120 (59Layer PREM + Atmosphere). + This is used for cached layer computation, where earth layer, which + are typically traversed twice (it's symmetric) are not recalculated + but rather cached.. + + """ + + # 3x3 complex + H_vac = cuda.local.array(shape=(3, 3), dtype=ctype) + H_decay = cuda.local.array(shape=(3, 3), dtype=ctype) + mix_nubar = cuda.local.array(shape=(3, 3), dtype=ctype) + mix_nubar_conj_transp = cuda.local.array(shape=(3, 3), dtype=ctype) + transition_product = cuda.local.array(shape=(3, 3), dtype=ctype) + transition_matrix = cuda.local.array(shape=(3, 3), dtype=ctype) + tmp = cuda.local.array(shape=(3, 3), dtype=ctype) + + clear_matrix(H_vac) + clear_matrix(H_decay) + clear_matrix(osc_probs) + + # 3-vector complex + raw_input_psi = cuda.local.array(shape=(3), dtype=ctype) + output_psi = cuda.local.array(shape=(3), dtype=ctype) + + use_mass_eigenstates = False + + cache = True + # cache = False + + # TODO: + # * ensure convention below is respected in MC reweighting + # (nubar > 0 for nu, < 0 for anti-nu) + # * nubar is passed in, so could already pass in the correct form + # of mixing matrix, i.e., possibly conjugated + if nubar > 0: + # in this case the mixing matrix is left untouched + copy_matrix(mix, mix_nubar) + + else: + # here we need to complex conjugate all entries + # (note that this only changes calculations with non-clear_matrix deltacp) + conjugate(mix, mix_nubar) + + conjugate_transpose(mix_nubar, mix_nubar_conj_transp) + + get_H_vac(mix_nubar, mix_nubar_conj_transp, dm, H_vac) + #Compute the decay matrix in the flavor basis + get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay) + + + if cache: + # allocate array to store all the transition matrices + # doesn't work in cuda...needs fixed shape + transition_matrices = cuda.local.array(shape=(120, 3, 3), dtype=ctype) + + # loop over layers + for i in range(distance_in_layer.shape[0]): + density = density_in_layer[i] + distance = distance_in_layer[i] + if distance > 0.0: + layer_matrix_index = -1 + # chaeck if exists + for j in range(i): + # if density_in_layer[j] == density and distance_in_layer[j] == distance: + if (abs(density_in_layer[j] - density) < 1e-5) and ( + abs(distance_in_layer[j] - distance) < 1e-5 + ): + layer_matrix_index = j + + # use from cached + if layer_matrix_index >= 0: + for j in range(3): + for k in range(3): + transition_matrices[i, j, k] = transition_matrices[ + layer_matrix_index, j, k + ] + + # only calculate if necessary + else: + get_transition_matrix( + nubar, + energy, + density, + distance, + mix_nubar, + mix_nubar_conj_transp, + mat_pot, + H_vac, + decay_flag, + H_decay, + dm, + transition_matrix, + ) + # copy + for j in range(3): + for k in range(3): + transition_matrices[i, j, k] = transition_matrix[j, k] + else: + # identity matrix + for j in range(3): + for k in range(3): + if j == k: + transition_matrix[j, k] = 0.0 + else: + transition_matrix[j, k] = 1.0 + + # now multiply them all + first_layer = True + for i in range(distance_in_layer.shape[0]): + distance = distance_in_layer[i] + if distance > 0.0: + for j in range(3): + for k in range(3): + transition_matrix[j, k] = transition_matrices[i, j, k] + if first_layer: + copy_matrix(transition_matrix, transition_product) + first_layer = False + else: + matrix_dot_matrix(transition_matrix, transition_product, tmp) + copy_matrix(tmp, transition_product) + + else: + # non-cache loop + first_layer = True + for i in range(distance_in_layer.shape[0]): + density = density_in_layer[i] + distance = distance_in_layer[i] + # only do something if distance > 0. + if distance > 0.0: + get_transition_matrix( + nubar, + energy, + density, + distance, + mix_nubar, + mix_nubar_conj_transp, + mat_pot, + H_vac, + decay_flag, + H_decay, + dm, + transition_matrix, + ) + if first_layer: + copy_matrix(transition_matrix, transition_product) + first_layer = False + else: + matrix_dot_matrix(transition_matrix, transition_product, tmp) + copy_matrix(tmp, transition_product) + + # convrt to flavour eigenstate basis + matrix_dot_matrix(transition_product, mix_nubar_conj_transp, tmp) + matrix_dot_matrix(mix_nubar, tmp, transition_product) + + # loop on neutrino types, and compute probability for neutrino i: + for i in range(3): + for j in range(3): + raw_input_psi[j] = 0.0 + + if use_mass_eigenstates: + convert_from_mass_eigenstate(i + 1, mix_nubar, raw_input_psi) + else: + raw_input_psi[i] = 1.0 + + matrix_dot_vector(transition_product, raw_input_psi, output_psi) + osc_probs[i][0] += output_psi[0].real ** 2 + output_psi[0].imag ** 2 + osc_probs[i][1] += output_psi[1].real ** 2 + output_psi[1].imag ** 2 + osc_probs[i][2] += output_psi[2].real ** 2 + output_psi[2].imag ** 2 + + +@myjit +def get_transition_matrix( + nubar, + energy, + rho, + baseline, + mix_nubar, + mix_nubar_conj_transp, + mat_pot, + H_vac, + decay_flag, + H_decay, + dm, + transition_matrix, +): + """ Calculate neutrino flavour transition amplitude matrix + + Parameters + ---------- + nubar : int + +1 for neutrinos, -1 for antineutrinos + + energy : real float + Neutrino energy, GeV + + rho : real float + Electron number density (in moles/cm^3) (numerically, this is just the + product of electron fraction and mass density in g/cm^3, since the + number of grams per cm^3 corresponds to the number of moles of nucleons + per cm^3) + + baseline : real float + Baseline, km + + mix_nubar : complex 2d array + Mixing matrix, already conjugated if antineutrino + + mix_nubar_conj_transp : complex conjugate 2d array + Conjugate transpose of mix_nubar + + mat_pot : complex 2d array + Generalised matter potential matrix without "a" factor (will be + multiplied with "a" factor); set to diag([1, 0, 0]) for only standard + oscillations + + H_vac : complex 2d array + Hamiltonian in vacuum, without the 1/2E term + + H_decay : complex 2d array + Decay matrix, without the 1/2E term + + decay_flag: int + +1 forstandard oscillations + decay, -1 for standard oscillations + + dm : real 2d array + Mass splitting matrix, eV^2 + + transition_matrix : complex 2d array (empty) + Transition matrix in mass eigenstate basis + + Notes + ----- + For neutrino (nubar > 0) or antineutrino (nubar < 0) with energy energy + traversing layer of matter of uniform density rho with thickness baseline + + """ + + H_mat = cuda.local.array(shape=(3, 3), dtype=ctype) + dm_mat = cuda.local.array(shape=(3, 3), dtype=ctype) + dm_mat_mat = cuda.local.array(shape=(3, 3), dtype=ctype) + #H_decay = cuda.local.array(shape=(3, 3), dtype=ctype) + H_full = cuda.local.array(shape=(3, 3), dtype=ctype) + tmp = cuda.local.array(shape=(3, 3), dtype=ctype) + H_full_mass_eigenstate_basis = cuda.local.array(shape=(3, 3), dtype=ctype) + + # Compute the matter potential including possible generalized interactions + # in the flavor basis + + get_H_mat(rho, mat_pot, nubar, H_mat) + + # Get the full Hamiltonian by adding together matter and vacuum parts + one_over_two_e = 0.5 / energy + + if (decay_flag == 1): + for i in range(3): + for j in range(3): + H_full[i, j] = (H_vac[i, j] + H_decay[i,j]) * one_over_two_e + H_mat[i, j] + + # Calculate modified mass eigenvalues in matter from the full Hamiltonian + get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat) + + else: + for i in range(3): + for j in range(3): + H_full[i, j] = H_vac[i, j] * one_over_two_e + H_mat[i, j] + + + # Calculate modified mass eigenvalues in matter from the full Hamiltonian and + # the vacuum mass splittings + get_dms(energy, H_full, dm, dm_mat_mat, dm_mat) + + # Now we transform the full Hamiltonian back into the + # mass eigenstate basis so we don't need to compute products of the effective + # mixing matrix elements explicitly + matrix_dot_matrix(H_full, mix_nubar, tmp) + matrix_dot_matrix(mix_nubar_conj_transp, tmp, H_full_mass_eigenstate_basis) + + # We can now proceed to calculating the transition amplitude from the Hamiltonian + # in the mass basis and the effective mass splittings + get_transition_matrix_massbasis( + baseline, + energy, + dm_mat, + dm_mat_mat, + H_full_mass_eigenstate_basis, + transition_matrix, + ) + + +@myjit +def get_transition_matrix_massbasis( + baseline, + energy, + dm_mat, + dm_mat_mat, + H_full_mass_eigenstate_basis, + transition_matrix, +): + """ + Calculate the transition amplitude matrix + + Parameters + ---------- + baseline : float + Baseline traversed, km + + energy : float + Neutrino energy, GeV + + dm_mat : complex 2d array + + dm_mat_mat : complex 2d array + + H_full_mass_eigenstate_basis : complex 2d array + + transition_matrix : complex 2d array (empty) + Transition matrix in mass eigenstate basis + + Notes + ----- + - corrsponds to matrix A (equation 10) in original Barger paper + - take into account generic potential matrix (=Hamiltonian) + + """ + + product = cuda.local.array(shape=(3, 3, 3), dtype=ctype) + + clear_matrix(transition_matrix) + + get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product) + + # (1/2)*(1/(h_bar*c)) in units of GeV/(eV^2 km) + hbar_c_factor = 2.534 + + for k in range(3): + arg = -dm_mat[k, 0] * (baseline / energy) * hbar_c_factor + c = cmath.exp(arg * 1.0j) + for i in range(3): + for j in range(3): + transition_matrix[i, j] += c * product[i, j, k] + + +@myjit +def get_H_vac(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac): + """ Calculate vacuum Hamiltonian in flavor basis for neutrino or antineutrino + + Parameters: + ----------- + mix_nubar : complex 2d array + Mixing matrix, already conjugated if antineutrino + + mix_nubar_conj_transp : conjugate 2d array + Conjugate transpose of mix_nubar + + dm_vac_vac : 2d array + Matrix of mass splittings + + H_vac : complex 2d array + Hamiltonian in vacuum, without the 1/2E term + + + Notes + ------ + The Hailtonian does not contain the energy dependent factor of + 1/(2 * E), as it will be added later + + """ + + dm_vac_diag = cuda.local.array(shape=(3, 3), dtype=ctype) + tmp = cuda.local.array(shape=(3, 3), dtype=ctype) + + clear_matrix(dm_vac_diag) + + dm_vac_diag[1, 1] = dm_vac_vac[1, 0] + 0j + dm_vac_diag[2, 2] = dm_vac_vac[2, 0] + 0j + + matrix_dot_matrix(dm_vac_diag, mix_nubar_conj_transp, tmp) + matrix_dot_matrix(mix_nubar, tmp, H_vac) + +@myjit +def get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay): + """ Calculate decay Hamiltonian in flavor basis for neutrino or antineutrino + + Parameters: + ----------- + mix_nubar : complex 2d array + Mixing matrix, already conjugated if antineutrino + + mix_nubar_conj_transp : conjugate 2d array + Conjugate transpose of mix_nubar + + H_decay : complex 2d array + Hamiltonian for decay, without the 1/2E term + + + Notes + ------ + The Hailtonian does not contain the energy dependent factor of + 1/(2 * E), as it will be added later + + """ + + #dm_decay_diag = cuda.local.array(shape=(3, 3), dtype=ctype) + tmp = cuda.local.array(shape=(3, 3), dtype=ctype) + + clear_matrix(H_decay) + + #dm_decay_diag[2, 2] = 0. - 1.0j*alpha3 + + #matrix_dot_matrix(dm_decay_diag, mix_nubar_conj_transp, tmp) + matrix_dot_matrix(mat_decay, mix_nubar_conj_transp, tmp) + matrix_dot_matrix(mix_nubar, tmp, H_decay) + +@myjit +def get_H_mat(rho, mat_pot, nubar, H_mat): + """ Calculate matter Hamiltonian in flavor basis + + Parameters: + ----------- + rho : real float + Electron number density (in moles/cm^3) (numerically, this is just the + product of electron fraction and mass density in g/cm^3, since the + number of grams per cm^3 corresponds to the number of moles of nucleons + per cm^3) + + mat_pot : complex 2d array + Generalised matter potential matrix without "a" factor (will be + multiplied with "a" factor); set to diag([1, 0, 0]) for only standard + oscillations + + nubar : int + +1 for neutrinos, -1 for antineutrinos + + H_mat : complex 2d array (empty) + matter hamiltonian + + Notes + ----- + In the following, `a` is just the standard effective matter potential + induced by charged-current weak interactions with electrons + + """ + + # 2*sqrt(2)*Gfermi in (eV^2 cm^3)/(mole GeV) + tworttwoGf = 1.52588e-4 + a = 0.5 * rho * tworttwoGf + + # standard matter interaction Hamiltonian + clear_matrix(H_mat) + + # formalism of Hamiltonian: not 1+epsilon_ee^f in [0,0] element but just epsilon... + # changed when fitting in Thomas' NSI branch -EL + + # Obtain effective non-standard matter interaction Hamiltonian + # changed when fitting in Thomas' NSI branch -EL + for i in range(3): + for j in range(3): + # matter potential V -> -V* for anti-neutrinos + if nubar == -1: + H_mat[i, j] = -a * mat_pot[i, j].conjugate() + elif nubar == 1: + H_mat[i, j] = a * mat_pot[i, j] + +@myjit +def get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat): + """Compute the matter-mass vector M, dM = M_i-M_j and dMimj + + Parameters + ---------- + energy : float + Neutrino energy, GeV + + H_full : complex 2d array + full hamiltonian + + dm_mat_mat : complex 2d array (empty) + + dm_mat : complex 2d array (empty) + + + Notes + ----- + Calculate mass eigenstates in matter + neutrino or anti-neutrino (type already taken into account in Hamiltonian) + of energy energy. + """ + + + m_mat = 2.0 * energy * np.linalg.eigvals(H_full) + + for i in range(3): + for j in range(3): + dm_mat_mat[i, j] = m_mat[i] - m_mat[j] + dm_mat[i, j] = m_mat[i] + +@myjit +def get_dms(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat): + """Compute the matter-mass vector M, dM = M_i-M_j and dMimj + + Parameters + ---------- + energy : float + Neutrino energy, GeV + + H_full : complex 2d array + full hamiltonian + + dm_vac_vac : 2d array + + dm_mat_mat : complex 2d array (empty) + + dm_mat : complex 2d array (empty) + + + Notes + ----- + Calculate mass eigenstates in matter + neutrino or anti-neutrino (type already taken into account in Hamiltonian) + of energy energy. + + - only god knows what happens in this function, somehow it seems to work + + """ + + # the following is for solving the characteristic polynomial of H_full: + # P(x) = x**3 + c2*x**2 + c1*x + c0 + real_product_a = (H_full[0, 1] * H_full[1, 2] * H_full[2, 0]).real + real_product_b = (H_full[0, 0] * H_full[1, 1] * H_full[2, 2]).real + + norm_H_e_mu_sq = H_full[0, 1].real ** 2 + H_full[0, 1].imag ** 2 + norm_H_e_tau_sq = H_full[0, 2].real ** 2 + H_full[0, 2].imag ** 2 + norm_H_mu_tau_sq = H_full[1, 2].real ** 2 + H_full[1, 2].imag ** 2 + + # c1 = H_{11} * H_{22} + H_{11} * H_{33} + H_{22} * H_{33} + # - |H_{12}|**2 - |H_{13}|**2 - |H_{23}|**2 + # given Hermiticity of Hamiltonian (real diagonal elements), + # this coefficient must be real + c1 = ( + (H_full[0, 0].real * (H_full[1, 1] + H_full[2, 2])).real + - (H_full[0, 0].imag * (H_full[1, 1] + H_full[2, 2])).imag + + (H_full[1, 1].real * H_full[2, 2]).real + - (H_full[1, 1].imag * H_full[2, 2]).imag + - norm_H_e_mu_sq + - norm_H_mu_tau_sq + - norm_H_e_tau_sq + ) + + # c0 = H_{11} * |H_{23}|**2 + H_{22} * |H_{13}|**2 + H_{33} * |H_{12}|**2 + # - H_{11} * H_{22} * H_{33} - 2*Re(H*_{13} * H_{12} * H_{23}) + # hence, this coefficient is also real + c0 = ( + H_full[0, 0].real * norm_H_mu_tau_sq + + H_full[1, 1].real * norm_H_e_tau_sq + + H_full[2, 2].real * norm_H_e_mu_sq + - 2.0 * real_product_a + - real_product_b + ) + + # c2 = -H_{11} - H_{22} - H_{33} + # hence, this coefficient is also real + c2 = -H_full[0, 0].real - H_full[1, 1].real - H_full[2, 2].real + + one_over_two_e = 0.5 / energy + one_third = 1.0 / 3.0 + two_third = 2.0 / 3.0 + + # we also have to perform the corresponding algebra + # for the vacuum case, where the relevant elements of the + # hamiltonian are mass differences + x = dm_vac_vac[1, 0] + y = dm_vac_vac[2, 0] + + c2_v = -one_over_two_e * (x + y) + + # p is real due to reality of c1 and c2 + p = c2 ** 2 - 3.0 * c1 + p_v = one_over_two_e ** 2 * (x ** 2 + y ** 2 - x * y) + p = max(0.0, p) + + # q is real + q = -13.5 * c0 - c2 ** 3 + 4.5 * c1 * c2 + q_v = one_over_two_e ** 3 * (x + y) * ((x + y) ** 2 - 4.5 * x * y) + + # we need the quantity p**3 - q**2 to obtain the eigenvalues, + # but let's prevent inaccuracies and instead write + tmp = 27 * (0.25 * c1 ** 2 * (p - c1) + c0 * (q + 6.75 * c0)) + # changed from p**3 - q**2 when fitting in Thomas' NSI branch -EL + # TODO: can we simplify this quantity to reduce numerical inaccuracies? + tmp_v = p_v ** 3 - q_v ** 2 + + tmp = max(0.0, tmp) + + theta = cuda.local.array(shape=(3), dtype=ftype) + theta_v = cuda.local.array(shape=(3), dtype=ftype) + m_mat = cuda.local.array(shape=(3), dtype=ftype) + m_mat_u = cuda.local.array(shape=(3), dtype=ftype) + m_mat_v = cuda.local.array(shape=(3), dtype=ftype) + + a = two_third * math.pi + # intermediate result, needed to calculate the three + # mass eigenvalues (theta0, theta1, theta2 are the three + # corresponding arguments of the cosine, see m_mat_u) + res = math.atan2(math.sqrt(tmp), q) * one_third + theta[0] = res + a + theta[1] = res - a + theta[2] = res + res_v = math.atan2(math.sqrt(tmp_v), q_v) * one_third + theta_v[0] = res_v + a + theta_v[1] = res_v - a + theta_v[2] = res_v + + b = two_third * math.sqrt(p) + b_v = two_third * math.sqrt(p_v) + + for i in range(3): + m_mat_u[i] = ( + 2.0 * energy * (b * math.cos(theta[i]) - c2 * one_third + dm_vac_vac[0, 0]) + ) + m_mat_v[i] = ( + 2.0 + * energy + * (b_v * math.cos(theta_v[i]) - c2_v * one_third + dm_vac_vac[0, 0]) + ) + + # Sort according to which reproduce the vaccum eigenstates + for i in range(3): + tmp_v = abs(dm_vac_vac[i, 0] - m_mat_v[0]) + k = 0 + for j in range(3): + tmp = abs(dm_vac_vac[i, 0] - m_mat_v[j]) + if tmp < tmp_v: + k = j + tmp_v = tmp + m_mat[i] = m_mat_u[k] + + for i in range(3): + for j in range(3): + dm_mat_mat[i, j] = m_mat[i] - m_mat[j] + dm_mat[i, j] = m_mat[i] + #dm_mat_vac[i, j] = m_mat[i] - dm_vac_vac[j, 0] + + +@myjit +def get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product): + """ + Parameters + ---------- + energy : float + Neutrino energy, GeV + + dm_mat : complex 2d array + + dm_mat_mat : complex 2d array + + H_full_mass_eigenstate_basis : complex 2d array + + product : complex 3d-array (empty) + + """ + + H_minus_M = cuda.local.array(shape=(3, 3, 3), dtype=ctype) + + for i in range(3): + for j in range(3): + for k in range(3): + H_minus_M[i, j, k] = 2.0 * energy * H_full_mass_eigenstate_basis[i, j] + if i == j: + H_minus_M[i, j, k] -= dm_mat[k, j] + # also, cler product + product[i, j, k] = 0.0 + + # Calculate the product in eq.(10) of H_minus_M for j!=k + for i in range(3): + for j in range(3): + for k in range(3): + product[i, j, 0] += H_minus_M[i, k, 1] * H_minus_M[k, j, 2] + product[i, j, 1] += H_minus_M[i, k, 2] * H_minus_M[k, j, 0] + product[i, j, 2] += H_minus_M[i, k, 0] * H_minus_M[k, j, 1] + product[i, j, 0] /= dm_mat_mat[0, 1] * dm_mat_mat[0, 2] + product[i, j, 1] /= dm_mat_mat[1, 2] * dm_mat_mat[1, 0] + product[i, j, 2] /= dm_mat_mat[2, 0] * dm_mat_mat[2, 1] + + +@myjit +def convert_from_mass_eigenstate(state, mix_nubar, psi): + """ + Parameters + ---------- + state : (un?)signed int + + mix_nubar : complex 2d array + Mixing matrix, already conjugated if antineutrino + + psi : complex 1d-array (empty) + + + Notes + ----- + this is untested! + + """ + + mass = cuda.local.array(shape=(3), dtype=ctype) + + lstate = state - 1 + for i in range(3): + mass[i] = 1.0 if lstate == i else 0.0 + + # note: mix_nubar is already taking into account whether we're considering + # nu or anti-nu + matrix_dot_vector(mix_nubar, mass, psi) diff --git a/pisa/stages/osc/prob3numba/numba_osc_kernels.py b/pisa/stages/osc/prob3numba/numba_osc_kernels.py index 76d1f9885..9dee2c540 100644 --- a/pisa/stages/osc/prob3numba/numba_osc_kernels.py +++ b/pisa/stages/osc/prob3numba/numba_osc_kernels.py @@ -120,7 +120,7 @@ @myjit def osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, density_in_layer, distance_in_layer, osc_probs + dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, density_in_layer, distance_in_layer, osc_probs ): """ Calculate oscillation probabilities @@ -144,6 +144,11 @@ def osc_probs_layers_kernel( mat_decay : complex 2d array decay matrix with -j*alpha3 = [2,2] element + + lri_pot : real 2d array + Potential contribution due to matter with the consideration of + Long Range Interaction. this potential not a generalised one, so + it passed as a separate matrix. nubar : real int, scalar or Nd array (broadcast dim) 1 for neutrinos, -1 for antineutrinos @@ -256,6 +261,7 @@ def osc_probs_layers_kernel( H_vac, decay_flag, H_decay, + lri_pot, dm, transition_matrix, ) @@ -306,6 +312,7 @@ def osc_probs_layers_kernel( H_vac, decay_flag, H_decay, + lri_pot, dm, transition_matrix, ) @@ -330,7 +337,9 @@ def osc_probs_layers_kernel( else: raw_input_psi[i] = 1.0 + matrix_dot_vector(transition_product, raw_input_psi, output_psi) + osc_probs[i][0] += output_psi[0].real ** 2 + output_psi[0].imag ** 2 osc_probs[i][1] += output_psi[1].real ** 2 + output_psi[1].imag ** 2 osc_probs[i][2] += output_psi[2].real ** 2 + output_psi[2].imag ** 2 @@ -348,6 +357,7 @@ def get_transition_matrix( H_vac, decay_flag, H_decay, + lri_pot, dm, transition_matrix, ): @@ -389,6 +399,11 @@ def get_transition_matrix( decay_flag: int +1 forstandard oscillations + decay, -1 for standard oscillations + + lri_pot : real 2d array + Potential contribution due to matter with the consideration of + Long Range Interaction. this potential not a generalised one, so + it passed as a separate matrix. dm : real 2d array Mass splitting matrix, eV^2 @@ -416,6 +431,14 @@ def get_transition_matrix( get_H_mat(rho, mat_pot, nubar, H_mat) + # Adding the LRI potential to H_mat(lri pot is 2d array with full of zeros in the absence of lri) + for i in range(3): + for j in range(3): + if nubar > 0: + H_mat[i, j] = H_mat[i, j] + lri_pot[i, j]*1e9 #eV/GeV + elif nubar < 0: + H_mat[i, j] = H_mat[i, j] - lri_pot[i, j]*1e9 #ev/Gev + # Get the full Hamiltonian by adding together matter and vacuum parts one_over_two_e = 0.5 / energy diff --git a/pisa/stages/osc/prob3numba/numba_osc_tests.py b/pisa/stages/osc/prob3numba/numba_osc_tests.py index 994d34947..6dfdeaee5 100755 --- a/pisa/stages/osc/prob3numba/numba_osc_tests.py +++ b/pisa/stages/osc/prob3numba/numba_osc_tests.py @@ -69,6 +69,8 @@ VacuumLikeNSIParams, ) +from pisa.stages.osc.lri_params import LRIParams + TEST_DATA_DIR = find_resource("osc/numba_osc_tests_data") FINFO_FTYPE = np.finfo(FTYPE) @@ -91,6 +93,12 @@ # NOTE: !!DO NOT CHANGE!! (unless a function is incorrect) tests rely on these # ---------------------------------------------------------------------------- # +#lri param +lri_params = LRIParams() +lri_params.v_lri = 1e-14 +lri_std_mat_pot = lri_params.potential_matrix_emu +# print(lri_std_mat_pot) + DEFAULTS = dict( energy=1, # GeV state=1, @@ -109,6 +117,7 @@ dm31=2.494e-3, decay_flag=-1, mat_decay=np.diag([0, 0, -1.0e-4j ]).astype(np.complex128), + lri_pot=lri_std_mat_pot, ) # define non-0 NSI parameters for non-vacuum NSI @@ -162,6 +171,9 @@ decay_flag=1, mat_decay=DEFAULTS["mat_decay"], ), + nufit32_lri_std_mat=dict( #nufit 3.2 lri potential with std matter potential + lri_pot = lri_std_mat_pot + ), ) @@ -249,6 +261,7 @@ def test_prob3numba(ignore_fails=False, define_as_ref=False): tc_["mat_pot"].astype(CX), tc_["decay_flag"], tc_["mat_decay"].astype(CX), + tc_["lri_pot"].astype(FX), nubars, energies, tc_["layer_densities"].astype(FX), @@ -396,6 +409,7 @@ def run_test_case(tc_name, tc, ignore_fails=False, define_as_ref=False): mat_pot=tc_["mat_pot"], decay_flag=tc_["decay_flag"], mat_decay=tc_["mat_decay"], + lri_pot=tc_["lri_pot"], nubar=tc_["nubar"], energy=tc_["energy"], densities=tc_["layer_densities"], @@ -439,6 +453,7 @@ def run_test_case(tc_name, tc, ignore_fails=False, define_as_ref=False): H_vac=H_vac_ref, decay_flag=tc_["decay_flag"], H_decay=H_decay_ref, + lri_pot=tc_["lri_pot"], dm=tc_["dm"], # output: transition_matrix=np.ones(shape=(3, 3), dtype=CX), diff --git a/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py b/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py new file mode 100755 index 000000000..994d34947 --- /dev/null +++ b/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py @@ -0,0 +1,791 @@ +#!/usr/bin/env python +# pylint: disable = invalid-name + + +""" +Tests for prob3numba code +""" + + +from __future__ import absolute_import, print_function, division + + +__all__ = [ + "TEST_DATA_DIR", + "FINFO_FTYPE", + "AC_KW", + "PRINTOPTS", + "A2S_KW", + "MAT_DOT_MAT_SUBSCR", + "DEFAULTS", + "TEST_CASES", + "auto_populate_test_case", + "test_prob3numba", + "run_test_case", + "stability_test", + "execute_func", + "compare_numeric", + "check", + "ary2str", + "main", +] + + +from argparse import ArgumentParser +from collections import OrderedDict +from collections.abc import Mapping +from copy import deepcopy +from inspect import getmodule, signature +from os.path import join + +import numpy as np +# import os +# os.environ['PISA_FTYPE'] = 'single' # for checking unit test on single precision +from pisa import FTYPE +from pisa.utils.comparisons import ALLCLOSE_KW +from pisa.utils.fileio import expand, from_file, to_file +from pisa.utils.log import Levels, logging, set_verbosity +from pisa.utils.numba_tools import WHERE +from pisa.utils.resources import find_resource +from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import ( + CX, + FX, + IX, + # propagate_scalar_vacuum, + propagate_scalar, + propagate_array, + get_transition_matrix_hostfunc, + get_transition_matrix_massbasis_hostfunc, + get_H_vac_hostfunc, + get_H_decay_hostfunc, + get_H_mat_hostfunc, + get_dms_hostfunc, + get_dms_numerical_hostfunc, + get_product_hostfunc, + convert_from_mass_eigenstate_hostfunc, +) +from pisa.stages.osc.nsi_params import ( + StdNSIParams, + VacuumLikeNSIParams, +) + +TEST_DATA_DIR = find_resource("osc/numba_osc_tests_data") + +FINFO_FTYPE = np.finfo(FTYPE) + +AC_KW = dict(atol=FINFO_FTYPE.resolution * 10, rtol=ALLCLOSE_KW["rtol"] * 100) + +PRINTOPTS = dict( + precision=FINFO_FTYPE.precision + 2, floatmode="fixed", sign=" ", linewidth=200 +) + +A2S_KW = dict(precision=PRINTOPTS["precision"], separator=", ") + +MAT_DOT_MAT_SUBSCR = "in,nj->ij" +"""matrix dot matrix subscripts for use by `numpy.einsum`""" + +# ---------------------------------------------------------------------------- # +# Define relevant values for testing purposes (from nufit3.2, from intermediate +# calculations performed here, or arbitary values). +# +# NOTE: !!DO NOT CHANGE!! (unless a function is incorrect) tests rely on these +# ---------------------------------------------------------------------------- # + +DEFAULTS = dict( + energy=1, # GeV + state=1, + nubar=1, + rho=1, # moles of electrons / cm^3 + baseline=1, # km + mat_pot=np.diag([1, 0, 0]).astype(np.complex128), + layer_distances=np.logspace(0, 2, 10), # km + layer_densities=np.linspace(0.5, 3, 10), # g/cm^3 + # osc params: defaults are nufit 3.2 normal ordering values + t12=np.deg2rad(33.62), + t23=np.deg2rad(47.2), + t13=np.deg2rad(8.54), + dcp=np.deg2rad(234), + dm21=7.40e-5, + dm31=2.494e-3, + decay_flag=-1, + mat_decay=np.diag([0, 0, -1.0e-4j ]).astype(np.complex128), +) + +# define non-0 NSI parameters for non-vacuum NSI +# roughly based on best fit params from Thomas Ehrhardts 3y DRAGON analysis +nsi_params = StdNSIParams() +nsi_params.eps_emu_magn = 0.07 +nsi_params.eps_emu_phase = np.deg2rad(340) +nsi_params.eps_etau_magn = 0.06 +nsi_params.eps_etau_phase = np.deg2rad(35) +nsi_params.eps_mutau_magn = 0.003 +nsi_params.eps_mutau_phase = np.deg2rad(175) +mat_pot_std_nsi_no = np.diag([1, 0, 0]).astype(np.complex128) + nsi_params.eps_matrix + +# Vacuum-like NSI parameters +nsi_params = VacuumLikeNSIParams() +nsi_params.eps_prime = 0.1 +mat_pot_vac_nsi_no = np.diag([1, 0, 0]).astype(np.complex128) + nsi_params.eps_matrix + +TEST_CASES = dict( + nufit32_no=dict(), # nufit 3.2 normal ordering (also overall) best-fit + nufit32_no_nubar=dict(nubar=-1), # NO but anti-neutrinos + nufit32_no_E1TeV=dict(energy=1e3), # NO but e=1 TeV + nufit32_no_blearth=dict( + baseline=6371e3 * 2, + layer_distances=( + 6371e3 + * 2 + * DEFAULTS["layer_distances"] + / np.sum(DEFAULTS["layer_distances"]) + ), + ), + nufit32_io=dict( # nufit 3.2 best-fit params for inverted ordering + t12=np.deg2rad(33.62), + t23=np.deg2rad(48.1), + t13=np.deg2rad(8.58), + dcp=np.deg2rad(278), + dm21=7.40e-5, + dm31=-2.465e-3, + ), + nufit32_std_nsi_no=dict( # nufit 3.2 normal ordering with non-0 standard NSI parameters + mat_pot=mat_pot_std_nsi_no, + ), + nufit32_vac_nsi_no=dict( # nufit 3.2 normal ordering with non-0 vacuum NSI parameters + mat_pot=mat_pot_vac_nsi_no, + ), + nufit32_std_decay_no=dict( # nufit 3.2 normal ordering with no neutrino decay + decay_flag=-1, + mat_decay=DEFAULTS["mat_decay"], + ), + nufit32_std_decay=dict( # nufit 3.2 normal ordering with neutrino decay + decay_flag=1, + mat_decay=DEFAULTS["mat_decay"], + ), +) + + +def auto_populate_test_case(tc, defaults): + """Populate defaults and construct dm / PMNS matrices if they aren't + present in a test case. + + Parameters + ---------- + test_case : mutable mapping + + defaults : mapping + + """ + for key, val in defaults.items(): + if key not in tc: + tc[key] = val + + # Construct dm and PMNS matrices derived from test case values, if + # these were not already specified + + if "dm" not in tc: # construct Delta m^2 matrix if not present + if "dm32" not in tc: # NO case; Delta m^2_32/eV^2 + tc["dm32"] = tc["dm31"] - tc["dm21"] + + if "dm31" not in tc: # IO case; Delta m^2_32/eV^2 + tc["dm31"] = tc["dm32"] + tc["dm21"] + + tc["dm"] = np.array( + [ + [0, -tc["dm21"], -tc["dm31"]], + [tc["dm21"], 0, -tc["dm32"]], + [tc["dm31"], tc["dm32"], 0], + ] + ) + + if "pmns" not in tc: # construct PMNS matrix if not present + c12, s12 = np.cos(tc["t12"]), np.sin(tc["t12"]) + c23, s23 = np.cos(tc["t23"]), np.sin(tc["t23"]) + c13, s13 = np.cos(tc["t13"]), np.sin(tc["t13"]) + + tc["pmns"] = ( + np.array([[1, 0, 0], [0, c23, s23], [0, -s23, c23]]) + @ np.array( + [ + [c13, 0, s13 * np.exp(-1j * tc["dcp"])], + [0, 1, 0], + [-s13 * np.exp(1j * tc["dcp"]), 0, c13], + ] + ) + @ np.array([[c12, s12, 0], [-s12, c12, 0], [0, 0, 1]]) + ) + + +def test_prob3numba(ignore_fails=False, define_as_ref=False): + """Run all unit test cases for prob3numba code""" + + # Pull first test case to test calling `propagate_array` + tc_name, tc = next(iter(TEST_CASES.items())) + tc_ = deepcopy(tc) + logging.info( + "Testing call and return shape of `propagate_array` with test case '%s'", + tc_name, + ) + + # Test simple broadcasting over `nubars` and `energies` where both have + # same shape, as this is the "typical" use case + input_shape = (4, 5) + + # Without broadcasting, a single probability matrix is 3x3 + prob_array_shape = (3, 3) + + # Broadcasted shape + out_shape = input_shape + prob_array_shape + + nubars = np.full(shape=input_shape, fill_value=tc_["nubar"], dtype=IX) + energies = np.full(shape=input_shape, fill_value=tc_["energy"], dtype=FX) + + # Fill with NaN to ensure all elements are assinged a value + probabilities = np.full(shape=out_shape, fill_value=np.nan, dtype=FX) + + propagate_array( + tc_["dm"].astype(FX), + tc_["pmns"].astype(CX), + tc_["mat_pot"].astype(CX), + tc_["decay_flag"], + tc_["mat_decay"].astype(CX), + nubars, + energies, + tc_["layer_densities"].astype(FX), + tc_["layer_distances"].astype(FX), + # output: + probabilities, + ) + + # Check that all probability matrices have no NaNs and are equal to one + # another + ref_probs = probabilities[0, 0] + for i in range(input_shape[0]): + for j in range(input_shape[1]): + probs = probabilities[i, j] + assert np.all(np.isfinite(probs)) + assert np.all(probs == ref_probs) + + # Run all test cases + for tc_name, tc in TEST_CASES.items(): + run_test_case( + tc_name, tc, ignore_fails=ignore_fails, define_as_ref=define_as_ref + ) + + +def run_test_case(tc_name, tc, ignore_fails=False, define_as_ref=False): + """Run one test case""" + logging.info("== TEST CASE : %s ==", tc_name) + + st_test_kw = dict(ignore_fails=ignore_fails, define_as_ref=define_as_ref) + tf_sfx = f"__{tc_name}__{FX}.pkl" + + # Copy contents of test case, so if a function modifies these + # (accidentally), it doesn't affect the outcome of further tests + tc_ = deepcopy(tc) + test, ref = stability_test( + func=convert_from_mass_eigenstate_hostfunc, + func_kw=dict( + state=tc_["state"], + mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, + # output: + psi=np.ones(shape=3, dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"convert_from_mass_eigenstate_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\npsi = %s", ary2str(test["psi"])) + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=get_H_vac_hostfunc, + func_kw=dict( + mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, + mix_nubar_conj_transp=( + tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] + ), + dm_vac_vac=tc_["dm"], + # output: + H_vac=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_H_vac_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\nH_vac = %s", ary2str(test["H_vac"])) + # keep for use by `get_transition_matrix_hostfunc` + H_vac_ref = ref["H_vac"] + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=get_H_decay_hostfunc, + func_kw=dict( + mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, + mix_nubar_conj_transp=( + tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] + ), + mat_decay=tc_["mat_decay"], + #output; + H_decay=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_H_decay_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\nH_decay = %s", ary2str(test["H_decay"])) + # keep for use by `get_transition_matrix_hostfunc` + H_decay_ref = ref["H_decay"] + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=get_H_mat_hostfunc, + func_kw=dict( + rho=tc_["rho"], + mat_pot=tc_["mat_pot"], + nubar=tc_["nubar"], + # output: + H_mat=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_H_mat_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\nH_mat = %s", ary2str(test["H_mat"])) + # keep for use by `get_dms_hostfunc`, + # `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc`` + H_mat_ref = ref["H_mat"] + + # tc_ = deepcopy(tc) + # test, ref = stability_test( + # func=propagate_scalar_vacuum, + # func_kw=dict( + # dm=tc_["dm"], + # mix=tc_["pmns"], + # nubar=tc_["nubar"], + # energy=tc_["energy"], + # distances=tc_["layer_distances"], + # # output: + # probability=np.ones(shape=(3, 3), dtype=FX), + # ), + # ref_path=join(TEST_DATA_DIR, f"propagate_scalar_vacuum{tf_sfx}"), + # **st_test_kw, + # ) + # logging.debug("\nvac_prob = %s", ary2str(test["probability"])) + # # check unitarity + # # TODO: << BUG? >> these fail even in double precision! + # check( + # test=np.sum(test["probability"], axis=0), + # ref=np.ones(3), + # label=( + # f"{tc_name} :: propagate_scalar_vacuum :: sum(vacuum probability, axis=0)" + # ), + # ignore_fails=True, + # ) + # check( + # test=np.sum(test["probability"], axis=1), + # ref=np.ones(3), + # label=( + # f"{tc_name} :: propagate_scalar_vacuum :: sum(vacuum probability, axis=1)" + # ), + # ignore_fails=True, + # ) + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=propagate_scalar, + func_kw=dict( + dm=tc_["dm"], + mix=tc_["pmns"], + mat_pot=tc_["mat_pot"], + decay_flag=tc_["decay_flag"], + mat_decay=tc_["mat_decay"], + nubar=tc_["nubar"], + energy=tc_["energy"], + densities=tc_["layer_densities"], + distances=tc_["layer_distances"], + # output: + probability=np.ones(shape=(3, 3), dtype=FX), + ), + ref_path=join(TEST_DATA_DIR, f"propagate_scalar{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\nmat_prob = %s", ary2str(test["probability"])) + + if (tc_["decay_flag"] != 1): + # check unitarity + check( + test=np.sum(test["probability"], axis=0), + ref=np.ones(3), + label=f"{tc_name} :: propagate_scalar:: sum(matter probability, axis=0)", + ignore_fails=ignore_fails, + ) + check( + test=np.sum(test["probability"], axis=1), + ref=np.ones(3), + label=f"{tc_name} :: propagate_scalar :: sum(matter probability, axis=1)", + ignore_fails=ignore_fails, + ) + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=get_transition_matrix_hostfunc, + func_kw=dict( + nubar=tc_["nubar"], + energy=tc_["energy"], + rho=tc_["rho"], + baseline=tc_["baseline"], + mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, + mix_nubar_conj_transp=( + tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] + ), + mat_pot=tc_["mat_pot"], + H_vac=H_vac_ref, + decay_flag=tc_["decay_flag"], + H_decay=H_decay_ref, + dm=tc_["dm"], + # output: + transition_matrix=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_transition_matrix_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\ntransition_matrix = %s", ary2str(test["transition_matrix"])) + # check unitarity + check( + test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=0), + ref=np.ones(3), + label=( + f"{tc_name}" + " :: get_transition_matrix_hostfunc" + ":: sum(|transition_matrix|^2, axis=0)" + ), + ignore_fails=ignore_fails, + ) + check( + test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=1), + ref=np.ones(3), + label=( + f"{tc_name}" + " :: get_transition_matrix_hostfunc" + " :: sum(|transition_matrix|^2, axis=1)" + ), + ignore_fails=ignore_fails, + ) + + tc_ = deepcopy(tc) + + # Compute H_full as used in `numba_osc_kernels` to call `get_dms` for SI case + # and det_dms_numnerical for decay case from `get_transition_matrix` + + if (tc_["decay_flag"] == 1): + + H_full_ref = (H_vac_ref + H_decay_ref)/ (2 * tc_["energy"]) + H_mat_ref + + test, ref = stability_test( + func=get_dms_numerical_hostfunc, + func_kw=dict( + energy=tc_["energy"], + H_full=H_full_ref, + # outputs: + dm_mat_mat=np.ones(shape=(3, 3), dtype=CX), + dm_mat=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_dms_numerical_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\ndm_mat_mat = %s", ary2str(test["dm_mat_mat"])) + logging.debug("\ndm_mat = %s", ary2str(test["dm_mat"])) + # keep for use by `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc` + dm_mat_mat_ref = ref["dm_mat_mat"] + dm_mat_ref = ref["dm_mat"] + + else: + H_full_ref = H_vac_ref / (2 * tc_["energy"]) + H_mat_ref + + test, ref = stability_test( + func=get_dms_hostfunc, + func_kw=dict( + energy=tc_["energy"], + H_full=H_full_ref, + dm_vac_vac=tc_["dm"], + # outputs: + dm_mat_mat=np.ones(shape=(3, 3), dtype=CX), + dm_mat=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"get_dms_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\ndm_mat_mat = %s", ary2str(test["dm_mat_mat"])) + logging.debug("\ndm_mat = %s", ary2str(test["dm_mat"])) + # keep for use by `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc` + dm_mat_mat_ref = ref["dm_mat_mat"] + dm_mat_ref = ref["dm_mat"] + + + tc_ = deepcopy(tc) + + # Compute same intermediate result `H_full_mass_eigenstate_basis` as in + # `numba_osc_kernels.get_transition_matrix` which calls + # `get_transition_matrix_massbasis` + mix_nubar = tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T + mix_nubar_conj_transp = tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] + tmp = np.einsum(MAT_DOT_MAT_SUBSCR, H_full_ref, mix_nubar) + H_full_mass_eigenstate_basis = np.einsum( + MAT_DOT_MAT_SUBSCR, mix_nubar_conj_transp, tmp + ) + + test, ref = stability_test( + func=get_transition_matrix_massbasis_hostfunc, + func_kw=dict( + baseline=tc_["baseline"], + energy=tc_["energy"], + dm_mat=dm_mat_ref, + dm_mat_mat=dm_mat_mat_ref, + H_full_mass_eigenstate_basis=H_full_mass_eigenstate_basis, + # output: + transition_matrix=np.ones(shape=(3, 3), dtype=CX), + ), + ref_path=join( + TEST_DATA_DIR, f"get_transition_matrix_massbasis_hostfunc{tf_sfx}" + ), + **st_test_kw, + ) + logging.debug("\ntransition_matrix_mb = %s", ary2str(test["transition_matrix"])) + check( + test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=0), + ref=np.ones(3), + label=( + f"{tc_name}" + " :: get_transition_matrix_massbasis_hostfunc" + " :: sum(|transition_matrix (mass basis)|^2), axis=0)" + ), + ignore_fails=ignore_fails, + ) + check( + test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=1), + ref=np.ones(3), + label=( + f"{tc_name}" + " :: get_transition_matrix_massbasis_hostfunc" + " :: sum(|transition_matrix (mass basis)|^2), axis=1)" + ), + ignore_fails=ignore_fails, + ) + + tc_ = deepcopy(tc) + test, ref = stability_test( + func=get_product_hostfunc, + func_kw=dict( + energy=tc_["energy"], + dm_mat=dm_mat_ref, + dm_mat_mat=dm_mat_mat_ref, + H_full_mass_eigenstate_basis=H_full_ref, + # output: + product=np.ones(shape=(3, 3, 3), dtype=CX), + ), + ref_path=join(TEST_DATA_DIR, f"product_hostfunc{tf_sfx}"), + **st_test_kw, + ) + logging.debug("\nproduct = %s", ary2str(test["product"])) + + +def stability_test(func, func_kw, ref_path, ignore_fails=False, define_as_ref=False): + """basic stability test of a Numba CPUDispatcher function (i.e., function + compiled via @jit / @njit)""" + func_name = func.py_func.__name__ + logging.info("stability testing `%s`", func_name) + ref_path = expand(ref_path) + + test = execute_func(func=func, func_kw=func_kw) + + if define_as_ref: + to_file(test, ref_path) + + # Even when we define the test case as ref, round-trip to/from file to + # ensure that doesn't corrupt the values + ref = from_file(ref_path) + + check(test=test, ref=ref, label=func_name, ignore_fails=ignore_fails) + + return test, ref + + +def execute_func(func, func_kw): + """Run `func` with *func_kw.values() where `outputs` specify names in + `func_kw` taken to be outputs of the function; for these, mark changed. + Retrieve both input and output values as Numpy arrays on the host and + aggregate together in a single dict before returning. + + Parameters + ---------- + func : numba CPUDispatcher or CUDADispatcher + func_kw : OrderedDict + + Returns + ------- + ret_dict : OrderedDict + Keys are arg names and vals are type-"correct" values; all arrays are + converted to host Numpy arrays + + """ + py_func = func.py_func + func_name = ".".join([getmodule(py_func).__name__, py_func.__name__]) + arg_names = list(signature(py_func).parameters.keys()) + if hasattr(func, "signatures"): + arg_types = func.signatures[0] + else: + arg_types = func.compiled.argument_types + + missing = set(arg_names).difference(func_kw.keys()) + excess = set(func_kw.keys()).difference(arg_names) + if missing or excess: + msgs = [] + if missing: + msgs.append(f"missing kwargs {missing}") + if excess: + msgs.append(f"excess kwargs {excess}") + raise KeyError(f"{func_name}:" + ", ".join(msgs)) + + typed_args = OrderedDict() + for arg_name, arg_type in zip(arg_names, arg_types): + val = func_kw[arg_name] + if arg_type.name.startswith("array"): + arg_val = val.astype(arg_type.dtype.key) + else: + arg_val = arg_type(val) + typed_args[arg_name] = arg_val + + # Call the host function with typed args + + try: + func(*list(typed_args.values())) + except Exception: + logging.error("Failed running `%s` with args %s", func_name, str(typed_args)) + raise + + # All arrays converted to Numpy host arrays + + ret_dict = OrderedDict() + for key, val in typed_args.items(): + ret_dict[key] = val + + return ret_dict + + +def compare_numeric(test, ref, label=None, ac_kw=deepcopy(AC_KW), ignore_fails=False): + """Compare scalars or numpy ndarrays. + + Parameters + ---------- + test : scalar or numpy.ndarray + ref : scalar or numpy.ndarray + label : str or None, optional + ac_kw : mapping, optional + Keyword args to pass via **ac_kw to `numpy.isclose` / `numpy.allclose` + ignore_fails : bool, optional + + Returns + ------- + rslt : bool + + """ + pfx = f"{label} :: " if label else "" + with np.printoptions(**PRINTOPTS): + if np.isscalar(test): + if np.isclose(test, ref, **ac_kw): + return True + + msg = f"{pfx}test: {test} != ref: {ref}" + if ignore_fails: + logging.warning(msg) + else: + logging.error(msg) + return False + + # Arrays + if np.allclose(test, ref, **ac_kw): + return True + + diff = test - ref + msg = f"{pfx}test:" f"\n{(test)}\n!= ref:\n{(ref)}" f"\ndiff:\n{(diff)}" + + if not np.all(ref == 1): + nzmask = ref != 0 + zmask = ref == 0 + fdiff = np.empty_like(ref) + fdiff[nzmask] = diff[nzmask] / ref[nzmask] + fdiff[zmask] = np.nan + msg += f"\nfractdiff:\n{(fdiff)}" + + if ignore_fails: + logging.warning(msg) + else: + logging.error(msg) + + return False + + +def check(test, ref, label=None, ac_kw=deepcopy(AC_KW), ignore_fails=False): + """Check that `test` matches `ref` (closely enough). + + Parameters + ---------- + test + ref + ac_kw : mapping, optional + Kwargs to `np.allclose`, as used by + `pisa.utils.comparisons.recursiveEquality` + ignore_fails : bool, optional + If True and comparison fails, do not raise AssertionError + + Raises + ------ + AssertionError + If `test` is not close enough to `ref` and ignore_fails is False + + """ + same = True + with np.printoptions(**PRINTOPTS): + if isinstance(test, Mapping): + if not label: + label = "" + else: + label = label + ": " + + for key, val in test.items(): + same &= compare_numeric( + test=val, + ref=ref[key], + label=label + f"key: '{key}'", + ac_kw=ac_kw, + ignore_fails=ignore_fails, + ) + else: + same &= compare_numeric( + test=test, ref=ref, label=label, ac_kw=ac_kw, ignore_fails=ignore_fails + ) + if not ignore_fails and not same: + assert False + return same + + +def ary2str(array): + """Convert a numpy ndarray to string easy to copy back into code""" + return "np.array(" + np.array2string(array, **A2S_KW) + ")" + + +# Augment test cases with defaults / contruct arrays where necessary +for TC in TEST_CASES.values(): + auto_populate_test_case(tc=TC, defaults=DEFAULTS) + + +def main(description=__doc__): + """Script interface for `test_prob3numba` function""" + parser = ArgumentParser(description=description) + parser.add_argument("--ignore-fails", action="store_true") + parser.add_argument("--define-as-ref", action="store_true") + parser.add_argument("-v", action="count", default=Levels.WARN) + kwargs = vars(parser.parse_args()) + set_verbosity(kwargs.pop("v")) + test_prob3numba(**kwargs) + + +if __name__ == "__main__": + main() diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/convert_from_mass_eigenstate_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/convert_from_mass_eigenstate_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..c252477253d7650779f77040b27dd19ce72f204a GIT binary patch literal 392 zcmZo*om$Jt00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqldM)B(WrQ zN{?t>X>LKKUUGg>s$OnsPDy5BQBh*$lpeO?uQUm0T#pcveL!Yv zQD!pGf$?cY`MF7@X=$k-8#ShRK>Wcl`9hFAhjVDTv#Ts=(576?aY@L@XHo7tPW sZAuTL^OO>htC2w3X@s04G(N-9!8rfetv#l|A7EZcr%ntNpj|Zx`km%8&F*jXKrRid|qi1(8wMkBs+o3 z)S}E}pd;hcit=-lO4HI(LAGj4n*i|(Ls;*c1fi%`_E0Xvi`LU+m8}oq{2Ime!3h!D z_RqBD(8xKs&;H)qiMLL!yt4n?^7Y(HPHnPZQ~YaNf$d{^(al>fnEYK~zu-!il90u7 zxc(Nu>)HC6ckS!4bdPRM+h~80+snnf;nDshHXcVFC@;77{^BI<+VB*veo6*c4|B4i z86pHS_%QwN&FszGHl>Hrc}fY$`^*K!nNu>vG^QEg@D;kN(OoqqL&}@EbxMXT$hcBH E0PGIJ>i_@% literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..109badb615665223f8b7094ebdd8ab2f748e1859 GIT binary patch literal 590 zcmZo*o$AcQ00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYs$vm!pP zG%2xYN{>)pX>LKKUUGg>sveM;T9la#)D)jql%JbanwFMY1k|K4%>!aK!{iG=_8{gN zZ5}Yqe0bB|MDB04{(ejMDTGe5ZF?43$j^ss{M1=1-cl?AClOL~}+Ev9sKv`+!L(3`o3(PoODpP$!%AOI8I3?);NyqUe3 z+otp|I!`G9D(ex2`y)O%KQAl3q$n}3xL`^KJBEWBpbl<;I(RSC!F!<&-V1eb1JuEL zp$^^)b#TL!3_h6UZV<~ka}!JAQ&N)?E5R13KtqTK3TPO24^x^6!nYaxxcr(S1okPr V2a;8=$RNThqLje>2C=ME4*(W0tm^;( literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_H_decay_hostfunc__nufit32_lri_std_mat__f8.pkl new file mode 100644 index 0000000000000000000000000000000000000000..18260525e0283d28a622e3c7e01be81c5e603bc0 GIT binary patch literal 843 zcmZo*o$A2M00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYs$vm!pP zG%2xYN{>)pX>LKKUUGg>sveM;T9la#)D)jql%JbanwFMY1k|K4Z34t>hOpi>2|`h? z?4ewS7pOGU4|Me1 literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..091f672f8beff97e28b2179deb713e0d536c8954 GIT binary patch literal 497 zcmbu5v1-FG5J2rXZb}0k3hhwno-Cd^WGG~ChfsqH{Q)D(GBVhaQRKm!r?iFJdhOh= z>u)qC8Y G*aSN}ZH){7 literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_H_mat_hostfunc__nufit32_lri_std_mat__f8.pkl new file mode 100644 index 0000000000000000000000000000000000000000..404e1fe014c1ce80d9e5ec19a640822ab5900123 GIT binary patch literal 650 zcmZo*om$Vt00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqldXDBY#Sd zXkKYWI@;opElYn8F!4A~Sl!+8}*z`JcFhczY2{~4ecwpE~$q+*X%BE#49~c-!22r3$ NL42Gc3-WKN9st0Mkd*)c literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_H_vac_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_H_vac_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..923f2dac72e7fccacb01094fce6a98eeda6b3556 GIT binary patch literal 589 zcmZo*o$AEI00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYs$vm!pP zG%2xYN{>)pX>LKKUUGg>sveM;T9la#)D)jql%JbanwFMY1k|K4%>!aK!{iG=_8{gN zZ5}Yqe0bB|MDB04{(ejMDTGe5ZF?43$j^ss{M1=1-cl?AClOL~}+Ev9sKv`+!L(3`o3(PoODpP$!%AOI8I3?);NyqUe3 z+otp|I!`G9D(ex2`y)O%KQAl3q$n}3xL`^KJBEWBpbl<;I(RSC!F!<&-V1eb1JuEL zp$^^)b#TL!3_h6UZV=14QgY+V5|cp~Y@-S^h(J`a!R#FgMasK@v;~NNcZS?pbFA#TI=4aPcS8AQ+1IB5)pX>LKKUUGg>sveM;T9la#)D)jql%JbanwFMY1k|K4Z34t>hOpi>2|`h? z?4ewS7pbj=B`cmnOWQ zv*#wWdjNzgS|j+k6`x`{M^R9U1M^_Ad-y#C=><%l;iyUKAp~zkp9A zIN(btT>r5%$MV-S1lylZTeHF4@Zvt0{sj>I`_bg>mUdr=dV4Dft{)zZ5GR-F0RZ7w BK3xC+ literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_dms_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_dms_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..0c1a250f1090897aed5047b39a86e39e4bd400d3 GIT binary patch literal 631 zcmZo*otneM00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYauFSRJW za!QYAUTJPYrCxG=QL0{UX--LIVo^~dP@JteIWZ@(2&le?6|4kEr<7C{q=H14(oCjw zcC=3cn&i#g!)P#PGgz}#1{;A-+izE8nQY`v}ec9Cl4*SK5X1!p?A>&%HPqn zNZ$&mR4bWP#>%dXa~o8C2b6D?eO(MB&oCu}tA{Ds0^#BeK1}y|GkY_)P3d8Do>BsG zGFM7&d|6^L2v5ln)0n0LbrXmxHkiF5p-6c*kOsPsf#L29x!pj%HITLj@uy_)LQKn$ zg`1U|SP~D!V6$NU1E~c8NeRK7AiCy-!fq(v5=6tm?`z^|i8!{n!ANTJxr j26Q{RU2kg*!EQyTr zC^H%8#Q3zL{M@9{w6s)^l^W9~Ky))aIjq=uTUO2<%4K*Tn-jIBNpQbdVzS$O8yWlK z2Q(cS?Zoyk3}3{3Tvf{c9aLTvBEP?YPbE0uiz{6Ju`|c=*EBfWpH5q|!QJrSKA8Rm z5dHhnQW$In(%(kp0B(}E0 z)elz>m51|T`e5{w3@H?MvjM{i-OUe7ZKHSSv%$?E#DB8={r*qu)Zyw0@jq86G@UC= NfvYFPM{{$j9suwW5C8xG literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_io__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_io__f4.pkl index edf55208cf2627b05a6420d7f91343d490b10f04..641defa15a2d66aeea6fb4f4fe737fe209478b03 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i|NpPIKNK;2r?XUrS=hwLoi`8WT0E58y8X6r&9-hJ{~u6({?BAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr__7;h~EAG|Gj;(rT+!-*GBsfEPuUq{ibtvjJv1KF*VZO?{Iw&^ZV_* z_6+9&`)@4#wl86h(J!}s=l2Elud*=&%3rDqxDC?(zo6XL_`dP}1N)=}M7J^9Gt8Z0 z)*A4^&SA>F4P483_cIvTvQPQ-ZJ$Gxp0UVb=KTzPng_DPKPcISHSM^04QRmHe~V>Z ZR#e(g$;UJ^BCty}kA{ehDi*!~G0A|AiOa&DwXXe)|1m zH`Mkoc=PP{rqvww3=IpIS=4UtV_>NV%io8}3%cFC2hyMQWnpCQBjf!O4jkFwoX%j+ zAmjSYQsB})2D3`;hGiW483K%$zWUzTXHfK7;B_*?euf8g`!=7uq_nRkcb?I8paJs# cz1UXuRN7C;(DG*XW^S93p<6N~XiBLb08lDbHvj+t diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..77e85f73436c00292c9d38e5642c1c9443ddf084 GIT binary patch literal 1180 zcmZo*ojQ?)0StP$lk;6DVnf>e+QQ>Mw3 z&W`peK!d!Qdl+q|`1$#H{RaXt;muGoCCQnE5$KdDZ9p|WY^iyvMd_7OGT1Y?dYICX z^k(p3)5*ZlU=PvBT$GUy(kA8%^chT$BPp>sH77GK6+?zIH?txh>M#cy!(Mo2WV`bOXxeY466UsNszCH~o4@_w=pFyqSN=Z#l ztc*{~NlXVN6JBQ)Q2Kj(doVGGS zK4Qv@oyN;O&FY>-?6j43v|v5H$eg1+cv}LH|L=dj{l%oYJB=?e oSWNq9Y;n8MVMkWZQgb=`|J(C!n^6DVnf>e+QQ>Mj~ z&W`peK!d!Qdl+q|`1$#H{RaXt;muGoCCQnC5#knxDQ!S?J#49YsYU6PQ!>~yxO$k< zko0HpVbcxK{J|cgo4F_>AEZ^xnFCdVBPp>sH77GK6{{>~Ze~S1)T2U39tAQ}i!zge zArqfgl%JbanwFLd@~XzP2~bCe^{z<}ih5-a*UHS`_C<3&%NZ-Ci^wTzqS?FKDHO#yyb$)-xc-?u4E|*Sv-g9Z}GdH zt*?34zAj7m=;pMI_7}OmT)Z0|?LT7UapZyWa(nMDPSUOoPvPpP07H#A+0YCT?AU_M zo7tPWZAuTL^OO=$n2W-_AD^6`mla=9l$cjsFeO6`r#~Ce{Mmr!&;4lr+>hqZ{b>Gd zK=bE*G=J_#^Jl}93@w-^+#sG{&rK|eFUT)JaW*Uwz@lfMX}Gg9bipCd>JeX-n2h4& zWGgGft=1y&7=0g`6SbyEalcq%vfF$c9sA=4G#wf3)b=k7U&MV}Rm=VzR9+M!zrTP_ zB{<+qC|v)sGsp7RGz8n9PFu6V-SFZ*nEnM2{rl16?Ur_5hKAOOjN>>lx86)2t>=x|53AKg>zIYpW9#Nz=~JzINW+w2xLTXfAZ zK2&}$jE{?+k|733tI$YiO34KW4$K~yzSU8e0{PN}_jC4q<@L!+*$?HzW5VQT+HC2| z$uN0)C?775ZjT^1A7_FJ_jq9DD9WrriIX|0(s%#=e{X-AW&NElA>;iFw(^mSKb)}p z<#%b~SsCs93?c6=jw^B4Gx+^DelzlmUBU~qeWHI(?z^R5%5VxOe^F-gZIJ#VxzeqX zYk>NX+-@n9WUy~Ye%iCPj?KP-Ls-p8fMY*{vDmKrYrfbqNG3_ndisAKgVWL&p4n{n ZLfb^cuL1RMna*4N;6mkov^Xi%0|3(y3aS7A literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_E1TeV__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_E1TeV__f4.pkl index 2caff49d01b299c539cb8d5eed6f800feb20f1fe..8dc61275c16f9b252997e68d2f3c6ea7eae19b9a 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i85kPuBd5&Rd9&HcY}>Plo!f6XSX^XZVs7IXvCR|6|M$P%o@v9I9lmG&n@_AZ uvDiP+VF%m9CFTYS3_DykO)Nn2*Do2HPsz~oX7*-oo06ehG9_qAsU869c{oA< delta 107 zcmbQkd6|Qyfpx0KMwZFUlOHfAM{+VW*qc>$?sR>6*M$3m+0N@NiWWg8kIln)Gq-~H zOhEbW`p-K!%&wZ}YR@v44pi7Nm*us2_PndxUtXAP&H$9>7OmVqB}2=b*_*j-N``L9 Jl%Of4dH~f4Dk1;? diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_E1TeV__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_E1TeV__f8.pkl index 45ea8dddb9023fa70423a0c89f94ffb6823ccbb6..b169110d245594ba6ee944f16acf76558a22ced9 100644 GIT binary patch delta 231 zcmdnPvyzvkfpzMAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|@9O5gqe|GoVm53Yw@LdN?U7)`xv;Y7A delta 180 zcmZ34fiuJFj(xAO4@fvCvP_6 zEw%j&77^B*{4Dki6ISi&SNdxQl#d3<-+;<%-}~|zq(A4z7PXK0a{Czqc5I%u=9eAA z0)g44@t^G&9_;%vk(XsZ!-ddeGt_|k8}?@}UiE7q!-Xa1o+N%&vfJb-{2ye%sj$M> ag-Y-DP07&mX7*-oo06ehG9_qAsU83@i&IGe diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no__f4.pkl index 6e8c7c4e01351d50cd1d56b281b16762d9bb874d..77e85f73436c00292c9d38e5642c1c9443ddf084 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i85kPuBc{yQX}sLitnOLFPFrb53)bU{%sJYFw_L?li&kmS^xikKQwuXxf!|=vAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|3$mA?D`|9kt}EbH%d2^sHau$7No{NaS%FTYC@&&p`;X9#(3aa@VR zp26?O@tcuf>=Is>?GycTa^EfeQifAN`HM1>Z-ev~$(3%6Tm#g9448c=G@Mdwb`dJ$uyj4fit?@W=_>&)TOQx2vD= zmfHS?-sh#qq*&}3JOkc}GySz=c*3Iql2?GrFZ+M{9!USjeXqV>ePq0!A@lRXm3w~K zG5pyi_S@sL9m6p@g+q)i`x!P`^n7*uYu9kuk)eFnuYC+6+UIyGJ}cRk-ul9D4QRkj dp-tD;y{_CpB}2=b*_*j-N``L9l%Of4dH@|tR;>U4 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_blearth__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_blearth__f4.pkl index fb4405461d2276f39547addc110352b65712dc5c..19168f0ee606607a67587d17f17ae5d944b24012 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i*)@3gYel5k3Ge9GXDR8l-{(u~ZponcyC=Vyvd{Xa{r>s>0{cT8oa}C%F5W$P yy3+36AFkH>rzh@ap2TER)+M#u#-L{ZjAQzCQ!=!?nZ23Yrex@rObMD&ss{j%gEuAs delta 107 zcmV-x0F?in3DXD!fCZIEu>_L?li&kmSt2iazaJ47y(MNByB3gVKB7CvyHM9px&mky zyC{%nJ~>o*zkqg%JWy*-yV$T%IiQRzI|FpJyQPOtIs;KjyW}k!zcHChKa^-HO9M*- NhLmV6bd*w*buP8IDc}GA diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_blearth__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_blearth__f8.pkl index 377e5468744f3ac8195a38e020421b77531f9a98..1cd3d60d1e9812593100bcd6bd9ac6a4b4835254 100644 GIT binary patch delta 231 zcmdnPvyzvkfpzMAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|p}jab1__He&b*NkI=-Dm9mw;wB;9CC4g(YD-W22n5eKRfh9VG~>7 z{#ogg>t3nW?f>Vm?>+g>rTuT*CQLUKe6hd&@A>TTI}i6~rdjWibva>wuko%U^X|O; zS3Ec8DQryHzhqg_-H-B__Dav*2Chxa+5i0mOO<6sgFRQjmBsy|$x8cw_E_HMtb4zo eY3{_;Dr=9~Psz~oX7*-oo06ehG9_qAsU85Rsb~%W delta 180 zcmZ3(^*bwos9$^OSDWoGI8NZT(ctI^Hn zf7ZT8^{>LJ+ZFo{etVXEt?`1#GWR=d zy7TuSgSNfSy{Db`XQtYVPvLYrx}a?Tio0Szcndu3uUTHVT*{HDv_H;VC~ax}t<8 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_nubar__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_nubar__f4.pkl index e652d4db026da8235f9a4b90de33a115b65b430a..2636e0e936d3d1f08da7a13afb717a7ac949731a 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i85kPu54BCU+-h^l+%T7W=ZD(k=G7N%EK4U{+AfyEz4PDydV8l0sylbyeZ1ZF yh1Cud_kG)Czu8zuE`4bJe4h1=9sla>!)Kk@JtafSo7tPWZAyl5$&{cerFsAs>^kcJ delta 107 zcmV-x0F?in3DXD!fCZIEu>_L?li&kmS^xikKRS7dIT0_QLxc~ouKTItcxnT;mw`4bFxV+x6x8pDvIo&w2H|M`*xVrv-KbK|9yOd}uO9M*- NhLmVCbd*w*buMzTGJOC5 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_nubar__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_no_nubar__f8.pkl index 89e251801f2605ad926a1ee80ea818326f41ed96..30dd469781ae27025fcec12cddd4d3c98914824b 100644 GIT binary patch delta 231 zcmdnPvyzvkfpzMAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr__sG`+4X8|M&J!Coj5xy}-!cK;CY;W56%Fmc4hKqjDAZGc;&se|!1E z&f&JJsO~Iv`-1tOtJ*w&?b8y-KQa|4|K{Pc&mjFeR})&tosXGpN|T&N%&jUxJft!QOf5_6$oWS!$U6v17=o`Lv~&U1@*H{GB4#fd(}1jC|}h ZxpM!M3@vYFZ|1fs8O9}3f~J(}0RU;)WIO-> delta 180 zcmV;l089U=4ZI8lfCZH?u>=qWlYs??f8N)~`2y>C>`|dB2osDoX=P1BR4nGjx_L?li&kmS^ocjKQwuXxe&ueF$mB$xu(`2IJzC^H(pFKH|Yu>xbZmZH-zKVw{XFmH?RJGKXqNtyOd}uO9M*- NhLmV6bd*w*buM^WGNJ$g diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_std_decay__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_std_decay__f8.pkl index dead1d19ce27c46bb8085a8a7d9b033270ced2b1..65b3fec04a5be35c6a8e321a8097b3c59007cf7e 100644 GIT binary patch delta 231 zcmdnPvyzvkfpzMAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|3*mA?D`|9kt}EbH%d2^sHaU{Hu${NaS%1HVfX&&p`;SMYsraa@VR zo}uo?@tcuf>=+o#_KE&ExsSo1l;IRm{*uh(+aUeLa-~}%*BI|_IDWgOP?EuZLDtir zwRLRv35>#OP68bJ9lXVM-Cy&?&f#IA YsN6p#L(7}lo4IXDhHlA}pedz#0C4hOaR2}S delta 180 zcmZ3448c=G@Mdwb`dJ$uyj4fiuJG{_6x&)TOIx2vD= zmfHRUlb)9zlVY)NNDX)^&h*y~D6asLSAxo~`hWW#NdM-2ufAV>WV}Bi_L?li&kmS^xikKQwuXxf!|=vAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|3$mA?D`|9kt}EbH%d2^sHau$7No{NaS%FTYC@&&p`;X9#(3aa@VR zp26?O@tcuf>=Is>?GycTa^EfeQifAN`HM1>Z-ev~$(3%6Tm#g9448c=G@Mdwb`dJ$uyj4fit?@W=_>&)TOQx2vD= zmfHS?-sh#qq*&}3JOkc}GySz=c*3Iql2?GrFZ+M{9!USjeXqV>ePq0!A@lRXm3w~K zG5pyi_S@sL9m6p@g+q)i`x!P`^n7*uYu9kuk)eFnuYC+6+UIyGJ}cRk-ul9D4QRkj dp-tD;y{_CpB}2=b*_*j-N``L9l%Of4dH@|tR;>U4 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_std_nsi_no__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_std_nsi_no__f4.pkl index 6e8c7c4e01351d50cd1d56b281b16762d9bb874d..77e85f73436c00292c9d38e5642c1c9443ddf084 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i85kPuBc{yQX}sLitnOLFPFrb53)bU{%sJYFw_L?li&kmS^xikKQwuXxf!|=vAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_|3$mA?D`|9kt}EbH%d2^sHau$7No{NaS%FTYC@&&p`;X9#(3aa@VR zp26?O@tcuf>=Is>?GycTa^EfeQifAN`HM1>Z-ev~$(3%6Tm#g9448c=G@Mdwb`dJ$uyj4fit?@W=_>&)TOQx2vD= zmfHS?-sh#qq*&}3JOkc}GySz=c*3Iql2?GrFZ+M{9!USjeXqV>ePq0!A@lRXm3w~K zG5pyi_S@sL9m6p@g+q)i`x!P`^n7*uYu9kuk)eFnuYC+6+UIyGJ}cRk-ul9D4QRkj dp-tD;y{_CpB}2=b*_*j-N``L9l%Of4dH@|tR;>U4 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_vac_nsi_no__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_vac_nsi_no__f4.pkl index 9d9e2ce42f17c8122bbdd11721fc0394debce7e7..22158002ca2378eefc71c27175dc98b6fa4d3e60 100644 GIT binary patch delta 168 zcmcc2F^7|-fpzM{jVzOyZP;^)GUE&KOQvL~X-rdLfPgzHCK^yCm;&-wAoD@&DH&pu zyO;|i|NgJHU)0zC~oeEPeEcgz;Gv9OFYwH0Z|36TEPVI!9E9QSSpIUUw wJol*i4oSa{=HC>*Z2zfw$9&trdi#R6XLnD@(DG*XW^S93p<6N~XiBLb0MFVzqW}N^ delta 107 zcmV-x0F?in3DXD!fCZIEu>_L?li&kmSpWclKLXiqxnx-ywvHArxxYskIGi0CIOcm8 zHeVJmx&QxvKY=llxjypAH@`Q{HAF=hxJ(oqI33Q(w@d2cHL?DGKjFj9yOd}uO9M*- NhLmV6bd*w*buREKEb0IN diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_vac_nsi_no__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_hostfunc__nufit32_vac_nsi_no__f8.pkl index 7057d29612051c66b323e02236b7ac9a2d336ad7..7dbf261efedbfc786df3500a085ac399f1b1883a 100644 GIT binary patch delta 231 zcmdnPvyzvkfpzMAE2*lZ^=GD~nkIeUQ| z1?DLkVw2ffr_=}id;aME|M&LKbM4Z;Trk+*aL2aXBV?JK8e`$8xFGcu?ex; zGrV%F@45Zgjv>2j-HNX(_X%#lGi4f3{$*Cn1CaiZD}tXk@f+=D_`UB($(Db13}@K_ zkNN$!Yxr6%W%Pq>Kf|4`GmN!>;uC+=J#G27kKsgod{@hFB|GD{W>c>L4QOyzy}NdA Y<^CxdTHegw%xzOLbW5fLO)1p_0084>n*aa+ delta 180 zcmV;l089U=4ZI8lfCZH?u>=qWlYs??f7r|3>Hq)lKZ2Q_iXhu9zo@j+AoIw>zBY^X zEz9C1zX09doC*p9KY&H}q~1XTKe)I5W=RplK651dBhKO_zgCZYB0dQr i5I@%c?>`_U*EoFGdB2osDoX=P1BR4nEp(Jplyxr13|6uL diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_massbasis_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/get_transition_matrix_massbasis_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..a4171be77779b294d197e4cc106fa4abda8e352b GIT binary patch literal 722 zcmZo*oqC*!0StP$lk;Mw=;qetus6fdEW+Gn7n8a%N#*Xt1Bs22|6-mYSDZlwLU{gFS;6 zMLZ=pJ~y#sN{l8JkPCWv8t$gUW9M@w3uaY`d~*u~m$Xu_n3;1WSq%^NKS;sR`ukqRa{u zYeBJWQQ5gOI``>jt_x;6CtOpoc>nLYd2!!TFdrnpNo@Yk3lFcFGy2Ukmk?Iip(g*z a{OQ#z+l4mIHsA2C-o7mM{BAUBOZ5Q21^vMw=;qetus6fdEW+Gn7n8a^`@z=7ar|HlVs5w$!}TqV&os8SELn z7;-7O@wtg5Q+kAu90O#g7G)*_JrJK(l%JbanwFLda*oEd2@sPR9+=uj@6cznhjIz> zpKO1>|I<2kxOzhT&lL(y=SoxH>Iw0uWN`H`CmWg}f+m9xGhn=#y_wsl^e{S4DFKBG z7sT&i|7S=cJO}d>jLykYSjnQSv>(QoY1zDhUp58Khswhp3sVQ9OKTOw?*B@HtB2}` ztB11 zX0x#ZxBXSq>3uA|o2(7$rg4~YGC^%xv)ke2v1N%2*Y{;CI-I$3#eF+1sQldf&6a)2 zhph#6tI6@C@Y^4_u-Mb&%H-V*#+x=JPdm5o!`sOV&hlKf^M~qpg6QA(1S$}nubp`E|Nnb?=bb%!)btJa zGZgU13Ej`yryaMepYfL3{)XP?rN^XL>=`@*-ikB*wPSd~qX3dufXXlXfBPOt|Hgf< zzF&Q0yq_WS^TL&Te%Uen*(CPct0vxM{|Cu9sv6PYYPAX literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/product_hostfunc__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/product_hostfunc__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..83e76b500ad779fec1c4d09af2bf0a2cab716b53 GIT binary patch literal 832 zcmZo*oodR=00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYauFSRJW za!QYAUTJPYrCxG=QL0{UX--LIVo^~dP@JteIWZ@(2&le?6|4kEr<7C{q=H14(oCjw zcC=3cn&i#g!)P#PGgz}#1{;2YYiLX{9+iK>Lb|<5M%!Q}c>T5=&C!lM;(F zi>DyF;O@H*z))#eog~_`dp0OD z?)`NChus(LhZZ37^>wyc?*y8^!gYzQFU0)4Q4sZerMq_QS+C4!3sQey`{9oC2W<8r dugv(n+XhHAB-B5%1M#83@6Fsk1rh(D_9ATPARD@NCk;7rCChr z>}a0?G|8K}htXz=pP!%Ce;@!8-V7yElAJjpZuww8r46XAhb<*HJ~y#sN{bG~faqp;U}_t^L!Zqa$|cBuvi<%3PwUj->Iv~bS12@{ zD@}o`C&ZtU!PUc@Y-on?eg+?==e?P|ncJrHFgi~u0r{N^;&-tBGsF;{gLw)@=VU3Y zWKmYy5978Knse{p_wTfZ)e}xPOlfLKP-d*n8lV;kV&2YU{R;p^F#T>}{ZSIl%TIU-D zBBodDpQ;jeY>w4g`>goilDl-a?Z4#2Rpe{AaeqM#^KAWX+wC)%E0-lF9k&_3*liC8}HN2U1KZrM+JaUk*JeugtJS08;K1B5MKFYp- dJXZhqK2m`IKa^+=O9M*-hLmU$bd*w*buM#M9`yhK diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_io__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_io__f8.pkl index ac3fd936c0d39b1cfc5261af4d3010192a7b7b8c..41e3c13521cd85169136cdbda7f04dfe94afad08 100644 GIT binary patch delta 159 zcmZ3=HHn+0fpu#8MwUKiS@xWw%=m)*k|`PN8q+)&AYiuk!nWBhY_^jhGD>hjIeUQ| z1?DLke4AUD-5Db}H#kl+-TKbn@|=5ldfsIFMS4vKp1zuDFB@5!_PK7RecJ^Vfkm(G z+uyu7d1qejJ^Lp2m5VlS>9=p*)K*b`@1Fe>Mx)Y?f)DJ026{7lGq+925Gt7xG^JDz E0IP#RApigX delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a>&(Hv2cT?mr+5zR2>1lRvXvd(aIMmOs|V z#?KU;nLoRh`b?$g-#@kd2xf_X-aj5h<&xL;k3VjRd4}uX-aoCTcit)x;6IdT4od?| O1BR4n5_FVOlyxqf4l?Ti diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f4.pkl new file mode 100644 index 0000000000000000000000000000000000000000..44b7b778167a9d1c2728033ebbdada80f9a9229d GIT binary patch literal 1009 zcmZo*o%)QK0StP$lk;RK`D)LXcJoZTZKIe(_Or=b+G;ZCfW@a|aQ84JTOgd1 z!H>%^8A4#Eu;(V0#24h3pg4ztp#f|h$VrrdDH)OwpF=I`f%=?2BaFfw0Hyjc(*OVf literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_lri_std_mat__f8.pkl new file mode 100644 index 0000000000000000000000000000000000000000..6480a09b10de733882bee385cacef36a7e9d0930 GIT binary patch literal 1426 zcmZo*o!ZXI00uqW$@w`ssmUdo`FX`tdbs_IQc{ajQ(Q8WOQuYoq8T(rqlYOacS?^? zUTJPYrCxG=QK}x0nOc;Y3{)MTR+OKcRGOBSS_D+BG0g*FD4bj!bt#Z9O?W?N&sScb zyp;V=zCBEV$x>8K`i8+uNY)Ec5!2nbdytkYUd$%7jNJEEmnzUf4Ffh0PYO z9gP8%--{-U!JU%9hY}6IJe33t;0(14O=k{9sB%c~v8Cpv7Nu8$jMKZ?=Wsd)}2TNQ0gNidcJOQe< z=wo(($g?^?jN@taD;>_?~!t z<(~;c4$H2cKUw-i&cX4g_^;hpG#wIq+5)4xO$H{KEX;png-;{4y#pT1B_Bhf77 zE-1=RN=(Yk$t(c|1+d_N#Yn`01qUv6zO%n4p34z1b&CD*J9UrC4QAOZtp7N5$-bHP zeixncpW&2U6()YyzSs3vs51Lh`&-@4hhJyjwZE6Xj``xE2li+&QmO|4e0`v` literal 0 HcmV?d00001 diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no_E1TeV__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no_E1TeV__f4.pkl index 2e06b77c096f7cb323989457faf547121cad3fec..3f7b43452b81f7b4619a417de8cae98bc2fd6419 100644 GIT binary patch delta 132 zcmZ3%{*j%ffpzM$jV!%PHtab?neherB~vokHKwUBK)@Xp6AdU6Oab{Tkoh3?lnlPf zEllrC|NXDGzb$aX$l&lRqX@I>M!$gk1;t{f$F{#P>Q5Fk{SD*;)p|2~Gq+925Gt7x IG^JDz0AK$pQUCw| delta 71 zcmey!zJi^lfpzNKjV!%PlW#FSGSy~iu+L$CVzj5@t`UFP1EVSb>+M$;iJ6u%KQWSx b6*HX-hjIeUQ| z1?DLke4AUD-5H(#1pavZ|NndYoant@7fMCj9pAlBYS)(-JJD52)vhwpb~+EI9{mTD lpL=-Yf&FI}+5IxKnJV@<#?Is3ojYfM{7cJLPUdG70{|A0J{SN1 delta 108 zcmbQly_Ackfpx0XMwUM2&HtIL7$c1%KRy2c|GmAEtk_G>v#a@{jm+1xHh+v*?bq#kVJB_5!Y0c9wOyMSt36PyH?uc$+msBU Kk|{w`O7#H#uO>tQ delta 71 zcmV-N0J#702dD=GfCZJEu>_3*liC8}HBDM{KgCC+JZo;yJe`2eJk&M?KTx{pKBG0K dJlsv|K4u>UKa^+=O9M*-hLmU$bd*w*buJ_#9Gw6F diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no__f8.pkl index 80cc3cf3d0b58d09af76a899f89a69d44f810532..6480a09b10de733882bee385cacef36a7e9d0930 100644 GIT binary patch delta 159 zcmZ3=HHn+0fpu#8MwUKiS@xWw%=m)*k|`PN8q+)&AYiuk!nWBhY_^jhGD>hjIeUQ| z1?DLke4AUD-5Dbz7A!b$vGbk%J@H(QfT>gLkKd_#Ty8MSUSa*msY~|FwD-FxH7|bI z1N#i8?5Z&FyY{`Vze1JSr`q4@c0T+%^RE59^mWV^7d@~C8tBdJ&D=I6L#Sj*(3Da= E0Ln8#0{{R3 delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a@rrGR;Yg?mxXUME&@Pls~o4a$L;~nm-{N z;6W9znLn&*AeY*v;6KzTp?ygg-9J#~))%qnls|Ci>Os3{-9Hi4OV;e8;6IdT4od?| O1BR4n5_FVOlyxplV=_3*liC8}H6uhFKaCMnK4aTAK1n0oJupc>KVZicJ{)b1 dK2IP&J?z#jKa^+=O9M*-hLmU$bd*w*buNsV82hjIeUQ| z1?DLke4AUD-5Dc!T&MipV*JW}HGj!*&C5saMW?M*R(;TFUm8$*-o5sSJ#Q$(Cyju| z_Wz_6-v;PA+0a=cqX$QKd<3E&DZv+@`%0Fl)>fQhK#y<}= zmO{FUyFY(rw+&F^=RY7(*Qa8w!awpahhW{$(Lb!z9@N delta 71 zcmV-N0J#702dD=GfCZJEu>_3*liC8}HD=3oKLt^=Jk>&yJZQL$Jcc_3KVYrwK7HG` dJlR9)J~kHxKa^+=O9M*-hLmU$bd*w*buJgq98LfL diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no_nubar__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_no_nubar__f8.pkl index e4ba68ed8a7c76be15662a1857a51443ac9b47dc..4aea8e530e84eaa731a787df34109e2bb242c9fa 100644 GIT binary patch delta 159 zcmZ3=HHn+0fpu#8MwUKiS@xWw%=m)*k|`PN8q+)&AYiuk!nWBhY_^jhGD>hjIeUQ| z1?DLke4AUD-5DdlZ}wj3zx$p2t2qm&@D)t6?|k~)@vrJ6d+9ZowOkfXwBK+yWohz? z2lh`mJBw7N-nEaowsrOKFVpOAXuUaY?s3=tk@cE8wo4w^0}b?M_GWIIk|9(wC1^^i F9spD-L(2dF delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a+K3nHX%j?mss=|70SHmOrba&WBn^l0Pnz z7%&aDkw1S5ni-pt<2ZBsIY LN~Q!&Db)i26B8&V delta 71 zcmV-N0J#702dD=GfCZJEu>_3*liC8}HEBt7KMr!4JP>=uJWa*IJn$p=KI7@yKE@oJ dJXn3*KD&YaK9pzhjIeUQ| z1?DLke4AUD-5DdFT`!$Gqy3%zgR4&jj#y5%4`7$wseWOGy@Js3z^t?x_FI;1*;uM} z-(KJAHmg_cZTl$C^2`SdC)+FC__6-v;PA+0a-oWUQ8^B?msNV++jj4lRvpfP2^#(m_JPP zVSuSvm_HVHheUiN-#_+yWluX@+dtWk6@8e4lRsXgx~ZwJ+drzPmQFRW-#?US4od?| O1BR4n5_FVOlyxq$+Ain- diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_decay_no__f4.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_decay_no__f4.pkl index 86d2cec225c28fdb0c5dd57c31e6e3e45a27c9e0..44b7b778167a9d1c2728033ebbdada80f9a9229d 100644 GIT binary patch delta 132 zcmZ3%{*j%ffpzM$jV!%PHtab?neherB~vokHKwUBK)@Xp6AdU6Oab{Tkoh3?lnlPf zEllrCC-{}vH>_G>v#a@{jm+1xHh+v*?bq#kVJB_5!Y0c9wOyMSt36PyH?uc$+msBU Kk|{w`O7#H#uO>tQ delta 71 zcmV-N0J#702dD=GfCZJEu>_3*liC8}HBDM{KgCC+JZo;yJe`2eJk&M?KTx{pKBG0K dJlsv|K4u>UKa^+=O9M*-hLmU$bd*w*buJ_#9Gw6F diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_decay_no__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_decay_no__f8.pkl index 80cc3cf3d0b58d09af76a899f89a69d44f810532..6480a09b10de733882bee385cacef36a7e9d0930 100644 GIT binary patch delta 159 zcmZ3=HHn+0fpu#8MwUKiS@xWw%=m)*k|`PN8q+)&AYiuk!nWBhY_^jhGD>hjIeUQ| z1?DLke4AUD-5Dbz7A!b$vGbk%J@H(QfT>gLkKd_#Ty8MSUSa*msY~|FwD-FxH7|bI z1N#i8?5Z&FyY{`Vze1JSr`q4@c0T+%^RE59^mWV^7d@~C8tBdJ&D=I6L#Sj*(3Da= E0Ln8#0{{R3 delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a@rrGR;Yg?mxXUME&@Pls~o4a$L;~nm-{N z;6W9znLn&*AeY*v;6KzTp?ygg-9J#~))%qnls|Ci>Os3{-9Hi4OV;e8;6IdT4od?| O1BR4n5_FVOlyxplV=_G>v#a@{jm+1xHh+v*?bq#kVJB_5!Y0c9wOyMSt36PyH?uc$+msBU Kk|{w`O7#H#uO>tQ delta 71 zcmV-N0J#702dD=GfCZJEu>_3*liC8}HBDM{KgCC+JZo;yJe`2eJk&M?KTx{pKBG0K dJlsv|K4u>UKa^+=O9M*-hLmU$bd*w*buJ_#9Gw6F diff --git a/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_nsi_no__f8.pkl b/pisa_examples/resources/osc/numba_osc_tests_data/propagate_scalar__nufit32_std_nsi_no__f8.pkl index 80cc3cf3d0b58d09af76a899f89a69d44f810532..6480a09b10de733882bee385cacef36a7e9d0930 100644 GIT binary patch delta 159 zcmZ3=HHn+0fpu#8MwUKiS@xWw%=m)*k|`PN8q+)&AYiuk!nWBhY_^jhGD>hjIeUQ| z1?DLke4AUD-5Dbz7A!b$vGbk%J@H(QfT>gLkKd_#Ty8MSUSa*msY~|FwD-FxH7|bI z1N#i8?5Z&FyY{`Vze1JSr`q4@c0T+%^RE59^mWV^7d@~C8tBdJ&D=I6L#Sj*(3Da= E0Ln8#0{{R3 delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a@rrGR;Yg?mxXUME&@Pls~o4a$L;~nm-{N z;6W9znLn&*AeY*v;6KzTp?ygg-9J#~))%qnls|Ci>Os3{-9Hi4OV;e8;6IdT4od?| O1BR4n5_FVOlyxplV=_3*liC8}HK1a2KeuhjIeUQ| z1?DLke4AUD-5Db}mY(Ff(e=*$fYbDwPiIfDUngbKE5SL-zSm`TjsM}9_P^e-FP^jf zf&HmmwL=p{?%F?;FnG1@;S~EPYr|qV((l^4X*5UWEq!1QG|-#bo4IXDhEU0rpedz# E0G#ST82|tP delta 108 zcmV-y0F(cc3#AGKfCZHru>_6-v;PA+0a<+|mqJvG?my4Y^GWxPlt2A|Cu4x-nLmRE z-Lej}nLmj^>c6G@r{e~G7-9Ji5KKf|Elt1r3jQwh6-9NnTXF0#7;6IdT4od?| O1BR4n5_FVOlyxqEVKO-Y From e6a2bc8c107b872586806a2f567389ad002fed32 Mon Sep 17 00:00:00 2001 From: Gopal Garg Date: Thu, 1 Feb 2024 09:00:38 -0600 Subject: [PATCH 2/2] Updated PISA to include LRI potential --- pisa/stages/osc/prob3_bk.py | 407 -------- .../osc/prob3numba/numba_osc_hostfuncs_bk.py | 218 ----- .../osc/prob3numba/numba_osc_kernel_bk.py | 879 ------------------ .../osc/prob3numba/numba_osc_tests_bk.py | 791 ---------------- 4 files changed, 2295 deletions(-) delete mode 100644 pisa/stages/osc/prob3_bk.py delete mode 100644 pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py delete mode 100644 pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py delete mode 100755 pisa/stages/osc/prob3numba/numba_osc_tests_bk.py diff --git a/pisa/stages/osc/prob3_bk.py b/pisa/stages/osc/prob3_bk.py deleted file mode 100644 index 391bb4d43..000000000 --- a/pisa/stages/osc/prob3_bk.py +++ /dev/null @@ -1,407 +0,0 @@ -""" -PISA pi stage for the calculation of earth layers and osc. probabilities - -Maybe it would amke sense to split this up into a separate earth layer stage -and an osc. stage....todo - -""" - -from __future__ import absolute_import, print_function, division - -import numpy as np -from numba import guvectorize - -from pisa import FTYPE, CTYPE, TARGET, ureg -from pisa.core.stage import Stage -from pisa.utils.log import logging -from pisa.utils.profiler import profile -from pisa.stages.osc.nsi_params import StdNSIParams, VacuumLikeNSIParams -from pisa.stages.osc.osc_params import OscParams -from pisa.stages.osc.decay_params import DecayParams -from pisa.stages.osc.layers import Layers -from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import propagate_array, fill_probs -from pisa.utils.numba_tools import WHERE -from pisa.utils.resources import find_resource - - -class prob3(Stage): - """ - Prob3-like oscillation PISA Pi class - - Parameters - ---------- - params - Expected params .. :: - - detector_depth : float - earth_model : PREM file path - prop_height : quantity (dimensionless) - YeI : quantity (dimensionless) - YeO : quantity (dimensionless) - YeM : quantity (dimensionless) - density_scale : quantity (dimensionless) - theta12 : quantity (angle) - theta13 : quantity (angle) - theta23 : quantity (angle) - deltam21 : quantity (mass^2) - deltam31 : quantity (mass^2) - deltacp : quantity (angle) - eps_scale : quantity(dimensionless) - eps_prime : quantity(dimensionless) - phi12 : quantity(angle) - phi13 : quantity(angle) - phi23 : quantity(angle) - alpha1 : quantity(angle) - alpha2 : quantity(angle) - deltansi : quantity(angle) - eps_ee : quantity (dimensionless) - eps_emu_magn : quantity (dimensionless) - eps_emu_phase : quantity (angle) - eps_etau_magn : quantity (dimensionless) - eps_etau_phase : quantity (angle) - eps_mumu : quantity(dimensionless) - eps_mutau_magn : quantity (dimensionless) - eps_mutau_phase : quantity (angle) - eps_tautau : quantity (dimensionless) - decay_alpha3 : quantity (energy^2) - - **kwargs - Other kwargs are handled by Stage - ----- - - """ - - def __init__( - self, - nsi_type=None, - reparam_mix_matrix=False, - neutrino_decay=False, - scale_density=False, - **std_kwargs, - ): - - expected_params = ( - 'detector_depth', - 'earth_model', - 'prop_height', - 'YeI', - 'YeO', - 'YeM', - 'theta12', - 'theta13', - 'theta23', - 'deltam21', - 'deltam31', - 'deltacp' - ) - - # Add (matter) density scale as free parameter? - if scale_density: - expected_params = expected_params + ('density_scale',) - - # Check whether and if so with which NSI parameters we are to work. - if nsi_type is not None: - choices = ['standard', 'vacuum-like'] - nsi_type = nsi_type.strip().lower() - if not nsi_type in choices: - raise ValueError( - 'Chosen NSI type "%s" not available! Choose one of %s.' - % (nsi_type, choices) - ) - self.nsi_type = nsi_type - """Type of NSI to assume.""" - - self.reparam_mix_matrix = reparam_mix_matrix - """Use a PMNS mixing matrix parameterisation that differs from - the standard one by an overall phase matrix - diag(e^(i*delta_CP), 1, 1). This has no impact on - oscillation probabilities in the *absence* of NSI.""" - - self.neutrino_decay = neutrino_decay - - if neutrino_decay: - self.decay_flag = 1 - else : - self.decay_flag = -1 - - """Invoke neutrino decay with neutrino oscillation.""" - - if self.nsi_type is None: - nsi_params = () - elif self.nsi_type == 'vacuum-like': - nsi_params = ('eps_scale', - 'eps_prime', - 'phi12', - 'phi13', - 'phi23', - 'alpha1', - 'alpha2', - 'deltansi' - ) - elif self.nsi_type == 'standard': - nsi_params = ('eps_ee', - 'eps_emu_magn', - 'eps_emu_phase', - 'eps_etau_magn', - 'eps_etau_phase', - 'eps_mumu', - 'eps_mutau_magn', - 'eps_mutau_phase', - 'eps_tautau' - ) - - if self.neutrino_decay : - decay_params = ('decay_alpha3',) - else: - decay_params = () - - expected_params = expected_params + nsi_params + decay_params - - # init base class - super().__init__( - expected_params=expected_params, - **std_kwargs, - ) - - - self.layers = None - self.osc_params = None - self.nsi_params = None - self.decay_params = None - self.decay_matrix = None - # Note that the interaction potential (Hamiltonian) just scales with the - # electron density N_e for propagation through the Earth, - # even(to very good approx.) in the presence of generalised interactions - # (NSI), which is why we can simply treat it as a constant here. - self.gen_mat_pot_matrix_complex = None - """Interaction Hamiltonian without the factor sqrt(2)*G_F*N_e.""" - self.YeI = None - self.YeO = None - self.YeM = None - - def setup_function(self): - - # object for oscillation parameters - self.osc_params = OscParams() - if self.reparam_mix_matrix: - logging.debug( - 'Working with reparameterizated version of mixing matrix.' - ) - else: - logging.debug( - 'Working with standard parameterization of mixing matrix.' - ) - if self.nsi_type == 'vacuum-like': - logging.debug('Working in vacuum-like NSI parameterization.') - self.nsi_params = VacuumLikeNSIParams() - elif self.nsi_type == 'standard': - logging.debug('Working in standard NSI parameterization.') - self.nsi_params = StdNSIParams() - - if self.neutrino_decay: - logging.debug('Working with neutrino decay') - self.decay_params = DecayParams() - - # setup the layers - #if self.params.earth_model.value is not None: - earth_model = find_resource(self.params.earth_model.value) - self.YeI = self.params.YeI.value.m_as('dimensionless') - self.YeO = self.params.YeO.value.m_as('dimensionless') - self.YeM = self.params.YeM.value.m_as('dimensionless') - prop_height = self.params.prop_height.value.m_as('km') - detector_depth = self.params.detector_depth.value.m_as('km') - self.layers = Layers(earth_model, detector_depth, prop_height) - self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) - - # --- calculate the layers --- - if self.is_map: - # speed up calculation by adding links - # as layers don't care about flavour - self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', - 'nue_nc', 'numu_nc', 'nutau_nc', - 'nuebar_cc', 'numubar_cc', 'nutaubar_cc', - 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) - - for container in self.data: - self.layers.calcLayers(container['true_coszen']) - container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) - container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) - - # don't forget to un-link everything again - self.data.unlink_containers() - - # --- setup empty arrays --- - if self.is_map: - self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', - 'nue_nc', 'numu_nc', 'nutau_nc']) - self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', - 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) - for container in self.data: - container['probability'] = np.empty((container.size, 3, 3), dtype=FTYPE) - self.data.unlink_containers() - - # setup more empty arrays - for container in self.data: - container['prob_e'] = np.empty((container.size), dtype=FTYPE) - container['prob_mu'] = np.empty((container.size), dtype=FTYPE) - - def calc_probs(self, nubar, e_array, rho_array, len_array, out): - ''' wrapper to execute osc. calc ''' - if self.reparam_mix_matrix: - mix_matrix = self.osc_params.mix_matrix_reparam_complex - else: - mix_matrix = self.osc_params.mix_matrix_complex - - logging.debug('mat pot:\n%s' - % self.gen_mat_pot_matrix_complex) - logging.debug('decay mat:\n%s' - % self.decay_matix) - - propagate_array(self.osc_params.dm_matrix, # pylint: disable = unexpected-keyword-arg, no-value-for-parameter - mix_matrix, - self.gen_mat_pot_matrix_complex, - self.decay_flag, - self.decay_matix, - nubar, - e_array, - rho_array, - len_array, - out=out - ) - - def compute_function(self): - - if self.is_map: - # speed up calculation by adding links - self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', - 'nue_nc', 'numu_nc', 'nutau_nc']) - self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', - 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) - - # this can be done in a more clever way (don't have to recalculate all paths) - YeI = self.params.YeI.value.m_as('dimensionless') - YeO = self.params.YeO.value.m_as('dimensionless') - YeM = self.params.YeM.value.m_as('dimensionless') - if YeI != self.YeI or YeO != self.YeO or YeM != self.YeM: - self.YeI = YeI; self.YeO = YeO; self.YeM = YeM - self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) - for container in self.data: - self.layers.calcLayers(container['true_coszen']) - container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) - container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) - - # Matter density scale is a free parameter? - if 'density_scale' in self.params.names: - density_scale = self.params.density_scale.value.m_as('dimensionless') - else: - density_scale = 1. - - # some safety checks on units - # trying to avoid issue of angles with no dimension being assumed to be radians - # here we enforce the user must speficy a valid angle unit - for angle_param in [self.params.theta12, self.params.theta13, self.params.theta23, self.params.deltacp] : - assert angle_param.value.units != ureg.dimensionless, "Param %s is dimensionless, but should have angle units [rad, degree]" % angle_param.name - - # --- update mixing params --- - self.osc_params.theta12 = self.params.theta12.value.m_as('rad') - self.osc_params.theta13 = self.params.theta13.value.m_as('rad') - self.osc_params.theta23 = self.params.theta23.value.m_as('rad') - self.osc_params.dm21 = self.params.deltam21.value.m_as('eV**2') - self.osc_params.dm31 = self.params.deltam31.value.m_as('eV**2') - self.osc_params.deltacp = self.params.deltacp.value.m_as('rad') - if self.nsi_type == 'vacuum-like': - self.nsi_params.eps_scale = self.params.eps_scale.value.m_as('dimensionless') - self.nsi_params.eps_prime = self.params.eps_prime.value.m_as('dimensionless') - self.nsi_params.phi12 = self.params.phi12.value.m_as('rad') - self.nsi_params.phi13 = self.params.phi13.value.m_as('rad') - self.nsi_params.phi23 = self.params.phi23.value.m_as('rad') - self.nsi_params.alpha1 = self.params.alpha1.value.m_as('rad') - self.nsi_params.alpha2 = self.params.alpha2.value.m_as('rad') - self.nsi_params.deltansi = self.params.deltansi.value.m_as('rad') - elif self.nsi_type == 'standard': - self.nsi_params.eps_ee = self.params.eps_ee.value.m_as('dimensionless') - self.nsi_params.eps_emu = ( - (self.params.eps_emu_magn.value.m_as('dimensionless'), - self.params.eps_emu_phase.value.m_as('rad')) - ) - self.nsi_params.eps_etau = ( - (self.params.eps_etau_magn.value.m_as('dimensionless'), - self.params.eps_etau_phase.value.m_as('rad')) - ) - self.nsi_params.eps_mumu = self.params.eps_mumu.value.m_as('dimensionless') - self.nsi_params.eps_mutau = ( - (self.params.eps_mutau_magn.value.m_as('dimensionless'), - self.params.eps_mutau_phase.value.m_as('rad')) - ) - self.nsi_params.eps_tautau = self.params.eps_tautau.value.m_as('dimensionless') - - if self.neutrino_decay: - self.decay_params.decay_alpha3 = self.params.decay_alpha3.value.m_as('eV**2') - - # now we can proceed to calculate the generalised matter potential matrix - std_mat_pot_matrix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) - std_mat_pot_matrix[0, 0] += 1.0 - - # add effective nsi coupling matrix - if self.nsi_type is not None: - logging.debug('NSI matrix:\n%s' % self.nsi_params.eps_matrix) - self.gen_mat_pot_matrix_complex = ( - std_mat_pot_matrix + self.nsi_params.eps_matrix - ) - logging.debug('Using generalised matter potential:\n%s' - % self.gen_mat_pot_matrix_complex) - else: - self.gen_mat_pot_matrix_complex = std_mat_pot_matrix - logging.debug('Using standard matter potential:\n%s' - % self.gen_mat_pot_matrix_complex) - - if self.neutrino_decay: - self.decay_matix = self.decay_params.decay_matrix - logging.debug('Decay matrix:\n%s' % self.decay_params.decay_matrix) - else : - self.decay_matix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE) - - - for container in self.data: - self.calc_probs(container['nubar'], - container['true_energy'], - container['densities']*density_scale, - container['distances'], - out=container['probability'], - ) - container.mark_changed('probability') - - # the following is flavour specific, hence unlink - self.data.unlink_containers() - - for container in self.data: - # initial electrons (0) - fill_probs(container['probability'], - 0, - container['flav'], - out=container['prob_e'], - ) - # initial muons (1) - fill_probs(container['probability'], - 1, - container['flav'], - out=container['prob_mu'], - ) - - container.mark_changed('prob_e') - container.mark_changed('prob_mu') - - - def apply_function(self): - - # maybe speed up like this? - #self.data.representation = self.calc_mode - #for container in self.data: - # container['oscillated_flux'] = (container['nu_flux'][:,0] * container['prob_e']) + (container['nu_flux'][:,1] * container['prob_mu']) - - #self.data.representation = self.apply_mode - - # update the outputted weights - for container in self.data: - container['weights'] *= (container['nu_flux'][:,0] * container['prob_e']) + (container['nu_flux'][:,1] * container['prob_mu']) - diff --git a/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py b/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py deleted file mode 100644 index 2a988f6bb..000000000 --- a/pisa/stages/osc/prob3numba/numba_osc_hostfuncs_bk.py +++ /dev/null @@ -1,218 +0,0 @@ -# pylint: disable = invalid-name - -""" -Host function wrappers for numba_osc_kernels. -""" - -from __future__ import absolute_import, print_function, division - -__all__ = ["FX", "CX", "IX", "propagate_array", "fill_probs"] - -import numpy as np -from numba import guvectorize, njit - -from pisa import FTYPE, ITYPE, TARGET -from pisa.stages.osc.prob3numba.numba_osc_kernels import ( - # osc_probs_vacuum_kernel, - osc_probs_layers_kernel, - get_transition_matrix, - get_transition_matrix_massbasis, - get_H_vac, - get_H_decay, - get_H_mat, - get_dms, - get_dms_numerical, - get_product, - convert_from_mass_eigenstate, -) - - -assert FTYPE in [np.float32, np.float64], str(FTYPE) - -FX = "f4" if FTYPE == np.float32 else "f8" -"""Float string code to use, understood by both Numba and Numpy""" - -CX = "c8" if FTYPE == np.float32 else "c16" -"""Complex string code to use, understood by both Numba and Numpy""" - -IX = "i4" if ITYPE == np.int32 else "i8" -"""Signed integer string code to use, understood by both Numba and Numpy""" - - -# @guvectorize( -# [f"({FX}[:,:], {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:,:])"], -# "(a,a), (a,a), (), (), (i) -> (a,a)", -# target=TARGET, -# ) -# def propagate_array_vacuum(dm, mix, nubar, energy, distances, probability): -# """wrapper to run `osc_probs_vacuum_kernel` from host (whether TARGET is -# "cuda" or "host")""" -# osc_probs_vacuum_kernel(dm, mix, nubar, energy, distances, probability) -# -# -# @njit([f"({FX}[:,:], {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:,:])"], target=TARGET) -# def propagate_scalar_vacuum(dm, mix, nubar, energy, distances, probability): -# """wrapper to run `osc_probs_vacuum_kernel` from host (whether TARGET is -# "cuda" or "host")""" -# osc_probs_vacuum_kernel(dm, mix, nubar, energy, distances, probability) - - -@guvectorize( - [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], - "(a,a), (a,a), (b,c), (), (b,c), (), (), (i), (i) -> (a,a)", - target=TARGET, -) -def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability): - """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET - is "cuda" or "host")""" - osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability - ) - - -@njit( - [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"], - parallel=TARGET == "parallel" -) -def propagate_scalar( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability -): - """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET - is "cuda" or "host")""" - osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability - ) - - -@njit( - [ - "(" - f"{IX}, " # nubar - f"{FX}, " # energy - f"{FX}, " # rho - f"{FX}, " # baseline - f"{CX}[:,:], " # mix_nubar - f"{CX}[:,:], " # mix_nubar_conj_transp - f"{CX}[:,:], " # mat_pot - f"{CX}[:,:], " # H_vac - f"{IX}, " # neutrino decay flag - f"{CX}[:,:], " # H_decay - f"{FX}[:,:], " # dm - f"{CX}[:,:], " # transition_matrix - ")" - ], - parallel=TARGET == "parallel" -) -def get_transition_matrix_hostfunc( - nubar, - energy, - rho, - baseline, - mix_nubar, - mix_nubar_conj_transp, - mat_pot, - H_vac, - decay_flag, - H_decay, - dm, - transition_matrix, -): - """wrapper to run `get_transition_matrix` from host (whether TARGET is - "cuda" or "host")""" - get_transition_matrix( - nubar, - energy, - rho, - baseline, - mix_nubar, - mix_nubar_conj_transp, - mat_pot, - H_vac, - decay_flag, - H_decay, - dm, - transition_matrix, - ) - - -@njit([f"({FX}, {FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") -def get_transition_matrix_massbasis_hostfunc( - baseline, - energy, - dm_mat, - dm_mat_mat, - H_full_mass_eigenstate_basis, - transition_matrix, -): - """wrapper to run `get_transition_matrix_massbasis` from host (whether - TARGET is "cuda" or "host")""" - get_transition_matrix_massbasis( - baseline, - energy, - dm_mat, - dm_mat_mat, - H_full_mass_eigenstate_basis, - transition_matrix, - ) - - -@njit([f"({CX}[:,:], {CX}[:,:], {FX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") -def get_H_vac_hostfunc(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac): - """wrapper to run `get_H_vac` from host (whether TARGET is "cuda" or "host")""" - get_H_vac(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac) - -@njit([f"({CX}[:,:], {CX}[:,:], {FX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") -def get_H_decay_hostfunc(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay): - """wrapper to run `get_H_decay` from host (whether TARGET is "cuda" or "host")""" - get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay) - -# @guvectorize( -# [f"({FX}, {CX}[:,:], {IX}, {CX}[:,:])"], "(), (m, m), () -> (m, m)" -# ) -@njit([f"({FX}, {CX}[:,:], {IX}, {CX}[:,:])"], parallel=TARGET == "parallel") -def get_H_mat_hostfunc(rho, mat_pot, nubar, H_mat): - """wrapper to run `get_H_mat` from host (whether TARGET is "cuda" or "host")""" - get_H_mat(rho, mat_pot, nubar, H_mat) - - -@njit([f"({FX}, {CX}[:,:], {FX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") -def get_dms_hostfunc(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat): - """wrapper to run `get_dms` from host (whether TARGET is "cuda" or "host")""" - get_dms(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat) - -@njit([f"({FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:])"], parallel=TARGET == "parallel") -def get_dms_numerical_hostfunc(energy, H_full, dm_mat_mat, dm_mat): - """wrapper to run `get_dms` from host (whether TARGET is "cuda" or "host")""" - get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat) - -@njit([f"({FX}, {CX}[:,:], {CX}[:,:], {CX}[:,:], {CX}[:,:,:])"], parallel=TARGET == "parallel") -def get_product_hostfunc( - energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product -): - """wrapper to run `get_product` from host (whether TARGET is "cuda" or "host")""" - get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product) - - -@njit([f"({IX}, {CX}[:,:], {CX}[:])"], parallel=TARGET == "parallel") -def convert_from_mass_eigenstate_hostfunc(state, mix_nubar, psi): - """wrapper to run `convert_from_mass_eigenstate` from host (whether TARGET - is "cuda" or "host")""" - convert_from_mass_eigenstate(state, mix_nubar, psi) - - -@guvectorize( - [f"({FX}[:,:], {IX}, {IX}, {FX}[:])"], "(a,b), (), () -> ()" -) -def fill_probs(probability, initial_flav, flav, out): - """Fill `out` with transition probabilities to go from `initial_flav` to - `flav`, from values in `probaility` matrix. - - Parameters - ---------- - probability : real 2d array - initial_flav : signed(?) int (int4 or int8) - flav : signed(?) int (int4 or int8) - out : real 1-d array - - """ - out[0] = probability[initial_flav, flav] diff --git a/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py b/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py deleted file mode 100644 index 76d1f9885..000000000 --- a/pisa/stages/osc/prob3numba/numba_osc_kernel_bk.py +++ /dev/null @@ -1,879 +0,0 @@ -# pylint: disable = not-callable, invalid-name, too-many-nested-blocks - - -""" -Neutrino flavour oscillation in matter calculation -Based on the original prob3++ implementation of Roger Wendell -http://www.phy.duke.edu/~raw22/public/Prob3++/ (2012) - -See `numba_osc_tests.py` for unit tests of functions in this module. -""" - - -from __future__ import absolute_import, print_function, division - -__all__ = [ - # "osc_probs_vacuum_kernel", - "osc_probs_layers_kernel", - "get_transition_matrix", -] - -__version__ = "0.2" - -import cmath -import math -import numpy as np - -from pisa.utils.numba_tools import ( - myjit, - conjugate_transpose, - conjugate, - matrix_dot_matrix, - matrix_dot_vector, - clear_matrix, - copy_matrix, - cuda, - ctype, - ftype, -) - - -# TODO/FIXME: osc_probs_vacuum_kernel produces non-unitary results. No one -# should use this until the issue is resolved. - -# @myjit -# def osc_probs_vacuum_kernel(dm, mix, nubar, energy, distance_in_layer, osc_probs): -# """ Calculate vacumm mixing probabilities -# -# Parameters -# ---------- -# dm : real 2d array -# Mass splitting matrix, eV^2 -# -# mix : complex 2d array -# PMNS mixing matrix -# -# nubar : int -# +1 for neutrinos, -1 for antineutrinos -# -# energy : float -# Neutrino energy, GeV -# -# distance_in_layer : real 1d-array -# Baselines (will be summed up), km -# -# osc_probs : real 2d array (empty) -# Returned oscillation probabilities in the form: -# osc_prob[i,j] = probability of flavor i to oscillate into flavor j -# with 0 = electron, 1 = muon, 3 = tau -# -# -# Notes -# ----- -# This is largely unvalidated so far -# -# """ -# -# # no need to conjugate mix matrix, as we anyway only need real part -# # can this be right? -# -# clear_matrix(osc_probs) -# osc_probs_local = cuda.local.array(shape=(3, 3), dtype=ftype) -# -# # sum up length from all layers -# baseline = 0.0 -# for i in range(distance_in_layer.shape[0]): -# baseline += distance_in_layer[i] -# -# # make more precise 20081003 rvw -# l_over_e = 1.26693281 * baseline / energy -# s21 = math.sin(dm[1, 0] * l_over_e) -# s32 = math.sin(dm[2, 0] * l_over_e) -# s31 = math.sin((dm[2, 1] + dm[3, 2]) * l_over_e) -# -# # does anybody understand this loop? -# # ista = abs(*nutype) - 1 -# for ista in range(3): -# for iend in range(2): -# osc_probs_local[ista, iend] = ( -# (mix[ista, 0].real * mix[ista, 1].real * s21) ** 2 -# + (mix[ista, 1].real * mix[ista, 2].real * s32) ** 2 -# + (mix[ista, 2].real * mix[ista, 0].real * s31) ** 2 -# ) -# if iend == ista: -# osc_probs_local[ista, iend] = 1.0 - 4.0 * osc_probs_local[ista, iend] -# else: -# osc_probs_local[ista, iend] = -4.0 * osc_probs_local[ista, iend] -# -# osc_probs_local[ista, 2] = ( -# 1.0 - osc_probs_local[ista, 0] - osc_probs_local[ista, 1] -# ) -# -# # is this necessary? -# if nubar > 0: -# copy_matrix(osc_probs_local, osc_probs) -# else: -# for i in range(3): -# for j in range(3): -# osc_probs[i, j] = osc_probs_local[j, i] - - -@myjit -def osc_probs_layers_kernel( - dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, density_in_layer, distance_in_layer, osc_probs -): - """ Calculate oscillation probabilities - - given layers of length and density - - Parameters - ---------- - dm : real 2d array - Mass splitting matrix, eV^2 - - mix : complex 2d array - PMNS mixing matrix - - mat_pot : complex 2d array - Generalised matter potential matrix without "a" factor (will be - multiplied with "a" factor); set to diag([1, 0, 0]) for only standard - oscillations - - decay_flag: int - +1 forstandard oscillations + decay, -1 for standard oscillations - - mat_decay : complex 2d array - decay matrix with -j*alpha3 = [2,2] element - - nubar : real int, scalar or Nd array (broadcast dim) - 1 for neutrinos, -1 for antineutrinos - - energy : real float, scalar or Nd array (broadcast dim) - Neutrino energy, GeV - - density_in_layer : real 1d array - Density of each layer, moles of electrons / cm^2 - - distance_in_layer : real 1d array - Distance of each layer traversed, km - - osc_probs : real (N+2)-d array (empty) - Returned oscillation probabilities in the form: - osc_prob[i,j] = probability of flavor i to oscillate into flavor j - with 0 = electron, 1 = muon, 3 = tau - - - Notes - ----- - !!! Right now, because of CUDA, the maximum number of layers - is hard coded and set to 120 (59Layer PREM + Atmosphere). - This is used for cached layer computation, where earth layer, which - are typically traversed twice (it's symmetric) are not recalculated - but rather cached.. - - """ - - # 3x3 complex - H_vac = cuda.local.array(shape=(3, 3), dtype=ctype) - H_decay = cuda.local.array(shape=(3, 3), dtype=ctype) - mix_nubar = cuda.local.array(shape=(3, 3), dtype=ctype) - mix_nubar_conj_transp = cuda.local.array(shape=(3, 3), dtype=ctype) - transition_product = cuda.local.array(shape=(3, 3), dtype=ctype) - transition_matrix = cuda.local.array(shape=(3, 3), dtype=ctype) - tmp = cuda.local.array(shape=(3, 3), dtype=ctype) - - clear_matrix(H_vac) - clear_matrix(H_decay) - clear_matrix(osc_probs) - - # 3-vector complex - raw_input_psi = cuda.local.array(shape=(3), dtype=ctype) - output_psi = cuda.local.array(shape=(3), dtype=ctype) - - use_mass_eigenstates = False - - cache = True - # cache = False - - # TODO: - # * ensure convention below is respected in MC reweighting - # (nubar > 0 for nu, < 0 for anti-nu) - # * nubar is passed in, so could already pass in the correct form - # of mixing matrix, i.e., possibly conjugated - if nubar > 0: - # in this case the mixing matrix is left untouched - copy_matrix(mix, mix_nubar) - - else: - # here we need to complex conjugate all entries - # (note that this only changes calculations with non-clear_matrix deltacp) - conjugate(mix, mix_nubar) - - conjugate_transpose(mix_nubar, mix_nubar_conj_transp) - - get_H_vac(mix_nubar, mix_nubar_conj_transp, dm, H_vac) - #Compute the decay matrix in the flavor basis - get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay) - - - if cache: - # allocate array to store all the transition matrices - # doesn't work in cuda...needs fixed shape - transition_matrices = cuda.local.array(shape=(120, 3, 3), dtype=ctype) - - # loop over layers - for i in range(distance_in_layer.shape[0]): - density = density_in_layer[i] - distance = distance_in_layer[i] - if distance > 0.0: - layer_matrix_index = -1 - # chaeck if exists - for j in range(i): - # if density_in_layer[j] == density and distance_in_layer[j] == distance: - if (abs(density_in_layer[j] - density) < 1e-5) and ( - abs(distance_in_layer[j] - distance) < 1e-5 - ): - layer_matrix_index = j - - # use from cached - if layer_matrix_index >= 0: - for j in range(3): - for k in range(3): - transition_matrices[i, j, k] = transition_matrices[ - layer_matrix_index, j, k - ] - - # only calculate if necessary - else: - get_transition_matrix( - nubar, - energy, - density, - distance, - mix_nubar, - mix_nubar_conj_transp, - mat_pot, - H_vac, - decay_flag, - H_decay, - dm, - transition_matrix, - ) - # copy - for j in range(3): - for k in range(3): - transition_matrices[i, j, k] = transition_matrix[j, k] - else: - # identity matrix - for j in range(3): - for k in range(3): - if j == k: - transition_matrix[j, k] = 0.0 - else: - transition_matrix[j, k] = 1.0 - - # now multiply them all - first_layer = True - for i in range(distance_in_layer.shape[0]): - distance = distance_in_layer[i] - if distance > 0.0: - for j in range(3): - for k in range(3): - transition_matrix[j, k] = transition_matrices[i, j, k] - if first_layer: - copy_matrix(transition_matrix, transition_product) - first_layer = False - else: - matrix_dot_matrix(transition_matrix, transition_product, tmp) - copy_matrix(tmp, transition_product) - - else: - # non-cache loop - first_layer = True - for i in range(distance_in_layer.shape[0]): - density = density_in_layer[i] - distance = distance_in_layer[i] - # only do something if distance > 0. - if distance > 0.0: - get_transition_matrix( - nubar, - energy, - density, - distance, - mix_nubar, - mix_nubar_conj_transp, - mat_pot, - H_vac, - decay_flag, - H_decay, - dm, - transition_matrix, - ) - if first_layer: - copy_matrix(transition_matrix, transition_product) - first_layer = False - else: - matrix_dot_matrix(transition_matrix, transition_product, tmp) - copy_matrix(tmp, transition_product) - - # convrt to flavour eigenstate basis - matrix_dot_matrix(transition_product, mix_nubar_conj_transp, tmp) - matrix_dot_matrix(mix_nubar, tmp, transition_product) - - # loop on neutrino types, and compute probability for neutrino i: - for i in range(3): - for j in range(3): - raw_input_psi[j] = 0.0 - - if use_mass_eigenstates: - convert_from_mass_eigenstate(i + 1, mix_nubar, raw_input_psi) - else: - raw_input_psi[i] = 1.0 - - matrix_dot_vector(transition_product, raw_input_psi, output_psi) - osc_probs[i][0] += output_psi[0].real ** 2 + output_psi[0].imag ** 2 - osc_probs[i][1] += output_psi[1].real ** 2 + output_psi[1].imag ** 2 - osc_probs[i][2] += output_psi[2].real ** 2 + output_psi[2].imag ** 2 - - -@myjit -def get_transition_matrix( - nubar, - energy, - rho, - baseline, - mix_nubar, - mix_nubar_conj_transp, - mat_pot, - H_vac, - decay_flag, - H_decay, - dm, - transition_matrix, -): - """ Calculate neutrino flavour transition amplitude matrix - - Parameters - ---------- - nubar : int - +1 for neutrinos, -1 for antineutrinos - - energy : real float - Neutrino energy, GeV - - rho : real float - Electron number density (in moles/cm^3) (numerically, this is just the - product of electron fraction and mass density in g/cm^3, since the - number of grams per cm^3 corresponds to the number of moles of nucleons - per cm^3) - - baseline : real float - Baseline, km - - mix_nubar : complex 2d array - Mixing matrix, already conjugated if antineutrino - - mix_nubar_conj_transp : complex conjugate 2d array - Conjugate transpose of mix_nubar - - mat_pot : complex 2d array - Generalised matter potential matrix without "a" factor (will be - multiplied with "a" factor); set to diag([1, 0, 0]) for only standard - oscillations - - H_vac : complex 2d array - Hamiltonian in vacuum, without the 1/2E term - - H_decay : complex 2d array - Decay matrix, without the 1/2E term - - decay_flag: int - +1 forstandard oscillations + decay, -1 for standard oscillations - - dm : real 2d array - Mass splitting matrix, eV^2 - - transition_matrix : complex 2d array (empty) - Transition matrix in mass eigenstate basis - - Notes - ----- - For neutrino (nubar > 0) or antineutrino (nubar < 0) with energy energy - traversing layer of matter of uniform density rho with thickness baseline - - """ - - H_mat = cuda.local.array(shape=(3, 3), dtype=ctype) - dm_mat = cuda.local.array(shape=(3, 3), dtype=ctype) - dm_mat_mat = cuda.local.array(shape=(3, 3), dtype=ctype) - #H_decay = cuda.local.array(shape=(3, 3), dtype=ctype) - H_full = cuda.local.array(shape=(3, 3), dtype=ctype) - tmp = cuda.local.array(shape=(3, 3), dtype=ctype) - H_full_mass_eigenstate_basis = cuda.local.array(shape=(3, 3), dtype=ctype) - - # Compute the matter potential including possible generalized interactions - # in the flavor basis - - get_H_mat(rho, mat_pot, nubar, H_mat) - - # Get the full Hamiltonian by adding together matter and vacuum parts - one_over_two_e = 0.5 / energy - - if (decay_flag == 1): - for i in range(3): - for j in range(3): - H_full[i, j] = (H_vac[i, j] + H_decay[i,j]) * one_over_two_e + H_mat[i, j] - - # Calculate modified mass eigenvalues in matter from the full Hamiltonian - get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat) - - else: - for i in range(3): - for j in range(3): - H_full[i, j] = H_vac[i, j] * one_over_two_e + H_mat[i, j] - - - # Calculate modified mass eigenvalues in matter from the full Hamiltonian and - # the vacuum mass splittings - get_dms(energy, H_full, dm, dm_mat_mat, dm_mat) - - # Now we transform the full Hamiltonian back into the - # mass eigenstate basis so we don't need to compute products of the effective - # mixing matrix elements explicitly - matrix_dot_matrix(H_full, mix_nubar, tmp) - matrix_dot_matrix(mix_nubar_conj_transp, tmp, H_full_mass_eigenstate_basis) - - # We can now proceed to calculating the transition amplitude from the Hamiltonian - # in the mass basis and the effective mass splittings - get_transition_matrix_massbasis( - baseline, - energy, - dm_mat, - dm_mat_mat, - H_full_mass_eigenstate_basis, - transition_matrix, - ) - - -@myjit -def get_transition_matrix_massbasis( - baseline, - energy, - dm_mat, - dm_mat_mat, - H_full_mass_eigenstate_basis, - transition_matrix, -): - """ - Calculate the transition amplitude matrix - - Parameters - ---------- - baseline : float - Baseline traversed, km - - energy : float - Neutrino energy, GeV - - dm_mat : complex 2d array - - dm_mat_mat : complex 2d array - - H_full_mass_eigenstate_basis : complex 2d array - - transition_matrix : complex 2d array (empty) - Transition matrix in mass eigenstate basis - - Notes - ----- - - corrsponds to matrix A (equation 10) in original Barger paper - - take into account generic potential matrix (=Hamiltonian) - - """ - - product = cuda.local.array(shape=(3, 3, 3), dtype=ctype) - - clear_matrix(transition_matrix) - - get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product) - - # (1/2)*(1/(h_bar*c)) in units of GeV/(eV^2 km) - hbar_c_factor = 2.534 - - for k in range(3): - arg = -dm_mat[k, 0] * (baseline / energy) * hbar_c_factor - c = cmath.exp(arg * 1.0j) - for i in range(3): - for j in range(3): - transition_matrix[i, j] += c * product[i, j, k] - - -@myjit -def get_H_vac(mix_nubar, mix_nubar_conj_transp, dm_vac_vac, H_vac): - """ Calculate vacuum Hamiltonian in flavor basis for neutrino or antineutrino - - Parameters: - ----------- - mix_nubar : complex 2d array - Mixing matrix, already conjugated if antineutrino - - mix_nubar_conj_transp : conjugate 2d array - Conjugate transpose of mix_nubar - - dm_vac_vac : 2d array - Matrix of mass splittings - - H_vac : complex 2d array - Hamiltonian in vacuum, without the 1/2E term - - - Notes - ------ - The Hailtonian does not contain the energy dependent factor of - 1/(2 * E), as it will be added later - - """ - - dm_vac_diag = cuda.local.array(shape=(3, 3), dtype=ctype) - tmp = cuda.local.array(shape=(3, 3), dtype=ctype) - - clear_matrix(dm_vac_diag) - - dm_vac_diag[1, 1] = dm_vac_vac[1, 0] + 0j - dm_vac_diag[2, 2] = dm_vac_vac[2, 0] + 0j - - matrix_dot_matrix(dm_vac_diag, mix_nubar_conj_transp, tmp) - matrix_dot_matrix(mix_nubar, tmp, H_vac) - -@myjit -def get_H_decay(mix_nubar, mix_nubar_conj_transp, mat_decay, H_decay): - """ Calculate decay Hamiltonian in flavor basis for neutrino or antineutrino - - Parameters: - ----------- - mix_nubar : complex 2d array - Mixing matrix, already conjugated if antineutrino - - mix_nubar_conj_transp : conjugate 2d array - Conjugate transpose of mix_nubar - - H_decay : complex 2d array - Hamiltonian for decay, without the 1/2E term - - - Notes - ------ - The Hailtonian does not contain the energy dependent factor of - 1/(2 * E), as it will be added later - - """ - - #dm_decay_diag = cuda.local.array(shape=(3, 3), dtype=ctype) - tmp = cuda.local.array(shape=(3, 3), dtype=ctype) - - clear_matrix(H_decay) - - #dm_decay_diag[2, 2] = 0. - 1.0j*alpha3 - - #matrix_dot_matrix(dm_decay_diag, mix_nubar_conj_transp, tmp) - matrix_dot_matrix(mat_decay, mix_nubar_conj_transp, tmp) - matrix_dot_matrix(mix_nubar, tmp, H_decay) - -@myjit -def get_H_mat(rho, mat_pot, nubar, H_mat): - """ Calculate matter Hamiltonian in flavor basis - - Parameters: - ----------- - rho : real float - Electron number density (in moles/cm^3) (numerically, this is just the - product of electron fraction and mass density in g/cm^3, since the - number of grams per cm^3 corresponds to the number of moles of nucleons - per cm^3) - - mat_pot : complex 2d array - Generalised matter potential matrix without "a" factor (will be - multiplied with "a" factor); set to diag([1, 0, 0]) for only standard - oscillations - - nubar : int - +1 for neutrinos, -1 for antineutrinos - - H_mat : complex 2d array (empty) - matter hamiltonian - - Notes - ----- - In the following, `a` is just the standard effective matter potential - induced by charged-current weak interactions with electrons - - """ - - # 2*sqrt(2)*Gfermi in (eV^2 cm^3)/(mole GeV) - tworttwoGf = 1.52588e-4 - a = 0.5 * rho * tworttwoGf - - # standard matter interaction Hamiltonian - clear_matrix(H_mat) - - # formalism of Hamiltonian: not 1+epsilon_ee^f in [0,0] element but just epsilon... - # changed when fitting in Thomas' NSI branch -EL - - # Obtain effective non-standard matter interaction Hamiltonian - # changed when fitting in Thomas' NSI branch -EL - for i in range(3): - for j in range(3): - # matter potential V -> -V* for anti-neutrinos - if nubar == -1: - H_mat[i, j] = -a * mat_pot[i, j].conjugate() - elif nubar == 1: - H_mat[i, j] = a * mat_pot[i, j] - -@myjit -def get_dms_numerical(energy, H_full, dm_mat_mat, dm_mat): - """Compute the matter-mass vector M, dM = M_i-M_j and dMimj - - Parameters - ---------- - energy : float - Neutrino energy, GeV - - H_full : complex 2d array - full hamiltonian - - dm_mat_mat : complex 2d array (empty) - - dm_mat : complex 2d array (empty) - - - Notes - ----- - Calculate mass eigenstates in matter - neutrino or anti-neutrino (type already taken into account in Hamiltonian) - of energy energy. - """ - - - m_mat = 2.0 * energy * np.linalg.eigvals(H_full) - - for i in range(3): - for j in range(3): - dm_mat_mat[i, j] = m_mat[i] - m_mat[j] - dm_mat[i, j] = m_mat[i] - -@myjit -def get_dms(energy, H_full, dm_vac_vac, dm_mat_mat, dm_mat): - """Compute the matter-mass vector M, dM = M_i-M_j and dMimj - - Parameters - ---------- - energy : float - Neutrino energy, GeV - - H_full : complex 2d array - full hamiltonian - - dm_vac_vac : 2d array - - dm_mat_mat : complex 2d array (empty) - - dm_mat : complex 2d array (empty) - - - Notes - ----- - Calculate mass eigenstates in matter - neutrino or anti-neutrino (type already taken into account in Hamiltonian) - of energy energy. - - - only god knows what happens in this function, somehow it seems to work - - """ - - # the following is for solving the characteristic polynomial of H_full: - # P(x) = x**3 + c2*x**2 + c1*x + c0 - real_product_a = (H_full[0, 1] * H_full[1, 2] * H_full[2, 0]).real - real_product_b = (H_full[0, 0] * H_full[1, 1] * H_full[2, 2]).real - - norm_H_e_mu_sq = H_full[0, 1].real ** 2 + H_full[0, 1].imag ** 2 - norm_H_e_tau_sq = H_full[0, 2].real ** 2 + H_full[0, 2].imag ** 2 - norm_H_mu_tau_sq = H_full[1, 2].real ** 2 + H_full[1, 2].imag ** 2 - - # c1 = H_{11} * H_{22} + H_{11} * H_{33} + H_{22} * H_{33} - # - |H_{12}|**2 - |H_{13}|**2 - |H_{23}|**2 - # given Hermiticity of Hamiltonian (real diagonal elements), - # this coefficient must be real - c1 = ( - (H_full[0, 0].real * (H_full[1, 1] + H_full[2, 2])).real - - (H_full[0, 0].imag * (H_full[1, 1] + H_full[2, 2])).imag - + (H_full[1, 1].real * H_full[2, 2]).real - - (H_full[1, 1].imag * H_full[2, 2]).imag - - norm_H_e_mu_sq - - norm_H_mu_tau_sq - - norm_H_e_tau_sq - ) - - # c0 = H_{11} * |H_{23}|**2 + H_{22} * |H_{13}|**2 + H_{33} * |H_{12}|**2 - # - H_{11} * H_{22} * H_{33} - 2*Re(H*_{13} * H_{12} * H_{23}) - # hence, this coefficient is also real - c0 = ( - H_full[0, 0].real * norm_H_mu_tau_sq - + H_full[1, 1].real * norm_H_e_tau_sq - + H_full[2, 2].real * norm_H_e_mu_sq - - 2.0 * real_product_a - - real_product_b - ) - - # c2 = -H_{11} - H_{22} - H_{33} - # hence, this coefficient is also real - c2 = -H_full[0, 0].real - H_full[1, 1].real - H_full[2, 2].real - - one_over_two_e = 0.5 / energy - one_third = 1.0 / 3.0 - two_third = 2.0 / 3.0 - - # we also have to perform the corresponding algebra - # for the vacuum case, where the relevant elements of the - # hamiltonian are mass differences - x = dm_vac_vac[1, 0] - y = dm_vac_vac[2, 0] - - c2_v = -one_over_two_e * (x + y) - - # p is real due to reality of c1 and c2 - p = c2 ** 2 - 3.0 * c1 - p_v = one_over_two_e ** 2 * (x ** 2 + y ** 2 - x * y) - p = max(0.0, p) - - # q is real - q = -13.5 * c0 - c2 ** 3 + 4.5 * c1 * c2 - q_v = one_over_two_e ** 3 * (x + y) * ((x + y) ** 2 - 4.5 * x * y) - - # we need the quantity p**3 - q**2 to obtain the eigenvalues, - # but let's prevent inaccuracies and instead write - tmp = 27 * (0.25 * c1 ** 2 * (p - c1) + c0 * (q + 6.75 * c0)) - # changed from p**3 - q**2 when fitting in Thomas' NSI branch -EL - # TODO: can we simplify this quantity to reduce numerical inaccuracies? - tmp_v = p_v ** 3 - q_v ** 2 - - tmp = max(0.0, tmp) - - theta = cuda.local.array(shape=(3), dtype=ftype) - theta_v = cuda.local.array(shape=(3), dtype=ftype) - m_mat = cuda.local.array(shape=(3), dtype=ftype) - m_mat_u = cuda.local.array(shape=(3), dtype=ftype) - m_mat_v = cuda.local.array(shape=(3), dtype=ftype) - - a = two_third * math.pi - # intermediate result, needed to calculate the three - # mass eigenvalues (theta0, theta1, theta2 are the three - # corresponding arguments of the cosine, see m_mat_u) - res = math.atan2(math.sqrt(tmp), q) * one_third - theta[0] = res + a - theta[1] = res - a - theta[2] = res - res_v = math.atan2(math.sqrt(tmp_v), q_v) * one_third - theta_v[0] = res_v + a - theta_v[1] = res_v - a - theta_v[2] = res_v - - b = two_third * math.sqrt(p) - b_v = two_third * math.sqrt(p_v) - - for i in range(3): - m_mat_u[i] = ( - 2.0 * energy * (b * math.cos(theta[i]) - c2 * one_third + dm_vac_vac[0, 0]) - ) - m_mat_v[i] = ( - 2.0 - * energy - * (b_v * math.cos(theta_v[i]) - c2_v * one_third + dm_vac_vac[0, 0]) - ) - - # Sort according to which reproduce the vaccum eigenstates - for i in range(3): - tmp_v = abs(dm_vac_vac[i, 0] - m_mat_v[0]) - k = 0 - for j in range(3): - tmp = abs(dm_vac_vac[i, 0] - m_mat_v[j]) - if tmp < tmp_v: - k = j - tmp_v = tmp - m_mat[i] = m_mat_u[k] - - for i in range(3): - for j in range(3): - dm_mat_mat[i, j] = m_mat[i] - m_mat[j] - dm_mat[i, j] = m_mat[i] - #dm_mat_vac[i, j] = m_mat[i] - dm_vac_vac[j, 0] - - -@myjit -def get_product(energy, dm_mat, dm_mat_mat, H_full_mass_eigenstate_basis, product): - """ - Parameters - ---------- - energy : float - Neutrino energy, GeV - - dm_mat : complex 2d array - - dm_mat_mat : complex 2d array - - H_full_mass_eigenstate_basis : complex 2d array - - product : complex 3d-array (empty) - - """ - - H_minus_M = cuda.local.array(shape=(3, 3, 3), dtype=ctype) - - for i in range(3): - for j in range(3): - for k in range(3): - H_minus_M[i, j, k] = 2.0 * energy * H_full_mass_eigenstate_basis[i, j] - if i == j: - H_minus_M[i, j, k] -= dm_mat[k, j] - # also, cler product - product[i, j, k] = 0.0 - - # Calculate the product in eq.(10) of H_minus_M for j!=k - for i in range(3): - for j in range(3): - for k in range(3): - product[i, j, 0] += H_minus_M[i, k, 1] * H_minus_M[k, j, 2] - product[i, j, 1] += H_minus_M[i, k, 2] * H_minus_M[k, j, 0] - product[i, j, 2] += H_minus_M[i, k, 0] * H_minus_M[k, j, 1] - product[i, j, 0] /= dm_mat_mat[0, 1] * dm_mat_mat[0, 2] - product[i, j, 1] /= dm_mat_mat[1, 2] * dm_mat_mat[1, 0] - product[i, j, 2] /= dm_mat_mat[2, 0] * dm_mat_mat[2, 1] - - -@myjit -def convert_from_mass_eigenstate(state, mix_nubar, psi): - """ - Parameters - ---------- - state : (un?)signed int - - mix_nubar : complex 2d array - Mixing matrix, already conjugated if antineutrino - - psi : complex 1d-array (empty) - - - Notes - ----- - this is untested! - - """ - - mass = cuda.local.array(shape=(3), dtype=ctype) - - lstate = state - 1 - for i in range(3): - mass[i] = 1.0 if lstate == i else 0.0 - - # note: mix_nubar is already taking into account whether we're considering - # nu or anti-nu - matrix_dot_vector(mix_nubar, mass, psi) diff --git a/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py b/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py deleted file mode 100755 index 994d34947..000000000 --- a/pisa/stages/osc/prob3numba/numba_osc_tests_bk.py +++ /dev/null @@ -1,791 +0,0 @@ -#!/usr/bin/env python -# pylint: disable = invalid-name - - -""" -Tests for prob3numba code -""" - - -from __future__ import absolute_import, print_function, division - - -__all__ = [ - "TEST_DATA_DIR", - "FINFO_FTYPE", - "AC_KW", - "PRINTOPTS", - "A2S_KW", - "MAT_DOT_MAT_SUBSCR", - "DEFAULTS", - "TEST_CASES", - "auto_populate_test_case", - "test_prob3numba", - "run_test_case", - "stability_test", - "execute_func", - "compare_numeric", - "check", - "ary2str", - "main", -] - - -from argparse import ArgumentParser -from collections import OrderedDict -from collections.abc import Mapping -from copy import deepcopy -from inspect import getmodule, signature -from os.path import join - -import numpy as np -# import os -# os.environ['PISA_FTYPE'] = 'single' # for checking unit test on single precision -from pisa import FTYPE -from pisa.utils.comparisons import ALLCLOSE_KW -from pisa.utils.fileio import expand, from_file, to_file -from pisa.utils.log import Levels, logging, set_verbosity -from pisa.utils.numba_tools import WHERE -from pisa.utils.resources import find_resource -from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import ( - CX, - FX, - IX, - # propagate_scalar_vacuum, - propagate_scalar, - propagate_array, - get_transition_matrix_hostfunc, - get_transition_matrix_massbasis_hostfunc, - get_H_vac_hostfunc, - get_H_decay_hostfunc, - get_H_mat_hostfunc, - get_dms_hostfunc, - get_dms_numerical_hostfunc, - get_product_hostfunc, - convert_from_mass_eigenstate_hostfunc, -) -from pisa.stages.osc.nsi_params import ( - StdNSIParams, - VacuumLikeNSIParams, -) - -TEST_DATA_DIR = find_resource("osc/numba_osc_tests_data") - -FINFO_FTYPE = np.finfo(FTYPE) - -AC_KW = dict(atol=FINFO_FTYPE.resolution * 10, rtol=ALLCLOSE_KW["rtol"] * 100) - -PRINTOPTS = dict( - precision=FINFO_FTYPE.precision + 2, floatmode="fixed", sign=" ", linewidth=200 -) - -A2S_KW = dict(precision=PRINTOPTS["precision"], separator=", ") - -MAT_DOT_MAT_SUBSCR = "in,nj->ij" -"""matrix dot matrix subscripts for use by `numpy.einsum`""" - -# ---------------------------------------------------------------------------- # -# Define relevant values for testing purposes (from nufit3.2, from intermediate -# calculations performed here, or arbitary values). -# -# NOTE: !!DO NOT CHANGE!! (unless a function is incorrect) tests rely on these -# ---------------------------------------------------------------------------- # - -DEFAULTS = dict( - energy=1, # GeV - state=1, - nubar=1, - rho=1, # moles of electrons / cm^3 - baseline=1, # km - mat_pot=np.diag([1, 0, 0]).astype(np.complex128), - layer_distances=np.logspace(0, 2, 10), # km - layer_densities=np.linspace(0.5, 3, 10), # g/cm^3 - # osc params: defaults are nufit 3.2 normal ordering values - t12=np.deg2rad(33.62), - t23=np.deg2rad(47.2), - t13=np.deg2rad(8.54), - dcp=np.deg2rad(234), - dm21=7.40e-5, - dm31=2.494e-3, - decay_flag=-1, - mat_decay=np.diag([0, 0, -1.0e-4j ]).astype(np.complex128), -) - -# define non-0 NSI parameters for non-vacuum NSI -# roughly based on best fit params from Thomas Ehrhardts 3y DRAGON analysis -nsi_params = StdNSIParams() -nsi_params.eps_emu_magn = 0.07 -nsi_params.eps_emu_phase = np.deg2rad(340) -nsi_params.eps_etau_magn = 0.06 -nsi_params.eps_etau_phase = np.deg2rad(35) -nsi_params.eps_mutau_magn = 0.003 -nsi_params.eps_mutau_phase = np.deg2rad(175) -mat_pot_std_nsi_no = np.diag([1, 0, 0]).astype(np.complex128) + nsi_params.eps_matrix - -# Vacuum-like NSI parameters -nsi_params = VacuumLikeNSIParams() -nsi_params.eps_prime = 0.1 -mat_pot_vac_nsi_no = np.diag([1, 0, 0]).astype(np.complex128) + nsi_params.eps_matrix - -TEST_CASES = dict( - nufit32_no=dict(), # nufit 3.2 normal ordering (also overall) best-fit - nufit32_no_nubar=dict(nubar=-1), # NO but anti-neutrinos - nufit32_no_E1TeV=dict(energy=1e3), # NO but e=1 TeV - nufit32_no_blearth=dict( - baseline=6371e3 * 2, - layer_distances=( - 6371e3 - * 2 - * DEFAULTS["layer_distances"] - / np.sum(DEFAULTS["layer_distances"]) - ), - ), - nufit32_io=dict( # nufit 3.2 best-fit params for inverted ordering - t12=np.deg2rad(33.62), - t23=np.deg2rad(48.1), - t13=np.deg2rad(8.58), - dcp=np.deg2rad(278), - dm21=7.40e-5, - dm31=-2.465e-3, - ), - nufit32_std_nsi_no=dict( # nufit 3.2 normal ordering with non-0 standard NSI parameters - mat_pot=mat_pot_std_nsi_no, - ), - nufit32_vac_nsi_no=dict( # nufit 3.2 normal ordering with non-0 vacuum NSI parameters - mat_pot=mat_pot_vac_nsi_no, - ), - nufit32_std_decay_no=dict( # nufit 3.2 normal ordering with no neutrino decay - decay_flag=-1, - mat_decay=DEFAULTS["mat_decay"], - ), - nufit32_std_decay=dict( # nufit 3.2 normal ordering with neutrino decay - decay_flag=1, - mat_decay=DEFAULTS["mat_decay"], - ), -) - - -def auto_populate_test_case(tc, defaults): - """Populate defaults and construct dm / PMNS matrices if they aren't - present in a test case. - - Parameters - ---------- - test_case : mutable mapping - - defaults : mapping - - """ - for key, val in defaults.items(): - if key not in tc: - tc[key] = val - - # Construct dm and PMNS matrices derived from test case values, if - # these were not already specified - - if "dm" not in tc: # construct Delta m^2 matrix if not present - if "dm32" not in tc: # NO case; Delta m^2_32/eV^2 - tc["dm32"] = tc["dm31"] - tc["dm21"] - - if "dm31" not in tc: # IO case; Delta m^2_32/eV^2 - tc["dm31"] = tc["dm32"] + tc["dm21"] - - tc["dm"] = np.array( - [ - [0, -tc["dm21"], -tc["dm31"]], - [tc["dm21"], 0, -tc["dm32"]], - [tc["dm31"], tc["dm32"], 0], - ] - ) - - if "pmns" not in tc: # construct PMNS matrix if not present - c12, s12 = np.cos(tc["t12"]), np.sin(tc["t12"]) - c23, s23 = np.cos(tc["t23"]), np.sin(tc["t23"]) - c13, s13 = np.cos(tc["t13"]), np.sin(tc["t13"]) - - tc["pmns"] = ( - np.array([[1, 0, 0], [0, c23, s23], [0, -s23, c23]]) - @ np.array( - [ - [c13, 0, s13 * np.exp(-1j * tc["dcp"])], - [0, 1, 0], - [-s13 * np.exp(1j * tc["dcp"]), 0, c13], - ] - ) - @ np.array([[c12, s12, 0], [-s12, c12, 0], [0, 0, 1]]) - ) - - -def test_prob3numba(ignore_fails=False, define_as_ref=False): - """Run all unit test cases for prob3numba code""" - - # Pull first test case to test calling `propagate_array` - tc_name, tc = next(iter(TEST_CASES.items())) - tc_ = deepcopy(tc) - logging.info( - "Testing call and return shape of `propagate_array` with test case '%s'", - tc_name, - ) - - # Test simple broadcasting over `nubars` and `energies` where both have - # same shape, as this is the "typical" use case - input_shape = (4, 5) - - # Without broadcasting, a single probability matrix is 3x3 - prob_array_shape = (3, 3) - - # Broadcasted shape - out_shape = input_shape + prob_array_shape - - nubars = np.full(shape=input_shape, fill_value=tc_["nubar"], dtype=IX) - energies = np.full(shape=input_shape, fill_value=tc_["energy"], dtype=FX) - - # Fill with NaN to ensure all elements are assinged a value - probabilities = np.full(shape=out_shape, fill_value=np.nan, dtype=FX) - - propagate_array( - tc_["dm"].astype(FX), - tc_["pmns"].astype(CX), - tc_["mat_pot"].astype(CX), - tc_["decay_flag"], - tc_["mat_decay"].astype(CX), - nubars, - energies, - tc_["layer_densities"].astype(FX), - tc_["layer_distances"].astype(FX), - # output: - probabilities, - ) - - # Check that all probability matrices have no NaNs and are equal to one - # another - ref_probs = probabilities[0, 0] - for i in range(input_shape[0]): - for j in range(input_shape[1]): - probs = probabilities[i, j] - assert np.all(np.isfinite(probs)) - assert np.all(probs == ref_probs) - - # Run all test cases - for tc_name, tc in TEST_CASES.items(): - run_test_case( - tc_name, tc, ignore_fails=ignore_fails, define_as_ref=define_as_ref - ) - - -def run_test_case(tc_name, tc, ignore_fails=False, define_as_ref=False): - """Run one test case""" - logging.info("== TEST CASE : %s ==", tc_name) - - st_test_kw = dict(ignore_fails=ignore_fails, define_as_ref=define_as_ref) - tf_sfx = f"__{tc_name}__{FX}.pkl" - - # Copy contents of test case, so if a function modifies these - # (accidentally), it doesn't affect the outcome of further tests - tc_ = deepcopy(tc) - test, ref = stability_test( - func=convert_from_mass_eigenstate_hostfunc, - func_kw=dict( - state=tc_["state"], - mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, - # output: - psi=np.ones(shape=3, dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"convert_from_mass_eigenstate_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\npsi = %s", ary2str(test["psi"])) - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=get_H_vac_hostfunc, - func_kw=dict( - mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, - mix_nubar_conj_transp=( - tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] - ), - dm_vac_vac=tc_["dm"], - # output: - H_vac=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_H_vac_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\nH_vac = %s", ary2str(test["H_vac"])) - # keep for use by `get_transition_matrix_hostfunc` - H_vac_ref = ref["H_vac"] - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=get_H_decay_hostfunc, - func_kw=dict( - mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, - mix_nubar_conj_transp=( - tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] - ), - mat_decay=tc_["mat_decay"], - #output; - H_decay=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_H_decay_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\nH_decay = %s", ary2str(test["H_decay"])) - # keep for use by `get_transition_matrix_hostfunc` - H_decay_ref = ref["H_decay"] - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=get_H_mat_hostfunc, - func_kw=dict( - rho=tc_["rho"], - mat_pot=tc_["mat_pot"], - nubar=tc_["nubar"], - # output: - H_mat=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_H_mat_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\nH_mat = %s", ary2str(test["H_mat"])) - # keep for use by `get_dms_hostfunc`, - # `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc`` - H_mat_ref = ref["H_mat"] - - # tc_ = deepcopy(tc) - # test, ref = stability_test( - # func=propagate_scalar_vacuum, - # func_kw=dict( - # dm=tc_["dm"], - # mix=tc_["pmns"], - # nubar=tc_["nubar"], - # energy=tc_["energy"], - # distances=tc_["layer_distances"], - # # output: - # probability=np.ones(shape=(3, 3), dtype=FX), - # ), - # ref_path=join(TEST_DATA_DIR, f"propagate_scalar_vacuum{tf_sfx}"), - # **st_test_kw, - # ) - # logging.debug("\nvac_prob = %s", ary2str(test["probability"])) - # # check unitarity - # # TODO: << BUG? >> these fail even in double precision! - # check( - # test=np.sum(test["probability"], axis=0), - # ref=np.ones(3), - # label=( - # f"{tc_name} :: propagate_scalar_vacuum :: sum(vacuum probability, axis=0)" - # ), - # ignore_fails=True, - # ) - # check( - # test=np.sum(test["probability"], axis=1), - # ref=np.ones(3), - # label=( - # f"{tc_name} :: propagate_scalar_vacuum :: sum(vacuum probability, axis=1)" - # ), - # ignore_fails=True, - # ) - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=propagate_scalar, - func_kw=dict( - dm=tc_["dm"], - mix=tc_["pmns"], - mat_pot=tc_["mat_pot"], - decay_flag=tc_["decay_flag"], - mat_decay=tc_["mat_decay"], - nubar=tc_["nubar"], - energy=tc_["energy"], - densities=tc_["layer_densities"], - distances=tc_["layer_distances"], - # output: - probability=np.ones(shape=(3, 3), dtype=FX), - ), - ref_path=join(TEST_DATA_DIR, f"propagate_scalar{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\nmat_prob = %s", ary2str(test["probability"])) - - if (tc_["decay_flag"] != 1): - # check unitarity - check( - test=np.sum(test["probability"], axis=0), - ref=np.ones(3), - label=f"{tc_name} :: propagate_scalar:: sum(matter probability, axis=0)", - ignore_fails=ignore_fails, - ) - check( - test=np.sum(test["probability"], axis=1), - ref=np.ones(3), - label=f"{tc_name} :: propagate_scalar :: sum(matter probability, axis=1)", - ignore_fails=ignore_fails, - ) - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=get_transition_matrix_hostfunc, - func_kw=dict( - nubar=tc_["nubar"], - energy=tc_["energy"], - rho=tc_["rho"], - baseline=tc_["baseline"], - mix_nubar=tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T, - mix_nubar_conj_transp=( - tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] - ), - mat_pot=tc_["mat_pot"], - H_vac=H_vac_ref, - decay_flag=tc_["decay_flag"], - H_decay=H_decay_ref, - dm=tc_["dm"], - # output: - transition_matrix=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_transition_matrix_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\ntransition_matrix = %s", ary2str(test["transition_matrix"])) - # check unitarity - check( - test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=0), - ref=np.ones(3), - label=( - f"{tc_name}" - " :: get_transition_matrix_hostfunc" - ":: sum(|transition_matrix|^2, axis=0)" - ), - ignore_fails=ignore_fails, - ) - check( - test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=1), - ref=np.ones(3), - label=( - f"{tc_name}" - " :: get_transition_matrix_hostfunc" - " :: sum(|transition_matrix|^2, axis=1)" - ), - ignore_fails=ignore_fails, - ) - - tc_ = deepcopy(tc) - - # Compute H_full as used in `numba_osc_kernels` to call `get_dms` for SI case - # and det_dms_numnerical for decay case from `get_transition_matrix` - - if (tc_["decay_flag"] == 1): - - H_full_ref = (H_vac_ref + H_decay_ref)/ (2 * tc_["energy"]) + H_mat_ref - - test, ref = stability_test( - func=get_dms_numerical_hostfunc, - func_kw=dict( - energy=tc_["energy"], - H_full=H_full_ref, - # outputs: - dm_mat_mat=np.ones(shape=(3, 3), dtype=CX), - dm_mat=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_dms_numerical_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\ndm_mat_mat = %s", ary2str(test["dm_mat_mat"])) - logging.debug("\ndm_mat = %s", ary2str(test["dm_mat"])) - # keep for use by `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc` - dm_mat_mat_ref = ref["dm_mat_mat"] - dm_mat_ref = ref["dm_mat"] - - else: - H_full_ref = H_vac_ref / (2 * tc_["energy"]) + H_mat_ref - - test, ref = stability_test( - func=get_dms_hostfunc, - func_kw=dict( - energy=tc_["energy"], - H_full=H_full_ref, - dm_vac_vac=tc_["dm"], - # outputs: - dm_mat_mat=np.ones(shape=(3, 3), dtype=CX), - dm_mat=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"get_dms_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\ndm_mat_mat = %s", ary2str(test["dm_mat_mat"])) - logging.debug("\ndm_mat = %s", ary2str(test["dm_mat"])) - # keep for use by `get_transition_matrix_massbasis_hostfunc`, `get_product_hostfunc` - dm_mat_mat_ref = ref["dm_mat_mat"] - dm_mat_ref = ref["dm_mat"] - - - tc_ = deepcopy(tc) - - # Compute same intermediate result `H_full_mass_eigenstate_basis` as in - # `numba_osc_kernels.get_transition_matrix` which calls - # `get_transition_matrix_massbasis` - mix_nubar = tc_["pmns"] if tc_["nubar"] > 0 else tc_["pmns"].conj().T - mix_nubar_conj_transp = tc_["pmns"].conj().T if tc_["nubar"] > 0 else tc_["pmns"] - tmp = np.einsum(MAT_DOT_MAT_SUBSCR, H_full_ref, mix_nubar) - H_full_mass_eigenstate_basis = np.einsum( - MAT_DOT_MAT_SUBSCR, mix_nubar_conj_transp, tmp - ) - - test, ref = stability_test( - func=get_transition_matrix_massbasis_hostfunc, - func_kw=dict( - baseline=tc_["baseline"], - energy=tc_["energy"], - dm_mat=dm_mat_ref, - dm_mat_mat=dm_mat_mat_ref, - H_full_mass_eigenstate_basis=H_full_mass_eigenstate_basis, - # output: - transition_matrix=np.ones(shape=(3, 3), dtype=CX), - ), - ref_path=join( - TEST_DATA_DIR, f"get_transition_matrix_massbasis_hostfunc{tf_sfx}" - ), - **st_test_kw, - ) - logging.debug("\ntransition_matrix_mb = %s", ary2str(test["transition_matrix"])) - check( - test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=0), - ref=np.ones(3), - label=( - f"{tc_name}" - " :: get_transition_matrix_massbasis_hostfunc" - " :: sum(|transition_matrix (mass basis)|^2), axis=0)" - ), - ignore_fails=ignore_fails, - ) - check( - test=np.sum(np.abs(test["transition_matrix"]) ** 2, axis=1), - ref=np.ones(3), - label=( - f"{tc_name}" - " :: get_transition_matrix_massbasis_hostfunc" - " :: sum(|transition_matrix (mass basis)|^2), axis=1)" - ), - ignore_fails=ignore_fails, - ) - - tc_ = deepcopy(tc) - test, ref = stability_test( - func=get_product_hostfunc, - func_kw=dict( - energy=tc_["energy"], - dm_mat=dm_mat_ref, - dm_mat_mat=dm_mat_mat_ref, - H_full_mass_eigenstate_basis=H_full_ref, - # output: - product=np.ones(shape=(3, 3, 3), dtype=CX), - ), - ref_path=join(TEST_DATA_DIR, f"product_hostfunc{tf_sfx}"), - **st_test_kw, - ) - logging.debug("\nproduct = %s", ary2str(test["product"])) - - -def stability_test(func, func_kw, ref_path, ignore_fails=False, define_as_ref=False): - """basic stability test of a Numba CPUDispatcher function (i.e., function - compiled via @jit / @njit)""" - func_name = func.py_func.__name__ - logging.info("stability testing `%s`", func_name) - ref_path = expand(ref_path) - - test = execute_func(func=func, func_kw=func_kw) - - if define_as_ref: - to_file(test, ref_path) - - # Even when we define the test case as ref, round-trip to/from file to - # ensure that doesn't corrupt the values - ref = from_file(ref_path) - - check(test=test, ref=ref, label=func_name, ignore_fails=ignore_fails) - - return test, ref - - -def execute_func(func, func_kw): - """Run `func` with *func_kw.values() where `outputs` specify names in - `func_kw` taken to be outputs of the function; for these, mark changed. - Retrieve both input and output values as Numpy arrays on the host and - aggregate together in a single dict before returning. - - Parameters - ---------- - func : numba CPUDispatcher or CUDADispatcher - func_kw : OrderedDict - - Returns - ------- - ret_dict : OrderedDict - Keys are arg names and vals are type-"correct" values; all arrays are - converted to host Numpy arrays - - """ - py_func = func.py_func - func_name = ".".join([getmodule(py_func).__name__, py_func.__name__]) - arg_names = list(signature(py_func).parameters.keys()) - if hasattr(func, "signatures"): - arg_types = func.signatures[0] - else: - arg_types = func.compiled.argument_types - - missing = set(arg_names).difference(func_kw.keys()) - excess = set(func_kw.keys()).difference(arg_names) - if missing or excess: - msgs = [] - if missing: - msgs.append(f"missing kwargs {missing}") - if excess: - msgs.append(f"excess kwargs {excess}") - raise KeyError(f"{func_name}:" + ", ".join(msgs)) - - typed_args = OrderedDict() - for arg_name, arg_type in zip(arg_names, arg_types): - val = func_kw[arg_name] - if arg_type.name.startswith("array"): - arg_val = val.astype(arg_type.dtype.key) - else: - arg_val = arg_type(val) - typed_args[arg_name] = arg_val - - # Call the host function with typed args - - try: - func(*list(typed_args.values())) - except Exception: - logging.error("Failed running `%s` with args %s", func_name, str(typed_args)) - raise - - # All arrays converted to Numpy host arrays - - ret_dict = OrderedDict() - for key, val in typed_args.items(): - ret_dict[key] = val - - return ret_dict - - -def compare_numeric(test, ref, label=None, ac_kw=deepcopy(AC_KW), ignore_fails=False): - """Compare scalars or numpy ndarrays. - - Parameters - ---------- - test : scalar or numpy.ndarray - ref : scalar or numpy.ndarray - label : str or None, optional - ac_kw : mapping, optional - Keyword args to pass via **ac_kw to `numpy.isclose` / `numpy.allclose` - ignore_fails : bool, optional - - Returns - ------- - rslt : bool - - """ - pfx = f"{label} :: " if label else "" - with np.printoptions(**PRINTOPTS): - if np.isscalar(test): - if np.isclose(test, ref, **ac_kw): - return True - - msg = f"{pfx}test: {test} != ref: {ref}" - if ignore_fails: - logging.warning(msg) - else: - logging.error(msg) - return False - - # Arrays - if np.allclose(test, ref, **ac_kw): - return True - - diff = test - ref - msg = f"{pfx}test:" f"\n{(test)}\n!= ref:\n{(ref)}" f"\ndiff:\n{(diff)}" - - if not np.all(ref == 1): - nzmask = ref != 0 - zmask = ref == 0 - fdiff = np.empty_like(ref) - fdiff[nzmask] = diff[nzmask] / ref[nzmask] - fdiff[zmask] = np.nan - msg += f"\nfractdiff:\n{(fdiff)}" - - if ignore_fails: - logging.warning(msg) - else: - logging.error(msg) - - return False - - -def check(test, ref, label=None, ac_kw=deepcopy(AC_KW), ignore_fails=False): - """Check that `test` matches `ref` (closely enough). - - Parameters - ---------- - test - ref - ac_kw : mapping, optional - Kwargs to `np.allclose`, as used by - `pisa.utils.comparisons.recursiveEquality` - ignore_fails : bool, optional - If True and comparison fails, do not raise AssertionError - - Raises - ------ - AssertionError - If `test` is not close enough to `ref` and ignore_fails is False - - """ - same = True - with np.printoptions(**PRINTOPTS): - if isinstance(test, Mapping): - if not label: - label = "" - else: - label = label + ": " - - for key, val in test.items(): - same &= compare_numeric( - test=val, - ref=ref[key], - label=label + f"key: '{key}'", - ac_kw=ac_kw, - ignore_fails=ignore_fails, - ) - else: - same &= compare_numeric( - test=test, ref=ref, label=label, ac_kw=ac_kw, ignore_fails=ignore_fails - ) - if not ignore_fails and not same: - assert False - return same - - -def ary2str(array): - """Convert a numpy ndarray to string easy to copy back into code""" - return "np.array(" + np.array2string(array, **A2S_KW) + ")" - - -# Augment test cases with defaults / contruct arrays where necessary -for TC in TEST_CASES.values(): - auto_populate_test_case(tc=TC, defaults=DEFAULTS) - - -def main(description=__doc__): - """Script interface for `test_prob3numba` function""" - parser = ArgumentParser(description=description) - parser.add_argument("--ignore-fails", action="store_true") - parser.add_argument("--define-as-ref", action="store_true") - parser.add_argument("-v", action="count", default=Levels.WARN) - kwargs = vars(parser.parse_args()) - set_verbosity(kwargs.pop("v")) - test_prob3numba(**kwargs) - - -if __name__ == "__main__": - main()