From 00cfc35699aee858d82329f2738e81ed6eb43f72 Mon Sep 17 00:00:00 2001
From: Wang Boyu <boyu.wby@gmail.com>
Date: Tue, 26 Jul 2022 23:45:54 +0800
Subject: [PATCH] update examples to create RasterLayer from file

---
 examples/rainfall/rainfall/space.py         | 16 ++++++++--------
 examples/uganda/uganda/space.py             | 19 +++++--------------
 examples/urban_growth/urban_growth/space.py |  2 +-
 3 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/examples/rainfall/rainfall/space.py b/examples/rainfall/rainfall/space.py
index 6e1f39f0..3dbcdf8f 100644
--- a/examples/rainfall/rainfall/space.py
+++ b/examples/rainfall/rainfall/space.py
@@ -3,7 +3,6 @@
 
 import mesa
 import numpy as np
-import rasterio as rio
 
 from mesa_geo import Cell, RasterLayer
 from mesa_geo.geospace import GeoSpace
@@ -34,13 +33,14 @@ def __init__(self, crs, water_height):
 
     def set_elevation_layer(self, elevation_gzip_file, crs):
         with gzip.open(elevation_gzip_file, "rb") as elevation_file:
-            with rio.open(elevation_file, "r") as dataset:
-                values = dataset.read()
-                _, height, width = values.shape
-                total_bounds = dataset.bounds
-        raster_layer = RasterLayer(width, height, crs, total_bounds, cell_cls=LakeCell)
-        raster_layer.apply_raster(data=values, name="elevation")
-        raster_layer.apply_raster(data=np.zeros(values.shape), name="water_level")
+            raster_layer = RasterLayer.from_file(
+                elevation_file, cell_cls=LakeCell, attr_name="elevation"
+            )
+            raster_layer.crs = crs
+        raster_layer.apply_raster(
+            data=np.zeros(shape=(1, raster_layer.height, raster_layer.width)),
+            attr_name="water_level",
+        )
         super().add_layer(raster_layer)
 
     @property
diff --git a/examples/uganda/uganda/space.py b/examples/uganda/uganda/space.py
index caf9db28..f634aab4 100644
--- a/examples/uganda/uganda/space.py
+++ b/examples/uganda/uganda/space.py
@@ -5,7 +5,6 @@
 
 import geopandas as gpd
 import mesa
-import rasterio as rio
 
 from mesa_geo.geoagent import GeoAgent
 from mesa_geo.geospace import GeoSpace
@@ -38,20 +37,12 @@ def __init__(self, crs):
     def load_data(self, population_gzip_file, lake_zip_file, world_zip_file):
         world_size = gpd.GeoDataFrame.from_file(world_zip_file)
         with gzip.open(population_gzip_file, "rb") as population_file:
-            with rio.open(population_file, "r") as population_data:
-                values = population_data.read()
-                _, height, width = values.shape
-        self.add_layer(
-            RasterLayer(
-                width,
-                height,
-                world_size.crs,
-                world_size.total_bounds,
-                cell_cls=UgandaCell,
+            raster_layer = RasterLayer.from_file(
+                population_file, cell_cls=UgandaCell, attr_name="population"
             )
-        )
-        self.population_layer.apply_raster(values, name="population")
-
+            raster_layer.crs = world_size.crs
+            raster_layer.total_bounds = world_size.total_bounds
+            self.add_layer(raster_layer)
         self.lake = gpd.GeoDataFrame.from_file(lake_zip_file).geometry[0]
         self.add_agents(GeoAgent(uuid.uuid4().int, None, self.lake, self.crs))
 
diff --git a/examples/urban_growth/urban_growth/space.py b/examples/urban_growth/urban_growth/space.py
index acd49028..3146450d 100644
--- a/examples/urban_growth/urban_growth/space.py
+++ b/examples/urban_growth/urban_growth/space.py
@@ -98,7 +98,7 @@ def load_datasets(
             with gzip.open(data_file, "rb") as f:
                 with rio.open(f, "r") as dataset:
                     values = dataset.read()
-            self.raster_layer.apply_raster(values, name=attribute_name)
+            self.raster_layer.apply_raster(values, attr_name=attribute_name)
 
         for cell in self.raster_layer:
             cell.urban = True if cell.urban == 2 else False