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