Merge pull request #849 from StingraySoftware/shift_and_add

Shift and add

docs/changes/849.feature.rst

@@ -0,0 +1,2 @@
+implementation of the shift-and-add technique for QPOs and other varying power spectral features
+

docs/timeseries.rst

@@ -6,3 +6,4 @@ Working with more generic time series
 
    notebooks/StingrayTimeseries/StingrayTimeseries Tutorial.ipynb
    notebooks/StingrayTimeseries/Working with weights and polarization.ipynb
+   notebooks/Modeling/GP_Modeling/GP_modeling_tutorial.ipynb

stingray/crossspectrum.py

@@ -2132,7 +2132,22 @@ class DynamicalCrossspectrum(AveragedCrossspectrum):
         The number of averaged powers in each spectral bin (initially 1, it changes after rebinning).
     """
 
-    def __init__(self, data1, data2, segment_size, norm="frac", gti=None, sample_time=None):
+    def __init__(
+        self, data1=None, data2=None, segment_size=None, norm="frac", gti=None, sample_time=None
+    ):
+        self.segment_size = segment_size
+        self.sample_time = sample_time
+        self.gti = gti
+        self.norm = norm
+
+        if segment_size is None and data1 is None and data2 is None:
+            self._initialize_empty()
+            self.dyn_ps = None
+            return
+
+        if segment_size is None or data1 is None or data2 is None:
+            raise TypeError("data1, data2, and segment_size must all be specified")
+
         if isinstance(data1, EventList) and sample_time is None:
             raise ValueError("To pass input event lists, please specify sample_time")
         elif isinstance(data1, Lightcurve):
@@ -2145,11 +2160,6 @@ def __init__(self, data1, data2, segment_size, norm="frac", gti=None, sample_tim
         if segment_size < 2 * sample_time:
             raise ValueError("Length of the segment is too short to form a light curve!")
 
-        self.segment_size = segment_size
-        self.sample_time = sample_time
-        self.gti = gti
-        self.norm = norm
-
         self._make_matrix(data1, data2)
 
     def _make_matrix(self, data1, data2):
@@ -2370,6 +2380,74 @@ def trace_maximum(self, min_freq=None, max_freq=None):
 
         return np.array(max_positions)
 
+    def shift_and_add(self, f0_list, nbins=100, output_obj_type=AveragedCrossspectrum, rebin=None):
+        """Shift and add the dynamical cross spectrum.
+
+        This is the basic operation for the shift-and-add operation used to track
+        kHz QPOs in X-ray binaries (e.g. Méndez et al. 1998, ApJ, 494, 65).
+
+        Parameters
+        ----------
+        freqs : np.array
+            Array of frequencies, the same for all powers. Must be sorted and on a uniform
+            grid.
+        power_list : list of np.array
+            List of power spectra. Each power spectrum must have the same length
+            as the frequency array.
+        f0_list : list of float
+            List of central frequencies
+
+        Other parameters
+        ----------------
+        nbins : int, default 100
+            Number of bins to extract
+        rebin : int, default None
+            Rebin the final spectrum by this factor. At the moment, the rebinning
+            is linear.
+        output_obj_type : class, default :class:`AveragedCrossspectrum`
+            The type of the output object. Can be, e.g. :class:`AveragedCrossspectrum` or
+            :class:`AveragedPowerspectrum`.
+
+        Returns
+        -------
+        output: :class:`AveragedPowerspectrum` or :class:`AveragedCrossspectrum`
+            The final averaged power spectrum.
+
+        Examples
+        --------
+        >>> power_list = [[2, 5, 2, 2, 2], [1, 1, 5, 1, 1], [3, 3, 3, 5, 3]]
+        >>> power_list = np.array(power_list).T
+        >>> freqs = np.arange(5) * 0.1
+        >>> f0_list = [0.1, 0.2, 0.3, 0.4]
+        >>> dps = DynamicalCrossspectrum()
+        >>> dps.dyn_ps = power_list
+        >>> dps.freq = freqs
+        >>> dps.df = 0.1
+        >>> dps.m = 1
+        >>> output = dps.shift_and_add(f0_list, nbins=5)
+        >>> assert np.array_equal(output.m, [2, 3, 3, 3, 2])
+        >>> assert np.array_equal(output.power, [2. , 2. , 5. , 2. , 1.5])
+        >>> assert np.allclose(output.freq, [0.05, 0.15, 0.25, 0.35, 0.45])
+        """
+        from .fourier import shift_and_add
+
+        final_freqs, final_powers, count = shift_and_add(
+            self.freq, self.dyn_ps.T, f0_list, nbins=nbins, M=self.m, df=self.df, rebin=rebin
+        )
+
+        output = output_obj_type()
+        good = ~np.isnan(final_powers)
+
+        output.freq = final_freqs[good]
+        output.power = final_powers[good]
+        output.m = count[good]
+        output.df = self.df
+        output.norm = self.norm
+        output.gti = self.gti
+        output.segment_size = self.segment_size
+
+        return output
+
     def power_colors(
         self,
         freq_edges=[1 / 256, 1 / 32, 0.25, 2.0, 16.0],

stingray/fourier.py

@@ -20,6 +20,8 @@
     show_progress,
     sum_if_not_none_or_initialize,
     fix_segment_size_to_integer_samples,
+    rebin_data,
+    njit,
 )
 
 
@@ -1358,6 +1360,225 @@ def fun(times, gti, segment_size):
         yield cts
 
 
+@njit()
+def _safe_array_slice_indices(input_size, input_center_idx, nbins):
+    """Calculate the indices needed to extract a n-bin slice of an array, centered at an index.
+
+    Let us say we have an array of size ``input_size`` and we want to extract a slice of
+    ``nbins`` centered at index ``input_center_idx``. We should be robust when the slice goes
+    beyond the edges of the input array, possibly leaving missing values in the output array.
+    This function calculates the indices needed to extract the slice from the input array, and
+    the indices in the output array that will be filled.
+
+    In the most common case, the slice is entirely contained within the input array, so that the
+    output slice will just be ``[0:nbins]`` and the input slice
+    ``[input_center_idx - nbins // 2: input_center_idx - nbins // 2 + nbins]``.
+
+    Parameters
+    ----------
+    input_size : int
+        Input array size
+    center_idx : int
+        Index of the center of the slice
+    nbins : int
+        Number of bins to extract
+
+    Returns
+    -------
+    input_slice : list
+        Indices to extract the slice from the input array
+    output_slice : list
+        Indices to fill the output array
+
+    Examples
+    --------
+    >>> _safe_array_slice_indices(input_size=10, input_center_idx=5, nbins=3)
+    ([4, 7], [0, 3])
+
+    If the slice goes beyond the right edge: the output slice will only cover
+    the first two bins of the output array, and up to the end of the input array.
+    >>> _safe_array_slice_indices(input_size=6, input_center_idx=5, nbins=3)
+    ([4, 6], [0, 2])
+
+    """
+
+    minbin = input_center_idx - nbins // 2
+    maxbin = minbin + nbins
+
+    if minbin < 0:
+        output_slice = [-minbin, min(nbins, input_size - minbin)]
+        input_slice = [0, minbin + nbins]
+    elif maxbin > input_size:
+        output_slice = [0, nbins - (maxbin - input_size)]
+        input_slice = [minbin, input_size]
+    else:
+        output_slice = [0, nbins]
+        input_slice = [minbin, maxbin]
+
+    return input_slice, output_slice
+
+
+@njit()
+def extract_pds_slice_around_freq(freqs, powers, f0, nbins=100):
+    """Extract a slice of PDS around a given frequency.
+
+    This function extracts a slice of the power spectrum around a given frequency.
+    The slice has a length of ``nbins``. If the slice goes beyond the edges of the
+    power spectrum, the missing values are filled with NaNs.
+
+    Parameters
+    ----------
+    freqs : np.array
+        Array of frequencies, the same for all powers
+    powers : np.array
+        Array of powers
+    f0 : float
+        Central frequency
+
+    Other parameters
+    ----------------
+    nbins : int, default 100
+        Number of bins to extract
+
+    Examples
+    --------
+    >>> freqs = np.arange(1, 100) * 0.1
+    >>> powers = 10 / freqs
+    >>> f0 = 0.3
+    >>> p = extract_pds_slice_around_freq(freqs, powers, f0)
+    >>> assert np.isnan(p[0])
+    >>> assert not np.any(np.isnan(p[48:]))
+    """
+    powers = np.asarray(powers)
+    chunk = np.zeros(nbins) + np.nan
+    # fchunk = np.zeros(nbins)
+
+    start_f_idx = np.searchsorted(freqs, f0)
+
+    input_slice, output_slice = _safe_array_slice_indices(powers.size, start_f_idx, nbins)
+    chunk[output_slice[0] : output_slice[1]] = powers[input_slice[0] : input_slice[1]]
+    return chunk
+
+
+@njit()
+def _shift_and_average_core(input_array_list, weight_list, center_indices, nbins):
+    """Core function to shift_and_add, JIT-compiled for your convenience.
+
+    Parameters
+    ----------
+    input_array_list : list of np.array
+        List of input arrays
+    weight_list : list of float
+        List of weights for each input array
+    center_indices : list of int
+        Central indices of the slice of each input array to be summed
+    nbins : int
+        Number of bins to extract around the central index of each input array
+
+    Returns
+    -------
+    output_array : np.array
+        Average of the input arrays, weighted by the weights
+    sum_of_weights : np.array
+        Sum of the weights at each output bin
+    """
+    input_size = input_array_list[0].size
+    output_array = np.zeros(nbins)
+    sum_of_weights = np.zeros(nbins)
+    for idx, array, weight in zip(center_indices, input_array_list, weight_list):
+        input_slice, output_slice = _safe_array_slice_indices(input_size, idx, nbins)
+
+        for i in range(input_slice[1] - input_slice[0]):
+            output_array[output_slice[0] + i] += array[input_slice[0] + i] * weight
+
+            sum_of_weights[output_slice[0] + i] += weight
+
+    output_array = output_array / sum_of_weights
+
+    return output_array, sum_of_weights
+
+
+def shift_and_add(freqs, power_list, f0_list, nbins=100, rebin=None, df=None, M=None):
+    """Add a list of power spectra, centered on different frequencies.
+
+    This is the basic operation for the shift-and-add operation used to track
+    kHz QPOs in X-ray binaries (e.g. Méndez et al. 1998, ApJ, 494, 65).
+
+    Parameters
+    ----------
+    freqs : np.array
+        Array of frequencies, the same for all powers. Must be sorted and on a uniform
+        grid.
+    power_list : list of np.array
+        List of power spectra. Each power spectrum must have the same length
+        as the frequency array.
+    f0_list : list of float
+        List of central frequencies
+
+    Other parameters
+    ----------------
+    nbins : int, default 100
+        Number of bins to extract
+    rebin : int, default None
+        Rebin the final spectrum by this factor. At the moment, the rebinning
+        is linear.
+    df : float, default None
+        Frequency resolution of the power spectra. If not given, it is calculated
+        from the input frequencies.
+    M : int or list of int, default None
+        Number of segments used to calculate each power spectrum. If a list is
+        given, it must have the same length as the power list.
+
+    Returns
+    -------
+    f : np.array
+        Array of output frequencies
+    p : np.array
+        Array of output powers
+    n : np.array
+        Number of contributing power spectra at each frequency
+
+    Examples
+    --------
+    >>> power_list = [[2, 5, 2, 2, 2], [1, 1, 5, 1, 1], [3, 3, 3, 5, 3]]
+    >>> freqs = np.arange(5) * 0.1
+    >>> f0_list = [0.1, 0.2, 0.3, 0.4]
+    >>> f, p, n = shift_and_add(freqs, power_list, f0_list, nbins=5)
+    >>> assert np.array_equal(n, [2, 3, 3, 3, 2])
+    >>> assert np.array_equal(p, [2. , 2. , 5. , 2. , 1.5])
+    >>> assert np.allclose(f, [0.05, 0.15, 0.25, 0.35, 0.45])
+    """
+
+    # Check if the input list of power contains numpy arrays
+    if not hasattr(power_list[0], "size"):
+        power_list = np.asarray(power_list)
+    # input_size = np.size(power_list[0])
+    freqs = np.asarray(freqs)
+
+    # mid_idx = np.searchsorted(freqs, np.mean(f0_list))
+    if M is None:
+        M = 1
+    if not isinstance(M, Iterable):
+        M = np.ones(len(power_list)) * M
+
+    center_f_indices = np.searchsorted(freqs, f0_list)
+
+    final_powers, count = _shift_and_average_core(power_list, M, center_f_indices, nbins)
+
+    if df is None:
+        df = freqs[1] - freqs[0]
+
+    final_freqs = np.arange(-nbins // 2, nbins // 2 + 1)[:nbins] * df
+    final_freqs = final_freqs - (final_freqs[0] + final_freqs[-1]) / 2 + np.mean(f0_list)
+
+    if rebin is not None:
+        _, count, _, _ = rebin_data(final_freqs, count, rebin * df)
+        final_freqs, final_powers, _, _ = rebin_data(final_freqs, final_powers, rebin * df)
+        final_powers = final_powers / rebin
+
+    return final_freqs, final_powers, count
+
+
 def avg_pds_from_iterable(
     flux_iterable, dt, norm="frac", use_common_mean=True, silent=False, return_subcs=False
 ):

stingray/lightcurve.py

@@ -1383,8 +1383,7 @@ def analyze_lc_chunks(self, segment_size, func, fraction_step=1, **kwargs):
     def to_lightkurve(self):
         """
         Returns a `lightkurve.LightCurve` object.
-        This feature requires `Lightkurve
-        <https://docs.lightkurve.org/>`_ to be installed
+        This feature requires ``Lightkurve`` to be installed
         (e.g. ``pip install lightkurve``).  An `ImportError` will
         be raised if this package is not available.