From 9aca13052b1c1c60889488105afe8128e6600ac1 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Thu, 3 Oct 2024 12:31:13 +0200
Subject: [PATCH 01/12] Basic functions for shift-and-add technique

---
 stingray/fourier.py | 133 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 133 insertions(+)

diff --git a/stingray/fourier.py b/stingray/fourier.py
index 613e687eb..092ca1414 100644
--- a/stingray/fourier.py
+++ b/stingray/fourier.py
@@ -1358,6 +1358,139 @@ def fun(times, gti, segment_size):
         yield cts
 
 
+def center_pds_on_f(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 = center_pds_on_f(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)
+
+    minbin = start_f_idx - nbins // 2
+    maxbin = minbin + nbins
+
+    if minbin < 0:
+        chunk[-minbin : min(nbins, powers.size - minbin)] = powers[: minbin + nbins]
+    elif maxbin > powers.size:
+        chunk[: nbins - (maxbin - powers.size)] = powers[minbin:]
+    else:
+        chunk[:] = powers[minbin:maxbin]
+
+    return chunk
+
+
+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])
+    """
+    final_powers = np.zeros(nbins)
+    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
+
+    count = np.zeros(nbins)
+    for f0, powers, m in zip(f0_list, power_list, M):
+        idx = np.searchsorted(freqs, f0_list)
+
+        powers = np.asarray(powers) * m
+        new_power = center_pds_on_f(freqs, powers, f0, nbins=nbins)
+        bad = np.isnan(new_power)
+        new_power[bad] = 0.0
+        final_powers += new_power
+        count += np.array(~bad, dtype=int) * m
+
+    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)
+    final_powers = final_powers / count
+
+    if rebin is not None:
+        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
 ):

From ea9ffd0ddc37bb8f47d68ba2e62a1e445a3ca751 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Thu, 3 Oct 2024 16:26:31 +0200
Subject: [PATCH 02/12] Eliminate lightkurve link

---
 stingray/lightcurve.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index 367082cfc..085b13b9f 100644
--- a/stingray/lightcurve.py
+++ b/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.
 

From cb4262e180380725fb71d97aa0833ac1759164cd Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Fri, 4 Oct 2024 12:18:38 +0200
Subject: [PATCH 03/12] Add methods for shift-and-add to DynamicalPowerspectrum
 and DynamicalCrossspectrum

---
 stingray/crossspectrum.py            | 84 ++++++++++++++++++++++++++--
 stingray/powerspectrum.py            | 23 ++++++--
 stingray/tests/test_crossspectrum.py | 15 +++++
 stingray/tests/test_powerspectrum.py | 15 +++++
 4 files changed, 125 insertions(+), 12 deletions(-)

diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 608dc764a..741f448c1 100644
--- a/stingray/crossspectrum.py
+++ b/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 RuntimeError("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,68 @@ 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):
+        """Shift and add the dynamical power 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
+
+        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
+        )
+
+        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],
diff --git a/stingray/powerspectrum.py b/stingray/powerspectrum.py
index d9d7dc175..fa6334f7a 100755
--- a/stingray/powerspectrum.py
+++ b/stingray/powerspectrum.py
@@ -1020,7 +1020,20 @@ class DynamicalPowerspectrum(DynamicalCrossspectrum):
         The number of averaged cross spectra.
     """
 
-    def __init__(self, lc, segment_size, norm="frac", gti=None, sample_time=None):
+    def __init__(self, lc=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 lc is None:
+            self._initialize_empty()
+            self.dyn_ps = None
+            return
+
+        if segment_size is None or lc is None:
+            raise RuntimeError("lc and segment_size must all be specified")
+
         if isinstance(lc, EventList) and sample_time is None:
             raise ValueError("To pass an input event lists, please specify sample_time")
         elif isinstance(lc, Lightcurve):
@@ -1033,13 +1046,11 @@ def __init__(self, lc, segment_size, norm="frac", gti=None, sample_time=None):
         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(lc)
 
+    def shift_and_add(self, f0_list, nbins=100):
+        return super().shift_and_add(f0_list, nbins=nbins, output_obj_type=AveragedPowerspectrum)
+
     def _make_matrix(self, lc):
         """
         Create a matrix of powers for each time step and each frequency step.
