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)
def convert_from_mass_eigenstate(state, mix_nubar, psi):
'''
Parameters
----------
state : int
mix_nubar : 2d-array
psi : 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. if lstate == i else 0.
# note: mix_nubar is already taking into account whether we're considering
# nu or anti-nu
matrix_dot_vector(mix_nubar, mass, psi)
def sum_row_kernel(mix, bla, inp, out):
C = cuda.local.array(shape=(3, 3), dtype=ftype)
D = cuda.local.array(shape=(3), dtype=ctype)
E = cuda.local.array(shape=(3), dtype=ctype)
matrix_dot_matrix(mix, mix, C)
D[0] = 0.0 + 2.0j
D[1] = 1.0 + 2.0j
D[2] = 1.0 + 2.0j
matrix_dot_vector(C, D, E)
bla *= 0.1
out[0] += E[1].real * bla.real + inp[0]
def osc_probs_layers_kernel(dm, mix, mat_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
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)
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(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)
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,
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,
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