fix geopotential_height chunks

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
 

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