Merge pull request #1 from Filimoa/graphing-utils

Adding Graphing Helpers

Author: Filimoa

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "polars-h3"
-version = "0.3.0"
+version = "0.4.0"  
 edition = "2021"
 
 [lib]

Cargo.toml

@@ -117,3 +117,21 @@ We are unable to support the functions that work with geometries.
 | `polygon_wkt_to_cells` | Convert polygon WKT to a set of cells | 🛑 |
 | `directed_edge_to_boundary_wkt` | Convert directed edge ID to linestring WKT | 🛑 |
 
+
+### Plotting
+
+The library also comes with helper functions to plot hexes on a Folium map.
+
+
+```python
+import polars_h3 as pl_h3
+import polars as pl
+
+hex_map = pl_h3.graphing.plot_hex_outlines(df, "h3_cell")
+display(hex_map)
+
+# or if you have a metric to plot
+
+hex_map = pl_h3.graphing.plot_hex_fills(df, "h3_cell", "metric_col")
+display(hex_map)
+```

README.md

@@ -1,3 +1,4 @@
+from . import graphing
 from .core.edge import (
     are_neighbor_cells,
     cells_to_directed_edge,
@@ -100,4 +101,5 @@
     "edge_length",
     "average_hexagon_edge_length",
     "get_num_cells",
+    "graphing",
 ]

polars_h3/__init__.py

@@ -0,0 +1,189 @@
+from typing import Any, Literal, Union
+
+import polars as pl
+
+from .core.indexing import cell_to_boundary
+
+
+def _hex_bounds(
+    df: pl.DataFrame, boundary_col: str = "boundary"
+) -> tuple[tuple[float, float], tuple[float, float]]:
+    df_flat = (
+        df.explode(boundary_col)
+        .with_columns(
+            [
+                pl.col(boundary_col).list.get(0).alias("lat"),
+                pl.col(boundary_col).list.get(1).alias("lng"),
+            ]
+        )
+        .drop(boundary_col)
+    )
+
+    min_lat = float(df_flat["lat"].min())  # type: ignore
+    max_lat = float(df_flat["lat"].max())  # type: ignore
+    min_lng = float(df_flat["lng"].min())  # type: ignore
+    max_lng = float(df_flat["lng"].max())  # type: ignore
+
+    return ((min_lat, min_lng), (max_lat, max_lng))
+
+
+def plot_hex_outlines(
+    df: pl.DataFrame,
+    hex_id_col: str,
+    map: Union[Any, None] = None,
+    outline_color: str = "red",
+    map_size: Literal["medium", "large"] = "medium",
+) -> Any:
+    """
+    Plot hexagon outlines on a Folium map.
+
+    Parameters
+    ----------
+    df : pl.DataFrame
+        A DataFrame that must contain a column of hex IDs.
+    hex_id_col : str
+        The name of the column in `df` that contains hexagon identifiers (H3 cell IDs).
+    map : folium.Map or None, optional
+        An existing Folium map object on which to plot. If None, a new map is created.
+    outline_color : str, optional
+        The color used to outline the hexagons. Defaults to "red".
+    map_size : {"medium", "large"}, optional
+        The size of the displayed map. "medium" fits a 50% view, "large" takes 100%. Defaults to "medium".
+
+    Returns
+    -------
+    folium.Map
+        A Folium map object with hexagon outlines added.
+
+    Raises
+    ------
+    ValueError
+        If the input DataFrame is empty.
+    ImportError
+        If Folium is not installed.
+    """
+    if df.height == 0:
+        raise ValueError("DataFrame is empty")
+
+    try:
+        import folium
+    except ImportError as e:
+        raise ImportError(
+            "folium is required to plot hex outlines. Install with `pip install folium`"
+        ) from e
+
+    if not map:
+        map = folium.Map(
+            zoom_start=13,
+            tiles="cartodbpositron",
+            width="50%" if map_size == "medium" else "100%",
+            height="50%" if map_size == "medium" else "100%",
+        )
+
+    df = df.with_columns(
+        [
+            cell_to_boundary(pl.col(hex_id_col)).alias("boundary"),
+        ]
+    ).filter(pl.col("boundary").is_not_null())
+
+    for hex_cord in df["boundary"].to_list():
+        folium.Polygon(locations=hex_cord, weight=5, color=outline_color).add_to(map)
+
+    map_bounds = _hex_bounds(df, "boundary")
+    map.fit_bounds(map_bounds)
+    return map
+
+
+def plot_hex_fills(
+    df: pl.DataFrame,
+    hex_id_col: str,
+    metric_col: str,
+    map: Union[Any, None] = None,
+    map_size: Literal["medium", "large"] = "medium",
+) -> Any:
+    """
+    Render filled hexagonal cells on a Folium map, colorized by a specified metric.
+
+    If no map is provided, a new Folium map is created. The map is automatically
+    fit to the bounds of the plotted polygons.
+
+    #### Parameters
+    - `df`: pl.DataFrame
+    - `hex_id_col`: str
+      Column name in `df` holding H3 cell indices.
+    - `metric_col`: str
+      Column name in `df` containing the metric values for colorization.
+    - `map`: folium.Map | None, default None
+      An existing Folium Map object. If None, a new map is created.
+    - `map_size`: Literal["medium", "large"], default "medium"
+      Controls the size of the Folium map. `"medium"` sets width/height to 50% while `"large"` sets it to 100%.
+
+    #### Returns
+    folium.Map
+      The Folium Map object with the rendered hexagon polygons.
+    """
+    if df.height == 0:
+        raise ValueError("DataFrame is empty")
+
+    try:
+        import folium
+        import matplotlib
+    except ImportError as e:
+        raise ImportError(
+            "folium and matplotlib are required to plot hex fills. Install with `pip install folium matplotlib`"
+        ) from e
+
+    if not map:
+        map = folium.Map(
+            zoom_start=13,
+            tiles="cartodbpositron",
+            width="50%" if map_size == "medium" else "100%",
+            height="50%" if map_size == "medium" else "100%",
+        )
+
+    df = df.with_columns(
+        [
+            cell_to_boundary(pl.col(hex_id_col)).alias("boundary"),
+            pl.col(metric_col).log1p().alias("normalized_metric"),
+        ]
+    ).filter(pl.col("boundary").is_not_null())
+
+    hexagons = df[hex_id_col].to_list()
+    metrics = df[metric_col].to_list()
+    compressed_metrics = df["normalized_metric"].to_list()
+    boundaries = df["boundary"].to_list()
+
+    min_val = min(compressed_metrics)
+    max_val = max(compressed_metrics)
+
+    if max_val == min_val:
+        normalized_metrics = [0.0] * len(compressed_metrics)
+    else:
+        normalized_metrics = [
+            (x - min_val) / (max_val - min_val) for x in compressed_metrics
+        ]
+
+    colormap = matplotlib.colormaps.get_cmap("plasma")
+
+    for (hexagon, metric, boundary), norm_metric in zip(
+        zip(hexagons, metrics, boundaries), normalized_metrics, strict=False
+    ):
+        rgba = colormap(norm_metric)
+        color = (
+            f"#{int(rgba[0] * 255):02x}{int(rgba[1] * 255):02x}{int(rgba[2] * 255):02x}"
+        )
+
+        folium.Polygon(
+            locations=boundary,
+            fill=True,
+            fill_opacity=0.6 + 0.4 * norm_metric,
+            fill_color=color,
+            color=color,
+            weight=1,
+            tooltip=f"{hexagon}<br>Value: {metric}",
+        ).add_to(map)
+
+    map_bounds = _hex_bounds(df, "boundary")
+    map.fit_bounds(map_bounds)
+
+    return map

