Adds graph filtering by polygon (#595)

* Adds graph filtering by polygon

* Add test

* documentation

---------

Co-authored-by: pveigadecamargo <pveigadecamargo@anl.gov>
Co-authored-by: Renata Imai <renataakemii@yahoo.co.uk>

aequilibrae/project/network/network.py

@@ -267,7 +267,7 @@ def export_to_gmns(self, path: str):
 
         self.logger.info("Network exported successfully")
 
-    def build_graphs(self, fields: list = None, modes: list = None) -> None:
+    def build_graphs(self, fields: list = None, modes: list = None, limit_to_area: Polygon = None) -> None:
         """Builds graphs for all modes currently available in the model
 
         When called, it overwrites all graphs previously created and stored in the networks'
@@ -280,6 +280,10 @@ def build_graphs(self, fields: list = None, modes: list = None) -> None:
             **modes** (:obj:`list`, *Optional*): When working with very large graphs with large number of fields in the
             database, it may be useful to generate only those we need
 
+            **limit_to_area** (:obj:`Polygon`, *Optional*): When working with a very large model area, you may want to
+            filter your database to a small area for your computation, which you can do by providing a polygon.
+            The search is limited to a spatial index search, so it is very fast but NOT PRECISE.
+
         To use the 'fields' parameter, a minimalistic option is the following
 
         .. code-block:: python
@@ -308,19 +312,36 @@ def build_graphs(self, fields: list = None, modes: list = None) -> None:
             elif isinstance(modes, str):
                 modes = [modes]
 
+            if limit_to_area is not None:
+                spatial_add = """ WHERE links.rowid in (
+                                        select rowid from SpatialIndex where f_table_name = 'links' and
+                                       search_frame = GeomFromWKB(?, 4326))"""
+
             sql = f"select {','.join(all_fields)} from links"
 
+            sql_centroids = "select node_id from nodes where is_centroid=1 order by node_id;"
+            centroids = np.array([i[0] for i in conn.execute(sql_centroids).fetchall()], np.uint32)
+            centroids = centroids if centroids.shape[0] else None
+
             with pd.option_context("future.no_silent_downcasting", True):
-                df = pd.read_sql(sql, conn).fillna(value=np.nan).infer_objects(False)
+                if limit_to_area is None:
+                    df = pd.read_sql(sql, conn).fillna(value=np.nan).infer_objects(False)
+                else:
+                    sql += spatial_add
+                    df = (
+                        pd.read_sql_query(sql, conn, params=(limit_to_area.wkb,))
+                        .fillna(value=np.nan)
+                        .infer_objects(False)
+                    )
+
+                    # We filter to centroids existing in our filtered area
+                    centroids = centroids[np.isin(centroids, df.a_node) | np.isin(centroids, df.b_node)]
+
             valid_fields = list(df.select_dtypes(np.number).columns) + ["modes"]
-            sql = "select node_id from nodes where is_centroid=1 order by node_id;"
-            centroids = np.array([i[0] for i in conn.execute(sql).fetchall()], np.uint32)
-            centroids = centroids if centroids.shape[0] else None
 
         lonlat = self.nodes.lonlat.set_index("node_id")
         data = df[valid_fields]
         for m in modes:
-
             # For any link in net that doesn't support mode 'm', set a_node = b_node (these will be culled when
             # the compressed graph representation is created)
             net = pd.DataFrame(data, copy=True)

conftest.py

@@ -10,6 +10,7 @@
 
 import pandas as pd
 import numpy as np
+from shapely.geometry import Polygon
 
 from aequilibrae import Project
 from aequilibrae.project.database_connection import database_connection
@@ -113,3 +114,4 @@ def doctest_fixtures(doctest_namespace, create_path, tmp_path_factory):
     doctest_namespace["os"] = os
     doctest_namespace["pd"] = pd
     doctest_namespace["np"] = np
+    doctest_namespace["Polygon"] = Polygon

docs/source/modeling_with_aequilibrae/static_traffic_assignment/assignment_mechanics.rst

@@ -79,6 +79,16 @@ API is a method call for excluding one or more links from the Graph, **which is
 
     >>> graph.exclude_links([123, 975])
 
+When working with very large networks, it is possible to filter the database to a small
+area for computation by providing a polygon that delimits the desired area, instead of
+selecting the links for deletion. The selection of links and nodes is limited to a spatial
+inidex search, which is very fast but not accurate.
+
+.. code-block:: python
+
+    >>> polygon = Polygon([(-71.35, -29.95), (-71.35, -29.90), (-71.30, -29.90), (-71.30, -29.95), (-71.35, -29.95)])
+    >>> project.network.build_graphs(limit_to_area=polygon) # doctest: +SKIP
+
 More sophisticated graph editing is also possible, but it is recommended that
 changes to be made in the network DataFrame. For example:
 

tests/aequilibrae/project/test_network.py

@@ -5,7 +5,7 @@
 from unittest import TestCase
 from warnings import warn
 
-from shapely.geometry import box
+from shapely.geometry import box, Polygon
 
 from aequilibrae.project import Project
 from ...data import siouxfalls_project
@@ -74,3 +74,29 @@ def test_count_links(self):
     def test_count_nodes(self):
         items = self.siouxfalls.network.count_nodes()
         self.assertEqual(24, items, "Wrong number of nodes found")
+
+    def test_build_graphs_with_polygons(self):
+        coords = ((-96.75, 43.50), (-96.75, 43.55), (-96.70, 43.55), (-96.70, 43.50), (-96.75, 43.50))
+        polygon = Polygon(coords)
+
+        fields = ["distance"]
+        modes = ["c"]
+
+        self.siouxfalls.network.build_graphs(fields, modes, polygon)
+        assert len(self.siouxfalls.network.graphs) == 1
+
+        g = self.siouxfalls.network.graphs["c"]
+        assert g.num_nodes == 19
+        assert g.num_links == 52
+
+        existing_nodes = [i for i in range(1, 25) if i not in [1, 2, 3, 6, 7]]
+        assert list(g.centroids) == existing_nodes
+
+    def test_build_graphs_without_polygons(self):
+        self.siouxfalls.network.build_graphs()
+        assert len(self.siouxfalls.network.graphs) == 3
+
+        g = self.siouxfalls.network.graphs["c"]
+        assert g.num_nodes == 24
+        assert g.num_links == 76
+        assert list(g.centroids) == list(range(1, 25))