Merge pull request #749 from nu-radio/feature/stokes

add function and method to electric_field to compute Stokes parameters

NuRadioReco/framework/electric_field.py

@@ -2,6 +2,8 @@
 import NuRadioReco.framework.base_trace
 import NuRadioReco.framework.parameters as parameters
 import NuRadioReco.framework.parameter_serialization
+import radiotools.coordinatesystems
+from NuRadioReco.utilities.trace_utilities import get_stokes
 try:
     import cPickle as pickle
 except ImportError:
@@ -127,6 +129,76 @@ def set_position(self, position):
         """
         self._position = position
 
+    def get_stokes_parameters(
+            self, window_samples=None, vxB_vxvxB=False, magnetic_field_vector=None,
+            site=None, filter_kwargs=None
+        ):
+        """
+        Return the stokes parameters for the electric field.
+
+        By default, the stokes parameters are returned in (eTheta, ePhi);
+        this assumes the 3d efield trace is stored in (eR, eTheta, ePhi).
+        To return the stokes parameters in (vxB, vxvxB) coordinates instead,
+        one has to specify the magnetic field vector.
+
+        Parameters
+        ----------
+        window_samples : int | None, default: None
+            If None, return the stokes parameters over the full traces.
+            If not None, returns a rolling average of the stokes parameters
+            over ``window_samples``. This may be more optimal if the duration
+            of the signal is much shorter than the length of the full trace.
+        vxB_vxvxB : bool, default: False
+            If False, returns the stokes parameters for the
+            (assumed) (eTheta, ePhi) coordinates of the electric field.
+            If True, convert to (vxB, vxvxB) first. In this case,
+            one has to additionally specify either the magnetic field vector
+            or the sit.
+        magnetic_field_vector : 3-tuple of floats | None, default: None
+            The direction of the magnetic field (in x,y,z)
+        site : string | None, default: None
+            The site of the detector. Can be used instead of the ``magnetic_field_vector``
+            if the magnetic field vector for this site is included in ``radiotools``
+        filter_kwargs : dict | None, default: None
+            Optional arguments to bandpass filter the trace
+            before computing the stokes parameters. They are passed on to
+            `get_filtered_trace(**filter_kwargs)`
+
+        Returns
+        -------
+        stokes : array of floats
+            The stokes parameters. If ``window_samples=None`` (default), the shape of
+            the returned array is ``(4,)`` and corresponds to the I, Q, U and V parameters.
+            Otherwise, the array will have shape ``(4, len(efield) - window_samples + 1)``
+            and correspond to the values of the stokes parameters over the specified
+            window sizes.
+
+        See Also
+        --------
+        NuRadioReco.utilities.trace_utilities.get_stokes : Function that computes the stokes parameters
+        """
+        if filter_kwargs:
+            trace = self.get_filtered_trace(**filter_kwargs)
+        else:
+            trace = self.get_trace()
+
+        if not vxB_vxvxB:
+            return get_stokes(trace[1], trace[2], window_samples=window_samples)
+        else:
+            try:
+                zenith = self.get_parameter(parameters.electricFieldParameters.zenith)
+                azimuth = self.get_parameter(parameters.electricFieldParameters.azimuth)
+                cs = radiotools.coordinatesystems.cstrafo(
+                    zenith, azimuth, magnetic_field_vector=magnetic_field_vector, site=site)
+                efield_trace_vxB_vxvxB = cs.transform_to_vxB_vxvxB(
+                    cs.transform_from_onsky_to_ground(trace)
+                )
+                return get_stokes(*efield_trace_vxB_vxvxB[:2], window_samples=window_samples)
+            except KeyError as e:
+                logger.error("Failed to compute stokes parameters in (vxB, vxvxB), electric field does not have a signal direction")
+                raise(e)
+
+
     def serialize(self, save_trace):
         if save_trace:
             base_trace_pkl = NuRadioReco.framework.base_trace.BaseTrace.serialize(self)

NuRadioReco/utilities/trace_utilities.py

@@ -112,6 +112,81 @@ def get_electric_field_energy_fluence(electric_field_trace, times, signal_window
 
     return f_signal * dt * conversion_factor_integrated_signal
 
+def get_stokes(trace_u, trace_v, window_samples=128, squeeze=True):
+    """
+    Compute the stokes parameters for electric field traces
+
+    Parameters
+    ----------
+    trace_u : 1d array (float)
+        The u component of the electric field trace
+    trace_v : 1d array (float)
+        The v component of the electric field trace.
+        The two components should have equal lengths,
+        and the (u, v) coordinates should be perpendicular.
+        Common choices are (theta, phi) or (vxB, vxvxB)
+    window_samples : int | None, default: 128
+        If not None, return a running average
+        of the stokes parameters over ``window_samples``.
+        If None, compute the stokes parameters over the
+        whole trace (equivalent to ``window_samples=len(trace_u)``).
+    squeeze : bool, default: True
+        Only relevant if ``window_samples=None``. Squeezes
+        out the second axis (which has a length of one)
+        and returns an array of shape (4,)
+
+    Returns
+    -------
+    stokes : 2d array of floats
+        The stokes parameters I, Q, U, V. The shape of
+        the returned array is ``(4, len(trace_u) - window_samples +1)``,
+        i.e. stokes[0] returns the I parameter,
+        stokes[1] corresponds to Q, and so on.
+
+    Examples
+    --------
+    For an electric field defined in (eR, eTheta, ePhi) components,
+    the stokes parameters can be given simply by:
+
+    .. code-block::
+
+        get_stokes(electric_field.get_trace()[1], electric_field.get_trace()[2])
+
+    To instead get the stokes parameters in vxB and vxvxB, we need to first obtain
+    the appropriate electric field components
+
+    .. code-block::
+
+        cs = radiotools.coordinatesystems.cstrafo(zenith, azimuth, magnetic_field_vector)
+
+        efield_trace_vxB_vxvxB = cs.transform_to_vxB_vxvxB(
+            cs.transform_from_onsky_to_ground(efield.get_trace())
+        )
+
+    """
+
+    assert len(trace_u) == len(trace_v)
+    h1 = scipy.signal.hilbert(trace_u)
+    h2 = scipy.signal.hilbert(trace_v)
+    stokes_i = np.abs(h1)**2 + np.abs(h2)**2
+    stokes_q = np.abs(h1)**2 - np.abs(h2)**2
+    uv = 2 * h1 * np.conjugate(h2)
+    stokes_u = np.real(uv)
+    stokes_v = np.imag(uv)
+    stokes = np.array([stokes_i, stokes_q, stokes_u, stokes_v])
+    if window_samples == 1: # no need to average
+        return stokes
+    elif window_samples is None:
+        window_samples = len(h1)
+
+    stokes = np.asarray([
+        scipy.signal.convolve(i, np.ones(window_samples), mode='valid') for i in stokes
+    ])
+    stokes /= window_samples
+
+    if squeeze:
+        return np.squeeze(stokes)
+    return stokes
 
 def upsampling_fir(trace, original_sampling_frequency, int_factor=2, ntaps=2 ** 7):
     """