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 + diff --git a/docs/notebooks b/docs/notebooks index 2db8e50be..65c1ba97f 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit 2db8e50be46e3d5cddebb2c15647c45d63de4a92 +Subproject commit 65c1ba97fa5f40085973b7684e4ff967463c2b09 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 diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py index 608dc764a..64d33b2a7 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 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], diff --git a/stingray/fourier.py b/stingray/fourier.py index 613e687eb..8c9aff1df 100644 --- a/stingray/fourier.py +++ b/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 ): 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 - `_ 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. diff --git a/stingray/powerspectrum.py b/stingray/powerspectrum.py index d9d7dc175..1c3093129 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 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") elif isinstance(lc, Lightcurve): @@ -1033,13 +1046,59 @@ 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, 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 + ) + 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..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) @@ -1660,3 +1664,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.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_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 diff --git a/stingray/tests/test_powerspectrum.py b/stingray/tests/test_powerspectrum.py index 205493266..fbe77fdca 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) @@ -1021,6 +1025,36 @@ 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]) + + 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