Preserve caller coords/res/crs in generate_terrain (#3474)

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:

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)