Move check of powers outside function and test

Author: matteobachetti

stingray/fourier.py

@@ -972,6 +972,48 @@ def _intrinsic_coherence_uncertainties(
     return uncertainty
 
 
+@vectorize(
+    ["UniTuple(float64[:, :], 2)(float64[:, :], float64[:, :], float64, float64, int64)"],
+    nopython=True,
+)
+def check_powers_for_intrinsic_coherence(
+    power1, power2, power1_noise, power2_noise, n_ave, threshold=3
+):
+    """Check if the powers are above the threshold for the intrinsic coherence to be well defined.
+
+    Parameters
+    ----------
+    power1 : float `np.array`
+        sub-band periodogram
+    power2 : float `np.array`
+        reference-band periodogram
+    power1_noise : float
+        Poisson noise level of the sub-band periodogram
+    power2_noise : float
+        Poisson noise level of the reference-band periodogram
+    n_ave : int
+        number of intervals that have been averaged to obtain the input spectra
+
+    Other Parameters
+    ----------------
+    threshold : float, default 3
+        The threshold in sigma above the noise level for the powers to be considered "high".
+
+    Returns
+    -------
+    low_power1 : bool `np.array`
+        Whether the powers of the first band are below the threshold.
+    low_power2 : bool `np.array`
+        Whether the powers of the second band are below the threshold.
+
+    """
+    # Consider low power when the power is less than 3 sigma above the noise level.
+    # This is somewhat arbitrary, better criteria suggestions are welcome.
+    low_power1 = (power1 - power1_noise) <= 3 * power1_noise / np.sqrt(n_ave)
+    low_power2 = (power2 - power2_noise) <= 3 * power2_noise / np.sqrt(n_ave)
+    return low_power1 or low_power2
+
+
 @vectorize(
     ["UniTuple(float64[:, :], 2)(float64, float64, float64, float64, float64, int64, float64)"],
     nopython=True,
@@ -1022,19 +1064,16 @@ def _intrinsic_coherence(
         The intrinsic coherence values at all frequencies. If ``return_uncertainty`` is True,
         also returns the uncertainty on the coherence.
     """
+
+    if check_powers_for_intrinsic_coherence(power1, power2, power1_noise, power2_noise, n_ave):
+        # If the powers are too close to the noise level, the coherence is not well defined.
+        return np.nan, np.nan
+
     bsq = _bias_term(power1, power2, power1_noise, power2_noise, n_ave, 1)
     num = (cross_power * np.conj(cross_power)).real - bsq
     if num < 0:
         num = 0
 
-    # Consider low power when the power is less than 3 sigma above the noise level.
-    # This is somewhat arbitrary, better criteria suggestions are welcome.
-    low_power1 = power1 - power1_noise <= 3 * power1_noise / np.sqrt(n_ave)
-    low_power2 = power2 - power2_noise <= 3 * power2_noise / np.sqrt(n_ave)
-    if low_power1 or low_power2:
-        # If the powers are too close to the noise level, the coherence is not well defined.
-        return np.nan, np.nan
-
     power1_sub = power1 - power1_noise
     power2_sub = power2 - power2_noise
 
@@ -1072,17 +1111,16 @@ def _intrinsic_coherence_with_adjusted_bias(
     and the coherence will be set to NaN.
 
     """
+
+    if check_powers_for_intrinsic_coherence(power1, power2, power1_noise, power2_noise, n_ave):
+        # If the powers are too close to the noise level, the coherence is not well defined.
+        return np.nan, np.nan
+
     # Consider low power when the power is less than 3 sigma above the noise level.
     # This is somewhat arbitrary, better criteria suggestions are welcome.
     power1_sub = power1 - power1_noise
     power2_sub = power2 - power2_noise
 
-    low_power1 = power1_sub <= 3 * power1_noise / np.sqrt(n_ave)
-    low_power2 = power2_sub <= 3 * power2_noise / np.sqrt(n_ave)
-    if low_power1 or low_power2:
-        # If the powers are too close to the noise level, the coherence is not well defined.
-        return np.nan, np.nan
-
     current_coherence = 1.0
 
     for _ in range(40):

stingray/tests/test_fourier.py

@@ -20,6 +20,7 @@
 from stingray.fourier import get_average_ctrate, normalize_leahy_from_variance
 from stingray.fourier import integrate_power_in_frequency_range
 from stingray.fourier import get_rms_from_rms_norm_periodogram, get_rms_from_unnorm_periodogram
+from stingray.fourier import check_powers_for_intrinsic_coherence
 
 from stingray.utils import check_allclose_and_print
 from astropy.modeling.models import Lorentz1D
@@ -227,6 +228,16 @@ def test_raw_high_bias(self):
         assert np.allclose(coh, 0)
         assert np.isclose(coh_sngl, 0)
 
+    def test_check_powers_for_intrinsic_coherence(self):
+        pow1 = np.array([10, 20, 30])
+        pow2 = np.array([10, 20, 30])
+        pow1_noise = 5
+        pow2_noise = 5
+        n_ave = np.array([1, 10, 10])
+        # Only one power is below the threshold, due to the low number of averaged powers.
+        res = check_powers_for_intrinsic_coherence(pow1, pow2, pow1_noise, pow2_noise, n_ave)
+        assert np.all(res == np.array([True, False, False]))
+
 
 class TestFourier(object):
     @classmethod