Addition of LRI potential

pisa/stages/osc/lri_params.py

@@ -0,0 +1,121 @@
+"""
+LRI Params: Charecterize Long Range Interaction mediator
+
+Developed by following osc_params.py and nsi_params.py
+"""
+
+from __future__ import absolute_import, division
+
+import numpy as np
+from pisa import CTYPE, FTYPE
+from pisa.utils.comparisons import ALLCLOSE_KW, isscalar, recursiveEquality
+from pisa.utils.log import Levels, set_verbosity, logging
+
+__all__ = ['LRIParams']
+
+class LRIParams(object):
+    """
+    Holds the mediator information of long range interaction:z'
+    Assumed three anamoly free symmetries, Le_mu, Le_tau, Lmu_tau(by mixing z and z')).
+
+    Attributes
+    ----------
+    potential matrix : Three 2d float array of shape (3,3), one for each symmetry
+
+    Potential matrix holding the potential term of three different symmetris, which is a 
+    function of mediator mass, and the coupling constant of the interaction.
+
+
+    """
+
+    def __init__(self):
+
+        self._v_lri = 0.
+        self._potential_matrix_emu = np.zeros((3, 3), dtype=FTYPE)
+        self._potential_matrix_etau = np.zeros((3, 3), dtype=FTYPE)
+        self._potential_matrix_mutau = np.zeros((3, 3), dtype=FTYPE)
+
+    # --- LRI potential ---
+
+    @property
+    def v_lri(self):
+        """Potential term of symmetry e mu"""
+        return self._v_lri
+
+    @v_lri.setter
+    def v_lri(self, value):
+        assert value <1.
+        self._v_lri = value
+
+    @property
+    def potential_matrix_emu(self):
+        """LRI matter interaction potential matrix e mu symmetry"""
+
+        v_matrix = np.zeros((3, 3), dtype=FTYPE)
+
+        v_matrix[0, 0] = self.v_lri
+        v_matrix[0, 1] = 0.
+        v_matrix[0, 2] = 0.
+        v_matrix[1, 0] = 0.
+        v_matrix[1, 1] = - self.v_lri
+        v_matrix[1, 2] = 0.
+        v_matrix[2, 0] = 0.
+        v_matrix[2, 1] = 0.
+        v_matrix[2, 2] = 0.
+
+        assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW)
+
+        return v_matrix
+
+    @property
+    def potential_matrix_etau(self):
+        """LRI matter interaction potential matrix e tau symmetry"""
+
+        v_matrix = np.zeros((3, 3), dtype=FTYPE)
+
+        v_matrix[0, 0] = self.v_lri
+        v_matrix[0, 1] = 0.
+        v_matrix[0, 2] = 0.
+        v_matrix[1, 0] = 0.
+        v_matrix[1, 1] = 0.
+        v_matrix[1, 2] = 0.
+        v_matrix[2, 0] = 0.
+        v_matrix[2, 1] = 0.
+        v_matrix[2, 2] = - self.v_lri
+
+        assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW)
+
+        return v_matrix
+
+    @property
+    def potential_matrix_mutau(self):
+        """LRI matter interaction potential matrix mu tau symmetry"""
+
+        v_matrix = np.zeros((3, 3), dtype=FTYPE)
+
+        v_matrix[0, 0] = 0.
+        v_matrix[0, 1] = 0.
+        v_matrix[0, 2] = 0.
+        v_matrix[1, 0] = 0.
+        v_matrix[1, 1] = self.v_lri
+        v_matrix[1, 2] = 0.
+        v_matrix[2, 0] = 0.
+        v_matrix[2, 1] = 0.
+        v_matrix[2, 2] = - self.v_lri
+
+        assert np.allclose(v_matrix, v_matrix.conj().T, **ALLCLOSE_KW)
+
+        return v_matrix
+
+
+def test_lri_params():
+    """
+    # TODO: implement me!
+    """
+    pass
+
+if __name__=='__main__':
+    from pisa import TARGET
+    from pisa.utils.log import set_verbosity, logging
+    set_verbosity(1)
+    test_lri_params()

pisa/stages/osc/prob3.py

@@ -18,6 +18,7 @@
 from pisa.stages.osc.nsi_params import StdNSIParams, VacuumLikeNSIParams
 from pisa.stages.osc.osc_params import OscParams
 from pisa.stages.osc.decay_params import DecayParams
+from pisa.stages.osc.lri_params import LRIParams
 from pisa.stages.osc.layers import Layers
 from pisa.stages.osc.prob3numba.numba_osc_hostfuncs import propagate_array, fill_probs
 from pisa.utils.numba_tools import WHERE
