"""Utilities for geospatial data and raster grids."""
import heapq
import math
from typing import Literal
import numpy as np
from affine import Affine
from numba import njit
_R = 6371e3 # Radius of earth in m. Use 3956e3 for miles
AREA_FACTORS = {"m2": 1.0, "ha": 1e4, "km2": 1e6, "cell": 1}
# changed to N->S orientation in v0.5 TODO check if used in hydromt?
_IDENTITY: np.ndarray = np.array([1.0, 0.0, 0.0, 0.0, -1.0, 0.0])
IDENTITY = Affine(*_IDENTITY) # Affine transformation for identity
__all__ = [
"affine_to_coords",
"array_bounds",
"get_edge",
"idxs_to_coords",
"reggrid_area",
"reggrid_dx",
"reggrid_dy",
"rowcol",
"spread2d",
"transform_from_bounds",
"transform_from_origin",
"xy",
]
[docs]
@njit(cache=True)
def spread2d(
obs: np.ndarray,
msk: np.ndarray | None = None,
nodata: float = 0,
frc: np.ndarray | None = None,
latlon: bool = False,
transform: np.ndarray = _IDENTITY,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Fill no-data cells with the nearest observation and return source and distance maps.
Distances are accumulated through valid cells. The default friction is 1 per cell;
diagonal steps use the hypotenuse of the horizontal and vertical distances.
Parameters
----------
obs : 2D array
Input observations. Cells equal to `nodata` are candidates for filling.
msk : 2D array of bool, optional
Valid-cell mask. Observations and fill paths are restricted to True cells.
nodata : int or float, optional
Missing-data value in `obs`, by default 0.
frc : 2D array of float, optional
Per-cell friction multiplier for distance accumulation, by default 1 everywhere.
latlon : bool, optional
Whether coordinates are geographic. If True, transform units are interpreted
as degrees and distances are converted to metres, by default False.
transform : np.ndarray, optional
Six affine transform coefficients mapping pixel coordinates to map coordinates.
Returns
-------
out: 2D array of obs.dtype
Copy of `obs` with fillable no-data cells assigned the nearest observation.
src: 2D array of int32
Linear index of the nearest observation, or -1 where no observation is reachable.
dst: 2D array of float32
Accumulated friction distance to the nearest observation.
"""
nrow, ncol = obs.shape
xres, yres, north = transform[0], abs(transform[4]), transform[5]
if latlon:
lats = north + (np.arange(nrow) + 0.5) * yres
dys = degree_metres_y(lats) * yres
dxs = degree_metres_x(lats) * xres
else:
dx, dy = xres, yres
# output
out = obs.copy()
src = np.full(obs.shape, -1, dtype=np.int32) # linear index of source
dst = np.full(obs.shape, 0, dtype=np.float32) # distance from source
# initiate queue with correct dtype
# heapq is faster when fifo loop not in order of ascending distance from source;
# otherwise a fixed length numpy array queue is up to ~2x faster
q = [(np.float32(0), np.uint32(0), np.uint32(0)) for _ in range(0)]
heapq.heapify(q)
for r in range(nrow):
for c in range(ncol):
if obs[r, c] != nodata:
if msk is None or msk[r, c]:
heapq.heappush(q, (np.float32(0), np.uint32(r), np.uint32(c)))
src[r, c] = r * ncol + c
obs = obs.ravel()
while len(q) > 0:
d0, r, c = heapq.heappop(q) # type: ignore[assignment]
if dst[r, c] < d0:
continue
f0 = 1.0 if frc is None else frc[r, c]
if latlon:
dx, dy = dxs[r], dys[r]
for dr in range(-1, 2):
for dc in range(-1, 2):
if dr == 0 and dc == 0:
continue
r1, c1 = int(r) + dr, int(c) + dc
outside = r1 < 0 or r1 >= nrow or c1 < 0 or c1 >= ncol
if outside or (msk is not None and not msk[r1, c1]):
continue
d = d0 + np.hypot(dr * dy, dc * dx) * f0
if src[r1, c1] == -1 or d < dst[r1, c1]:
idx0 = src[r, c]
src[r1, c1] = idx0
dst[r1, c1] = d
out[r1, c1] = obs[idx0]
heapq.heappush(q, (np.float32(d), np.uint32(r1), np.uint32(c1)))
return out, src, dst
[docs]
def get_edge(a: np.ndarray, structure: np.ndarray | None = None) -> np.ndarray:
"""Get edge of valid cells.
Parameters
----------
a: 2D array of bool
Boolean array valid cells.
structure: 2D array with shape (3,3) of bool, optional
Structuring element used to define which cells are neighbors.
If None, a 3x3 square is used. The center cell is ignored.
Returns
-------
edge: 2D array of bool
Boolean array edge cells.
"""
if structure is None:
struct = np.ones((3, 3), dtype=np.bool_)
elif (
not isinstance(structure, np.ndarray)
or structure.shape != (3, 3)
or structure.dtype != bool
):
raise ValueError("structure must be a 3x3 boolean array")
else:
struct = structure
return _get_edge(a, struct)
@njit(cache=True)
def _get_edge(a: np.ndarray, struct: np.ndarray) -> np.ndarray:
s = np.where(struct.ravel())[0]
edge = a.copy()
nrow, ncol = a.shape
for r in range(nrow):
for c in range(ncol):
if not a[r, c] or r == 0 or r == nrow - 1 or c == 0 or c == ncol - 1:
continue
a0 = a[slice(r - 1, r + 2), slice(c - 1, c + 2)].ravel()
if np.all(a0[s]):
edge[r, c] = False
return edge
## TRANSFORM
# Adapted from https://github.com/rasterio/rasterio/blob/main/rasterio/transform.py
# changed xy and rowcol to work directly on numpy arrays
# avoid gdal dependency
def transform_from_origin(
west: float, north: float, xsize: float, ysize: float
) -> Affine:
"""Return an affine transform from the upper-left corner and pixel sizes.
Parameters
----------
west, north : float
Coordinates of the upper-left corner.
xsize, ysize : float
Pixel width and height in coordinate units.
Returns
-------
Affine
Transform mapping pixel coordinates to map coordinates.
"""
return Affine(xsize, 0.0, west, 0.0, -ysize, north)
def transform_from_bounds(
west: float,
south: float,
east: float,
north: float,
width: int,
height: int,
) -> Affine:
"""Return an affine transform from raster bounds and dimensions.
Parameters
----------
west, south, east, north : float
Raster bounds in coordinate units.
width, height : int
Raster dimensions in pixels.
Returns
-------
Affine
Transform mapping pixel coordinates to map coordinates.
"""
return Affine(
(east - west) / width,
0.0,
west,
0.0,
(south - north) / height,
north,
)
def array_bounds(
height: int, width: int, transform: Affine
) -> tuple[float, float, float, float]:
"""Return the west, south, east, and north bounds of an array.
Parameters
----------
height, width : int
Array dimensions in pixels.
transform : Affine
Transform mapping pixel coordinates to map coordinates.
Returns
-------
tuple of float
Bounds in `(west, south, east, north)` order.
"""
w, n = transform.xoff, transform.yoff
e = transform.a * width + transform.b * height + transform.c
s = transform.d * width + transform.e * height + transform.f
return w, s, e, n
def xy(
transform: Affine,
rows: np.ndarray | int,
cols: np.ndarray | int,
offset: Literal["center", "ul", "ur", "ll", "lr"] = "center",
) -> tuple[np.ndarray, np.ndarray]:
"""Return the x and y coordinates of pixels at `rows` and `cols`.
The pixel's center is returned by default, but a corner can be returned
by setting `offset` to one of `ul, ur, ll, lr`.
Parameters
----------
transform : Affine
Transform mapping pixel coordinates to map coordinates.
rows : ndarray or int
Pixel rows.
cols : ndarray or int
Pixel columns.
offset : {'center', 'ul', 'ur', 'll', 'lr'}
Determines if the returned coordinates are for the center of the
pixel or for a corner.
Returns
-------
xs : ndarray of float
x coordinates in coordinate reference system
ys : ndarray of float
y coordinates in coordinate reference system
"""
rows, cols = np.asarray(rows), np.asarray(cols)
if offset == "center":
coff, roff = (0.5, 0.5)
elif offset == "ul":
coff, roff = (0, 0)
elif offset == "ur":
coff, roff = (1, 0)
elif offset == "ll":
coff, roff = (0, 1)
elif offset == "lr":
coff, roff = (1, 1)
else:
raise ValueError("Invalid offset")
xs = transform.a * (cols + coff) + transform.b * (rows + roff) + transform.c
ys = transform.d * (cols + coff) + transform.e * (rows + roff) + transform.f
return xs, ys
def rowcol(
transform: Affine,
xs: np.ndarray | float,
ys: np.ndarray | float,
op=np.floor,
precision: int | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""
Returns the rows and cols of the pixels containing (x, y) given a
coordinate reference system.
Use an epsilon, magnitude determined by the precision parameter
and sign determined by the op function:
positive for floor, negative for ceil.
Parameters
----------
transform: Affine
Coefficients mapping pixel coordinates to coordinate reference system.
xs : ndarray or float
x values in coordinate reference system
ys : ndarray or float
y values in coordinate reference system
op : function {numpy.floor, numpy.ceil, numpy.round}
Function to convert fractional pixels to whole numbers
precision : int, optional
Decimal places of precision in indexing, as in `round()`.
Returns
-------
rows : ndarray of ints
array of row indices
cols : ndarray of ints
array of column indices
"""
xs, ys = np.asarray(xs), np.asarray(ys)
if precision is None:
eps = 0.0
else:
eps = 10.0**-precision * (1.0 - 2.0 * op(0.1))
invtransform = ~transform
fcols = invtransform.a * (xs + eps) + invtransform.b * (ys - eps) + invtransform.c
frows = invtransform.d * (xs + eps) + invtransform.e * (ys - eps) + invtransform.f
cols, rows = op(fcols).astype(int), op(frows).astype(int)
return rows, cols
def idxs_to_coords(
idxs: np.ndarray,
transform: Affine,
shape: tuple[int, int],
offset: Literal["center", "ul", "ur", "ll", "lr"] = "center",
) -> tuple[np.ndarray, np.ndarray]:
"""Return map coordinates for linear raster indices.
Parameters
----------
idxs : ndarray of int
linear indices
transform : Affine
Transform mapping pixel coordinates to map coordinates.
shape : tuple of int
Raster dimensions as `(height, width)`.
offset : {'center', 'ul', 'ur', 'll', 'lr'}
Determines if the returned coordinates are for the center of the
pixel or for a corner.
Returns
-------
xs : ndarray of float
x coordinates in coordinate reference system
ys : ndarray of float
y coordinates in coordinate reference system
Raises
------
IndexError
if any linear index outside domain.
"""
idxs = np.asarray(idxs).astype(int)
size = np.multiply(*shape)
if np.any(np.logical_or(idxs < 0, idxs >= size)):
raise IndexError("idxs coordinates outside domain")
ncol = shape[1]
rows = idxs // ncol
cols = idxs % ncol
return xy(transform, rows, cols, offset=offset)
def coords_to_idxs(
xs: np.ndarray | float,
ys: np.ndarray | float,
transform: Affine,
shape: tuple[int, int],
op=np.floor,
precision: int | None = None,
) -> np.ndarray:
"""Return linear raster indices for map coordinates.
Parameters
----------
xs : ndarray or float
x values in coordinate reference system
ys : ndarray or float
y values in coordinate reference system
transform : Affine
Transform mapping pixel coordinates to map coordinates.
shape : tuple of int
Raster dimensions as `(height, width)`.
op : function {numpy.floor, numpy.ceil, numpy.round}
Function to convert fractional pixels to whole numbers
precision : int, optional
Decimal places of precision in indexing, as in `round()`.
Returns
-------
idxs : ndarray of ints
array of linear indices
Raises
------
IndexError
if any coordinate outside domain.
"""
nrow, ncol = shape
rows, cols = rowcol(transform, xs, ys, op=op, precision=precision)
if not np.all(
np.logical_and(
np.logical_and(rows >= 0, rows < nrow),
np.logical_and(cols >= 0, cols < ncol),
)
):
raise IndexError("XY coordinates outside domain")
return rows * ncol + cols
# TODO: rename to transform_to_coords & correct upstream use in pyflwdir and hydromt
def affine_to_coords(
affine: Affine, shape: tuple[int, int]
) -> tuple[np.ndarray, np.ndarray]:
"""Return the x and y pixel-center coordinate arrays for an affine transform.
Parameters
----------
affine : Affine
Transform mapping pixel coordinates to map coordinates.
shape : tuple of int
Raster dimensions as `(height, width)`.
Returns
-------
tuple of 1D arrays
The x coordinates for columns and y coordinates for rows.
Returns
-------
x, y coordinate arrays : tuple of ndarray of float
"""
height, width = shape
x_coords = affine.a * (np.arange(width) + 0.5) + affine.b * 0.5 + affine.c
y_coords = affine.d * 0.5 + affine.e * (np.arange(height) + 0.5) + affine.f
return x_coords, y_coords
## DISTANCES // AREAS
def reggrid_dx(lats: np.ndarray, lons: np.ndarray) -> np.ndarray:
"""Return cell widths [m] for a regular geographic grid.
Parameters
----------
lats, lons : 1D arrays of float
Latitude and longitude coordinates at cell centers, in degrees.
Returns
-------
2D array of float
Cell widths, with shape `(len(lats), len(lons))`.
"""
xres = np.abs(np.mean(np.diff(lons)))
dx = degree_metres_x(lats) * xres
return dx[:, None] * np.ones((lats.size, lons.size), dtype=lats.dtype)
def reggrid_dy(lats: np.ndarray, lons: np.ndarray) -> np.ndarray:
"""Return cell heights [m] for a regular geographic grid.
Parameters
----------
lats, lons : 1D arrays of float
Latitude and longitude coordinates at cell centers, in degrees.
Returns
-------
2D array of float
Cell heights, with shape `(len(lats), len(lons))`.
"""
yres = np.abs(np.mean(np.diff(lats)))
dy = degree_metres_y(lats) * yres
return dy[:, None] * np.ones((lats.size, lons.size), dtype=lats.dtype)
[docs]
def reggrid_area(lats: np.ndarray, lons: np.ndarray) -> np.ndarray:
"""Return cell areas [m2] for a regular geographic grid.
Parameters
----------
lats, lons : 1D arrays of float
Latitude and longitude coordinates at cell centers, in degrees.
Returns
-------
2D array of float
Cell areas, with shape `(len(lats), len(lons))`.
"""
xres = np.abs(np.mean(np.diff(lons)))
yres = np.abs(np.mean(np.diff(lats)))
area = np.ones((lats.size, lons.size), dtype=np.float32)
return cellarea(lats, xres, yres)[:, None] * area
def area_grid(
transform: Affine,
shape: tuple[int, int],
latlon: bool = False,
unit: str = "m2",
) -> np.ndarray:
"""Return a raster of cell areas.
Parameters
----------
transform : Affine
Affine transform for the raster.
shape : tuple of int
Raster dimensions `(height, width)`.
latlon : bool, optional
Whether coordinates are geographic, by default False.
unit : {'m2', 'ha', 'km2', 'cell'}, optional
Area units, by default 'm2'.
Returns
-------
2D array
Cell areas in `unit`.
"""
unit = str(unit).lower()
if unit not in AREA_FACTORS:
fstr = '", "'.join(AREA_FACTORS.keys())
raise ValueError(f'Unknown unit: {unit}, select from "{fstr}".')
area: np.ndarray
if unit == "cell":
area = np.ones(shape, dtype=np.int32)
elif latlon:
lon, lat = affine_to_coords(transform, shape)
area = reggrid_area(lat, lon) / AREA_FACTORS[unit]
elif not latlon:
area0 = abs(transform[0] * transform[4]) / AREA_FACTORS[unit]
area = np.full(shape, area0, dtype=np.float32)
return area
@njit(cache=True)
def cellarea(lat: np.ndarray | float, xres: float, yres: float) -> np.ndarray | float:
"""Return the area [m2] of a cell at a given center latitude.
Parameters
----------
lat : float or array-like
Cell-center latitude in degrees.
xres, yres : float
Cell width and height in degrees.
Returns
-------
float or array-like
Cell area in square metres.
"""
l1 = np.radians(lat - np.abs(yres) / 2.0)
l2 = np.radians(lat + np.abs(yres) / 2.0)
dx = np.radians(np.abs(xres))
return _R**2 * dx * (np.sin(l2) - np.sin(l1))
@njit(cache=True)
def degree_metres_y(lat: np.ndarray | float) -> np.ndarray | float:
"""Return the north-south length of one degree [m] at a given latitude.
Parameters
----------
lat : float or array-like
Latitude in degrees.
Returns
-------
float or array-like
Length of one degree of latitude in metres.
"""
m1 = 111132.92 # latitude calculation term 1
m2 = -559.82 # latitude calculation term 2
m3 = 1.175 # latitude calculation term 3
m4 = -0.0023 # latitude calculation term 4
# # Calculate the length of a degree of latitude and longitude in meters
radlat = np.radians(lat)
latlen = (
m1
+ (m2 * np.cos(2.0 * radlat))
+ (m3 * np.cos(4.0 * radlat))
+ (m4 * np.cos(6.0 * radlat))
)
return latlen
@njit(cache=True)
def degree_metres_x(lat: np.ndarray | float) -> np.ndarray | float:
"""Return the east-west length of one degree [m] at a given latitude.
Parameters
----------
lat : float or array-like
Latitude in degrees.
Returns
-------
float or array-like
Length of one degree of longitude in metres.
"""
p1 = 111412.84 # longitude calculation term 1
p2 = -93.5 # longitude calculation term 2
p3 = 0.118 # longitude calculation term 3
# # Calculate the length of a degree of latitude and longitude in meters
radlat = np.radians(lat)
longlen = (
(p1 * np.cos(radlat))
+ (p2 * np.cos(3.0 * radlat))
+ (p3 * np.cos(5.0 * radlat))
)
return longlen
@njit(cache=True)
def distance(
idx0: int,
idx1: int,
ncol: int,
latlon: bool = False,
transform: np.ndarray = _IDENTITY,
) -> float:
"""Return the distance between two linear raster indices.
The raster is assumed to have regular spacing defined by `transform`.
Parameters
----------
idx0, idx1 : int
index of start, end cell
ncol : int
number of columns in raster
latlon : bool, optional
True for geographic CRS, False for projected CRS.
If True, the transform units are assumed to be degrees and converted to metric distances.
transform : Affine, optional
Transform mapping pixel coordinates to map coordinates.
Returns
-------
float
Distance in map units, or metres for geographic coordinates.
"""
xres, yres, north = transform[0], transform[4], transform[5]
# compute delta row, col
r0 = int(idx0 // ncol)
r1 = int(idx1 // ncol)
dr = abs(r1 - r0)
dc = abs(int(idx1 % ncol) - int(idx0 % ncol))
if latlon: # calculate cell size in metres
lat = north + (r0 + r1) / 2.0 * yres
dy = 0.0 if dr == 0 else degree_metres_y(lat) * yres
dx = 0.0 if dc == 0 else degree_metres_x(lat) * xres
else:
dy = xres
dx = yres
return math.hypot(dy * dr, dx * dc) # length
## VECTORIZE
def features(
flowpaths: list[np.ndarray],
xs: np.ndarray | None = None,
ys: np.ndarray | None = None,
transform: Affine | None = None,
shape: tuple[int, int] | None = None,
**kwargs,
) -> list[dict]:
"""Return one LineString feature for each flow path.
Parameters
----------
flowpaths : list of 1D-arrays of intp
linear indices of flowpaths
xs, ys : 1D-array of float
x, y coordinates
transform : Affine, optional
Transform mapping pixel coordinates to map coordinates. Required when `xs` or
`ys` is omitted.
shape : tuple of int
The height, width of the raster.
**kwargs : 2D array-like
Additional maps sampled at each flow path's most downstream cell and included
as feature properties, for example `strord=flw.stream_order()`.
Returns
-------
feats : list of dict
Geofeatures, to be parsed by e.g. geopandas.GeoDataFrame.from_features
"""
if xs is None or ys is None:
if transform is None or shape is None:
raise ValueError(
"transform and shape should be provided if xs and ys are None"
)
_size = shape[0] * shape[1]
else:
_size = xs.size
for key, value in kwargs.items():
if not isinstance(value, np.ndarray) or value.size != _size:
raise ValueError(
f'Kwargs map "{key}" should be ndarrays of same size as coordinates'
)
feats = []
for j, idxs in enumerate(flowpaths):
n = len(idxs)
if n < 2:
continue
idx0 = idxs[0]
pit = idxs[-1] == idxs[-2]
props = {key: kwargs[key].flat[idx0] for key in kwargs}
if xs is None or ys is None:
xi, yi = idxs_to_coords(idxs, transform, shape) # type: ignore[arg-type]
coordinates = list(zip(xi, yi))
else:
coordinates = [(xs[i], ys[i]) for i in idxs]
feats.append(
{
"type": "Feature",
"geometry": {
"type": "LineString",
"coordinates": coordinates,
},
"properties": {"idx": idx0, "idx_ds": idxs[-1], "pit": pit, **props},
}
)
return feats
