Repository: StingraySoftware/stingray
Pull request #808 (2024-03-11 16:29:21)

Speed up fourier py
In `generate_indices_of_segment_boundaries_unbinned`:

- [x] Avoid unnecessary check for array sortedness. This saves ~1/3 of computing time in large, memory mapped arrays
- [x] Use `numpy.searchsorted` only once

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/StingraySoftware/stingray/pull/808?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 79.38%. Comparing base [(`143d7e5`)](https://app.codecov.io/gh/StingraySoftware/stingray/commit/143d7e5601803bdcf050283a715fd5570ce15e55?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware) to head [(`533c372`)](https://app.codecov.io/gh/StingraySoftware/stingray/pull/808?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware).


<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main     #808       +/-   ##
===========================================
- Coverage   96.41%   79.38%   -17.03%     
===========================================
  Files          45       45               
  Lines        8989     8990        +1     
===========================================
- Hits         8667     7137     -1530     
- Misses        322     1853     +1531     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/StingraySoftware/stingray/pull/808?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=StingraySoftware).


Patch:
diff --git a/docs/changes/808.feature.rst b/docs/changes/808.feature.rst
new file mode 100644
index 00000000..d33b9ae0
--- /dev/null
+++ b/docs/changes/808.feature.rst
@@ -0,0 +1,1 @@
+Speed up computation of pds for large arrays
diff --git a/stingray/base.py b/stingray/base.py
index 6d9451eb..717630d7 100644
--- a/stingray/base.py
+++ b/stingray/base.py
@@ -1127,6 +1127,10 @@ class StingrayTimeseries(StingrayObject):
     ephem : str
         The JPL ephemeris used to barycenter the data, if any (e.g. DE430)
 
+    skip_checks : bool
+        Skip checks on the time array. Useful when the user is reasonably sure that the
+        input data are valid.
+
     **other_kw :
         Used internally. Any other keyword arguments will be set as attributes of the object.
 
@@ -1176,6 +1180,7 @@ class StingrayTimeseries(StingrayObject):
         ephem: str = None,
         timeref: str = None,
         timesys: str = None,
+        skip_checks: bool = False,
         **other_kw,
     ):
         StingrayObject.__init__(self)
@@ -1198,6 +1203,12 @@ class StingrayTimeseries(StingrayObject):
             if self.time.shape[0] != new_arr.shape[0]:
                 raise ValueError(f"Lengths of time and {kw} must be equal.")
             setattr(self, kw, new_arr)
+        from .utils import is_sorted
+
+        if not skip_checks:
+            if self.time is not None and not is_sorted(self.time):
+                warnings.warn("The time array is not sorted. Sorting it now.")
+                self.sort(inplace=True)
 
     @property
     def time(self):
@@ -2151,7 +2162,7 @@ class StingrayTimeseries(StingrayObject):
         --------
         >>> time = [2, 1, 3]
         >>> count = [200, 100, 300]
-        >>> ts = StingrayTimeseries(time, array_attrs={"counts": count}, dt=1)
+        >>> ts = StingrayTimeseries(time, array_attrs={"counts": count}, dt=1, skip_checks=True)
         >>> ts_new = ts.sort()
         >>> ts_new.time
         array([1, 2, 3])
diff --git a/stingray/events.py b/stingray/events.py
index 0ee2b011..ff98ba0e 100644
--- a/stingray/events.py
+++ b/stingray/events.py
@@ -145,6 +145,9 @@ class EventList(StingrayTimeseries):
     rmf_file : str, default None
         The file name of the RMF file to use for calibration.
 
+    skip_checks : bool, default False
+        Skip checks for the validity of the event list. Use with caution.
+
     **other_kw :
         Used internally. Any other keyword arguments will be ignored
 
@@ -211,6 +214,7 @@ class EventList(StingrayTimeseries):
         timeref=None,
         timesys=None,
         rmf_file=None,
