Repository: kraina-ai/srai (language: python)

Issue:
Add S2Regionizer



Discussion / hints:

Fix patch:
diff --git a/examples/regionizers/s2_regionizer.ipynb b/examples/regionizers/s2_regionizer.ipynb
new file mode 100644
index 00000000..5df73b4c
--- /dev/null
+++ b/examples/regionizers/s2_regionizer.ipynb
@@ -0,0 +1,136 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# S2 Regionizer Example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import geopandas as gpd\n",
+    "import matplotlib.pyplot as plt\n",
+    "from shapely import geometry\n",
+    "\n",
+    "from srai.regionizers import S2Regionizer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define starting polygons"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gdf = gpd.GeoDataFrame(\n",
+    "    geometry=[\n",
+    "        geometry.Polygon(\n",
+    "            shell=[\n",
+    "                (-1, 0),\n",
+    "                (0, 0.5),\n",
+    "                (1, 0),\n",
+    "                (1, 1),\n",
+    "                (0, 1),\n",
+    "            ],\n",
+    "            holes=[\n",
+    "                [\n",
+    "                    (0.8, 0.9),\n",
+    "                    (0.9, 0.55),\n",
+    "                    (0.8, 0.3),\n",
+    "                    (0.5, 0.4),\n",
+    "                ]\n",
+    "            ],\n",
+    "        ),\n",
+    "        geometry.Polygon(shell=[(-0.25, 0), (0.25, 0), (0, 0.2)]),\n",
+    "    ],\n",
+    "    crs=\"epsg:4326\",\n",
+    ")\n",
+    "resolution = 11\n",
+    "gdf.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Regionize using S2 wihtout buffer\n",
+    "\n",
+    "All cells are fully contained in the original polygon."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "regionizer = S2Regionizer(resolution, buffer=False)\n",
+    "gdf_s2 = regionizer.transform(gdf)\n",
+    "\n",
+    "ax = gdf.plot()\n",
+    "gdf_s2.plot(ax=ax, color=\"red\", alpha=0.5)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Regionizer using S2 with buffer (default)\n",
+    "\n",
+    "Whole polygon is covered by cells. Some are not fully contained in the original polygon."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "regionizer = S2Regionizer(resolution, buffer=True)\n",
+    "gdf_s2 = regionizer.transform(gdf)\n",
+    "\n",
+    "ax = gdf.plot()\n",
+    "gdf_s2.plot(ax=ax, color=\"red\", alpha=0.5)\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.14"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "2102df1d74c890a541339243f341c4aa4ffb38ff8671f11b79e33b00adbba0dd"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/pdm.lock b/pdm.lock
index b6cbda29..dc25881b 100644
--- a/pdm.lock
+++ b/pdm.lock
@@ -351,6 +351,12 @@ version = "4.38.0"
 requires_python = ">=3.7"
 summary = "Tools to manipulate font files"
 
+[[package]]
+name = "future"
+version = "0.18.2"
+requires_python = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*"
+summary = "Clean single-source support for Python 3 and 2"
+
 [[package]]
 name = "geojson"
 version = "2.5.0"
@@ -1443,6 +1449,24 @@ version = "1.0.1"
 requires_python = ">=3.7"
 summary = "R-Tree spatial index for Python GIS"
 
+[[package]]
+name = "s2"
+version = "0.1.9"
+requires_python = ">=3"
+summary = "Python bindings for S2, a hierarchical square geospatial indexing system."
+dependencies = [
+    "s2sphere",
+    "shapely",
+]
+
+[[package]]
+name = "s2sphere"
+version = "0.2.5"
+summary = "Python implementation of the S2 Geometry Library"
+dependencies = [
+    "future>=0.15",
+]
+
 [[package]]
 name = "scalene"
 version = "1.5.15"
@@ -1741,7 +1765,7 @@ summary = "Backport of pathlib-compatible object wrapper for zip files"
 
 [metadata]
 lock_version = "4.0"
-content_hash = "sha256:002b304234e9c082e87eb419fd00d8f1c857fa0d6b3783fa9d1cf1cb4d85085e"
+content_hash = "sha256:d1a12925fc90261fe9f1548db0256cfb03c4528a5b92c81c6b4782925319413c"
 
 [metadata.files]
 "anyio 3.6.2" = [
@@ -2189,6 +2213,9 @@ content_hash = "sha256:002b304234e9c082e87eb419fd00d8f1c857fa0d6b3783fa9d1cf1cb4
     {url = "https://files.pythonhosted.org/packages/55/5c/a4a25cf6db42d113d8f626901bb156b2f7cf7c7564a6bbc7b5cd6f7cb484/fonttools-4.38.0.zip", hash = "sha256:2bb244009f9bf3fa100fc3ead6aeb99febe5985fa20afbfbaa2f8946c2fbdaf1"},
     {url = "https://files.pythonhosted.org/packages/e3/d9/e9bae85e84737e76ebbcbea13607236da0c0699baed0ae4f1151b728a608/fonttools-4.38.0-py3-none-any.whl", hash = "sha256:820466f43c8be8c3009aef8b87e785014133508f0de64ec469e4efb643ae54fb"},
 ]
+"future 0.18.2" = [
+    {url = "https://files.pythonhosted.org/packages/45/0b/38b06fd9b92dc2b68d58b75f900e97884c45bedd2ff83203d933cf5851c9/future-0.18.2.tar.gz", hash = "sha256:b1bead90b70cf6ec3f0710ae53a525360fa360d306a86583adc6bf83a4db537d"},
+]
 "geojson 2.5.0" = [
     {url = "https://files.pythonhosted.org/packages/b6/8d/c42d3fe6f9b5e5bd6a55d9f03813d674d65d853cb12e6bc56f154a2ceca0/geojson-2.5.0.tar.gz", hash = "sha256:6e4bb7ace4226a45d9c8c8b1348b3fc43540658359f93c3f7e03efa9f15f658a"},
     {url = "https://files.pythonhosted.org/packages/e4/8d/9e28e9af95739e6d2d2f8d4bef0b3432da40b7c3588fbad4298c1be09e48/geojson-2.5.0-py2.py3-none-any.whl", hash = "sha256:ccbd13368dd728f4e4f13ffe6aaf725b6e802c692ba0dde628be475040c534ba"},
@@ -3322,6 +3349,14 @@ content_hash = "sha256:002b304234e9c082e87eb419fd00d8f1c857fa0d6b3783fa9d1cf1cb4
     {url = "https://files.pythonhosted.org/packages/ee/a0/481d67a0f52b23781f6f99543a9a4a2e6a5cffb165a6e50b85b7dd23af0b/Rtree-1.0.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:6db6a0a93e41594ffc14b053f386dd414ab5a82535bbd9aedafa6ac8dc0650d8"},
     {url = "https://files.pythonhosted.org/packages/fc/52/f77645283cc12539e39381ee00f1802ae6af8baa9357bd9b12d23ff61fbd/Rtree-1.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:becd711fe97c2e09b1b7969e83080a3c8012bce2d30f6db879aade255fcba5c1"},
 ]
+"s2 0.1.9" = [
+    {url = "https://files.pythonhosted.org/packages/a3/cd/149942fbf9cd80e2ef62bb4ccac8a016e6ba2ce383ab01ae8dbf8953fe59/s2-0.1.9-py2.py3-none-any.whl", hash = "sha256:8f550c1f65a16a5e2efb64867d12356f20092a2d87b10575104563ef6379c192"},
+    {url = "https://files.pythonhosted.org/packages/b9/b0/efaa760f81b02e11a74d00d9810e1f3b704ee8894be7cf996e4f74759a3f/s2-0.1.9.tar.gz", hash = "sha256:ad7a90d78f00dd9edf5002ebfa16163a6cbea582bba4ad8ea9d22df07a60851e"},
+]
+"s2sphere 0.2.5" = [
+    {url = "https://files.pythonhosted.org/packages/01/8e/08260a5f03c642d3f6f040756425dc3f466842c4d5e978f6c2be960dde2f/s2sphere-0.2.5.tar.gz", hash = "sha256:c2478c1ff7c601a59a7151a57b605435897514578fa6bdb8730721c182adbbaf"},
+    {url = "https://files.pythonhosted.org/packages/49/fc/be2251ac7ee75e725027b107ab14d161d829aaac22362f75eaabc859bfba/s2sphere-0.2.5-py2.py3-none-any.whl", hash = "sha256:d2340c9cf458ddc9a89afd1d8048a4195ce6fa6b0095ab900d4be5271e537401"},
+]
 "scalene 1.5.15" = [
     {url = "https://files.pythonhosted.org/packages/02/f7/2a33b6c6280ab15ee729ec5fa851c29c13f52f78e5a266c8705157b05559/scalene-1.5.15-cp39-cp39-macosx_10_15_universal2.whl", hash = "sha256:d4b71a8671d8856ae1fa54ded79875b44a9f5b7ba4c8dc13921be7cb560f283f"},
     {url = "https://files.pythonhosted.org/packages/16/e9/fb009403f4f6d5b698920ce7e8d92e94cda9d526de58a501f5dc9e1da3bc/scalene-1.5.15-cp39-cp39-win_amd64.whl", hash = "sha256:f5be6865d56e025f551886fd1768548b156696d6022d87ba1f11c4995266ba48"},
diff --git a/pyproject.toml b/pyproject.toml
index d812b3f8..031387d0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,6 +17,7 @@ dependencies = [
     "pyarrow>=10.0.0",
     "topojson>=1.5",
     "tqdm>=4.64.1",
+    "s2>=0.1.9",
 ]
 requires-python = ">=3.8"
 readme = "README.md"
diff --git a/srai/regionizers/__init__.py b/srai/regionizers/__init__.py
index 80608d87..ce27fd70 100644
--- a/srai/regionizers/__init__.py
+++ b/srai/regionizers/__init__.py
@@ -3,11 +3,13 @@
 from .administrative_boundary_regionizer import AdministrativeBoundaryRegionizer
 from .base import BaseRegionizer
 from .h3_regionizer import H3Regionizer
+from .s2_regionizer import S2Regionizer
 from .voronoi_regionizer import VoronoiRegionizer
 
 __all__ = [
     "BaseRegionizer",
     "AdministrativeBoundaryRegionizer",
     "H3Regionizer",
+    "S2Regionizer",
     "VoronoiRegionizer",
 ]
diff --git a/srai/regionizers/s2_regionizer.py b/srai/regionizers/s2_regionizer.py
new file mode 100644
index 00000000..124e2aca
--- /dev/null
+++ b/srai/regionizers/s2_regionizer.py
@@ -0,0 +1,106 @@
+"""
+S2 regionizer.
+
+This module exposes Google's S2 Geospatial Indexing System [1] as a regionizer.
+It uses the Python API [2].
+
+References:
+    [1] https://s2geometry.io/
+    [2] https://github.com/JoaoCarabetta/s2-py
+
+"""
+
+import json
+from typing import Any, Dict
+
+import geopandas as gpd
+from functional import seq
+from functional.pipeline import Sequence
+from s2 import s2
+from shapely.geometry import Polygon
+
+from . import BaseRegionizer
+
+
+class S2Regionizer(BaseRegionizer):
+    """
+    S2 Regionizer.
+
+    S2 Regionizer gives an opportunity to divide the given geometries into square S2 cells.
+
+    """
+
+    def __init__(self, resolution: int, buffer: bool = True) -> None:
+        """
+        Init S2 Regionizer.
+
+        Args:
+            resolution (int): Resolution of the cells (S2 supports 0-30). See [1] for
+            a full comparison.
+            buffer (bool, optional): If True then fully cover geometries with S2 cells.
+            Otherwise only use those cells that fully fit into them. Defaults to True.
+
+        Raises:
+            ValueError: If resolution is not between 0 and 30.
+
+        References:
+            [1] https://s2geometry.io/resources/s2cell_statistics.html
+
+        """
+        if not (0 <= resolution <= 30):
+            raise ValueError(f"Resolution {resolution} is not between 0 and 30.")
+
+        self.resolution = resolution
+        self.buffer = buffer
+
+    def transform(self, gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
+        """
+        Regionize a given GeoDataFrame.
+
+        Args:
+            gdf (gpd.GeoDataFrame): GeoDataFrame to be regionized.
+
+        Returns:
+            gpd.GeoDataFrame: GeoDataFrame with regionized geometries.
+
+        """
+        gdf_wgs84 = gdf.to_crs(epsg=4326)
+
+        # transform multipolygons to multiple polygons
+        gdf_exploded = gdf_wgs84.explode(index_parts=True).reset_index(drop=True)
+
+        s2_gdf = self._fill_with_s2_cells(gdf_exploded)
+
+        # s2 library fills also holes in Polygons, so here we remove redundant cells
+        res: gpd.GeoDataFrame = gpd.sjoin(
+            s2_gdf,
+            gdf_wgs84,
+            how="inner",
+            predicate="intersects" if self.buffer else "within",
+        ).drop(columns=["index_right"])
+
+        res = res[~res.index.duplicated(keep="first")]
+
+        return res
+
+    def _fill_with_s2_cells(self, gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
+        geo_json = json.loads(gdf.to_json())
+        cells = (
+            seq(geo_json["features"])
+            .flat_map(lambda f: self._geojson_to_cells(f["geometry"], self.resolution))
+            .to_dict()
+        )
+        cells_gdf = gpd.GeoDataFrame(
+            None,
+            index=cells.keys(),
+            geometry=list(cells.values()),
+            crs="epsg:4326",
+        )
+
+        return cells_gdf
+
+    def _geojson_to_cells(self, geo_json: Dict[str, Any], res: int) -> Sequence:
+        raw_cells = s2.polyfill(geo_json, res, with_id=True, geo_json_conformant=True)
+        cells: Sequence = seq(raw_cells).map(lambda c: (c["id"], Polygon(c["geometry"])))
+
+        return cells


Test patch:
diff --git a/tests/regionizers/test_s2_regionizer.py b/tests/regionizers/test_s2_regionizer.py
new file mode 100644
index 00000000..98ef70f1
--- /dev/null
+++ b/tests/regionizers/test_s2_regionizer.py
@@ -0,0 +1,58 @@
+"""Tests for H3Regionizer."""
+from contextlib import nullcontext as does_not_raise
+from typing import Any, List
+from unittest import TestCase
+
+import geopandas as gpd
+import pytest
+
+from srai.regionizers import S2Regionizer
+
+ut = TestCase()
+S2_RESOLUTION = 7
+
+
+@pytest.fixture  # type: ignore
+def expected_s2_indexes() -> List[str]:
+    """Get expected s2 indexes."""
+    return [
+        "0555c",
+        "0ffe4",
+        "0fff4",
+        "0fffc",
+        "10004",
+        "1000c",
+        "10014",
+        "1001c",
+        "1aaa4",
+        "05554",
+        "1aaac",
+    ]
+
+
+@pytest.mark.parametrize(  # type: ignore
+    "gdf_fixture,expected_s2_indexes_fixture,resolution,expectation",
+    [
+        ("gdf_polygons", "expected_s2_indexes", S2_RESOLUTION, does_not_raise()),
+        ("gdf_multipolygon", "expected_s2_indexes", S2_RESOLUTION, does_not_raise()),
+        ("gdf_empty", "expected_s2_indexes", S2_RESOLUTION, pytest.raises(AttributeError)),
+        ("gdf_polygons", "expected_s2_indexes", -1, pytest.raises(ValueError)),
+        ("gdf_polygons", "expected_s2_indexes", 31, pytest.raises(ValueError)),
+        ("gdf_no_crs", "expected_s2_indexes", S2_RESOLUTION, pytest.raises(ValueError)),
+    ],
+)
+def test_transform(
+    gdf_fixture: str,
+    expected_s2_indexes_fixture: str,
+    resolution: int,
+    expectation: Any,
+    request: Any,
+) -> None:
+    """Test transform of H3Regionizer."""
+    gdf: gpd.GeoDataFrame = request.getfixturevalue(gdf_fixture)
+    s2_indexes: List[str] = request.getfixturevalue(expected_s2_indexes_fixture)
+    with expectation:
+        gdf_s2 = S2Regionizer(resolution).transform(gdf)
+
+        ut.assertCountEqual(first=gdf_s2.index.to_list(), second=s2_indexes)
+        assert "geometry" in gdf_s2