diff --git a/xrspatial/terrain.py b/xrspatial/terrain.py index 27c984d06..234762c93 100644 --- a/xrspatial/terrain.py +++ b/xrspatial/terrain.py @@ -594,6 +594,10 @@ def generate_terrain(agg: xr.DataArray, ---------- agg : xr.DataArray 2D template array (determines height, width, and backend). + If it carries its own ``y``/``x`` coordinates and attrs (e.g. + ``crs``, ``res``, ``units``), they are preserved on the output and + ``x_range``/``y_range`` only drive the noise field. A bare template + with no coordinates gets a synthetic grid from ``x_range``/``y_range``. x_range : tuple, default=(0, 500) Range of x values. y_range : tuple, default=(0, 500) @@ -695,16 +699,39 @@ def generate_terrain(agg: xr.DataArray, octaves, lacunarity, persistence, noise_mode, warp_strength, warp_octaves, worley_blend, worley_seed) - # Build pixel-center coordinates (matches datashader Canvas convention) - dx = (x_range[1] - x_range[0]) / width - dy = (y_range[1] - y_range[0]) / height - xs = np.linspace(x_range[0] + dx / 2, x_range[1] - dx / 2, width) - ys = np.linspace(y_range[0] + dy / 2, y_range[1] - dy / 2, height) + # When the caller's array carries its own georeference (coords / attrs + # such as crs, res, units), preserve it so generate_terrain fills the + # caller's grid instead of replacing it with a synthetic (x_range, + # y_range) one. A bare template with no coords keeps the old behaviour + # of synthesizing pixel-center coordinates from x_range / y_range. + if 'x' in agg.coords and 'y' in agg.coords: + coords = {'y': agg.coords['y'], 'x': agg.coords['x']} + else: + # pixel-center coordinates (matches datashader Canvas convention) + dx = (x_range[1] - x_range[0]) / width + dy = (y_range[1] - y_range[0]) / height + xs = np.linspace(x_range[0] + dx / 2, x_range[1] - dx / 2, width) + ys = np.linspace(y_range[0] + dy / 2, y_range[1] - dy / 2, height) + coords = {'y': ys, 'x': xs} + + # Carry the caller's attrs (crs, units, ...) verbatim. Only synthesize a + # res when the caller did not provide one, deriving it from the output + # coord spacing so it stays consistent with the coordinates above. + attrs = dict(agg.attrs) + if 'res' not in attrs: + xs_arr = np.asarray(coords['x']) + ys_arr = np.asarray(coords['y']) + res_x = (float(xs_arr[1] - xs_arr[0]) if xs_arr.size > 1 + else float(x_range[1] - x_range[0])) + res_y = (float(ys_arr[1] - ys_arr[0]) if ys_arr.size > 1 + else float(y_range[1] - y_range[0])) + attrs['res'] = (res_x, res_y) + result = xr.DataArray(out, name=name, - coords={'y': ys, 'x': xs}, + coords=coords, dims=['y', 'x'], - attrs={'res': (dx, dy)}) + attrs=attrs) # --- hydraulic erosion --- if erode: diff --git a/xrspatial/tests/test_terrain.py b/xrspatial/tests/test_terrain.py index facb18156..74d99d710 100644 --- a/xrspatial/tests/test_terrain.py +++ b/xrspatial/tests/test_terrain.py @@ -366,3 +366,107 @@ def test_persistence_rejected(self, per): from xrspatial.terrain import generate_terrain with pytest.raises(ValueError, match="persistence"): generate_terrain(self._template(), persistence=per) + + +# ===================================================================== +# Issue #3474: preserve the caller's coords / chunks / res / crs +# ===================================================================== + +def _georef_arr(H=40, W=30, res=30, backend='numpy'): + """A georeferenced template: real y/x coords plus crs / res / units.""" + ys = np.arange(H) * res + xs = np.arange(W) * res + data = np.zeros((H, W), dtype=np.float32) + raster = xr.DataArray( + data, coords={'y': ys, 'x': xs}, dims=('y', 'x'), + name='study_area', + attrs={'res': (res, res), 'crs': 'EPSG:5070', 'units': 'meters'}, + ) + if has_cuda_and_cupy() and 'cupy' in backend: + import cupy + raster.data = cupy.asarray(raster.data) + if 'dask' in backend and da is not None: + raster.data = da.from_array(raster.data, chunks=(13, 17)) + return raster + + +def test_generate_terrain_preserves_caller_coords_and_attrs(): + src = _georef_arr() + terrain = generate_terrain(src) + + np.testing.assert_array_equal(terrain.y.data, src.y.data) + np.testing.assert_array_equal(terrain.x.data, src.x.data) + assert terrain.attrs['res'] == (30, 30) + assert terrain.attrs['crs'] == 'EPSG:5070' + assert terrain.attrs['units'] == 'meters' + + +def test_generate_terrain_accessor_preserves_georeference(): + src = _georef_arr() + terrain = src.xrs.generate_terrain() + + np.testing.assert_array_equal(terrain.y.data, src.y.data) + np.testing.assert_array_equal(terrain.x.data, src.x.data) + assert terrain.attrs['res'] == (30, 30) + assert terrain.attrs['crs'] == 'EPSG:5070' + + +def test_generate_terrain_bare_template_keeps_synthetic_grid(): + """No coords / attrs on the input -> old (x_range, y_range) behaviour.""" + bare = create_test_arr() # 50x50, dims only, no coords / attrs + terrain = generate_terrain(bare, x_range=(0, 500), y_range=(0, 500)) + + assert 'crs' not in terrain.attrs + dx = 500 / 50 + np.testing.assert_allclose(terrain.x.data[0], dx / 2) + np.testing.assert_allclose(terrain.attrs['res'], (dx, dx)) + + +def test_generate_terrain_coords_without_res_attr(): + """Coords but no res attr -> res derived from coord spacing.""" + ys = np.arange(40) * 25.0 + xs = np.arange(30) * 25.0 + src = xr.DataArray(np.zeros((40, 30), dtype=np.float32), + coords={'y': ys, 'x': xs}, dims=('y', 'x')) + terrain = generate_terrain(src) + + np.testing.assert_array_equal(terrain.x.data, xs) + np.testing.assert_allclose(terrain.attrs['res'], (25.0, 25.0)) + + +@dask_array_available +def test_generate_terrain_dask_preserves_chunks_coords_attrs(): + src = _georef_arr(backend='dask') + terrain = generate_terrain(src) + + assert isinstance(terrain.data, da.Array) + assert terrain.data.chunks == src.data.chunks + np.testing.assert_array_equal(terrain.y.data, src.y.data) + np.testing.assert_array_equal(terrain.x.data, src.x.data) + assert terrain.attrs['crs'] == 'EPSG:5070' + assert terrain.attrs['res'] == (30, 30) + + +@cuda_and_cupy_available +def test_generate_terrain_cupy_preserves_coords_attrs(): + src = _georef_arr(backend='cupy') + terrain = generate_terrain(src) + + np.testing.assert_array_equal(terrain.y.data, src.y.data) + np.testing.assert_array_equal(terrain.x.data, src.x.data) + assert terrain.attrs['crs'] == 'EPSG:5070' + assert terrain.attrs['res'] == (30, 30) + + +@cuda_and_cupy_available +@dask_array_available +def test_generate_terrain_dask_cupy_preserves_coords_attrs(): + src = _georef_arr(backend='dask+cupy') + terrain = generate_terrain(src) + + assert isinstance(terrain.data, da.Array) + assert terrain.data.chunks == src.data.chunks + np.testing.assert_array_equal(terrain.y.data, src.y.data) + np.testing.assert_array_equal(terrain.x.data, src.x.data) + assert terrain.attrs['crs'] == 'EPSG:5070' + assert terrain.attrs['res'] == (30, 30)