Exemplo n.º 1
0
    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()
Exemplo n.º 2
0
    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()
Exemplo n.º 3
0
    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()