Turbulence kinetic energy (tke) calculation on 4D xarray throwing ValueError

[blue-jaye-121]

What went wrong?

When attempting to calculate the metpy.calc.tke(u, v, w, perturbation=False, axis=-1) function with a 4D xarray (pressure, lat, lon, time), the error "ValueError: different number of dimensions on data and dims: 3 vs 4" is thrown. It seems the preprocess and wrap decorator expects the function to return a 4D array if given one; the result is 3D. When removing the preprocess and wrap decorator from the function, it works as intended.

Operating System

Windows

Version

1.7.0

Python Version

3.12.10

Code to Reproduce

import xarray as xr
import numpy as np 
from metpy.calc import tke
from metpy.units import units

#Make a simple 4D array 

# Define coordinate values
pressure = [1000, 850]  # hPa
lat = [10, 20]          # degrees North
lon = [30, 40]          # degrees East
time = np.array(['2025-01-01T00:00', '2025-01-01T06:00'], dtype='datetime64')

# Define dimensions
dims = ("pressure", "lat", "lon", "time")

# Generate random data
uwind_data = np.random.uniform(-30, 30, size=(2, 2, 2, 2)) * units('m/s')
vwind_data = np.random.uniform(-30, 30, size=(2, 2, 2, 2)) * units('m/s')
wwind_data = np.random.uniform(-1, 1, size=(2, 2, 2, 2)) * units('m/s')

# Create the Dataset
ds = xr.Dataset(
    {
        "uwind": (dims, uwind_data),
        "vwind": (dims, vwind_data),
        "wwind": (dims, wwind_data),
    },
    coords={
        "pressure": pressure,
        "lat": lat,
        "lon": lon,
        "time": time
    }
)

#Attempt to calculate turbulence kinetic energy
print(tke(ds.uwind, ds.vwind, ds.wwind))

Errors, Traceback, and Logs

ValueError                                Traceback (most recent call last)
File c:\users\savan\unidata\metpy\untitled1.py:45
     30 ds = xr.Dataset(
     31     {
     32         "uwind": (dims, uwind_data),
   (...)
     41     }
     42 )
     44 #Attempt to calculate turbulence kinetic energy
---> 45 print(tke(ds.uwind, ds.vwind, ds.wwind))

File ~\miniforge3\envs\asv_env\Lib\site-packages\metpy\xarray.py:1342, in preprocess_and_wrap.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
   1340     return tuple(wrapping(*args) for args in zip(result, match, strict=False))
   1341 else:
-> 1342     return wrapping(result, match)

File ~\miniforge3\envs\asv_env\Lib\site-packages\metpy\xarray.py:1379, in _wrap_output_like_not_matching_units(result, match)
   1372 if (
   1373     not is_quantity(result) and (
   1374         is_quantity(match) or (output_xarray and is_quantity(match.data))
   1375     )
   1376 ):
   1377     result = units.Quantity(result)
   1378 return (
-> 1379     xr.DataArray(result, coords=match.coords, dims=match.dims)
   1380     if output_xarray and result is not None
   1381     else result
   1382 )

File ~\miniforge3\envs\asv_env\Lib\site-packages\xarray\core\dataarray.py:479, in DataArray.__init__(self, data, coords, dims, name, attrs, indexes, fastpath)
    477 data = _check_data_shape(data, coords, dims)
    478 data = as_compatible_data(data)
--> 479 coords, dims = _infer_coords_and_dims(data.shape, coords, dims)
    480 variable = Variable(dims, data, attrs, fastpath=True)
    482 if not isinstance(coords, Coordinates):

File ~\miniforge3\envs\asv_env\Lib\site-packages\xarray\core\dataarray.py:184, in _infer_coords_and_dims(shape, coords, dims)
    182 dims_tuple = tuple(dims)
    183 if len(dims_tuple) != len(shape):
--> 184     raise ValueError(
    185         "different number of dimensions on data "
    186         f"and dims: {len(shape)} vs {len(dims_tuple)}"
    187     )
    188 for d in dims_tuple:
    189     if not hashable(d):

ValueError: different number of dimensions on data and dims: 3 vs 4

[bvmcode]

Using the values attribute to use the underlying numpy arrays works

tke(ds.uwind.values, ds.vwind.values, ds.wwind.values)

[dopplershift]

@bvmcode That makes sense since it looks like the issue is in MetPy's xarray handling combined with what tke is doing.