Add PulsarPeriodogram class (Fixes #243)

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 *

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()

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