feat: Introduce ResamplingSlicingReader and enhance transform handling in GeoAnimatorManager

Author: dwr-psandhu

dvue/animator/__init__.py

@@ -22,6 +22,7 @@
     BufferedSlicingReader,
     TransformSpec,
     StreamingTransformedSlicingReader,
+    ResamplingSlicingReader,
     DiffSlicingReader,
     RawSequentialBuffer,
 )
@@ -32,6 +33,7 @@
     "BufferedSlicingReader",
     "TransformSpec",
     "StreamingTransformedSlicingReader",
+    "ResamplingSlicingReader",
     "DiffSlicingReader",
     "RawSequentialBuffer",
     "GeoAnimatorManager",

dvue/animator/multi_ui.py

@@ -37,7 +37,8 @@
 )
 from bokeh.plotting import figure as bk_figure
 
-from .reader import SlicingReader, DiffSlicingReader, BufferedSlicingReader, RawSequentialBuffer
+from .reader import SlicingReader, DiffSlicingReader, BufferedSlicingReader, RawSequentialBuffer, \
+    ResamplingSlicingReader
 from .ui import (
     CURATED_COLORMAPS,
     CURATED_COLORMAPS_WITH_SEP,
@@ -827,57 +828,62 @@ def _setup_reader(self, base: SlicingReader, transform_name: str) -> SlicingRead
         chunk_size = self._buffer_chunk_size
         min_chunk = 50
         max_chunk = 2000
+        target_buf_s = 10.0
         if (
             transform_name
             and transform_name != "none"
             and transform_name in self._transform_options
         ):
             spec_or_fn = self._transform_options[transform_name]
-            if isinstance(spec_or_fn, TransformSpec):
-                freq_nanos = 0
+            if not isinstance(spec_or_fn, TransformSpec):
+                raise TypeError(
+                    f"transform_options value for {transform_name!r} must be a "
+                    f"TransformSpec, got {type(spec_or_fn).__name__}. "
+                    "Use make_resample_transform(), make_moving_average_transform(), "
+                    "or make_godin_transform() from dsm2ui.animate."
+                )
+            freq_nanos = 0
+            try:
+                freq_nanos = int(
+                    pd.tseries.frequencies.to_offset(
+                        reader.time_index.freq
+                    ).nanos
+                )
+            except (AttributeError, TypeError):
+                pass
+
+            if spec_or_fn.kind == "aggregate":
+                filter_spec = spec_or_fn.filter_spec
+                if filter_spec is not None:
+                    try:
+                        raw_overlap = filter_spec.get_overlap(freq_nanos)
+                    except (AttributeError, TypeError):
+                        raw_overlap = 0
+                    if raw_overlap > 0:
+                        reader = RawSequentialBuffer(reader)
+                    reader = StreamingTransformedSlicingReader(reader, filter_spec)
+                reader = ResamplingSlicingReader(
+                    reader,
+                    spec_or_fn.output_freq,
+                    spec_or_fn.resample_agg,
+                )
+                chunk_size = max(5, self._buffer_chunk_size // 20)
+                min_chunk = 2
+                max_chunk = 365
+                target_buf_s = 30.0
+            else:
                 try:
-                    freq_nanos = int(
-                        pd.tseries.frequencies.to_offset(
-                            reader.time_index.freq
-                        ).nanos
-                    )
                     raw_overlap = spec_or_fn.get_overlap(freq_nanos)
                 except (AttributeError, TypeError):
                     raw_overlap = 0
                 if raw_overlap > 0:
                     reader = RawSequentialBuffer(reader)
                 reader = StreamingTransformedSlicingReader(reader, spec_or_fn)
-                # Scale chunk_size for coarse aggregate transforms so the number
-                # of raw steps per chunk stays roughly constant.  See the same
-                # logic in GeoAnimatorManager._setup_reader for the rationale.
-                if (
-                    spec_or_fn.kind == "aggregate"
-                    and spec_or_fn.output_freq is not None
-                    and freq_nanos > 0
-                ):
-                    try:
-                        out_ns = int(
-                            pd.tseries.frequencies.to_offset(
-                                spec_or_fn.output_freq
-                            ).nanos
-                        )
-                        if out_ns > freq_nanos:
-                            ratio = out_ns // freq_nanos
-                            chunk_size = max(5, self._buffer_chunk_size // ratio)
-                            min_chunk = max(2, 50 // ratio)
-                            max_chunk = max(chunk_size, 2000 // ratio)
-                    except Exception:
-                        pass
-            else:
-                raise TypeError(
-                    f"transform_options value for {transform_name!r} must be a "
-                    f"TransformSpec, got {type(spec_or_fn).__name__}. "
-                    "Use make_resample_transform(), make_moving_average_transform(), "
-                    "or make_godin_transform() from dsm2ui.animate."
-                )
+
         return BufferedSlicingReader(
             reader, chunk_size=chunk_size, prefetch=True,
             adaptive=True, min_chunk_size=min_chunk, max_chunk_size=max_chunk,
+            target_buffer_seconds=target_buf_s,
         )
 
     def _get_diff_reader(self) -> SlicingReader:
@@ -1537,6 +1543,27 @@ def _compute() -> None:
             for _tcb in self._extra_transform_callbacks:
                 _tcb(_new_spec)
 
+            # ── Pre-load both maps in the background thread ──────────────────
+            # _apply_frame (and its _update_map_a / _update_map_b calls)
+            # triggers cold-start HDF5 reads on the Tornado IOLoop, which
+            # freezes the browser for the duration of the chunk load.
+            # Pre-fetch both maps here, off the IOLoop, then pass the
+            # pre-computed values to _apply for pure Bokeh model mutations.
+            ts = ti[nearest_idx]
+            ts_str = ts.strftime("%Y-%m-%d %H:%M")
+            _geo_ids_a_snap = self._geo_ids_a
+            _geo_ids_b_snap = self._geo_ids_b
+            try:
+                _sa = new_reader_a.get_slice_nearest(ts)
+                vals_a = _sa.reindex(_geo_ids_a_snap).fillna(np.nan).tolist()
+            except Exception:
+                vals_a = [float("nan")] * len(_geo_ids_a_snap)
+            try:
+                _sb = new_reader_b.get_slice_nearest(ts)
+                vals_b = _sb.reindex(_geo_ids_b_snap).fillna(np.nan).tolist()
+            except Exception:
+                vals_b = [float("nan")] * len(_geo_ids_b_snap)
+
             def _apply() -> None:
                 self._reader_a = new_reader_a
                 self._reader_b = new_reader_b
@@ -1552,9 +1579,12 @@ def _apply() -> None:
                 finally:
                     self._syncing = False
 
-                ts_str = ti[nearest_idx].strftime("%Y-%m-%d %H:%M")
                 self._time_div.text = f"<b>{ts_str}</b>"
-                self._apply_frame(nearest_idx, ts_str)
+                # Pure Bokeh mutations — no I/O on the IOLoop.
+                self._render_map_a_vals(ts, ts_str, vals_a)
+                self._render_map_b_vals(ts, ts_str, vals_b)
+                for _cb in self._extra_frame_callbacks:
+                    _cb(ts)
 
                 for pane in (self._pane_a, self._pane_b, self._pane_diff):
                     pane.loading = False

dvue/animator/reader.py

@@ -710,13 +710,19 @@ class TransformSpec:
     :func:`~dsm2ui.animate.make_godin_transform` for production examples.
     """
 
-    def __init__(self, transform_fn, kind: str, get_overlap, output_freq=None):
+    def __init__(self, transform_fn, kind: str, get_overlap, output_freq=None,
+                 filter_spec=None, resample_agg="mean"):
         if kind not in ("convolution", "aggregate"):
             raise ValueError(f"kind must be 'convolution' or 'aggregate', got {kind!r}")
         self.transform_fn = transform_fn
         self.kind = kind
         self.get_overlap = get_overlap   # callable(freq_nanos: int) -> int
         self.output_freq = output_freq
+        # Optional: for kind="aggregate", carry the convolution pre-filter spec
+        # so _setup_reader can stack STSR(filter_spec) → ResamplingSlicingReader
+        # instead of using STSR's slow aggregate code path.
+        self.filter_spec = filter_spec   # TransformSpec or None
+        self.resample_agg = resample_agg # aggregation for ResamplingSlicingReader
 
 
 class StreamingTransformedSlicingReader(SlicingReader):
@@ -894,6 +900,113 @@ def __exit__(self, *args):
         self.close()
 
 
+# ---------------------------------------------------------------------------
+# Resampling reader
+# ---------------------------------------------------------------------------
+
+class ResamplingSlicingReader(SlicingReader):
+    """Wraps any :class:`SlicingReader` and resamples its output to a coarser frequency.
+
+    Unlike :class:`StreamingTransformedSlicingReader` with ``kind="aggregate"``,
+    this reader requests only the inner steps needed for each output window —
+    never a large overlap-padded chunk.  It is the efficient outer layer for
+    composed transforms::
+
+        BufferedSlicingReader(chunk_size=20)
+          └── ResamplingSlicingReader(output_freq="D")
+              └── StreamingTransformedSlicingReader(godin, kind="convolution")
+                  └── RawSequentialBuffer
+                      └── BaseReader
+
+    Parameters
+    ----------
+    inner : SlicingReader
+        The pre-filtered source (Godin STSR, rolling STSR, or a raw reader).
+    output_freq : str
+        Pandas offset string for the desired output step (e.g. ``"D"``, ``"h"``).
+    agg : {"mean", "max", "min", "sum"}, optional
+        Aggregation applied when resampling.  Default ``"mean"``.
+    """
+
+    def __init__(self, inner: SlicingReader, output_freq: str, agg: str = "mean") -> None:
+        self._inner = inner
+        self._output_freq = output_freq
+        self._agg = agg
+
+        raw_ti = inner.time_index
+        out_freq = pd.tseries.frequencies.to_offset(output_freq)
+        out_start = raw_ti[0].floor(output_freq)
+        out_end = raw_ti[-1].floor(output_freq)
+        out_ti = pd.date_range(out_start, out_end, freq=output_freq)
+        if out_ti.freq is None:
+            out_ti = out_ti.copy()
+            out_ti.freq = out_freq
+
+        self._vmin = inner.vmin
+        self._vmax = inner.vmax
+
+        super().__init__(out_ti)
+
+    @property
+    def vmin(self) -> float:
+        return self._vmin
+
+    @property
+    def vmax(self) -> float:
+        return self._vmax
+
+    def get_slice_range(self, start_out: int, end_out: int) -> pd.DataFrame:
+        n_out = len(self._time_index)
+        start_out = max(0, min(start_out, n_out))
+        end_out = max(start_out, min(end_out, n_out))
+        if start_out >= end_out:
+            return pd.DataFrame(index=self._time_index[start_out:end_out])
+
+        out_freq = pd.tseries.frequencies.to_offset(self._output_freq)
+        out_start_ts = self._time_index[start_out]
+        last_out_ts = self._time_index[end_out - 1]
+        raw_upper_ts = last_out_ts + out_freq
+
+        raw_ti = self._inner.time_index
+        raw_start = int(raw_ti.searchsorted(out_start_ts, side="left"))
+        raw_end = int(raw_ti.searchsorted(raw_upper_ts, side="left"))
+        raw_end = min(len(raw_ti), raw_end)
+
+        if raw_start >= raw_end:
+            return pd.DataFrame(index=self._time_index[start_out:end_out])
+
+        raw_df = self._inner.get_slice_range(raw_start, raw_end)
+        if raw_df.empty:
+            return pd.DataFrame(index=self._time_index[start_out:end_out])
+
+        resampled = getattr(raw_df.resample(self._output_freq), self._agg)()
+        if resampled.index.freq is None:
+            resampled.index.freq = out_freq
+
+        want = self._time_index[start_out:end_out]
+        result = resampled.reindex(want)
+        if result.index.freq is None:
+            result.index.freq = self._time_index.freq
+        return result
+
+    def get_slice(self, timestamp: pd.Timestamp) -> pd.Series:
+        i = int(self._time_index.get_indexer([timestamp], method="nearest")[0])
+        chunk = self.get_slice_range(i, i + 1)
+        if len(chunk) == 0:
+            return pd.Series(dtype=float)
+        return chunk.iloc[0].astype(float)
+
+    def close(self) -> None:
+        if hasattr(self._inner, "close"):
+            self._inner.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
 # ---------------------------------------------------------------------------
 # Raw sequential look-ahead buffer
 # ---------------------------------------------------------------------------