From d3d11a10979ec804a16869abeb9c4ce6c9d13954 Mon Sep 17 00:00:00 2001 From: Chris Arderne Date: Sun, 18 Feb 2024 10:24:20 +0000 Subject: [PATCH] bump to Python 3.9, simplify API --- Makefile | 4 ++ examples/example.ipynb | 15 ++++--- gridfinder/gridfinder.py | 51 ++++------------------ gridfinder/post.py | 76 ++++++++++++++------------------- gridfinder/prepare.py | 91 ++++++++++++++++++---------------------- gridfinder/util.py | 44 ++++++++----------- pyproject.toml | 12 ++++-- tests/test_gridfinder.py | 4 +- tests/test_post.py | 6 +-- tests/test_prepare.py | 21 ++++++---- 10 files changed, 135 insertions(+), 189 deletions(-) diff --git a/Makefile b/Makefile index 85a1925..1466a4f 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,10 @@ lint: ruff format . ruff check . +.PHONY: check +check: + pyright + .PHONY: test test: pytest diff --git a/examples/example.ipynb b/examples/example.ipynb index 58c8142..3f8e98c 100644 --- a/examples/example.ipynb +++ b/examples/example.ipynb @@ -23,7 +23,7 @@ "import matplotlib.animation as animation\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "import rasterio\n", + "import rasterio as rs\n", "import seaborn as sns\n", "from IPython.display import Markdown, display\n", "from matplotlib import cm\n", @@ -96,7 +96,8 @@ "metadata": {}, "outputs": [], "source": [ - "targets_clean = gf.drop_zero_pop(targets_out, pop_in, aoi_in)\n", + "aoi = gpd.read_file(aoi_in)\n", + "targets_clean = gf.drop_zero_pop(targets_out, pop_in, aoi)\n", "save_raster(targets_clean_out, targets_clean, affine)\n", "print(\"Removed zero pop\")\n", "plt.imshow(ntl_thresh, cmap=\"viridis\")" @@ -115,7 +116,11 @@ "metadata": {}, "outputs": [], "source": [ - "roads_raster, affine = gf.prepare_roads(roads_in, aoi_in, targets_out)\n", + "roads = gpd.read_file(roads_in)\n", + "with rs.open(targets_out) as ds:\n", + " shape = ds.shape\n", + " affine = ds.transform\n", + "roads_raster = gf.prepare_roads(roads, aoi, shape, affine)\n", "save_raster(costs_out, roads_raster, affine, nodata=-1)\n", "print(\"Costs prepared\")\n", "plt.imshow(roads_raster, cmap=\"viridis\", vmin=0, vmax=1)" @@ -134,9 +139,7 @@ "metadata": {}, "outputs": [], "source": [ - "targets, costs, start, affine = gf.get_targets_costs(targets_clean_out, costs_out)\n", - "est_mem = gf.estimate_mem_use(targets, costs)\n", - "print(f\"Estimated memory usage: {est_mem:.2f} GB\")" + "targets, costs, start, affine = gf.get_targets_costs(targets_clean_out, costs_out)" ] }, { diff --git a/gridfinder/gridfinder.py b/gridfinder/gridfinder.py index 1be32f4..c6f14cf 100644 --- a/gridfinder/gridfinder.py +++ b/gridfinder/gridfinder.py @@ -2,14 +2,12 @@ Implements Dijkstra's algorithm on a cost-array to create an MST. """ -import pickle from heapq import heapify, heappop, heappush from math import sqrt -from typing import List, Optional, Tuple import numba as nb import numpy as np -import rasterio +import rasterio as rs from affine import Affine from gridfinder.util import Loc, Pathy @@ -18,7 +16,7 @@ def get_targets_costs( targets_in: Pathy, costs_in: Pathy, -) -> Tuple[np.ndarray, np.ndarray, Loc, Affine]: +) -> tuple[np.ndarray, np.ndarray, Loc, Affine]: """Load the targets and costs arrays from the given file paths. Parameters @@ -34,12 +32,12 @@ def get_targets_costs( affine : Affine transformation for the rasters. """ - targets_ra = rasterio.open(targets_in) - affine = targets_ra.transform - targets = targets_ra.read(1) + with rs.open(targets_in) as ds: + affine = ds.transform + targets = ds.read(1) - costs_ra = rasterio.open(costs_in) - costs = costs_ra.read(1) + with rs.open(costs_in) as ds: + costs = ds.read(1) target_list = np.argwhere(targets == 1.0) start = tuple(target_list[0].tolist()) @@ -50,40 +48,12 @@ def get_targets_costs( return targets, costs, start, affine -def estimate_mem_use(targets: np.ndarray, costs: np.ndarray) -> float: - """Estimate memory usage in GB, probably not very accurate. - - Parameters - ---------- - targets : 2D array of targets. - costs : 2D array of costs. - - Returns - ------- - est_mem : Estimated memory requirement in GB. - """ - - # make sure these match the ones used in optimise below - visited = np.zeros_like(targets, dtype=np.int8) - dist = np.full_like(costs, np.nan, dtype=np.float32) - prev = np.full_like(costs, np.nan, dtype=object) - - est_mem_arr = [targets, costs, visited, dist, prev] - est_mem = len(pickle.dumps(est_mem_arr, -1)) - - return est_mem / 1e9 - - @nb.njit def optimise( targets: np.ndarray, costs: np.ndarray, start: Loc, silent: bool = False, - jupyter: bool = False, - animate: bool = False, - affine: Optional[Affine] = None, - animate_path: str = "", ) -> np.ndarray: """Run the Dijkstra algorithm for the supplied arrays. @@ -101,11 +71,6 @@ def optimise( on-grid point. Values of 0 imply that cell is part of an MV grid line. """ - if jupyter or animate or affine or animate_path: - print( - "Warning: the following parameters are ignored: jupyter, animate, affine, animate_path" # NoQA - ) - shape = costs.shape max_i = shape[0] max_j = shape[1] @@ -116,7 +81,7 @@ def optimise( dist[start] = 0 - queue: List[Tuple[float, Loc]] = [(0.0, start)] + queue: list[tuple[float, Loc]] = [(0.0, start)] heapify(queue) counter = 0 diff --git a/gridfinder/post.py b/gridfinder/post.py index 30f3ef3..122fec7 100644 --- a/gridfinder/post.py +++ b/gridfinder/post.py @@ -2,12 +2,11 @@ Post-processing for gridfinder package. """ -from typing import Optional, Tuple, Union import geopandas as gpd import numpy as np import pandas as pd -import rasterio +import rasterio as rs import shapely.wkt from affine import Affine from rasterio.features import rasterize @@ -18,9 +17,7 @@ from gridfinder.util import Pathy -def threshold( - dists_in: Union[Pathy, np.ndarray], cutoff: float = 0.0 -) -> Tuple[np.ndarray, Optional[Affine]]: +def threshold(dists_in: Pathy, cutoff: float = 0.0) -> tuple[np.ndarray, Affine]: """Convert distance array into binary array of connected locations. Parameters @@ -33,26 +30,18 @@ def threshold( guess : Binary representation of input array. affine: Affine transformation for raster. """ - if isinstance(dists_in, np.ndarray): - guess = dists_in.copy() - guess[dists_in > cutoff] = 0 - guess[dists_in <= cutoff] = 1 + with rs.open(dists_in) as ds: + dists_r = ds.read(1) + affine = ds.transform - return guess, None + guess = dists_r.copy() + guess[dists_r > cutoff] = 0 + guess[dists_r <= cutoff] = 1 - else: - dists_rd = rasterio.open(dists_in) - dists_r = dists_rd.read(1) - affine = dists_rd.transform - - guess = dists_r.copy() - guess[dists_r > cutoff] = 0 - guess[dists_r <= cutoff] = 1 - - return guess, affine + return guess, affine -def thin(guess_in: Union[Pathy, np.ndarray]) -> Tuple[np.ndarray, Optional[Affine]]: +def thin(guess_in: Pathy) -> tuple[np.ndarray, Affine]: """ Use scikit-image skeletonize to 'thin' the guess raster. @@ -63,22 +52,16 @@ def thin(guess_in: Union[Pathy, np.ndarray]) -> Tuple[np.ndarray, Optional[Affin Returns ------- guess_skel : Thinned version. - affine : Only if path-like supplied. + affine : affine """ + with rs.open(guess_in) as ds: + guess_arr = ds.read(1) + affine = ds.transform - if isinstance(guess_in, np.ndarray): - guess_skel = skeletonize(guess_in) - guess_skel = guess_skel.astype("int32") - return guess_skel, None - else: - guess_rd = rasterio.open(guess_in) - guess_arr = guess_rd.read(1) - affine = guess_rd.transform - - guess_skel = skeletonize(guess_arr) - guess_skel = guess_skel.astype("int32") + guess_skel = skeletonize(guess_arr) + guess_skel = guess_skel.astype("int32") - return guess_skel, affine + return guess_skel, affine def raster_to_lines(guess_skel_in: Pathy) -> gpd.GeoDataFrame: @@ -94,9 +77,10 @@ def raster_to_lines(guess_skel_in: Pathy) -> gpd.GeoDataFrame: guess_gdf : Converted to geometry. """ - rast = rasterio.open(guess_skel_in) - arr = rast.read(1) - affine = rast.transform + with rs.open(guess_skel_in) as ds: + arr = ds.read(1) + rast_crs = ds.crs + affine = ds.transform max_row = arr.shape[0] max_col = arr.shape[1] @@ -141,7 +125,7 @@ def raster_to_lines(guess_skel_in: Pathy) -> gpd.GeoDataFrame: guess_gdf = pd.DataFrame(shapes) geometry = guess_gdf[0].map(shapely.wkt.loads) guess_gdf = guess_gdf.drop(0, axis=1) - guess_gdf = gpd.GeoDataFrame(guess_gdf, crs=rast.crs, geometry=geometry) + guess_gdf = gpd.GeoDataFrame(guess_gdf, crs=rast_crs, geometry=geometry) guess_gdf["same"] = 0 guess_gdf = guess_gdf.dissolve(by="same") @@ -155,7 +139,7 @@ def accuracy( guess_in: Pathy, aoi_in: Pathy, buffer_amount: float = 0.01, -) -> Tuple[float, float]: +) -> tuple[float, float]: """Measure accuracy against a specified grid 'truth' file. Parameters @@ -175,25 +159,27 @@ def accuracy( grid = gpd.read_file(grid_in, mask=aoi) grid_buff = grid.buffer(buffer_amount) - guesses_reader = rasterio.open(guess_in) - guesses = guesses_reader.read(1) + with rs.open(guess_in) as ds: + guesses = ds.read(1) + out_shape = ds.shape + affine = ds.transform grid_for_raster = [(row.geometry) for _, row in grid.iterrows()] grid_raster = rasterize( grid_for_raster, - out_shape=guesses_reader.shape, + out_shape=out_shape, fill=1, default_value=0, all_touched=True, - transform=guesses_reader.transform, + transform=affine, ) grid_buff_raster = rasterize( grid_buff, - out_shape=guesses_reader.shape, + out_shape=out_shape, fill=1, default_value=0, all_touched=True, - transform=guesses_reader.transform, + transform=affine, ) grid_raster = flip_arr_values(grid_raster) diff --git a/gridfinder/prepare.py b/gridfinder/prepare.py index 1085026..c3007ca 100644 --- a/gridfinder/prepare.py +++ b/gridfinder/prepare.py @@ -6,12 +6,12 @@ import os from math import sqrt from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Optional import fiona import geopandas as gpd import numpy as np -import rasterio +import rasterio as rs from geopandas.geodataframe import GeoDataFrame from rasterio import Affine from rasterio.features import rasterize @@ -26,7 +26,7 @@ def clip_rasters( folder_in: Pathy, folder_out: Pathy, - aoi_in: Pathy, + aoi: GeoDataFrame, debug: bool = False, ) -> None: """Read continental rasters one at a time, clip to AOI and save @@ -38,30 +38,28 @@ def clip_rasters( aoi_in : Path to an AOI file (readable by Fiona) to use for clipping. """ - if isinstance(aoi_in, gpd.GeoDataFrame): - aoi = aoi_in - else: - aoi = gpd.read_file(aoi_in) + folder_in = Path(folder_in) + folder_out = Path(folder_out) coords = [json.loads(aoi.to_json())["features"][0]["geometry"]] - for file_path in os.listdir(folder_in): - if file_path.endswith(".tif"): + for file_path in folder_in.iterdir(): + if file_path.suffix == ".tif": if debug: print(f"Doing {file_path}") - ntl_rd = rasterio.open(os.path.join(folder_in, file_path)) - ntl, affine = mask(dataset=ntl_rd, shapes=coords, crop=True, nodata=0) + with rs.open(file_path) as ds: + ntl, affine = mask(dataset=ds, shapes=coords, crop=True, nodata=0) if ntl.ndim == 3: ntl = ntl[0] - save_raster(Path(folder_out) / file_path, ntl, affine) + save_raster(folder_out / file_path.name, ntl, affine) def merge_rasters( folder: Pathy, percentile: int = 70, -) -> Tuple[np.ndarray, Optional[Affine]]: +) -> tuple[np.ndarray, Affine]: """Merge a set of monthly rasters keeping the nth percentile value. Used to remove transient features from time-series data. @@ -83,7 +81,7 @@ def merge_rasters( for file in os.listdir(folder): if file.endswith(".tif"): - ntl_rd = rasterio.open(os.path.join(folder, file)) + ntl_rd = rs.open(os.path.join(folder, file)) rasters.append(ntl_rd.read(1)) if not affine: @@ -93,6 +91,8 @@ def merge_rasters( raster_merged = np.percentile(raster_arr, percentile, axis=0) + if not isinstance(affine, Affine): + raise Exception("No affine found") return raster_merged, affine @@ -125,7 +125,7 @@ def prepare_ntl( ntl_filter: Optional[np.ndarray] = None, threshold: float = 0.1, upsample_by: int = 2, -) -> Tuple[np.ndarray, Affine]: +) -> tuple[np.ndarray, Affine]: """Convert the supplied NTL raster and output an array of electrified cells as targets for the algorithm. @@ -154,7 +154,7 @@ def prepare_ntl( if ntl_filter is None: ntl_filter = create_filter() - ntl_big = rasterio.open(ntl_in) + ntl_big = rs.open(ntl_in) coords = [json.loads(aoi.to_json())["features"][0]["geometry"]] ntl, affine = mask(dataset=ntl_big, shapes=coords, crop=True, nodata=0) @@ -183,7 +183,7 @@ def prepare_ntl( affine.f, ) with fiona.Env(): - with rasterio.Env(): + with rs.Env(): reproject( ntl_filtered, ntl_interp, @@ -207,7 +207,7 @@ def prepare_ntl( def drop_zero_pop( targets_in: Pathy, pop_in: Pathy, - aoi: Union[Pathy, GeoDataFrame], + aoi: GeoDataFrame, ) -> np.ndarray: """Drop electrified cells with no other evidence of human activity. @@ -222,21 +222,19 @@ def drop_zero_pop( Array with zero population sites dropped. """ - aoi_gdf = gpd.read_file(aoi) if isinstance(aoi, (str, Path)) else aoi - # Clip population layer to AOI - clipped, affine, crs = clip_raster(pop_in, aoi_gdf) + clipped, affine, crs = clip_raster(pop_in, aoi) # We need to warp the population layer to exactly overlap with targets # First get array, affine and crs from targets (which is what we) - targets_rd = rasterio.open(targets_in) - targets = targets_rd.read(1) + with rs.open(targets_in) as ds: + targets = ds.read(1) + dest_affine = ds.transform + dest_crs = ds.crs ghs_proj = np.empty_like(targets) - dest_affine = targets_rd.transform - dest_crs = targets_rd.crs # Then use reproject - with rasterio.Env(): + with rs.Env(): reproject( source=clipped, destination=ghs_proj, @@ -296,38 +294,27 @@ def add_around(blob: list, cell: Loc) -> list: def prepare_roads( - roads_in: Pathy, - aoi_in: Union[Pathy, GeoDataFrame], - ntl_in: Pathy, + roads: GeoDataFrame, + aoi: GeoDataFrame, + shape: tuple[int, int], + affine: Affine, nodata: float = 1, -) -> Tuple[np.ndarray, Affine]: +) -> np.ndarray: """Prepare a roads feature layer for use in algorithm. Parameters ---------- - roads_in : Path to a roads feature layer. This implementation is specific to + roads: Roads feature layer. This implementation is specific to OSM data and won't assign proper weights to other data inputs. - aoi_in : AOI to clip roads. - ntl_in : Path to a raster file, only used for correct shape and - affine of roads raster. + aoi: AOI to clip roads. + shape: shape of resultant raster + affine: affine of resultant raster Returns ------- roads_raster : Roads as a raster array with the value being the cost of traversing. - affine : Affine raster transformation for the new raster (same as ntl_in). """ - ntl_rd = rasterio.open(ntl_in) - shape = ntl_rd.read(1).shape - affine = ntl_rd.transform - - if isinstance(aoi_in, gpd.GeoDataFrame): - aoi = aoi_in - else: - aoi = gpd.read_file(aoi_in) - - roads = gpd.read_file(roads_in) - roads["weight"] = 1.0 roads.loc[roads["highway"] == "motorway", "weight"] = 1 / 10 roads.loc[roads["highway"] == "trunk", "weight"] = 1 / 9 @@ -342,13 +329,15 @@ def prepare_roads( if "power" in roads: roads.loc[roads["power"] == "line", "weight"] = 0 - roads = roads[roads.weight != 1] + roads = roads.loc[roads.weight != 1] # sort by weight descending so that lower weight (bigger roads) are # processed last and overwrite higher weight roads - roads = roads.sort_values(by="weight", ascending=False) + roads_sorted = roads.sort_values(by="weight", ascending=False) - roads_for_raster = [(row.geometry, row.weight) for _, row in roads.iterrows()] + roads_for_raster = [ + (row.geometry, row.weight) for _, row in roads_sorted.iterrows() + ] rast = rasterize( roads_for_raster, out_shape=shape, @@ -360,7 +349,7 @@ def prepare_roads( # Clip resulting raster by the AOI if not aoi.crs == roads.crs: - aoi: GeoDataFrame = aoi.to_crs(crs=roads.crs) # type: ignore + aoi = aoi.to_crs(crs=roads.crs) # type: ignore coords = [json.loads(aoi.to_json())["features"][0]["geometry"]] with MemoryFile() as f: @@ -381,4 +370,4 @@ def prepare_roads( if len(clipped.shape) >= 3: clipped = clipped[0] - return clipped, affine + return clipped diff --git a/gridfinder/util.py b/gridfinder/util.py index 9fc26a7..197b872 100644 --- a/gridfinder/util.py +++ b/gridfinder/util.py @@ -4,20 +4,18 @@ import json from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Optional, Union -import geopandas as gpd import numpy as np -import rasterio +import rasterio as rs from affine import Affine from geopandas import GeoDataFrame from pyproj import Proj -from rasterio.io import DatasetReader from rasterio.mask import mask Pathy = Union[str, Path] -Loc = Tuple[int, int] +Loc = tuple[int, int] def save_raster( @@ -45,7 +43,7 @@ def save_raster( if not crs: crs = "+proj=latlong" - filtered_out = rasterio.open( + with rs.open( p_path, "w", driver="GTiff", @@ -56,16 +54,14 @@ def save_raster( crs=crs, transform=affine, nodata=nodata, - ) - filtered_out.write(raster, 1) - filtered_out.close() + ) as filtered_out: + filtered_out.write(raster, 1) def clip_raster( - raster: Union[Pathy, DatasetReader], - boundary: Union[Pathy, GeoDataFrame], - boundary_layer: Optional[str] = None, -) -> Tuple[ + raster: Pathy, + boundary: GeoDataFrame, +) -> tuple[ np.ndarray, Affine, dict, @@ -76,7 +72,6 @@ def clip_raster( ---------- raster : Location of or already opened raster. boundary : The polygon by which to clip the raster. - boundary_layer : For multi-layer files (like GeoPackage), specify layer to be used. Returns ------- @@ -85,19 +80,12 @@ def clip_raster( crs : form {'init': 'epsg:4326'} """ - raster_ds = raster if isinstance(raster, DatasetReader) else rasterio.open(raster) + raster_ds = rs.open(raster) + raster_crs = raster_ds.crs - boundary_gdf: GeoDataFrame = ( - boundary - if isinstance(boundary, GeoDataFrame) - else gpd.read_file(boundary, layer=boundary_layer) - ) - - if not ( - boundary_gdf.crs == raster_ds.crs or boundary_gdf.crs == raster_ds.crs.data - ): - boundary_gdf = boundary_gdf.to_crs(crs=raster_ds.crs) # type: ignore - coords = [json.loads(boundary_gdf.to_json())["features"][0]["geometry"]] + if not (boundary.crs == raster_crs or boundary.crs == raster_crs.data): + boundary = boundary.to_crs(crs=raster_crs) # type: ignore + coords = [json.loads(boundary.to_json())["features"][0]["geometry"]] # mask/clip the raster using rasterio.mask clipped, affine = mask(dataset=raster_ds, shapes=coords, crop=True) @@ -105,4 +93,6 @@ def clip_raster( if len(clipped.shape) >= 3: clipped = clipped[0] - return clipped, affine, raster_ds.crs + raster_ds.close() + + return clipped, affine, raster_crs diff --git a/pyproject.toml b/pyproject.toml index 945407a..10a09bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ name = "gridfinder" description = "Algorithm for guessing MV grid network based on night time lights" readme = "README.md" license = {text = "MIT License"} -requires-python = ">=3.8" +requires-python = ">=3.9" keywords = ["ntl", "electricity", "grid"] authors = [ @@ -25,6 +25,10 @@ classifiers = [ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", ] dependencies = [ @@ -69,7 +73,7 @@ exclude = ["docs*", "tests*"] [tool.setuptools_scm] [tool.ruff] -target-version = "py38" +target-version = "py39" line-length = 88 exclude = [] @@ -91,9 +95,9 @@ select = [ known-first-party = ["gridfinder"] [tool.pyright] -include = ["gridfinder"] +include = ["gridfinder", "tests"] reportMissingImports = true reportMissingParameterType = true reportUnnecessaryTypeIgnoreComment = true reportDeprecated = true -pythonVersion = "3.8" +pythonVersion = "3.9" diff --git a/tests/test_gridfinder.py b/tests/test_gridfinder.py index 3b30c3f..e3bd9d1 100644 --- a/tests/test_gridfinder.py +++ b/tests/test_gridfinder.py @@ -1,7 +1,7 @@ from pathlib import Path import numpy as np -import rasterio +import rasterio as rs import gridfinder as gf @@ -9,6 +9,6 @@ def test_optimise(p_targets_clean: Path, p_costs: Path, p_dist: Path) -> None: targets, costs, start, _ = gf.get_targets_costs(p_targets_clean, p_costs) res = gf.optimise(targets, costs, start) - with rasterio.open(p_dist) as ds: + with rs.open(p_dist) as ds: exp = ds.read(1) assert np.array_equal(res, exp) diff --git a/tests/test_post.py b/tests/test_post.py index ec3c04e..1a3eff1 100644 --- a/tests/test_post.py +++ b/tests/test_post.py @@ -1,14 +1,14 @@ from pathlib import Path import numpy as np -import rasterio +import rasterio as rs import gridfinder as gf def test_threshold(p_dist: Path, p_guess: Path) -> None: res, affine = gf.threshold(p_dist, cutoff=0.0) - with rasterio.open(p_guess) as ds: + with rs.open(p_guess) as ds: exp = ds.read(1) exp_affine = ds.transform assert np.array_equal(res, exp) @@ -17,7 +17,7 @@ def test_threshold(p_dist: Path, p_guess: Path) -> None: def test_thin(p_guess: Path, p_guess_thin: Path) -> None: res, affine = gf.thin(p_guess) - with rasterio.open(p_guess_thin) as ds: + with rs.open(p_guess_thin) as ds: exp = ds.read(1) exp_affine = ds.transform assert np.array_equal(res, exp) diff --git a/tests/test_prepare.py b/tests/test_prepare.py index 678f909..a0b5d7b 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -1,14 +1,15 @@ from pathlib import Path import numpy as np -import rasterio +import rasterio as rs +from geopandas import gpd import gridfinder as gf def test_prepare_ntl(p_ntl: Path, p_aoi: Path, p_targets: Path) -> None: res, res_affine = gf.prepare_ntl(p_ntl, p_aoi) - with rasterio.open(p_targets) as ds: + with rs.open(p_targets) as ds: exp = ds.read(1) exp_affine = ds.transform assert np.array_equal(res, exp) @@ -18,8 +19,9 @@ def test_prepare_ntl(p_ntl: Path, p_aoi: Path, p_targets: Path) -> None: def test_drop_zero_pop( p_targets: Path, p_pop: Path, p_aoi: Path, p_targets_clean: Path ) -> None: - res = gf.drop_zero_pop(p_targets, p_pop, p_aoi) - with rasterio.open(p_targets_clean) as ds: + aoi = gpd.read_file(p_aoi) + res = gf.drop_zero_pop(p_targets, p_pop, aoi) + with rs.open(p_targets_clean) as ds: exp = ds.read(1) assert np.array_equal(res, exp) @@ -30,9 +32,12 @@ def test_prepare_roads( p_targets_clean: Path, p_costs: Path, ) -> None: - res, res_affine = gf.prepare_roads(p_roads, p_aoi, p_targets_clean) - with rasterio.open(p_costs) as ds: + with rs.open(p_targets_clean) as ds: + shape = ds.shape + affine = ds.transform + aoi = gpd.read_file(p_aoi) + roads = gpd.read_file(p_roads) + res = gf.prepare_roads(roads, aoi, shape, affine) + with rs.open(p_costs) as ds: exp = ds.read(1) - exp_affine = ds.transform assert np.array_equal(res, exp) - assert res_affine == exp_affine