@@ -64,6 +65,8 @@ class prob3(Stage):
             eps_mutau_phase : quantity (angle)
             eps_tautau : quantity (dimensionless)
             decay_alpha3 : quantity (energy^2)
+            v_lri : quantity (eV)
+
 
     **kwargs
         Other kwargs are handled by Stage
@@ -77,6 +80,7 @@ def __init__(
       reparam_mix_matrix=False,
       neutrino_decay=False,
       scale_density=False,
+      lri_type=None,
       **std_kwargs,
     ):
 
@@ -155,7 +159,22 @@ def __init__(
         else: 
             decay_params = ()
 
-        expected_params = expected_params + nsi_params + decay_params
+        if lri_type is not None:
+            choices = ['emu-symmetry', 'etau-symmetry', 'mutau-symmetry']
+            lri_type = lri_type.strip().lower()
+            if not lri_type in choices:
+                raise ValueError(
+                    'Chosen LRI symmetry type "%s" not available! Choose one of %s.'
+                    % (lri_type, choices)
+                )
+        self.lri_type = lri_type
+
+        if self.lri_type is None:
+            lri_params = ()
+        else:
+            lri_params = ('v_lri',)
+
+        expected_params = expected_params + nsi_params + decay_params + lri_params
 
         # init base class
         super().__init__(
@@ -169,6 +188,8 @@ def __init__(
         self.nsi_params = None
         self.decay_params = None
         self.decay_matrix = None
+        self.lri_params = None
+        self.lri_pot = None
         # Note that the interaction potential (Hamiltonian) just scales with the
         # electron density N_e for propagation through the Earth,
         # even(to very good approx.) in the presence of generalised interactions
@@ -201,6 +222,10 @@ def setup_function(self):
         if self.neutrino_decay:
             logging.debug('Working with neutrino decay')
             self.decay_params = DecayParams()
+
+        if self.lri_type is not None:
+            logging.debug('Working with LRI')
+            self.lri_params = LRIParams()
 
         # setup the layers
         #if self.params.earth_model.value is not None:
@@ -262,6 +287,7 @@ def calc_probs(self, nubar, e_array, rho_array, len_array, out):
                         self.gen_mat_pot_matrix_complex,
                         self.decay_flag,
                         self.decay_matix,
+                        self.lri_pot,
                         nubar,
                         e_array,
                         rho_array,
@@ -337,6 +363,9 @@ def compute_function(self):
 
         if self.neutrino_decay:
             self.decay_params.decay_alpha3 = self.params.decay_alpha3.value.m_as('eV**2')
+
+        if self.lri_type is not None:
+            self.lri_params.v_lri = self.params.v_lri.value.m_as('eV')
 
         # now we can proceed to calculate the generalised matter potential matrix
         std_mat_pot_matrix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE)
@@ -361,6 +390,18 @@ def compute_function(self):
         else :
             self.decay_matix = np.zeros((3, 3), dtype=FTYPE) + 1.j * np.zeros((3, 3), dtype=FTYPE)
 
+        self.lri_pot = np.zeros((3, 3), dtype=FTYPE)
+        types_lri = ['emu-symmetry', 'etau-symmetry', 'etau-symmetry']
+        if self.lri_type is not None:               
+            if self.lri_type == 'emu-symmetry':
+                self.lri_pot = self.lri_params.potential_matrix_emu
+            elif self.lri_type == 'etau-symmetry':
+                self.lri_pot = self.lri_params.potential_matrix_etau
+            elif self.lri_type == 'mutau-symmetry':
+                self.lri_pot = self.lri_params.potential_matrix_mutau    
+            else:
+                raise Exception("Implemented symmetries are" % types_lri)
+
 
         for container in self.data:
             self.calc_probs(container['nubar'],

pisa/stages/osc/prob3numba/numba_osc_hostfuncs.py

@@ -58,29 +58,29 @@
 
 
 @guvectorize(
-    [f"({FX}[:,:], {CX}[:,:], {CX}[:,:],  {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"],
-    "(a,a), (a,a), (b,c), (), (b,c), (), (), (i), (i) -> (a,a)",
+    [f"({FX}[:,:], {CX}[:,:], {CX}[:,:],  {IX}, {CX}[:,:], {FX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"],
+    "(a,a), (a,a), (b,c), (), (b,c), (b,c), (), (), (i), (i) -> (a,a)",
     target=TARGET,
 )
-def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability):
+def propagate_array(dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability):
     """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET
     is "cuda" or "host")"""
     osc_probs_layers_kernel(
-        dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability
+        dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability
     )
 
 
 @njit(
-    [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"],
+    [f"({FX}[:,:], {CX}[:,:], {CX}[:,:], {IX}, {CX}[:,:], {FX}[:,:], {IX}, {FX}, {FX}[:], {FX}[:], {FX}[:,:])"],
     parallel=TARGET == "parallel"
 )
 def propagate_scalar(
-    dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability
+    dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability
 ):
     """wrapper to run `osc_probs_layers_kernel` from host (whether TARGET
     is "cuda" or "host")"""
     osc_probs_layers_kernel(
-        dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, densities, distances, probability
+        dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, densities, distances, probability
     )
 
 
@@ -97,6 +97,7 @@ def propagate_scalar(
         f"{CX}[:,:], "  # H_vac
         f"{IX}, "       # neutrino decay flag
         f"{CX}[:,:], "  # H_decay
+        f"{FX}[:,:], "  # lri_pot
         f"{FX}[:,:], "  # dm
         f"{CX}[:,:], "  # transition_matrix
         ")"
@@ -114,6 +115,7 @@ def get_transition_matrix_hostfunc(
     H_vac,
     decay_flag,
     H_decay,
+    lri_pot,
     dm,
     transition_matrix,
 ):
@@ -130,6 +132,7 @@ def get_transition_matrix_hostfunc(
         H_vac,
         decay_flag,
         H_decay,
+        lri_pot,
         dm,
         transition_matrix,
     )

pisa/stages/osc/prob3numba/numba_osc_kernels.py

@@ -120,7 +120,7 @@
 
 @myjit
 def osc_probs_layers_kernel(
-    dm, mix, mat_pot, decay_flag, mat_decay, nubar, energy, density_in_layer, distance_in_layer, osc_probs
+    dm, mix, mat_pot, decay_flag, mat_decay, lri_pot, nubar, energy, density_in_layer, distance_in_layer, osc_probs
 ):
     """ Calculate oscillation probabilities
 
@@ -144,6 +144,11 @@ def osc_probs_layers_kernel(
 
     mat_decay : complex 2d array
         decay matrix with -j*alpha3 = [2,2] element
+
+    lri_pot : real 2d array
+        Potential contribution due to matter with the consideration of 
+        Long Range Interaction. this potential not a generalised one, so 
+        it passed as a separate matrix.
 
     nubar : real int, scalar or Nd array (broadcast dim)
         1 for neutrinos, -1 for antineutrinos
@@ -256,6 +261,7 @@ def osc_probs_layers_kernel(
                         H_vac,
                         decay_flag,
                         H_decay,
+                        lri_pot,
                         dm,
                         transition_matrix,
                     )
@@ -306,6 +312,7 @@ def osc_probs_layers_kernel(
                     H_vac,
                     decay_flag,
                     H_decay,
+                    lri_pot,
                     dm,
                     transition_matrix,
                 )
@@ -330,7 +337,9 @@ def osc_probs_layers_kernel(
         else:
             raw_input_psi[i] = 1.0
 
+
         matrix_dot_vector(transition_product, raw_input_psi, output_psi)
+
         osc_probs[i][0] += output_psi[0].real ** 2 + output_psi[0].imag ** 2
         osc_probs[i][1] += output_psi[1].real ** 2 + output_psi[1].imag ** 2
         osc_probs[i][2] += output_psi[2].real ** 2 + output_psi[2].imag ** 2
@@ -348,6 +357,7 @@ def get_transition_matrix(
     H_vac,
     decay_flag,
     H_decay,
+    lri_pot,
     dm,
     transition_matrix,
 ):
@@ -389,6 +399,11 @@ def get_transition_matrix(
 
     decay_flag: int
         +1 forstandard oscillations + decay, -1 for standard oscillations
+
+    lri_pot : real 2d array
+        Potential contribution due to matter with the consideration of 
+        Long Range Interaction. this potential not a generalised one, so 
+        it passed as a separate matrix.
 
     dm : real 2d array
         Mass splitting matrix, eV^2
@@ -416,6 +431,14 @@ def get_transition_matrix(
 
     get_H_mat(rho, mat_pot, nubar, H_mat)
 
+    # Adding the LRI potential to H_mat(lri pot is 2d array with full of zeros in the absence of lri)
+    for i in range(3):
+        for j in range(3):
+            if nubar > 0:
+                H_mat[i, j] = H_mat[i, j] + lri_pot[i, j]*1e9 #eV/GeV
+            elif nubar < 0:
+                H_mat[i, j] = H_mat[i, j] - lri_pot[i, j]*1e9 #ev/Gev
+
     # Get the full Hamiltonian by adding together matter and vacuum parts
     one_over_two_e = 0.5 / energy