fixes NDVI with target_band

openeo_processes_dask/process_implementations/cubes/indices.py

@@ -1,3 +1,4 @@
+import numpy as np
 import xarray as xr
 
 from openeo_processes_dask.process_implementations.data_model import RasterCube
@@ -12,6 +13,47 @@
 __all__ = ["ndvi"]
 
 
+def _dummy_value(dtype):
+    """A sensible 'missing' fill value for a given numpy dtype."""
+    if np.issubdtype(dtype, np.floating):
+        return np.nan
+    if np.issubdtype(dtype, np.datetime64):
+        return np.datetime64("NaT")
+    if np.issubdtype(dtype, np.timedelta64):
+        return np.timedelta64("NaT")
+    if np.issubdtype(dtype, np.integer):
+        return -1  # ints can't hold NaN
+    if np.issubdtype(dtype, np.bool_):
+        return False
+    return ""  # str / object
+
+
+def _add_missing_coords(target, reference, dim):
+    """
+    Ensure `target` has every coordinate that `reference` has.
+    Missing coordinates are created with a dummy value
+    """
+    missing = {}
+    for name, coord in reference.coords.items():
+        if name in target.coords:
+            continue
+
+        fill = _dummy_value(coord.dtype)
+
+        if coord.dims == ():  # scalar coordinate
+            missing[name] = ((), np.array(fill, dtype=coord.dtype))
+
+        elif coord.dims == (dim,):  # coordinate along concat dim
+            missing[name] = (dim, np.full(target.sizes[dim], fill, dtype=coord.dtype))
+
+        elif all(d in target.dims for d in coord.dims):  # other dims
+            shape = tuple(target.sizes[d] for d in coord.dims)
+            missing[name] = (coord.dims, np.full(shape, fill, dtype=coord.dtype))
+        # else: dims don't exist in target -> can't sensibly create it, skip
+
+    return target.assign_coords(**missing) if missing else target
+
+
 def ndvi(data: RasterCube, nir="nir", red="red", target_band=None):
     if len(data.openeo.band_dims) == 0:
         raise DimensionAmbiguous(
@@ -54,6 +96,10 @@ def ndvi(data: RasterCube, nir="nir", red="red", target_band=None):
         if target_band in data.coords:
             raise BandExists("A band with the specified target name exists.")
         nd = nd.expand_dims(band_dim).assign_coords({band_dim: [target_band]})
-        nd = xr.merge([data, nd])
+
+        # add potentially missing coords from data to nd, so that xr.concat works
+        nd = _add_missing_coords(nd, data, band_dim)
+        nd = xr.concat([data, nd], dim=band_dim)
+
     nd.attrs = data.attrs
     return nd

tests/test_indices.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+import xarray as xr
 
 from openeo_processes_dask.process_implementations.cubes.indices import ndvi
 from openeo_processes_dask.process_implementations.exceptions import (
@@ -73,11 +74,14 @@ def test_ndvi(temporal_interval, bounding_box, random_raster_data, process_regis
 
     target_band = "yay"
     output_with_extra_dim = ndvi(input_cube, target_band=target_band)
+    assert isinstance(output_with_extra_dim, xr.DataArray)
     assert len(output_with_extra_dim.dims) == len(output.dims) + 1
     assert (
         len(output_with_extra_dim.coords[band_dim])
         == len(input_cube.coords[band_dim]) + 1
     )
+    assert output_with_extra_dim.coords[band_dim].values[-1] == target_band
+    assert output_with_extra_dim.coords["common_name"][-1] == ""
 
     with pytest.raises(BandExists):
         output_with_extra_dim = ndvi(input_cube, target_band="t")