From 59b2aa5f4aeb2b6f8f5a91e8068780235aa3e50d Mon Sep 17 00:00:00 2001 From: pupperemeritus Date: Sat, 29 Jul 2023 23:53:51 +0530 Subject: [PATCH] Added basic tests, corrected algorithms wrt papers, typos in docstrings --- docs/changes/737.feature.rst | 2 +- stingray/fourier.py | 138 +++++++++++++--------------- stingray/lombscargle.py | 141 ++++++++++++++++++++--------- stingray/tests/test_lombscargle.py | 118 ++++++++++++++++++++++++ 4 files changed, 282 insertions(+), 117 deletions(-) create mode 100644 stingray/tests/test_lombscargle.py diff --git a/docs/changes/737.feature.rst b/docs/changes/737.feature.rst index 534f4c495..70e8eb97b 100644 --- a/docs/changes/737.feature.rst +++ b/docs/changes/737.feature.rst @@ -1,2 +1,2 @@ -- Implemented the Lomb Scargle Fourier Transform +- Implemented the Lomb Scargle Fourier Transform (fast and slow versions) - Using which wrote the corresponding :class:`LombScargleCrossspectrum` and :class:`LombScarglePowerspectrum` \ No newline at end of file diff --git a/stingray/fourier.py b/stingray/fourier.py index d88080218..7931c3b58 100644 --- a/stingray/fourier.py +++ b/stingray/fourier.py @@ -1967,64 +1967,69 @@ def lsft_fast( raise ValueError("sign must be 1 or -1") # Constants initialization - const1 = np.sqrt(0.5) - const2 = const1 * sign - sum_xx = np.sum(y) - num_xt = len(y) num_ww = len(freqs) + const1 = np.sqrt(0.5) * np.sqrt(num_ww) + const2 = const1 # Arrays initialization ft_real = ft_imag = np.zeros((num_ww)) - ft_res = np.zeros((num_ww), dtype=np.complex128) f0, df, N = freqs[0], freqs[1] - freqs[0], len(freqs) # Sum (y_i * cos(wt - wtau)) - sumr, sumi = trig_sum(t, y, df, N, f0, oversampling=oversampling) + Sh, Ch = trig_sum(t, y, df, N, f0, oversampling=oversampling) # Summation of (cos(wt - wtau))^2 and (sin(wt - wtau))^2 - s2cos, c2cos = trig_sum( - t, np.ones_like(y) / len(y), df, N, f0, freq_factor=2, oversampling=oversampling + _, scos2x = trig_sum( + t, + np.ones_like(y) / len(y), + df, + N, + f0, + freq_factor=2, + oversampling=oversampling, ) # cos^2(x) = (1 + cos(2x))/2 # sin^2(x) = (1 - cos(2x))/2 - scos2, ssin2 = np.abs(1 + c2cos) / 2, np.abs(1 - s2cos) / 2 + C2, S2 = (1 + scos2x) / 2, (1 - scos2x) / 2 + # Angular Frequency + w = freqs * 2 * np.pi - freqs_new = f0 + df * np.arange(N) + # Summation of cos(2wt) and sin(2wt) + csum, ssum = trig_sum( + t, np.ones_like(y) / len(y), df, N, f0, freq_factor=2, oversampling=oversampling + ) - fft_freqs = np.fft.ifft(freqs_new, N) + wtau = 0.5 * np.arctan2(ssum, csum) - # Summation of cos(2wt) and sin(2wt) - csum, ssum = fft_freqs.real, fft_freqs.imag + coswtau = np.cos(wtau) + sinwtau = np.sin(wtau) - # Looping through all the frequencies - for i in range(num_ww): - if i == 0: - ft_real = sum_xx / np.sqrt(num_xt) - ft_imag = 0 - phase_this = 0 - else: - # Angular Frequency - wrun = freqs[i] * 2 * np.pi + sumr = Ch * coswtau + Sh * sinwtau + sumi = Sh * coswtau - Ch * sinwtau - # arctan(ssum / csum) - watan = np.arctan2(ssum[i], csum[i]) - wtau = 0.5 * watan + cos2wtau = np.cos(2 * wtau) + sin2wtau = np.sin(2 * wtau) - # Calculating the real and imaginary components of the Fourier transform - ft_real = const1 * (sumr[i] / np.sqrt(scos2[i])) - ft_imag = const2 * (sumi[i] / np.sqrt(ssin2[i])) + scos2 = 0.5 * (N + C2 * cos2wtau + S2 * sin2wtau) + ssin2 = 0.5 * (N - C2 * cos2wtau - S2 * sin2wtau) - # Phase of current angular frequency - phase_this = wtau - wrun * t[0] - # Resultant fourier transform for current angular frequency - ft_res[i] = np.multiply(np.complex128(complex(ft_real, ft_imag)), np.exp(1j * phase_this)) + # Calculating the real and imaginary components of the Fourier transform + ft_real = const1 * sumr / np.sqrt(scos2) + ft_imag = const2 * sumi / np.sqrt(ssin2) - if fullspec: - return ft_res - else: - return ft_res[freqs > 0] + # Looping through all the frequencies + ft_real[0] = sum_xx / np.sqrt(num_xt) + ft_imag[0] = 0 + + # Resultant fourier transform for current angular frequency + ft_res = np.complex128(ft_real + (1j * ft_imag)) * np.exp(-1j * w * t[0]) + + if sign == -1: + ft_res = np.conjugate(ft_res) + + return ft_res if fullspec else ft_res[freqs > 0] def lsft_slow( @@ -2059,58 +2064,43 @@ def lsft_slow( ft_res : numpy.ndarray An array of Fourier transformed data. """ - if sign not in [1, -1]: + if sign not in (1, -1): raise ValueError("sign must be 1 or -1") # Constants initialization - const1 = np.sqrt(0.5) - const2 = const1 * sign - - sum_xx = np.sum(y) - - num_xt = len(y) - num_ww = len(freqs) + const1 = np.sqrt(0.5) * np.sqrt(len(y)) + const2 = const1 - # Arrays initialization - ft_real = ft_imag = np.zeros((num_ww)) - ft_res = np.zeros((num_ww), dtype=np.complex128) + ft_real = np.zeros_like(freqs) + ft_imag = np.zeros_like(freqs) + ft_res = np.zeros_like(freqs, dtype=np.complex128) - # Looping through all the frequencies - for i in range(num_ww): + for i, freq in enumerate(freqs): + wrun = freq * 2 * np.pi if i == 0: - ft_real = sum_xx / np.sqrt(num_xt) - ft_imag = 0 - phase_this = 0 + ft_real[i] = sum(y) / np.sqrt(len(y)) + ft_imag[i] = 0 else: - # Angular Frequency - wrun = freqs[i] * 2 * np.pi - - # Summation of cos(2wt) and sin(2wt) csum = np.sum(np.cos(2.0 * wrun * t)) ssum = np.sum(np.sin(2.0 * wrun * t)) - # arctan(ssum / csum) watan = np.arctan2(ssum, csum) wtau = 0.5 * watan - # Sum (y_i * cos(wt - wtau)) - sumr = np.sum(np.multiply(y, np.cos(wrun * t - wtau))) - sumi = np.sum(np.multiply(y, np.sin(wrun * t - wtau))) + cos_wt_wtau = np.cos(wrun * t - wtau) + sin_wt_wtau = np.sin(wrun * t - wtau) - # Summation of (cos(wt - wtau))^2 and (sin(wt - wtau))^2 - scos2 = np.sum((np.power(np.cos(wrun * t - wtau), 2))) - ssin2 = np.sum((np.power(np.sin(wrun * t - wtau), 2))) + sumr = np.sum(y * cos_wt_wtau) + sumi = np.sum(y * sin_wt_wtau) - # Calculating the real and imaginary components of the Fourier transform - ft_real = np.multiply(const1, np.divide(sumr, np.sqrt(scos2))) - ft_imag = np.multiply(const2, np.divide(sumi, np.sqrt(ssin2))) + scos2 = np.sum(np.power(cos_wt_wtau, 2)) + ssin2 = np.sum(np.power(sin_wt_wtau, 2)) - # Phase of current angular frequency - phase_this = wtau - wrun * t[0] - # Resultant fourier transform for current angular frequency - ft_res[i] = np.multiply(np.complex128(complex(ft_real, ft_imag)), np.exp(1j * phase_this)) + ft_real[i] = const1 * sumr / np.sqrt(scos2) + ft_imag[i] = const2 * sumi / np.sqrt(ssin2) - if fullspec: - return ft_res - else: - return ft_res[freqs > 0] + ft_res[i] = np.complex128(ft_real[i] + (1j * ft_imag[i])) * np.exp(-1j * wrun * t[0]) + if sign == -1: + ft_res = np.conjugate(ft_res) + + return ft_res if fullspec else ft_res[freqs > 0] diff --git a/stingray/lombscargle.py b/stingray/lombscargle.py index 8435b8454..60d6e4498 100644 --- a/stingray/lombscargle.py +++ b/stingray/lombscargle.py @@ -30,10 +30,10 @@ class LombScargleCrossspectrum(Crossspectrum): data2: :class:`stingray.lightcurve.Lightcurve` or :class:`stingray.events.EventList`, optional, default ``None`` The dataset for the second, or "reference", band. - norm: {``frac``, ``abs``, ``leahy``, ``none``}, default ``none`` - The normalization of the (real part of the) cross spectrum. + norm: {``frac``, ``abs``, ``leahy``, ``none``}, string, optional, default ``none`` + The normalization of the cross spectrum. - power_type: string, optional, default ``real`` + power_type: {``real``, ``absolute``, ``all`}, string, optional, default ``real`` Parameter to choose among complete, real part and magnitude of the cross spectrum. fullspec: boolean, optional, default ``False`` @@ -69,7 +69,7 @@ class LombScargleCrossspectrum(Crossspectrum): Attributes ---------- - freqs: numpy.ndarray + freq: numpy.ndarray The array of mid-bin frequencies that the Fourier transform samples power: numpy.ndarray @@ -101,6 +101,13 @@ class LombScargleCrossspectrum(Crossspectrum): nphots2: float The total number of photons in light curve 2 + References + ---------- + .. [1] Scargle, J. D. , "Studies in astronomical time series analysis. III - Fourier + transforms, autocorrelation functions, and cross-correlation + functions of unevenly spaced data". ApJ 1:343, p874-887, 1989 + .. [2] Press W.H. and Rybicki, G.B, "Fast algorithm for spectral analysis + of unevenly sampled data". ApJ 1:338, p277, 1989 """ def __init__( @@ -120,6 +127,23 @@ def __init__( ): self._type = None good_input = data1 is not None and data2 is not None + if data1 is None and data2 is None: + if not skip_checks: + good_input = self.initial_checks( + data1=data1, + data2=data2, + norm=norm, + power_type=power_type, + dt=dt, + fullspec=fullspec, + min_freq=min_freq, + max_freq=max_freq, + df=df, + method=method, + oversampling=oversampling, + ) + self._initialize_empty() + return if not skip_checks: good_input = self.initial_checks( data1=data1, @@ -134,7 +158,6 @@ def __init__( method=method, oversampling=oversampling, ) - if dt is None: if isinstance(data1, Lightcurve): dt = data1.dt @@ -208,10 +231,12 @@ def initial_checks( if power_type not in ["all", "absolute", "real"]: raise ValueError("power_type must be one of ['all','absolute','real']") - if data1 is None or data2 is None: + if np.logical_xor(data1 is None, data2 is None): raise ValueError("You can't do a cross spectrum with just one lightcurve") + if min_freq < 0: raise ValueError("min_freq must be non-negative") + if max_freq is not None: if max_freq < min_freq or max_freq < 0: raise ValueError("max_freq must be non-negative and greater than min_freq") @@ -222,17 +247,40 @@ def initial_checks( if not isinstance(oversampling, int): raise TypeError("oversampling must be an integer") - # if not isinstance(df, float) and df is not None: - # raise TypeError("df must be a float") - if not isinstance(fullspec, bool): raise TypeError("fullspec must be a boolean") - if type(data1) not in [EventList, Lightcurve] or type(data2) not in [ - EventList, - Lightcurve, - ]: - raise TypeError("One of the arguments is not of type eventlist or lightcurve") + if np.logical_xor( + not (isinstance(data1, EventList) or isinstance(data1, Lightcurve) or data1 is None), + not (isinstance(data2, EventList) or isinstance(data2, Lightcurve) or data2 is None), + ): + raise TypeError("One of the arguments is not of type Eventlist or Lightcurve or None") + + if not ( + isinstance(data1, EventList) or isinstance(data1, Lightcurve) or data1 is None + ) and ( + not (isinstance(data2, EventList) or isinstance(data2, Lightcurve) or data2 is None), + ): + raise TypeError("Both the events are not of type Eventlist or Lightcurve or None") + + if type(data1) == type(data2): + if isinstance(data1, Lightcurve): + if len(data1.time) != len(data2.time): + raise ValueError("data1 and data2 must have the same length") + if isinstance(data1, EventList): + lc1 = data1.to_lc() + lc2 = data2.to_lc() + if len(lc1.time) != len(lc2.time): + raise ValueError("data1 and data2 must have the same length") + else: + if isinstance(data1, EventList): + lc1 = data1.to_lc() + if len(lc1.time) != len(data2.time): + raise ValueError("data1 and data2 must have the same length") + elif isinstance(data2, EventList): + lc2 = data2.to_lc() + if len(lc2.time) != len(data1.time): + raise ValueError("data1 and data2 must have the same length") return True @@ -298,7 +346,7 @@ def _make_crossspectrum(self, lc1, lc2, fullspec, method, oversampling): self.m = 1 - self.freqs, self.unnorm_power = self._ls_cross( + self.freq, self.unnorm_power = self._ls_cross( self.lc1, self.lc2, fullspec=fullspec, @@ -309,8 +357,8 @@ def _make_crossspectrum(self, lc1, lc2, fullspec, method, oversampling): self.power = self._normalize_crossspectrum(self.unnorm_power) if lc1.err_dist.lower() != lc2.err_dist.lower(): simon( - "Your lightcurves have different statistics.", - "The errors in the Crossspectrum will be incorrect.", + "Your lightcurves have different statistics." + "The errors in the Crossspectrum will be incorrect." ) elif lc1.err_dist.lower() != "poisson": @@ -363,7 +411,7 @@ def _make_auxil_pds(self, lc1, lc2): oversampling=self.oversampling, ) - def _ls_cross(self, lc1, lc2, freqs=None, fullspec=False, method="fast", oversampling=5): + def _ls_cross(self, lc1, lc2, freq=None, fullspec=False, method="fast", oversampling=5): """ Lomb-Scargle Fourier transform the two light curves, then compute the cross spectrum. Computed as CS = lc1 x lc2* (where lc2 is the one that gets @@ -390,15 +438,15 @@ def _ls_cross(self, lc1, lc2, freqs=None, fullspec=False, method="fast", oversam Returns ------- - freqs: numpy.ndarray + freq: numpy.ndarray The frequency grid at which the LSFT was evaluated cross: numpy.ndarray The cross spectrum value at each frequency. """ - if not freqs: - freqs = ( + if not freq: + freq = ( LombScargle( lc1.time, lc1.counts, @@ -416,36 +464,36 @@ def _ls_cross(self, lc1, lc2, freqs=None, fullspec=False, method="fast", oversam normalization="psd", ).autofrequency(minimum_frequency=self.min_freq, maximum_frequency=self.max_freq), )[0] - if max(freqs2) > max(freqs): - freqs = freqs2 + if max(freqs2) > max(freq): + freq = freqs2 if method == "slow": - lsft1 = lsft_slow(lc1.counts, lc1.time, freqs, sign=1, fullspec=fullspec) - lsft2 = lsft_slow(lc2.counts, lc2.time, freqs, sign=-1, fullspec=fullspec) + lsft1 = lsft_slow(lc1.counts, lc1.time, freq, sign=-1, fullspec=fullspec) + lsft2 = lsft_slow(lc2.counts, lc2.time, freq, sign=1, fullspec=fullspec) elif method == "fast": lsft1 = lsft_fast( lc1.counts, lc1.time, - freqs, + freq, + sign=-1, fullspec=fullspec, oversampling=oversampling, ) lsft2 = lsft_fast( lc2.counts, lc2.time, - freqs, + freq, fullspec=fullspec, - sign=-1, oversampling=oversampling, ) cross = np.multiply(lsft1, lsft2) if not fullspec: - freqs = freqs[freqs > 0] - cross = cross[freqs > 0] - return freqs, cross + freq = freq[freq > 0] + cross = cross[freq > 0] + return freq, cross def _initialize_empty(self): - self.freqs = None + self.freq = None self.power = None self.power_err = None self.unnorm_power = None @@ -458,6 +506,13 @@ def _initialize_empty(self): self.n = None self.fullspec = None self.k = 1 + self.err_dist = None + self.method = None + self.meancounts1 = None + self.meancounts2 = None + self.oversampling = None + self.variance1 = None + self.variance2 = None return def time_lag(self): @@ -473,7 +528,7 @@ def time_lag(self): if self.__class__ == LombScargleCrossspectrum: ph_lag = self.phase_lag() - return ph_lag / (2 * np.pi * self.freqs) + return ph_lag / (2 * np.pi * self.freq) else: raise AttributeError("Object has no attribute named 'time_lag' !") @@ -494,9 +549,14 @@ class LombScarglePowerspectrum(LombScargleCrossspectrum): data: :class:`stingray.lightcurve.Lightcurve` or :class:`stingray.events.EventList` object, optional, default ``None`` The light curve data to be Fourier-transformed. - norm: {"leahy" | "frac" | "abs" | "none" }, optional, default "frac" - The normaliation of the power spectrum to be used. Options are - "leahy", "frac", "abs" and "none", default is "frac". + norm: {``frac``, ``abs``, ``leahy``, ``none``}, string, optional, default ``none`` + The normalization of the cross spectrum. + + power_type: {``real``, ``absolute``, ``all`}, string, optional, default ``real`` + Parameter to choose among complete, real part and magnitude of the cross spectrum. + + fullspec: boolean, optional, default ``False`` + If False, keep only the positive frequencies, or if True, keep all of them . Other Parameters ---------------- @@ -528,10 +588,7 @@ class LombScarglePowerspectrum(LombScargleCrossspectrum): Attributes ---------- - norm: {"leahy" | "frac" | "abs" | "none" } - The normalization of the power spectrum. - - freqs: numpy.ndarray + freq: numpy.ndarray The array of mid-bin frequencies that the Fourier transform samples. power: numpy.ndarray @@ -562,8 +619,8 @@ class LombScarglePowerspectrum(LombScargleCrossspectrum): def __init__( self, data: Optional[Union[Lightcurve, EventList]] = None, - norm: Optional[str] = "none", - power_type: Optional[str] = "real", + norm: Optional[str] = "frac", + power_type: Optional[str] = "all", dt: Optional[float] = None, fullspec: Optional[bool] = False, skip_checks: Optional[bool] = False, diff --git a/stingray/tests/test_lombscargle.py b/stingray/tests/test_lombscargle.py new file mode 100644 index 000000000..d4d845ca1 --- /dev/null +++ b/stingray/tests/test_lombscargle.py @@ -0,0 +1,118 @@ +import numpy as np +import pytest +import copy +from stingray.lombscargle import LombScargleCrossspectrum, LombScarglePowerspectrum +from stingray.lightcurve import Lightcurve +from stingray.events import EventList +from stingray.simulator import Simulator +from stingray.exceptions import StingrayError +from scipy.interpolate import interp1d + + +class TestLombScargleCrossspectrum: + def setup_class(self): + sim = Simulator(0.0001, 10000, 100, 1, random_state=42, tstart=0) + lc1 = sim.simulate(0) + lc2 = sim.simulate(0) + self.rate1 = lc1.countrate + self.rate2 = lc2.countrate + low, high = lc1.time.min(), lc1.time.max() + s1 = lc1.counts + s2 = lc2.counts + t = lc1.time + t_new = t.copy() + t_new[1:-1] = t[1:-1] + (np.random.rand(len(t) - 2) / (high - low)) + s1_new = interp1d(t, s1, fill_value="extrapolate")(t_new) + s2_new = interp1d(t, s2, fill_value="extrapolate")(t_new) + self.lc1 = Lightcurve(t, s1_new, dt=lc1.dt) + self.lc2 = Lightcurve(t, s2_new, dt=lc2.dt) + with pytest.warns(UserWarning) as record: + self.lscs = LombScargleCrossspectrum(lc1, lc2) + + @pytest.mark.parametrize("skip_checks", [True, False]) + def test_initialize_empty(self, skip_checks): + lscs = LombScargleCrossspectrum(skip_checks=skip_checks) + assert lscs.freq is None + assert lscs.power is None + + def test_make_empty_crossspectrum(self): + lscs = LombScargleCrossspectrum() + assert lscs.freq is None + assert lscs.power is None + assert lscs.df is None + assert lscs.nphots1 is None + assert lscs.nphots2 is None + assert lscs.m == 1 + assert lscs.n is None + assert lscs.power_err is None + assert lscs.dt is None + assert lscs.err_dist is None + assert lscs.variance1 is None + assert lscs.variance2 is None + assert lscs.meancounts1 is None + assert lscs.meancounts2 is None + assert lscs.oversampling is None + assert lscs.method is None + + def test_init_with_one_lc_none(self): + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(self.lc1) + + def test_init_with_norm_not_str(self): + with pytest.raises(TypeError): + lscs = LombScargleCrossspectrum(norm=1) + + def test_init_with_invalid_norm(self): + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(norm="frabs") + + def test_init_with_power_type_not_str(self): + with pytest.raises(TypeError): + lscs = LombScargleCrossspectrum(power_type=3) + + def test_init_with_invalid_power_type(self): + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(power_type="reel") + + def test_init_with_wrong_lc_instance(self): + lc1_ = {"a": 1, "b": 2} + lc2_ = {"a": 1, "b": 2} + with pytest.raises(TypeError): + lscs = LombScargleCrossspectrum(lc1_, lc2_) + + def test_init_with_wrong_lc2_instance(self): + lc_ = {"a": 1, "b": 2} + with pytest.raises(TypeError): + lscs = LombScargleCrossspectrum(self.lc1, lc_) + + def test_init_with_wrong_lc1_instance(self): + lc_ = {"a": 1, "b": 2} + with pytest.raises(TypeError): + lscs = LombScargleCrossspectrum(lc_, self.lc2) + + def test_init_with_invalid_min_freq(self): + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(self.lc1, self.lc2, min_freq=-1) + + def test_init_with_invalid_max_freq(self): + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(self.lc1, self.lc2, max_freq=1, min_freq=3) + + def test_make_crossspectrum_diff_lc_counts_shape(self): + lc_ = Simulator(0.0001, 10423, 100, 1, random_state=42, tstart=0).simulate(0) + with pytest.raises(ValueError): + lscs = LombScargleCrossspectrum(self.lc1, lc_) + + def test_make_crossspectrum_diff_lc_stat(self): + lc_ = copy.deepcopy(self.lc1) + lc_.err_dist = "gauss" + with pytest.warns(UserWarning) as record: + cs = LombScargleCrossspectrum(self.lc1, lc_) + assert np.any(["different statistics" in r.message.args[0] for r in record]) + + def test_make_crossspectrum_diff_dt(self): + lc_ = Simulator(0.0002, 10000, 100, 1, random_state=42, tstart=0).simulate(0) + with pytest.raises( + StingrayError, match="Lightcurves do not have the same time binning dt." + ): + lscs = LombScargleCrossspectrum(self.lc1, lc_)