Sync view state between JS and Python (#448)

Closes #112



In #444 @batpad
mentioned

> It does seem to make the map a bit slow for me, and we should
investigate that, but in principle, this for me works as expected

While playing around with that PR, I also saw some view state lagging at
low zooms with lots of data rendered. In this screencast I'm smoothly
zooming in, but the rendering is very jagged, and sometimes zooms back
out slightly. **This screencast is of the previous PR**


https://github.com/developmentseed/lonboard/assets/15164633/bbe0a168-b76e-4563-bc30-828d648210db

My hypothesis of what's happening is that when rendering lots of data
there's a slight lag for deck to call the `onViewStateChange` handler.
And so deck will occasionally call the handler with "old" data and
that's where you sometimes get the "snap back" in the scrolling.

I think this is a really bad UX, and it seems not to happen when deck
manages the view state internally. So I think potentially the best UX
and performance solution is to decouple the view state from the model,
while still syncing from deck to Python in the debounced model handler.

This PR tries to implement this.

---------

Co-authored-by: Sanjay Bhangar <[email protected]>

Author: kylebarron

Diff for: docs/api/map.md

@@ -8,3 +8,6 @@ https://mkdocstrings.github.io/python/usage/configuration/members/#filters
       group_by_category: false
       show_bases: false
       filters:
+
+
+::: lonboard.models.ViewState

Diff for: lonboard/_map.py

@@ -13,6 +13,7 @@
 from lonboard._layer import BaseLayer
 from lonboard._viewport import compute_view
 from lonboard.basemap import CartoBasemap
+from lonboard.traits import DEFAULT_INITIAL_VIEW_STATE, ViewStateTrait
 from lonboard.types.map import MapKwargs
 
 if TYPE_CHECKING:
@@ -95,32 +96,28 @@ def __init__(
     _esm = bundler_output_dir / "index.js"
     _css = bundler_output_dir / "index.css"
 
-    _initial_view_state = traitlets.Dict().tag(sync=True)
+    view_state = ViewStateTrait()
     """
-    The initial view state of the map.
+    The view state of the map.
 
-    - Type: `dict`, optional
+    - Type: [`ViewState`][lonboard.models.ViewState]
     - Default: Automatically inferred from the data passed to the map.
 
-    The keys _must_ include:
+    You can initialize the map to a specific view state using this property:
 
-    - `longitude`: longitude at the map center.
-    - `latitude`: latitude at the map center.
-    - `zoom`: zoom level.
-
-    Keys may additionally include:
+    ```py
+    Map(
+        layers,
+        view_state={"longitude": -74.0060, "latitude": 40.7128, "zoom": 7}
+    )
+    ```
 
-    - `pitch` (float, optional) - pitch angle in degrees. Default `0` (top-down).
-    - `bearing` (float, optional) - bearing angle in degrees. Default `0` (north).
-    - `maxZoom` (float, optional) - max zoom level. Default `20`.
-    - `minZoom` (float, optional) - min zoom level. Default `0`.
-    - `maxPitch` (float, optional) - max pitch angle. Default `60`.
-    - `minPitch` (float, optional) - min pitch angle. Default `0`.
+    !!! note
 
-    Note that currently no camel-case/snake-case translation occurs for this method, and
-    so keys must be in camel case.
+        The properties of the view state are immutable. Use
+        [`set_view_state`][lonboard.Map.set_view_state] to modify a map's view state
+        once it's been initially rendered.
 
-    This API is not yet stabilized and may change in the future.
     """
 
     _height = traitlets.Int(default_value=DEFAULT_HEIGHT, allow_none=True).tag(
@@ -268,12 +265,51 @@ def __init__(
       global `parameters` when that layer is rendered.
     """
 
+    def set_view_state(
+        self,
+        *,
+        longitude: Optional[float] = None,
+        latitude: Optional[float] = None,
+        zoom: Optional[float] = None,
+        pitch: Optional[float] = None,
+        bearing: Optional[float] = None,
+    ) -> None:
+        """Set the view state of the map.
+
+        Any parameters that are unset will not be changed.
+
+        Other Args:
+            longitude: the new longitude to set on the map. Defaults to None.
+            latitude: the new latitude to set on the map. Defaults to None.
+            zoom: the new zoom to set on the map. Defaults to None.
+            pitch: the new pitch to set on the map. Defaults to None.
+            bearing: the new bearing to set on the map. Defaults to None.
+        """
+        view_state = (
+            self.view_state._asdict()  # type: ignore
+            if self.view_state is not None
+            else DEFAULT_INITIAL_VIEW_STATE
+        )
+
+        if longitude is not None:
+            view_state["longitude"] = longitude
+        if latitude is not None:
+            view_state["latitude"] = latitude
+        if zoom is not None:
+            view_state["zoom"] = zoom
+        if pitch is not None:
+            view_state["pitch"] = pitch
+        if bearing is not None:
+            view_state["bearing"] = bearing
+
+        self.view_state = view_state
+
     def fly_to(
         self,
         *,
         longitude: Union[int, float],
         latitude: Union[int, float],
-        zoom: int,
+        zoom: float,
         duration: int = 4000,
         pitch: Union[int, float] = 0,
         bearing: Union[int, float] = 0,
@@ -333,6 +369,6 @@ def to_html(
             drop_defaults=False,
         )
 
-    @traitlets.default("_initial_view_state")
+    @traitlets.default("view_state")
     def _default_initial_view_state(self):
         return compute_view(self.layers)

Diff for: lonboard/_serialization.py

@@ -1,13 +1,15 @@
 import math
 from io import BytesIO
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
 from numpy.typing import NDArray
 from traitlets import TraitError
 
+from lonboard.models import ViewState
+
 DEFAULT_PARQUET_COMPRESSION = "ZSTD"
 DEFAULT_PARQUET_COMPRESSION_LEVEL = 7
 DEFAULT_PARQUET_CHUNK_SIZE = 2**16
@@ -88,5 +90,12 @@ def validate_accessor_length_matches_table(accessor, table):
         raise TraitError("accessor must have same length as table")
 
 
+def serialize_view_state(data: Optional[ViewState], obj):
+    if data is None:
+        return None
+
+    return data._asdict()
+
+
 ACCESSOR_SERIALIZATION = {"to_json": serialize_accessor}
 TABLE_SERIALIZATION = {"to_json": serialize_table}

Diff for: lonboard/_viewport.py

@@ -62,6 +62,18 @@ def compute_view(layers: List[BaseLayer]):
     # When no geo column is found, bbox will have inf values
     try:
         zoom = bbox_to_zoom_level(bbox)
-        return {"longitude": center.x, "latitude": center.y, "zoom": zoom}
+        return {
+            "longitude": center.x,
+            "latitude": center.y,
+            "zoom": zoom,
+            "pitch": 0,
+            "bearing": 0,
+        }
     except OverflowError:
-        return {"longitude": center.x or 0, "latitude": center.y or 0, "zoom": 0}
+        return {
+            "longitude": center.x or 0,
+            "latitude": center.y or 0,
+            "zoom": 0,
+            "pitch": 0,
+            "bearing": 0,
+        }

Diff for: lonboard/models.py

@@ -0,0 +1,18 @@
+from typing import NamedTuple
+
+
+class ViewState(NamedTuple):
+    longitude: float
+    """Longitude at the map center"""
+
+    latitude: float
+    """Latitude at the map center."""
+
+    zoom: float
+    """Zoom level."""
+
+    pitch: float
+    """Pitch angle in degrees. `0` is top-down."""
+
+    bearing: float
+    """Bearing angle in degrees. `0` is north."""

Diff for: lonboard/traits.py

@@ -21,8 +21,18 @@
 from lonboard._serialization import (
     ACCESSOR_SERIALIZATION,
     TABLE_SERIALIZATION,
+    serialize_view_state,
 )
 from lonboard._utils import get_geometry_column_index
+from lonboard.models import ViewState
+
+DEFAULT_INITIAL_VIEW_STATE = {
+    "latitude": 10,
+    "longitude": 0,
+    "zoom": 0.5,
+    "bearing": 0,
+    "pitch": 0,
+}
 
 
 # This is a custom subclass of traitlets.TraitType because its `error` method ignores
@@ -822,3 +832,31 @@ def validate(
 
         self.error(obj, value)
         assert False
+
+
+class ViewStateTrait(FixedErrorTraitType):
+    allow_none = True
+    default_value = DEFAULT_INITIAL_VIEW_STATE
+
+    def __init__(
+        self: TraitType,
+        *args,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.tag(sync=True, to_json=serialize_view_state)
+
+    def validate(self, obj, value):
+        if value is None:
+            return None
+
+        if isinstance(value, ViewState):
+            return value
+
+        if isinstance(value, dict):
+            value = {**DEFAULT_INITIAL_VIEW_STATE, **value}
+            return ViewState(**value)
+
+        self.error(obj, value)
+        assert False

Diff for: src/actions/fly-to.ts

@@ -1,11 +1,10 @@
-import { SetStateAction } from "react";
 import { FlyToMessage } from "../types";
 import { FlyToInterpolator, MapViewState } from "@deck.gl/core/typed";
 import { isDefined } from "../util";
 
 export function flyTo(
   msg: FlyToMessage,
-  setInitialViewState: (value: SetStateAction<MapViewState>) => void,
+  setInitialViewState: (value: MapViewState) => void,
 ) {
   const {
     longitude,

Diff for: src/index.tsx

@@ -13,6 +13,7 @@ import { isDefined, loadChildModels } from "./util.js";
 import { v4 as uuidv4 } from "uuid";
 import { Message } from "./types.js";
 import { flyTo } from "./actions/fly-to.js";
+import { useViewStateDebounced } from "./state";
 
 await initParquetWasm();
 
@@ -64,25 +65,30 @@ async function getChildModelState(
 function App() {
   let model = useModel();
 
-  let [pythonInitialViewState] = useModelState<MapViewState>(
-    "_initial_view_state",
-  );
   let [mapStyle] = useModelState<string>("basemap_style");
   let [mapHeight] = useModelState<number>("_height");
   let [showTooltip] = useModelState<boolean>("show_tooltip");
   let [pickingRadius] = useModelState<number>("picking_radius");
   let [useDevicePixels] = useModelState<number | boolean>("use_device_pixels");
   let [parameters] = useModelState<object>("parameters");
 
-  let [initialViewState, setInitialViewState] = useState(
-    pythonInitialViewState,
-  );
+  // initialViewState is the value of view_state on the Python side. This is
+  // called `initial` here because it gets passed in to deck's
+  // `initialViewState` param, as deck manages its own view state. Further
+  // updates to `view_state` from Python are set on the deck `initialViewState`
+  // property, which can set new camera state, as described here:
+  // https://deck.gl/docs/developer-guide/interactivity
+  //
+  // `setViewState` is a debounced way to update the model and send view
+  // state information back to Python.
+  const [initialViewState, setViewState] =
+    useViewStateDebounced<MapViewState>("view_state");
 
   // Handle custom messages
   model.on("msg:custom", (msg: Message, buffers) => {
     switch (msg.type) {
       case "fly-to":
-        flyTo(msg, setInitialViewState);
+        flyTo(msg, setViewState);
         break;
 
       default:
@@ -165,6 +171,17 @@ function App() {
           overAlloc: 1,
           poolSize: 0,
         }}
+        onViewStateChange={(event) => {
+          const { viewState } = event;
+          const { longitude, latitude, zoom, pitch, bearing } = viewState;
+          setViewState({
+            longitude,
+            latitude,
+            zoom,
+            pitch,
+            bearing,
+          });
+        }}
         parameters={parameters || {}}
       >
         <Map mapStyle={mapStyle || DEFAULT_MAP_STYLE} />

Diff for: src/model/layer.ts

@@ -584,9 +584,6 @@ export class PolygonModel extends BaseArrowLayerModel {
   }
 
   layerProps(): Omit<GeoArrowPolygonLayerProps, "id"> {
-    console.log("table", this.table);
-    console.log("filled", this.filled);
-
     return {
       data: this.table,
       ...(isDefined(this.stroked) && { stroked: this.stroked }),

Diff for: src/state.ts

@@ -0,0 +1,47 @@
+import * as React from "react";
+import { useModel } from "@anywidget/react";
+import type { AnyModel } from "@anywidget/types";
+import { debounce } from "./util";
+
+const debouncedModelSaveViewState = debounce((model: AnyModel) => {
+  // TODO: this and below is hard-coded to the view_state model property!
+  const viewState = model.get("view_state");
+
+  // transitionInterpolator is sometimes a key in the view state while panning
+  // This is a function object and so can't be serialized via JSON.
+  //
+  // In the future anywidget may support custom serializers for sending data
+  // back from JS to Python. Until then, we need to clean the object ourselves.
+  // Because this is in a debounce it shouldn't often mess with deck's internal
+  // transition state it expects, because hopefully the transition will have
+  // finished in the 300ms that the user has stopped panning.
+  if ("transitionInterpolator" in viewState) {
+    console.debug("Deleting transitionInterpolator!");
+    delete viewState.transitionInterpolator;
+    model.set("view_state", viewState);
+  }
+
+  model.save_changes();
+}, 300);
+
+// TODO: add a `wait` parameter here, instead of having it hard-coded?
+export function useViewStateDebounced<T>(key: string): [T, (value: T) => void] {
+  let model = useModel();
+  let [value, setValue] = React.useState(model.get(key));
+  React.useEffect(() => {
+    let callback = () => {
+      setValue(model.get(key));
+    };
+    model.on(`change:${key}`, callback);
+    return () => model.off(`change:${key}`, callback);
+  }, [model, key]);
+  return [
+    value,
+    (value) => {
+      model.set(key, value);
+      // Note: I think this has to be defined outside of this function so that
+      // you're calling debounce on the same function object?
+      debouncedModelSaveViewState(model);
+    },
+  ];
+}

Diff for: src/util.ts

@@ -21,3 +21,13 @@ export async function loadChildModels(
 export function isDefined<T>(value: T | undefined | null): value is T {
   return value !== undefined && value !== null;
 }
+
+// From https://gist.github.com/ca0v/73a31f57b397606c9813472f7493a940
+export function debounce<T extends Function>(cb: T, wait = 20) {
+  let h: ReturnType<typeof setTimeout> | undefined;
+  let callable = (...args: any) => {
+    clearTimeout(h);
+    h = setTimeout(() => cb(...args), wait);
+  };
+  return <T>(<any>callable);
+}