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 <[email protected]> Co-authored-by: Leander Fischer <[email protected]> Co-authored-by: Jan Weldert <[email protected]>
icecube · Jul 28, 2024 · 1f28fb3 · 1f28fb3
1 parent ee8cc50
commit 1f28fb3
Show file tree

Hide file tree

Showing 2 changed files with 264 additions and 7 deletions.
diff --git a/pisa/stages/osc/external.py 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
@@ -53,6 +53,7 @@ def extCalcLayers(cz,
         prop_height,
         detector_depth,
         rhos,
+        rhos_neutron_weighted,
         coszen_limit,
         radii,
         max_layers):
@@ -67,14 +68,16 @@ 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)
 
     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():