From ec8679ad74a9316bbf544c86b5be1aceddaa60c4 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Mon, 23 Jan 2023 20:22:45 -0500
Subject: [PATCH] Add `Viz.median_from_range()`

A super snappy `numpy.median()` calculator (per input range) which we
slap an `lru_cache` on thanks to handy dunder method hacks for such
things on mutable types XD
---
 piker/ui/_dataviz.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/piker/ui/_dataviz.py b/piker/ui/_dataviz.py
index 494ef985..32f400a1 100644
--- a/piker/ui/_dataviz.py
+++ b/piker/ui/_dataviz.py
@@ -19,6 +19,7 @@ Data vizualization APIs
 
 '''
 from __future__ import annotations
+from functools import lru_cache
 from math import (
     ceil,
     floor,
@@ -282,6 +283,21 @@ class Viz(msgspec.Struct):  # , frozen=True):
         tuple[float, float],
     ] = {}
 
+    # cache of median calcs from input read slice hashes
+    # see `.median()`
+    _meds: dict[
+        int,
+        float,
+    ] = {}
+
+    # to make lru_cache-ing work, see
+    # https://docs.python.org/3/faq/programming.html#how-do-i-cache-method-calls
+    def __eq__(self, other):
+        return self._shm._token == other._shm._token
+
+    def __hash__(self):
+        return hash(self._shm._token)
+
     @property
     def shm(self) -> ShmArray:
         return self._shm
@@ -462,6 +478,19 @@ class Viz(msgspec.Struct):  # , frozen=True):
             mxmn,
         )
 
+    @lru_cache(maxsize=6116)
+    def median_from_range(
+        self,
+        start: int,
+        stop: int,
+
+    ) -> float:
+        in_view = self.shm.array[start:stop]
+        if self.is_ohlc:
+            return np.median(in_view['close'])
+        else:
+            return np.median(in_view[self.name])
+
     def view_range(self) -> tuple[int, int]:
         '''
         Return the start and stop x-indexes for the managed ``ViewBox``.