Repository: Unidata/MetPy

Issue:
Turbulence kinetic energy (tke) calculation on 4D xarray throwing ValueError
### 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

```python
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

```pytb
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
```


Discussion:
Using the `values` attribute to use the underlying numpy arrays works

```python
tke(ds.uwind.values, ds.vwind.values, ds.wwind.values)
```
@bvmcode That makes sense since it looks like the issue is in MetPy's xarray handling combined with what `tke` is doing.



Patch:
diff --git a/src/metpy/calc/turbulence.py b/src/metpy/calc/turbulence.py
index 56acc4fd736..bedd01991db 100644
--- a/src/metpy/calc/turbulence.py
+++ b/src/metpy/calc/turbulence.py
@@ -46,7 +46,7 @@ def get_perturbation(ts, axis=-1):
 
 
 @exporter.export
-@preprocess_and_wrap(wrap_like='u')
+@preprocess_and_wrap()
 def tke(u, v, w, perturbation=False, axis=-1):
     r"""Compute turbulence kinetic energy.
 
@@ -180,7 +180,7 @@ def kinematic_flux(vel, b, perturbation=False, axis=-1):
 
 
 @exporter.export
-@preprocess_and_wrap(wrap_like='u')
+@preprocess_and_wrap()
 def friction_velocity(u, w, v=None, perturbation=False, axis=-1):
     r"""Compute the friction velocity from the time series of velocity components.
 
@@ -245,5 +245,4 @@ def friction_velocity(u, w, v=None, perturbation=False, axis=-1):
     # the friction velocity is the 4th root of the kinematic momentum flux
     # As an optimization, first do inplace square root, then return the
     # square root of that. This is faster than np.power(..., 0.25)
-    np.sqrt(kf, out=kf)
-    return np.sqrt(kf)
+    return np.sqrt(np.sqrt(kf))


Test patch:
diff --git a/tests/calc/test_turbulence.py b/tests/calc/test_turbulence.py
index 7a3626aaa2e..a0e9b00dbbc 100644
--- a/tests/calc/test_turbulence.py
+++ b/tests/calc/test_turbulence.py
@@ -4,10 +4,12 @@
 """Test the `turbulence` module."""
 
 import numpy as np
-from numpy.testing import assert_almost_equal, assert_array_equal
 import pytest
+import xarray as xr
 
 from metpy.calc.turbulence import friction_velocity, get_perturbation, kinematic_flux, tke
+from metpy.testing import assert_almost_equal, assert_array_almost_equal, assert_array_equal
+from metpy.units import units
 
 
 #
@@ -30,6 +32,47 @@ def uvw_and_known_tke():
     return u, v, w, e_true
 
 
+@pytest.fixture()
+def uvw_and_known_tke_xarray():
+    """Provide a set of u, v, w with a known tke value as an xarray."""
+    # 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 16 linearly spaced values between -30 and 30
+    uwind_values = np.linspace(0, 30, num=16).reshape(2, 2, 2, 2)
+    vwind_values = np.linspace(-30, 0, num=16).reshape(2, 2, 2, 2)
+    wwind_values = np.linspace(-1, 2, num=16).reshape(2, 2, 2, 2)
+
+    # Apply units
+    u = uwind_values * units('m/s')
+    v = vwind_values * units('m/s')
+    w = wwind_values * units('m/s')
+
+    # Create the Dataset
+    ds = xr.Dataset(
+        {
+            'uwind': (dims, u),
+            'vwind': (dims, v),
+            'wwind': (dims, w),
+        },
+        coords={
+            'pressure': pressure,
+            'lat': lat,
+            'lon': lon,
+            'time': time
+        }
+    )
+
+    e_true = np.full((2, 2, 2), 1.005) * units('m^2/s^2')
+    return ds, e_true
+
+
 def test_no_tke_1d():
     """Test tke calculation where the expected value is 0."""
     observations = 5
@@ -71,6 +114,12 @@ def test_known_tke(uvw_and_known_tke):
     assert_array_equal(e_true, tke(u, v, w))
 
 
+def test_known_tke_xarray(uvw_and_known_tke_xarray):
+    """Test basic behavior of tke with known xarray values."""
+    data, e_true = uvw_and_known_tke_xarray
+    assert_array_almost_equal(tke(data.uwind, data.vwind, data.wwind), e_true)
+
+
 def test_known_tke_using_perturbation_velocities(uvw_and_known_tke):
     """Test basic behavior of tke with known values."""
     u, v, w, e_true = uvw_and_known_tke
@@ -334,6 +383,46 @@ def uvw_and_known_u_star_zero_mean():
     return u, v, w, u_star_true
 
 
+@pytest.fixture()
+def uvw_and_known_friction_velocity_xarray():
+    """Provide a set of u, v, w with a known tke value as an xarray."""
+    # 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 16 linearly spaced values between -30 and 30
+    uwind_values = np.linspace(0, 30, num=16).reshape(2, 2, 2, 2)
+    vwind_values = np.linspace(-30, 0, num=16).reshape(2, 2, 2, 2)
+    wwind_values = np.linspace(-1, 2, num=16).reshape(2, 2, 2, 2)
+
+    # Apply units
+    u = uwind_values * units('m/s')
+    v = vwind_values * units('m/s')
+    w = wwind_values * units('m/s')
+
+    # Create the Dataset
+    ds = xr.Dataset(
+        {
+            'uwind': (dims, u),
+            'vwind': (dims, v),
+            'wwind': (dims, w),
+        },
+        coords={
+            'pressure': pressure,
+            'lat': lat,
+            'lon': lon,
+            'time': time
+        }
+    )
+    expected = np.full((2, 2, 2), .3760603) * units('meter / second')
+    return ds, expected
+
+
 @pytest.fixture()
 def uvw_and_known_u_star_nonzero_mean():
     """Return components and friction velocity for a non-zero-mean time series."""
@@ -362,6 +451,12 @@ def test_u_star_1d_nonzero_mean(uvw_and_known_u_star_nonzero_mean):
                         u_star_true['uwvw'])
 
 
+def test_friction_velocity_nonzero_xarray(uvw_and_known_friction_velocity_xarray):
+    """Test friction velocity in 1d with an xarray."""
+    data, expected = uvw_and_known_friction_velocity_xarray
+    assert_array_almost_equal(friction_velocity(data.uwind, data.wwind, data.vwind), expected)
+
+
 def test_u_star_2d_axis_last_zero_mean(uvw_and_known_u_star_zero_mean):
     """Test friction velocity in 2D with a zero-mean time series along the last axis."""
     u, v, w, u_star_true = uvw_and_known_u_star_zero_mean