Repository: StingraySoftware/stingray (python)

Issue:
analyze_segments fails when some interval gives invalid data
### Description of the Bug

Inside `StingrayTimeseries.analyze_segments`, there is a `results = np.array(results)` operation that fails whenever the rows of results have different shapes. In particular, it fails when the segment has not enough data and the row is given `nan` by default. 

### Steps/Code to Replicate the Bug

NA

### Expected Results

NA

### Actual Results

NA

Pull request description:
Fix operation on segment when result is NaN
This is a quick fix to #821

Resolve #821


Patch:
diff --git a/docs/changes/822.bugfix.rst b/docs/changes/822.bugfix.rst
new file mode 100644
index 00000000..06f526c5
--- /dev/null
+++ b/docs/changes/822.bugfix.rst
@@ -0,0 +1,1 @@
+Fix case when analyze_segments has an invalid segment
diff --git a/docs/notebooks b/docs/notebooks
index 0e34107a..2db8e50b 160000
--- a/docs/notebooks
+++ b/docs/notebooks
@@ -1,1 +1,1 @@
-Subproject commit 0e34107ac872db8ed9b4acfcc75b97edf07f3ccb
+Subproject commit 2db8e50be46e3d5cddebb2c15647c45d63de4a92
diff --git a/setup.cfg b/setup.cfg
index 82852c01..9a480a5e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -116,6 +116,7 @@ filterwarnings =
     ignore:.*datetime.datetime.utcfromtimestamp.*:DeprecationWarning
     ignore:.*__array_wrap__ must accept context and return_scalar arguments:DeprecationWarning
     ignore:.*Pyarrow:
+    ignore:.*Creating an ndarray from ragged nested sequences:
 
 ;addopts = --disable-warnings
 
diff --git a/stingray/base.py b/stingray/base.py
index 7e8a3533..49c4dc4f 100644
--- a/stingray/base.py
+++ b/stingray/base.py
@@ -2582,6 +2582,8 @@ class StingrayTimeseries(StingrayObject):
     def analyze_segments(self, func, segment_size, fraction_step=1, **kwargs):
         """Analyze segments of the light curve with any function.
 
+        Intervals with less than one data point are skipped.
+
         Parameters
         ----------
         func : function
@@ -2609,8 +2611,10 @@ class StingrayTimeseries(StingrayObject):
             Lower time boundaries of all time segments.
         stop_times : array
             upper time boundaries of all segments.
-        result : array of N elements
-            The result of ``func`` for each segment of the light curve
+        result : list of N elements
+            The result of ``func`` for each segment of the light curve. If the function
+            returns multiple outputs, they are returned as a list of arrays.
+            If a given interval has not enough data for a calculation, ``None`` is returned.
 
         Examples
         --------
@@ -2649,23 +2653,37 @@ class StingrayTimeseries(StingrayObject):
             stop = np.searchsorted(self.time, stop_times)
 
         results = []
+
+        n_outs = 1
         for i, (st, sp, tst, tsp) in enumerate(zip(start, stop, start_times, stop_times)):
             if sp - st <= 1:
                 warnings.warn(
-                    f"Segment {i} ({tst}--{tsp}) has one data point or less. Skipping it."
+                    f"Segment {i} ({tst}--{tsp}) has one data point or less. Skipping it "
                 )
-                res = np.nan
-            else:
-                lc_filt = self[st:sp]
-                lc_filt.gti = np.asanyarray([[tst, tsp]])
+                continue
+            lc_filt = self[st:sp]
+            lc_filt.gti = np.asanyarray([[tst, tsp]])
 
-                res = func(lc_filt, **kwargs)
+            res = func(lc_filt, **kwargs)
             results.append(res)
+            if isinstance(res, Iterable) and not isinstance(res, str):
+                n_outs = len(res)
+
+        # If the function returns multiple outputs, we need to separate them
+
+        if n_outs > 1:
+            outs = [[] for _ in range(n_outs)]
+            for res in results:
+                for i in range(n_outs):
+                    outs[i].append(res[i])
+            results = outs
 
-        results = np.array(results)
+        # Try to transform into a (possibly multi-dimensional) numpy array
+        try:
+            results = np.array(results)
+        except ValueError:  # pragma: no cover
+            pass
 
-        if len(results.shape) == 2:
-            results = [results[:, i] for i in range(results.shape[1])]
         return start_times, stop_times, results
 
     def analyze_by_gti(self, func, fraction_step=1, **kwargs):
diff --git a/stingray/events.py b/stingray/events.py
index 6cd11e42..c59ec49c 100644
--- a/stingray/events.py
+++ b/stingray/events.py
@@ -797,6 +797,9 @@ class EventList(StingrayTimeseries):
             en1_ct = np.count_nonzero(mask1)
             en2_ct = np.count_nonzero(mask2)
 
+            if en1_ct == 0 or en2_ct == 0:
+                warnings.warn("No counts in one of the energy ranges. Returning NaN")
+                return np.nan, np.nan
             color = en2_ct / en1_ct
             color_err = color * (np.sqrt(en1_ct) / en1_ct + np.sqrt(en2_ct) / en2_ct)
             return color, color_err


Test patch:
diff --git a/stingray/tests/test_base.py b/stingray/tests/test_base.py
index 145d194a..d43f7911 100644
--- a/stingray/tests/test_base.py
+++ b/stingray/tests/test_base.py
@@ -1497,17 +1497,18 @@ class TestAnalyzeChunks(object):
 
         assert ts.estimate_segment_size(100, min_samples=40) == 8.0
 
-    def test_analyze_segments_bad_intv(self):
+    @pytest.mark.parametrize("n_outs", [0, 1, 2, 3])
+    def test_analyze_segments_bad_intv(self, n_outs):
         ts = StingrayTimeseries(time=np.arange(10), dt=1, gti=[[-0.5, 0.5], [1.5, 10.5]])
 
         def func(x):
-            return np.size(x.time)
+            if n_outs == 0:
+                return np.size(x.time)
+            return [np.size(x.time) for _ in range(n_outs)]
 
         # I do not specify the segment_size, which means results will be calculated per-GTI
         with pytest.warns(UserWarning, match="has one data point or less."):
-            _, _, results = ts.analyze_segments(func, segment_size=None)
-        # the first GTI contains only one bin, the result will be invalid
-        assert np.isnan(results[0])
+            ts.analyze_segments(func, segment_size=None)
 
     def test_analyze_segments_by_gti(self):
         ts = StingrayTimeseries(time=np.arange(11), dt=1, gti=[[-0.5, 5.5], [6.5, 10.5]])
diff --git a/stingray/tests/test_events.py b/stingray/tests/test_events.py
index 2ae34d40..cfc272e2 100644
--- a/stingray/tests/test_events.py
+++ b/stingray/tests/test_events.py
@@ -644,6 +644,12 @@ class TestColors(object):
         assert np.allclose(start, np.arange(10) * 10000)
         assert np.allclose(stop, np.arange(1, 11) * 10000)
 
+    def test_colors_missing_energies(self):
+        events = copy.deepcopy(self.events)
+        events.filter_energy_range([0, 3], inplace=True)
+        with pytest.warns(UserWarning, match="No counts in one of the energy ranges"):
+            events.get_color_evolution([[0, 3], [4, 6]], 10000)
+
     def test_colors_no_segment(self):
         start, stop, colors, color_errs = self.events.get_color_evolution([[0, 3], [4, 6]])
         # 50000 / 50000 = 1