pisa.stages.osc.prob3numba package

Submodules

pisa.stages.osc.prob3numba.numba_osc_hostfuncs module

Host function wrappers for numba_osc_kernels.

pisa.stages.osc.prob3numba.numba_osc_hostfuncs.CX = 'c16'

Complex string code to use, understood by both Numba and Numpy

pisa.stages.osc.prob3numba.numba_osc_hostfuncs.FX = 'f8'

Float string code to use, understood by both Numba and Numpy

pisa.stages.osc.prob3numba.numba_osc_hostfuncs.IX = 'i8'

Signed integer string code to use, understood by both Numba and Numpy

pisa.stages.osc.prob3numba.numba_osc_kernels module

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.

pisa.stages.osc.prob3numba.numba_osc_kernels.get_transition_matrix(nubar, energy, rho, baseline, mix_nubar, mix_nubar_conj_transp, mat_pot, H_vac, decay_flag, H_decay, lri_pot, 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

• 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

• 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

pisa.stages.osc.prob3numba.numba_osc_kernels.osc_probs_layers_kernel(dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, 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

• 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

• 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..

pisa.stages.osc.prob3numba.numba_osc_tests module

Tests for prob3numba code

pisa.stages.osc.prob3numba.numba_osc_tests.MAT_DOT_MAT_SUBSCR = 'in,nj->ij'

matrix dot matrix subscripts for use by numpy.einsum

pisa.stages.osc.prob3numba.numba_osc_tests.ary2str(array)

Convert a numpy ndarray to string easy to copy back into code

pisa.stages.osc.prob3numba.numba_osc_tests.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)

pisa.stages.osc.prob3numba.numba_osc_tests.check(test, ref, label=None, ac_kw={'atol': 1.0000000000000002e-14, 'rtol': 1e-10}, 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

pisa.stages.osc.prob3numba.numba_osc_tests.compare_numeric(test, ref, label=None, ac_kw={'atol': 1.0000000000000002e-14, 'rtol': 1e-10}, 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

Return type:

bool

pisa.stages.osc.prob3numba.numba_osc_tests.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 – Keys are arg names and vals are type-“correct” values; all arrays are converted to host Numpy arrays

Return type:

OrderedDict

pisa.stages.osc.prob3numba.numba_osc_tests.main(description='\nTests for prob3numba code\n')

Script interface for test_prob3numba function

pisa.stages.osc.prob3numba.numba_osc_tests.run_test_case(tc_name, tc, ignore_fails=False, define_as_ref=False)

Run one test case

pisa.stages.osc.prob3numba.numba_osc_tests.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)

pisa.stages.osc.prob3numba.numba_osc_tests.test_prob3numba(ignore_fails=False, define_as_ref=False)

Run all unit test cases for prob3numba code

pisa.stages.osc.prob3numba.test_numba module

Test functions used during developpement of the osc. code Please ignore unless you are the author

pisa.stages.osc.prob3numba.test_numba.main()

Module contents