|
1 | 1 | import numpy as np |
| 2 | +import xarray as xr |
2 | 3 |
|
3 | 4 | from xrspatial.surface_distance import ( |
4 | 5 | surface_distance, surface_allocation, surface_direction, |
5 | 6 | ) |
| 7 | +from xrspatial.utils import has_cuda_and_cupy |
6 | 8 |
|
7 | 9 | from .common import get_xr_dataarray |
8 | 10 |
|
9 | 11 |
|
| 12 | +def _sparse_source_raster(ny, nx, type): |
| 13 | + """Source raster with a handful of scattered target pixels. |
| 14 | +
|
| 15 | + ``get_xr_dataarray(is_int=True)`` draws integers over ``[-nx, nx)``, |
| 16 | + and surface_distance treats every non-zero finite pixel as a source, |
| 17 | + so all but roughly 1 in ``2 * nx`` pixels seed the search at distance |
| 18 | + zero. The Dijkstra relaxation body then never runs and the benchmark |
| 19 | + times a heap drain instead of distance propagation. Scattered point |
| 20 | + sources make the frontier cross the whole grid, which is what these |
| 21 | + functions are for. |
| 22 | + """ |
| 23 | + rng = np.random.default_rng(71942) |
| 24 | + z = np.zeros((ny, nx), dtype=np.float32) |
| 25 | + n_sources = max(4, (ny * nx) // 20000) |
| 26 | + rows = rng.integers(0, ny, n_sources) |
| 27 | + cols = rng.integers(0, nx, n_sources) |
| 28 | + z[rows, cols] = np.arange(1, n_sources + 1, dtype=np.float32) |
| 29 | + |
| 30 | + chunks = (max(1, ny // 2), max(1, nx // 2)) |
| 31 | + if type == "cupy": |
| 32 | + if not has_cuda_and_cupy(): |
| 33 | + raise NotImplementedError() |
| 34 | + import cupy |
| 35 | + z = cupy.asarray(z) |
| 36 | + elif type == "dask": |
| 37 | + import dask.array as da |
| 38 | + z = da.from_array(z, chunks=chunks) |
| 39 | + elif type == "dask+cupy": |
| 40 | + if not has_cuda_and_cupy(): |
| 41 | + raise NotImplementedError() |
| 42 | + import cupy |
| 43 | + import dask.array as da |
| 44 | + z = da.from_array(cupy.asarray(z), chunks=chunks) |
| 45 | + elif type != "numpy": |
| 46 | + raise RuntimeError(f"Unrecognised type {type}") |
| 47 | + |
| 48 | + y = np.linspace(-90, 90, ny) |
| 49 | + x = np.linspace(-180, 180, nx) |
| 50 | + return xr.DataArray(z, coords=dict(y=y, x=x), dims=["y", "x"]) |
| 51 | + |
| 52 | + |
| 53 | +def _compute(result): |
| 54 | + if hasattr(result.data, "compute"): |
| 55 | + result.data.compute() |
| 56 | + |
| 57 | + |
10 | 58 | class SurfaceDistance: |
11 | | - params = ([100, 300, 1000], ["numpy", "dask"]) |
| 59 | + params = ([100, 300, 1000], ["numpy", "cupy", "dask", "dask+cupy"]) |
12 | 60 | param_names = ("nx", "type") |
13 | 61 |
|
14 | 62 | def setup(self, nx, type): |
15 | 63 | ny = nx // 2 |
16 | | - self.agg = get_xr_dataarray((ny, nx), type, is_int=True) |
| 64 | + self.agg = _sparse_source_raster(ny, nx, type) |
17 | 65 | self.elev = get_xr_dataarray((ny, nx), type) |
18 | 66 |
|
| 67 | + # A finite max_distance whose pixel radius stays inside one chunk |
| 68 | + # (chunks are ny//2 x nx//2) routes the dask backends through the |
| 69 | + # bounded map_overlap branch instead of the iterative tile one. |
| 70 | + cellsize = min(360.0 / (nx - 1), 180.0 / (ny - 1)) |
| 71 | + self.max_distance = 20 * cellsize |
| 72 | + |
19 | 73 | def time_surface_distance(self, nx, type): |
20 | | - surface_distance(self.agg, self.elev) |
| 74 | + _compute(surface_distance(self.agg, self.elev)) |
| 75 | + |
| 76 | + def time_surface_distance_bounded(self, nx, type): |
| 77 | + _compute(surface_distance( |
| 78 | + self.agg, self.elev, max_distance=self.max_distance)) |
21 | 79 |
|
22 | 80 | def time_surface_allocation(self, nx, type): |
23 | | - surface_allocation(self.agg, self.elev) |
| 81 | + _compute(surface_allocation(self.agg, self.elev)) |
24 | 82 |
|
25 | 83 | def time_surface_direction(self, nx, type): |
26 | | - surface_direction(self.agg, self.elev) |
| 84 | + _compute(surface_direction(self.agg, self.elev)) |
| 85 | + |
| 86 | + |
| 87 | +class SurfaceDistanceGeodesic: |
| 88 | + """Great-circle horizontal distances from lat/lon coordinates. |
| 89 | +
|
| 90 | + Runs a separate numba kernel (``_dijkstra_geodesic``) behind a |
| 91 | + precomputed per-pixel neighbour-distance grid, and costs about twice |
| 92 | + the planar path. numpy only: the module raises NotImplementedError |
| 93 | + for geodesic on cupy, dask, and dask+cupy. |
| 94 | + """ |
| 95 | + |
| 96 | + params = [300, 1000] |
| 97 | + param_names = ("nx",) |
| 98 | + |
| 99 | + def setup(self, nx): |
| 100 | + ny = nx // 2 |
| 101 | + self.agg = _sparse_source_raster(ny, nx, "numpy") |
| 102 | + self.elev = get_xr_dataarray((ny, nx), "numpy") |
| 103 | + |
| 104 | + def time_surface_distance_geodesic(self, nx): |
| 105 | + surface_distance(self.agg, self.elev, method="geodesic") |
0 commit comments