forum3894-append-doy-udf for https://forum.dataspace.copernicus.eu/t/feasibility-of-temporal-feature-encoding-doy-as-sin-cos-in-openeo-as-a-process-or-as-a-udf/3894

forum3894-append-doy-udf.ipynb

forum3894-append-doy-udf.py

import datetime
import calendar
import xarray
import numpy
import math
import openeo
from openeo.metadata import CubeMetadata


def broadcast_as_band(
    ref: xarray.DataArray,
    value: float = 0.0,
    label: str = "new",
    dim_name: str = "bands",
) -> xarray.DataArray:
    """
    Broadcast a scalar value to a new DataArray,
    based on a given DataArray,
    so that it can be concatenated as a new band to the given DataArray
    """
    assert dim_name in ref.dims
    shape = [(1 if d == dim_name else s) for (d, s) in zip(ref.dims, ref.shape)]
    data = numpy.full(shape=shape, fill_value=value)
    coords = {k: ([label] if k == dim_name else v) for (k, v) in ref.coords.items()}
    return xarray.DataArray(data=data, dims=ref.dims, coords=coords)


def append_scalars_as_bands(
    da: xarray.DataArray,
    values: dict,
    dim_name: str = "bands",
) -> xarray.DataArray:
    """
    Append scalar values
    (given as dictionary mapping band labels to values)
    as new bands to a given DataArray
    """
    objs = [da] + [
        broadcast_as_band(ref=da, value=v, label=k, dim_name=dim_name)
        for (k, v) in values.items()
    ]
    return xarray.concat(objs=objs, dim=dim_name)


def apply_datacube(cube: xarray.DataArray, context: dict) -> xarray.DataArray:
    assert not {"t", "temporal", "time"}.intersection(cube.dims), (
        "No temporal dimension expected in cube"
    )
    assert "t" in cube.attrs, "temporal label expected in cube attributes"

    date = cube.attrs["t"]
    day_of_year = date.timetuple().tm_yday
    days_in_year = 365 + calendar.isleap(date.year)
    angle = 2 * numpy.pi * (day_of_year - 1) / days_in_year

    cube = append_scalars_as_bands(
        da=cube,
        values={
            "sin_doy": math.sin(angle),
            "cos_doy": math.cos(angle),
        },
    )

    return cube


def apply_metadata(metadata: CubeMetadata, context: dict) -> CubeMetadata:
    return metadata.rename_labels(
        dimension="bands",
        target=metadata.band_names + ["sin_doy", "cos_doy"],
    )