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 = 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 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 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()