update resolution change sample udf (was broken)

Open-EO/openeo-geopyspark-driver#1397

Author: jdries

docs/examples/udf/udf_modify_spatial.py

@@ -2,48 +2,41 @@
 import xarray
 
 from openeo.metadata import CollectionMetadata
-from openeo.udf import XarrayDataCube
 from openeo.udf.debug import inspect
 
 
 def apply_metadata(input_metadata: CollectionMetadata, context: dict) -> CollectionMetadata:
-    xstep = input_metadata.get("x", "step")
-    ystep = input_metadata.get("y", "step")
-    new_metadata = {
-        "x": {"type": "spatial", "axis": "x", "step": xstep / 2.0, "reference_system": 4326},
-        "y": {"type": "spatial", "axis": "y", "step": ystep / 2.0, "reference_system": 4326},
-        "t": {"type": "temporal"},
-    }
-    return CollectionMetadata(new_metadata)
+    res= [ d.step / 2.0 for d in input_metadata.spatial_dimensions]
+    return input_metadata.resample_spatial(resolution=res)
 
 
 def fancy_upsample_function(array: np.array, factor: int = 2) -> np.array:
     assert array.ndim == 3
     return array.repeat(factor, axis=-1).repeat(factor, axis=-2)
 
 
-def apply_datacube(cube: XarrayDataCube, context: dict) -> XarrayDataCube:
-    cubearray: xarray.DataArray = cube.get_array().copy() + 60
+def apply_datacube(cube: xarray.DataArray, context: dict) -> xarray.DataArray:
+    cubearray: xarray.DataArray = cube.copy() + 60
 
     # We make prediction and transform numpy array back to datacube
 
     # Pixel size of the original image
-    init_pixel_size_x = cubearray.coords["x"][-1] - cubearray.coords["x"][-2]
-    init_pixel_size_y = cubearray.coords["y"][-1] - cubearray.coords["y"][-2]
+    init_pixel_size_x = cubearray.coords["x"][-1].item() - cubearray.coords["x"][-2].item()
+    init_pixel_size_y = cubearray.coords["y"][-1].item() - cubearray.coords["y"][-2].item()
 
     if cubearray.data.ndim == 4 and cubearray.data.shape[0] == 1:
         cubearray = cubearray[0]
     predicted_array = fancy_upsample_function(cubearray.data, 2)
-    inspect(data=predicted_array, message="predicted array")
+    inspect(data=predicted_array, message=f"predicted array")
     coord_x = np.linspace(
-        start=cube.get_array().coords["x"].min(),
-        stop=cube.get_array().coords["x"].max() + init_pixel_size_x,
+        start=cube.coords["x"][0].item(),
+        stop=cube.coords["x"][-1].item() + init_pixel_size_x,
         num=predicted_array.shape[-2],
         endpoint=False,
     )
     coord_y = np.linspace(
-        start=cube.get_array().coords["y"].min(),
-        stop=cube.get_array().coords["y"].max() + init_pixel_size_y,
+        start=cube.coords["y"].min().item(),
+        stop=cube.coords["y"].max().item() + init_pixel_size_y,
         num=predicted_array.shape[-1],
         endpoint=False,
     )
@@ -53,4 +46,4 @@ def apply_datacube(cube: XarrayDataCube, context: dict) -> XarrayDataCube:
         coords=dict(x=coord_x, y=coord_y),
     )
 
-    return XarrayDataCube(predicted_cube)
+    return predicted_cube