Skip to content

Commit

Permalink
create RasterLayer from file
Browse files Browse the repository at this point in the history
  • Loading branch information
wang-boyu authored and rht committed Jul 27, 2022
1 parent 93e8d60 commit bab4758
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 11 deletions.
33 changes: 26 additions & 7 deletions mesa_geo/raster_layers.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
Sequence,
Iterator,
Iterable,
Type,
)

import numpy as np
Expand Down Expand Up @@ -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 = []
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions tests/test_RasterLayer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand Down

0 comments on commit bab4758

Please sign in to comment.