diff --git a/stingray/pulse/__init__.py b/stingray/pulse/__init__.py index 6e1695989..bba4a5569 100644 --- a/stingray/pulse/__init__.py +++ b/stingray/pulse/__init__.py @@ -23,3 +23,4 @@ from stingray.pulse.pulsar import * from stingray.pulse.search import * from stingray.pulse.modeling import * + from .periodogram import * diff --git a/stingray/pulse/periodogram.py b/stingray/pulse/periodogram.py new file mode 100644 index 000000000..badc4a025 --- /dev/null +++ b/stingray/pulse/periodogram.py @@ -0,0 +1,80 @@ +import numpy as np +import matplotlib.pyplot as plt +from stingray.pulse.search import epoch_folding_search, z_n_search + +class PulsarPeriodogram: + """ + A class to generate and store Epoch Folding and Z^2_n periodograms. + + Parameters + ---------- + events : stingray.EventList + The event list to search for pulsations. + frequencies : array-like + The frequencies to search. + nbin : int, optional, default 32 + The number of bins for the pulse profile (used in Epoch Folding). + nharm : int, optional, default 1 + The number of harmonics for the Z^2_n search. + + Attributes + ---------- + ef_freq : numpy.ndarray + The frequencies used for Epoch Folding. + ef_stat : numpy.ndarray + The Epoch Folding statistics. + z_freq : numpy.ndarray + The frequencies used for Z^2_n search. + z_stat : numpy.ndarray + The Z^2_n statistics. + """ + def __init__(self, events, frequencies, nbin=32, nharm=1): + self.events = events + self.frequencies = frequencies + self.nbin = nbin + self.nharm = nharm + + # Run Epoch Folding Search + # epoch_folding_search returns (freqs, stats) + self.ef_freq, self.ef_stat = epoch_folding_search( + events.time, frequencies, nbin=nbin + ) + + # Run Z_n Search + # z_n_search returns (freqs, stats) + self.z_freq, self.z_stat = z_n_search( + events.time, frequencies, nbin=nbin, nharm=nharm + ) + + def plot(self, ax=None, show=True): + """ + Plot the periodograms. + + Parameters + ---------- + ax : matplotlib.axes.Axes, optional + If provided, plot on this axis. Otherwise, create new figures. + show : bool, optional + If True, show the plot. + """ + if ax is None: + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True) + else: + # If user provides axes, we expect a list/array of 2 axes + ax1, ax2 = ax + + # Plot Epoch Folding + ax1.plot(self.ef_freq, self.ef_stat, label='Epoch Folding') + ax1.axhline(self.nbin - 1, ls='--', color='red', label='Noise Level (Approx)') + ax1.set_ylabel('EF Statistics') + ax1.set_title('Pulsar Periodogram') + ax1.legend() + + # Plot Z_n + ax2.plot(self.z_freq, self.z_stat, color='orange', label=f'Z_{self.nharm} Statistics') + ax2.set_ylabel(f'Z_{self.nharm} Statistics') + ax2.set_xlabel('Frequency (Hz)') + ax2.legend() + + if show: + plt.show() \ No newline at end of file diff --git a/stingray/tests/test_phase_lag_logic.py b/stingray/tests/test_phase_lag_logic.py new file mode 100644 index 000000000..8246b193e --- /dev/null +++ b/stingray/tests/test_phase_lag_logic.py @@ -0,0 +1,51 @@ +import numpy as np +import pytest +from unittest.mock import patch +from astropy.modeling import models +from stingray.spectroscopy import get_phase_lag + +# A simple Mock object to simulate the Crossspectrum class +class MockCrossspectrum: + def __init__(self, freqs, power): + self.freq = freqs + self.power = power + +def test_get_phase_lag_logic(): + # 1. Setup Frequency Domain (0 to 5 Hz) + freqs = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + + # 2. Setup Power (Complex numbers to simulate phase) + # At freq=1.0 (Index 1): Phase = 0 (Real number) + # At freq=3.0 (Index 3): Phase = pi/4 (Real == Imaginary) + power = np.zeros_like(freqs, dtype=complex) + power[1] = 10.0 + 0j # Angle = 0.0 + power[3] = 10.0 + 10.0j # Angle = pi/4 (approx 0.785) + + cs = MockCrossspectrum(freqs, power) + + # 3. Setup Model (Astropy Compound Model) + # Fundamental at 1.0 Hz, Harmonic at 3.0 Hz + # FIX: Changed Lorentzian1D to Lorentz1D + m1 = models.Lorentz1D(amplitude=1, x_0=1.0, fwhm=0.5) + m2 = models.Lorentz1D(amplitude=1, x_0=3.0, fwhm=0.5) + model = m1 + m2 + + # 4. Mock the internal dependency 'get_mean_phase_difference' + # We force it to return a fixed value (0.1) so we can verify the math formula strictly. + with patch('stingray.spectroscopy.get_mean_phase_difference', return_value=(0.1, 0.0)): + + # 5. Call the function + phi1, phi2, avg_psi = get_phase_lag(cs, model) + + # 6. Verify the Math manually + # Formula from source code: cap_phi_1 = pi/2 + delta_E_1 + expected_phi1 = np.pi/2 + 0.0 + + # Formula from source code: cap_phi_2 = 2 * (cap_phi_1 + avg_psi) + delta_E_2 + # avg_psi is mocked as 0.1 + expected_phi2 = 2 * (expected_phi1 + 0.1) + (np.pi/4) + + # Assertions + assert np.isclose(phi1, expected_phi1), f"Phi1 mismatch: Got {phi1}, Expected {expected_phi1}" + assert np.isclose(phi2, expected_phi2), f"Phi2 mismatch: Got {phi2}, Expected {expected_phi2}" + assert avg_psi == 0.1 \ No newline at end of file