polars_h3/graphing.py

@@ -7,7 +7,7 @@ homepage = "https://github.com/Filimonov/polars-h3"
 
 [project]
 name = "polars-h3"  
-version = "0.3.0"  
+version = "0.4.0"  
 description = "H3 bindings for Polars"
 readme = "README.md"
 requires-python = ">=3.9"

pyproject.toml

@@ -109,21 +109,27 @@ pub fn cell_to_latlng(cell_series: &Series) -> PolarsResult<Series> {
 }
 
 pub fn cell_to_boundary(cell_series: &Series) -> PolarsResult<Series> {
-    use crate::engine::utils::parse_cell_indices;
     let cells = parse_cell_indices(cell_series)?;
 
     let coords: ListChunked = cells
         .into_par_iter()
         .map(|cell| {
             cell.map(|idx| {
                 let boundary = idx.boundary();
-                // Convert boundary vertices into a flat Vec<f64> of [lat, lng, lat, lng, ...]
-                let mut latlngs = Vec::with_capacity(boundary.len() * 2);
-                for vertex in boundary.iter() {
-                    latlngs.push(vertex.lat());
-                    latlngs.push(vertex.lng());
-                }
-                Series::new(PlSmallStr::from(""), &latlngs)
+
+                // Create a Vec<Vec<f64>> for the boundary: each inner vec is [lat, lng]
+                let latlng_pairs: Vec<Vec<f64>> = boundary
+                    .iter()
+                    .map(|vertex| vec![vertex.lat(), vertex.lng()])
+                    .collect();
+
+                // Convert each [lat, lng] pair into its own Series
+                let inner_series: Vec<Series> = latlng_pairs
+                    .into_iter()
+                    .map(|coords| Series::new(PlSmallStr::from(""), coords))
+                    .collect();
+
+                Series::new(PlSmallStr::from(""), inner_series)
             })
         })
         .collect();

src/engine/indexing.rs

@@ -312,7 +312,7 @@ fn is_valid_vertex(inputs: &[Series]) -> PolarsResult<Series> {
 fn boundary_list_dtype(input_fields: &[Field]) -> PolarsResult<Field> {
     Ok(Field::new(
         input_fields[0].name.clone(),
-        DataType::List(Box::new(DataType::Float64)),
+        DataType::List(Box::new(DataType::List(Box::new(DataType::Float64)))),
     ))
 }
 

src/expressions.rs

@@ -162,7 +162,6 @@ def test_cell_to_boundary_known(test_params):
     )
     boundary = df["boundary"][0]
     for i, (exp_lat, exp_lng) in enumerate(test_params["output_boundary"]):
-        lat = boundary[i * 2]
-        lng = boundary[i * 2 + 1]
+        lat, lng = boundary[i]
         assert pytest.approx(lat, abs=1e-7) == exp_lat
         assert pytest.approx(lng, abs=1e-7) == exp_lng