+        skip_checks=False,
         **other_kw,
     ):
         if ncounts is not None:
@@ -243,6 +247,7 @@ class EventList(StingrayTimeseries):
             timeref=timeref,
             timesys=timesys,
             rmf_file=rmf_file,
+            skip_checks=skip_checks,
             **other_kw,
         )
 
@@ -502,7 +507,8 @@ class EventList(StingrayTimeseries):
 
         Examples
         --------
-        >>> events = EventList(time=[0, 2, 1], energy=[0.3, 2, 0.5], pi=[3, 20, 5])
+        >>> events = EventList(time=[0, 2, 1], energy=[0.3, 2, 0.5], pi=[3, 20, 5],
+        ...                    skip_checks=True)
         >>> e1 = events.sort()
         >>> assert np.allclose(e1.time, [0, 1, 2])
         >>> assert np.allclose(e1.energy, [0.3, 0.5, 2])
diff --git a/stingray/fourier.py b/stingray/fourier.py
index fec7500c..8c52d384 100644
--- a/stingray/fourier.py
+++ b/stingray/fourier.py
@@ -1200,6 +1200,15 @@ def _which_segment_idx_fun(binned=False, dt=None):
     # Make function interface equal (fluxes gets ignored)
     if not binned:
         fun = generate_indices_of_segment_boundaries_unbinned
+
+        # Define a new function, make sure that, by default, the sort check
+        # is disabled.
+        def fun(*args, **kwargs):
+            check_sorted = kwargs.pop("check_sorted", False)
+            return generate_indices_of_segment_boundaries_unbinned(
+                *args, check_sorted=check_sorted, **kwargs
+            )
+
     else:
         # Define a new function, so that we can pass the correct dt as an
         # argument.
