start geopandas extension

lonboard/geopandas.py

@@ -0,0 +1,323 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import geopandas as gpd
+import numpy as np
+import pandas as pd
+
+from lonboard import Map, viz
+from lonboard.basemap import CartoBasemap
+from lonboard.colormap import apply_categorical_cmap, apply_continuous_cmap
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+__all__ = ["LonboardAccessor"]
+
+QUERY_NAME_TRANSLATION = str.maketrans(dict.fromkeys("., -_/", ""))
+BASEMAP_PROVIDERS = {
+    "CartoDB Positron": CartoBasemap.Positron,
+    "CartoDB Positron No Label": CartoBasemap.PositronNoLabels,
+    "CartoDB Darkmatter": CartoBasemap.DarkMatter,
+    "CartoDB Darkmatter No Label": CartoBasemap.DarkMatterNoLabels,
+    "CartoDB Voyager": CartoBasemap.Voyager,
+    "CartoDB Voyager No Label": CartoBasemap.VoyagerNoLabels,
+}
+# Convert keys to lower case without spaces
+BASEMAP_PROVIDERS = {
+    k.translate(QUERY_NAME_TRANSLATION).lower(): v for k, v in BASEMAP_PROVIDERS.items()
+}
+
+
+@pd.api.extensions.register_dataframe_accessor("lb")
+class LonboardAccessor:
+    """Geopandas Extension class to provide the `explore` method."""
+
+    def __init__(self, pandas_obj) -> None:  # noqa: ANN001, D107
+        self._validate(pandas_obj)
+        self._obj = pandas_obj
+
+    @staticmethod
+    def _validate(obj) -> None:  # noqa: ANN001
+        if not isinstance(obj, gpd.GeoDataFrame):
+            raise TypeError("must be a geodataframe")
+
+    def explore(  # noqa: PLR0913
+        self,
+        column: str | None = None,
+        cmap: str | None = None,
+        scheme: str | None = None,
+        k: int | None = 6,
+        categorical: bool = False,  # noqa: FBT001, FBT002
+        elevation: str | np.ndarray | None = None,
+        elevation_scale: float | None = 1,
+        alpha: float | None = 1,
+        layer_kwargs: dict[str, Any] | None = None,
+        map_kwargs: dict[str, Any] | None = None,
+        classification_kwds: dict[str, Any] | None = None,
+        nan_color: list[int] | NDArray[np.number] | None = None,
+        color: str | None = None,
+        vmin: float | None = None,
+        vmax: float | None = None,
+        wireframe: bool = False,  # noqa: FBT001, FBT002
+        tiles: str | None = None,
+        highlight: bool = False,  # noqa: FBT001, FBT002
+        m: Map | None = None,
+    ) -> Map:
+        """Explore a dataframe using lonboard and deckgl.
+
+        Keyword Args:
+            column : Name of column on dataframe to visualize on map.
+            cmap : Name of matplotlib colormap to use.
+            scheme : Name of a classification scheme defined by mapclassify.Classifier.
+            k : Number of classes to generate. Defaults to 6.
+            categorical : Whether the data should be treated as categorical or
+                continuous.
+            elevation : Name of column on the dataframe used to extrude each geometry or
+                an array-like in the same order as observations. Defaults to None.
+            elevation_scale : Constant scaler multiplied by elevation value.
+            alpha : Alpha (opacity) parameter in the range (0,1) passed to
+                mapclassify.util.get_color_array.
+            layer_kwargs : Additional keyword arguments passed to lonboard.viz layer
+                arguments (either polygon_kwargs, scatterplot_kwargs, or path_kwargs,
+                depending on input geometry type).
+            map_kwargs : Additional keyword arguments passed to lonboard.viz map_kwargs.
+            classification_kwds : Additional keyword arguments passed to
+                `mapclassify.classify`.
+            nan_color : Color used to shade NaN observations formatted as an RGBA list.
+                Defaults to [255, 255, 255, 255]. If no alpha channel is passed it is
+                assumed to be 255.
+            color : single or array of colors passed to Layer.get_fill_color
+                or a lonboard.basemap object, or a string to a maplibre style basemap.
+            vmin : Minimum value for color mapping.
+            vmax : Maximum value for color mapping.
+            wireframe : Whether to use wireframe styling in deckgl.
+            tiles : Either a known string {"CartoDB Positron",
+                "CartoDB Positron No Label", "CartoDB Darkmatter",
+                "CartoDB Darkmatter No Label", "CartoDB Voyager",
+                "CartoDB Voyager No Label"}
+            highlight : Whether to highlight each feature on mouseover (passed to
+                lonboard.Layer's auto_highlight). Defaults to False.
+            m: An existing Map object to plot onto.
+
+        Returns:
+        lonboard.Map
+            a lonboard map with geodataframe included as a Layer object.
+
+        """
+        return _dexplore(
+            self._obj,
+            column=column,
+            cmap=cmap,
+            scheme=scheme,
+            k=k,
+            categorical=categorical,
+            elevation=elevation,
+            elevation_scale=elevation_scale,
+            alpha=alpha,
+            layer_kwargs=layer_kwargs,
+            map_kwargs=map_kwargs,
+            classification_kwds=classification_kwds,
+            nan_color=nan_color,
+            color=color,
+            vmin=vmin,
+            vmax=vmax,
+            wireframe=wireframe,
+            tiles=tiles,
+            highlight=highlight,
+            m=m,
+        )
+
+
+def _dexplore(  # noqa: C901, PLR0912, PLR0913, PLR0915
+    gdf,  # noqa: ANN001
+    *,
+    column,  # noqa: ANN001
+    cmap,  # noqa: ANN001
+    scheme,  # noqa: ANN001
+    k,  # noqa: ANN001
+    categorical,  # noqa: ANN001
+    elevation,  # noqa: ANN001
+    elevation_scale,  # noqa: ANN001
+    alpha,  # noqa: ANN001
+    layer_kwargs,  # noqa: ANN001
+    map_kwargs,  # noqa: ANN001
+    classification_kwds,  # noqa: ANN001
+    nan_color,  # noqa: ANN001
+    color,  # noqa: ANN001
+    vmin,  # noqa: ANN001
+    vmax,  # noqa: ANN001
+    wireframe,  # noqa: ANN001
+    tiles,  # noqa: ANN001
+    highlight,  # noqa: ANN001
+    m,  # noqa: ANN001
+) -> Map:
+    """Explore a dataframe using lonboard and deckgl.
+
+    See the public docstring for detailed parameter information
+
+    Returns
+    -------
+    lonboard.Map
+        a lonboard map with geodataframe included as a Layer object.
+
+    """
+    if map_kwargs is None:
+        map_kwargs = {}
+    if classification_kwds is None:
+        classification_kwds = {}
+    if layer_kwargs is None:
+        layer_kwargs = {}
+    if isinstance(elevation, str):
+        if elevation in gdf.columns:
+            elevation = gdf[elevation]
+        else:
+            raise ValueError(
+                f"the designated height column {elevation} is not in the dataframe",
+            )
+        if not pd.api.types.is_numeric_dtype(elevation):
+            raise ValueError("elevation must be a numeric data type")
+    if elevation is not None:
+        layer_kwargs["extruded"] = True
+    if nan_color is None:
+        nan_color = [255, 255, 255, 255]
+    if not pd.api.types.is_list_like(nan_color):
+        raise ValueError("nan_color must be an iterable of 3 or 4 values")
+    if len(nan_color) != 4:
+        if len(nan_color) == 3:
+            nan_color = np.append(nan_color, [255])
+        else:
+            raise ValueError("nan_color must be an iterable of 3 or 4 values")
+
+    # only polygons have z
+    if ["Polygon", "MultiPolygon"] in gdf.geometry.geom_type.unique():
+        layer_kwargs["get_elevation"] = elevation
+        layer_kwargs["elevation_scale"] = elevation_scale
+        layer_kwargs["wireframe"] = wireframe
+        layer_kwargs["auto_highlight"] = highlight
+
+    line = False  # set color of lines, not fill_color
+    if ["LineString", "MultiLineString"] in gdf.geometry.geom_type.unique():
+        line = True
+    if color:
+        if line:
+            layer_kwargs["get_color"] = color
+        else:
+            layer_kwargs["get_fill_color"] = color
+    if column is not None:
+        try:
+            from matplotlib import colormaps
+        except ImportError as e:
+            raise ImportError(
+                "you must have matplotlib installed to style by a column",
+            ) from e
+
+        if column not in gdf.columns:
+            raise ValueError(f"the designated column {column} is not in the dataframe")
+        if gdf[column].dtype in ["O", "category"]:
+            categorical = True
+        if cmap is not None and cmap not in colormaps:
+            raise ValueError(
+                f"`cmap` must be one of {list(colormaps.keys())} but {cmap} was passed",
+            )
+        if cmap is None:
+            cmap = "tab20" if categorical else "viridis"
+        if categorical:
+            color_array = _get_categorical_cmap(gdf[column], cmap, nan_color, alpha)
+        elif scheme is None:
+            if vmin is None:
+                vmin = np.nanmin(gdf[column])
+            if vmax is None:
+                vmax = np.nanmax(gdf[column])
+            # minmax scale the column first, matplotlib needs 0-1
+            transformed = (gdf[column] - vmin) / (vmax - vmin)
+            color_array = apply_continuous_cmap(
+                values=transformed,
+                cmap=colormaps[cmap],
+                alpha=alpha,
+            )
+        else:
+            try:
+                from mapclassify._classify_API import _classifiers
+                from mapclassify.util import get_color_array
+
+                _klasses = list(_classifiers.keys())
+                _klasses.append("userdefined")
+            except ImportError as e:
+                raise ImportError(
+                    "you must have the `mapclassify` package installed to use the "
+                    "`scheme` keyword",
+                ) from e
+            if scheme.replace("_", "") not in _klasses:
+                raise ValueError(
+                    "the classification scheme must be a valid mapclassify"
+                    f"classifier in {_klasses},"
+                    f"but {scheme} was passed instead",
+                )
+            if k is not None and "k" in classification_kwds:
+                # k passed directly takes precedence
+                classification_kwds.pop("k")
+            color_array = get_color_array(
+                gdf[column],
+                scheme=scheme,
+                k=k,
+                cmap=cmap,
+                alpha=alpha,
+                nan_color=nan_color,
+                **classification_kwds,
+            )
+
+        if line:
+            layer_kwargs["get_color"] = color_array
+
+        else:
+            layer_kwargs["get_fill_color"] = color_array
+    if tiles:
+        map_kwargs["basemap_style"] = _query_name(tiles)
+    new_m = viz(
+        gdf,
+        polygon_kwargs=layer_kwargs,
+        scatterplot_kwargs=layer_kwargs,
+        path_kwargs=layer_kwargs,
+        map_kwargs=map_kwargs,
+    )
+    if m is not None:
+        new_m = m.add_layer(new_m)
+
+    return new_m
+
+
+def _get_categorical_cmap(categories, cmap, nan_color, alpha):  # noqa: ANN001, ANN202
+    try:
+        from matplotlib import colormaps
+    except ImportError as e:
+        raise ImportError(
+            "this function requires the `lonboard` package to be installed",
+        ) from e
+
+    cat_codes = pd.Series(pd.Categorical(categories).codes, dtype="category")
+    # nans are encoded as -1 OR largest category depending on input type
+    # re-encode to always be last category
+    cat_codes = cat_codes.cat.rename_categories({-1: len(cat_codes.unique()) - 1})
+    unique_cats = categories.dropna().unique()
+    n_cats = len(unique_cats)
+    colors = colormaps[cmap].resampled(n_cats)(list(range(n_cats)), alpha, bytes=True)
+    colors = np.vstack([colors, nan_color])
+    temp_cmap = dict(zip(range(n_cats + 1), colors))
+    return apply_categorical_cmap(cat_codes, temp_cmap)
+
+
+def _query_name(name: str) -> CartoBasemap:
+    """Return basemap URL based on the name query (mimicking behavior from xyzservices).
+
+    Returns a matching basemap from name contains the same letters in the same
+    order as the provider's name irrespective of the letter case, spaces, dashes
+    and other characters. See examples for details.
+    """
+    name_clean = name.translate(QUERY_NAME_TRANSLATION).lower()
+    if name_clean in BASEMAP_PROVIDERS:
+        return BASEMAP_PROVIDERS[name_clean]
+
+    raise ValueError(f"No matching provider found for the query '{name}'.")

pyproject.toml

@@ -78,9 +78,11 @@ dev = [
     "geoarrow-rust-core>=0.3.0",
     "geodatasets>=2024.8.0",
     "jupyterlab>=4.3.3",
+    "mapclassify>=2.8.1",
     "matplotlib>=3.7.5",
     "movingpandas>=0.20.0",
     "palettable>=3.3.3",
+    "pandas-stubs>=2.2.2.240807",
     "pre-commit>=3.5.0",
     "pyarrow>=17.0.0",
     "pyogrio>=0.9.0",
@@ -115,7 +117,13 @@ artifacts = ["lonboard/static/*.js", "lonboard/static/*.css"]
 "MANIFEST.in" = "MANIFEST.in"
 
 [[tool.mypy.overrides]]
-module = ["pyogrio.*", "pyarrow.*", "geodatasets.*", "ipywidgets.*"]
+module = [
+    "pyogrio.*",
+    "pyarrow.*",
+    "geodatasets.*",
+    "ipywidgets.*",
+    "geopandas.*",
+]
 ignore_missing_imports = true
 
 [tool.ruff.lint]