From c72e6795ef0134201d33a5b30f34e034466a7cad Mon Sep 17 00:00:00 2001
From: Wang Boyu <boyu.wby@gmail.com>
Date: Sun, 7 Aug 2022 20:20:26 +0800
Subject: [PATCH] extract geoagents as geodataframe

---
 mesa_geo/geospace.py   | 19 +++++++++++++++++++
 tests/test_GeoSpace.py | 18 ++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/mesa_geo/geospace.py b/mesa_geo/geospace.py
index 38e6d3e1..ab9a1e09 100644
--- a/mesa_geo/geospace.py
+++ b/mesa_geo/geospace.py
@@ -226,6 +226,9 @@ def get_neighbors(self, agent):
         """Get (touching) neighbors of an agent."""
         return self._agent_layer.get_neighbors(agent)
 
+    def get_agents_as_GeoDataFrame(self, agent_cls=GeoAgent) -> gpd.GeoDataFrame:
+        return self._agent_layer.get_agents_as_GeoDataFrame(agent_cls)
+
 
 class _AgentLayer:
     """Layer that contains the GeoAgents. Mainly for internal usage within `GeoSpace`.
@@ -368,3 +371,19 @@ def get_neighbors(self, agent):
         neighbors_idx = self._neighborhood.neighbors[idx]
         neighbors = [self.agents[i] for i in neighbors_idx]
         return neighbors
+
+    def get_agents_as_GeoDataFrame(self, agent_cls=GeoAgent) -> gpd.GeoDataFrame:
+        agents_list = []
+        crs = None
+        for agent in self.agents:
+            if isinstance(agent, agent_cls):
+                crs = agent.crs
+                agent_dict = {
+                    attr: value
+                    for attr, value in vars(agent).items()
+                    if attr not in {"model", "pos", "_crs"}
+                }
+                agents_list.append(agent_dict)
+        agents_gdf = gpd.GeoDataFrame.from_records(agents_list, index="unique_id")
+        agents_gdf.crs = crs
+        return agents_gdf
diff --git a/tests/test_GeoSpace.py b/tests/test_GeoSpace.py
index a473d88f..460b997a 100644
--- a/tests/test_GeoSpace.py
+++ b/tests/test_GeoSpace.py
@@ -4,6 +4,7 @@
 import warnings
 
 import numpy as np
+import pandas as pd
 import geopandas as gpd
 from shapely.geometry import Point
 
@@ -116,3 +117,20 @@ def test_get_neighbors_within_distance(self):
             self.geo_space.get_neighbors_within_distance(agent_to_check, distance=1.0)
         )
         self.assertEqual(len(neighbors), 7)
+
+    def test_get_agents_as_GeoDataFrame(self):
+        self.geo_space.add_agents(self.agents)
+
+        agents_list = [
+            {"geometry": agent.geometry, "unique_id": agent.unique_id}
+            for agent in self.agents
+        ]
+        agents_gdf = gpd.GeoDataFrame.from_records(agents_list, index="unique_id")
+        agents_gdf.crs = self.geo_space.crs
+
+        pd.testing.assert_frame_equal(
+            self.geo_space.get_agents_as_GeoDataFrame(), agents_gdf
+        )
+        self.assertEqual(
+            self.geo_space.get_agents_as_GeoDataFrame().crs, agents_gdf.crs
+        )