Fix powerspectrum from lightcurve - wrong size, float casting

docs/changes/760.bugfix.rst

@@ -0,0 +1,2 @@
++ Fix a bug with segment sizes not exact multiples of dt when dealing with light curves
++ Fix a bug when light curve segments contain complex values

stingray/fourier.py

@@ -1028,7 +1028,9 @@ def get_average_ctrate(times, gti, segment_size, counts=None):
     return n_ph / (n_intvs * segment_size)
 
 
-def get_flux_iterable_from_segments(times, gti, segment_size, n_bin=None, fluxes=None, errors=None):
+def get_flux_iterable_from_segments(
+    times, gti, segment_size, n_bin=None, dt=None, fluxes=None, errors=None
+):
     """
     Get fluxes from different segments of the observation.
 
@@ -1059,6 +1061,8 @@ def get_flux_iterable_from_segments(times, gti, segment_size, n_bin=None, fluxes
         Array of fluxes.
     errors : float `np.array`, default None
         Array of error bars corresponding to the flux values above.
+    dt : float, default None
+        Time resolution of the light curve used to produce periodograms
 
     Yields
     ------
@@ -1073,10 +1077,14 @@ def get_flux_iterable_from_segments(times, gti, segment_size, n_bin=None, fluxes
             "At least one between fluxes (if light curve) and " "n_bin (if events) has to be set"
         )
 
-    dt = None
     binned = fluxes is not None
-    if binned:
+    cast_kind = float
+    if dt is None and binned:
         dt = np.median(np.diff(times[:100]))
+    if binned:
+        fluxes = np.asarray(fluxes)
+        if np.iscomplexobj(fluxes):
+            cast_kind = complex
 
     fun = _which_segment_idx_fun(binned, dt)
 
@@ -1093,9 +1101,9 @@ def get_flux_iterable_from_segments(times, gti, segment_size, n_bin=None, fluxes
             ).astype(float)
             cts = np.array(cts)
         else:
-            cts = fluxes[idx0:idx1].astype(float)
+            cts = fluxes[idx0:idx1].astype(cast_kind)
             if errors is not None:
-                cts = cts, errors[idx0:idx1]
+                cts = cts, errors[idx0:idx1].astype(cast_kind)
 
         yield cts
 
@@ -1842,11 +1850,13 @@ def avg_pds_from_events(
     """
     if segment_size is None:
         segment_size = gti.max() - gti.min()
-    n_bin = np.rint(segment_size / dt).astype(int)
-    dt = segment_size / n_bin
-
+    n_bin = int(segment_size / dt)
+    if fluxes is None:
+        dt = segment_size / n_bin
+    else:
+        segment_size = n_bin * dt
     flux_iterable = get_flux_iterable_from_segments(
-        times, gti, segment_size, n_bin, fluxes=fluxes, errors=errors
+        times, gti, segment_size, n_bin, dt=dt, fluxes=fluxes, errors=errors
     )
     cross = avg_pds_from_iterable(
         flux_iterable,
@@ -1944,15 +1954,18 @@ def avg_cs_from_events(
     """
     if segment_size is None:
         segment_size = gti.max() - gti.min()
-    n_bin = np.rint(segment_size / dt).astype(int)
+    n_bin = int(segment_size / dt)
     # adjust dt
-    dt = segment_size / n_bin
-
+    # dt = segment_size / n_bin
+    if fluxes1 is None and fluxes2 is None:
+        dt = segment_size / n_bin
+    else:
+        segment_size = n_bin * dt
     flux_iterable1 = get_flux_iterable_from_segments(
-        times1, gti, segment_size, n_bin, fluxes=fluxes1, errors=errors1
+        times1, gti, segment_size, n_bin, dt=dt, fluxes=fluxes1, errors=errors1
     )
     flux_iterable2 = get_flux_iterable_from_segments(
-        times2, gti, segment_size, n_bin, fluxes=fluxes2, errors=errors2
+        times2, gti, segment_size, n_bin, dt=dt, fluxes=fluxes2, errors=errors2
     )
 
     is_events = np.all([val is None for val in (fluxes1, fluxes2, errors1, errors2)])

stingray/tests/test_fourier.py

@@ -48,6 +48,48 @@ def test_norm():
     assert np.isclose(pdsfrac[good].mean(), pois_frac, rtol=0.01)
 
 
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128])
+def test_flux_iterables(dtype):
+    times = np.arange(4)
+    fluxes = np.ones(4).astype(dtype)
+    errors = np.ones(4).astype(dtype) * np.sqrt(2)
+    gti = np.asarray([[-0.5, 3.5]])
+    iter = get_flux_iterable_from_segments(times, gti, 2, n_bin=None, fluxes=fluxes, errors=errors)
+    cast_kind = float
+    if np.iscomplexobj(fluxes):
+        cast_kind = complex
+    for it, er in iter:
+        assert np.allclose(it, 1, rtol=0.01)
+        assert np.allclose(er, np.sqrt(2), rtol=0.01)
+        assert isinstance(it[0], cast_kind)
+        assert isinstance(er[0], cast_kind)
+
+
+def test_avg_pds_imperfect_lc_size():
+    times = np.arange(100)
+    fluxes = np.ones(100).astype(float)
+    gti = np.asarray([[-0.5, 99.5]])
+    segment_size = 5.99
+    dt = 1
+    res = avg_pds_from_events(times, gti, segment_size, dt, fluxes=fluxes)
+    assert res.meta["segment_size"] == 5
+    assert res.meta["dt"] == 1
+
+
+def test_avg_cs_imperfect_lc_size():
+    times1 = times2 = np.arange(100)
+    fluxes1 = np.ones(100).astype(float)
+    fluxes2 = np.ones(100).astype(float)
+    gti = np.asarray([[-0.5, 99.5]])
+    segment_size = 5.99
+    dt = 1
+    res = avg_cs_from_events(
+        times1, times2, gti, segment_size, dt, fluxes1=fluxes1, fluxes2=fluxes2
+    )
+    assert res.meta["segment_size"] == 5
+    assert res.meta["dt"] == 1
+
+
 class TestCoherence(object):
     @classmethod
     def setup_class(cls):