From 4e1ae8e42539fe350afda9647cfe7ed39047ac2e Mon Sep 17 00:00:00 2001 From: Wang Boyu Date: Tue, 26 Jul 2022 23:14:38 +0800 Subject: [PATCH 1/2] create RasterLayer from file --- mesa_geo/raster_layers.py | 33 ++++++++++++++++++++++++++------- tests/test_RasterLayer.py | 8 ++++---- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/mesa_geo/raster_layers.py b/mesa_geo/raster_layers.py index 618416a1..8b9ae3e6 100644 --- a/mesa_geo/raster_layers.py +++ b/mesa_geo/raster_layers.py @@ -15,6 +15,7 @@ Sequence, Iterator, Iterable, + Type, ) import numpy as np @@ -154,7 +155,7 @@ class RasterLayer(RasterBase): cells: List[List[Cell]] _neighborhood_cache: Dict[Any, List[Coordinate]] - def __init__(self, width, height, crs, total_bounds, cell_cls=Cell): + def __init__(self, width, height, crs, total_bounds, cell_cls: Type[Cell] = Cell): super().__init__(width, height, crs, total_bounds) self.cell_cls = cell_cls self.cells = [] @@ -232,13 +233,13 @@ def coord_iter(self) -> Iterator[Tuple[Cell, int, int]]: for col in range(self.height): yield self.cells[row][col], row, col # cell, x, y - def apply_raster(self, data: np.ndarray, name: str = None) -> None: + def apply_raster(self, data: np.ndarray, attr_name: str = None) -> None: assert data.shape == (1, self.height, self.width) - if name is None: - name = f"attribute_{len(self.cell_cls.__dict__)}" + if attr_name is None: + attr_name = f"attribute_{len(self.cell_cls.__dict__)}" for x in range(self.width): for y in range(self.height): - setattr(self.cells[x][y], name, data[0, self.height - y - 1, x]) + setattr(self.cells[x][y], attr_name, data[0, self.height - y - 1, x]) def iter_neighborhood( self, @@ -396,6 +397,24 @@ def to_image(self, colormap) -> ImageLayer: values[:, row, col] = colormap(cell) return ImageLayer(values=values, crs=self.crs, total_bounds=self.total_bounds) + @classmethod + def from_file( + cls, raster_file, cell_cls: Type[Cell] = Cell, attr_name: str = None + ) -> RasterLayer: + with rio.open(raster_file, "r") as dataset: + values = dataset.read() + _, height, width = values.shape + total_bounds = [ + dataset.bounds.left, + dataset.bounds.bottom, + dataset.bounds.right, + dataset.bounds.top, + ] + obj = cls(width, height, dataset.crs, total_bounds, cell_cls) + obj._transform = dataset.transform + obj.apply_raster(values, attr_name=attr_name) + return obj + class ImageLayer(RasterBase): _values: np.ndarray @@ -458,8 +477,8 @@ def to_crs(self, crs, inplace=False) -> ImageLayer | None: return layer @classmethod - def from_file(cls, raster_file: str) -> ImageLayer: - with rio.open(raster_file, "r") as dataset: + def from_file(cls, image_file) -> ImageLayer: + with rio.open(image_file, "r") as dataset: values = dataset.read() total_bounds = [ dataset.bounds.left, diff --git a/tests/test_RasterLayer.py b/tests/test_RasterLayer.py index 6a33650a..4e8c116a 100644 --- a/tests/test_RasterLayer.py +++ b/tests/test_RasterLayer.py @@ -38,12 +38,12 @@ def test_apple_raster(self): """ self.assertEqual(self.raster_layer.cells[0][1].attribute_5, 3) - self.raster_layer.apply_raster(raster_data, name="elevation") + self.raster_layer.apply_raster(raster_data, attr_name="elevation") self.assertEqual(self.raster_layer.cells[0][1].elevation, 3) def test_get_min_cell(self): self.raster_layer.apply_raster( - np.array([[[1, 2], [3, 4], [5, 6]]]), name="elevation" + np.array([[[1, 2], [3, 4], [5, 6]]]), attr_name="elevation" ) min_cell = min( @@ -63,7 +63,7 @@ def test_get_min_cell(self): self.assertEqual(min_cell.elevation, 1) self.raster_layer.apply_raster( - np.array([[[1, 2], [3, 4], [5, 6]]]), name="water_level" + np.array([[[1, 2], [3, 4], [5, 6]]]), attr_name="water_level" ) min_cell = min( self.raster_layer.get_neighboring_cells( @@ -77,7 +77,7 @@ def test_get_min_cell(self): def test_get_max_cell(self): self.raster_layer.apply_raster( - np.array([[[1, 2], [3, 4], [5, 6]]]), name="elevation" + np.array([[[1, 2], [3, 4], [5, 6]]]), attr_name="elevation" ) max_cell = max( From 94cf2d067ffc0c87d0c673cf54def41cea184373 Mon Sep 17 00:00:00 2001 From: Wang Boyu Date: Tue, 26 Jul 2022 23:45:54 +0800 Subject: [PATCH 2/2] 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