From b5cf3c808f56a5614dcbc8195d87eda90e59a637 Mon Sep 17 00:00:00 2001
From: pveigadecamargo <pveigadecamargo@anl.gov>
Date: Sun, 8 Sep 2024 19:03:16 +1000
Subject: [PATCH] brings GeoPandas to the fold

---
 aequilibrae/project/data_loader.py   | 19 +++++++++++--------
 aequilibrae/project/network/links.py |  6 +++---
 aequilibrae/project/network/nodes.py |  5 +++--
 aequilibrae/project/zoning.py        |  9 ++++-----
 4 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/aequilibrae/project/data_loader.py b/aequilibrae/project/data_loader.py
index 23178e618..370e7be68 100644
--- a/aequilibrae/project/data_loader.py
+++ b/aequilibrae/project/data_loader.py
@@ -1,7 +1,8 @@
 from os import PathLike
+from typing import Union
 
+import geopandas as gpd
 import pandas as pd
-import shapely.wkb
 
 from aequilibrae.utils.db_utils import commit_and_close
 from aequilibrae.utils.spatialite_utils import connect_spatialite
@@ -12,17 +13,19 @@ def __init__(self, path_to_file: PathLike, table_name: str):
         self.__pth_file = path_to_file
         self.table_name = table_name
 
-    def load_table(self) -> pd.DataFrame:
+    def load_table(self) -> Union[gpd.GeoDataFrame, pd.DataFrame]:
         with commit_and_close(connect_spatialite(self.__pth_file)) as conn:
             fields, _, geo_field = self.__find_table_fields()
             fields = [f'"{x}"' for x in fields]
+            keys = ','.join(fields)
             if geo_field is not None:
-                fields.append('ST_AsBinary("geometry") geometry')
-            keys = ",".join(fields)
-            df = pd.read_sql_query(f"select {keys} from '{self.table_name}'", conn)
-        if geo_field is not None:
-            df.geometry = df.geometry.apply(shapely.wkb.loads)
-        return df
+                keys += 'Hex(ST_AsBinary("geometry")) as geometry'
+
+            sql = f"select {keys} from '{self.table_name}'"
+            if geo_field is None:
+                return pd.read_sql_query(sql, conn)
+            else:
+                return gpd.GeoDataFrame.from_postgis(sql, conn, geom_col="geometry", crs="EPSG:4326")
 
     def __find_table_fields(self):
         with commit_and_close(connect_spatialite(self.__pth_file)) as conn:
diff --git a/aequilibrae/project/network/links.py b/aequilibrae/project/network/links.py
index 23558743d..f1997824c 100644
--- a/aequilibrae/project/network/links.py
+++ b/aequilibrae/project/network/links.py
@@ -1,6 +1,6 @@
 from copy import deepcopy
 
-import pandas as pd
+import geopandas as gpd
 import shapely.wkb
 
 from aequilibrae.project.basic_table import BasicTable
@@ -136,11 +136,11 @@ def refresh_fields(self) -> None:
         self.__fields = deepcopy(tl.fields)
 
     @property
-    def data(self) -> pd.DataFrame:
+    def data(self) -> gpd.GeoDataFrame:
         """Returns all links data as a Pandas DataFrame
 
         :Returns:
-            **table** (:obj:`DataFrame`): Pandas dataframe with all the links, complete with Geometry
+            **table** (:obj:`GeoDataFrame`): GeoPandas GeoDataFrame with all the nodes
         """
         dl = DataLoader(self.project.path_to_file, "links")
         return dl.load_table()
diff --git a/aequilibrae/project/network/nodes.py b/aequilibrae/project/network/nodes.py
index 9fcae40ec..98dc1bb25 100644
--- a/aequilibrae/project/network/nodes.py
+++ b/aequilibrae/project/network/nodes.py
@@ -1,5 +1,6 @@
 from copy import deepcopy
 
+import geopandas as gpd
 import pandas as pd
 
 from aequilibrae.project.basic_table import BasicTable
@@ -111,11 +112,11 @@ def save(self):
             item.save()
 
     @property
-    def data(self) -> pd.DataFrame:
+    def data(self) -> gpd.GeoDataFrame:
         """Returns all nodes data as a Pandas DataFrame
 
         :Returns:
-            **table** (:obj:`DataFrame`): Pandas DataFrame with all the nodes, complete with Geometry
+            **table** (:obj:`GeoDataFrame`): GeoPandas GeoDataFrame with all the nodes
         """
         dl = DataLoader(self.project.path_to_file, "nodes")
         return dl.load_table()
diff --git a/aequilibrae/project/zoning.py b/aequilibrae/project/zoning.py
index f50167be8..0970c591c 100644
--- a/aequilibrae/project/zoning.py
+++ b/aequilibrae/project/zoning.py
@@ -2,15 +2,14 @@
 from os.path import join, realpath
 from typing import Union, Dict
 
-import pandas as pd
-
+import geopandas as gpd
 import shapely.wkb
 from shapely.geometry import Point, Polygon, LineString, MultiLineString
 from shapely.ops import unary_union
 
 from aequilibrae.project.basic_table import BasicTable
-from aequilibrae.project.project_creation import run_queries_from_sql_file
 from aequilibrae.project.data_loader import DataLoader
+from aequilibrae.project.project_creation import run_queries_from_sql_file
 from aequilibrae.project.table_loader import TableLoader
 from aequilibrae.project.zone import Zone
 from aequilibrae.utils.db_utils import commit_and_close
@@ -166,11 +165,11 @@ def __create_return_zone(self, data):
         return zone
 
     @property
-    def data(self) -> pd.DataFrame:
+    def data(self) -> gpd.GeoDataFrame:
         """Returns all zones data as a Pandas DataFrame
 
         :Returns:
-            **table** (:obj:`DataFrame`): Pandas DataFrame with all the zones, complete with Geometry
+            **table** (:obj:`GeoDataFrame`): GeoPandas GeoDataFrame with all the nodes
         """
         dl = DataLoader(self.project.path_to_file, "zones")
         return dl.load_table()