Merge branch 'Unidata:main' into benchmarking

Author: blue-jaye-121

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))

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