feat: add adjacency neighbourhood

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - OSMTagLoader
 - Neighbourhood
 - H3Neighbourhood
+- AdjacencyNeighbourhood
 
 ### Changed
 - Change embedders and joiners interface to have `.transform` method

pyproject.toml

@@ -157,13 +157,13 @@ using = "requirements"
 zero = false
 ignore_licenses = ["UNKNOWN"]
 ignore_packages = [
-    'srai',           # skip self
-    'scalene',        # uses Apache-2.0 license, takes time to analyse
-    'docformatter',   # uses MIT license, has mismatched license in analysis
-    'pymap3d',        # uses BSD-2 license, has mismatched license in analysis
-    'mkdocs-jupyter', # uses Apache-2.0 license, has mismatched license in analysis
+    'srai',                     # skip self
+    'scalene',                  # uses Apache-2.0 license, takes time to analyse
+    'docformatter',             # uses MIT license, has mismatched license in analysis
+    'pymap3d',                  # uses BSD-2 license, has mismatched license in analysis
+    'mkdocs-jupyter',           # uses Apache-2.0 license, has mismatched license in analysis
     'nvidia-cuda-runtime-cu11', # uses NVIDIA license
-    'nvidia-cudnn-cu11', # uses NVIDIA license
-    'nvidia-cublas-cu11', # uses NVIDIA license
-    'nvidia-cuda-nvrtc-cu11', # uses NVIDIA license
+    'nvidia-cudnn-cu11',        # uses NVIDIA license
+    'nvidia-cublas-cu11',       # uses NVIDIA license
+    'nvidia-cuda-nvrtc-cu11',   # uses NVIDIA license
 ]

srai/neighbourhoods/__init__.py

@@ -1,5 +1,6 @@
 """Neighbourhoods."""
-from .h3 import H3Neighbourhood
-from .neighbourhood import Neighbourhood
+from ._base import Neighbourhood
+from .adjacency_neighbourhood import AdjacencyNeighbourhood
+from .h3_neighbourhood import H3Neighbourhood
 
-__all__ = ["Neighbourhood", "H3Neighbourhood"]
+__all__ = ["Neighbourhood", "AdjacencyNeighbourhood", "H3Neighbourhood"]

srai/neighbourhoods/adjacency_neighbourhood.py

@@ -0,0 +1,82 @@
+"""
+Adjacency neighbourhood.
+
+This module contains the AdjacencyNeighbourhood class, that allows to get the neighbours of any
+region based on its borders.
+"""
+from typing import Dict, Hashable, Set
+
+import geopandas as gpd
+
+from srai.neighbourhoods import Neighbourhood
+
+
+class AdjacencyNeighbourhood(Neighbourhood[Hashable]):
+    """
+    Adjacency Neighbourhood.
+
+    This class allows to get the neighbours of any region based on common border. Additionally, a
+    lookup table is implemented to accelerate repeated queries.
+
+    By default, a lookup table will be populated lazily based on queries. A dedicated function
+    `generate_neighbourhoods` allows for precalculation of all the neighbourhoods at once.
+    """
+
+    def __init__(self, regions_gdf: gpd.GeoDataFrame) -> None:
+        """
+        Init AdjacencyNeighbourhood.
+
+        Args:
+            regions_gdf (gpd.GeoDataFrame): regions for which a neighbourhood will be calculated.
+
+        Raises:
+            AttributeError: If regions_gdf doesn't have geometry column.
+        """
+        if "geometry" not in regions_gdf.columns:
+            raise AttributeError("Regions must have a geometry column.")
+        self.regions_gdf = regions_gdf
+        self.lookup: Dict[Hashable, Set[Hashable]] = {}
+
+    def generate_neighbourhoods(self) -> None:
+        """Generate the lookup table for all regions."""
+        for region_id in self.regions_gdf.index:
+            if region_id not in self.lookup:
+                self.lookup[region_id] = self._get_adjacent_neighbours(region_id)
+
+    def get_neighbours(self, index: Hashable) -> Set[Hashable]:
+        """
+        Get the direct neighbours of any region using its index.
+
+        Args:
+            index (Hashable): Unique identifier of the region.
+
+        Returns:
+            Set[Hashable]: Indexes of the neighbours.
+        """
+        if self._index_incorrect(index):
+            return set()
+
+        if index not in self.lookup:
+            self.lookup[index] = self._get_adjacent_neighbours(index)
+
+        return self.lookup[index]
+
+    def _get_adjacent_neighbours(self, index: Hashable) -> Set[Hashable]:
+        """
+        Get the direct neighbours of a region using `touches` [1] operator from the Shapely library.
+
+        Args:
+            index (Hashable): Unique identifier of the region.
+
+        Returns:
+            Set[Hashable]: Indexes of the neighbours.
+
+        References:
+            1. https://shapely.readthedocs.io/en/stable/reference/shapely.touches.html
+        """
+        current_region = self.regions_gdf.loc[index]
+        neighbours = self.regions_gdf[self.regions_gdf.geometry.touches(current_region["geometry"])]
+        return set(neighbours.index)
+
+    def _index_incorrect(self, index: Hashable) -> bool:
+        return index not in self.regions_gdf.index

