From 8be0f3097c314fe7343f38a47c56c2bdf20e2dc6 Mon Sep 17 00:00:00 2001
From: Wang Boyu <boyu.wby@gmail.com>
Date: Mon, 8 Aug 2022 19:57:07 +0800
Subject: [PATCH] update examples to export data after simulation

---
 examples/geo_schelling/model.py      | 11 ++++++++++-
 examples/geo_schelling/server.py     |  1 +
 examples/rainfall/rainfall/model.py  | 17 ++++++++++++++++-
 examples/rainfall/rainfall/server.py |  2 ++
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/examples/geo_schelling/model.py b/examples/geo_schelling/model.py
index 486ebf42..3e50b314 100644
--- a/examples/geo_schelling/model.py
+++ b/examples/geo_schelling/model.py
@@ -56,9 +56,10 @@ def __repr__(self):
 class GeoSchelling(Model):
     """Model class for the Schelling segregation model."""
 
-    def __init__(self, density=0.6, minority_pc=0.2):
+    def __init__(self, density=0.6, minority_pc=0.2, export_data=False):
         self.density = density
         self.minority_pc = minority_pc
+        self.export_data = export_data
 
         self.schedule = RandomActivation(self)
         self.space = GeoSpace(warn_crs_conversion=False)
@@ -82,6 +83,11 @@ def __init__(self, density=0.6, minority_pc=0.2):
                     agent.atype = 0
                 self.schedule.add(agent)
 
+    def export_agents_to_file(self) -> None:
+        self.space.get_agents_as_GeoDataFrame(agent_cls=SchellingAgent).to_crs(
+            "epsg:4326"
+        ).to_file("data/schelling_agents.geojson", driver="GeoJSON")
+
     def step(self):
         """Run one step of the model.
 
@@ -93,3 +99,6 @@ def step(self):
 
         if self.happy == self.schedule.get_agent_count():
             self.running = False
+
+        if not self.running and self.export_data:
+            self.export_agents_to_file()
diff --git a/examples/geo_schelling/server.py b/examples/geo_schelling/server.py
index a172d86b..2807eb67 100644
--- a/examples/geo_schelling/server.py
+++ b/examples/geo_schelling/server.py
@@ -21,6 +21,7 @@ def render(self, model):
 model_params = {
     "density": mesa.visualization.Slider("Agent density", 0.6, 0.1, 1.0, 0.1),
     "minority_pc": mesa.visualization.Slider("Fraction minority", 0.2, 0.00, 1.0, 0.05),
+    "export_data": mesa.visualization.Checkbox("Export data after simulation", False),
 }
 
 
diff --git a/examples/rainfall/rainfall/model.py b/examples/rainfall/rainfall/model.py
index 55b9f7cf..a16a5161 100644
--- a/examples/rainfall/rainfall/model.py
+++ b/examples/rainfall/rainfall/model.py
@@ -56,10 +56,12 @@ def step(self):
 
 
 class Rainfall(mesa.Model):
-    def __init__(self, rain_rate=500, water_height=5):
+    def __init__(self, rain_rate=500, water_height=5, export_data=False, num_steps=20):
         super().__init__()
         self.rain_rate = rain_rate
         self.water_amount = 0
+        self.export_data = export_data
+        self.num_steps = num_steps
 
         self.space = CraterLake(crs="epsg:4326", water_height=water_height)
         self.schedule = mesa.time.RandomActivation(self)
@@ -81,6 +83,13 @@ def contained(self):
     def outflow(self):
         return self.space.outflow
 
+    def export_water_level_to_file(self):
+        self.space.raster_layer.to_file(
+            raster_file="data/water_level.asc",
+            attr_name="water_level",
+            driver="AAIGrid",
+        )
+
     def step(self):
         for _ in range(self.rain_rate):
             random_x = np.random.randint(0, self.space.raster_layer.width)
@@ -96,3 +105,9 @@ def step(self):
 
         self.schedule.step()
         self.datacollector.collect(self)
+
+        self.num_steps -= 1
+        if self.num_steps == 0:
+            self.running = False
+        if not self.running and self.export_data:
+            self.export_water_level_to_file()
diff --git a/examples/rainfall/rainfall/server.py b/examples/rainfall/rainfall/server.py
index c07cbdd4..cab7e0ac 100644
--- a/examples/rainfall/rainfall/server.py
+++ b/examples/rainfall/rainfall/server.py
@@ -10,6 +10,8 @@
 model_params = {
     "rain_rate": mesa.visualization.Slider("rain rate", 500, 0, 500, 5),
     "water_height": mesa.visualization.Slider("water height", 5, 1, 5, 1),
+    "num_steps": mesa.visualization.Slider("total number of steps", 20, 1, 100, 1),
+    "export_data": mesa.visualization.Checkbox("export data after simulation", False),
 }