diff --git a/stingray/tests/test_crossspectrum.py b/stingray/tests/test_crossspectrum.py
index b7e1919a5..fb5b6a9b6 100644
--- a/stingray/tests/test_crossspectrum.py
+++ b/stingray/tests/test_crossspectrum.py
@@ -1660,3 +1660,18 @@ def test_rebin_frequency_mean_method(self, method):
         assert np.allclose(new_dps.freq, rebin_freq)
         assert np.allclose(new_dps.dyn_ps, rebin_dps, atol=0.00001)
         assert np.isclose(new_dps.df, df_new)
+
+    def test_shift_and_add(self):
+        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])
\ No newline at end of file
diff --git a/stingray/tests/test_powerspectrum.py b/stingray/tests/test_powerspectrum.py
index 205493266..062571d18 100644
--- a/stingray/tests/test_powerspectrum.py
+++ b/stingray/tests/test_powerspectrum.py
@@ -1021,6 +1021,21 @@ def test_rebin_frequency_mean_method(self, method):
         assert np.allclose(new_dps.dyn_ps, rebin_dps, atol=0.00001)
         assert np.isclose(new_dps.df, df_new)
 
+    def test_shift_and_add(self):
+        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 = DynamicalPowerspectrum()
+        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.0, 2.0, 5.0, 2.0, 1.5])
+        assert np.allclose(output.freq, [0.05, 0.15, 0.25, 0.35, 0.45])
+
 
 class TestRoundTrip:
     @classmethod

From 9f4ecdc77b08fc7b693e925c7de8063ac0e371e7 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Fri, 4 Oct 2024 12:23:11 +0200
Subject: [PATCH 04/12] Test for bad input to dyn*spectrum

---
 stingray/crossspectrum.py            | 2 +-
 stingray/powerspectrum.py            | 2 +-
 stingray/tests/test_crossspectrum.py | 8 ++++++--
 stingray/tests/test_powerspectrum.py | 4 ++++
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 741f448c1..322763962 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -2146,7 +2146,7 @@ def __init__(
             return
 
         if segment_size is None or data1 is None or data2 is None:
-            raise RuntimeError("data1, data2, and segment_size must all be specified")
+            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")
diff --git a/stingray/powerspectrum.py b/stingray/powerspectrum.py
index fa6334f7a..b0e138046 100755
--- a/stingray/powerspectrum.py
+++ b/stingray/powerspectrum.py
@@ -1032,7 +1032,7 @@ def __init__(self, lc=None, segment_size=None, norm="frac", gti=None, sample_tim
             return
 
         if segment_size is None or lc is None:
-            raise RuntimeError("lc and segment_size must all be specified")
+            raise TypeError("lc and segment_size must all be specified")
 
         if isinstance(lc, EventList) and sample_time is None:
             raise ValueError("To pass an input event lists, please specify sample_time")
diff --git a/stingray/tests/test_crossspectrum.py b/stingray/tests/test_crossspectrum.py
index fb5b6a9b6..e5ea543e0 100644
--- a/stingray/tests/test_crossspectrum.py
+++ b/stingray/tests/test_crossspectrum.py
@@ -1430,6 +1430,10 @@ def setup_class(cls):
         test_counts = [2, 3, 1, 3, 1, 5, 2, 1, 4, 2, 2, 2, 3, 4, 1, 7]
         cls.lc_test = Lightcurve(test_times, test_counts)
 
+    def test_bad_args(self):
+        with pytest.raises(TypeError, match=".must all be specified"):
+            _ = DynamicalCrossspectrum(1)
+
     def test_with_short_seg_size(self):
         with pytest.raises(ValueError):
             dps = DynamicalCrossspectrum(self.lc, self.lc, segment_size=0)
@@ -1673,5 +1677,5 @@ def test_shift_and_add(self):
         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])
\ No newline at end of file
+        assert np.array_equal(output.power, [2.0, 2.0, 5.0, 2.0, 1.5])
+        assert np.allclose(output.freq, [0.05, 0.15, 0.25, 0.35, 0.45])
diff --git a/stingray/tests/test_powerspectrum.py b/stingray/tests/test_powerspectrum.py
index 062571d18..d05d12a7f 100644
--- a/stingray/tests/test_powerspectrum.py
+++ b/stingray/tests/test_powerspectrum.py
@@ -878,6 +878,10 @@ def setup_class(cls):
         test_counts = [2, 3, 1, 3, 1, 5, 2, 1, 4, 2, 2, 2, 3, 4, 1, 7]
         cls.lc_test = Lightcurve(test_times, test_counts)
 
