diff --git a/CHANGELOG.md b/CHANGELOG.md index a6e033b0..87539680 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Fixes - Fix the `ERA5` interface when making a pressure-level request with a single pressure level. This change accommodates CDS-Beta server behavior. Previously, a ValueError was raised in this case. +- Bypass the ValueError raised by `dask.array.gradient` when the input array is not correctly chunk along the level dimension. Previously, `Cocip` would raise an error when computing tau cirrus in the case that the `met` data had single chunks along the level dimension. ## v0.54.3 diff --git a/pycontrails/models/tau_cirrus.py b/pycontrails/models/tau_cirrus.py index b189d95e..847c60da 100644 --- a/pycontrails/models/tau_cirrus.py +++ b/pycontrails/models/tau_cirrus.py @@ -80,7 +80,14 @@ def tau_cirrus(met: MetDataset) -> xr.DataArray: met.data["air_pressure"], ) - dz = -dask.array.gradient(geopotential_height, axis=geopotential_height.get_axis_num("level")) + # dask.array.gradient expects at least 2 elements in each chunk + level_axis = geopotential_height.get_axis_num("level") + if geopotential_height.chunks: + level_chunks = geopotential_height.chunks[level_axis] # type: ignore[call-overload, index] + if any(chunk < 2 for chunk in level_chunks): + geopotential_height = geopotential_height.chunk(level=-1) + + dz = -dask.array.gradient(geopotential_height, axis=level_axis) dz = xr.DataArray(dz, dims=geopotential_height.dims) da = beta_e * dz