bump to Python 3.9, simplify API

Makefile

@@ -3,6 +3,10 @@ lint:
 	ruff format .
 	ruff check .
 
+.PHONY: check
+check:
+	pyright
+
 .PHONY: test
 test:
 	pytest

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)"
    ]
   },
   {

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

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)