_hexbin_mapbox.py

from plotly.express._core import build_dataframe
from plotly.express._doc import make_docstring
from plotly.express._chart_types import choropleth_mapbox, scatter_mapbox
import narwhals.stable.v1 as nw
import numpy as np


def _project_latlon_to_wgs84(lat, lon):
    """
    Projects lat and lon to WGS84, used to get regular hexagons on a mapbox map
    """
    x = lon * np.pi / 180
    y = np.arctanh(np.sin(lat * np.pi / 180))
    return x, y


def _project_wgs84_to_latlon(x, y):
    """
    Projects WGS84 to lat and lon, used to get regular hexagons on a mapbox map
    """
    lon = x * 180 / np.pi
    lat = (2 * np.arctan(np.exp(y)) - np.pi / 2) * 180 / np.pi
    return lat, lon


def _getBoundsZoomLevel(lon_min, lon_max, lat_min, lat_max, mapDim):
    """
    Get the mapbox zoom level given bounds and a figure dimension
    Source: https://stackoverflow.com/questions/6048975/google-maps-v3-how-to-calculate-the-zoom-level-for-a-given-bounds
    """

    scale = (
        2  # adjustment to reflect MapBox base tiles are 512x512 vs. Google's 256x256
    )
    WORLD_DIM = {"height": 256 * scale, "width": 256 * scale}
    ZOOM_MAX = 18

    def latRad(lat):
        sin = np.sin(lat * np.pi / 180)
        radX2 = np.log((1 + sin) / (1 - sin)) / 2
        return max(min(radX2, np.pi), -np.pi) / 2

    def zoom(mapPx, worldPx, fraction):
        return 0.95 * np.log(mapPx / worldPx / fraction) / np.log(2)

    latFraction = (latRad(lat_max) - latRad(lat_min)) / np.pi

    lngDiff = lon_max - lon_min
    lngFraction = ((lngDiff + 360) if lngDiff < 0 else lngDiff) / 360

    latZoom = zoom(mapDim["height"], WORLD_DIM["height"], latFraction)
    lngZoom = zoom(mapDim["width"], WORLD_DIM["width"], lngFraction)

    return min(latZoom, lngZoom, ZOOM_MAX)


