fixes NDVI with target_band

openeo_processes_dask/process_implementations/cubes/indices.py

@@ -1,3 +1,6 @@
+from typing import Any
+
+import numpy as np
 import xarray as xr
 
 from openeo_processes_dask.process_implementations.data_model import RasterCube
@@ -12,7 +15,63 @@
 __all__ = ["ndvi"]
 
 
-def ndvi(data: RasterCube, nir="nir", red="red", target_band=None):
+def _dummy_value(dtype: np.dtype) -> Any:
+    """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: xr.DataArray, reference: xr.DataArray, dim: str
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """
+    Adds missing secondary coordinate to dimension dim in 'reference' to 'target'
+    Returns the target DataArray with added coordiantes, and the reference with
+    potentially altered coordinate data type
+    :param target: DataArray to which secondary coordinate is added
+    :param reference: DataArray to be checked for secondary coordiantes to dimension dim
+    :param dim: Name of dimension to check for secondary coordiante in dim
+    :return: DataArray `target` with potentially added coordinate dim,
+        DataArray `reference` with potentially altered coordinate dtype
+    """
+    missing = {}
+
+    # get all secondary coordinates to dim in reference DataArray
+    secondary_coord_to_dim = [
+        (name, coord)
+        for name, coord in reference.coords.items()
+        if name != dim and coord.dims == (dim,)
+    ]
+
+    for name, coord in secondary_coord_to_dim:
+        # special case: dim is "common_name"
+        if coord.dtype.kind == "U" and name == "common_name":
+            fill = "NDVI"
+
+            # convert coord dtype to U<4 to fit "NDVI", if it is less than 4 characters
+            if coord.dtype.kind == "U" and int(coord.dtype.str[2:]) < 4:
+                reference = reference.assign_coords(common_name=coord.astype("<U4"))
+
+        else:
+            fill = _dummy_value(coord.dtype)
+
+        missing[name] = (dim, np.array([fill]))
+
+    return target.assign_coords(**missing) if missing else target, reference
+
+
+def ndvi(
+    data: RasterCube, nir: str = "nir", red: str = "red", target_band: str | None = None
+) -> xr.DataArray:
     if len(data.openeo.band_dims) == 0:
         raise DimensionAmbiguous(
             "Dimension of type `bands` is not available or is ambiguous."
@@ -54,6 +113,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, data = _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,6 +1,8 @@
 import numpy as np
 import pytest
+import xarray as xr
 
+import openeo_processes_dask.process_implementations.cubes.indices as indices
 from openeo_processes_dask.process_implementations.cubes.indices import ndvi
 from openeo_processes_dask.process_implementations.exceptions import (
     BandExists,
@@ -71,13 +73,56 @@ def test_ndvi(temporal_interval, bounding_box, random_raster_data, process_regis
     with pytest.raises(KeyError):
         ndvi(cube_with_nothing_resolvable)
 
-    target_band = "yay"
+    target_band = "yayyyy"
     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"].dtype == np.dtype("<U4")
+    assert output_with_extra_dim.coords["common_name"][-1] == "NDVI"
+
+    output_with_other_dim = ndvi(
+        input_cube.rename({"common_name": "renamed"}), target_band=target_band
+    )
+    assert isinstance(output_with_other_dim, xr.DataArray)
+    assert len(output_with_other_dim.dims) == len(output.dims) + 1
+    assert (
+        len(output_with_other_dim.coords[band_dim])
+        == len(input_cube.coords[band_dim]) + 1
+    )
+    assert output_with_other_dim.coords[band_dim].values[-1] == target_band
+    assert output_with_other_dim.coords["renamed"][-1] == ""
+
+    # unnamed cube caused problems in the past
+    input_cube_noname = input_cube.rename(None)
+    out_noname = ndvi(input_cube_noname, target_band="ndvi")
+    assert (
+        len(out_noname.coords[band_dim]) == len(input_cube_noname.coords[band_dim]) + 1
+    )
 
     with pytest.raises(BandExists):
         output_with_extra_dim = ndvi(input_cube, target_band="t")
+
+
+def test_dummy_value():
+    out = indices._dummy_value(np.float32)
+    assert np.isnan(out)
+
+    out = indices._dummy_value(np.datetime64("2026-09-01").dtype)
+    assert np.isnat(out)
+
+    out = indices._dummy_value(np.timedelta64(100).dtype)
+    assert np.isnat(out)
+
+    out = indices._dummy_value(np.int64)
+    assert out == -1
+
+    out = indices._dummy_value(np.dtype("bool"))
+    assert out is False
+
+    out = indices._dummy_value(np.dtype("<U5"))
+    assert out == ""