srai/neighbourhoods/h3.py → srai/neighbourhoods/h3_neighbourhood.py

@@ -7,7 +7,7 @@
 
 import h3
 
-from .neighbourhood import Neighbourhood
+from srai.neighbourhoods import Neighbourhood
 
 
 class H3Neighbourhood(Neighbourhood[str]):

tests/neighbourhoods/test_adjacency_neighbourhood.py

@@ -0,0 +1,214 @@
+from typing import Set
+
+import geopandas as gpd
+import pytest
+from shapely import geometry
+
+from srai.neighbourhoods import AdjacencyNeighbourhood
+from srai.utils.constants import WGS84_CRS
+
+
+@pytest.fixture  # type: ignore
+def no_geometry_gdf() -> gpd.GeoDataFrame:
+    """Get GeoDataFrame with no geometry."""
+    return gpd.GeoDataFrame()
+
+
+@pytest.fixture  # type: ignore
+def empty_gdf() -> gpd.GeoDataFrame:
+    """Get empty GeoDataFrame."""
+    return gpd.GeoDataFrame(geometry=[])
+
+
+@pytest.fixture  # type: ignore
+def squares_regions_fixture() -> gpd.GeoDataFrame:
+    """
+    This GeoDataFrame represents 9 squares on a cartesian plane in a 3 by 3 grid pattern. Squares
+    are adjacent by edges and by vertices. Squares are given compass directions. Visually it looks
+    like this ("C" means "CENTER"):
+
+    [["NW", "N", "NE"],
+     [ "W", "C",  "E"],
+     ["SW", "S", "SE"]]
+
+    Returns:
+        GeoDataFrame: A GeoDataFrame containing square regions.
+    """
+    regions = gpd.GeoDataFrame(
+        geometry=[
+            geometry.box(minx=0, maxx=1, miny=0, maxy=1),
+            geometry.box(minx=1, maxx=2, miny=0, maxy=1),
+            geometry.box(minx=2, maxx=3, miny=0, maxy=1),
+            geometry.box(minx=0, maxx=1, miny=1, maxy=2),
+            geometry.box(minx=1, maxx=2, miny=1, maxy=2),
+            geometry.box(minx=2, maxx=3, miny=1, maxy=2),
+            geometry.box(minx=0, maxx=1, miny=2, maxy=3),
+            geometry.box(minx=1, maxx=2, miny=2, maxy=3),
+            geometry.box(minx=2, maxx=3, miny=2, maxy=3),
+        ],
+        index=["SW", "S", "SE", "W", "CENTER", "E", "NW", "N", "NE"],  # compass directions
+        crs=WGS84_CRS,
+    )
+    return regions
+
+
+@pytest.fixture  # type: ignore
+def rounded_regions_fixture() -> gpd.GeoDataFrame:
+    """
+    This GeoDataFrame represents 9 small squares with buffer on a cartesian plane in a 3 by 3 grid
+    pattern. Regions are adjacent by edges, but not by vertices. Regions are given compass
+    directions. Visually it looks like this ("C" means "CENTER"):
+
+    [["NW", "N", "NE"],
+     [ "W", "C",  "E"],
+     ["SW", "S", "SE"]]
+
+    Returns:
+        GeoDataFrame: A GeoDataFrame containing rounded regions.
+    """
+    regions = gpd.GeoDataFrame(
+        geometry=[
+            geometry.box(minx=0, maxx=0.5, miny=0, maxy=0.5).buffer(0.25),
+            geometry.box(minx=1, maxx=1.5, miny=0, maxy=0.5).buffer(0.25),
+            geometry.box(minx=2, maxx=2.5, miny=0, maxy=0.5).buffer(0.25),
+            geometry.box(minx=0, maxx=0.5, miny=1, maxy=1.5).buffer(0.25),
+            geometry.box(minx=1, maxx=1.5, miny=1, maxy=1.5).buffer(0.25),
+            geometry.box(minx=2, maxx=2.5, miny=1, maxy=1.5).buffer(0.25),
+            geometry.box(minx=0, maxx=0.5, miny=2, maxy=2.5).buffer(0.25),
+            geometry.box(minx=1, maxx=1.5, miny=2, maxy=2.5).buffer(0.25),
+            geometry.box(minx=2, maxx=2.5, miny=2, maxy=2.5).buffer(0.25),
+        ],
+        index=["SW", "S", "SE", "W", "CENTER", "E", "NW", "N", "NE"],  # compass directions
+        crs=WGS84_CRS,
+    )
+    return regions
+
+
+def test_no_geometry_gdf_attribute_error(no_geometry_gdf: gpd.GeoDataFrame) -> None:
+    """Test checks if GeoDataFrames without geometry are disallowed."""
+    with pytest.raises(AttributeError):
+        AdjacencyNeighbourhood(no_geometry_gdf)
+
+
+def test_empty_gdf_empty_set(empty_gdf: gpd.GeoDataFrame) -> None:
+    """Test checks if empty GeoDataFrames return empty neighbourhoods."""
+    neighbourhood = AdjacencyNeighbourhood(empty_gdf)
+    assert neighbourhood.get_neighbours(1) == set()
+
+
+def test_lazy_loading_empty_set(squares_regions_fixture: gpd.GeoDataFrame) -> None:
+    """Test checks if lookup table is empty after init."""
+    neighbourhood = AdjacencyNeighbourhood(squares_regions_fixture)
+    assert neighbourhood.lookup == {}
+
+
+def test_adjacency_lazy_loading(rounded_regions_fixture: gpd.GeoDataFrame) -> None:
+    """Test checks if lookup table is lazily populated."""
+    neighbourhood = AdjacencyNeighbourhood(rounded_regions_fixture)
+    neighbours = neighbourhood.get_neighbours("SW")
+    assert neighbours == {"W", "S"}
+    assert neighbourhood.lookup == {
+        "SW": {"W", "S"},
+    }
+
+
+def test_generate_all_neighbourhoods_rounded_regions(
+    rounded_regions_fixture: gpd.GeoDataFrame,
+) -> None:
+    """Test checks `generate_neighbourhoods` function with rounded regions."""
+    neighbourhood = AdjacencyNeighbourhood(rounded_regions_fixture)
+    neighbourhood.generate_neighbourhoods()
+    assert neighbourhood.lookup == {
+        "SW": {"W", "S"},
+        "S": {"SW", "CENTER", "SE"},
+        "SE": {"E", "S"},
+        "W": {"SW", "CENTER", "NW"},
+        "CENTER": {"N", "W", "E", "S"},
+        "E": {"CENTER", "NE", "SE"},
+        "NW": {"N", "W"},
+        "N": {"CENTER", "NE", "NW"},
+        "NE": {"N", "E"},
+    }
+
+
+def test_generate_all_neighbourhoods_squares_regions(
+    squares_regions_fixture: gpd.GeoDataFrame,
+) -> None:
+    """Test checks `generate_neighbourhoods` function with square regions."""
+    neighbourhood = AdjacencyNeighbourhood(squares_regions_fixture)
+    neighbourhood.generate_neighbourhoods()
+    assert neighbourhood.lookup == {
+        "SW": {"W", "S", "CENTER"},
+        "S": {"SW", "W", "CENTER", "SE", "E"},
+        "SE": {"E", "S", "CENTER"},
+        "W": {"N", "SW", "S", "CENTER", "NW"},
+        "CENTER": {"SW", "N", "W", "S", "SE", "E", "NW", "NE"},
+        "E": {"N", "S", "CENTER", "SE", "NE"},
+        "NW": {"W", "N", "CENTER"},
+        "N": {"W", "CENTER", "E", "NW", "NE"},
+        "NE": {"E", "N", "CENTER"},
+    }
+
+
+@pytest.mark.parametrize(  # type: ignore
+    "index, distance, neighbours_expected",
+    [
+        ("SW", -2, set()),
+        ("SW", -1, set()),
+        ("SW", 0, set()),
+        ("SW", 1, {"S", "W"}),
+        ("SW", 2, {"CENTER", "SE", "NW"}),
+        ("SW", 3, {"N", "E"}),
+        ("SW", 4, {"NE"}),
+        ("SW", 5, set()),
+        ("CENTER", 1, {"N", "E", "S", "W"}),
+        ("CENTER", 2, {"NW", "NE", "SW", "SE"}),
+        ("CENTER", 3, set()),
+        ("N", 1, {"NW", "NE", "CENTER"}),
+        ("N", 2, {"E", "S", "W"}),
+        ("N", 3, {"SE", "SW"}),
+        ("N", 4, set()),
+    ],
+)
+def test_adjacency_lazy_loading_at_distance(
+    index: str,
+    distance: int,
+    neighbours_expected: Set[str],
+    rounded_regions_fixture: gpd.GeoDataFrame,
+) -> None:
+    """Test checks `get_neighbours_at_distance` function with rounded regions."""
+    neighbourhood = AdjacencyNeighbourhood(rounded_regions_fixture)
+    neighbours = neighbourhood.get_neighbours_at_distance(index, distance)
+    assert neighbours == neighbours_expected
+
+
+@pytest.mark.parametrize(  # type: ignore
+    "index, distance, neighbours_expected",
+    [
+        ("SW", -2, set()),
+        ("SW", -1, set()),
+        ("SW", 0, set()),
+        ("SW", 1, {"S", "W"}),
+        ("SW", 2, {"S", "W", "CENTER", "SE", "NW"}),
+        ("SW", 3, {"S", "W", "CENTER", "SE", "NW", "N", "E"}),
+        ("SW", 4, {"S", "W", "CENTER", "SE", "NW", "N", "E", "NE"}),
+        ("SW", 5, {"S", "W", "CENTER", "SE", "NW", "N", "E", "NE"}),
+        ("CENTER", 1, {"N", "E", "S", "W"}),
+        ("CENTER", 2, {"N", "E", "S", "W", "NW", "NE", "SW", "SE"}),
+        ("CENTER", 3, {"N", "E", "S", "W", "NW", "NE", "SW", "SE"}),
+        ("N", 1, {"NW", "NE", "CENTER"}),
+        ("N", 2, {"NW", "NE", "CENTER", "E", "S", "W"}),
+        ("N", 3, {"NW", "NE", "CENTER", "E", "S", "W", "SE", "SW"}),
+        ("N", 4, {"NW", "NE", "CENTER", "E", "S", "W", "SE", "SW"}),
+    ],
+)
+def test_adjacency_lazy_loading_up_to_distance(
+    index: str,
+    distance: int,
+    neighbours_expected: Set[str],
+    rounded_regions_fixture: gpd.GeoDataFrame,
+) -> None:
+    """Test checks `get_neighbours_up_to_distance` function with rounded regions."""
+    neighbourhood = AdjacencyNeighbourhood(rounded_regions_fixture)
+    neighbours = neighbourhood.get_neighbours_up_to_distance(index, distance)
+    assert neighbours == neighbours_expected