@@ -1335,11 +1344,10 @@ def get_flux_iterable_from_segments(
             yield None
             continue
         if not binned:
-            event_times = times[idx0:idx1]
             # astype here serves to avoid integer rounding issues in Windows,
             # where long is a 32-bit integer.
             cts = histogram(
-                (event_times - s).astype(float), bins=n_bin, range=[0, segment_size]
+                (times[idx0:idx1] - s).astype(float), bins=n_bin, range=[0, segment_size]
             ).astype(float)
             cts = np.array(cts)
         else:
diff --git a/stingray/gti.py b/stingray/gti.py
index 4f633080..007638e6 100644
--- a/stingray/gti.py
+++ b/stingray/gti.py
@@ -1565,7 +1565,7 @@ def generate_indices_of_gti_boundaries(times, gti, dt=0):
         yield s, e, idx0, idx1
 
 
-def generate_indices_of_segment_boundaries_unbinned(times, gti, segment_size):
+def generate_indices_of_segment_boundaries_unbinned(times, gti, segment_size, check_sorted=True):
     """
     Get the indices of events from different segments of the observation.
 
@@ -1581,6 +1581,11 @@ def generate_indices_of_segment_boundaries_unbinned(times, gti, segment_size):
     segment_size : float
         Length of segments.
 
+    Other Parameters
+    ----------------
+    check_sorted : bool, default True
+        If True, checks that the time array is sorted.
+
     Yields
     ------
     t0: float
@@ -1598,7 +1603,8 @@ def generate_indices_of_segment_boundaries_unbinned(times, gti, segment_size):
     --------
     >>> times = [0.1, 0.2, 0.5, 0.8, 1.1]
     >>> gtis = [[0, 0.55], [0.6, 2.1]]
-    >>> vals = generate_indices_of_segment_boundaries_unbinned(times, gtis, 0.5)
+    >>> vals = generate_indices_of_segment_boundaries_unbinned(
+    ...    times, gtis, 0.5)
     >>> v0 = next(vals)
     >>> assert np.allclose(v0[:2], [0, 0.5])
     >>> # Note: 0.5 is not included in the interval
@@ -1613,10 +1619,23 @@ def generate_indices_of_segment_boundaries_unbinned(times, gti, segment_size):
 
     start, stop = time_intervals_from_gtis(gti, segment_size)
 
-    assert is_sorted(times), "Array is not sorted"
+    if check_sorted:
+        assert is_sorted(times), "Array is not sorted"
+
+    all_times = np.sort(
+        np.array(  # Wrap in a numpy array
+            list(  # Transform into a proper iterable. Set is not recognized by np.array
+                set(  # Only unique values. Start and stop have a lot of overlap
+                    np.concatenate([start, stop])  # Concatenate start and stop
+                )
+            )
+        )
+    )
 
-    startidx = np.asarray(np.searchsorted(times, start))
-    stopidx = np.asarray(np.searchsorted(times, stop))
+    idxs = np.searchsorted(times, all_times)
+    idx_dict = dict([(s, a) for s, a in zip(all_times, idxs)])
+    startidx = np.asarray([idx_dict[s] for s in start])
+    stopidx = np.asarray([idx_dict[s] for s in stop])
 
     for s, e, idx0, idx1 in zip(start, stop, startidx, stopidx):
         yield s, e, idx0, idx1
diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index ca597f18..ad7fed6c 100644
--- a/stingray/lightcurve.py
+++ b/stingray/lightcurve.py
@@ -4,6 +4,7 @@ Definition of :class::class:`Lightcurve`.
 :class::class:`Lightcurve` is used to create light curves out of photon counting data
 or to save existing light curves in a class that's easy to use.
 """
+
 import os
 import logging
 import warnings


Test patch:
diff --git a/stingray/tests/test_base.py b/stingray/tests/test_base.py
index 1766bec7..cb3c7143 100644
--- a/stingray/tests/test_base.py
+++ b/stingray/tests/test_base.py
@@ -850,9 +850,10 @@ class TestStingrayTimeseries:
         bleh = [4, 1, 2, 0.5]
         mjdref = 57000
 
-        lc = StingrayTimeseries(
-            times, array_attrs={"blah": blah, "_bleh": bleh}, dt=1, mjdref=mjdref
-        )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            lc = StingrayTimeseries(
+                times, array_attrs={"blah": blah, "_bleh": bleh}, dt=1, mjdref=mjdref
+            )
 
         lc_new = lc.sort()
 
@@ -1223,12 +1224,14 @@ class TestJoinEvents:
 
     def test_overlapping_join_infer(self):
         """Join two non-overlapping event lists."""
-        ts = StingrayTimeseries(
-            time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]]
-        )
-        ts_other = StingrayTimeseries(
-            time=[5.1, 7, 6.1, 6.11, 10.1], energy=[2, 3, 8, 1, 2], gti=[[5, 7], [8, 10]]
-        )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ts = StingrayTimeseries(
+                time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]]
+            )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ts_other = StingrayTimeseries(
+                time=[5.1, 7, 6.1, 6.11, 10.1], energy=[2, 3, 8, 1, 2], gti=[[5, 7], [8, 10]]
+            )
         ts_new = ts.join(ts_other, strategy="infer")
 
         assert (ts_new.time == np.array([1, 1.1, 5, 5.1, 6, 6.1, 6.11, 7, 10, 10.1])).all()
@@ -1237,15 +1240,20 @@ class TestJoinEvents:
 
     def test_overlapping_join_change_mjdref(self):
         """Join two non-overlapping event lists."""
