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 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 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 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 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()
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()
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"))
from cross_section_test import get_diff_xs
import nuSQUIDSpy as nsq
# specialty-made utility functions
from nus_utils import get_flavor, get_neut, get_curr
from utils import bhist, get_exp_std, get_width, get_nearest_entry_to
from utils import Data, get_index, get_loc, sci
# tau stuff
from tau_funcs import TauData
# reconstruction data
from deporeco import DataReco
const = nsq.Const()
# colormap thing
cmap = plt.get_cmap('coolwarm')
n_colors = 6
def get_color(which, how_many=n_colors):
return (cmap(float(which) / how_many))
# load the data using the default filename, 'atmosphere.txt'
data = Data()
tauData = TauData()
scale_e = np.array(data.energies)
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 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
from pisa.stages.osc.pi_osc_params import OscParams
from pisa.utils.fileio import from_file
from pisa.utils.log import logging
__all__ = [
# constants
'NSI_AVAIL', 'DECOH_AVAIL', 'ATM_AVAIL', 'NSI_CLASS', 'DECOH_CLASS', 'ATM_CLASS', 'NSQ_CONST',
# functions
'validate_calc_grid', 'compute_binning_constants', 'init_nusquids_prop',
'evolve_states', 'osc_probs', 'earth_model'
]
__version__ = '0.1'
NSQ_CONST = nsq.Const()
PRIMARIES = ['numu', 'numubar', 'nue', 'nuebar']
# flavor codes used internally by nuSQuIDS
FLAV_INDS = {'nue': 0, 'nuebar': 0, 'numu': 1, 'numubar': 1, 'nutau': 2,
'nutaubar': 2}
FLAV_INDS = OrderedDict(sorted(FLAV_INDS.items(), key=lambda t: t[0]))
def validate_calc_grid(calc_grid):
"""Check whether a multi-dimensional binning is suitable for use as
the grid on which oscillations are calculated for event-by-event
reweighting."""
calc_grid = MultiDimBinning(calc_grid)
dim_names = set(calc_grid.names)
if not dim_names == set(['true_energy', 'true_coszen']):
raise ValueError('Oscillation grid must contain "true_energy" and'
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 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()
def _compute_outputs(self, inputs=None):
"""Compute histograms for output channels."""
logging.debug('Entering nusquids._compute_outputs')
if not isinstance(inputs, MapSet):
raise AssertionError('inputs is not a MapSet object, instead '
'is type {0}'.format(type(inputs)))
# TODO(shivesh): oversampling
# TODO(shivesh): more options
# TODO(shivesh): static function
# TODO(shivesh): hashing
binning = self.input_binning.basename_binning
binning = binning.reorder_dimensions(('coszen', 'energy'),
use_basenames=True)
cz_binning = binning['coszen']
en_binning = binning['energy']
units = nsq.Const()
interactions = False
cz_min = cz_binning.bin_edges.min().m_as('radian')
cz_max = cz_binning.bin_edges.max().m_as('radian')
en_min = en_binning.bin_edges.min().m_as('GeV') * units.GeV
en_max = en_binning.bin_edges.max().m_as('GeV') * units.GeV
cz_centers = cz_binning.weighted_centers.m_as('radian')
en_centers = en_binning.weighted_centers.m_as('GeV') * units.GeV
cz_grid = np.array([cz_min] + cz_centers.tolist() + [cz_max])
en_grid = np.array([en_min] + en_centers.tolist() + [en_max])
nu_flavours = 3
nuSQ = nsq.nuSQUIDSAtm(cz_grid, en_grid, nu_flavours,
nsq.NeutrinoType.both, interactions)
nuSQ.Set_EvalThreads(multiprocessing.cpu_count())
theta12 = self.params['theta12'].value.m_as('radian')
theta13 = self.params['theta13'].value.m_as('radian')
theta23 = self.params['theta23'].value.m_as('radian')
deltam21 = self.params['deltam21'].value.m_as('eV**2')
deltam31 = self.params['deltam21'].value.m_as('eV**2')
# TODO(shivesh): check if deltacp should be in radians
deltacp = self.params['deltacp'].value.m_as('radian')
nuSQ.Set_MixingAngle(0, 1, theta12)
nuSQ.Set_MixingAngle(0, 2, theta13)
nuSQ.Set_MixingAngle(1, 2, theta23)
nuSQ.Set_SquareMassDifference(1, deltam21)
nuSQ.Set_SquareMassDifference(2, deltam31)
nuSQ.Set_CPPhase(0, 2, deltacp)
nuSQ.Set_rel_error(1.0e-10)
nuSQ.Set_abs_error(1.0e-10)
# Pad the edges of the energy, coszen space to cover the entire grid range
cz_shape = cz_binning.shape[0] + 2
en_shape = en_binning.shape[0] + 2
shape = (cz_shape, en_shape) + (2, 3)
initial_state = np.full(shape, np.nan)
def pad_inputs(x):
return np.pad(unp.nominal_values(x.hist), 1, 'edge')
# Third index is selecting nu(0), nubar(1)
# Fourth index is selecting flavour nue(0), numu(1), nutau(2)
initial_state[:, :, 0, 0] = pad_inputs(inputs['nue'])
initial_state[:, :, 1, 0] = pad_inputs(inputs['nuebar'])
initial_state[:, :, 0, 1] = pad_inputs(inputs['numu'])
initial_state[:, :, 1, 1] = pad_inputs(inputs['numubar'])
initial_state[:, :, 0, 2] = np.zeros(pad_inputs(inputs['nue']).shape)
initial_state[:, :, 1, 2] = np.zeros(pad_inputs(inputs['nue']).shape)
if np.any(np.isnan(initial_state)):
raise AssertionError('nan entries in initial_state: '
'{0}'.format(initial_state))
nuSQ.Set_initial_state(initial_state, nsq.Basis.flavor)
# TODO(shivesh): use verbosity level to set this
nuSQ.Set_ProgressBar(True)
nuSQ.EvolveState()
os = self.params['oversample'].value.m
os_binning = binning.oversample(os)
os_cz_binning = os_binning['coszen']
os_en_binning = os_binning['energy']
os_cz_centers = os_cz_binning.weighted_centers.m_as('radians')
os_en_centers = os_en_binning.weighted_centers.m_as('GeV')
fs = {}
for nu in self.output_names:
fs[nu] = np.full(os_binning.shape, np.nan)
for icz, cz_bin in enumerate(os_cz_centers):
for ie, en_bin in enumerate(os_en_centers):
en_bin_u = en_bin * units.GeV
fs['nue'][icz][ie] = nuSQ.EvalFlavor(0, cz_bin, en_bin_u, 0)
fs['nuebar'][icz][ie] = nuSQ.EvalFlavor(0, cz_bin, en_bin_u, 1)
fs['numu'][icz][ie] = nuSQ.EvalFlavor(1, cz_bin, en_bin_u, 0)
fs['numubar'][icz][ie] = nuSQ.EvalFlavor(
1, cz_bin, en_bin_u, 1)
fs['nutau'][icz][ie] = nuSQ.EvalFlavor(2, cz_bin, en_bin_u, 0)
fs['nutaubar'][icz][ie] = nuSQ.EvalFlavor(
2, cz_bin, en_bin_u, 1)
out_binning = self.input_binning.reorder_dimensions(
('coszen', 'energy'), use_basenames=True)
os_out_binning = out_binning.oversample(os)
outputs = []
for key in fs.iterkeys():
if np.any(np.isnan(fs[key])):
raise AssertionError(
'Invalid value computed for {0} oscillated output: '
'\n{1}'.format(key, fs[key]))
map = Map(name=key, binning=os_out_binning, hist=fs[key])
map = map.downsample(os) / float(os)
map = map.reorder_dimensions(self.input_binning)
outputs.append(map)
return MapSet(outputs)
"""
This script will calcualte the DIS cross sections for neutrinos and convolve that with an expected flux in the ice
"""
import sys
import nuSQUIDSpy as nsq
import numpy as np # useful for energy ranges
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
constants = nsq.Const()
nBins = 100
eMin = 1. * constants.GeV
eMax = (10**12) * constants.GeV
energies = np.logspace(np.log10(eMin), np.log10(eMax), nBins)
# I use these dictionaries since the keys will be helpful in making plot labels
flavors = {
'electron': nsq.NeutrinoCrossSections_NeutrinoFlavor.electron,
'muon': nsq.NeutrinoCrossSections_NeutrinoFlavor.muon,
'tau': nsq.NeutrinoCrossSections_NeutrinoFlavor.tau
}
currents = {
'CC': nsq.NeutrinoCrossSections_Current.CC,
'NC': nsq.NeutrinoCrossSections_Current.NC
}
neut_types = {
# this notebook we will demostrate some of the functionalities of the $\nu$-SQuIDS' python bindings. All of the calculations performed here can also be done in the C++ interface. Enjoy :)! Carlos, Jordi & Chris. # # The Basics: single energy mode # #### Basic definitions # To start, like in the C++ case, we need to create a $\nu$-SQuIDS object. To begin this demonstration we will create a simple single energy three flavor neutrino oscillation calculator. Thus we just need to specify the number of neutrinos (3) and if we are dealing with neutrinos or antineutrinos. nuSQ = nsq.nuSQUIDS(3,nsq.NeutrinoType.neutrino) # nuSQuIDS inputs should be given in natural units. In order to make this convenient we have define a units class called *Const*. We can instanciate it as follows units = nsq.Const() # As in the C++ $\nu$-SQuIDS interface one can propagate the neutrinos in various environments (see the documentation for further details), and the user can create and include their own environments. To start a simple case, lets consider oscillactions in <strong> Vacuum </strong> nuSQ.Set_Body(nsq.Vacuum()) # Since we have specify that we are considering vacuum propagation, we must construct - as in the C++ interface - a *trayectory* inside that object. This can be done using the `Track` property of the given `Body`. Each `Body` will have its on `Track` subclass and its constructors. We can set and construct a <strong>vacuum trayectory</strong> in the following way: nuSQ.Set_Track(nsq.Vacuum.Track(100.0*units.km)) # Next we have to set the <strong>neutrino energy</strong>, which can be done as follows nuSQ.Set_E(1.0*units.GeV)
"""
This was used to make some plots comparing cascade and track flux rates
"""
from MCEq.core import MCEqRun
import crflux.models as crf
import numpy as np
import nuSQUIDSpy as nsq
from cascade.raw_fluxes import get_key
from cascade.utils import get_closest
from cascade.cross_section_test import xs_obj as xs
from cascade.cross_section_test import flavors, currents, neut_types, get_total_flux
un = nsq.Const()
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
def make_plots():
Emin = 1 * un.GeV
Emax = 10 * un.PeV
energies = nsq.logspace(Emin, Emax, 100)
angles = nsq.linspace(-0.99, -0.98, 2)
# just look at straight up/down
