from __future__ import annotations
from functools import partial
from math import atan2
from typing import Optional
try:
import dask.array as da
except ImportError:
da = None
import numpy as np
import xarray as xr
from numba import cuda
from xrspatial.dataset_support import supports_dataset
from xrspatial.geodesic import (INV_2R, WGS84_A2, WGS84_B2,
_check_geodesic_memory_backend_aware,
_cpu_geodesic_aspect, _run_gpu_geodesic_aspect)
from xrspatial.utils import (Z_UNITS, ArrayTypeFunctionMapping, _boundary_to_dask,
_extract_latlon_coords, _pad_array, _validate_boundary,
_validate_raster, cuda_args, get_dataarray_resolution,
ngjit, warn_if_unit_mismatch)
def _geodesic_cuda_dims(shape):
"""Smaller thread block for register-heavy geodesic kernels."""
tpb = (16, 16)
bpg = (
(shape[0] + tpb[0] - 1) // tpb[0],
(shape[1] + tpb[1] - 1) // tpb[1],
)
return bpg, tpb
# 3rd-party
try:
import cupy
except ImportError:
class cupy(object):
ndarray = False
RADIAN = 180 / np.pi
# =====================================================================
# Planar backend functions (unchanged)
# =====================================================================
@ngjit
def _cpu(data: np.ndarray, cellsize_x, cellsize_y):
data = data.astype(np.float32)
out = np.empty(data.shape, dtype=np.float32)
out[:] = np.nan
rows, cols = data.shape
for y in range(1, rows-1):
for x in range(1, cols-1):
a = data[y-1, x-1]
b = data[y-1, x]
c = data[y-1, x+1]
d = data[y, x-1]
f = data[y, x+1]
g = data[y+1, x-1]
h = data[y+1, x]
i = data[y+1, x+1]
dz_dx = ((c + 2 * f + i) - (a + 2 * d + g)) / (8 * cellsize_x)
dz_dy = ((g + 2 * h + i) - (a + 2 * b + c)) / (8 * cellsize_y)
if dz_dx == 0 and dz_dy == 0:
# flat surface, slope = 0, thus invalid aspect
out[y, x] = -1.
else:
_aspect = np.arctan2(dz_dy, -dz_dx) * RADIAN
# convert to compass direction values (0-360 degrees)
if _aspect < 0:
out[y, x] = 90.0 - _aspect
elif _aspect > 90.0:
out[y, x] = 360.0 - _aspect + 90.0
else:
out[y, x] = 90.0 - _aspect
return out
def _run_numpy(data: np.ndarray,
cellsize_x,
cellsize_y,
boundary: str = 'nan') -> np.ndarray:
if boundary == 'nan':
return _cpu(data, cellsize_x, cellsize_y)
padded = _pad_array(data, 1, boundary)
result = _cpu(padded, cellsize_x, cellsize_y)
return result[1:-1, 1:-1]
@cuda.jit(device=True)
def _gpu(arr, cellsize_x, cellsize_y):
a = arr[0, 0]
b = arr[0, 1]
c = arr[0, 2]
d = arr[1, 0]
f = arr[1, 2]
g = arr[2, 0]
h = arr[2, 1]
i = arr[2, 2]
dz_dx = ((c + 2 * f + i) - (a + 2 * d + g)) / (8 * cellsize_x[0])
dz_dy = ((g + 2 * h + i) - (a + 2 * b + c)) / (8 * cellsize_y[0])
if dz_dx == 0 and dz_dy == 0:
# flat surface, slope = 0, thus invalid aspect
_aspect = -1
else:
# Reuse the numpy kernel's RADIAN constant (180 / pi) so the branch
# below selects the same way on both backends; a coarser constant can
# push _aspect across the 90 boundary and yield 360 where numpy yields
# 0 (issue #2827).
_aspect = atan2(dz_dy, -dz_dx) * RADIAN
# convert to compass direction values (0-360 degrees)
if _aspect < 0:
_aspect = 90 - _aspect
elif _aspect > 90:
_aspect = 360 - _aspect + 90
else:
_aspect = 90 - _aspect
# Keep the output in [0, 360) to match the numpy kernel. The numpy
# kernel needs no equivalent wrap: its `elif _aspect > 90` excludes an
# exact-90 tie, which then yields 0 via the else branch. The GPU's
# 450 - 90 = 360 case is folded back to 0 here so the two agree.
if _aspect >= 360.0:
_aspect -= 360.0
return _aspect
@cuda.jit
def _run_gpu(arr, cellsize_x_arr, cellsize_y_arr, out):
i, j = cuda.grid(2)
di = 1
dj = 1
if (i-di >= 0 and
i+di < out.shape[0] and
j-dj >= 0 and
j+dj < out.shape[1]):
out[i, j] = _gpu(arr[i-di:i+di+1, j-dj:j+dj+1],
cellsize_x_arr,
cellsize_y_arr)
def _run_cupy(data: cupy.ndarray,
cellsize_x,
cellsize_y,
boundary: str = 'nan') -> cupy.ndarray:
if boundary != 'nan':
padded = _pad_array(data, 1, boundary)
result = _run_cupy(padded, cellsize_x, cellsize_y)
return result[1:-1, 1:-1]
cellsize_x_arr = cupy.array([float(cellsize_x)], dtype='f4')
cellsize_y_arr = cupy.array([float(cellsize_y)], dtype='f4')
data = data.astype(cupy.float32)
griddim, blockdim = cuda_args(data.shape)
out = cupy.empty(data.shape, dtype='f4')
out[:] = cupy.nan
_run_gpu[griddim, blockdim](data, cellsize_x_arr, cellsize_y_arr, out)
return out
def _run_dask_numpy(data: da.Array,
cellsize_x,
cellsize_y,
boundary: str = 'nan') -> da.Array:
data = data.astype(np.float32)
_func = partial(_cpu, cellsize_x=cellsize_x, cellsize_y=cellsize_y)
out = data.map_overlap(_func,
depth=(1, 1),
boundary=_boundary_to_dask(boundary),
meta=np.array((), dtype=np.float32))
return out
def _run_dask_cupy(data: da.Array,
cellsize_x,
cellsize_y,
boundary: str = 'nan') -> da.Array:
data = data.astype(cupy.float32)
_func = partial(_run_cupy, cellsize_x=cellsize_x, cellsize_y=cellsize_y)
out = data.map_overlap(_func,
depth=(1, 1),
boundary=_boundary_to_dask(boundary, is_cupy=True),
meta=cupy.array((), dtype=cupy.float32))
return out
# =====================================================================
# Geodesic backend functions
# =====================================================================
def _run_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
if boundary != 'nan':
data_p = _pad_array(data.astype(np.float64), 1, boundary)
lat_p = _pad_array(lat_2d, 1, boundary)
lon_p = _pad_array(lon_2d, 1, boundary)
stacked = np.stack([data_p, lat_p, lon_p], axis=0)
result = _cpu_geodesic_aspect(stacked, a2, b2, z_factor)
return result[1:-1, 1:-1]
stacked = np.stack([
data.astype(np.float64),
lat_2d,
lon_2d,
], axis=0)
return _cpu_geodesic_aspect(stacked, a2, b2, z_factor)
def _run_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
if boundary != 'nan':
data_p = _pad_array(data.astype(cupy.float64), 1, boundary)
lat_p = _pad_array(cupy.asarray(lat_2d, dtype=cupy.float64), 1, boundary)
lon_p = _pad_array(cupy.asarray(lon_2d, dtype=cupy.float64), 1, boundary)
stacked = cupy.stack([data_p, lat_p, lon_p], axis=0)
H, W = data_p.shape
out = cupy.full((H, W), cupy.nan, dtype=cupy.float32)
a2_arr = cupy.array([a2], dtype=cupy.float64)
b2_arr = cupy.array([b2], dtype=cupy.float64)
zf_arr = cupy.array([z_factor], dtype=cupy.float64)
inv_2r_arr = cupy.array([INV_2R], dtype=cupy.float64)
griddim, blockdim = _geodesic_cuda_dims((H, W))
_run_gpu_geodesic_aspect[griddim, blockdim](
stacked, a2_arr, b2_arr, zf_arr, inv_2r_arr, out)
return out[1:-1, 1:-1]
lat_2d_gpu = cupy.asarray(lat_2d, dtype=cupy.float64)
lon_2d_gpu = cupy.asarray(lon_2d, dtype=cupy.float64)
stacked = cupy.stack([
data.astype(cupy.float64),
lat_2d_gpu,
lon_2d_gpu,
], axis=0)
H, W = data.shape
out = cupy.full((H, W), cupy.nan, dtype=cupy.float32)
a2_arr = cupy.array([a2], dtype=cupy.float64)
b2_arr = cupy.array([b2], dtype=cupy.float64)
zf_arr = cupy.array([z_factor], dtype=cupy.float64)
inv_2r_arr = cupy.array([INV_2R], dtype=cupy.float64)
griddim, blockdim = _geodesic_cuda_dims((H, W))
_run_gpu_geodesic_aspect[griddim, blockdim](stacked, a2_arr, b2_arr, zf_arr, inv_2r_arr, out)
return out
def _dask_geodesic_aspect_chunk(stacked_chunk, a2, b2, z_factor):
"""Returns (3, h, w) to preserve shape for map_overlap."""
result_2d = _cpu_geodesic_aspect(stacked_chunk, a2, b2, z_factor)
out = np.empty_like(stacked_chunk, dtype=np.float32)
out[0] = result_2d
out[1] = 0.0
out[2] = 0.0
return out
def _dask_geodesic_aspect_chunk_cupy(stacked_chunk, a2, b2, z_factor):
H, W = stacked_chunk.shape[1], stacked_chunk.shape[2]
result_2d = cupy.full((H, W), cupy.nan, dtype=cupy.float32)
a2_arr = cupy.array([a2], dtype=cupy.float64)
b2_arr = cupy.array([b2], dtype=cupy.float64)
zf_arr = cupy.array([z_factor], dtype=cupy.float64)
inv_2r_arr = cupy.array([INV_2R], dtype=cupy.float64)
griddim, blockdim = _geodesic_cuda_dims((H, W))
_run_gpu_geodesic_aspect[griddim, blockdim](
stacked_chunk, a2_arr, b2_arr, zf_arr, inv_2r_arr, result_2d)
out = cupy.zeros_like(stacked_chunk, dtype=cupy.float32)
out[0] = result_2d
return out
def _run_dask_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
lat_dask = da.from_array(lat_2d, chunks=data.chunksize)
lon_dask = da.from_array(lon_2d, chunks=data.chunksize)
stacked = da.stack([
data.astype(np.float64),
lat_dask,
lon_dask,
], axis=0).rechunk({0: 3})
_func = partial(_dask_geodesic_aspect_chunk, a2=a2, b2=b2, z_factor=z_factor)
dask_bnd = _boundary_to_dask(boundary)
out = stacked.map_overlap(
_func,
depth=(0, 1, 1),
boundary={0: np.nan, 1: dask_bnd, 2: dask_bnd},
meta=np.array((), dtype=np.float32),
)
return out[0]
def _to_cupy_f64(block):
# Only reached from the dask+cupy path, so `cupy` is the real module here,
# never the import-time fallback class.
return cupy.asarray(block, dtype=cupy.float64)
def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
# Keep lat/lon as dask-of-numpy on the (zero-stride) broadcast views, then
# convert each block to cupy lazily. Converting up front with
# cupy.asarray(lat_2d) would densify the full (H, W) grid onto a single GPU
# at graph-construction time and OOM on large rasters.
lat_dask = da.from_array(lat_2d, chunks=data.chunksize).map_blocks(
_to_cupy_f64, dtype=np.float64)
lon_dask = da.from_array(lon_2d, chunks=data.chunksize).map_blocks(
_to_cupy_f64, dtype=np.float64)
stacked = da.stack([
data.astype(cupy.float64),
lat_dask,
lon_dask,
], axis=0).rechunk({0: 3})
_func = partial(_dask_geodesic_aspect_chunk_cupy, a2=a2, b2=b2, z_factor=z_factor)
dask_bnd = _boundary_to_dask(boundary, is_cupy=True)
out = stacked.map_overlap(
_func,
depth=(0, 1, 1),
boundary={0: cupy.nan, 1: dask_bnd, 2: dask_bnd},
meta=cupy.array((), dtype=cupy.float32),
)
return out[0]
# =====================================================================
# Public API
# =====================================================================
[docs]
@supports_dataset
def aspect(agg: xr.DataArray,
name: Optional[str] = 'aspect',
method: str = 'planar',
z_unit: str = 'meter',
boundary: str = 'nan') -> xr.DataArray:
"""
Calculates the aspect value of an elevation aggregate.
Calculates, for all cells in the array, the downward slope direction
of each cell based on the elevation of its neighbors in a 3x3 grid.
The value is measured clockwise in degrees with 0 (due north), and 360
(again due north). Values along the edges are not calculated.
Direction of the aspect can be determined by its value:
From 0 to 22.5: North
From 22.5 to 67.5: Northeast
From 67.5 to 112.5: East
From 112.5 to 157.5: Southeast
From 157.5 to 202.5: South
From 202.5 to 247.5: Southwest
From 247.5 to 292.5: West
From 292.5 to 337.5: Northwest
From 337.5 to 360: North
Note that values of -1 denote flat areas.
Parameters
----------
agg : xarray.DataArray or xr.Dataset
2D NumPy, CuPy, or Dask with NumPy-backed xarray DataArray
of elevation values.
If a Dataset is passed, the operation is applied to each
data variable independently.
name : str, default='aspect'
Name of ouput DataArray.
method : str, default='planar'
``'planar'`` uses the classic Horn algorithm, scaling the gradients
by the x and y cell sizes so non-square cells are handled correctly.
``'geodesic'`` converts cells to Earth-Centered Earth-Fixed (ECEF)
coordinates and fits a 3D plane, yielding accurate results for
geographic (lat/lon) coordinate systems.
z_unit : str, default='meter'
Unit of the elevation values. Only used when ``method='geodesic'``.
Accepted values: ``'meter'``, ``'foot'``, ``'kilometer'``, ``'mile'``
(and common aliases).
boundary : str, default='nan'
How to handle edges where the kernel extends beyond the raster.
``'nan'`` — fill missing neighbours with NaN (default).
``'nearest'`` — repeat edge values.
``'reflect'`` — mirror at boundary.
``'wrap'`` — periodic / toroidal.
Returns
-------
aspect_agg : xarray.DataArray or xr.Dataset
If `agg` is a DataArray, returns a DataArray of the same type.
If `agg` is a Dataset, returns a Dataset with aspect computed
for each data variable.
2D aggregate array of calculated aspect values.
All other input attributes are preserved.
Notes
-----
The ``'planar'`` method uses the coordinate spacing directly as the cell
size. If the coordinates are in degrees (lat/lon) but the elevation values
are in meters, the result is wrong by orders of magnitude. When this
mismatch is detected, a ``UserWarning`` is emitted suggesting you reproject
to a projected CRS or use ``method='geodesic'``.
References
----------
Examples
--------
Aspect works with NumPy backed xarray DataArray
.. sourcecode:: python
>>> import numpy as np
>>> import xarray as xr
>>> from xrspatial import aspect
>>> data = np.array([
[1, 1, 1, 1, 1],
[1, 1, 1, 2, 0],
[1, 1, 1, 0, 0],
[4, 4, 9, 2, 4],
[1, 5, 0, 1, 4],
[1, 5, 0, 5, 5]
], dtype=np.float32)
>>> raster = xr.DataArray(data, dims=['y', 'x'], name='raster')
>>> aspect_agg = aspect(raster)
"""
_validate_raster(agg, func_name='aspect', name='agg')
if method not in ('planar', 'geodesic'):
raise ValueError(
f"method must be 'planar' or 'geodesic', got {method!r}"
)
_validate_boundary(boundary)
if method == 'planar':
warn_if_unit_mismatch(agg)
cellsize_x, cellsize_y = get_dataarray_resolution(agg)
mapper = ArrayTypeFunctionMapping(
numpy_func=_run_numpy,
dask_func=_run_dask_numpy,
cupy_func=_run_cupy,
dask_cupy_func=_run_dask_cupy,
)
out = mapper(agg)(agg.data, cellsize_x, cellsize_y, boundary=boundary)
else: # geodesic
if z_unit not in Z_UNITS:
raise ValueError(
f"z_unit must be one of {sorted(Z_UNITS)}, "
f"got {z_unit!r}"
)
z_factor = Z_UNITS[z_unit]
_check_geodesic_memory_backend_aware(agg, func_name='aspect')
lat_2d, lon_2d = _extract_latlon_coords(agg)
mapper = ArrayTypeFunctionMapping(
numpy_func=_run_numpy_geodesic,
cupy_func=_run_cupy_geodesic,
dask_func=_run_dask_numpy_geodesic,
dask_cupy_func=_run_dask_cupy_geodesic,
)
out = mapper(agg)(agg.data, lat_2d, lon_2d, WGS84_A2, WGS84_B2, z_factor, boundary)
result = xr.DataArray(out,
name=name,
coords=agg.coords,
dims=agg.dims,
attrs=agg.attrs)
# On dask backends, xr.DataArray keeps the dask array's internal graph
# token as .name when name=None, so reset it post-construction to match
# the numpy/cupy backends. (Same fix as zonal #2611, focal #2733,
# slope #2838.)
result.name = name
return result
[docs]
@supports_dataset
def northness(agg: xr.DataArray,
name: Optional[str] = 'northness',
method: str = 'planar',
z_unit: str = 'meter',
boundary: str = 'nan') -> xr.DataArray:
"""
Computes the north-south component of aspect.
Returns ``cos(aspect)`` for each cell, ranging from +1 (due north)
to -1 (due south). Flat cells (where ``aspect()`` returns -1) are
set to NaN.
This is the standard way to encode aspect for use in regression,
clustering, and other models that assume linear inputs. Raw aspect
in degrees is circular (0 and 360 are the same direction), so
feeding it directly into a linear model gives wrong results.
Parameters
----------
agg : xarray.DataArray or xr.Dataset
2D elevation raster (NumPy, CuPy, Dask, or Dask+CuPy backed).
If a Dataset is passed, the operation is applied to each
data variable independently.
name : str, default='northness'
Name of output DataArray.
method : str, default='planar'
Passed to :func:`aspect`. ``'planar'`` or ``'geodesic'``.
z_unit : str, default='meter'
Passed to :func:`aspect`. Only used when ``method='geodesic'``.
boundary : str, default='nan'
Passed to :func:`aspect`. ``'nan'``, ``'nearest'``,
``'reflect'``, or ``'wrap'``.
Returns
-------
northness_agg : xarray.DataArray or xr.Dataset
Values in [-1, +1]. NaN where the input has NaN or where
the surface is flat.
References
----------
Stage, A.R. (1976). "An Expression for the Effect of Aspect, Slope,
and Habitat Type on Tree Growth." *Forest Science* 22(4): 457-460.
Examples
--------
.. sourcecode:: python
>>> import numpy as np
>>> import xarray as xr
>>> from xrspatial import northness
>>> data = np.array([
[1, 1, 1, 1, 1],
[1, 1, 1, 2, 0],
[1, 1, 1, 0, 0],
[4, 4, 9, 2, 4],
[1, 5, 0, 1, 4],
[1, 5, 0, 5, 5]
], dtype=np.float32)
>>> raster = xr.DataArray(data, dims=['y', 'x'])
>>> north = northness(raster)
"""
asp = aspect(agg, name='_aspect', method=method, z_unit=z_unit,
boundary=boundary)
asp_data = asp.data
if da is not None and isinstance(asp_data, da.Array):
trig = da.cos(da.deg2rad(asp_data))
out = da.where(asp_data == -1, np.nan, trig)
else:
trig = np.cos(np.deg2rad(asp_data))
out = np.where(asp_data == -1, np.nan, trig)
result = xr.DataArray(out,
name=name,
coords=agg.coords,
dims=agg.dims,
attrs=agg.attrs)
# Reset .name post-construction so dask backends don't leak the graph
# token when name=None, matching aspect()/slope() (#2841, #2838).
result.name = name
return result
[docs]
@supports_dataset
def eastness(agg: xr.DataArray,
name: Optional[str] = 'eastness',
method: str = 'planar',
z_unit: str = 'meter',
boundary: str = 'nan') -> xr.DataArray:
"""
Computes the east-west component of aspect.
Returns ``sin(aspect)`` for each cell, ranging from +1 (due east)
to -1 (due west). Flat cells (where ``aspect()`` returns -1) are
set to NaN.
This is the standard way to encode aspect for use in regression,
clustering, and other models that assume linear inputs. Raw aspect
in degrees is circular (0 and 360 are the same direction), so
feeding it directly into a linear model gives wrong results.
Parameters
----------
agg : xarray.DataArray or xr.Dataset
2D elevation raster (NumPy, CuPy, Dask, or Dask+CuPy backed).
If a Dataset is passed, the operation is applied to each
data variable independently.
name : str, default='eastness'
Name of output DataArray.
method : str, default='planar'
Passed to :func:`aspect`. ``'planar'`` or ``'geodesic'``.
z_unit : str, default='meter'
Passed to :func:`aspect`. Only used when ``method='geodesic'``.
boundary : str, default='nan'
Passed to :func:`aspect`. ``'nan'``, ``'nearest'``,
``'reflect'``, or ``'wrap'``.
Returns
-------
eastness_agg : xarray.DataArray or xr.Dataset
Values in [-1, +1]. NaN where the input has NaN or where
the surface is flat.
References
----------
Stage, A.R. (1976). "An Expression for the Effect of Aspect, Slope,
and Habitat Type on Tree Growth." *Forest Science* 22(4): 457-460.
Examples
--------
.. sourcecode:: python
>>> import numpy as np
>>> import xarray as xr
>>> from xrspatial import eastness
>>> data = np.array([
[1, 1, 1, 1, 1],
[1, 1, 1, 2, 0],
[1, 1, 1, 0, 0],
[4, 4, 9, 2, 4],
[1, 5, 0, 1, 4],
[1, 5, 0, 5, 5]
], dtype=np.float32)
>>> raster = xr.DataArray(data, dims=['y', 'x'])
>>> east = eastness(raster)
"""
asp = aspect(agg, name='_aspect', method=method, z_unit=z_unit,
boundary=boundary)
asp_data = asp.data
if da is not None and isinstance(asp_data, da.Array):
trig = da.sin(da.deg2rad(asp_data))
out = da.where(asp_data == -1, np.nan, trig)
else:
trig = np.sin(np.deg2rad(asp_data))
out = np.where(asp_data == -1, np.nan, trig)
result = xr.DataArray(out,
name=name,
coords=agg.coords,
dims=agg.dims,
attrs=agg.attrs)
# Reset .name post-construction so dask backends don't leak the graph
# token when name=None, matching aspect()/slope() (#2841, #2838).
result.name = name
return result