-        ts = StingrayTimeseries(
-            time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]], mjdref=57001
-        )
-        ts_other = StingrayTimeseries(
-            time=np.asarray([5.1, 7, 6.1, 6.11, 10.1]) + 86400,
-            energy=[2, 3, 8, 1, 2],
-            gti=np.asarray([[5, 7], [8, 10]]) + 86400,
-            mjdref=57000,
-        )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ts = StingrayTimeseries(
+                time=[1, 1.1, 10, 6, 5],
+                energy=[10, 6, 3, 11, 2],
+                gti=[[1, 3], [5, 6]],
+                mjdref=57001,
+            )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ts_other = StingrayTimeseries(
+                time=np.asarray([5.1, 7, 6.1, 6.11, 10.1]) + 86400,
+                energy=[2, 3, 8, 1, 2],
+                gti=np.asarray([[5, 7], [8, 10]]) + 86400,
+                mjdref=57000,
+            )
         with pytest.warns(UserWarning, match="Attribute mjdref is different"):
             ts_new = ts.join(ts_other, strategy="intersection")
 
diff --git a/stingray/tests/test_events.py b/stingray/tests/test_events.py
index d0fbb681..3f909edb 100644
--- a/stingray/tests/test_events.py
+++ b/stingray/tests/test_events.py
@@ -488,10 +488,12 @@ class TestJoinEvents:
 
     def test_overlapping_join_infer(self):
         """Join two non-overlapping event lists."""
-        ev = EventList(time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]])
-        ev_other = EventList(
-            time=[5.1, 7, 6.1, 6.11, 10.1], energy=[2, 3, 8, 1, 2], gti=[[5, 7], [8, 10]]
-        )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ev = EventList(time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]])
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ev_other = EventList(
+                time=[5.1, 7, 6.1, 6.11, 10.1], energy=[2, 3, 8, 1, 2], gti=[[5, 7], [8, 10]]
+            )
         ev_new = ev.join(ev_other, strategy="infer")
 
         assert (ev_new.time == np.array([1, 1.1, 5, 5.1, 6, 6.1, 6.11, 7, 10, 10.1])).all()
@@ -500,15 +502,20 @@ class TestJoinEvents:
 
     def test_overlapping_join_change_mjdref(self):
         """Join two non-overlapping event lists."""
-        ev = EventList(
-            time=[1, 1.1, 10, 6, 5], energy=[10, 6, 3, 11, 2], gti=[[1, 3], [5, 6]], mjdref=57001
-        )
-        ev_other = EventList(
-            time=np.asarray([5.1, 7, 6.1, 6.11, 10.1]) + 86400,
-            energy=[2, 3, 8, 1, 2],
-            gti=np.asarray([[5, 7], [8, 10]]) + 86400,
-            mjdref=57000,
-        )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ev = EventList(
+                time=[1, 1.1, 10, 6, 5],
+                energy=[10, 6, 3, 11, 2],
+                gti=[[1, 3], [5, 6]],
+                mjdref=57001,
+            )
+        with pytest.warns(UserWarning, match="The time array is not sorted."):
+            ev_other = EventList(
+                time=np.asarray([5.1, 7, 6.1, 6.11, 10.1]) + 86400,
+                energy=[2, 3, 8, 1, 2],
+                gti=np.asarray([[5, 7], [8, 10]]) + 86400,
+                mjdref=57000,
+            )
         with pytest.warns(UserWarning, match="Attribute mjdref is different"):
             ev_new = ev.join(ev_other, strategy="intersection")
 
diff --git a/stingray/tests/test_lombscargle.py b/stingray/tests/test_lombscargle.py
index 4a8f76f5..25f0d192 100644
--- a/stingray/tests/test_lombscargle.py
+++ b/stingray/tests/test_lombscargle.py
@@ -213,8 +213,8 @@ class TestLombScargleCrossspectrum:
             func()
 
     def test_no_dt(self):
-        el1 = EventList(self.lc1.counts, self.lc1.time, dt=None)
-        el2 = EventList(self.lc2.counts, self.lc2.time, dt=None)
+        el1 = EventList(self.lc1.time, self.lc1.counts, dt=None)
+        el2 = EventList(self.lc2.time, self.lc2.counts, dt=None)
         with pytest.raises(ValueError):
             lscs = LombScargleCrossspectrum(el1, el2)