Zonal mpsi function

build_envs/docs.yml

@@ -18,3 +18,4 @@ dependencies:
     - nbsphinx
     - xskillscore>=0.0.17
     - nc-time-axis
+    - uxarray

build_envs/environment.yml

@@ -14,6 +14,7 @@ dependencies:
   - pandas
   - xarray
   - xskillscore>=0.0.17
+  - uxarray>=2025.12.0
 # Packages listed below are for testing
   - geocat-datafiles
   - netcdf4

build_envs/min-deps.yml

@@ -23,3 +23,4 @@ dependencies:
   - pytest
   - pytest-cov
   - geocat-datafiles
+  - uxarray

docs/release-notes.rst

@@ -8,6 +8,17 @@
 Release Notes
 =============
 
+
+vYYYY.MM.## (unreleased)
+------------------------
+This release has updates to our internal structure to ingest UXarray datasets in a new zonal meridional stream function for unstructured grid datafiles using UXarray in ``meteorology.zonal_meridional_psi``. There are also two new helper functioins ``util._find_coord`` and ``util._find_optional_coord``
+
+New Features
+^^^^^^^^^^^^
+* Zonal Meridional Stream Function for unstructured grid datafiles using
+UXarray in ``meteorology.zonal_meridional_psi`` by `Julia Kent`_ in (:pr:`773`)
+
+
 v2026.01.0 (January 27, 2026)
 -----------------------------
 This release updates ``calendar_average`` and ``climate_anomaly`` to allow for

docs/user_api/index.rst

@@ -70,6 +70,7 @@ Meteorology
    relhum_water
    saturation_vapor_pressure
    saturation_vapor_pressure_slope
+   zonal_meridional_psi
 
 Spherical Harmonics
 ^^^^^^^^^^^^^^^^^^^

geocat/comp/__init__.py

@@ -33,6 +33,8 @@
     saturation_vapor_pressure_slope,
     delta_pressure,
     dpres_plev,
+    zonal_meridional_psi,
+    zonal_mpsi,
 )
 from .spherical import decomposition, recomposition, scale_voronoi
 from .stats import eofunc, eofunc_eofs, eofunc_pcs, eofunc_ts, pearson_r, nmse

geocat/comp/gc_util.py

@@ -1,4 +1,5 @@
 import typing
+import warnings
 
 
 def _generate_wrapper_docstring(
@@ -27,3 +28,128 @@ def _generate_wrapper_docstring(
 
     # assign docstring to wrapper function
     setattr(wrapper_fcn, '__doc__', wrapper_docstring)
+
+
+def _find_var(
+    ds,
+    standard_name=None,
+    long_name=None,
+    possible_names=None,
+    units=None,
+    description='variable',
+):
+    """
+    Find a variable using CF-compliant checks.
+
+    Searches in priority order:
+    1. CF standard_name attribute match (if standard_names provided)
+    2. Name attribute match (if long_names provided)
+    3. Direct variable name match (if possible_names provided)
+    4. Units match (if units provided)
+
+    Parameters
+    ----------
+    ds : xr.Dataset or ux.UxDataset
+        The dataset to search
+    standard_name : str, optional
+        CF standard_name to check in attrs
+    long_name : str, optional
+        long_name to check in attrs
+    possible_names : list of str, optional
+        List of possible variable names to check first
+    units : str, optional
+        Possible units to check in attrs
+    description : str, optional
+        String
+
+    Returns
+    -------
+    str
+        The name of the found variable
+
+    Raises
+    ------
+    KeyError
+        If no matching variable is found
+    """
+    error_parts = []
+
+    # First try CF standard_name attribute match
+    if standard_name:
+        for var_name in ds.data_vars:
+            var_attrs = ds[var_name].attrs
+            if var_name == standard_name or standard_name in var_attrs.get(
+                'standard_name', ''
+            ):
+                return var_name
+        error_parts.append(f"Tried standard_name: {standard_name}. ")
+
+    # Then try long_name attribute match
+    if long_name:
+        for var_name in ds.data_vars:
+            var_attrs = ds[var_name].attrs
+            if long_name in var_attrs.get('long_name', ''):
+                return var_name
+        error_parts.append(f"Tried long_name: {long_name}. ")
+
+    # Then try direct name match
+    if possible_names:
+        for name in possible_names:
+            if name in ds:
+                return name
+        error_parts.append(f"Tried names: {possible_names}. ")
+
+    # Finally try units match (less reliable)
+    if units:
+        for var_name in ds.data_vars:
+            var_attrs = ds[var_name].attrs
+            if 'units' in var_attrs:
+                if units in var_attrs.get('units', ''):
+                    warnings.warn(
+                        f"Found {description} '{var_name}' using units attribute only. "
+                        f"This is unreliable - multiple variables may share the same units. "
+                        f"Please verify this is correct and add CF standard_name. {error_parts}",
+                        UserWarning,
+                        stacklevel=3,
+                    )
+                    return var_name
+        error_parts.append(f"Tried units: {units}. ")
+
+    raise KeyError(f"Could not find {description} in dataset. {' '.join(error_parts)}")
+
+
+def _find_optional_var(
+    ds,
+    standard_name=None,
+    long_name=None,
+    possible_names=None,
+    units=None,
+    description=None,
+):
+    """
+    Find an optional variable using CF-compliant checks.
+
+    Parameters
+    ----------
+    ds : xr.Dataset or ux.UxDataset
+        The dataset to search
+    standard_name : str, optional
+        CF standard_name to check in attrs
+    long_name : str, optional
+        Long_name to check in attrs
+    possible_names : list of str, optional
+        List of possible variable names
+    units : str, optional
+        Possible units to check in attrs
+
+    Returns
+    -------
+    str or None
+        The name of the found coordinate, or None if not found
+    """
+    try:
+        return _find_var(
+            ds, standard_name, long_name, possible_names, units, description
+        )
+    except KeyError:
+        return None

geocat/comp/interpolation.py

@@ -6,6 +6,7 @@
 import metpy.interpolate
 import numpy as np
 import xarray as xr
+import uxarray as ux
 
 supported_types = typing.Union[xr.DataArray, np.ndarray]
 
@@ -383,7 +384,12 @@ def pressure_at_hybrid_levels(psfc, hya, hyb, p0=100000.0):
     """
     # check input types
     in_types = [type(psfc), type(hya), type(hyb)]
-    if not set(in_types).issubset({xr.DataArray, np.ndarray}) or len(set(in_types)) > 1:
+    if (
+        not set(in_types).issubset(
+            {xr.DataArray, np.ndarray, ux.core.dataarray.UxDataArray}
+        )
+        or len(set(in_types)) > 1
+    ):
         raise TypeError(
             "psfc, hya, and hyb must be xarray DataArrays or all numpy arrays"
         )