def _compute_hexbin(x, y, x_range, y_range, color, nx, agg_func, min_count):
    """
    Computes the aggregation at hexagonal bin level.
    Also defines the coordinates of the hexagons for plotting.
    The binning is inspired by matplotlib's implementation.

    Parameters
    ----------
    x : np.ndarray
        Array of x values (shape N)
    y : np.ndarray
        Array of y values (shape N)
    x_range : np.ndarray
        Min and max x (shape 2)
    y_range : np.ndarray
        Min and max y (shape 2)
    color : np.ndarray
        Metric to aggregate at hexagon level (shape N)
    nx : int
        Number of hexagons horizontally
    agg_func : function
        Numpy compatible aggregator, this function must take a one-dimensional
        np.ndarray as input and output a scalar
    min_count : int
        Minimum number of points in the hexagon for the hexagon to be displayed

    Returns
    -------
    np.ndarray
        X coordinates of each hexagon (shape M x 6)
    np.ndarray
        Y coordinates of each hexagon (shape M x 6)
    np.ndarray
        Centers of the hexagons (shape M x 2)
    np.ndarray
        Aggregated value in each hexagon (shape M)

    """
    xmin = x_range.min()
    xmax = x_range.max()
    ymin = y_range.min()
    ymax = y_range.max()

    # In the x-direction, the hexagons exactly cover the region from
    # xmin to xmax. Need some padding to avoid roundoff errors.
    padding = 1.0e-9 * (xmax - xmin)
    xmin -= padding
    xmax += padding

    Dx = xmax - xmin
    Dy = ymax - ymin
    if Dx == 0 and Dy > 0:
        dx = Dy / nx
    elif Dx == 0 and Dy == 0:
        dx, _ = _project_latlon_to_wgs84(1, 1)
    else:
        dx = Dx / nx
    dy = dx * np.sqrt(3)
    ny = np.ceil(Dy / dy).astype(int)

    # Center the hexagons vertically since we only want regular hexagons
    ymin -= (ymin + dy * ny - ymax) / 2

    x = (x - xmin) / dx
    y = (y - ymin) / dy
    ix1 = np.round(x).astype(int)
    iy1 = np.round(y).astype(int)
    ix2 = np.floor(x).astype(int)
    iy2 = np.floor(y).astype(int)

    nx1 = nx + 1
    ny1 = ny + 1
    nx2 = nx
    ny2 = ny
    n = nx1 * ny1 + nx2 * ny2

    d1 = (x - ix1) ** 2 + 3.0 * (y - iy1) ** 2
    d2 = (x - ix2 - 0.5) ** 2 + 3.0 * (y - iy2 - 0.5) ** 2
    bdist = d1 < d2

    if color is None:
        lattice1 = np.zeros((nx1, ny1))
        lattice2 = np.zeros((nx2, ny2))
        c1 = (0 <= ix1) & (ix1 < nx1) & (0 <= iy1) & (iy1 < ny1) & bdist
        c2 = (0 <= ix2) & (ix2 < nx2) & (0 <= iy2) & (iy2 < ny2) & ~bdist
        np.add.at(lattice1, (ix1[c1], iy1[c1]), 1)
        np.add.at(lattice2, (ix2[c2], iy2[c2]), 1)
        if min_count is not None:
            lattice1[lattice1 < min_count] = np.nan
            lattice2[lattice2 < min_count] = np.nan
        accum = np.concatenate([lattice1.ravel(), lattice2.ravel()])
        good_idxs = ~np.isnan(accum)
    else:
        if min_count is None:
            min_count = 1

        # create accumulation arrays
        lattice1 = np.empty((nx1, ny1), dtype=object)
        for i in range(nx1):
            for j in range(ny1):
                lattice1[i, j] = []
        lattice2 = np.empty((nx2, ny2), dtype=object)
        for i in range(nx2):
            for j in range(ny2):
                lattice2[i, j] = []

        for i in range(len(x)):
            if bdist[i]:
                if 0 <= ix1[i] < nx1 and 0 <= iy1[i] < ny1:
                    lattice1[ix1[i], iy1[i]].append(color[i])
            else:
                if 0 <= ix2[i] < nx2 and 0 <= iy2[i] < ny2:
                    lattice2[ix2[i], iy2[i]].append(color[i])

        for i in range(nx1):
            for j in range(ny1):
                vals = lattice1[i, j]
                if len(vals) >= min_count:
                    lattice1[i, j] = agg_func(vals)
                else:
                    lattice1[i, j] = np.nan
        for i in range(nx2):
            for j in range(ny2):
                vals = lattice2[i, j]
                if len(vals) >= min_count:
                    lattice2[i, j] = agg_func(vals)
                else:
                    lattice2[i, j] = np.nan

        accum = np.hstack(
            (lattice1.astype(float).ravel(), lattice2.astype(float).ravel())
        )
        good_idxs = ~np.isnan(accum)

    agreggated_value = accum[good_idxs]

    centers = np.zeros((n, 2), float)
    centers[: nx1 * ny1, 0] = np.repeat(np.arange(nx1), ny1)
    centers[: nx1 * ny1, 1] = np.tile(np.arange(ny1), nx1)
    centers[nx1 * ny1 :, 0] = np.repeat(np.arange(nx2) + 0.5, ny2)
    centers[nx1 * ny1 :, 1] = np.tile(np.arange(ny2), nx2) + 0.5
    centers[:, 0] *= dx
    centers[:, 1] *= dy
    centers[:, 0] += xmin
    centers[:, 1] += ymin
    centers = centers[good_idxs]

    # Define normalised regular hexagon coordinates
    hx = [0, 0.5, 0.5, 0, -0.5, -0.5]
    hy = [
        -0.5 / np.cos(np.pi / 6),
        -0.5 * np.tan(np.pi / 6),
        0.5 * np.tan(np.pi / 6),
        0.5 / np.cos(np.pi / 6),
        0.5 * np.tan(np.pi / 6),
        -0.5 * np.tan(np.pi / 6),
    ]

    # Number of hexagons needed
    m = len(centers)

    # Coordinates for all hexagonal patches
    hxs = np.array([hx] * m) * dx + np.vstack(centers[:, 0])
    hys = np.array([hy] * m) * dy / np.sqrt(3) + np.vstack(centers[:, 1])

    return hxs, hys, centers, agreggated_value


