From 1f28fb3347874b9c004f27502f383f91ad523f59 Mon Sep 17 00:00:00 2001 From: Philipp Eller Date: Sun, 28 Jul 2024 11:52:47 +0200 Subject: [PATCH] adding stage for external osc calc (#756) * adding stage for external osc calc * Minor corrections to the external osc stage (#759) * Compute neutron-weighted densities * Computing neutron densities and changing nubar boolean flag * Do not scale cross sections * Do not scale non-unitary cross sections (neither CC nor NC) by default * Remove SM flux/xsec flags so the stage is universally usable by all osc functions in neurthino * Update external.py small cleanup * trigger GitHub actions (squash later) --------- Co-authored-by: Tetiana Kozynets <37784144+kotania@users.noreply.github.com> Co-authored-by: Leander Fischer Co-authored-by: Jan Weldert <32642322+JanWeldert@users.noreply.github.com> --- pisa/stages/osc/external.py | 184 ++++++++++++++++++++++++++++++++++++ pisa/stages/osc/layers.py | 87 +++++++++++++++-- 2 files changed, 264 insertions(+), 7 deletions(-) create mode 100644 pisa/stages/osc/external.py diff --git a/pisa/stages/osc/external.py b/pisa/stages/osc/external.py new file mode 100644 index 000000000..9554b3bec --- /dev/null +++ b/pisa/stages/osc/external.py @@ -0,0 +1,184 @@ +""" +PISA pi stage for the calculation of earth layers and osc. probabilities +""" + +from __future__ import absolute_import, print_function, division + +import numpy as np + +from pisa import FTYPE, TARGET, ureg +from pisa.core.stage import Stage +from pisa.utils.log import logging +from pisa.utils.profiler import profile +from pisa.stages.osc.layers import Layers +from pisa.utils.resources import find_resource + +class external(Stage): + """ + Use an external function to calculate oscillation probabilities + + Parameters + ---------- + params + + osc_prob : callable + the external function + + Expected params .. :: + + detector_depth : float + earth_model : PREM file path + prop_height : quantity (dimensionless) + YeI : quantity (dimensionless) + YeO : quantity (dimensionless) + YeM : quantity (dimensionless) + **kwargs + Other kwargs are handled by Stage + ----- + """ + + def __init__( + self, + **std_kwargs, + ): + + expected_params = ( + 'detector_depth', + 'earth_model', + 'prop_height', + 'YeI', + 'YeO', + 'YeM', + ) + + + # init base class + super().__init__( + expected_params=expected_params, + **std_kwargs, + ) + + self.osc_prob = None + self.external_params = None + self.layers = None + self.YeI = None + self.YeO = None + self.YeM = None + + + def setup_function(self): + + # setup the layers + earth_model = find_resource(self.params.earth_model.value) + self.YeI = self.params.YeI.value.m_as('dimensionless') + self.YeO = self.params.YeO.value.m_as('dimensionless') + self.YeM = self.params.YeM.value.m_as('dimensionless') + prop_height = self.params.prop_height.value.m_as('km') + detector_depth = self.params.detector_depth.value.m_as('km') + self.layers = Layers(earth_model, detector_depth, prop_height) + self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) + + + # --- calculate the layers --- + if self.is_map: + # speed up calculation by adding links + # as layers don't care about flavour + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc', + 'nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + + for container in self.data: + self.layers.calcLayers(container['true_coszen']) + container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) + container['densities_neutron_weighted'] = self.layers.density_neutron_weighted.reshape((container.size, self.layers.max_layers)) + container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) + + # don't forget to un-link everything again + self.data.unlink_containers() + + # --- setup empty arrays --- + if self.is_map: + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc']) + self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + for container in self.data: + container['probability'] = np.empty((container.size, 3, 3), dtype=FTYPE) + self.data.unlink_containers() + + # setup more empty arrays + for container in self.data: + container['prob_e'] = np.empty((container.size), dtype=FTYPE) + container['prob_mu'] = np.empty((container.size), dtype=FTYPE) + + def compute_function(self): + + + assert self.is_map + + if self.is_map: + # speed up calculation by adding links + self.data.link_containers('nu', ['nue_cc', 'numu_cc', 'nutau_cc', + 'nue_nc', 'numu_nc', 'nutau_nc']) + self.data.link_containers('nubar', ['nuebar_cc', 'numubar_cc', 'nutaubar_cc', + 'nuebar_nc', 'numubar_nc', 'nutaubar_nc']) + + # this can be done in a more clever way (don't have to recalculate all paths) + YeI = self.params.YeI.value.m_as('dimensionless') + YeO = self.params.YeO.value.m_as('dimensionless') + YeM = self.params.YeM.value.m_as('dimensionless') + if YeI != self.YeI or YeO != self.YeO or YeM != self.YeM: + self.YeI = YeI; self.YeO = YeO; self.YeM = YeM + self.layers.setElecFrac(self.YeI, self.YeO, self.YeM) + for container in self.data: + self.layers.calcLayers(container['true_coszen']) + container['densities'] = self.layers.density.reshape((container.size, self.layers.max_layers)) + container['densities_neutron_weighted'] = self.layers.density_neutron_weighted.reshape((container.size, self.layers.max_layers)) + container['distances'] = self.layers.distance.reshape((container.size, self.layers.max_layers)) + + for container in self.data: + energy_idx = self.data.representation.names.index('true_energy') + + energies = self.data.representation.dims[energy_idx].weighted_centers.m + distances = container['distances'].reshape(*self.data.representation.shape, -1) + densities = container['densities'].reshape(*self.data.representation.shape, -1) + densities_neutron_weighted = container['densities_neutron_weighted'].reshape(*self.data.representation.shape, -1) + if energy_idx == 0: + distances = distances[0, :] + densities = densities[0, :] + densities_neutron_weighted = densities_neutron_weighted[0, :] + else: + distances = distances[:, 0] + densities = densities[:, 0] + densities_neutron_weighted = densities_neutron_weighted[:, 0] + + if container['nubar'] == 1: + is_anti = False + elif container['nubar'] == -1: + is_anti = True + + p = self.osc_prob(energies, distances, self.external_params, is_anti, densities, densities_neutron_weighted) + + if energy_idx == 0: + container['probability'] = p[:, :, :3, :3].reshape(-1, 3, 3) + else: + container['probability'] = np.swapaxes(p[:, :, :3, :3], 0, 1).reshape(-1, 3, 3) + + container.mark_changed('probability') + + # the following is flavour specific, hence unlink + self.data.unlink_containers() + + for container in self.data: + container['prob_e'] = container['probability'][:, 0, container['flav']] + container['prob_mu'] = container['probability'][:, 1, container['flav']] + container.mark_changed('prob_e') + container.mark_changed('prob_mu') + + def apply_function(self): + + # update the outputted weights + for container in self.data: + container['weights'] *= (container['nu_flux'][:,0] * container['prob_e']) + (container['nu_flux'][:,1] * container['prob_mu']) + diff --git a/pisa/stages/osc/layers.py b/pisa/stages/osc/layers.py index 495b727c3..1999397c2 100644 --- a/pisa/stages/osc/layers.py +++ b/pisa/stages/osc/layers.py @@ -53,6 +53,7 @@ def extCalcLayers(cz, prop_height, detector_depth, rhos, + rhos_neutron_weighted, coszen_limit, radii, max_layers): @@ -67,6 +68,7 @@ def extCalcLayers(cz, prop_height : height at which neutrinos are assumed to be produced (float) detector_depth : depth at which the detector is buried (float) rhos : densities (already weighted by electron fractions) (ndarray) + rhos_neutron_weighted : densities (already weighted by neutron fractions) (ndarray) radii : radii defining the Earth's layer (ndarray) coszen : coszen values corresponding to the radii above (ndarray) max_layers : maximum number of layers it is possible to cross (int) @@ -74,7 +76,8 @@ def extCalcLayers(cz, Returns ------- n_layers : int number of layers - density : array of densities, flattened from (cz, max_layers) + density : array of electron-weighted densities, flattened from (cz, max_layers) + density_neutron_weighted : array of neutron-weighted densities, flattened from (cz, max_layers) distance : array of distances per layer, flattened from (cz, max_layers) """ @@ -84,6 +87,7 @@ def extCalcLayers(cz, # in the pi_prob3 module densities = np.zeros((len(cz), max_layers), dtype=FTYPE) + densities_neutron_weighted = np.zeros((len(cz), max_layers), dtype=FTYPE) distances = np.zeros((len(cz), max_layers), dtype=FTYPE) number_of_layers = np.zeros(len(cz)) @@ -108,6 +112,7 @@ def extCalcLayers(cz, segments_lengths = np.concatenate((segments_lengths, np.zeros(radii.shape[0] - idx))) density = rhos*(segments_lengths > 0.) + density_neutron_weighted = rhos_neutron_weighted * (segments_lengths > 0.) else: # @@ -154,18 +159,26 @@ def extCalcLayers(cz, # start by removing the layers not crossed from rhos inner_layer_mask = coszen_limit>coszen density = np.concatenate((rhos[inner_layer_mask],rhos[inner_layer_mask][1:-1][::-1])) + density_neutron_weighted = np.concatenate( + ( + rhos_neutron_weighted[inner_layer_mask], + rhos_neutron_weighted[inner_layer_mask][1:-1][::-1], + ) + ) # As an extra precaution, set all densities that are not crossed to zero density*=(segments_lengths>0) + density_neutron_weighted *= segments_lengths > 0 number_of_layers[i] = np.sum(segments_lengths > 0.) # append to the large list for j in range(len(density)): # index j may not run all the way to max_layers, unreached indices stay zero densities[i, j] = density[j] + densities_neutron_weighted[i, j] = density_neutron_weighted[j] distances[i, j] = segments_lengths[j] - return number_of_layers, densities, distances + return number_of_layers, densities, densities_neutron_weighted, distances class Layers(object): @@ -196,7 +209,10 @@ class Layers(object): number of layers crossed for every CZ value density : 1d float array of length (max_layers * len(cz)) - containing density values and filled up with 0s otherwise + containing electron-weighted density values and filled up with 0s otherwise + + density_neutron_weighted : 1d float array of length (max_layers * len(cz)) + containing neutron-weighted density values and filled up with 0s otherwise distance : 1d float array of length (max_layers * len(cz)) containing distance values and filled up with 0s otherwise @@ -219,16 +235,24 @@ def __init__(self, prem_file, detector_depth=1., prop_height=2.): # w.r.t the file. The first elements of the arrays below corresponds # the Earth's surface, and the following numbers go deeper toward the # planet's core + self.rhos_unweighted = self.prem[..., 1][::-1].astype(FTYPE) self.rhos = self.prem[..., 1][::-1].astype(FTYPE) + self.rhos_neutron_weighted = self.prem[..., 1][::-1].astype(FTYPE) self.radii = self.prem[..., 0][::-1].astype(FTYPE) r_earth = self.prem[-1][0] self.default_elec_frac = 0.5 - - - + # Add an external layer corresponding to the atmosphere / production boundary self.radii = np.concatenate((np.array([r_earth+prop_height]), self.radii)) self.rhos = np.concatenate((np.ones(1, dtype=FTYPE), self.rhos)) + self.rhos_unweighted = np.concatenate( + (np.ones(1, dtype=FTYPE), self.rhos_unweighted) + ) + self.rhos_neutron_weighted = np.concatenate( + (np.ones(1, dtype=FTYPE), self.rhos_neutron_weighted) + ) + + self.max_layers = 2 * (len(self.radii)) @@ -270,9 +294,11 @@ def setElecFrac(self, YeI, YeO, YeM): raise ValueError("Cannot set electron fraction when not using an Earth model") self.YeFrac = np.array([YeI, YeO, YeM], dtype=FTYPE) + self.YnFrac = np.array([1 - YeI, 1 - YeO, 1 - YeM], dtype=FTYPE) # re-weight the layer densities accordingly self.weight_density_to_YeFrac() + self.weight_density_to_YnFrac() def scaling(self, scaling_array): """ @@ -336,12 +362,13 @@ def calcLayers(self, cz): raise ValueError("Cannot calculate layers when not using an Earth model") # run external function - self._n_layers, self._density, self._distance = extCalcLayers( + self._n_layers, self._density, self._density_neutron_weighted, self._distance = extCalcLayers( cz=cz, r_detector=self.r_detector, prop_height=self.prop_height, detector_depth=self.detector_depth, rhos=self.rhos, + rhos_neutron_weighted=self.rhos_neutron_weighted, coszen_limit=self.coszen_limit, radii=self.radii, max_layers=self.max_layers, @@ -359,6 +386,12 @@ def density(self): raise ValueError("Cannot get density when not using an Earth model") return self._density + @property + def density_neutron_weighted(self): + if not self.using_earth_model: + raise ValueError("Cannot get density when not using an Earth model") + return self._density_neutron_weighted + @property def distance(self): return self._distance @@ -415,6 +448,46 @@ def weight_density_to_YeFrac(self): self.rhos = weighted_densities + def weight_density_to_YnFrac(self): + """ + Adjust the densities of the provided earth model layers + for the different neutron fractions in the inner core, + outer core and mantle. + """ + + # TODO make this generic + R_INNER = 1221.5 + R_OUTER = 3480.0 + R_MANTLE = 6371.0 # the crust is assumed to have the same electron fraction as the mantle + + assert ( + isinstance(self.YnFrac, np.ndarray) and self.YnFrac.shape[0] == 3 + ), "ERROR: YnFrac must be an array of size 3" + # + # TODO: insert extra radii is the electron density boundaries + # don't match the current layer boundaries + + # + # Weight the density properly + # + density_inner = self.rhos * self.YnFrac[0] * (self.radii <= R_INNER) + density_outer = ( + self.rhos_unweighted + * self.YnFrac[1] + * (self.radii <= R_OUTER) + * (self.radii > R_INNER) + ) + density_mantle = ( + self.rhos_unweighted + * self.YnFrac[2] + * (self.radii <= R_MANTLE) + * (self.radii > R_OUTER) + ) + + weighted_densities = density_inner + density_outer + density_mantle + + self.rhos_neutron_weighted = weighted_densities + def test_layers_1():