def setup_function(self):
# object for oscillation parameters
self.osc_params = OscParams()
# setup the layers
#if self.params.earth_model.value is not None:
earth_model = find_resource(self.params.earth_model.value)
YeI = self.params.YeI.value.m_as('dimensionless')
YeO = self.params.YeO.value.m_as('dimensionless')
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(YeI, YeO, YeM)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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'].get('host'))
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.calc_mode == 'binned':
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 setup_function(self):
# setup Earth model
if self.params.earth_model.value is not None:
earth_model = find_resource(self.params.earth_model.value)
YeI = self.params.YeI.value.m_as('dimensionless')
YeO = self.params.YeO.value.m_as('dimensionless')
YeM = self.params.YeM.value.m_as('dimensionless')
else:
earth_model = None
# setup the layers
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)
if earth_model is not None:
self.layers.setElecFrac(YeI, YeO, YeM)
# set the correct data mode
self.data.representation = self.calc_mode
# --- calculate the layers ---
if self.data.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:
if self.params.earth_model.value is not None:
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))
else:
self.layers.calcPathLength(container['true_coszen'])
container['distances'] = self.layers.distance
# don't forget to un-link everything again
self.data.unlink_containers()
# --- setup empty arrays ---
if self.data.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 setup_function(self):
import ROOT
# setup the layers
earth_model = find_resource(self.earth_model)
self.layers = Layers(earth_model, self.detector_depth,
self.prop_height)
# This is a bit hacky, but setting the electron density to 1.
# gives us the total density of matter, which is what we want.
self.layers.setElecFrac(1., 1., 1.)
# setup cross-sections
self.xsroot = ROOT.TFile(self.xsec_file)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# layers don't care about flavor
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'].get(WHERE))
container['densities'] = self.layers.density.reshape(
(container.size, self.layers.max_layers))
container['distances'] = self.layers.distance.reshape(
(container.size, self.layers.max_layers))
container['rho_int'] = np.empty((container.size), dtype=FTYPE)
# don't forget to un-link everything again
self.data.unlink_containers()
# --- setup cross section and survival probability ---
if self.calc_mode == 'binned':
# The cross-sections do not depend on nc/cc, so we can at least link those containers
self.data.link_containers('nue', ['nue_cc', 'nue_nc'])
self.data.link_containers('nuebar', ['nuebar_cc', 'nuebar_nc'])
self.data.link_containers('numu', ['numu_cc', 'numu_nc'])
self.data.link_containers('numubar', ['numubar_cc', 'numubar_nc'])
self.data.link_containers('nutau', ['nutau_cc', 'nutau_nc'])
self.data.link_containers('nutaubar',
['nutaubar_cc', 'nutaubar_nc'])
for container in self.data:
container['xsection'] = np.empty((container.size), dtype=FTYPE)
container['survival_prob'] = np.empty((container.size),
dtype=FTYPE)
self.data.unlink_containers()
def setup_function(self):
sys.path.append(self.globes_wrapper)
import GLoBES
### you need to start GLoBES from the folder containing a dummy experiment
# therefore we go to the folder, load GLoBES and then go back
curdir = os.getcwd()
os.chdir(self.globes_wrapper)
self.globes_calc = GLoBES.GLoBESCalculator("calc")
os.chdir(curdir)
self.globes_calc.InitSteriles(2)
# object for oscillation parameters
self.osc_params = OscParams()
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height.m_as('km')
detector_depth = self.detector_depth.m_as('km')
self.layers = Layers(earth_model, detector_depth, prop_height)
# The electron fractions are taken into account internally by GLoBES/SNU.
# See the SNU patch for details. It uses the density to decide
# whether it is in the core or in the mantle. Therefore, we just multiply by
# one to give GLoBES the raw densities.
self.layers.setElecFrac(1., 1., 1.)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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'].get('host'))
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 probability containers
for container in self.data:
container['prob_e'] = np.empty((container.size), dtype=FTYPE)
container['prob_mu'] = np.empty((container.size), dtype=FTYPE)
container['prob_nonsterile'] = np.empty((container.size), dtype=FTYPE)
class decoherence(PiStage):
"""
PISA Pi stage representing oscillations in the presence of decoherence
Paramaters
----------
Uses the standard parameters as required by a PISA pi stage (see `pisa/core/pi_stage.py`)
Expected contents of `params` ParamSet:
detector_depth : float
earth_model : PREM file path
prop_height : quantity (dimensionless)
YeI : quantity (dimensionless)
YeO : quantity (dimensionless)
YeM : quantity (dimensionless)
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltacp : quantity (angle)
gamma12 : quantity (energy)
gamma13 : quantity (energy)
gamma23 : quantity (energy)
"""
def __init__(
self,
data=None,
params=None,
input_names=None,
output_names=None,
debug_mode=None,
input_specs=None,
calc_specs=None,
output_specs=None,
):
expected_params = (
'detector_depth',
'earth_model',
'prop_height',
'YeI',
'YeO',
'YeM',
'theta12',
'theta13',
'theta23',
'deltam21',
'deltam31',
'deltacp',
'gamma21',
'gamma31',
'gamma32',
)
input_names = ()
output_names = ()
# what are the keys used from the inputs during apply
input_apply_keys = (
'weights',
'sys_flux',
)
# what are keys added or altered in the calculation used during apply
output_calc_keys = (
'prob_e',
'prob_mu',
)
# what keys are added or altered for the outputs during apply
output_apply_keys = ('weights', )
# init base class
super(decoherence, self).__init__(
data=data,
params=params,
expected_params=expected_params,
input_names=input_names,
output_names=output_names,
debug_mode=debug_mode,
input_specs=input_specs,
calc_specs=calc_specs,
output_specs=output_specs,
input_apply_keys=input_apply_keys,
output_calc_keys=output_calc_keys,
output_apply_keys=output_apply_keys,
)
#Have not yet implemented matter effects
if self.params.earth_model.value is not None:
raise ValueError(
"Matter effects not yet implemented for decoherence, must set 'earth_model' to None"
)
assert self.input_mode is not None
assert self.calc_mode is not None
assert self.output_mode is not None
self.layers = None
#Toggle between 2-flavor and 3-flavor models
self.two_flavor = False
def setup_function(self):
# setup Earth model
if self.params.earth_model.value is not None:
earth_model = find_resource(self.params.earth_model.value)
YeI = self.params.YeI.value.m_as('dimensionless')
YeO = self.params.YeO.value.m_as('dimensionless')
YeM = self.params.YeM.value.m_as('dimensionless')
else:
earth_model = None
# setup the layers
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)
if earth_model is not None:
self.layers.setElecFrac(YeI, YeO, YeM)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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:
if self.params.earth_model.value is not None:
self.layers.calcLayers(container['true_coszen'].get('host'))
container['densities'] = self.layers.density.reshape(
(container.size, self.layers.max_layers))
container['distances'] = self.layers.distance.reshape(
(container.size, self.layers.max_layers))
else:
self.layers.calcPathLength(
container['true_coszen'].get('host'))
container['distances'] = self.layers.distance
# don't forget to un-link everything again
self.data.unlink_containers()
# --- setup empty arrays ---
if self.calc_mode == 'binned':
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)
@profile
def compute_function(self):
# set the correct data mode
self.data.data_specs = self.calc_specs
if self.calc_mode == 'binned':
# 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'
])
# --- update params ---
self.decoh_params = DecoherenceParams(
deltam21=self.params.deltam21.value,
deltam31=self.params.deltam31.value,
theta12=self.params.theta12.value,
theta13=self.params.theta13.value,
theta23=self.params.theta23.value,
deltacp=self.params.deltacp.value,
gamma21=self.params.gamma21.value,
gamma31=self.params.gamma31.value,
gamma32=self.params.gamma32.value)
# Calculate oscillation probabilities
for container in self.data:
self.calc_probs(
container['nubar'],
container['true_energy'],
#container['densities'],
container['distances'],
out=container['probability'],
)
# the following is flavour specific, hence unlink
self.data.unlink_containers()
for container in self.data:
# initial electrons (0)
fill_probs(
container['probability'].get(WHERE),
0, # electron
container['flav'],
out=container['prob_e'].get(WHERE),
)
# initial muons (1)
fill_probs(
container['probability'].get(WHERE),
1, # muon
container['flav'],
out=container['prob_mu'].get(WHERE),
)
container['prob_e'].mark_changed(WHERE)
container['prob_mu'].mark_changed(WHERE)
@profile
def apply_function(self):
# update the outputted weights
for container in self.data:
apply_probs(container['sys_flux'].get(WHERE),
container['prob_e'].get(WHERE),
container['prob_mu'].get(WHERE),
out=container['weights'].get(WHERE))
container['weights'].mark_changed(WHERE)
def calc_probs(self, nubar, e_array, len_array, out):
#Get the probability values output array
prob_array = out.get(WHERE)
#Attach units
L = len_array.get(WHERE) * ureg["km"]
E = e_array.get(WHERE) * ureg["GeV"]
#nue
calc_decoherence_probs(decoh_params=self.decoh_params,
flav="nue",
energy=E,
baseline=L,
prob_e=prob_array[:, 0, 0],
prob_mu=prob_array[:, 0, 1],
prob_tau=prob_array[:, 0, 2],
two_flavor=self.two_flavor)
#numu
calc_decoherence_probs(decoh_params=self.decoh_params,
flav="numu",
energy=E,
baseline=L,
prob_e=prob_array[:, 1, 0],
prob_mu=prob_array[:, 1, 1],
prob_tau=prob_array[:, 1, 2],
two_flavor=self.two_flavor)
#nutau (basically just the inverse of the numu case)
np.copyto(dst=prob_array[:, 2, 0], src=prob_array[:, 1, 0])
np.copyto(dst=prob_array[:, 2, 1], src=prob_array[:, 1, 2])
np.copyto(dst=prob_array[:, 2, 2], src=prob_array[:, 1, 1])
#Register that arrays have changed
out.mark_changed(WHERE)
class pi_earth_absorption(PiStage):
"""
earth absorption PISA Pi class
Paramaters
----------
earth_model : str
PREM file path
xsec_file : str
path to ROOT file containing cross-sections
detector_depth : quantity (distance), optional
detector depth
prop_height : quantity (distance), optional
height of neutrino production in the atmosphere
Notes
-----
"""
def __init__(self,
earth_model,
xsec_file,
data=None,
params=None,
input_names=None,
output_names=None,
debug_mode=None,
input_specs=None,
calc_specs=None,
output_specs=None,
detector_depth=2. * ureg.km,
prop_height=20. * ureg.km):
expected_params = ()
input_names = ()
output_names = ()
input_apply_keys = ('weights', )
# The weights are simply scaled by the earth survival probability
output_calc_keys = ('survival_prob', )
output_apply_keys = ('weights', )
# init base class
super(pi_earth_absorption, self).__init__(
data=data,
params=params,
expected_params=expected_params,
input_names=input_names,
output_names=output_names,
debug_mode=debug_mode,
input_specs=input_specs,
calc_specs=calc_specs,
output_specs=output_specs,
input_apply_keys=input_apply_keys,
output_calc_keys=output_calc_keys,
output_apply_keys=output_apply_keys,
)
assert self.input_mode is not None
assert self.calc_mode is not None
assert self.output_mode is not None
self.layers = None
self.xsroot = None
self.earth_model = earth_model
self.xsec_file = xsec_file
self.detector_depth = detector_depth.m_as('km')
self.prop_height = prop_height.m_as('km')
# this does nothing for speed, but makes for convenient numpy style broadcasting
# TODO: Use numba vectorization (not sure how that works with splines)
self.calculate_xsections = np.vectorize(self.calculate_xsections)
def setup_function(self):
import ROOT
# setup the layers
earth_model = find_resource(self.earth_model)
self.layers = Layers(earth_model, self.detector_depth,
self.prop_height)
# This is a bit hacky, but setting the electron density to 1.
# gives us the total density of matter, which is what we want.
self.layers.setElecFrac(1., 1., 1.)
# setup cross-sections
self.xsroot = ROOT.TFile(self.xsec_file)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# layers don't care about flavor
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'].get(WHERE))
container['densities'] = self.layers.density.reshape(
(container.size, self.layers.max_layers))
container['distances'] = self.layers.distance.reshape(
(container.size, self.layers.max_layers))
container['rho_int'] = np.empty((container.size), dtype=FTYPE)
# don't forget to un-link everything again
self.data.unlink_containers()
# --- setup cross section and survival probability ---
if self.calc_mode == 'binned':
# The cross-sections do not depend on nc/cc, so we can at least link those containers
self.data.link_containers('nue', ['nue_cc', 'nue_nc'])
self.data.link_containers('nuebar', ['nuebar_cc', 'nuebar_nc'])
self.data.link_containers('numu', ['numu_cc', 'numu_nc'])
self.data.link_containers('numubar', ['numubar_cc', 'numubar_nc'])
self.data.link_containers('nutau', ['nutau_cc', 'nutau_nc'])
self.data.link_containers('nutaubar',
['nutaubar_cc', 'nutaubar_nc'])
for container in self.data:
container['xsection'] = np.empty((container.size), dtype=FTYPE)
container['survival_prob'] = np.empty((container.size),
dtype=FTYPE)
self.data.unlink_containers()
@profile
def compute_function(self):
# --- calculate the integrated density in the layers ---
if self.calc_mode == 'binned':
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:
calculate_integrated_rho(container['distances'].get(WHERE),
container['densities'].get(WHERE),
out=container['rho_int'].get(WHERE))
# don't forget to un-link everything again
self.data.unlink_containers()
# --- calculate survival probability ---
if self.calc_mode == 'binned':
# The cross-sections do not depend on nc/cc, so we can at least link those containers
self.data.link_containers('nue', ['nue_cc', 'nue_nc'])
self.data.link_containers('nuebar', ['nuebar_cc', 'nuebar_nc'])
self.data.link_containers('numu', ['numu_cc', 'numu_nc'])
self.data.link_containers('numubar', ['numubar_cc', 'numubar_nc'])
self.data.link_containers('nutau', ['nutau_cc', 'nutau_nc'])
self.data.link_containers('nutaubar',
['nutaubar_cc', 'nutaubar_nc'])
for container in self.data:
container['xsection'] = self.calculate_xsections(
container['flav'], container['nubar'],
container['true_energy'].get(WHERE))
calculate_survivalprob(container['rho_int'].get(WHERE),
container['xsection'].get(WHERE),
out=container['survival_prob'].get(WHERE))
container['survival_prob'].mark_changed(WHERE)
self.data.unlink_containers()
@profile
def apply_function(self):
for container in self.data:
vectorizer.multiply(container['survival_prob'],
out=container['weights'])
def calculate_xsections(self, flav, nubar, energy):
'''Calculates the cross-sections on isoscalar targets.
The result is returned in cm^2. The xsection on one
target is calculated by taking the xsection for O16
and dividing it by 16.
'''
flavor = FLAV_BAR_STR_MAPPING[(flav, nubar)]
return (self.xsroot.Get('nu_' + flavor +
'_O16').Get('tot_cc').Eval(energy) +
self.xsroot.Get('nu_' + flavor + '_O16').Get('tot_nc').Eval(
energy)) * 10**(-38) / 16. # this gives cm^2
class pi_prob3(PiStage):
"""
prob3 osc PISA Pi class
Paramaters
----------
detector_depth : float
earth_model : PREM file path
prop_height : quantity (dimensionless)
YeI : quantity (dimensionless)
YeO : quantity (dimensionless)
YeM : quantity (dimensionless)
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltacp : quantity (angle)
None
Notes
-----
"""
def __init__(
self,
data=None,
params=None,
input_names=None,
output_names=None,
debug_mode=None,
input_specs=None,
calc_specs=None,
output_specs=None,
):
expected_params = (
'detector_depth',
'earth_model',
'prop_height',
'YeI',
'YeO',
'YeM',
'theta12',
'theta13',
'theta23',
'deltam21',
'deltam31',
'deltacp',
)
input_names = ()
output_names = ()
# what are the keys used from the inputs during apply
input_apply_keys = (
'weights',
'sys_flux',
)
# what are keys added or altered in the calculation used during apply
output_calc_keys = (
'prob_e',
'prob_mu',
)
# what keys are added or altered for the outputs during apply
output_apply_keys = ('weights', )
# init base class
super(pi_prob3, self).__init__(
data=data,
params=params,
expected_params=expected_params,
input_names=input_names,
output_names=output_names,
debug_mode=debug_mode,
input_specs=input_specs,
calc_specs=calc_specs,
output_specs=output_specs,
input_apply_keys=input_apply_keys,
output_calc_keys=output_calc_keys,
output_apply_keys=output_apply_keys,
)
assert self.input_mode is not None
assert self.calc_mode is not None
assert self.output_mode is not None
self.layers = None
self.osc_params = None
def setup_function(self):
# object for oscillation parameters
self.osc_params = OscParams()
# setup the layers
#if self.params.earth_model.value is not None:
earth_model = find_resource(self.params.earth_model.value)
YeI = self.params.YeI.value.m_as('dimensionless')
YeO = self.params.YeO.value.m_as('dimensionless')
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(YeI, YeO, YeM)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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'].get('host'))
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.calc_mode == 'binned':
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 '''
propagate_array(
self.osc_params.
dm_matrix, # pylint: disable = unexpected-keyword-arg, no-value-for-parameter
self.osc_params.mix_matrix_complex,
self.osc_params.nsi_eps,
nubar,
e_array.get(WHERE),
rho_array.get(WHERE),
len_array.get(WHERE),
out=out.get(WHERE))
out.mark_changed(WHERE)
@profile
def compute_function(self):
# set the correct data mode
self.data.data_specs = self.calc_specs
if self.calc_mode == 'binned':
# 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'
])
# --- 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')
for container in self.data:
self.calc_probs(
container['nubar'],
container['true_energy'],
container['densities'],
container['distances'],
out=container['probability'],
)
# the following is flavour specific, hence unlink
self.data.unlink_containers()
for container in self.data:
# initial electrons (0)
fill_probs(
container['probability'].get(WHERE),
0,
container['flav'],
out=container['prob_e'].get(WHERE),
)
# initial muons (1)
fill_probs(
container['probability'].get(WHERE),
1,
container['flav'],
out=container['prob_mu'].get(WHERE),
)
container['prob_e'].mark_changed(WHERE)
container['prob_mu'].mark_changed(WHERE)
@profile
def apply_function(self):
# update the outputted weights
for container in self.data:
apply_probs(container['sys_flux'].get(WHERE),
container['prob_e'].get(WHERE),
container['prob_mu'].get(WHERE),
out=container['weights'].get(WHERE))
container['weights'].mark_changed(WHERE)
def setup_function(self):
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height
detector_depth = self.detector_depth
self.layers = Layers(earth_model, detector_depth, prop_height)
# We must treat densities and electron fractions correctly here, so we set them
# to 1 in the Layers module to get unweighted densities.
self.layers.setElecFrac(1, 1, 1)
nsq_units = nsq.Const() # natural units for nusquids
# Because we don't want to extrapolate, we check that all points at which we
# want to evaluate probabilities are fully contained within the node specs. This
# is of course not necessary in events mode.
if isinstance(self.node_mode, MultiDimBinning) and not self.exact_mode:
logging.debug("setting up nuSQuIDS nodes in binned mode")
# we can prepare the calculator like this only in binned mode, see
# compute_function for node_mode == "events"
self.data.representation = self.calc_mode
for container in self.data:
for var in ["true_coszen", "true_energy"]:
unit = "dimensionless" if var == "true_coszen" else "GeV"
upper_bound = np.max(self.node_mode[var].bin_edges.m_as(unit))
lower_bound = np.min(self.node_mode[var].bin_edges.m_as(unit))
err_msg = (
"The outer edges of the node_mode must encompass "
"the entire range of calc_specs to avoid extrapolation"
)
if np.any(container[var] > upper_bound):
maxval = np.max(container[var])
raise ValueError(
err_msg + f"\nmax input: {maxval}, upper "
f"bound: {upper_bound}"
)
if np.any(container[var] < lower_bound):
minval = np.max(container[var])
raise ValueError(
err_msg + f"\nmin input: {minval}, lower "
f"bound: {lower_bound}"
)
# Layers in nuSQuIDS are special: We need all the individual distances and
# densities for the nodes to solve the interaction picture states, but on
# the final calculation grid (or events) we only need the *total* traversed
# distance. Because we are placing nodes at the bin edges rather than the
# bin middle, this doesn't really fit with how containers store data, so we
# are making arrays as variables that never go into the container.
# These are stored because we need them later during interpolation
self.coszen_node_mode = self.node_mode["true_coszen"].bin_edges.m_as(
"dimensionless"
)
self.e_node_mode = self.node_mode["true_energy"].bin_edges.m_as("GeV")
logging.debug(
f"Setting up nodes at\n"
f"cos_zen = \n{self.coszen_node_mode}\n"
f"energy = \n{self.e_node_mode}\n"
)
# things are getting a bit meshy from here...
self.e_mesh, self.cosz_mesh = np.meshgrid(
self.e_node_mode, self.coszen_node_mode
)
e_nodes = self.e_mesh.ravel()
coszen_nodes = self.cosz_mesh.ravel()
# The lines below should not be necessary because we will always get at
# least two numbers from the bin edges. However, if either energy or coszen
# somehow was just a scalar, we would need to broadcast it out to the same
# size. Keeping the code in here in case you want to use the stage in 1D.
# convert lists to ndarrays and scalars to ndarrays with length 1
e_nodes = np.atleast_1d(e_nodes)
coszen_nodes = np.atleast_1d(coszen_nodes)
# broadcast against each other and make a copy
# (see https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html)
e_nodes, coszen_nodes = [
np.array(a) for a in np.broadcast_arrays(e_nodes, coszen_nodes)
]
assert len(e_nodes) == len(coszen_nodes)
assert coszen_nodes.ndim == 1
assert e_nodes.ndim == 1
self.layers.calcLayers(coszen_nodes)
distances = np.reshape(
self.layers.distance, (len(e_nodes), self.layers.max_layers)
)
densities = np.reshape(
self.layers.density, (len(e_nodes), self.layers.max_layers)
)
# HACK: We need the correct electron densities for each layer. We can
# determine whether we are in the core or mantle based on the density.
# Needless to say it isn't optimal to have these numbers hard-coded.
ye = np.zeros_like(densities)
ye[densities < 10] = self.YeM
ye[(densities >= 10) & (densities < 13)] = self.YeO
ye[densities >= 13] = self.YeI
self.nus_layer = self.nusquids_layers_class(
distances * nsq_units.km,
densities,
ye,
e_nodes * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.both,
)
self.apply_prop_settings(self.nus_layer)
# Now that we have our nusquids calculator set up on the node grid, we make
# container output space for the probability output which may be on a finer grid
# than the nodes or even working in events mode.
self.data.representation = self.calc_mode
# --- calculate the layers ---
if isinstance(self.calc_mode, MultiDimBinning):
# 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",
],
)
# calculate the distance difference between minimum and maximum production
# height, if applicable
if self.avg_height:
layers_min = Layers(
earth_model,
detector_depth,
self.prop_height - self.prop_height_range / 2.0,
)
layers_min.setElecFrac(1, 1, 1)
layers_max = Layers(
earth_model,
detector_depth,
self.prop_height + self.prop_height_range / 2.0,
)
layers_max.setElecFrac(1, 1, 1)
for container in self.data:
self.layers.calcLayers(container["true_coszen"])
distances = self.layers.distance.reshape((container.size, -1))
tot_distances = np.sum(distances, axis=1)
if self.avg_height:
layers_min.calcLayers(container["true_coszen"])
dists_min = layers_min.distance.reshape((container.size, -1))
min_tot_dists = np.sum(dists_min, axis=1)
layers_max.calcLayers(container["true_coszen"])
dists_max = layers_max.distance.reshape((container.size, -1))
max_tot_dists = np.sum(dists_max, axis=1)
# nuSQuIDS assumes the original distance is the longest distance and
# the averaging range is the difference between the minimum and maximum
# distance.
avg_ranges = max_tot_dists - min_tot_dists
tot_distances = max_tot_dists
assert np.all(avg_ranges > 0)
# If the low-pass cutoff is zero, nusquids will not evaluate the filter.
container["lowpass_cutoff"] = self.eval_lowpass_cutoff * np.ones(
container.size
)
if not self.apply_lowpass_above_hor:
container["lowpass_cutoff"] = np.where(
container["true_coszen"] >= 0, 0, container["lowpass_cutoff"]
)
if isinstance(self.node_mode, MultiDimBinning) and not self.exact_mode:
# To project out probabilities we only need the *total* distance
container["tot_distances"] = tot_distances
if self.avg_height:
container["avg_ranges"] = avg_ranges
else:
container["avg_ranges"] = np.zeros(container.size, dtype=FTYPE)
if not self.apply_height_avg_below_hor:
container["avg_ranges"] = np.where(
container["true_coszen"] >= 0, container["avg_ranges"], 0.0
)
elif self.node_mode == "events" or self.exact_mode:
# in any other mode (events or exact) we store all densities and
# distances in the container in calc_specs
densities = self.layers.density.reshape((container.size, -1))
container["densities"] = densities
container["distances"] = distances
self.data.unlink_containers()
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar", ["nutaubar_cc", "nutaubar_nc"])
# 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)
if self.use_taus:
container["prob_tau"] = np.empty((container.size), dtype=FTYPE)
self.data.unlink_containers()
if self.exact_mode:
return
# --- containers for interpolated states ---
# This is not needed in exact mode
if isinstance(self.calc_mode, MultiDimBinning):
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["interp_states_e"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
container["interp_states_mu"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
container["interp_states_tau"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
self.data.unlink_containers()
self.interpolation_warning_issued = False
class nusquids(Stage):
"""
PISA Pi stage for weighting events due to the effect of neutrino oscillations, using
nuSQuIDS as the oscillation probability calculator. One specialty here is that we
have to specify an additional binning to determine where to place nodes for the
exact calculation. The points where the actual probability is evaluated is
determined by calc_mode as usual and may be much finer than node_mode or even
event-wise since the interpolation step is fast.
Parameters
----------
Uses the standard parameters as required by a PISA pi stage
(see `pisa/core/stage.py`)
node_mode : MultiDimBinning
Binning to determine where to place nodes at which the evaluation of interaction
states occurs. The nodes are places at the _corners_ of the binning to avoid
extrapolation.
use_decoherence : bool
set to true to include neutrino decoherence in the oscillation probability
calculation
num_decoherence_gamma : int
number of decoherence gamma parameters to be considered in the decoherence model
must be either 1 or 3
use_nsi : bool
set to true to include Non-Standard Interactions (NSI) in the oscillation
probability calculation
num_neutrinos : int
Number of neutrino flavors to include. This stage supports 3 or 4 flavors, but
nuSQuIDS allows up to 6 such that this stage could easily be expanded.
earth_model : str
Path to Earth model (PREM) file.
detector_depth : quantity (distance)
prop_height : quantity (distance) or str
Height at which neutrinos are produced. If a quantity is given, the height is
assumed to be the same for all neutrinos. An alternative is to pass
`from_container`. In that case, the stage will search the input container for a
key `prop_height` and take the height from there on a bin-wise or event-wise
basis depending on `calc_specs`.
prop_height_range : quantity (distance)
Production height is averaged around the mean set by `prop_height` assuming
a uniform distribution in [mean - range/2, mean + range/2]. The production
heights are projected onto the direction of the neutrino, such that the
averaging range is longer for shallow angles above the horizon.
apply_lowpass_above_hor : bool
Whether to apply the low-pass filter for evaluations above the horizon. If
`True` (default), the low-pass filter is applied everywhere. If `False`, the
filter is applied only below the horizon. Because propagation distances are
very short above the horizon, fast oscillations no longer average out and the
filter might wash out important features.
apply_height_avg_below_hor : bool
Whether to apply the production height averaging below the horizon. If `True`
(default), the production height averaging is applied everywhere if a
`prop_height_range` is set. If `False`, the height averaging is only applied
above the horizon. Since the production height is only a very small fraction
of the total propagation distance below the horizon, the height averaging is
no longer important and a little bit of time can be saved by computing the
slightly cheaper non-averaged probabilities.
YeI : quantity (dimensionless)
Inner electron fraction.
YeO : quantity (dimensionless)
Outer electron fraction.
YeM : quantity (dimensionless)
Mantle electron fraction.
rel_err : float
Relative error of the numerical integration
abs_err : float
Absolute error of the numerical integration
prop_lowpass_cutoff : quantity (1/distance)
Frequency cutoff for fast oscillations applied during numerical integration
of the interaction state. The frequency is passed as oscillations per distance.
A reasonable order of magnitude would allow ~100 oscillations over 12000 km.
prop_lowpass_frac : quantity (dimensionless)
This number determines how harsh the cut-off of the low-pass filter is applied
during numerical integration is. A value of 0.1 would mean that the filter would
begin to kick in when 90% of the cutoff frequency is reached and linearly
decrease oscillation amplitudes until the cutoff is reached.
eval_lowpass_cutoff : quantity (1/distance)
Same as `prop_lowpass_cutoff`, but applied during evaluation of interpolated
states, not during integration.
eval_lowpass_frac : quantity (1/distance)
Same as `prop_lowpass_frac`, but applied during evaluation of interpolated
states, not during integration.
suppress_interpolation_warning : bool
Suppress warning about negative probabilities that can indicate insufficient
nodes in a problematic region of energy and coszen. Set this option only at your
own risk after you optimized nodes and are sure that remaining negative
probabilities won't be a problem!
exact_mode : bool
With this turned on, the probabilities are evaluated using the exact calculation
for constant densities in every layer without numerical integration. This method
is much faster than the numerical integration for a node, but you lose the
option to project out probabilities from interaction picture states. In this
mode, nuSQuIDS behaves essentially like GLoBES with the same speed trade-off.
You cannot apply filters in this mode either. Its only recommended use is for
pseudo-data generation, where you may want an exact event-by-event calculation
that is allowed to take several minutes.
vacuum : bool
Do not include matter effects. Greatly increases evaluation speed.
params : ParamSet or sequence with which to instantiate a ParamSet.
Expected params .. ::
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltacp : quantity (angle)
Additional expected params if `num_neutrinos == 4` .. ::
theta14 : quantity (angle)
theta24 : quantity (angle)
deltam41 : quantity (mass^2)
deltacp14 : quantity (angle)
deltacp24 : quantity (angle)
Additional ParamSet params expected when using the `use_decoherence` argument:
n_energy : quantity (dimensionless)
* If using `num_decoherence_gamma` == 1:
gamma : quantity (energy)
* If using `num_decoherence_gamma` == 3:
gamma12 : quantity (energy)
gamma13 : quantity (energy)
gamma23 : quantity (energy)
"""
def __init__(
self,
earth_model=None,
detector_depth=None,
prop_height=None,
prop_height_range=None,
YeI=None,
YeO=None,
YeM=None,
rel_err=None,
abs_err=None,
prop_lowpass_cutoff=None,
prop_lowpass_frac=None,
eval_lowpass_cutoff=None,
eval_lowpass_frac=None,
apply_lowpass_above_hor=True,
apply_height_avg_below_hor=True,
suppress_interpolation_warning=False,
node_mode=None,
use_decoherence=False,
num_decoherence_gamma=1,
use_nsi=False,
num_neutrinos=3,
use_taus=False,
exact_mode=False,
vacuum=False,
**std_kwargs,
):
# Checks
if use_nsi:
raise NotImplementedError("NSI not implemented")
if type(prop_height) is not ureg.Quantity:
raise NotImplementedError(
"Getting propagation heights from containers is "
"not yet implemented, saw {} type".format(type(prop_height))
)
# Store args
self.num_neutrinos = int(num_neutrinos)
assert (
self.num_neutrinos < 5
), "currently only supports up to 4 flavor oscillations"
self.use_nsi = use_nsi
self.use_decoherence = use_decoherence
self.num_decoherence_gamma = num_decoherence_gamma
self.node_mode = node_mode
self.vacuum = vacuum
self.use_taus = use_taus
self.earth_model = earth_model
self.YeI = YeI.m_as("dimensionless")
self.YeO = YeO.m_as("dimensionless")
self.YeM = YeM.m_as("dimensionless")
self.detector_depth = detector_depth.m_as("km")
self.prop_height = prop_height.m_as("km")
self.avg_height = False
self.concurrent_threads = PISA_NUM_THREADS if TARGET == "parallel" else 1
self.prop_height_range = None
self.apply_height_avg_below_hor = apply_height_avg_below_hor
if prop_height_range is not None: # this is optional
self.prop_height_range = prop_height_range.m_as("km")
self.avg_height = True
self.layers = None
self.rel_err = rel_err.m_as("dimensionless") if rel_err is not None else 1.0e-10
self.abs_err = abs_err.m_as("dimensionless") if abs_err is not None else 1.0e-10
self.prop_lowpass_cutoff = (
prop_lowpass_cutoff.m_as("1/km") if prop_lowpass_cutoff is not None else 0.0
)
self.prop_lowpass_frac = (
prop_lowpass_frac.m_as("dimensionless")
if prop_lowpass_frac is not None
else 0.0
)
self.eval_lowpass_cutoff = (
eval_lowpass_cutoff.m_as("1/km") if eval_lowpass_cutoff is not None else 0.0
)
self.eval_lowpass_frac = (
eval_lowpass_frac.m_as("dimensionless")
if eval_lowpass_frac is not None
else 0.0
)
if self.prop_lowpass_frac > 1.0 or self.eval_lowpass_frac > 1.0:
raise ValueError("lowpass filter fraction cannot be greater than one")
if self.prop_lowpass_frac < 0.0 or self.eval_lowpass_frac < 0.0:
raise ValueError("lowpass filter fraction cannot be smaller than zero")
self.apply_lowpass_above_hor = apply_lowpass_above_hor
self.nus_layer = None
self.nus_layerbar = None
# Define the layers class
self.nusquids_layers_class = nsq.nuSQUIDSLayers
# Define standard params
expected_params = [
"theta12",
"theta13",
"theta23",
"deltam21",
"deltam31",
"deltacp",
]
# Add decoherence parameters
if self.use_decoherence:
# Use derived nuSQuIDS classes
import nuSQUIDSDecohPy
self.nusquids_layers_class = nuSQUIDSDecohPy.nuSQUIDSDecohLayers
# Checks
assert (
self.num_neutrinos == 3
), "Decoherence only supports 3 neutrinos currently"
# Add decoherence params
expected_params.extend(["gamma0"])
expected_params.extend(["n"])
expected_params.extend(["E0"])
# We may want to reparametrize this with the difference between deltacp14 and
# deltacp24, as the absolute value seems to play a small role (see
# https://arxiv.org/pdf/2010.06321.pdf)
if self.num_neutrinos == 4:
expected_params.extend(
[
"theta14",
"theta24",
"theta34",
"deltam41",
"deltacp14",
"deltacp24",
]
)
# init base class
super().__init__(
expected_params=expected_params,
**std_kwargs,
)
# This is special: We have an additional "binning" to account for. It is in
# principle possible to work in event mode even for the nodes, which would mean
# that the full oscillation problem is solved for all events individually.
# Together with the constant oscillation mode, this can be used to calculate
# probabilities in exact mode in a time that is reasonable at least for
# generating pseudodata.
assert not (self.use_nsi and self.use_decoherence), (
"NSI and decoherence not " "suported together, must use one or the other"
)
self.exact_mode = exact_mode
if exact_mode:
# No interpolation is happening in exact mode so any passed node_mode
# will be ignored. Probabilities are calculated at calc_specs.
if self.node_mode is not None:
logging.warn(
"nuSQuIDS is configured in exact mode, the passed "
f"`node_mode`\n({self.node_mode})\n will be ignored!"
)
if self.prop_lowpass_cutoff > 0 or self.eval_lowpass_cutoff > 0:
logging.warn(
"nuSQuIDS is configured in exact mode, low-pass filters "
"will be ignored"
)
else:
if isinstance(self.calc_mode, MultiDimBinning):
assert isinstance(self.node_mode, MultiDimBinning), (
"cannot use " "event-wise nodes with binned calculation"
)
self.e_node_mode = None
self.e_mesh = None
self.coszen_node_mode = None
self.cosz_mesh = None
# We don't want to spam the user with repeated warnings about the same issue.
self.interpolation_warning_issued = suppress_interpolation_warning
def set_osc_parameters(self, nus_layer):
# nuSQuIDS uses zero-index for mixing angles
nus_layer.Set_MixingAngle(0, 1, self.params.theta12.value.m_as("rad"))
nus_layer.Set_MixingAngle(0, 2, self.params.theta13.value.m_as("rad"))
nus_layer.Set_MixingAngle(1, 2, self.params.theta23.value.m_as("rad"))
# mass differences in nuSQuIDS are always w.r.t. m_1
nus_layer.Set_SquareMassDifference(1, self.params.deltam21.value.m_as("eV**2"))
nus_layer.Set_SquareMassDifference(2, self.params.deltam31.value.m_as("eV**2"))
nus_layer.Set_CPPhase(0, 2, self.params.deltacp.value.m_as("rad"))
# set decoherence parameters
if self.use_decoherence:
nsq_units = nsq.Const() # TODO Once only (make into a member)
gamma0 = self.params.gamma0.value.m_as("eV") * nsq_units.eV
gamma0_matrix_diagonal = np.array(
[0.0, gamma0, gamma0, gamma0, gamma0, gamma0, gamma0, gamma0, gamma0]
) # "State selection" case (see arXiv:2007.00068 eqn 11) #TODO implement other models
nus_layer.Set_DecoherenceGammaMatrixDiagonal(gamma0_matrix_diagonal)
nus_layer.Set_DecoherenceGammaEnergyDependence(
self.params.n.value.m_as("dimensionless")
)
nus_layer.Set_DecoherenceGammaEnergyScale(
self.params.E0.value.m_as("eV") * nsq_units.eV
)
if self.num_neutrinos == 3:
return
nus_layer.Set_MixingAngle(0, 3, self.params.theta14.value.m_as("rad"))
nus_layer.Set_MixingAngle(1, 3, self.params.theta24.value.m_as("rad"))
nus_layer.Set_MixingAngle(2, 3, self.params.theta34.value.m_as("rad"))
nus_layer.Set_SquareMassDifference(3, self.params.deltam41.value.m_as("eV**2"))
nus_layer.Set_CPPhase(0, 3, self.params.deltacp14.value.m_as("rad"))
nus_layer.Set_CPPhase(1, 3, self.params.deltacp24.value.m_as("rad"))
# TODO: Implement NSI, decoherence
def apply_prop_settings(self, nus_layer):
nsq_units = nsq.Const()
nus_layer.Set_rel_error(self.rel_err)
nus_layer.Set_abs_error(self.abs_err)
nus_layer.Set_EvolLowPassCutoff(self.prop_lowpass_cutoff / nsq_units.km)
# The ramp of the low-pass filter starts to drop at (cutoff - scale)
scale = self.prop_lowpass_frac * self.prop_lowpass_cutoff / nsq_units.km
nus_layer.Set_EvolLowPassScale(scale)
nus_layer.Set_AllowConstantDensityOscillationOnlyEvolution(self.exact_mode)
nus_layer.Set_EvalThreads(self.concurrent_threads)
def setup_function(self):
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height
detector_depth = self.detector_depth
self.layers = Layers(earth_model, detector_depth, prop_height)
# We must treat densities and electron fractions correctly here, so we set them
# to 1 in the Layers module to get unweighted densities.
self.layers.setElecFrac(1, 1, 1)
nsq_units = nsq.Const() # natural units for nusquids
# Because we don't want to extrapolate, we check that all points at which we
# want to evaluate probabilities are fully contained within the node specs. This
# is of course not necessary in events mode.
if isinstance(self.node_mode, MultiDimBinning) and not self.exact_mode:
logging.debug("setting up nuSQuIDS nodes in binned mode")
# we can prepare the calculator like this only in binned mode, see
# compute_function for node_mode == "events"
self.data.representation = self.calc_mode
for container in self.data:
for var in ["true_coszen", "true_energy"]:
unit = "dimensionless" if var == "true_coszen" else "GeV"
upper_bound = np.max(self.node_mode[var].bin_edges.m_as(unit))
lower_bound = np.min(self.node_mode[var].bin_edges.m_as(unit))
err_msg = (
"The outer edges of the node_mode must encompass "
"the entire range of calc_specs to avoid extrapolation"
)
if np.any(container[var] > upper_bound):
maxval = np.max(container[var])
raise ValueError(
err_msg + f"\nmax input: {maxval}, upper "
f"bound: {upper_bound}"
)
if np.any(container[var] < lower_bound):
minval = np.max(container[var])
raise ValueError(
err_msg + f"\nmin input: {minval}, lower "
f"bound: {lower_bound}"
)
# Layers in nuSQuIDS are special: We need all the individual distances and
# densities for the nodes to solve the interaction picture states, but on
# the final calculation grid (or events) we only need the *total* traversed
# distance. Because we are placing nodes at the bin edges rather than the
# bin middle, this doesn't really fit with how containers store data, so we
# are making arrays as variables that never go into the container.
# These are stored because we need them later during interpolation
self.coszen_node_mode = self.node_mode["true_coszen"].bin_edges.m_as(
"dimensionless"
)
self.e_node_mode = self.node_mode["true_energy"].bin_edges.m_as("GeV")
logging.debug(
f"Setting up nodes at\n"
f"cos_zen = \n{self.coszen_node_mode}\n"
f"energy = \n{self.e_node_mode}\n"
)
# things are getting a bit meshy from here...
self.e_mesh, self.cosz_mesh = np.meshgrid(
self.e_node_mode, self.coszen_node_mode
)
e_nodes = self.e_mesh.ravel()
coszen_nodes = self.cosz_mesh.ravel()
# The lines below should not be necessary because we will always get at
# least two numbers from the bin edges. However, if either energy or coszen
# somehow was just a scalar, we would need to broadcast it out to the same
# size. Keeping the code in here in case you want to use the stage in 1D.
# convert lists to ndarrays and scalars to ndarrays with length 1
e_nodes = np.atleast_1d(e_nodes)
coszen_nodes = np.atleast_1d(coszen_nodes)
# broadcast against each other and make a copy
# (see https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html)
e_nodes, coszen_nodes = [
np.array(a) for a in np.broadcast_arrays(e_nodes, coszen_nodes)
]
assert len(e_nodes) == len(coszen_nodes)
assert coszen_nodes.ndim == 1
assert e_nodes.ndim == 1
self.layers.calcLayers(coszen_nodes)
distances = np.reshape(
self.layers.distance, (len(e_nodes), self.layers.max_layers)
)
densities = np.reshape(
self.layers.density, (len(e_nodes), self.layers.max_layers)
)
# HACK: We need the correct electron densities for each layer. We can
# determine whether we are in the core or mantle based on the density.
# Needless to say it isn't optimal to have these numbers hard-coded.
ye = np.zeros_like(densities)
ye[densities < 10] = self.YeM
ye[(densities >= 10) & (densities < 13)] = self.YeO
ye[densities >= 13] = self.YeI
self.nus_layer = self.nusquids_layers_class(
distances * nsq_units.km,
densities,
ye,
e_nodes * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.both,
)
self.apply_prop_settings(self.nus_layer)
# Now that we have our nusquids calculator set up on the node grid, we make
# container output space for the probability output which may be on a finer grid
# than the nodes or even working in events mode.
self.data.representation = self.calc_mode
# --- calculate the layers ---
if isinstance(self.calc_mode, MultiDimBinning):
# 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",
],
)
# calculate the distance difference between minimum and maximum production
# height, if applicable
if self.avg_height:
layers_min = Layers(
earth_model,
detector_depth,
self.prop_height - self.prop_height_range / 2.0,
)
layers_min.setElecFrac(1, 1, 1)
layers_max = Layers(
earth_model,
detector_depth,
self.prop_height + self.prop_height_range / 2.0,
)
layers_max.setElecFrac(1, 1, 1)
for container in self.data:
self.layers.calcLayers(container["true_coszen"])
distances = self.layers.distance.reshape((container.size, -1))
tot_distances = np.sum(distances, axis=1)
if self.avg_height:
layers_min.calcLayers(container["true_coszen"])
dists_min = layers_min.distance.reshape((container.size, -1))
min_tot_dists = np.sum(dists_min, axis=1)
layers_max.calcLayers(container["true_coszen"])
dists_max = layers_max.distance.reshape((container.size, -1))
max_tot_dists = np.sum(dists_max, axis=1)
# nuSQuIDS assumes the original distance is the longest distance and
# the averaging range is the difference between the minimum and maximum
# distance.
avg_ranges = max_tot_dists - min_tot_dists
tot_distances = max_tot_dists
assert np.all(avg_ranges > 0)
# If the low-pass cutoff is zero, nusquids will not evaluate the filter.
container["lowpass_cutoff"] = self.eval_lowpass_cutoff * np.ones(
container.size
)
if not self.apply_lowpass_above_hor:
container["lowpass_cutoff"] = np.where(
container["true_coszen"] >= 0, 0, container["lowpass_cutoff"]
)
if isinstance(self.node_mode, MultiDimBinning) and not self.exact_mode:
# To project out probabilities we only need the *total* distance
container["tot_distances"] = tot_distances
if self.avg_height:
container["avg_ranges"] = avg_ranges
else:
container["avg_ranges"] = np.zeros(container.size, dtype=FTYPE)
if not self.apply_height_avg_below_hor:
container["avg_ranges"] = np.where(
container["true_coszen"] >= 0, container["avg_ranges"], 0.0
)
elif self.node_mode == "events" or self.exact_mode:
# in any other mode (events or exact) we store all densities and
# distances in the container in calc_specs
densities = self.layers.density.reshape((container.size, -1))
container["densities"] = densities
container["distances"] = distances
self.data.unlink_containers()
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar", ["nutaubar_cc", "nutaubar_nc"])
# 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)
if self.use_taus:
container["prob_tau"] = np.empty((container.size), dtype=FTYPE)
self.data.unlink_containers()
if self.exact_mode:
return
# --- containers for interpolated states ---
# This is not needed in exact mode
if isinstance(self.calc_mode, MultiDimBinning):
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["interp_states_e"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
container["interp_states_mu"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
container["interp_states_tau"] = np.empty(
(container.size, self.num_neutrinos ** 2),
dtype=FTYPE,
)
self.data.unlink_containers()
self.interpolation_warning_issued = False
# @line_profile
def calc_node_probs(self, nus_layer, flav_in, flav_out, n_nodes):
"""
Evaluate oscillation probabilities at nodes. This does not require any
interpolation.
"""
ini_state = np.array([0] * self.num_neutrinos)
ini_state[flav_in] = 1
nus_layer.Set_initial_state(ini_state, nsq.Basis.flavor)
if not self.vacuum:
nus_layer.EvolveState()
prob_nodes = nus_layer.EvalFlavorAtNodes(flav_out)
return prob_nodes
def calc_interpolated_states(self, evolved_states, e_out, cosz_out):
"""
Calculate interpolated states at the energies and zenith angles requested.
"""
nsq_units = nsq.Const()
interp_states = np.zeros((e_out.size, evolved_states.shape[1]))
assert np.all(e_out <= np.max(self.e_node_mode * nsq_units.GeV))
assert np.all(e_out >= np.min(self.e_node_mode * nsq_units.GeV))
assert np.all(cosz_out <= np.max(self.coszen_node_mode))
assert np.all(cosz_out >= np.min(self.coszen_node_mode))
for i in range(evolved_states.shape[1]):
z = evolved_states[:, i].reshape(self.e_mesh.shape).T
assert np.all(np.isfinite(z))
# RectBivariateSpline takes in the 1D node position and assumes that they
# are on a mesh.
f = RectBivariateSpline(
np.log10(self.e_node_mode * nsq_units.GeV),
self.coszen_node_mode,
z,
kx=2,
ky=2,
)
interp_states[..., i] = f(np.log10(e_out), cosz_out, grid=False)
return interp_states
def calc_probs_interp(
self,
flav_out,
nubar,
interp_states,
out_distances,
e_out,
avg_ranges=0,
lowpass_cutoff=0,
):
"""
Project out probabilities from interpolated interaction picture states.
"""
nsq_units = nsq.Const()
prob_interp = np.zeros(e_out.size)
scale = self.eval_lowpass_frac * lowpass_cutoff
prob_interp = self.nus_layer.EvalWithState(
flav_out,
out_distances,
e_out,
interp_states,
rho=int(nubar),
avg_cutoff=0.0,
avg_scale=0.0,
# Range averaging is only computed in the places where t_range > 0, so
# we don't need to introduce switches for averaged and non-averaged regions.
lowpass_cutoff=lowpass_cutoff,
lowpass_scale=scale,
t_range=avg_ranges,
)
return prob_interp
def compute_function_no_interpolation(self):
"""
Version of the compute function that does not use any interpolation between
nodes.
"""
nsq_units = nsq.Const()
# it is possible to work in binned calc mode while being in exact mode
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar", ["nutaubar_cc", "nutaubar_nc"])
for container in self.data:
nubar = container["nubar"] < 0
flav = container["flav"]
# HACK: We need the correct electron densities for each layer. We can
# determine whether we are in the core or mantle based on the density.
ye = np.zeros_like(container["densities"])
ye[container["densities"] < 10] = self.YeM
ye[
(container["densities"] >= 10) & (container["densities"] < 13)
] = self.YeO
ye[container["densities"] >= 13] = self.YeI
nus_layer = self.nusquids_layers_class(
container["distances"] * nsq_units.km,
container["densities"],
ye,
container["true_energy"] * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.antineutrino if nubar else nsq.NeutrinoType.neutrino,
)
self.apply_prop_settings(nus_layer)
self.set_osc_parameters(nus_layer)
container["prob_e"] = self.calc_node_probs(
nus_layer, 0, flav, container.size
)
container["prob_mu"] = self.calc_node_probs(
nus_layer, 1, flav, container.size
)
container.mark_changed("prob_e")
container.mark_changed("prob_mu")
if self.use_taus:
container["prob_tau"] = self.calc_node_probs(
nus_layer, 2, flav, container.size
)
container.mark_changed("prob_tau")
self.data.unlink_containers()
# @line_profile
def compute_function_interpolated(self):
"""
Version of the compute function that does use interpolation between nodes.
"""
nsq_units = nsq.Const()
# We need to make two evolutions, one for numu and the other for nue.
# These produce neutrino and antineutrino states at the same time thanks to
# the "both" neutrino mode of nuSQuIDS.
self.apply_prop_settings(self.nus_layer)
self.set_osc_parameters(self.nus_layer)
ini_state_nue = np.array([1, 0, 0] + [0] * (self.num_neutrinos - 3))
ini_state_numu = np.array([0, 1, 0] + [0] * (self.num_neutrinos - 3))
ini_state_nutau = np.array([0, 0, 1] + [0] * (self.num_neutrinos - 3))
self.nus_layer.Set_initial_state(ini_state_nue, nsq.Basis.flavor)
if not self.vacuum:
self.nus_layer.EvolveState()
evolved_states_nue = self.nus_layer.GetStates(0)
evolved_states_nuebar = self.nus_layer.GetStates(1)
self.nus_layer.Set_initial_state(ini_state_numu, nsq.Basis.flavor)
if not self.vacuum:
self.nus_layer.EvolveState()
evolved_states_numu = self.nus_layer.GetStates(0)
evolved_states_numubar = self.nus_layer.GetStates(1)
if self.use_taus:
self.nus_layer.Set_initial_state(ini_state_nutau, nsq.Basis.flavor)
if not self.vacuum:
self.nus_layer.EvolveState()
evolved_states_nutau = self.nus_layer.GetStates(0)
evolved_states_nutaubar = self.nus_layer.GetStates(1)
# Now comes the step where we interpolate the interaction picture states
# and project out oscillation probabilities. This can be done in either events
# or binned mode.
if isinstance(self.calc_mode, MultiDimBinning):
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:
nubar = container["nubar"] < 0
container["interp_states_e"] = self.calc_interpolated_states(
evolved_states_nuebar if nubar else evolved_states_nue,
container["true_energy"] * nsq_units.GeV,
container["true_coszen"],
)
container["interp_states_mu"] = self.calc_interpolated_states(
evolved_states_numubar if nubar else evolved_states_numu,
container["true_energy"] * nsq_units.GeV,
container["true_coszen"],
)
if self.use_taus:
container["interp_states_tau"] = self.calc_interpolated_states(
evolved_states_nutaubar if nubar else evolved_states_nutau,
container["true_energy"] * nsq_units.GeV,
container["true_coszen"],
)
self.data.unlink_containers()
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar", ["nutaubar_cc", "nutaubar_nc"])
for container in self.data:
nubar = container["nubar"] < 0
flav_out = container["flav"]
input_flavs = ["e", "mu", "tau"] if self.use_taus else ["e", "mu"]
for flav_in in input_flavs:
container["prob_" + flav_in] = self.calc_probs_interp(
flav_out=flav_out,
nubar=nubar,
interp_states=container["interp_states_" + flav_in],
out_distances=container["tot_distances"] * nsq_units.km,
e_out=container["true_energy"] * nsq_units.GeV,
avg_ranges=container["avg_ranges"] * nsq_units.km,
lowpass_cutoff=container["lowpass_cutoff"] / nsq_units.km,
)
# It is possible to get slightly negative probabilities from imperfect
# state interpolation between nodes.
# It's impractical to avoid any probability dipping below zero in every
# conceivable situation because that would require very dense node
# spacing. We get around this by flooring the probability at zero.
# However, dipping below zero by more than 1% may indicate that nodes
# aren't spaced tightly enough to achieve an acceptable accuracy, so we
# issue a warning.
if (
np.any(container["prob_" + flav_in] < -0.01)
and not self.interpolation_warning_issued
):
mask = container["prob_" + flav_in] < -0.01
en_med = np.median(container["true_energy"][mask])
cz_med = np.median(container["true_coszen"][mask])
logging.warn(
f"Some probabilities in nu_{flav_in} -> {container.name} dip "
"below zero by more than 1%! This may indicate too few nodes "
f"in the problematic region. Median energy: {en_med}, median "
f"coszen: {cz_med}. This warning is only issued once."
)
self.interpolation_warning_issued = True
container["prob_" + flav_in][container["prob_" + flav_in] < 0] = 0.0
container.mark_changed("prob_e")
container.mark_changed("prob_mu")
if self.use_taus:
container.mark_changed("prob_tau")
self.data.unlink_containers()
def compute_function(self):
if self.node_mode == "events" or self.exact_mode:
self.compute_function_no_interpolation()
else:
self.compute_function_interpolated()
@profile
def apply_function(self):
for container in self.data:
scales = (
container["nu_flux"][:, 0] * container["prob_e"]
+ container["nu_flux"][:, 1] * container["prob_mu"]
)
if self.use_taus:
scales += container["nu_flux"][:, 2] * container["prob_tau"]
container["weights"] = container["weights"] * scales
class nusquids(Stage):
"""
PISA Pi stage for weighting events due to the effect of neutrino oscillations, using
nuSQuIDS as the oscillation probability calculator. One specialty here is that we
have to specify an additional binning to determine where to place nodes for the
exact calculation. The points where the actual probability is evaluated is
determined by calc_mode as usual and may be much finer than node_mode or even
event-wise since the interpolation step is fast.
Parameters
----------
Uses the standard parameters as required by a PISA pi stage
(see `pisa/core/stage.py`)
node_mode : MultiDimBinning
Binning to determine where to place nodes at which the evaluation of interaction
states occurs. The nodes are places at the _corners_ of the binning to avoid
extrapolation.
use_decoherence : bool
set to true to include neutrino decoherence in the oscillation probability
calculation
num_decoherence_gamma : int
number of decoherence gamma parameters to be considered in the decoherence model
must be either 1 or 3
use_nsi : bool
set to true to include Non-Standard Interactions (NSI) in the oscillation
probability calculation
num_neutrinos : int
Number of neutrino flavors to include. This stage supports 3 or 4 flavors, but
nuSQuIDS allows up to 6 such that this stage could easily be expanded.
earth_model : str
Path to Earth model (PREM) file.
detector_depth : quantity (distance)
prop_height : quantity (distance) or str
Height at which neutrinos are produced. If a quantity is given, the height is
assumed to be the same for all neutrinos. An alternative is to pass
`from_container`. In that case, the stage will search the input container for a
key `prop_height` and take the height from there on a bin-wise or event-wise
basis depending on `calc_specs`.
prop_height_min : quantity (distance)
Minimum production height (optional). If this value is passed probabilities are
averaged between the maximum production height in `prop_height` and this value
under the assumption of a uniform production height distribution.
YeI : quantity (dimensionless)
Inner electron fraction.
YeO : quantity (dimensionless)
Outer electron fraction.
YeM : quantity (dimensionless)
Mantle electron fraction.
rel_err : float
Relative error of the numerical integration
abs_err : float
Absolute error of the numerical integration
prop_lowpass_cutoff : quantity (1/distance)
Frequency cutoff for fast oscillations applied during numerical integration
of the interaction state. The frequency is passed as oscillations per distance.
A reasonable order of magnitude would allow ~100 oscillations over 12000 km.
prop_lowpass_frac : quantity (dimensionless)
This number determines how harsh the cut-off of the low-pass filter is applied
during numerical integration is. A value of 0.1 would mean that the filter would
begin to kick in when 90% of the cutoff frequency is reached and linearly
decrease oscillation amplitudes until the cutoff is reached.
eval_lowpass_cutoff : quantity (1/distance)
Same as `prop_lowpass_cutoff`, but applied during evaluation of interpolated
states, not during integration.
eval_lowpass_frac : quantity (1/distance)
Same as `prop_lowpass_frac`, but applied during evaluation of interpolated
states, not during integration.
exact_mode : bool
With this turned on, the probabilities are evaluated using the exact calculation
for constant densities in every layer without numerical integration. This method
is much faster than the numerical integration for a node, but you lose the
option to project out probabilities from interaction picture states. In this
mode, nuSQuIDS behaves essentially like GLoBES with the same speed trade-off.
You cannot apply filters in this mode either. Its only recommended use is for
pseudo-data generation, where you may want an exact event-by-event calculation
that is allowed to take several minutes.
params : ParamSet or sequence with which to instantiate a ParamSet.
Expected params .. ::
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltacp : quantity (angle)
Additional expected params if `num_neutrinos == 4` .. ::
theta14 : quantity (angle)
theta24 : quantity (angle)
deltam41 : quantity (mass^2)
deltacp14 : quantity (angle)
deltacp24 : quantity (angle)
Additional ParamSet params expected when using the `use_decoherence` argument:
n_energy : quantity (dimensionless)
* If using `num_decoherence_gamma` == 1:
gamma : quantity (energy)
* If using `num_decoherence_gamma` == 3:
gamma12 : quantity (energy)
gamma13 : quantity (energy)
gamma23 : quantity (energy)
"""
def __init__(
self,
earth_model=None,
detector_depth=None,
prop_height=None,
prop_height_min=None,
YeI=None,
YeO=None,
YeM=None,
rel_err=None,
abs_err=None,
prop_lowpass_cutoff=None,
prop_lowpass_frac=None,
eval_lowpass_cutoff=None,
eval_lowpass_frac=None,
node_mode=None,
use_decoherence=False,
num_decoherence_gamma=1,
use_nsi=False,
num_neutrinos=3,
exact_mode=False,
**std_kwargs,
):
if use_nsi:
raise NotImplementedError("NSI not implemented")
if use_decoherence:
raise NotImplementedError("Decoherence not implemented")
if type(prop_height) is not ureg.Quantity:
raise NotImplementedError(
"Getting propagation heights from containers is "
"not yet implemented")
self.num_neutrinos = int(num_neutrinos)
assert self.num_neutrinos < 5, "currently only supports up to 4 flavor oscillations"
self.use_nsi = use_nsi
self.use_decoherence = use_decoherence
self.num_decoherence_gamma = num_decoherence_gamma
self.node_mode = node_mode
self.earth_model = earth_model
self.YeI = YeI.m_as("dimensionless")
self.YeO = YeO.m_as("dimensionless")
self.YeM = YeM.m_as("dimensionless")
self.detector_depth = detector_depth.m_as("km")
self.prop_height = prop_height.m_as("km")
self.avg_height = False
self.prop_height_min = None
if prop_height_min is not None: # this is optional
self.prop_height_min = prop_height_min.m_as("km")
self.avg_height = True
self.layers = None
self.rel_err = rel_err.m_as(
"dimensionless") if rel_err is not None else 1.0e-10
self.abs_err = abs_err.m_as(
"dimensionless") if abs_err is not None else 1.0e-10
self.prop_lowpass_cutoff = (prop_lowpass_cutoff.m_as("1/km")
if prop_lowpass_cutoff is not None else 0.)
self.prop_lowpass_frac = (prop_lowpass_frac.m_as("dimensionless")
if prop_lowpass_frac is not None else 0.)
self.eval_lowpass_cutoff = (eval_lowpass_cutoff.m_as("1/km")
if eval_lowpass_cutoff is not None else 0.)
self.eval_lowpass_frac = (eval_lowpass_frac.m_as("dimensionless")
if eval_lowpass_frac is not None else 0.)
if self.prop_lowpass_frac > 1. or self.eval_lowpass_frac > 1.:
raise ValueError(
"lowpass filter fraction cannot be greater than one")
if self.prop_lowpass_frac < 0. or self.eval_lowpass_frac < 0.:
raise ValueError(
"lowpass filter fraction cannot be smaller than zero")
self.nus_layer = None
self.nus_layerbar = None
# Define standard params
expected_params = [
"theta12",
"theta13",
"theta23",
"deltam21",
"deltam31",
"deltacp",
]
# Add decoherence parameters
assert self.num_decoherence_gamma in [
1, 3
], ("Must choose either 1 or 3 "
"decoherence gamma parameters")
if self.use_decoherence:
if self.num_decoherence_gamma == 1:
expected_params.extend(["gamma"])
elif self.num_decoherence_gamma == 3:
expected_params.extend(["gamma21", "gamma31", "gamma32"])
expected_params.extend(["n_energy"])
# We may want to reparametrize this with the difference between deltacp14 and
# deltacp24, as the absolute value seems to play a small role (see
# https://arxiv.org/pdf/2010.06321.pdf)
if self.num_neutrinos == 4:
expected_params.extend([
"theta14",
"theta24",
"theta34",
"deltam41",
"deltacp14",
"deltacp24",
])
# init base class
super().__init__(
expected_params=expected_params,
**std_kwargs,
)
# This is special: We have an additional "binning" to account for. It is in
# principle possible to work in event mode even for the nodes, which would mean
# that the full oscillation problem is solved for all events individually.
# Together with the constant oscillation mode, this can be used to calculate
# probabilities in exact mode in a time that is reasonable at least for
# generating pseudodata.
assert not (self.use_nsi and self.use_decoherence), (
"NSI and decoherence not "
"suported together, must use one or the other")
self.exact_mode = exact_mode
if exact_mode:
# No interpolation is happening in exact mode so any passed node_mode
# will be ignored. Probabilities are calculated at calc_specs.
if self.node_mode is not None:
logging.warn(
"nuSQuIDS is configured in exact mode, the passed "
f"`node_mode`\n({self.node_mode})\n will be ignored!")
if self.prop_lowpass_cutoff > 0 or self.eval_lowpass_cutoff > 0:
logging.warn(
"nuSQuIDS is configured in exact mode, low-pass filters "
"will be ignored")
else:
if isinstance(self.calc_mode, MultiDimBinning):
assert isinstance(self.node_mode, MultiDimBinning), (
"cannot use "
"event-wise nodes with binned calculation")
self.e_node_mode = None
self.e_mesh = None
self.coszen_node_mode = None
self.cosz_mesh = None
def set_osc_parameters(self, nus_layer):
# nuSQuIDS uses zero-index for mixing angles
nus_layer.Set_MixingAngle(0, 1, self.params.theta12.value.m_as("rad"))
nus_layer.Set_MixingAngle(0, 2, self.params.theta13.value.m_as("rad"))
nus_layer.Set_MixingAngle(1, 2, self.params.theta23.value.m_as("rad"))
# mass differences in nuSQuIDS are always w.r.t. m_1
nus_layer.Set_SquareMassDifference(
1, self.params.deltam21.value.m_as("eV**2"))
nus_layer.Set_SquareMassDifference(
2, self.params.deltam31.value.m_as("eV**2"))
nus_layer.Set_CPPhase(0, 2, self.params.deltacp.value.m_as("rad"))
if self.num_neutrinos == 3: return
nus_layer.Set_MixingAngle(0, 3, self.params.theta14.value.m_as("rad"))
nus_layer.Set_MixingAngle(1, 3, self.params.theta24.value.m_as("rad"))
nus_layer.Set_MixingAngle(2, 3, self.params.theta34.value.m_as("rad"))
nus_layer.Set_SquareMassDifference(
3, self.params.deltam41.value.m_as("eV**2"))
nus_layer.Set_CPPhase(0, 3, self.params.deltacp14.value.m_as("rad"))
nus_layer.Set_CPPhase(1, 3, self.params.deltacp24.value.m_as("rad"))
# TODO: Implement NSI, decoherence
def apply_prop_settings(self, nus_layer):
nsq_units = nsq.Const()
nus_layer.Set_rel_error(self.rel_err)
nus_layer.Set_abs_error(self.abs_err)
nus_layer.Set_EvolLowPassCutoff(self.prop_lowpass_cutoff /
nsq_units.km)
# The ramp of the low-pass filter starts to drop at (cutoff - scale)
scale = self.prop_lowpass_frac * self.prop_lowpass_cutoff / nsq_units.km
nus_layer.Set_EvolLowPassScale(scale)
nus_layer.Set_AllowConstantDensityOscillationOnlyEvolution(
self.exact_mode)
def setup_function(self):
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height
detector_depth = self.detector_depth
self.layers = Layers(earth_model, detector_depth, prop_height)
self.layers.setElecFrac(self.YeI, self.YeO, self.YeM)
nsq_units = nsq.Const() # natural units for nusquids
# Because we don't want to extrapolate, we check that all points at which we
# want to evaluate probabilities are fully contained within the node specs. This
# is of course not necessary in events mode.
if isinstance(self.node_mode, MultiDimBinning) and not self.exact_mode:
logging.debug("setting up nuSQuIDS nodes in binned mode")
# we can prepare the calculator like this only in binned mode, see
# compute_function for node_mode == "events"
self.data.representation = self.calc_mode
for container in self.data:
for var in ["true_coszen", "true_energy"]:
upper_bound = np.max(self.node_mode[var].bin_edges)
lower_bound = np.min(self.node_mode[var].bin_edges)
err_msg = (
"The outer edges of the node_mode must encompass "
"the entire range of calc_specs to avoid extrapolation"
)
if np.any(container[var] > upper_bound):
maxval = np.max(container[var])
raise ValueError(err_msg +
f"\nmax input: {maxval}, upper "
f"bound: {upper_bound}")
if np.any(container[var] < lower_bound):
minval = np.max(container[var])
raise ValueError(err_msg +
f"\nmin input: {minval}, lower "
f"bound: {lower_bound}")
# Layers in nuSQuIDS are special: We need all the individual distances and
# densities for the nodes to solve the interaction picture states, but on
# the final calculation grid (or events) we only need the *total* traversed
# distance. Because we are placing nodes at the bin edges rather than the
# bin middle, this doesn't really fit with how containers store data, so we
# are making arrays as variables that never go into the container.
# These are stored because we need them later during interpolation
self.coszen_node_mode = self.node_mode[
"true_coszen"].bin_edges.m_as("dimensionless")
self.e_node_mode = self.node_mode["true_energy"].bin_edges.m_as(
"GeV")
logging.debug(f"Setting up nodes at\n"
f"cos_zen = \n{self.coszen_node_mode}\n"
f"energy = \n{self.e_node_mode}\n")
# things are getting a bit meshy from here...
self.e_mesh, self.cosz_mesh = np.meshgrid(self.e_node_mode,
self.coszen_node_mode)
e_nodes = self.e_mesh.ravel()
coszen_nodes = self.cosz_mesh.ravel()
# The lines below should not be necessary because we will always get at
# least two numbers from the bin edges. However, if either energy or coszen
# somehow was just a scalar, we would need to broadcast it out to the same
# size. Keeping the code in here in case you want to use the stage in 1D.
# convert lists to ndarrays and scalars to ndarrays with length 1
e_nodes = np.atleast_1d(e_nodes)
coszen_nodes = np.atleast_1d(coszen_nodes)
# broadcast against each other and make a copy
# (see https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html)
e_nodes, coszen_nodes = [
np.array(a)
for a in np.broadcast_arrays(e_nodes, coszen_nodes)
]
assert len(e_nodes) == len(coszen_nodes)
assert coszen_nodes.ndim == 1
assert e_nodes.ndim == 1
self.layers.calcLayers(coszen_nodes)
distances = np.reshape(self.layers.distance,
(len(e_nodes), self.layers.max_layers))
densities = np.reshape(self.layers.density,
(len(e_nodes), self.layers.max_layers))
# electron fraction is already included by multiplying the densities with
# them in the Layers module, so we pass 1. to nuSQuIDS (unless energies are
# very high, this should be equivalent).
ye = np.broadcast_to(np.array([1.]),
(len(e_nodes), self.layers.max_layers))
self.nus_layer = nsq.nuSQUIDSLayers(
distances * nsq_units.km,
densities,
ye,
e_nodes * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.both,
)
self.apply_prop_settings(self.nus_layer)
# Now that we have our nusquids calculator set up on the node grid, we make
# container output space for the probability output which may be on a finer grid
# than the nodes or even working in events mode.
self.data.representation = self.calc_mode
# --- calculate the layers ---
if isinstance(self.calc_mode, MultiDimBinning):
# 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"
])
# calculate the distance difference between minimum and maximum production
# height, if applicable
if self.avg_height:
layers_min = Layers(earth_model, detector_depth,
self.prop_height_min)
layers_min.setElecFrac(self.YeI, self.YeO, self.YeM)
for container in self.data:
self.layers.calcLayers(container["true_coszen"])
distances = self.layers.distance.reshape(
(container.size, self.layers.max_layers))
tot_distances = np.sum(distances, axis=1)
if self.avg_height:
layers_min.calcLayers(container["true_coszen"])
dists_min = layers_min.distance.reshape(
(container.size, self.layers.max_layers))
min_tot_dists = np.sum(dists_min, axis=1)
# nuSQuIDS assumes the original distance is the longest distance and
# the averaging range is the difference between the minimum and maximum
# distance.
avg_ranges = tot_distances - min_tot_dists
assert np.all(avg_ranges > 0)
if isinstance(self.node_mode,
MultiDimBinning) and not self.exact_mode:
# To project out probabilities we only need the *total* distance
container["tot_distances"] = tot_distances
# for the binned node_mode we already calculated layers above
if self.avg_height:
container["avg_ranges"] = avg_ranges
elif self.node_mode == "events" or self.exact_mode:
# in any other mode (events or exact) we store all densities and
# distances in the container in calc_specs
densities = self.layers.density.reshape(
(container.size, self.layers.max_layers))
container["densities"] = densities
container["distances"] = distances
self.data.unlink_containers()
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar",
["nutaubar_cc", "nutaubar_nc"])
# 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)
self.data.unlink_containers()
if self.exact_mode: return
# --- containers for interpolated states ---
# This is not needed in exact mode
if isinstance(self.calc_mode, MultiDimBinning):
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["interp_states_e"] = np.empty(
(container.size, self.num_neutrinos**2),
dtype=FTYPE,
)
container["interp_states_mu"] = np.empty(
(container.size, self.num_neutrinos**2),
dtype=FTYPE,
)
self.data.unlink_containers()
# @line_profile
def calc_node_probs(self, nus_layer, flav_in, flav_out, n_nodes):
"""
Evaluate oscillation probabilities at nodes. This does not require any
interpolation.
"""
ini_state = np.array([0] * self.num_neutrinos)
ini_state[flav_in] = 1
nus_layer.Set_initial_state(ini_state, nsq.Basis.flavor)
nus_layer.EvolveState()
prob_nodes = nus_layer.EvalFlavorAtNodes(flav_out)
return prob_nodes
def calc_interpolated_states(self, evolved_states, e_out, cosz_out):
"""
Calculate interpolated states at the energies and zenith angles requested.
"""
nsq_units = nsq.Const()
interp_states = np.zeros((e_out.size, evolved_states.shape[1]))
assert np.all(e_out <= np.max(self.e_node_mode * nsq_units.GeV))
assert np.all(e_out >= np.min(self.e_node_mode * nsq_units.GeV))
assert np.all(cosz_out <= np.max(self.coszen_node_mode))
assert np.all(cosz_out >= np.min(self.coszen_node_mode))
for i in range(evolved_states.shape[1]):
z = evolved_states[:, i].reshape(self.e_mesh.shape).T
assert np.all(np.isfinite(z))
# RectBivariateSpline takes in the 1D node position and assumes that they
# are on a mesh.
f = RectBivariateSpline(
np.log10(self.e_node_mode * nsq_units.GeV),
self.coszen_node_mode,
z,
kx=2,
ky=2,
)
interp_states[..., i] = f(np.log10(e_out), cosz_out, grid=False)
return interp_states
def calc_probs_interp(self,
flav_out,
nubar,
interp_states,
out_distances,
e_out,
avg_ranges=0):
"""
Project out probabilities from interpolated interaction picture states.
"""
nsq_units = nsq.Const()
prob_interp = np.zeros(e_out.size)
scale = self.eval_lowpass_frac * self.eval_lowpass_cutoff / nsq_units.km
prob_interp = self.nus_layer.EvalWithState(
flav_out,
out_distances,
e_out,
interp_states,
avr_scale=0.,
rho=int(nubar),
lowpass_cutoff=self.eval_lowpass_cutoff / nsq_units.km,
lowpass_scale=scale,
t_range=avg_ranges)
return prob_interp
def compute_function_no_interpolation(self):
"""
Version of the compute function that does not use any interpolation between
nodes.
"""
nsq_units = nsq.Const()
# it is possible to work in binned calc mode while being in exact mode
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar",
["nutaubar_cc", "nutaubar_nc"])
for container in self.data:
nubar = container["nubar"] < 0
flav = container["flav"]
# electron fraction is already included by multiplying the densities
# with them in the Layers module, so we pass 1. to nuSQuIDS (unless
# energies are very high, this should be equivalent).
ye = np.broadcast_to(np.array([1.]),
(container.size, self.layers.max_layers))
nus_layer = nsq.nuSQUIDSLayers(
container["distances"] * nsq_units.km,
container["densities"],
ye,
container["true_energy"] * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.antineutrino
if nubar else nsq.NeutrinoType.neutrino,
)
self.apply_prop_settings(nus_layer)
self.set_osc_parameters(nus_layer)
container["prob_e"] = self.calc_node_probs(nus_layer, 0, flav,
container.size)
container["prob_mu"] = self.calc_node_probs(
nus_layer, 1, flav, container.size)
container.mark_changed("prob_e")
container.mark_changed("prob_mu")
self.data.unlink_containers()
@profile
def compute_function_interpolated(self):
"""
Version of the compute function that does use interpolation between nodes.
"""
nsq_units = nsq.Const()
# We need to make two evolutions, one for numu and the other for nue.
# These produce neutrino and antineutrino states at the same time thanks to
# the "both" neutrino mode of nuSQuIDS.
self.apply_prop_settings(self.nus_layer)
self.set_osc_parameters(self.nus_layer)
ini_state_nue = np.array([1, 0] + [0] * (self.num_neutrinos - 2))
ini_state_numu = np.array([0, 1] + [0] * (self.num_neutrinos - 2))
self.nus_layer.Set_initial_state(ini_state_nue, nsq.Basis.flavor)
self.nus_layer.EvolveState()
evolved_states_nue = self.nus_layer.GetStates(0)
evolved_states_nuebar = self.nus_layer.GetStates(1)
self.nus_layer.Set_initial_state(ini_state_numu, nsq.Basis.flavor)
self.nus_layer.EvolveState()
evolved_states_numu = self.nus_layer.GetStates(0)
evolved_states_numubar = self.nus_layer.GetStates(1)
# Now comes the step where we interpolate the interaction picture states
# and project out oscillation probabilities. This can be done in either events
# or binned mode.
if isinstance(self.calc_mode, MultiDimBinning):
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:
nubar = container["nubar"] < 0
container["interp_states_e"] = self.calc_interpolated_states(
evolved_states_nuebar if nubar else evolved_states_nue,
container["true_energy"] * nsq_units.GeV,
container["true_coszen"])
container["interp_states_mu"] = self.calc_interpolated_states(
evolved_states_numubar if nubar else evolved_states_numu,
container["true_energy"] * nsq_units.GeV,
container["true_coszen"])
self.data.unlink_containers()
if isinstance(self.calc_mode, MultiDimBinning):
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar",
["nutaubar_cc", "nutaubar_nc"])
for container in self.data:
nubar = container["nubar"] < 0
flav_out = container["flav"]
for flav_in in ["e", "mu"]:
container["prob_" + flav_in] = self.calc_probs_interp(
flav_out, nubar, container["interp_states_" + flav_in],
container["tot_distances"] * nsq_units.km,
container["true_energy"] * nsq_units.GeV,
container["avg_ranges"] *
nsq_units.km if self.avg_height else 0.)
container.mark_changed("prob_e")
container.mark_changed("prob_mu")
self.data.unlink_containers()
def compute_function(self):
if self.node_mode == "events" or self.exact_mode:
self.compute_function_no_interpolation()
else:
self.compute_function_interpolated()
@profile
def apply_function(self):
for container in self.data:
scales = container['nu_flux'][:, 0] * container[
'prob_e'] + container['nu_flux'][:, 1] * container['prob_mu']
container['weights'] = container["weights"] * scales
def setup_function(self):
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height
detector_depth = self.detector_depth
self.layers = Layers(earth_model, detector_depth, prop_height)
self.layers.setElecFrac(self.YeI, self.YeO, self.YeM)
nsq_units = nsq.Const() # natural units for nusquids
# Because we don't want to extrapolate, we check that all points at which we
# want to evaluate probabilities are fully contained within the node specs. This
# is of course not necessary in events mode.
if self.node_mode == "binned" and not self.exact_mode:
logging.debug("setting up nuSQuIDS nodes in binned mode")
# we can prepare the calculator like this only in binned mode, see
# compute_function for node_mode == "events"
self.data.data_specs = self.calc_specs
for container in self.data:
for var in ["true_coszen", "true_energy"]:
upper_bound = np.max(self.node_specs[var].bin_edges)
lower_bound = np.min(self.node_specs[var].bin_edges)
err_msg = (
"The outer edges of the node_specs must encompass "
"the entire range of calc_specs to avoid extrapolation"
)
if np.any(container[var].get(WHERE) > upper_bound):
maxval = np.max(container[var].get(WHERE))
raise ValueError(err_msg +
f"\nmax input: {maxval}, upper "
f"bound: {upper_bound}")
if np.any(container[var].get(WHERE) < lower_bound):
minval = np.max(container[var].get(WHERE))
raise ValueError(err_msg +
f"\nmin input: {minval}, lower "
f"bound: {lower_bound}")
# Layers in nuSQuIDS are special: We need all the individual distances and
# densities for the nodes to solve the interaction picture states, but on
# the final calculation grid (or events) we only need the *total* traversed
# distance. Because we are placing nodes at the bin edges rather than the
# bin middle, this doesn't really fit with how containers store data, so we
# are making arrays as variables that never go into the container.
# These are stored because we need them later during interpolation
self.coszen_node_specs = self.node_specs[
"true_coszen"].bin_edges.m_as("dimensionless")
self.e_node_specs = self.node_specs["true_energy"].bin_edges.m_as(
"GeV")
logging.debug(f"Setting up nodes at\n"
f"cos_zen = \n{self.coszen_node_specs}\n"
f"energy = \n{self.e_node_specs}\n")
# things are getting a bit meshy from here...
self.e_mesh, self.cosz_mesh = np.meshgrid(self.e_node_specs,
self.coszen_node_specs)
e_nodes = self.e_mesh.ravel()
coszen_nodes = self.cosz_mesh.ravel()
# The lines below should not be necessary because we will always get at
# least two numbers from the bin edges. However, if either energy or coszen
# somehow was just a scalar, we would need to broadcast it out to the same
# size. Keeping the code in here in case you want to use the stage in 1D.
# convert lists to ndarrays and scalars to ndarrays with length 1
e_nodes = np.atleast_1d(e_nodes)
coszen_nodes = np.atleast_1d(coszen_nodes)
# broadcast against each other and make a copy
# (see https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html)
e_nodes, coszen_nodes = [
np.array(a)
for a in np.broadcast_arrays(e_nodes, coszen_nodes)
]
assert len(e_nodes) == len(coszen_nodes)
assert coszen_nodes.ndim == 1
assert e_nodes.ndim == 1
self.layers.calcLayers(coszen_nodes)
distances = np.reshape(self.layers.distance,
(len(e_nodes), self.layers.max_layers))
densities = np.reshape(self.layers.density,
(len(e_nodes), self.layers.max_layers))
# electron fraction is already included by multiplying the densities with
# them in the Layers module, so we pass 1. to nuSQuIDS (unless energies are
# very high, this should be equivalent).
ye = np.broadcast_to(np.array([1.]),
(len(e_nodes), self.layers.max_layers))
self.nus_layer = nsq.nuSQUIDSLayers(
distances * nsq_units.km,
densities,
ye,
e_nodes * nsq_units.GeV,
self.num_neutrinos,
nsq.NeutrinoType.both,
)
self.apply_prop_settings(self.nus_layer)
# Now that we have our nusquids calculator set up on the node grid, we make
# container output space for the probability output which may be on a finer grid
# than the nodes or even working in events mode.
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == "binned":
# 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"
])
# calculate the distance difference between minimum and maximum production
# height, if applicable
if self.avg_height:
layers_min = Layers(earth_model, detector_depth,
self.prop_height_min)
layers_min.setElecFrac(self.YeI, self.YeO, self.YeM)
for container in self.data:
self.layers.calcLayers(container["true_coszen"].get("host"))
distances = self.layers.distance.reshape(
(container.size, self.layers.max_layers))
tot_distances = np.sum(distances, axis=1)
if self.avg_height:
layers_min.calcLayers(container["true_coszen"].get("host"))
dists_min = layers_min.distance.reshape(
(container.size, self.layers.max_layers))
min_tot_dists = np.sum(dists_min, axis=1)
# nuSQuIDS assumes the original distance is the longest distance and
# the averaging range is the difference between the minimum and maximum
# distance.
avg_ranges = tot_distances - min_tot_dists
assert np.all(avg_ranges > 0)
if self.node_mode == "binned" and not self.exact_mode:
# To project out probabilities we only need the *total* distance
container["tot_distances"] = tot_distances
# for the binned node_mode we already calculated layers above
if self.avg_height:
container["avg_ranges"] = avg_ranges
elif self.node_mode == "events" or self.exact_mode:
# in any other mode (events or exact) we store all densities and
# distances in the container in calc_specs
densities = self.layers.density.reshape(
(container.size, self.layers.max_layers))
container["densities"] = densities
container["distances"] = distances
self.data.unlink_containers()
if self.calc_mode == "binned":
self.data.link_containers("nue", ["nue_cc", "nue_nc"])
self.data.link_containers("numu", ["numu_cc", "numu_nc"])
self.data.link_containers("nutau", ["nutau_cc", "nutau_nc"])
self.data.link_containers("nuebar", ["nuebar_cc", "nuebar_nc"])
self.data.link_containers("numubar", ["numubar_cc", "numubar_nc"])
self.data.link_containers("nutaubar",
["nutaubar_cc", "nutaubar_nc"])
# 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)
self.data.unlink_containers()
if self.exact_mode: return
# --- containers for interpolated states ---
# This is not needed in exact mode
if self.calc_mode == "binned":
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["interp_states_e"] = np.empty(
(container.size, self.num_neutrinos**2),
dtype=FTYPE,
)
container["interp_states_mu"] = np.empty(
(container.size, self.num_neutrinos**2),
dtype=FTYPE,
)
self.data.unlink_containers()
class globes(Stage):
"""
GLoBES PISA Pi class
Parameters
----------
earth_model : PREM file path
globes_wrapper : path to globes wrapper
detector_depth : float
prop_height : quantity (dimensionless)
params : ParamSet or sequence with which to instantiate a ParamSet.
Expected params .. ::
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltam41 : quantity (mass^2)
theta24 : quantity (angle)
theta34 : quantity (angle)
deltacp : quantity (angle)
"""
def __init__(
self,
earth_model,
globes_wrapper,
detector_depth=2.*ureg.km,
prop_height=20.*ureg.km,
**std_kwargs,
):
expected_params = (
'theta12',
'theta13',
'theta23',
'deltam21',
'deltam31',
'deltam41',
'theta24',
'theta34',
'deltacp',
)
# init base class
super().__init__(
expected_params=expected_params,
**std_kwargs,
)
self.layers = None
self.osc_params = None
self.earth_model = earth_model
self.globes_wrapper = globes_wrapper
self.detector_depth = detector_depth
self.prop_height = prop_height
self.globes_calc = None
@profile
def setup_function(self):
sys.path.append(self.globes_wrapper)
import GLoBES
### you need to start GLoBES from the folder containing a dummy experiment
# therefore we go to the folder, load GLoBES and then go back
curdir = os.getcwd()
os.chdir(self.globes_wrapper)
self.globes_calc = GLoBES.GLoBESCalculator("calc")
os.chdir(curdir)
self.globes_calc.InitSteriles(2)
# object for oscillation parameters
self.osc_params = OscParams()
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height.m_as('km')
detector_depth = self.detector_depth.m_as('km')
self.layers = Layers(earth_model, detector_depth, prop_height)
# The electron fractions are taken into account internally by GLoBES/SNU.
# See the SNU patch for details. It uses the density to decide
# whether it is in the core or in the mantle. Therefore, we just multiply by
# one to give GLoBES the raw densities.
self.layers.setElecFrac(1., 1., 1.)
# set the correct data mode
self.data.representation = self.calc_mode
# --- calculate the layers ---
if self.data.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 probability containers
for container in self.data:
container['prob_e'] = np.empty((container.size), dtype=FTYPE)
container['prob_mu'] = np.empty((container.size), dtype=FTYPE)
container['prob_nonsterile'] = np.empty((container.size), dtype=FTYPE)
if '_cc' in container.name:
container['prob_nonsterile'] = np.ones(container.size)
elif '_nc' in container.name:
if 'nue' in container.name:
container['prob_e'] = np.ones(container.size)
container['prob_mu'] = np.zeros(container.size)
elif 'numu' in container.name:
container['prob_e'] = np.zeros(container.size)
container['prob_mu'] = np.ones(container.size)
elif 'nutau' in container.name:
container['prob_e'] = np.zeros(container.size)
container['prob_mu'] = np.zeros(container.size)
else:
raise Exception('unknown container name: %s' % container.name)
def calc_prob_e_mu(self, flav, nubar, energy, rho_array, len_array):
'''Calculates probability for an electron/muon neutrino to oscillate into
the flavour of a given event, including effects from sterile neutrinos.
'''
# We use the layers module to calculate lengths and densities.
# The output must be converted into a regular python list.
self.globes_calc.SetManualDensities(list(len_array), list(rho_array))
# this calls the calculator without the calculation of layers
# The flavour convention in GLoBES is that
# e = 1, mu = 2, tau = 3
# while in PISA it's
# e = 0, mu = 1, tau = 2
# which is why we add +1 to the flavour.
# Nubar follows the same convention in PISA and GLoBES:
# +1 = particle, -1 = antiparticle
nue_to_nux = self.globes_calc.MatterProbabilityPrevBaseline(1, flav+1, nubar, energy)
numu_to_nux = self.globes_calc.MatterProbabilityPrevBaseline(2, flav+1, nubar, energy)
return (nue_to_nux, numu_to_nux)
def calc_prob_nonsterile(self, flav, nubar, energy, rho_array, len_array):
'''Calculates the probability of a given neutrino to oscillate into
another non-sterile flavour.
'''
# We use the layers module to calculate lengths and densities.
# The output must be converted into a regular python list.
self.globes_calc.SetManualDensities(list(len_array), list(rho_array))
# this calls the calculator without the calculation of layers
# The flavour convention in GLoBES is that
# e = 1, mu = 2, tau = 3
# while in PISA it's
# e = 0, mu = 1, tau = 2
# which is why we add +1 to the flavour.
# Nubar follows the same convention in PISA and GLoBES:
# +1 = particle, -1 = antiparticle
nux_to_nue = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 1, nubar, energy)
nux_to_numu = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 2, nubar, energy)
nux_to_nutau = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 3, nubar, energy)
nux_to_nonsterile = nux_to_nue + nux_to_numu + nux_to_nutau
return nux_to_nonsterile
@profile
def compute_function(self):
# --- update mixing params ---
params = [self.params.theta12.value.m_as('rad'),
self.params.theta13.value.m_as('rad'),
self.params.theta23.value.m_as('rad'),
self.params.deltacp.value.m_as('rad'),
self.params.deltam21.value.m_as('eV**2'),
self.params.deltam31.value.m_as('eV**2'),
self.params.deltam41.value.m_as('eV**2'),
0.0,
self.params.theta24.value.m_as('rad'),
self.params.theta34.value.m_as('rad'),
0.0,
0.0
]
self.globes_calc.SetParametersArr(params)
# set the correct data mode
self.data.representation = self.calc_mode
for container in self.data:
# standard oscillations are only applied to charged current events,
# while the loss due to oscillation into sterile neutrinos is only
# applied to neutral current events.
# Accessing single entries from containers is very slow.
# For this reason, we make a copy of the content we need that is
# a simple numpy array.
flav = container['flav']
nubar = container['nubar']
energies = np.array(container['true_energy'])
densities = np.array(container['densities'])
distances = np.array(container['distances'])
prob_e = np.zeros(container.size)
prob_mu = np.zeros(container.size)
prob_nonsterile = np.zeros(container.size)
if '_cc' in container.name:
for i in range(container.size):
prob_e[i], prob_mu[i] = self.calc_prob_e_mu(flav,
nubar,
energies[i],
densities[i],
distances[i]
)
container['prob_e'] = prob_e
container['prob_mu'] = prob_mu
elif '_nc' in container.name:
for i in range(container.size):
prob_nonsterile[i] = self.calc_prob_nonsterile(flav,
nubar,
energies[i],
densities[i],
distances[i]
)
container['prob_nonsterile'] = prob_nonsterile
else:
raise Exception('unknown container name: %s' % container.name)
container.mark_changed('prob_e')
container.mark_changed('prob_mu')
container.mark_changed('prob_nonsterile')
@profile
def apply_function(self):
# update the outputted weights
for container in self.data:
apply_probs(container['nu_flux'],
container['prob_e'],
container['prob_mu'],
container['prob_nonsterile'],
out=container['weights'])
container.mark_changed('weights')
class pi_prob3(PiStage):
"""
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)
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)
**kwargs
Other kwargs are handled by PiStage
-----
"""
def __init__(
self,
nsi_type=None,
reparam_mix_matrix=False,
data=None,
params=None,
input_names=None,
output_names=None,
debug_mode=None,
input_specs=None,
calc_specs=None,
output_specs=None,
):
expected_params = (
'detector_depth',
'earth_model',
'prop_height',
'YeI',
'YeO',
'YeM',
'theta12',
'theta13',
'theta23',
'deltam21',
'deltam31',
'deltacp'
)
# 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."""
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'
)
expected_params = expected_params + nsi_params
input_names = ()
output_names = ()
# what are the keys used from the inputs during apply
input_apply_keys = ('weights', 'nu_flux')
# what are keys added or altered in the calculation used during apply
output_calc_keys = ('prob_e', 'prob_mu')
# what keys are added or altered for the outputs during apply
output_apply_keys = ('weights',)
# init base class
super().__init__(
data=data,
params=params,
expected_params=expected_params,
input_names=input_names,
output_names=output_names,
debug_mode=debug_mode,
input_specs=input_specs,
calc_specs=calc_specs,
output_specs=output_specs,
input_apply_keys=input_apply_keys,
output_calc_keys=output_calc_keys,
output_apply_keys=output_apply_keys,
)
assert self.input_mode is not None
assert self.calc_mode is not None
assert self.output_mode is not None
self.layers = None
self.osc_params = None
self.nsi_params = 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()
# 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)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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'].get('host'))
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.calc_mode == 'binned':
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
propagate_array(self.osc_params.dm_matrix, # pylint: disable = unexpected-keyword-arg, no-value-for-parameter
mix_matrix,
self.gen_mat_pot_matrix_complex,
nubar,
e_array.get(WHERE),
rho_array.get(WHERE),
len_array.get(WHERE),
out=out.get(WHERE)
)
out.mark_changed(WHERE)
@profile
def compute_function(self):
# set the correct data mode
self.data.data_specs = self.calc_specs
if self.calc_mode == 'binned':
# 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'].get('host'))
container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers))
container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers))
# 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')
# 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)
for container in self.data:
self.calc_probs(container['nubar'],
container['true_energy'],
container['densities'],
container['distances'],
out=container['probability'],
)
# the following is flavour specific, hence unlink
self.data.unlink_containers()
for container in self.data:
# initial electrons (0)
fill_probs(container['probability'].get(WHERE),
0,
container['flav'],
out=container['prob_e'].get(WHERE),
)
# initial muons (1)
fill_probs(container['probability'].get(WHERE),
1,
container['flav'],
out=container['prob_mu'].get(WHERE),
)
container['prob_e'].mark_changed(WHERE)
container['prob_mu'].mark_changed(WHERE)
@profile
def apply_function(self):
# update the outputted weights
for container in self.data:
apply_probs(container['nu_flux'].get(WHERE),
container['prob_e'].get(WHERE),
container['prob_mu'].get(WHERE),
out=container['weights'].get(WHERE))
container['weights'].mark_changed(WHERE)
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()
# 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)
class pi_globes(PiStage):
"""
GLoBES PISA Pi class
Paramaters
----------
detector_depth : float
earth_model : PREM file path
globes_wrapper : path to globes wrapper
prop_height : quantity (dimensionless)
params : ParamSet or sequence with which to instantiate a ParamSet.
Expected params are:
theta12 : quantity (angle)
theta13 : quantity (angle)
theta23 : quantity (angle)
deltam21 : quantity (mass^2)
deltam31 : quantity (mass^2)
deltam41 : quantity (mass^2)
theta24 : quantity (angle)
theta34 : quantity (angle)
deltacp : quantity (angle)
Notes
-----
"""
def __init__(self,
earth_model,
globes_wrapper,
data=None,
params=None,
input_names=None,
output_names=None,
debug_mode=None,
input_specs=None,
calc_specs=None,
output_specs=None,
detector_depth=2.*ureg.km,
prop_height=20.*ureg.km
):
expected_params = (
'theta12',
'theta13',
'theta23',
'deltam21',
'deltam31',
'deltam41',
'theta24',
'theta34',
'deltacp',
)
input_names = ()
output_names = ()
# what are the keys used from the inputs during apply
input_apply_keys = ('weights',
'sys_flux',
)
# what are keys added or altered in the calculation used during apply
output_calc_keys = ('prob_e',
'prob_mu',
'prob_nonsterile',
)
# what keys are added or altered for the outputs during apply
output_apply_keys = ('weights',
)
# init base class
super().__init__(
data=data,
params=params,
expected_params=expected_params,
input_names=input_names,
output_names=output_names,
debug_mode=debug_mode,
input_specs=input_specs,
calc_specs=calc_specs,
output_specs=output_specs,
input_apply_keys=input_apply_keys,
output_calc_keys=output_calc_keys,
output_apply_keys=output_apply_keys,
)
assert self.input_mode is not None
assert self.calc_mode is not None
assert self.output_mode is not None
self.layers = None
self.osc_params = None
self.earth_model = earth_model
self.globes_wrapper = globes_wrapper
self.detector_depth = detector_depth
self.prop_height = prop_height
self.globes_calc = None
# This does nothing for speed, but just allows one to use numpy broadcasting on the
# arguments of the function. The internal implementation is basically a for-loop.
# The signature is chosen so that the function for a single event expects an array
# of n layer distances and densities.
self.calc_prob_e_mu = np.vectorize(self.calc_prob_e_mu, signature='(),(),(),(n),(n)->(),()')
self.calc_prob_nonsterile = np.vectorize(self.calc_prob_nonsterile,
signature='(),(),(),(n),(n)->()')
def setup_function(self):
sys.path.append(self.globes_wrapper)
import GLoBES
### you need to start GLoBES from the folder containing a dummy experiment
# therefore we go to the folder, load GLoBES and then go back
curdir = os.getcwd()
os.chdir(self.globes_wrapper)
self.globes_calc = GLoBES.GLoBESCalculator("calc")
os.chdir(curdir)
self.globes_calc.InitSteriles(2)
# object for oscillation parameters
self.osc_params = OscParams()
earth_model = find_resource(self.earth_model)
prop_height = self.prop_height.m_as('km')
detector_depth = self.detector_depth.m_as('km')
self.layers = Layers(earth_model, detector_depth, prop_height)
# The electron fractions are taken into account internally by GLoBES/SNU.
# See the SNU patch for details. It uses the density to decide
# whether it is in the core or in the mantle. Therefore, we just multiply by
# one to give GLoBES the raw densities.
self.layers.setElecFrac(1., 1., 1.)
# set the correct data mode
self.data.data_specs = self.calc_specs
# --- calculate the layers ---
if self.calc_mode == 'binned':
# 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'].get('host'))
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 probability containers
for container in self.data:
container['prob_e'] = np.empty((container.size), dtype=FTYPE)
container['prob_mu'] = np.empty((container.size), dtype=FTYPE)
container['prob_nonsterile'] = np.empty((container.size), dtype=FTYPE)
def calc_prob_e_mu(self, flav, nubar, energy, rho_array, len_array):
'''Calculates probability for an electron/muon neutrino to oscillate into
the flavour of a given event, including effects from sterile neutrinos.
'''
# We use the layers module to calculate lengths and densities.
# The output must be converted into a regular python list.
self.globes_calc.SetManualDensities(list(len_array), list(rho_array))
# this calls the calculator without the calculation of layers
# The flavour convention in GLoBES is that
# e = 1, mu = 2, tau = 3
# while in PISA it's
# e = 0, mu = 1, tau = 2
# which is why we add +1 to the flavour.
# Nubar follows the same convention in PISA and GLoBES:
# +1 = particle, -1 = antiparticle
nue_to_nux = self.globes_calc.MatterProbabilityPrevBaseline(1, flav+1, nubar, energy)
numu_to_nux = self.globes_calc.MatterProbabilityPrevBaseline(2, flav+1, nubar, energy)
return (nue_to_nux, numu_to_nux)
def calc_prob_nonsterile(self, flav, nubar, energy, rho_array, len_array):
'''Calculates the probability of a given neutrino to oscillate into
another non-sterile flavour.
'''
# We use the layers module to calculate lengths and densities.
# The output must be converted into a regular python list.
self.globes_calc.SetManualDensities(list(len_array), list(rho_array))
# this calls the calculator without the calculation of layers
# The flavour convention in GLoBES is that
# e = 1, mu = 2, tau = 3
# while in PISA it's
# e = 0, mu = 1, tau = 2
# which is why we add +1 to the flavour.
# Nubar follows the same convention in PISA and GLoBES:
# +1 = particle, -1 = antiparticle
nux_to_nue = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 1, nubar, energy)
nux_to_numu = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 2, nubar, energy)
nux_to_nutau = self.globes_calc.MatterProbabilityPrevBaseline(flav+1, 3, nubar, energy)
nux_to_nonsterile = nux_to_nue + nux_to_numu + nux_to_nutau
return nux_to_nonsterile
@profile
def compute_function(self):
# --- update mixing params ---
params = np.array([self.params.theta12.value.m_as('rad'),
self.params.theta13.value.m_as('rad'),
self.params.theta23.value.m_as('rad'),
self.params.deltacp.value.m_as('rad'),
self.params.deltam21.value.m_as('eV**2'),
self.params.deltam31.value.m_as('eV**2'),
self.params.deltam41.value.m_as('eV**2'),
0.0,
self.params.theta24.value.m_as('rad'),
self.params.theta34.value.m_as('rad'),
0.0,
0.0
], dtype=float)
self.globes_calc.SetParametersArr(params)
# set the correct data mode
self.data.data_specs = self.calc_specs
for container in self.data:
# standard oscillations are only applied to charged current events,
# while the loss due to oscillation into sterile neutrinos is only
# applied to neutral current events.
if '_cc' in container.name:
prob_e, prob_mu = self.calc_prob_e_mu(container['flav'],
container['nubar'],
container['true_energy'],
container['densities'],
container['distances']
)
container['prob_e'] = prob_e
container['prob_mu'] = prob_mu
container['prob_nonsterile'] = np.ones_like(prob_e)
elif '_nc' in container.name:
prob_nonsterile = self.calc_prob_nonsterile(container['flav'],
container['nubar'],
container['true_energy'],
container['densities'],
container['distances']
)
if 'nue' in container.name:
container['prob_e'] = np.ones_like(prob_nonsterile)
container['prob_mu'] = np.zeros_like(prob_nonsterile)
elif 'numu' in container.name:
container['prob_e'] = np.zeros_like(prob_nonsterile)
container['prob_mu'] = np.ones_like(prob_nonsterile)
elif 'nutau' in container.name:
container['prob_e'] = np.zeros_like(prob_nonsterile)
container['prob_mu'] = np.zeros_like(prob_nonsterile)
else:
raise Exception('unknown container name: %s' % container.name)
container['prob_nonsterile'] = prob_nonsterile
else:
raise Exception('unknown container name: %s' % container.name)
container['prob_e'].mark_changed(WHERE)
container['prob_mu'].mark_changed(WHERE)
container['prob_nonsterile'].mark_changed(WHERE)
@profile
def apply_function(self):
# update the outputted weights
for container in self.data:
apply_probs(container['sys_flux'].get(WHERE),
container['prob_e'].get(WHERE),
container['prob_mu'].get(WHERE),
container['prob_nonsterile'].get(WHERE),
out=container['weights'].get(WHERE))
container['weights'].mark_changed(WHERE)