def _compute_wgs84_hexbin(
    lat=None,
    lon=None,
    lat_range=None,
    lon_range=None,
    color=None,
    nx=None,
    agg_func=None,
    min_count=None,
    native_namespace=None,
):
    """
    Computes the lat-lon aggregation at hexagonal bin level.
    Latitude and longitude need to be projected to WGS84 before aggregating
    in order to display regular hexagons on the map.

    Parameters
    ----------
    lat : np.ndarray
        Array of latitudes (shape N)
    lon : np.ndarray
        Array of longitudes (shape N)
    lat_range : np.ndarray
        Min and max latitudes (shape 2)
    lon_range : np.ndarray
        Min and max longitudes (shape 2)
    color : np.ndarray
        Metric to aggregate at hexagon level (shape N)
    nx : int
        Number of hexagons horizontally
    agg_func : function
        Numpy compatible aggregator, this function must take a one-dimensional
        np.ndarray as input and output a scalar
    min_count : int
        Minimum number of points in the hexagon for the hexagon to be displayed

    Returns
    -------
    np.ndarray
        Lat coordinates of each hexagon (shape M x 6)
    np.ndarray
        Lon coordinates of each hexagon (shape M x 6)
    nw.Series
        Unique id for each hexagon, to be used in the geojson data (shape M)
    np.ndarray
        Aggregated value in each hexagon (shape M)

    """
    # Project to WGS 84
    x, y = _project_latlon_to_wgs84(lat, lon)

    if lat_range is None:
        lat_range = np.array([lat.min(), lat.max()])
    if lon_range is None:
        lon_range = np.array([lon.min(), lon.max()])

    x_range, y_range = _project_latlon_to_wgs84(lat_range, lon_range)

    hxs, hys, centers, agreggated_value = _compute_hexbin(
        x, y, x_range, y_range, color, nx, agg_func, min_count
    )

    # Convert back to lat-lon
    hexagons_lats, hexagons_lons = _project_wgs84_to_latlon(hxs, hys)

    # Create unique feature id based on hexagon center
    centers = centers.astype(str)
    hexagons_ids = (
        nw.from_dict(
            {"x1": centers[:, 0], "x2": centers[:, 1]},
            native_namespace=native_namespace,
        )
        .select(hexagons_ids=nw.concat_str([nw.col("x1"), nw.col("x2")], separator=","))
        .get_column("hexagons_ids")
    )

    return hexagons_lats, hexagons_lons, hexagons_ids, agreggated_value


def _hexagons_to_geojson(hexagons_lats, hexagons_lons, ids=None):
    """
    Creates a geojson of hexagonal features based on the outputs of
    _compute_wgs84_hexbin
    """
    features = []
    if ids is None:
        ids = np.arange(len(hexagons_lats))
    for lat, lon, idx in zip(hexagons_lats, hexagons_lons, ids):
        points = np.array([lon, lat]).T.tolist()
        points.append(points[0])
        features.append(
            dict(
                type="Feature",
                id=idx,
                geometry=dict(type="Polygon", coordinates=[points]),
            )
        )
    return dict(type="FeatureCollection", features=features)