+    def test_bad_args(self):
+        with pytest.raises(TypeError, match=".must all be specified"):
+            _ = DynamicalPowerspectrum(1)
+
     def test_with_short_seg_size(self):
         with pytest.raises(ValueError):
             dps = DynamicalPowerspectrum(self.lc, segment_size=0)

From 55e6ac3c4b9b2965f94a4d16752d09fdd3bef660 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Fri, 4 Oct 2024 12:26:10 +0200
Subject: [PATCH 05/12] Add changelog

---
 docs/changes/849.feature.rst | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 docs/changes/849.feature.rst

diff --git a/docs/changes/849.feature.rst b/docs/changes/849.feature.rst
new file mode 100644
index 000000000..e5880abe5
--- /dev/null
+++ b/docs/changes/849.feature.rst
@@ -0,0 +1,2 @@
+implementation of the shift-and-add technique for QPOs and other varying power spectral features
+

From 6f489458ab25620799fac5e9004a605f11844710 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Fri, 4 Oct 2024 13:48:29 +0200
Subject: [PATCH 06/12] Fix and test rebin option

---
 stingray/crossspectrum.py            |  8 +++++---
 stingray/fourier.py                  |  2 ++
 stingray/powerspectrum.py            |  6 ++++--
 stingray/tests/test_powerspectrum.py | 15 +++++++++++++++
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 322763962..476d36355 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -2380,7 +2380,7 @@ 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):
+    def shift_and_add(self, f0_list, nbins=100, output_obj_type=AveragedCrossspectrum, rebin=None):
         """Shift and add the dynamical power spectrum.
 
         This is the basic operation for the shift-and-add operation used to track
@@ -2401,7 +2401,9 @@ def shift_and_add(self, f0_list, nbins=100, output_obj_type=AveragedCrossspectru
         ----------------
         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.
         Returns
         -------
         output: :class:`AveragedPowerspectrum` or :class:`AveragedCrossspectrum`
@@ -2426,7 +2428,7 @@ def shift_and_add(self, f0_list, nbins=100, output_obj_type=AveragedCrossspectru
         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
+            self.freq, self.dyn_ps.T, f0_list, nbins=nbins, M=self.m, df=self.df, rebin=rebin
         )
 
         output = output_obj_type()
diff --git a/stingray/fourier.py b/stingray/fourier.py
index 092ca1414..6f1148697 100644
--- a/stingray/fourier.py
+++ b/stingray/fourier.py
@@ -20,6 +20,7 @@
     show_progress,
     sum_if_not_none_or_initialize,
     fix_segment_size_to_integer_samples,
+    rebin_data,
 )
 
 
@@ -1485,6 +1486,7 @@ def shift_and_add(freqs, power_list, f0_list, nbins=100, rebin=None, df=None, M=
     final_powers = final_powers / count
 
     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
 
diff --git a/stingray/powerspectrum.py b/stingray/powerspectrum.py
index b0e138046..08cb85984 100755
--- a/stingray/powerspectrum.py
+++ b/stingray/powerspectrum.py
@@ -1048,8 +1048,10 @@ def __init__(self, lc=None, segment_size=None, norm="frac", gti=None, sample_tim
 
         self._make_matrix(lc)
 
-    def shift_and_add(self, f0_list, nbins=100):
-        return super().shift_and_add(f0_list, nbins=nbins, output_obj_type=AveragedPowerspectrum)
+    def shift_and_add(self, f0_list, nbins=100, rebin=None):
+        return super().shift_and_add(
+            f0_list, nbins=nbins, output_obj_type=AveragedPowerspectrum, rebin=rebin
+        )
 
     def _make_matrix(self, lc):
         """
diff --git a/stingray/tests/test_powerspectrum.py b/stingray/tests/test_powerspectrum.py
index d05d12a7f..fbe77fdca 100644
--- a/stingray/tests/test_powerspectrum.py
+++ b/stingray/tests/test_powerspectrum.py
@@ -1040,6 +1040,21 @@ def test_shift_and_add(self):
         assert np.array_equal(output.power, [2.0, 2.0, 5.0, 2.0, 1.5])
         assert np.allclose(output.freq, [0.05, 0.15, 0.25, 0.35, 0.45])
 
+    def test_shift_and_add_rebin(self):
+        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 = DynamicalPowerspectrum()
+        dps.dyn_ps = power_list
+        dps.freq = freqs
+        dps.df = 0.1
+        dps.m = 1
+        output = dps.shift_and_add(f0_list, nbins=5, rebin=2)
+        assert np.array_equal(output.m, [5, 6])
+        assert np.array_equal(output.power, [2.0, 3.5])
+        assert np.allclose(output.freq, [0.1, 0.3])
+
 
 class TestRoundTrip:
     @classmethod

