Following @guillaume-eb suggestion, I removed the dask.delayed option and used only simple Dask Arrays. The code seems to work just fine now. Thank you.
import datetime
import numpy as np
import xarray as xr
import dask.array as dask_array
def uv2mag(u: xr.DataArray,
v: xr.DataArray
) -> dask_array:
return (u.data**2 + v.data**2)**0.5
def uv2dir(u: xr.DataArray,
v: xr.DataArray
) -> dask_array:
return np.rad2deg(np.arctan2(u.data, v.data))
def open_dataset(*args, **kwargs) -> xr.Dataset:
uv = kwargs.pop("uv", None)
ds = xr.open_dataset(*args, **kwargs)
if uv:
uvar, vvar = uv
ds["magnitude"] = (
ds[uvar].dims,
uv2mag(ds[uvar], ds[vvar]),
{"long_name": "magnitude"},
)
ds["direction"] = (
ds[uvar].dims,
uv2dir(ds[uvar], ds[vvar]),
{"long_name": "direction"},
)
return ds
url = "https://tds.hycom.org/thredds/dodsC/FMRC_ESPC-D-V02_uv3z/FMRC_ESPC-D-V02_uv3z_best.ncd"
uvar = "water_u"
vvar = "water_v"
ds = open_dataset(url, drop_variables="tau", chunks=10, uv=[uvar, vvar])
ds2 = ds.isel(time=slice(0, 5), depth=0, lat=1000, lon=1000)
ds2.load()