From 6cacd7d18b1714b647fbf7d5daf32770beeae9fb Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Mon, 28 Nov 2022 10:09:45 -0500
Subject: [PATCH] Make `Viz.slice_from_time()` take input array

Probably means it doesn't need to be a `Flume` method but it's
convenient to expect the caller to pass in the `np.ndarray` with
a `'time'` field instead of a `timeframe: str` arg; also, return the
slice mask instead of the sliced array as output (again allowing the
caller to do any slicing). Also, handle the slice-outside-time-range
case by just returning the entire index range with a `None` mask.

Adjust `Viz.view_data()` to instead do timeframe (for rt vs. hist shm
array) lookup and equiv array slicing with the returned mask.
---
 piker/data/flows.py | 88 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 64 insertions(+), 24 deletions(-)

diff --git a/piker/data/flows.py b/piker/data/flows.py
index 9bb27230..655f9ff7 100644
--- a/piker/data/flows.py
+++ b/piker/data/flows.py
@@ -236,29 +236,43 @@ class Flume(Struct):
 
     def slice_from_time(
         self,
-        array: np.ndarray,
+        arr: np.ndarray,
         start_t: float,
         stop_t: float,
-        timeframe_s: int = 1,
-        return_data: bool = False,
 
-    ) -> np.ndarray:
+    ) -> tuple[
+        slice,
+        slice,
+        np.ndarray | None,
+    ]:
         '''
-        Slice an input struct array providing only datums
-        "in view" of this chart.
+        Slice an input struct array to a time range and return the absolute
+        and "readable" slices for that array as well as the indexing mask
+        for the caller to use to slice the input array if needed.
 
         '''
-        arr = {
-            1: self.rt_shm.array,
-            60: self.hist_shm.arry,
-        }[timeframe_s]
-
         times = arr['time']
-        index = array['index']
+        index = arr['index']
+
+        if (
+            start_t < 0
+            or start_t >= stop_t
+        ):
+            return (
+                slice(
+                    index[0],
+                    index[-1],
+                ),
+                slice(
+                    0,
+                    len(arr),
+                ),
+                None,
+            )
 
         # use advanced indexing to map the
         # time range to the index range.
-        mask = (
+        mask: np.ndarray = (
             (times >= start_t)
             &
             (times < stop_t)
@@ -273,7 +287,24 @@ class Flume(Struct):
         # ]
 
         i_by_t = index[mask]
-        i_0 = i_by_t[0]
+        try:
+            i_0 = i_by_t[0]
+        except IndexError:
+            if (
+                start_t < times[0]
+                or stop_t >= times[-1]
+            ):
+                return (
+                    slice(
+                        index[0],
+                        index[-1],
+                    ),
+                    slice(
+                        0,
+                        len(arr),
+                    ),
+                    None,
+                )
 
         abs_slc = slice(
             i_0,
@@ -285,17 +316,12 @@ class Flume(Struct):
             0,
             i_by_t[-1] - i_0,
         )
-        if not return_data:
-            return (
-                abs_slc,
-                read_slc,
-            )
 
         # also return the readable data from the timerange
         return (
             abs_slc,
             read_slc,
-            arr[mask],
+            mask,
         )
 
     def view_data(
@@ -304,18 +330,32 @@ class Flume(Struct):
         timeframe_s: int = 1,
 
     ) -> np.ndarray:
+        '''
+        Return sliced-to-view source data along with absolute
+        (``ShmArray._array['index']``) and read-relative
+        (``ShmArray.array``) slices.
 
+        '''
         # get far-side x-indices plot view
         vr = plot.viewRect()
 
+        if timeframe_s > 1:
+            arr = self.hist_shm.array
+        else:
+            arr = self.rt_shm.array
+
         (
             abs_slc,
-            buf_slc,
-            iv_arr,
+            read_slc,
+            mask,
         ) = self.slice_from_time(
+            arr,
             start_t=vr.left(),
             stop_t=vr.right(),
             timeframe_s=timeframe_s,
-            return_data=True,
         )
-        return iv_arr
+        return (
+            abs_slc,
+            read_slc,
+            arr[mask] if mask is not None else arr,
+        )