From 7b7585d24fc37602c21c2f01cb622e566a9a1cbe Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Fri, 4 Oct 2024 14:38:17 +0200
Subject: [PATCH 07/12] Additional test with orbital motion

---
 stingray/tests/test_fourier.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/stingray/tests/test_fourier.py b/stingray/tests/test_fourier.py
index acd3fcfd8..96076a113 100644
--- a/stingray/tests/test_fourier.py
+++ b/stingray/tests/test_fourier.py
@@ -938,3 +938,29 @@ def test_deprecation_rms_calculation(self):
             M=M,
         )
         assert np.isclose(rms, rms_from_unnorm, atol=3 * rmse_from_unnorm)
+
+
+@pytest.mark.parametrize("ntimes", [100, 1000])
+def test_shift_and_add_orbit(ntimes):
+    # This time correct for orbital motion
+    from stingray.fourier import shift_and_add
+
+    fmid = 0.7
+    freqs = np.linspace(0.699, 0.701, 1001)
+    porb = 2.52 * 86400
+    asini = 22.5
+    t0 = porb / 2
+    times = np.linspace(0, porb, ntimes + 1)[:-1]
+    power_list = np.zeros((times.size, freqs.size))
+    omega = 2 * np.pi / porb
+    orbit_freqs = fmid * (1 - asini * omega * np.cos(omega * (times - t0)))
+
+    idx = np.searchsorted(freqs, orbit_freqs)
+    for i_t, power in zip(idx, power_list):
+        power[i_t] = 1
+
+    f, p, n = shift_and_add(freqs, power_list, orbit_freqs, nbins=5)
+    # If we corrected well, the power should be the average of all max powers in the
+    # original series
+    assert np.max(p) == 1
+    assert np.max(n) == times.size

From e867a1c203fa95c8670cd6dd86e963ab241b5aa6 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Mon, 7 Oct 2024 16:46:59 +0200
Subject: [PATCH 08/12] Specify how the output frequency array is created [docs
 only]

---
 stingray/fourier.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/stingray/fourier.py b/stingray/fourier.py