def create_hexbin_mapbox(
    data_frame=None,
    lat=None,
    lon=None,
    color=None,
    nx_hexagon=5,
    agg_func=None,
    animation_frame=None,
    color_discrete_sequence=None,
    color_discrete_map={},
    labels={},
    color_continuous_scale=None,
    range_color=None,
    color_continuous_midpoint=None,
    opacity=None,
    zoom=None,
    center=None,
    mapbox_style=None,
    title=None,
    template=None,
    width=None,
    height=None,
    min_count=None,
    show_original_data=False,
    original_data_marker=None,
):
    """
    Returns a figure aggregating scattered points into connected hexagons
    """
    args = build_dataframe(args=locals(), constructor=None)
    native_namespace = nw.get_native_namespace(args["data_frame"])
    if agg_func is None:
        agg_func = np.mean

    lat_range = (
        args["data_frame"]
        .select(
            nw.min(args["lat"]).name.suffix("_min"),
            nw.max(args["lat"]).name.suffix("_max"),
        )
        .to_numpy()
        .squeeze()
    )

    lon_range = (
        args["data_frame"]
        .select(
            nw.min(args["lon"]).name.suffix("_min"),
            nw.max(args["lon"]).name.suffix("_max"),
        )
        .to_numpy()
        .squeeze()
    )

    hexagons_lats, hexagons_lons, hexagons_ids, count = _compute_wgs84_hexbin(
        lat=args["data_frame"].get_column(args["lat"]).to_numpy(),
        lon=args["data_frame"].get_column(args["lon"]).to_numpy(),
        lat_range=lat_range,
        lon_range=lon_range,
        color=None,
        nx=nx_hexagon,
        agg_func=agg_func,
        min_count=min_count,
        native_namespace=native_namespace,
    )

    geojson = _hexagons_to_geojson(hexagons_lats, hexagons_lons, hexagons_ids)

    if zoom is None:
        if height is None and width is None:
            mapDim = dict(height=450, width=450)
        elif height is None and width is not None:
            mapDim = dict(height=450, width=width)
        elif height is not None and width is None:
            mapDim = dict(height=height, width=height)
        else:
            mapDim = dict(height=height, width=width)
        zoom = _getBoundsZoomLevel(
            lon_range[0], lon_range[1], lat_range[0], lat_range[1], mapDim
        )

    if center is None:
        center = dict(lat=lat_range.mean(), lon=lon_range.mean())

    if args["animation_frame"] is not None:
        groups = dict(args["data_frame"].group_by(args["animation_frame"]).__iter__())
    else:
        groups = {(0,): args["data_frame"]}

    agg_data_frame_list = []
    for key, df in groups.items():
        _, _, hexagons_ids, aggregated_value = _compute_wgs84_hexbin(
            lat=df.get_column(args["lat"]).to_numpy(),
            lon=df.get_column(args["lon"]).to_numpy(),
            lat_range=lat_range,
            lon_range=lon_range,
            color=df.get_column(args["color"]).to_numpy() if args["color"] else None,
            nx=nx_hexagon,
            agg_func=agg_func,
            min_count=min_count,
            native_namespace=native_namespace,
        )
        agg_data_frame_list.append(
            nw.from_dict(
                {
                    "frame": [key[0]] * len(hexagons_ids),
                    "locations": hexagons_ids,
                    "color": aggregated_value,
                },
                native_namespace=native_namespace,
            )
        )

    agg_data_frame = nw.concat(agg_data_frame_list, how="vertical").with_columns(
        color=nw.col("color").cast(nw.Int64)
    )

    if range_color is None:
        range_color = [agg_data_frame["color"].min(), agg_data_frame["color"].max()]

    fig = choropleth_mapbox(
        data_frame=agg_data_frame.to_native(),
        geojson=geojson,
        locations="locations",
        color="color",
        hover_data={"color": True, "locations": False, "frame": False},
        animation_frame=("frame" if args["animation_frame"] is not None else None),
        color_discrete_sequence=color_discrete_sequence,
        color_discrete_map=color_discrete_map,
        labels=labels,
        color_continuous_scale=color_continuous_scale,
        range_color=range_color,
        color_continuous_midpoint=color_continuous_midpoint,
        opacity=opacity,
        zoom=zoom,
        center=center,
        mapbox_style=mapbox_style,
        title=title,
        template=template,
        width=width,
        height=height,
    )

    if show_original_data:
        original_fig = scatter_mapbox(
            data_frame=(
                args["data_frame"].sort(
                    by=args["animation_frame"], descending=False, nulls_last=True
                )
                if args["animation_frame"] is not None
                else args["data_frame"]
            ),
            lat=args["lat"],
            lon=args["lon"],
            animation_frame=args["animation_frame"],
        )
        original_fig.data[0].hoverinfo = "skip"
        original_fig.data[0].hovertemplate = None
        original_fig.data[0].marker = original_data_marker

        fig.add_trace(original_fig.data[0])

        if args["animation_frame"] is not None:
            for i in range(len(original_fig.frames)):
                original_fig.frames[i].data[0].hoverinfo = "skip"
                original_fig.frames[i].data[0].hovertemplate = None
                original_fig.frames[i].data[0].marker = original_data_marker

                fig.frames[i].data = [
                    fig.frames[i].data[0],
                    original_fig.frames[i].data[0],
                ]

    return fig


create_hexbin_mapbox.__doc__ = make_docstring(
    create_hexbin_mapbox,
    override_dict=dict(
        nx_hexagon=["int", "Number of hexagons (horizontally) to be created"],
        agg_func=[
            "function",
            "Numpy array aggregator, it must take as input a 1D array",
            "and output a scalar value.",
        ],
        min_count=[
            "int",
            "Minimum number of points in a hexagon for it to be displayed.",
            "If None and color is not set, display all hexagons.",
            "If None and color is set, only display hexagons that contain points.",
        ],
        show_original_data=[
            "bool",
            "Whether to show the original data on top of the hexbin aggregation.",
        ],
        original_data_marker=["dict", "Scattermapbox marker options."],
    ),
)