index 6f1148697..ad4021d75 100644
--- a/stingray/fourier.py
+++ b/stingray/fourier.py
@@ -1442,7 +1442,9 @@ def shift_and_add(freqs, power_list, f0_list, nbins=100, rebin=None, df=None, M=
     Returns
     -------
     f : np.array
-        Array of output frequencies
+        Array of output frequencies. This will be centered on the mean of the
+        input ``f0_list``, and have the same frequency resolution as the original
+        frequency array.
     p : np.array
         Array of output powers
     n : np.array

From 9d06733c26079c52c2817682327e767a0c3d5435 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Mon, 7 Oct 2024 16:51:42 +0200
Subject: [PATCH 09/12] Fix docstrings [docs only]

---
 stingray/crossspectrum.py |  6 ++++-
 stingray/powerspectrum.py | 46 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 476d36355..64d33b2a7 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -2381,7 +2381,7 @@ 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 power spectrum.
+        """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).
@@ -2404,6 +2404,10 @@ def shift_and_add(self, f0_list, nbins=100, output_obj_type=AveragedCrossspectru
         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`
diff --git a/stingray/powerspectrum.py b/stingray/powerspectrum.py
index 08cb85984..1c3093129 100755
--- a/stingray/powerspectrum.py
+++ b/stingray/powerspectrum.py
@@ -1049,6 +1049,52 @@ def __init__(self, lc=None, segment_size=None, norm="frac", gti=None, sample_tim
         self._make_matrix(lc)
 
     def shift_and_add(self, f0_list, nbins=100, rebin=None):
+        """Shift-and-add the dynamical power 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.
+
+        Returns
+        -------
+        output: :class:`AveragedPowerspectrum`
+            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 = DynamicalPowerspectrum()
+        >>> 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 isinstance(output, AveragedPowerspectrum)
+        >>> 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])
+        """
         return super().shift_and_add(
             f0_list, nbins=nbins, output_obj_type=AveragedPowerspectrum, rebin=rebin
         )

From fc6d5680084e006598a7b5029bab95edb65bf584 Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Tue, 8 Oct 2024 14:02:24 +0200
Subject: [PATCH 10/12] Speed up functions by eliminating needless copies and
 `asarray`s, jit-ing functions

---
 stingray/fourier.py | 138 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 111 insertions(+), 27 deletions(-)

diff --git a/stingray/fourier.py b/stingray/fourier.py
index ad4021d75..8c9aff1df 100644
--- a/stingray/fourier.py
+++ b/stingray/fourier.py
@@ -21,6 +21,7 @@
     sum_if_not_none_or_initialize,
     fix_segment_size_to_integer_samples,
     rebin_data,
+    njit,
 )
 
 
@@ -1359,7 +1360,66 @@ def fun(times, gti, segment_size):
         yield cts
 
 
-def center_pds_on_f(freqs, powers, f0, nbins=100):
+@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.
@@ -1385,27 +1445,57 @@ def center_pds_on_f(freqs, powers, f0, nbins=100):
     >>> freqs = np.arange(1, 100) * 0.1
     >>> powers = 10 / freqs
     >>> f0 = 0.3
-    >>> p = center_pds_on_f(freqs, powers, f0)
+    >>> 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)
+    # fchunk = np.zeros(nbins)
 
     start_f_idx = np.searchsorted(freqs, f0)
 
-    minbin = start_f_idx - nbins // 2
-    maxbin = minbin + nbins
+    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
 
-    if minbin < 0:
-        chunk[-minbin : min(nbins, powers.size - minbin)] = powers[: minbin + nbins]
-    elif maxbin > powers.size:
-        chunk[: nbins - (maxbin - powers.size)] = powers[minbin:]
-    else:
-        chunk[:] = powers[minbin:maxbin]
 
-    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):
@@ -1442,9 +1532,7 @@ def shift_and_add(freqs, power_list, f0_list, nbins=100, rebin=None, df=None, M=
     Returns
     -------
     f : np.array
-        Array of output frequencies. This will be centered on the mean of the
-        input ``f0_list``, and have the same frequency resolution as the original
-        frequency array.
+        Array of output frequencies
     p : np.array
         Array of output powers
     n : np.array
@@ -1460,32 +1548,28 @@ def shift_and_add(freqs, power_list, f0_list, nbins=100, rebin=None, df=None, M=
     >>> 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])
     """
-    final_powers = np.zeros(nbins)
+
+    # 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))
+    # 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
 
-    count = np.zeros(nbins)
-    for f0, powers, m in zip(f0_list, power_list, M):
-        idx = np.searchsorted(freqs, f0_list)
+    center_f_indices = np.searchsorted(freqs, f0_list)
 
-        powers = np.asarray(powers) * m
-        new_power = center_pds_on_f(freqs, powers, f0, nbins=nbins)
-        bad = np.isnan(new_power)
-        new_power[bad] = 0.0
-        final_powers += new_power
-        count += np.array(~bad, dtype=int) * m
+    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)
-    final_powers = final_powers / count
 
     if rebin is not None:
         _, count, _, _ = rebin_data(final_freqs, count, rebin * df)

From a4a8b8a022766a47149c8a762b4c25d5af43055b Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Tue, 8 Oct 2024 22:56:34 +0200
Subject: [PATCH 11/12] Update docs with latest fixes and additions

---
 docs/notebooks      | 2 +-
 docs/timeseries.rst | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/notebooks b/docs/notebooks
index 2db8e50be..29620d646 160000
--- a/docs/notebooks
+++ b/docs/notebooks
@@ -1 +1 @@
-Subproject commit 2db8e50be46e3d5cddebb2c15647c45d63de4a92
+Subproject commit 29620d646b22266dc097c455fe615f9a0827269a
diff --git a/docs/timeseries.rst b/docs/timeseries.rst
index 1a62b5860..54322bc1f 100644
--- a/docs/timeseries.rst
+++ b/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

From 69bf6c564dce2a664e78d4e40286b164c5b02c8e Mon Sep 17 00:00:00 2001
From: Matteo Bachetti <matteo@matteobachetti.it>
Date: Wed, 9 Oct 2024 14:18:16 +0200
Subject: [PATCH 12/12] Update docs with dynspec notebooks

---
 docs/notebooks | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/notebooks b/docs/notebooks
index 29620d646..65c1ba97f 160000
--- a/docs/notebooks
+++ b/docs/notebooks
@@ -1 +1 @@
-Subproject commit 29620d646b22266dc097c455fe615f9a0827269a
+Subproject commit 65c1ba97fa5f40085973b7684e4ff967463c2b09