Repository: zdlawrence/pyzome (language: python)

Issue:
`qglogp.plumb_wave_activity_flux` requires unneeded coordinates for some combos of WAF components
The `plumb_wave_activity_flux` function allows one to specify which WAF components to compute to allow for scenarios in which you may have only a lat/lon or lon/lev cross-section of data (or simply may desire a single component).

As currently written, the function requires coords for lon, lat, and lev, regardless of components desired. The checks should be more sophisticated.

Have to handle carefully, since, e.g., Fz doesn't require latitude for derivatives, but still uses the coriolis parameter in the calculation.
`qglogp.plumb_wave_activity_flux` requires unneeded coordinates for some combos of WAF components
The `plumb_wave_activity_flux` function allows one to specify which WAF components to compute to allow for scenarios in which you may have only a lat/lon or lon/lev cross-section of data (or simply may desire a single component).

As currently written, the function requires coords for lon, lat, and lev, regardless of components desired. The checks should be more sophisticated.

Have to handle carefully, since, e.g., Fz doesn't require latitude for derivatives, but still uses the coriolis parameter in the calculation.


Discussion / hints:

Fix patch:
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
index 93b6363..3b339d5 100644
--- a/CONTRIBUTING.rst
+++ b/CONTRIBUTING.rst
@@ -89,16 +89,3 @@ Ready to contribute? Here's how to set up `pyzome` for local development.
 
 7. Submit a pull request through the GitHub website.
 
-Pull Request Guidelines
------------------------
-
-Before you submit a pull request, check that it meets these guidelines:
-
-1. The pull request should include tests.
-2. If the pull request adds functionality, the docs should be updated. Put
-   your new functionality into a function with a docstring, and add the
-   feature to the list in README.rst.
-3. The pull request should work for Python 2.7, 3.3, 3.4, 3.5 and for PyPy. Check
-   https://travis-ci.org/zdlawrence/pyzome/pull_requests
-   and make sure that the tests pass for all supported Python versions.
-
diff --git a/docs/requirements.txt b/docs/requirements.txt
index bfc3f80..080cf40 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,1 +1,2 @@
-nbsphinx
\ No newline at end of file
+nbsphinx
+sphinx-autosummary-accessors
\ No newline at end of file
diff --git a/docs/source/basic.rst b/docs/source/basic.rst
index 4ead2f0..8a36900 100644
--- a/docs/source/basic.rst
+++ b/docs/source/basic.rst
@@ -11,11 +11,18 @@ convenience, and safety. For instance, while::
 will take the mean over the "lon" dimension regardless of 
 what ``lon`` contains::
 
-   >>> pzm.zonal_mean(ds, strict=True)
+   >>> import pyzome as pzm
+   >>> pzm.zonal_mean(ds)
 
 will infer the dimension name, and check to ensure that the 
 corresponding longitude dimension is regularly spaced and spans 
-a full 360 degrees. 
+a full 360 degrees. The check on the longitudes can be disabled
+by passing ``strict=False`` as a keyword argument.
+
+The functions in this module are also provided as pyzome xarray 
+accessor methods. The above example could also be achieved with:
+
+   >>> ds.pzm.zonal_mean()
 
 Functions
 ---------
diff --git a/docs/source/conf.py b/docs/source/conf.py
index cc72565..642225e 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -1,3 +1,5 @@
+import sphinx_autosummary_accessors
+
 # Configuration file for the Sphinx documentation builder.
 #
 # For the full list of built-in configuration values, see the documentation:
@@ -9,12 +11,18 @@
 project = "pyzome"
 copyright = "2023, Zachary D. Lawrence"
 author = "Zachary D. Lawrence"
-release = "0.0.3"
+release = "2023.09.0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-extensions = ["sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx.ext.todo"]
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.todo",
+    "sphinx_autosummary_accessors",
+]
+autodoc_member_order = "bysource"
 
 templates_path = ["_templates"]
 exclude_patterns = []
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 95103cd..603308c 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -11,10 +11,12 @@ Welcome to pyzome's documentation!
    :maxdepth: 1
    :caption: Contents:
 
+   whats-new
    install
    intro
    basic
    zonal_waves
+   xarray_accessor
    qglogp
    tem
    mock_data 
diff --git a/docs/source/install.rst b/docs/source/install.rst
index 82f44ae..336cbb5 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -6,9 +6,12 @@ Installation instructions
 Dependencies
 ------------
 
-``pyzome`` currently only requires `numpy <http://www.numpy.org/>`__, 
-`scipy <https://www.scipy.org/>`__, `xarray <http://xarray.pydata.org/en/stable/>`__, 
-and `xrft <https://xrft.readthedocs.io/en/latest/>`__.
+``pyzome`` requires 
+    * `numpy <http://www.numpy.org/>`__
+    * `scipy <https://www.scipy.org/>`__
+    * `xarray <http://xarray.pydata.org/en/stable/>`__
+    * `dask <https://docs.dask.org/en/stable/>`__
+    * `xrft <https://xrft.readthedocs.io/en/latest/>`__
 
 Installation
 ------------
diff --git a/docs/source/intro.rst b/docs/source/intro.rst
index ccc7e72..f068dde 100644
--- a/docs/source/intro.rst
+++ b/docs/source/intro.rst
@@ -13,8 +13,8 @@ quasi-geostrophic diagnostics.
 
 ``pyzome`` is built on top of the ``xarray`` package, 
 and is designed to work with ``xarray`` DataArrays and 
-Datasets. It is generally designed to work with regularly 
-gridded atmospheric data, as is commonly output fro many 
+Datasets. It is primarily intended to work with regularly 
+gridded atmospheric data, as is commonly output from many 
 models and reanalysis datasets. ``pyzome`` leverages coordinate 
 names and attributes to automatically perform operations 
 such as means and derivatives across geophysical dimensions. 
@@ -30,9 +30,9 @@ e.g., ``lon``, ``lons``, or ``longitude``, and ``lat``,
 ``lats`` or ``latitude`` (or other reasonably similar names), 
 and ``pyzome`` will automatically determine the correct dimensions.
 
-For explicit diagnostic calculations, ``pyzome`` expects that 
-quantities are expressed in SI units. This means that data on 
-pressure levels should have their pressures in Pascals 
+For explicit diagnostic calculations, ``pyzome`` currently expects 
+that  quantities are expressed in SI units. This means that, e.g., 
+data on pressure levels should have their pressures in Pascals 
 (with an appropriate coordinate name and units attribute). 
 These assumptions may be relaxed in the future by building 
 on top of other python modules such as `pint <https://pint.readthedocs.io/en/stable/>`_
diff --git a/docs/source/whats-new.rst b/docs/source/whats-new.rst
new file mode 100644
index 0000000..ae39aca
--- /dev/null
+++ b/docs/source/whats-new.rst
@@ -0,0 +1,29 @@
+.. _whats-new.2023.09.0:
+
+What's New
+==========
+
+v2023.09.0 (26 Sep 2023)
+-------------------------
+
+New Features
+~~~~~~~~~~~~
+- Added xarray accessor classes to use variable-agnostic pyzome functions directly on xarray objects.
+- Updated the "scipy" fft backend to use lazy evaluation, consistent with the "xrft" backend 
+
+Breaking Changes
+~~~~~~~~~~~~~~~~
+- pyzome.basic.zonal_mean and pyzome.basic.meridional_mean now use ``strict = True`` by default 
+
+Bug Fixes
+~~~~~~~~~
+- Addressed a bug in pyzome.qglogp.plumb_wave_activity_flux that caused the function to not properly check for required coordinates
+- Fixed the obtrusive logger from the recipes.zmd module
+
+Documentation
+~~~~~~~~~~~~~
+- Updated docs to include xarray accessor classes
+
+Internal Changes
+~~~~~~~~~~~~~~~~
+- switch to `CALVER <https://calver.org/>`_ versioning scheme
\ No newline at end of file
diff --git a/docs/source/xarray_accessor.rst b/docs/source/xarray_accessor.rst
new file mode 100644
index 0000000..5ea83b8
--- /dev/null
+++ b/docs/source/xarray_accessor.rst
@@ -0,0 +1,41 @@
+.. _xarray_accessors:
+
+xarray accessors (``pyzome.accessor``)
+=======================================
+pyzome provides a set of `xarray <http://xarray.pydata.org/en/stable/>`__ 
+accessors that link variable-agnostic pyzome functions to xarray object 
+methods. 
+
+pyzome coordinates
+------------------
+The pyzome coordinate mappings, which identify the relevant coordinates in 
+an xarray object, are stored internally as a dictionary. pyzome always uses 
+the convention that `lon` refers to longitude, `lat` to latitude, `plev` to 
+pressure levels, and `zonal_wavenum` to zonal wavenumbers. Thus, to get the 
+coordinate names for a given xarray object, you can use the `coord_map` 
+method of the xarray accessor:
+
+    >>> import pyzome as pzm
+    >>> # assume ds is an xarray dataset with coordinates called "longitude" and "pressure"
+    >>> ds.pzm.coord_map("lon")
+    >>> ds.pzm.coord_map("plev")
+
+Alternatively, you can directly refer to `lon`, `lat`, `plev`, and 
+`zonal_wavenum` as accessor _atrributes_ to directly obtain the underlying
+xarray coordinates:
+
+    >>> ds.pzm.lon
+    >>> ds.pzm.plev
+
+If any of these pyzome coordinates are not present or cannot be identified in 
+the xarray object, trying to access them will raise a KeyError.
+
+PyzomeAccessor
+--------------
+.. autoclass:: pyzome.accessor.PyzomeAccessor
+    :members:
+
+PyzomeDataArrayAccessor
+-----------------------
+.. autoclass:: pyzome.accessor.PyzomeDataArrayAccessor
+    :members:
\ No newline at end of file
diff --git a/docs/source/zonal_waves.rst b/docs/source/zonal_waves.rst
index e5899a1..6a00eac 100644
--- a/docs/source/zonal_waves.rst
+++ b/docs/source/zonal_waves.rst
@@ -11,18 +11,22 @@ These coefficients can then be used in the other functions
 described here to, e.g., filter the original field, find 
 the amplitude/phases of the waves, etc.
 
+:py:func:`pyzome.zonal_wave_coeffs` requires that the 
+underlying input have longitudes that are regularly 
+spaced and span the globe. This way the FFTs can be 
+done with no windowing/tapering, since the domain is 
+periodic. 
+
 :py:func:`pyzome.zonal_wave_coeffs` and some of the other 
 functions herein can use one of two modules to perform the 
-forward/inverse Fourier transforms, either ``"scipy"`` or 
-``"xrft"``. These are provided as separate options because:
-
-   * ``scipy`` FFTs output a dtype consistent with its input (e.g., ``complex64`` for ``float32`` input, and ``complex128`` for ``float64`` input) BUT works in a memory-eager manner (i.e., it loads the input xarray data)
-   * ``xrft`` FFTs always output ``complex128`` BUT they are able to operate lazily and in parallel using ``xarray`` and ``dask``.
+forward/inverse Fourier transforms, either ``scipy`` or 
+``xrft``. ``scipy`` is used by default since its 
+FFT functions output dtypes consistent with the input 
+(e.g., ``complex64`` for ``float32`` input, and ``complex128``
+for ``float64`` input). Both backends are able to operate 
+lazily by leveraging ``dask``. These options may be deprecated
+in the future in favor of using a single backend alone. 
 
-``scipy`` is used by default. In the future, these options 
-may be deprecated when/if a better solution is developed 
-that combines the benefits of both without explicit 
-loading or type-conversion. 
 
 Functions
 ---------
diff --git a/pyproject.toml b/pyproject.toml
index 394c279..5ca43ac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyzome"
-version = "0.0.3.1"
+version = "2023.09.0"
 authors = [
     {name = "Zachary D. Lawrence", email = "zachary.lawrence@colorado.edu"},
 ]
@@ -25,6 +25,7 @@ dependencies = [
     "numpy",
     "scipy",
     "xarray",
+    "dask",
     "xrft",
 ]
 
diff --git a/pyzome/__init__.py b/pyzome/__init__.py
index f5e4256..77a6c37 100644
--- a/pyzome/__init__.py
+++ b/pyzome/__init__.py
@@ -1,4 +1,5 @@
+from .accessor import PyzomeDataArrayAccessor, PyzomeDatasetAccessor
 from .basic import zonal_mean, meridional_mean
 from .zonal_waves import zonal_wave_coeffs
 
-__version__ = "0.0.3.1"
+__version__ = "2023.09.0"
diff --git a/pyzome/accessor.py b/pyzome/accessor.py
new file mode 100644
index 0000000..2e79160
--- /dev/null
+++ b/pyzome/accessor.py
@@ -0,0 +1,204 @@
+from typing import Sequence, Optional
+
+import xarray as xr
+
+from .constants import SCALE_HEIGHT, PREF
+from .checks import infer_xr_coord_names
+from .basic import zonal_mean, meridional_mean
+from .qglogp import add_logp_altitude
+from .zonal_waves import (
+    zonal_wave_coeffs,
+    inflate_zonal_wave_coeffs,
+    filter_by_zonal_wave_truncation,
+    zonal_wave_contributions,
+)
+
+
+class PyzomeAccessor:
+    r"""Base class for pyzome xarray accessors. Provides the common coordinate
+    mapping functionality for selecting correct xarray coordinate names for
+    use in pyzome functions. The properties and methods of this base class
+    apply to both xarray DataArray and Dataset objects.
+
+    Examples
+    --------
+    >>> import xarray as xr
+    >>> import pyzome as pzm
+    >>> ds = xr.open_dataset("...")
+    >>> ds.pzm.lon
+    >>> ds.pzm.zonal_mean()
+    """
+
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+        self._coord_map = infer_xr_coord_names(xarray_obj)
+
+    def coord_map(self, coord: str):
+        r"""Provides the mapping between the pyzome standard coordinate names
+        (lon, lat, plev, zonal_wavenum) and the actual coordinate names in the
+        xarray object.
+
+        Parameters
+        ----------
+        coord : str
+            The pyzome standard coordinate name (lon, lat, plev, zonal_wavenum)
+
+        Returns
+        -------
+        str
+            The actual coordinate name in the xarray object
+
+        Raises
+        ------
+        KeyError
+            If the underlying coordinate name is not found in the xarray object
+        """
+        if coord not in self._coord_map:
+            raise KeyError(
+                f"pyzome coordinate '{coord}' not identified in original data"
+            )
+        return self._coord_map[coord]
+
+    @property
+    def lon(self):
+        r"""The longitude coordinate of the xarray object, if it exists/was identified"""
+        return self._obj[self.coord_map("lon")]
+
+    @property
+    def lat(self):
+        r"""The latitude coordinate of the xarray object, if it exists/was identified"""
+        return self._obj[self.coord_map("lat")]
+
+    @property
+    def plev(self):
+        r"""The pressure level coordinate of the xarray object, if it exists/was identified"""
+        return self._obj[self.coord_map("plev")]
+
+    @property
+    def zonal_wavenum(self):
+        r"""The zonal wavenumber coordinate of the xarray object, if it exists/was identified"""
+        return self._obj[self.coord_map("zonal_wavenum")]
+
+    def add_logp_altitude(self, H: float = SCALE_HEIGHT, p0: float = PREF):
+        r"""Adds a log-pressure altitude coordinate to the xarray object. Requires
+        that the pyzome `plev` coordinate be identified, and that it be in units of Pa.
+
+        See Also
+        --------
+        pyzome.qglogp.add_logp_altitude
+        """
+        return add_logp_altitude(self._obj, plev_coord=self.plev.name, H=H, p0=p0)
+
+    def zonal_mean(self, strict: bool = True):
+        r"""Computes the zonal mean of the xarray object, if a longitude coordinate
+        was identifed.
+
+        See Also
+        --------
+        pyzome.basic.zonal_mean
+        """
+        return zonal_mean(self._obj, lon_coord=self.lon.name, strict=strict)
+
+    def meridional_mean(self, lat1: float, lat2: float, strict: bool = True):
+        r"""Computes the meridional mean of the xarray object, if a latitude coordinate
+        was identifed.
+
+        See Also
+        --------
+        pyzome.basic.meridional_mean
+        """
+        return meridional_mean(
+            self._obj, lat1, lat2, lat_coord=self.lat.name, strict=strict
+        )
+
+
+@xr.register_dataarray_accessor("pzm")
+class PyzomeDataArrayAccessor(PyzomeAccessor):
+    r"""xarray DataArray accessor for pyzome functions. Inherits from
+    PyzomeAccessor for maintaining the same coordinate mapping functionality.
+
+    Examples
+    --------
+    >>> import xarray as xr
+    >>> import pyzome as pzm
+    >>> da = xr.open_dataarray("...") # must be a DataArray
+    >>> da.pzm.lon
+    >>> da.pzm.zonal_wave_coeffs().pzm.filter_by_zonal_wave_truncation(waves=[1,2,3])
+    """
+
+    def __init__(self, xarray_obj: xr.DataArray):
+        self._obj = xarray_obj
+        self._coord_map = infer_xr_coord_names(xarray_obj)
+
+    def zonal_wave_coeffs(
+        self, waves: Optional[Sequence[int]] = None, fftpkg: str = "scipy"
+    ):
+        r"""Computes the zonal wave coefficients of a DataArray, if a longitude
+        coordinate was identifed.
+
+        See Also
+        --------
+        pyzome.zonal_waves.zonal_wave_coeffs
+        """
+        return zonal_wave_coeffs(
+            self._obj, waves=waves, fftpkg=fftpkg, lon_coord=self.lon.name
+        )
+
+    def inflate_zonal_wave_coeffs(self):
+        r"""Inflates the zonal wave coefficients of a DataArray to the full
+        spectrum expected for inverse transforms.
+
+        See Also
+        --------
+        pyzome.zonal_waves.inflate_zonal_wave_coeffs
+        """
+        return inflate_zonal_wave_coeffs(self._obj, wave_coord=self.zonal_wavenum.name)
+
+    def filter_by_zonal_wave_truncation(
+        self,
+        waves: Sequence[int],
+        fftpkg: str = "scipy",
+        lons: Optional[xr.DataArray] = None,
+    ):
+        r"""Filters the input DataArray by truncating to include only the specified
+        zonal wavenumbers.
+
+        See Also
+        --------
+        pyzome.zonal_waves.filter_by_zonal_wave_truncation
+        """
+        return filter_by_zonal_wave_truncation(
+            self._obj,
+            waves=waves,
+            fftpkg=fftpkg,
+            wave_coord=self.zonal_wavenum.name,
+            lons=lons,
+        )
+
+    def zonal_wave_contributions(
+        self,
+        waves: Sequence[int],
+        fftpkg: str = "scipy",
+        lons: Optional[xr.DataArray] = None,
+    ):
+        r"""Computes the individual contributions of each given zonal wavenumber
+        to the input DataArray.
+
+        See Also
+        --------
+        pyzome.zonal_waves.zonal_wave_contributions
+        """
+        return zonal_wave_contributions(
+            self._obj,
+            waves=waves,
+            fftpkg=fftpkg,
+            wave_coord=self.zonal_wavenum.name,
+            lons=lons,
+        )
+
+
+@xr.register_dataset_accessor("pzm")
+class PyzomeDatasetAccessor(PyzomeAccessor):
+    def __init__(self, xarray_obj: xr.Dataset):
+        self._obj = xarray_obj
+        self._coord_map = infer_xr_coord_names(xarray_obj)
diff --git a/pyzome/basic.py b/pyzome/basic.py
index 982094a..c2b66c2 100644
--- a/pyzome/basic.py
+++ b/pyzome/basic.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import overload
 
 import numpy as np
 import xarray as xr
@@ -6,11 +7,29 @@
 from .checks import has_global_regular_lons, infer_xr_coord_names
 
 
+@overload
 def zonal_mean(
-    dat: xr.DataArray | xr.Dataset,
+    dat: xr.Dataset,
     lon_coord: str = "",
     strict: bool = False,
-) -> xr.DataArray | xr.Dataset:
+) -> xr.Dataset:
+    ...
+
+
+@overload
+def zonal_mean(
+    dat: xr.DataArray,
+    lon_coord: str = "",
+    strict: bool = False,
+) -> xr.DataArray:
+    ...
+
+
+def zonal_mean(
+    dat,
+    lon_coord: str = "",
+    strict: bool = True,
+):
     r"""Compute the zonal mean.
 
     This is primarily a convenience function that will make other
@@ -19,7 +38,7 @@ def zonal_mean(
 
     Parameters
     ----------
-    dat : `xarray.DataArray` or `xarray.Dataset`
+    dat : ``xarray.DataArray`` or ``xarray.Dataset``
         data containing a dimension named longitude that spans all 360 degrees
 
     lon_coord : str, optional
@@ -28,13 +47,12 @@ def zonal_mean(
         corresponds to the longitude
 
     strict : bool, optional
-        If True, the function will check whether the longitudes span 360
-        degrees with regular spacing. If False (the default), this check is
-        skipped.
+        If True (the default), the function will check whether the longitudes
+        span 360 degrees with regular spacing. If False this check is skipped.
 
     Returns
     -------
-    zonal average: `xarray.DataArray` or `xarray.Dataset`
+    zonal average: ``xarray.DataArray`` or ``xarray.Dataset``
         The mean across the longitude dimension
 
     """
@@ -49,19 +67,42 @@ def zonal_mean(
     return dat.mean(lon_coord)
 
 
+@overload
 def meridional_mean(
-    dat: xr.DataArray | xr.Dataset,
+    dat: xr.Dataset,
     lat1: float,
     lat2: float,
     lat_coord: str = "",
-) -> xr.DataArray | xr.Dataset:
+    strict: bool = False,
+) -> xr.Dataset:
+    ...
+
+
+@overload
+def meridional_mean(
+    dat: xr.DataArray,
+    lat1: float,
+    lat2: float,
+    lat_coord: str = "",
+    strict: bool = False,
+) -> xr.DataArray:
+    ...
+
+
+def meridional_mean(
+    dat,
+    lat1: float,
+    lat2: float,
+    lat_coord: str = "",
+    strict: bool = True,
+):
     r"""Compute the cos(lat) weighted mean of data between two latitudes.
 
     This function is imported at the top level of the package by default.
 
     Parameters
     ----------
-    dat : `xarray.DataArray` or `xarray.Dataset`
+    dat : ``xarray.DataArray`` or ``xarray.Dataset``
         data containing a latitude dimension that spans
         lat1 and lat2. The cos(lat) weighting assumes that the
         latitudes are equally spaced. If given a dataset, the
@@ -79,9 +120,14 @@ def meridional_mean(
         string (the default), the function tries to infer which coordinate
         corresponds to the latitude
 
+    strict : bool, optional
+        If True (the default), the function will check whether the latitudes
+        on `dat` span `lat1` and `lat2` inclusive. If False this check is
+        skipped.
+
     Returns
     -------
-    meridional average: `xarray.DataArray` or `xarray.Dataset`
+    meridional average: ``xarray.DataArray`` or ``xarray.Dataset``
         the weighted mean across the latitude dimension limited
         by lat1 and lat2
 
@@ -97,7 +143,7 @@ def meridional_mean(
 
     min_lat = float(dat[lat_coord].min())
     max_lat = float(dat[lat_coord].max())
-    if not ((min_lat <= lat1 <= max_lat) and (min_lat <= lat2 <= max_lat)):
+    if strict and not ((min_lat <= lat1 <= max_lat) and (min_lat <= lat2 <= max_lat)):
         msg = (
             f"data only contains lats in range of {min_lat} to {max_lat} "
             + f"(chose lat1={lat1}, lat2={lat2})"
@@ -108,4 +154,5 @@ def meridional_mean(
     ixs = {lat_coord: np.logical_and(lats >= lat1, lats <= lat2)}
     wgts = np.cos(np.deg2rad(dat[lat_coord].isel(ixs)))
 
-    return dat.isel(ixs).weighted(wgts).mean(lat_coord)  # type: ignore
+    with xr.set_options(keep_attrs=True):
+        return dat.isel(ixs).weighted(wgts).mean(lat_coord)  # type: ignore
diff --git a/pyzome/checks.py b/pyzome/checks.py
index 9315e55..2789d4b 100644
--- a/pyzome/checks.py
+++ b/pyzome/checks.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
+
 import re
+from typing import Iterable
 
 import numpy as np
 import xarray as xr
@@ -10,7 +12,7 @@
 coord_regex = {
     "lon": re.compile("lon[a-z]*"),
     "lat": re.compile("lat[a-z]*"),
-    "lev": re.compile("(p?lev|pre?s|isobaric)[a-z]*"),
+    "plev": re.compile("(p?lev|pre?s|isobaric)[a-z]*"),
     "zonal_wavenum": re.compile("(zonal_wavenum|wavenum_lon|lon_wavenum)[a-z]*"),
 }
 
@@ -54,7 +56,14 @@
     },
     "pressure": {"Pa", "Pascals"},
     "temperature": {"K", "degK", "Kelvin"},
-    "prs_vvel": {"Pa s-1", "Pa*s-1", "Pa * s-1", "Pa/s", "Pa / s", "Pascals per second"},
+    "prs_vvel": {
+        "Pa s-1",
+        "Pa*s-1",
+        "Pa * s-1",
+        "Pa/s",
+        "Pa / s",
+        "Pascals per second",
+    },
     "wind": {"m s-1", "m*s-1", "m * s-1", "m/s", "m / s", "meters per second"},
 }
 
@@ -68,7 +77,7 @@ def has_global_regular_lons(
 
     Parameters
     ----------
-    lons : `xarray.DataArray` or `np.array`
+    lons : ``xarray.DataArray`` or ``np.array``
         data containing the longitudes, in degrees
 
     enforce : bool, optional
@@ -107,8 +116,8 @@ def has_global_regular_lons(
 
 
 def infer_xr_coord_names(
-    dat: xr.DataArray | xr.Dataset, required: list[str] = []
-) -> dict:
+    dat: xr.DataArray | xr.Dataset, required: Iterable[str] = []
+) -> dict[str, str]:
     r"""A convenience function that identifies commonly used coordinate names
     for gridded earth datasets. This function enables other functions in
     pyzome to perform operations across coordinates without the user having
@@ -118,8 +127,8 @@ def infer_xr_coord_names(
 
     Parameters
     ----------
-    dat : `xarray.DataArray` or `xr.Dataset`
-        The data containing the coordinates of lat, lon, lev, etc.
+    dat : ``xarray.DataArray`` or ``xr.Dataset``
+        The data containing the coordinates of lat, lon, plev, etc.
 
     required: List of strings, optional
         If this kwarg is defined, this function will throw errors if it is
@@ -129,7 +138,7 @@ def infer_xr_coord_names(
 
     Returns
     -------
-    coord_names: dict 
+    coord_names: dict
         string keys and values, with the keys being the
         coordinate name category (e.g., lat, lon), and the values being the
         actual coordinate name in the given xarray data. E.g,
@@ -142,7 +151,7 @@ def infer_xr_coord_names(
     dat_coords = list(dat.coords)
     for dc in dat_coords:
         for coord in coord_regex:
-            if coord_regex[coord].match(dc.lower()): # type: ignore
+            if coord_regex[coord].match(dc.lower()):  # type: ignore
                 # Check if more than one coord in dat matches the same pattern
                 if coord in coord_names:
                     msg = (
@@ -175,7 +184,7 @@ def check_var_SI_units(dat: xr.DataArray, var: str, enforce: bool = False) -> bo
 
     Parameters
     ----------
-    dat : `xarray.DataArray`
+    dat : ``xarray.DataArray``
         The data with attributes to check for SI units.
 
     var: string
@@ -209,14 +218,14 @@ def check_var_SI_units(dat: xr.DataArray, var: str, enforce: bool = False) -> bo
     return units_SI
 
 
-def check_for_logp_coord(dat: xr.DataArray, enforce: bool = False) -> bool:
+def check_for_logp_coord(dat: xr.DataArray | xr.Dataset, enforce: bool = False) -> bool:
     r"""A function that checks whether a log-pressure altitude coordinate
     (assumed to be created by pyzome) exists in the given DataArray. Uses
     a combination of units and long_name to check.
 
     Parameters
     ----------
-    dat : `xarray.DataArray`
+    dat : ``xarray.DataArray``
         The data to check.
 
     enforce : bool, optional
diff --git a/pyzome/mock_data.py b/pyzome/mock_data.py
index b17d166..563e19b 100644
--- a/pyzome/mock_data.py
+++ b/pyzome/mock_data.py
@@ -73,8 +73,8 @@ def lat_coord(
 
     Returns
     -------
-    latitudes: `xarray.DataArray` 
-        Regularly spaced latitudes with the given resolution 
+    latitudes: ``xarray.DataArray``
+        Regularly spaced latitudes with the given resolution
         (and name/limits, if given)
 
     """
@@ -116,7 +116,7 @@ def lon_coord(
 
     Returns
     -------
-    longitudes: `xarray.DataArray` 
+    longitudes: ``xarray.DataArray``
         Regularly spaced longitudes with the given resolution
         (and name/limits, if given)
 
@@ -145,7 +145,7 @@ def plev_coord(
     levels_per_decade: int,
     left_lim_exponent: int = 3,
     right_lim_exponent: int = 0,
-    name: str = "lev",
+    name: str = "plev",
     units: str = "hPa",
 ) -> xr.DataArray:
     """
@@ -154,13 +154,13 @@ def plev_coord(
     Parameters
     ----------
     levels_per_decade : int
-        The number of levels per decade (i.e., between 10^3 and 10^2, 10^2 
+        The number of levels per decade (i.e., between 10^3 and 10^2, 10^2
         and 10^1, etc.)
     left_lim_exponent : int, optional
-        The base-10 exponent for the left-most pressure level, which defaults 
+        The base-10 exponent for the left-most pressure level, which defaults
         to 3 (for 1000 hPa)
     right_lim_exponent : int, optional
-        The base-10 exponent for the right-most pressure level, which defaults 
+        The base-10 exponent for the right-most pressure level, which defaults
         to 0 (for 1 hPa)
     name : str, optional
         The name of the coordinate, which defaults to "lev"
@@ -169,7 +169,7 @@ def plev_coord(
 
     Returns
     -------
-    pressures: `xarray.DataArray`
+    pressures: ``xarray.DataArray``
         Pressure-like coordinate with logarithmic spacing having the given
         number of levels per decade (and name/limits/units, if given)
 
@@ -217,7 +217,7 @@ def time_coord(
 
     Returns
     -------
-    times: `xarray.DataArray`
+    times: ``xarray.DataArray``
         Time-like coordinate with the given start/end dates and frequency
 
     """
@@ -248,13 +248,13 @@ def create_dummy_geo_field(
 
     Parameters
     ----------
-    lons : `xarray.DataArray`
+    lons : ``xarray.DataArray``
         The longitudes of the data. Defaults to None (longitudes excluded)
-    lats : `xarray.DataArray`
+    lats : ``xarray.DataArray``
         The latitudes of the data. Defaults to None (latitudes excluded)
-    levs : `xarray.DataArray`, optional
+    levs : ``xarray.DataArray``, optional
         The vertical levels of the data. Defaults to None (vertical levels excluded)
-    times : `xarray.DataArray`, optional
+    times : ``xarray.DataArray``, optional
         The times of the data. Defaults to None (times excluded)
     name : str, optional
         The name of the data. Defaults to "dummy"
@@ -263,7 +263,7 @@ def create_dummy_geo_field(
 
     Returns
     -------
-    dummy_data: `xarray.DataArray` 
+    dummy_data: ``xarray.DataArray``
         Random data with corresponding geophysical coordinates
 
     """
@@ -303,22 +303,22 @@ def create_dummy_geo_dataset(
     ----------
     field_names : list[str]
         The names of the fields
-    lons : `xarray.DataArray`
+    lons : ``xarray.DataArray``
         The longitudes of the data
-    lats : `xarray.DataArray`
+    lats : ``xarray.DataArray``
         The latitudes of the data
-    levs : `xarray.DataArray`, optional
+    levs : ``xarray.DataArray``, optional
         The vertical levels of the data. Defaults to None (vertical levels excluded)
-    times : `xarray.DataArray`, optional
+    times : ``xarray.DataArray``, optional
         The times of the data. Defaults to None (times excluded)
     field_attrs : dict, optional
         The attributes of the data. Defaults to None. If provided, this should be
         a nested dictionary with the outer keys corresponding to the field names,
         and inner keys/dicts corresponding to the attributes for each field.
-    
+
     Returns
     -------
-    dummy_dataset: `xarray.Dataset` 
+    dummy_dataset: ``xarray.Dataset``
         Dataset with fields having random data matching the given geophysical coordinates
 
     """
diff --git a/pyzome/qglogp.py b/pyzome/qglogp.py
index 1362630..33a7ced 100644
--- a/pyzome/qglogp.py
+++ b/pyzome/qglogp.py
@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import overload
 
 import numpy as np
 import xarray as xr
@@ -15,20 +16,40 @@
 )
 
 
+@overload
 def add_logp_altitude(
-    dat: xr.Dataset | xr.DataArray,
-    lev_coord: str = "",
+    dat: xr.Dataset,
+    plev_coord: str = "",
     H: float = SCALE_HEIGHT,
     p0: float = PREF,
-) -> xr.Dataset | xr.DataArray:
+) -> xr.Dataset:
+    ...
+
+
+@overload
+def add_logp_altitude(
+    dat: xr.DataArray,
+    plev_coord: str = "",
+    H: float = SCALE_HEIGHT,
+    p0: float = PREF,
+) -> xr.DataArray:
+    ...
+
+
+def add_logp_altitude(
+    dat,
+    plev_coord: str = "",
+    H: float = SCALE_HEIGHT,
+    p0: float = PREF,
+):
     r"""A convenience function for adding a log-pressure altitude coordinate
     to an xarray Dataset or DataArray
 
     Parameters
     ----------
-    dat : `xarray.Dataset` or `xarray.DataArray`
+    dat : ``xarray.Dataset`` or ``xarray.DataArray``
         data containing a pressure coordinate in SI units (Pa)
-    lev_coord : string, optional
+    plev_coord : string, optional
         The name of the pressure coordinate in the input data. Defaults
         to an empty string, for which the function will try to infer the
         pressure coordinate
@@ -40,17 +61,17 @@ def add_logp_altitude(
 
     Returns
     -------
-    `xarray.Dataset` or `xarray.DataArray`
+    ``xarray.Dataset`` or ``xarray.DataArray``
         The input data with an added log-pressure altitude coordinate named 'z'
 
     """
 
-    if lev_coord == "":
-        coords = infer_xr_coord_names(dat, required=["lev"])
-        lev_coord = coords["lev"]
-    check_var_SI_units(dat[lev_coord], "pressure", enforce=True)
+    if plev_coord == "":
+        coords = infer_xr_coord_names(dat, required=["plev"])
+        plev_coord = coords["plev"]
+    check_var_SI_units(dat[plev_coord], "pressure", enforce=True)
 
-    z = -H * np.log(dat[lev_coord] / p0)
+    z = -H * np.log(dat[plev_coord] / p0)
     z.attrs["units"] = "m"  # type: ignore
     z.attrs["long_name"] = "log-pressure altitude"  # type: ignore
     z.attrs["note"] = "added by pyzome"  # type: ignore
@@ -63,13 +84,12 @@ def buoyancy_frequency_squared(
     Rs: float = GAS_CONST_DRY_AIR,
     Cp: float = SPEC_HEAT_DRY_AIR,
     H: float = SCALE_HEIGHT,
-    p0: float = PREF,
 ) -> xr.DataArray:
     r"""Calculates the buoyancy frequency squared given temperature.
 
     Parameters
     ----------
-    T : `xarray.DataArray`
+    T : ``xarray.DataArray``
         The temperature data in units of Kelvin
     Rs : float, optional
         Specific gas constant. Defaults to 287.058 J/kg/K for dry
@@ -80,13 +100,10 @@ def buoyancy_frequency_squared(
     H : float, optional
         The scale height used to calculate the log-pressure altitude.
         Defaults to 7000 m.
-    p0 : float, optional
-        Reference pressure used to calculate the log-pressure altitude.
-        Defaults to 100000 Pa for Earth.
 
     Returns
     -------
-    Nsq: `xarray.DataArray`
+    Nsq: ``xarray.DataArray``
         The buoyancy frequency squared, in units of s-2.
 
     """
@@ -118,9 +135,9 @@ def merid_grad_qgpv(
 
     Parameters
     ----------
-    u : `xarray.DataArray`
+    u : ``xarray.DataArray``
         The zonal mean zonal wind data.
-    Nsq : `xarray.DataArray` or float
+    Nsq : ``xarray.DataArray`` or float
         The squared buoyancy frequency. Nsq need not have the same dimensions
         as u, but it should be consistent with u in the sense that it must be
         able to be properly broadcasted when used in computations with u (it is
@@ -150,7 +167,7 @@ def merid_grad_qgpv(
 
     Returns
     -------
-    qgpv_grad: `xarray.DataArray` or tuple of DataArrays
+    qgpv_grad: ``xarray.DataArray`` or tuple of DataArrays
         Depending on the value of the terms keyword argument, either the full
         meridional QGPV gradient field, or a tuple of the individual terms
         that make up the total, in units of s-1.
@@ -218,11 +235,11 @@ def refractive_index(
 
     Parameters
     ----------
-    u : `xarray.DataArray`
+    u : ``xarray.DataArray``
         The zonal mean zonal wind data.
-    q_phi : `xarray.DataArray`
+    q_phi : ``xarray.DataArray``
         The meridional QGPV gradient
-    Nsq : `xarray.DataArray` or float
+    Nsq : ``xarray.DataArray`` or float
         The squared buoyancy frequency. Nsq need not have the same dimensions
         as u, but it should be consistent with u and q_phi in the sense that it
         must be able to be properly broadcasted when used in computations (it is
@@ -262,7 +279,7 @@ def refractive_index(
 
     Returns
     -------
-    RIsq: `xarray.DataArray` or tuple of DataArrays
+    RIsq: ``xarray.DataArray`` or tuple of 3 DataArrays
         Depending on the value of the terms keyword argument, either the full
         squared refractive index field, or a tuple of the individual terms
         that make the total, in units of m-2.
@@ -316,9 +333,9 @@ def plumb_wave_activity_flux(
 
     Parameters
     ----------
-    psip : `xarray.DataArray`
+    psip : ``xarray.DataArray``
         The eddy streamfunction in units of m+2 s-1
-    Nsq : `xarray.DataArray` or float
+    Nsq : ``xarray.DataArray`` or float
         The squared buoyancy frequency. Nsq need not have the same dimensions
         as psip, but it should be consistent with psip in the sense that it
         must be able to be properly broadcasted when used in computations (it is
@@ -342,51 +359,54 @@ def plumb_wave_activity_flux(
 
     Returns
     -------
-    waf: list of `xarray.DataArray`s
+    waf: list of ``xarray.DataArray``
         The wave activity flux vector components consistent with the
         desired components keyword argument.
 
     Todo
     -----
     * Unit checks on psip and Nsq (if Nsq is a DataArray)
-    * Add units on output terms
 
     """
 
     check_for_logp_coord(psip, enforce=True)
-    coords = infer_xr_coord_names(psip, required=["lev"])
-
-    if ("y" in components) and ("lat" not in coords):
-        msg = "A latitude coordinate must be available"
-        raise ValueError(msg)
-    elif (("x" in components) or ("z" in components)) and ("lon" not in coords):
-        msg = "A longitude coordinate must be available"
-        raise ValueError(msg)
+    coords = infer_xr_coord_names(psip, required=["lon", "lat", "plev"])
 
+    plev_coord = coords["plev"]
     lat_coord = coords["lat"]
     lon_coord = coords["lon"]
+
     r2d = 180.0 / np.pi
     lats = np.deg2rad(psip[lat_coord])
     cosphi = np.cos(lats)
     f = 2 * Omega * np.sin(lats)
 
-    p = psip[coords["lev"]] / 100000.0
+    p = psip[plev_coord] / 100000.0
 
     waf = []
     dpsi_dlam = r2d * psip.differentiate(lon_coord, edge_order=2)
     if "x" in components:
         d2psi_dlam2 = r2d * dpsi_dlam.differentiate(lon_coord, edge_order=2)
         wafx = (p / (2 * a * a * cosphi)) * ((dpsi_dlam) ** 2 - psip * d2psi_dlam2)
+        wafx.name = "plumb_wafx"
+        wafx.attrs["long_name"] = "Longitude component of Plumb Wave Activity Flux"
+        wafx.attrs["units"] = "m+2 s-2"
         waf.append(wafx)
     if "y" in components:
         dpsi_dphi = r2d * psip.differentiate(lat_coord, edge_order=2)
         d2psi_dphidlam = r2d * dpsi_dphi.differentiate(lon_coord, edge_order=2)
         wafy = (p / (2 * a * a)) * (dpsi_dlam * dpsi_dphi - psip * d2psi_dphidlam)
+        wafy.name = "plumb_wafy"
+        wafy.attrs["long_name"] = "Latitude component of Plumb Wave Activity Flux"
+        wafy.attrs["units"] = "m+2 s-2"
         waf.append(wafy)
     if "z" in components:
         dpsi_dz = psip.differentiate("z", edge_order=2)
         d2psi_dlamdz = dpsi_dlam.differentiate("z", edge_order=2)
         wafz = (p * f * f / (2 * Nsq * a)) * (dpsi_dlam * dpsi_dz - psip * d2psi_dlamdz)
+        wafz.name = "plumb_wafz"
+        wafz.attrs["long_name"] = "Vertical component of Plumb Wave Activity Flux"
+        wafz.attrs["units"] = "m+2 s-2"
         waf.append(wafz)
 
     return waf
diff --git a/pyzome/recipes/zmd.py b/pyzome/recipes/zmd.py
index b12b406..e2a5685 100644
--- a/pyzome/recipes/zmd.py
+++ b/pyzome/recipes/zmd.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
+
 import logging
-from typing import Iterable
+from typing import Optional, Sequence
 
 import xarray as xr
 
@@ -8,7 +9,8 @@
 from ..zonal_waves import zonal_wave_coeffs, zonal_wave_covariance
 from ..checks import infer_xr_coord_names
 
-logging.basicConfig(level=logging.NOTSET)
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
 
 LONG_NAMES = {
     "u": "zonal mean zonal wind",
@@ -53,9 +55,8 @@
 
 def create_zonal_mean_dataset(
     ds: xr.Dataset,
-    verbose: bool = False,
     include_waves: bool = False,
-    waves: None | Iterable[int] = None,
+    waves: Optional[Sequence[int]] = None,
     fftpkg: str = "scipy",
     lon_coord: str = "",
 ):
@@ -67,7 +68,7 @@ def create_zonal_mean_dataset(
 
     Parameters
     ----------
-    ds : `xarray.Dataset`
+    ds : ``xarray.Dataset``
         Dataset containing full fields (i.e., containing latitude &
         longitude dimensions) of basic state variables. This function
         currently assumes specific names and units:
@@ -85,9 +86,6 @@ def create_zonal_mean_dataset(
         Note that ds need not contain all of these variables, and this
         function will still provide as many diagnostics as possible.
 
-    verbose : bool, optional
-        Whether to print out progress information as the function proceeds.
-        Defaults to False.
     include_waves : bool, optional
         Whether to include diagnostics as a function of zonal wavenumber such
         as eddy covariances and fourier coefficients. Defaults to False.
@@ -96,21 +94,21 @@ def create_zonal_mean_dataset(
         kwarg is only considered if include_waves is True.
     fftpkg : string, optional
         String that specifies how to perform the FFT on the data. Options are
-        'scipy' or 'xrft'. Specifying scipy uses some operations that are memory-eager
-        and leverages scipy.fft.rfft. Specifying xrft should leverage the benefits
-        of xarray/dask for large datasets by using xrft.fft. Defaults to scipy.
-        This kwarg is only considered if include_waves is True.
+        'scipy' or 'xrft'. Defaults to scipy. This kwarg is only considered
+        if include_waves is True.
 
     Returns
     -------
-    `xarray.Dataset`
+    ``xarray.Dataset``
         An xarray Dataset containing the possible zonal mean diagnostics.
 
     Notes
     -----
+    This function writes progress messages to a module level logger
+
     Please see https://essd.copernicus.org/articles/10/1925/2018/ for
     a description of a different zonal mean dataset compiled for the
-    SPARC Reanalysis Intercomparison Project. This function does *not*
+    SPARC Reanalysis Intercomparison Project. This function does `not`
     provide all the same diagnostics as listed in that publication.
     However, if this function is provided with all of u, v, w, and T,
     it will return all the terms necessary to compute further diagnostics
@@ -118,9 +116,6 @@ def create_zonal_mean_dataset(
 
     """
 
-    logger = logging.getLogger("create_zonal_mean_dataset")
-    logger.setLevel(logging.INFO if verbose else logging.WARNING)
-
     if lon_coord == "":
         coords = infer_xr_coord_names(ds, required=["lon"])
         lon_coord = coords["lon"]
@@ -137,6 +132,8 @@ def create_zonal_mean_dataset(
         for v1, v2 in valid_covs
         if v1 in vars_available and v2 in vars_available
     ]
+    logger.info(f" *** Provided with {vars_available}")
+    logger.info(f" *** Will compute {cov_pairs}")
 
     inter = {}
     out_coords = None
@@ -183,12 +180,15 @@ def create_zonal_mean_dataset(
             inter.pop(var)
 
     # add attributes
+    logger.info(" *** Adding DataArray attributes")
     out_vars = inter.keys()
     for var in out_vars:
         inter[var].name = var
         inter[var].attrs["long_name"] = LONG_NAMES[var]
         inter[var].attrs["units"] = UNITS[var]
 
+    # create dataset
+    logger.info(" *** Compiling final dataset")
     out_ds = xr.Dataset(inter, coords=out_coords)
     out_ds.attrs["nlons"] = ds[lon_coord].size
     out_ds.attrs["lon0"] = ds[lon_coord].values[0]
diff --git a/pyzome/tem.py b/pyzome/tem.py
index 18f4e4d..8cc4ec3 100644
--- a/pyzome/tem.py
+++ b/pyzome/tem.py
@@ -27,13 +27,13 @@ def resid_vel(
 
     Parameters
     ----------
-    v : `xarray.DataArray`
+    v : ``xarray.DataArray``
         data containing the zonal mean of the meridional wind
-    w : `xarray.DataArray`
+    w : ``xarray.DataArray``
         data containing the zonal mean of the vertical *pressure* velocity
-    T : `xarray.DataArray`
+    T : ``xarray.DataArray``
         data containing the zonal mean of the air temperature
-    vT : `xarray.DataArray`
+    vT : ``xarray.DataArray``
         data containing the zonal mean meridional eddy heat flux, nominally
         defined as zonal_mean((v - zonal_mean(v))*(T - zonal_mean(T)))
     lat_coord : str, optional
@@ -58,7 +58,7 @@ def resid_vel(
 
     Returns
     -------
-    residual velocities: tuple of two `xarray.DataArray`
+    residual velocities: tuple of two ``xarray.DataArray``
         The meridional and vertical residual velocity components.
 
     Notes
@@ -74,9 +74,9 @@ def resid_vel(
 
     """
 
-    coords = infer_xr_coord_names(v, required=["lat", "lev"])
+    coords = infer_xr_coord_names(v, required=["lat", "plev"])
     lat = coords["lat"]
-    lev = coords["lev"]
+    lev = coords["plev"]
 
     check_var_SI_units(v[lev], "pressure", enforce=True)
     check_var_SI_units(v, "wind", enforce=True)
@@ -119,15 +119,15 @@ def epflux_vector(
 
     Parameters
     ----------
-    u : `xarray.DataArray`
+    u : ``xarray.DataArray``
         data containing the zonal mean of the zonal wind
-    T : `xarray.DataArray`
+    T : ``xarray.DataArray``
         data containing the zonal mean of the air temperature
-    uv : `xarray.DataArray`
+    uv : ``xarray.DataArray``
         data containing the zonal mean meridional momentum flux
-    vT : `xarray.DataArray`
+    vT : ``xarray.DataArray``
         data containing the zonal mean meridional heat flux
-    uw : `xarray.DataArray`
+    uw : ``xarray.DataArray``
         data containing the zonal mean vertical momentum flux
         (consistent with w, the vertical pressure velocity field)
     p0 : float, optional
@@ -146,7 +146,7 @@ def epflux_vector(
 
     Returns
     -------
-    ep_flux: tuple of two `xarray.DataArray`
+    ep_flux: tuple of two ``xarray.DataArray``
         The meridional and vertical components of the EP-Flux (F_lat, F_prs)
 
     Notes
@@ -162,9 +162,9 @@ def epflux_vector(
 
     """
 
-    coords = infer_xr_coord_names(u, required=["lat", "lev"])
+    coords = infer_xr_coord_names(u, required=["lat", "plev"])
     lat = coords["lat"]
-    lev = coords["lev"]
+    lev = coords["plev"]
 
     check_var_SI_units(u[lev], "pressure", enforce=True)
     check_var_SI_units(T, "temperature", enforce=True)
@@ -176,23 +176,12 @@ def epflux_vector(
     cos_lats = np.cos(np.deg2rad(u[lat]))
     to_theta = (p0 / u[lev]) ** (Rs / Cp)
 
-    Tht = T * to_theta
-    vTht = vT * to_theta
-
-    dTht_dp = Tht.differentiate(lev, edge_order=2)
+    psi = (vT * to_theta) / (T * to_theta).differentiate(lev, edge_order=2)
     du_dp = u.differentiate(lev, edge_order=2)
     ducos_dphi = (180.0 / np.pi) * (u * cos_lats).differentiate(lat, edge_order=2)
 
-    F_lat = a * cos_lats * ((vTht / dTht_dp) * du_dp - uv)
-    F_prs = (
-        a
-        * cos_lats
-        * (
-            ((-vTht / dTht_dp) * (1 / (a * cos_lats)) * ducos_dphi)
-            + (f * vTht / dTht_dp)
-            - uw
-        )
-    )
+    F_lat = a * cos_lats * (psi * du_dp - uv)
+    F_prs = a * cos_lats * (psi * (f - (ducos_dphi / (a * cos_lats))) - uw)
 
     F_lat.attrs["units"] = "m+3 s-2"  # type: ignore
     F_prs.attrs["units"] = "Pa m+2 s-2"  # type: ignore
@@ -215,11 +204,11 @@ def qg_epflux_vector(
 
     Parameters
     ----------
-    T : `xarray.DataArray`
+    T : ``xarray.DataArray``
         data containing the full field air temperature
-    uv : `xarray.DataArray`
+    uv : ``xarray.DataArray``
         data containing the zonal mean meridional momentum flux
-    vT : `xarray.DataArray`
+    vT : ``xarray.DataArray``
         data containing the zonal mean meridional heat flux
     p0 : float, optional
         Reference pressure for computation of potential temperature. Defaults
@@ -237,7 +226,7 @@ def qg_epflux_vector(
 
     Returns
     -------
-    qg_ep_flux: tuple of two `xarray.DataArray`
+    qg_ep_flux: tuple of two ``xarray.DataArray``
         The meridional and vertical components of the quasi-geostrophic EP-Flux
         (F_lat, F_prs).
 
@@ -254,9 +243,9 @@ def qg_epflux_vector(
 
     """
 
-    coords = infer_xr_coord_names(T, required=["lat", "lev"])
+    coords = infer_xr_coord_names(T, required=["lat", "plev"])
     lat = coords["lat"]
-    lev = coords["lev"]
+    lev = coords["plev"]
 
     check_var_SI_units(T[lev], "pressure", enforce=True)
     check_var_SI_units(T, "temperature", enforce=True)
@@ -267,13 +256,10 @@ def qg_epflux_vector(
     cos_lats = np.cos(np.deg2rad(T[lat]))
     to_theta = (p0 / T[lev]) ** (Rs / Cp)
 
-    Tht = T * to_theta
-    vTht = vT * to_theta
-
-    dTht_dp = Tht.differentiate(lev, edge_order=2)
+    psi = (vT * to_theta) / (T * to_theta).differentiate(lev, edge_order=2)
 
     F_lat = -a * cos_lats * uv
-    F_prs = a * cos_lats * f * (vTht / dTht_dp)
+    F_prs = a * cos_lats * f * psi
 
     F_lat.attrs["units"] = "m+3 s-2"  # type: ignore
     F_prs.attrs["units"] = "Pa m+2 s-2"  # type: ignore
@@ -293,9 +279,9 @@ def epflux_div(
 
     Parameters
     ----------
-    F_lat : `xarray.DataArray`
+    F_lat : ``xarray.DataArray``
         data containing the meridional EP-Flux component
-    F_prs : `xarray.DataArray`
+    F_prs : ``xarray.DataArray``
         data containing the vertical EP-Flux component
     accel : bool, optional
         If True, will scale the output by 1 / (a*cos(lat)) so
@@ -310,9 +296,10 @@ def epflux_div(
 
     Returns
     -------
-    epflux_divergence: `xarray.DataArray` or tuple of (`xarray.DataArray`, `xarray.DataArray`)
-        The total EP-Flux divergence or
-        the individual terms from the meridional and vertical divergence (if terms=True)
+    epflux_divergence: ``xarray.DataArray`` or tuple of two ``xr.DataArray``
+        The total EP-Flux divergence or a tuple containing the
+        the individual terms from the meridional and vertical
+        divergence (if terms=True)
 
     See Also
     --------
@@ -325,9 +312,9 @@ def epflux_div(
 
     """
 
-    coords = infer_xr_coord_names(F_lat, required=["lat", "lev"])
+    coords = infer_xr_coord_names(F_lat, required=["lat", "plev"])
     lat = coords["lat"]
-    lev = coords["lev"]
+    lev = coords["plev"]
 
     check_var_SI_units(F_prs[lev], "pressure", enforce=True)
 
diff --git a/pyzome/zonal_waves.py b/pyzome/zonal_waves.py
index e6984ef..e97f1d3 100644
--- a/pyzome/zonal_waves.py
+++ b/pyzome/zonal_waves.py
@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import warnings
-from typing import Iterable
+from typing import Optional, Sequence
 
-import numpy as np
-import xarray as xr
 import scipy
 import xrft
+import numpy as np
+import xarray as xr
+import dask.array.fft as daskfft
 
 from .checks import has_global_regular_lons, infer_xr_coord_names
 from .mock_data import lon_coord
@@ -15,7 +16,7 @@
 def zonal_wave_coeffs(
     dat: xr.DataArray,
     *,
-    waves: None | Iterable[int] = None,
+    waves: Optional[Sequence[int]] = None,
     fftpkg: str = "scipy",
     lon_coord: str = "",
 ) -> xr.DataArray:
@@ -25,7 +26,7 @@ def zonal_wave_coeffs(
 
     Parameters
     ----------
-    dat : `xarray.DataArray`
+    dat : ``xarray.DataArray``
         input data containing a longitude dimension that spans all 360 degrees
     waves : array-like, optional
         The zonal wavenumbers to maintain in the output. Defaults to None
@@ -43,7 +44,7 @@ def zonal_wave_coeffs(
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
         Output of an rFFT along the longitude dimension, for specified
         waves only. The output is a complex-valued DataArray containing
         at least the `zonal_wavenum` dimension, with attributes `nlons`
@@ -52,7 +53,7 @@ def zonal_wave_coeffs(
 
     """
 
-    if fftpkg not in ["scipy", "xrft"]:
+    if fftpkg not in {"scipy", "xrft"}:
         msg = "fftpkg must be 'scipy' or 'xrft'"
         raise ValueError(msg)
 
@@ -79,29 +80,27 @@ def _zonal_wave_coeffs_scipy(dat: xr.DataArray, lon_coord: str) -> xr.DataArray:
 
     Parameters
     ----------
-    dat : `xarray.DataArray`
+    dat : ``xarray.DataArray``
         data containing a dimension named longitude that spans all 360 degrees
     lon_coord : string
         name of the dimension/coordinate corresponding to the longitudes
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
         Output of the rFFT along the longitude dimension.
 
     """
     nlons = dat[lon_coord].size
-    lon_ax = dat.get_axis_num(lon_coord)
-
-    new_dims = list(dat.dims)
-    new_dims[lon_ax] = "zonal_wavenum"  # type: ignore
-
-    new_coords = dict(dat.coords)
-    new_coords.pop(lon_coord)
-    new_coords["zonal_wavenum"] = np.arange(0, nlons // 2 + 1)
 
-    fc = scipy.fft.rfft(dat.values, axis=lon_ax)
-    fc = xr.DataArray(fc, coords=new_coords, dims=new_dims)
+    scipy_rfft = daskfft.fft_wrap(scipy.fft.rfft, kind="rfft")
+    fc = xr.apply_ufunc(
+        scipy_rfft,
+        dat.chunk({lon_coord: -1}),
+        input_core_dims=[[lon_coord]],
+        output_core_dims=[["zonal_wavenum"]],
+        dask="allowed",
+    ).assign_coords({"zonal_wavenum": np.arange(0, nlons // 2 + 1)})
 
     return fc
 
@@ -113,14 +112,14 @@ def _zonal_wave_coeffs_xrft(dat: xr.DataArray, lon_coord: str) -> xr.DataArray:
 
     Parameters
     ----------
-    dat : `xarray.DataArray`
+    dat : ``xarray.DataArray``
         data containing a dimension named longitude that spans all 360 degrees
     lon_coord : string
         name of the dimension/coordinate corresponding to the longitudes
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
         Output of the rFFT along the longitude dimension.
 
     """
@@ -150,7 +149,7 @@ def inflate_zonal_wave_coeffs(
 
     Parameters
     ----------
-    fc_subset : `xarray.DataArray`
+    fc_subset : ``xarray.DataArray``
         Fourier coefficients as a function of zonal wavenumber, as
         returned by `zonal_wave_coeffs`.
     wave_coord : str, optional
@@ -160,7 +159,7 @@ def inflate_zonal_wave_coeffs(
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
         Fourier coefficients as a function of zonal wavenumber, with the
         spectrum size determined by the "nlons" attribute of the input
         DataArray.
@@ -171,14 +170,20 @@ def inflate_zonal_wave_coeffs(
         wave_coord = coords["zonal_wavenum"]
 
     if "nlons" not in fc_subset.attrs:
-        msg = "input DataArray must have an 'nlons' attribute specifying the number of longitudes in the source data"
+        msg = (
+            "input DataArray must have an 'nlons' attribute specifying "
+            + "the number of longitudes in the source data"
+        )
         raise KeyError(msg)
 
     expected_wavenums = np.arange(fc_subset.attrs["nlons"] // 2 + 1)
     if np.array_equiv(fc_subset[wave_coord], expected_wavenums):
         return fc_subset
     elif np.any(~np.in1d(fc_subset[wave_coord], expected_wavenums, assume_unique=True)):
-        msg = "input DataArray wavenumbers are not a subset of the expected wavenumbers based on the 'nlons' attribute"
+        msg = (
+            "input DataArray wavenumbers are not a subset of the expected "
+            + "wavenumbers based on the 'nlons' attribute"
+        )
         raise ValueError(msg)
 
     return fc_subset.reindex({wave_coord: expected_wavenums}, fill_value=0j)  # type: ignore
@@ -194,7 +199,7 @@ def zonal_wave_ampl_phase(
 
     Parameters
     ----------
-    fc : `xarray.DataArray`
+    fc : ``xarray.DataArray``
         Fourier coefficients as a function of zonal wavenumber, as
         returned by `zonal_wave_coeffs`.
     phase_deg : boolean, optional
@@ -206,7 +211,7 @@ def zonal_wave_ampl_phase(
 
     Returns
     -------
-    tuple of two `xarray.DataArray`
+    tuple of two ``xarray.DataArray``
         Contains the amplitudes and phases for the input field zonal waves.
 
     See Also
@@ -220,7 +225,10 @@ def zonal_wave_ampl_phase(
         wave_coord = coords["zonal_wavenum"]
 
     if "nlons" not in fc.attrs:
-        msg = "input DataArray must have an 'nlons' attribute specifying the number of longitudes in the source data"
+        msg = (
+            "input DataArray must have an 'nlons' attribute specifying "
+            + "the number of longitudes in the source data"
+        )
         raise KeyError(msg)
 
     # where the longitudinal wavenumber is 0, `where' will
@@ -236,18 +244,17 @@ def zonal_wave_ampl_phase(
 
     with xr.set_options(keep_attrs=True):
         ampl = np.absolute(fc) * mult_mask / fc.nlons
-        # phas = np.angle(fc, deg=phase_deg) # returns a np.ndarray instead of xr.DataArray/Dataset
-        phas = xr.apply_ufunc(np.angle, fc, kwargs={"deg": phase_deg})
+        phas = xr.apply_ufunc(np.angle, fc, kwargs={"deg": phase_deg}, dask="allowed")
 
     return (ampl.astype(fc.dtype), phas.astype(fc.dtype))
 
 
 def filter_by_zonal_wave_truncation(
     fc: xr.DataArray,
-    waves: Iterable[int],
+    waves: Sequence[int],
     fftpkg: str = "scipy",
     wave_coord: str = "",
-    lons: None | xr.DataArray = None,
+    lons: Optional[xr.DataArray] = None,
 ) -> xr.DataArray:
     """Filters a field by truncating the zonal wavenumbers. This is done
     by taking an inverse rFFT of the Fourier coefficients, with the unwanted
@@ -255,7 +262,7 @@ def filter_by_zonal_wave_truncation(
 
     Parameters
     ----------
-    fc : `xarray.DataArray`
+    fc : ``xarray.DataArray``
         Fourier coefficients as a function of zonal wavenumber, as
         returned by `zonal_wave_coeffs`.
     waves : iterable of int or slice
@@ -268,25 +275,39 @@ def filter_by_zonal_wave_truncation(
         The coordinate name of the wavenumber dimension. If given an empty
         string (the default), the function tries to infer which coordinate
         corresponds to the wavenumbers
-    lons: `xarray.DataArray`, optional
+    lons: ``xarray.DataArray``, optional
         The longitude coordinate of the input field. If not given, the
         function will attempt to reconstruct the coordinate from the input
         Fourier coefficients, by looking for 'nlons' and 'lon0' attributes.
 
+    Returns
+    -------
+    ``xarray.DataArray``
+        The field filtered to retain only the given zonal wavenumbers.
+
     """
+    if fftpkg not in {"scipy", "xrft"}:
+        msg = "fftpkg must be 'scipy' or 'xrft'"
+        raise ValueError(msg)
+
     if wave_coord == "":
         coords = infer_xr_coord_names(fc, required=["zonal_wavenum"])
         wave_coord = coords["zonal_wavenum"]
 
     if "nlons" not in fc.attrs:
-        msg = "input DataArray must have an 'nlons' attribute specifying the number of longitudes in the source data"
+        msg = (
+            "input DataArray must have an 'nlons' attribute specifying "
+            + "the number of longitudes in the source data"
+        )
         raise KeyError(msg)
 
     if "lon0" not in fc.attrs:
-        msg = "input DataArray must have a 'lon0' attribute specifying the starting longitude of the source data"
+        msg = (
+            "input DataArray must have a 'lon0' attribute specifying the "
+            + "starting longitude of the source data"
+        )
         raise KeyError(msg)
 
-    wavenum_ax = fc.get_axis_num(wave_coord)
     if lons is None:
         warnings.warn("attempting to infer the input longitude coordinate", UserWarning)
         lons = lon_coord(
@@ -295,44 +316,75 @@ def filter_by_zonal_wave_truncation(
             right_lim=fc.attrs["lon0"] + 360.0,
         )
 
-    if fftpkg == "scipy":
-        new_dims = list(fc.dims)
-        new_dims[wavenum_ax] = lons.name  # type: ignore
-        new_coords = dict(fc.coords)
-        new_coords.pop(wave_coord, None)
-        new_coords[lons.name] = lons
+    funcs = {"scipy": _inverse_transform_scipy, "xrft": _inverse_transform_xrft}
+    fc_limited = fc.where(fc[wave_coord].isin(waves), other=0j)
 
-        filtered = scipy.fft.irfft(
-            fc.where(fc[wave_coord].isin(waves), other=0j).values, axis=wavenum_ax
-        )
-        filtered = xr.DataArray(filtered, dims=new_dims, coords=new_coords)
-
-    elif fftpkg == "xrft":
-        filtered = xrft.ifft(
-            fc.where(fc[wave_coord].isin(waves), other=0j),
-            dim=wave_coord,
-            real_dim=wave_coord,
-            true_amplitude=False,
-            true_phase=False,
-            prefix="",
-        )
-        filtered = filtered.rename({wave_coord: lons.name}).assign_coords(
-            {lons.name: lons}
-        )
+    return funcs[fftpkg](fc_limited, wave_coord, lons)
 
-    else:
-        msg = "Invalid fftpkg: must be 'scipy' or 'xrft'"
-        raise ValueError(msg)
 
-    return filtered
+def _inverse_transform_scipy(fc: xr.DataArray, wave_coord: str, lons: xr.DataArray):
+    r"""Calculates the inverse rFFT of the Zonal Fourier coefficients.
+    Uses scipy.fft.irfft to perform the calculation. This function is
+    intended to be used internally by filter_by_zonal_wave_truncation.
+
+    Parameters
+    ----------
+    fc : ``xarray.DataArray``
+        Fourier coefficients as a function of zonal wavenumber, as
+        returned by `zonal_wave_coeffs`.
+    wave_coord : string
+        name of the dimension/coordinate corresponding to the wavenumbers
+    lons: ``xarray.DataArray``
+        The longitude coordinate of the original data prior to FFT
+
+    """
+    scipy_irfft = daskfft.fft_wrap(scipy.fft.irfft, kind="irfft")
+    recons = xr.apply_ufunc(
+        scipy_irfft,
+        fc.chunk({wave_coord: -1}),
+        input_core_dims=[[wave_coord]],
+        output_core_dims=[[lons.name]],
+        dask="allowed",
+    ).assign_coords({lons.name: lons})
+
+    return recons
+
+
+def _inverse_transform_xrft(fc: xr.DataArray, wave_coord: str, lons: xr.DataArray):
+    r"""Calculates the inverse rFFT of the Zonal Fourier coefficients.
+    Uses xrft.ifft to perform the calculation. This function is
+    intended to be used internally by filter_by_zonal_wave_truncation.
+
+    Parameters
+    ----------
+    fc : ``xarray.DataArray``
+        Fourier coefficients as a function of zonal wavenumber, as
+        returned by `zonal_wave_coeffs`.
+    wave_coord : string
+        name of the dimension/coordinate corresponding to the wavenumbers
+    lons: ``xarray.DataArray``
+        The longitude coordinate of the original data prior to FFT
+
+    """
+    recons = xrft.ifft(
+        fc,
+        dim=wave_coord,
+        real_dim=wave_coord,
+        true_amplitude=False,
+        true_phase=False,
+        prefix="",
+    )
+    recons = recons.rename({wave_coord: lons.name}).assign_coords({lons.name: lons})
+
+    return recons
 
 
 def zonal_wave_contributions(
     fc: xr.DataArray,
-    waves: Iterable[int],
+    waves: Sequence[int],
     fftpkg: str = "scipy",
     wave_coord: str = "",
-    lons: None | xr.DataArray = None,
+    lons: Optional[xr.DataArray] = None,
 ) -> xr.DataArray:
     r"""Computes the individual contributions of each zonal wavenumber
     to the input field. This is done by taking an inverse FFT of the
@@ -340,7 +392,7 @@ def zonal_wave_contributions(
 
     Parameters
     ----------
-    fc : `xarray.DataArray`
+    fc : ``xarray.DataArray``
         Fourier coefficients as a function of zonal wavenumber, as
         returned by `zonal_wave_coeffs`.
     waves : iterable of int, optional
@@ -355,14 +407,14 @@ def zonal_wave_contributions(
         The coordinate name of the wavenumber dimension. If given an empty
         string (the default), the function tries to infer which coordinate
         corresponds to the wavenumbers
-    lons: `xarray.DataArray`, optional
+    lons: ``xarray.DataArray``, optional
         The longitude coordinate of the input field. If not given, the
         function will attempt to reconstruct the coordinate from the input
         Fourier coefficients, by looking for 'nlons' and 'lon0' attributes.
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
         The individual contributions of each wavenumber to the original
         input field.
 
@@ -377,16 +429,15 @@ def zonal_wave_contributions(
         coords = infer_xr_coord_names(fc, required=["zonal_wavenum"])
         wave_coord = coords["zonal_wavenum"]
 
-    contributions = []
-    for wave in waves:
-        contributions.append(
+    contributions = xr.concat(
+        [
             filter_by_zonal_wave_truncation(
                 fc, [wave], fftpkg=fftpkg, wave_coord=wave_coord, lons=lons
             )
-        )
-    contributions = xr.concat(contributions, dim=wave_coord).assign_coords(
-        {wave_coord: waves}
-    )
+            for wave in waves
+        ],
+        dim=wave_coord,
+    ).assign_coords({wave_coord: waves})
 
     return contributions
 
@@ -395,19 +446,19 @@ def zonal_wave_covariance(
     fc1: xr.DataArray,
     fc2: xr.DataArray,
     wave_coord: str = "",
-    nlons: int | None = None,
+    nlons: Optional[int] = None,
 ) -> xr.DataArray:
     r"""Calculates the covariance of two fields partititioned into
     zonal wavenumbers.
 
     Parameters
     ----------
-    fc1 : `xarray.DataArray`
+    fc1 : ``xarray.DataArray``
         The zonal Fourier coefficients of the first field, as returned by
-        `zonal_wave_coeffs`.
-    fc2 : `xarray.DataArray`
-        The zonal Fourier coefficients of the second field, as returned by
-        `zonal_wave_coeffs`.
+        ``zonal_wave_coeffs``.
+    fc2 : ``xarray.DataArray``
+        The zonal Fourier c`oefficients of the second field, as returned by
+        ``zonal_wave_coeffs``.
     wave_coord : str optional
         The coordinate name of the wavenumber dimension. If given an empty
         string (the default), the function tries to infer which coordinate
@@ -418,7 +469,7 @@ def zonal_wave_covariance(
 
     Returns
     -------
-    `xarray.DataArray`
+    ``xarray.DataArray``
 
     See Also
     --------
diff --git a/requirements.txt b/requirements.txt
index cff5831..f1efcd4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,5 @@
 numpy
 scipy
 xarray
+dask
 xrft
\ No newline at end of file


Test patch:
diff --git a/tests/test_accessor.py b/tests/test_accessor.py
new file mode 100644
index 0000000..afba3e4
--- /dev/null
+++ b/tests/test_accessor.py
@@ -0,0 +1,107 @@
+import pytest
+import numpy as np
+
+from pyzome.accessor import PyzomeDataArrayAccessor, PyzomeDatasetAccessor
+from pyzome.checks import check_for_logp_coord
+from pyzome.mock_data import (
+    create_dummy_geo_field,
+    create_dummy_geo_dataset,
+    lon_coord,
+    lat_coord,
+    plev_coord,
+)
+
+
+def test_pzm_accessor_coord_map():
+    """Test that the internal coordinate map works correctly"""
+    da = create_dummy_geo_field(
+        lon_coord(10, name="longitude"), lat_coord(10), plev_coord(3, name="pressures")
+    )
+
+    assert da.pzm.coord_map("lon") == "longitude"
+    assert da.pzm.coord_map("lat") == "lat"
+    assert da.pzm.coord_map("plev") == "pressures"
+
+    with pytest.raises(KeyError) as e:
+        da.pzm.coord_map("foo")
+    assert "pyzome coordinate 'foo' not identified in original data" in str(e.value)
+
+
+def test_pzm_accessor_coord_props():
+    """Test that the coordinate properties lead to the right references"""
+    da = create_dummy_geo_field(
+        lon_coord(10, name="lons"),
+        lat_coord(10, name="lats"),
+        plev_coord(3, name="pres"),
+    )
+    da = da.assign_coords({"zonal_wavenum": 1}).expand_dims("zonal_wavenum")
+
+    assert da["lons"].equals(da.pzm.lon)
+    assert da["lats"].equals(da.pzm.lat)
+    assert da["pres"].equals(da.pzm.plev)
+    assert da["zonal_wavenum"].equals(da.pzm.zonal_wavenum)
+
+
+def test_pzm_dataset_accessor():
+    ds = create_dummy_geo_dataset(
+        field_names=["u", "v", "T"],
+        lons=lon_coord(2.5),
+        lats=lat_coord(2.5),
+        levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
+    )
+    assert isinstance(ds.pzm, PyzomeDatasetAccessor)
+
+    # zonal mean should remove lon coord
+    ds_zm = ds.pzm.zonal_mean(strict=True)
+    with pytest.raises(KeyError):
+        ds_zm.pzm.coord_map("lon")
+
+    # meridional mean should remove lat coord
+    ds_zm_mm = ds_zm.pzm.meridional_mean(-90, 90, strict=True)
+    with pytest.raises(KeyError):
+        ds_zm_mm.pzm.coord_map("lat")
+
+    # check that logp altitude coord gets added
+    assert check_for_logp_coord(ds.pzm.add_logp_altitude())
+
+
+def test_pzm_dataarray_accessor():
+    da = create_dummy_geo_field(
+        lon_coord(2.5),
+        lat_coord(2.5),
+        plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
+    )
+    assert isinstance(da.pzm, PyzomeDataArrayAccessor)
+
+    # zonal mean should remove lon coord
+    da_zm = da.pzm.zonal_mean(strict=True)
+    with pytest.raises(KeyError):
+        da_zm.pzm.coord_map("lon")
+
+    # meridional mean should remove lat coord
+    da_zm_mm = da_zm.pzm.meridional_mean(-90, 90, strict=True)
+    with pytest.raises(KeyError):
+        da_zm_mm.pzm.coord_map("lat")
+
+    # check that logp altitude coord gets added
+    assert check_for_logp_coord(da_zm_mm.pzm.add_logp_altitude())
+
+    # zonal wave coeffs adds zonal_wavenum and removes lon
+    fc = da.pzm.zonal_wave_coeffs(waves=[1, 2, 3])
+    assert fc.pzm.coord_map("zonal_wavenum") == "zonal_wavenum"
+    assert np.allclose(fc.pzm.zonal_wavenum.values, [1, 2, 3])
+    with pytest.raises(KeyError):
+        fc.pzm.coord_map("lon")
+
+    # inflate zonal wave coeffs and reconstruct; should add lons back
+    # and remove zonal_wavenums
+    fc_inflated = fc.pzm.inflate_zonal_wave_coeffs()
+    recons = fc_inflated.pzm.filter_by_zonal_wave_truncation([1, 2, 3])
+    assert recons.pzm.lon.equals(da.pzm.lon)
+    with pytest.raises(KeyError):
+        recons.pzm.coord_map("zonal_wavenum")
+
+    # zonal wave contributions should add back lon but keep wavenums
+    contribs = fc_inflated.pzm.zonal_wave_contributions([1, 2, 3])
+    assert contribs.pzm.lon.equals(da.pzm.lon)
+    assert contribs.pzm.zonal_wavenum.equals(fc.pzm.zonal_wavenum)
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 18fa14f..c5a361e 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -2,7 +2,6 @@
 import numpy as np
 
 from pyzome.basic import zonal_mean, meridional_mean
-from pyzome.exceptions import AttrError, CoordinateError, LongitudeError, UnitError
 from pyzome.mock_data import create_dummy_geo_field, lon_coord, lat_coord, plev_coord
 
 
@@ -14,7 +13,7 @@ def test_zonal_mean():
     # Basic tests that the function works as expected
     assert np.allclose(da_zonal_mean.values, np.mean(da.values, axis=2))
     assert da_zonal_mean.shape == (10, 19)
-    assert da_zonal_mean.dims == ("lev", "lat")
+    assert da_zonal_mean.dims == ("plev", "lat")
 
 
 def test_meridional_mean():
@@ -28,20 +27,20 @@ def test_meridional_mean():
     )
     assert np.allclose(da_meridional_mean.values, global_mean)
     assert da_meridional_mean.shape == (10, 36)
-    assert da_meridional_mean.dims == ("lev", "lon")
+    assert da_meridional_mean.dims == ("plev", "lon")
 
 
 def test_meridional_mean_wrong_lat_order_fail():
     """Test that meridional mean fails with lats in wrong order"""
     da = create_dummy_geo_field(lon_coord(10), lat_coord(10), plev_coord(3))
     with pytest.raises(ValueError) as e:
-        da_meridional_mean = meridional_mean(da, 90, -90)
+        _ = meridional_mean(da, 90, -90)
     assert "lat1 must be less than lat2" in str(e.value)
 
 
-def test_meridional_mean_too_few_lats_fail():
-    """Test that meridional mean fails with too few lats"""
+def test_meridional_mean_strict_too_few_lats_fail():
+    """Test that meridional mean fails with too few lats in strict mode"""
     da = create_dummy_geo_field(lon_coord(10), lat_coord(10), plev_coord(3))
     with pytest.raises(ValueError) as e:
-        da_meridional_mean = meridional_mean(da.sel(lat=slice(30, 60)), 0, 90)
+        _ = meridional_mean(da.sel(lat=slice(30, 60)), 0, 90, strict=True)
     assert "data only contains lats in range of" in str(e.value)
diff --git a/tests/test_checks.py b/tests/test_checks.py
index f7b7810..b3e2cd8 100644
--- a/tests/test_checks.py
+++ b/tests/test_checks.py
@@ -51,17 +51,16 @@ def test_infer_xr_coord_names():
     """Test that standard coordinates can be inferred/detected"""
 
     da = create_dummy_geo_field(lon_coord(10), lat_coord(10), plev_coord(3))
-    mapper = {"lev": "lev", "lat": "lat", "lon": "lon"}
-    assert infer_xr_coord_names(da, required=["lev", "lat", "lon"]) == mapper
+    mapper = {"plev": "plev", "lat": "lat", "lon": "lon"}
+    assert infer_xr_coord_names(da, required=["plev", "lat", "lon"]) == mapper
 
     da = create_dummy_geo_field(
         lon_coord(10, name="longitude"),
         lat_coord(10, name="latitude"),
         plev_coord(3, name="level"),
     )
-    print(da)
-    mapper = {"lev": "level", "lat": "latitude", "lon": "longitude"}
-    assert infer_xr_coord_names(da, required=["lev", "lat", "lon"]) == mapper
+    mapper = {"plev": "level", "lat": "latitude", "lon": "longitude"}
+    assert infer_xr_coord_names(da, required=["plev", "lat", "lon"]) == mapper
 
 
 def test_infer_multiple_coords_fail():
@@ -117,17 +116,17 @@ def test_check_SI_units_wrong():
 def test_check_for_logp_coord():
     """Test that check_for_logp_coord properly detects a log-p altitude coordinate"""
     da = create_dummy_geo_field(
-        lon_coord(10), 
-        lat_coord(10), 
+        lon_coord(10),
+        lat_coord(10),
         plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
     )
 
-    assert check_for_logp_coord(da) is False 
+    assert check_for_logp_coord(da) is False
     with pytest.raises(CoordinateError) as e:
         check_for_logp_coord(da, enforce=True)
-    assert "z is not a coordinate in the data" in str(e.value) 
+    assert "z is not a coordinate in the data" in str(e.value)
 
-    da = da.assign_coords({"z": np.log(da.lev)})
+    da = da.assign_coords({"z": np.log(da.plev)})
     assert check_for_logp_coord(da) is False
     with pytest.raises(AttrError) as e:
         check_for_logp_coord(da, enforce=True)
@@ -155,6 +154,3 @@ def test_check_for_logp_coord():
 
     da.z.attrs["units"] = "m"
     assert check_for_logp_coord(da, enforce=True) is True
-
-
-    
\ No newline at end of file
diff --git a/tests/test_mock_data.py b/tests/test_mock_data.py
index 59ee8dc..041d105 100644
--- a/tests/test_mock_data.py
+++ b/tests/test_mock_data.py
@@ -138,7 +138,7 @@ def test_invalid_longitude_limits():
 def test_plev_default_values():
     """Test that plev_coord creates a pressure-like coordinate correctly for default values"""
     da = plev_coord(3)
-    assert da.name == "lev"
+    assert da.name == "plev"
     np.testing.assert_allclose(da.values, np.logspace(3, 0, 3 * 3 + 1))
     assert da.attrs["standard_name"] == "pressure"
     assert da.attrs["units"] == "hPa"
@@ -220,12 +220,12 @@ def test_create_dummy_geo_field():
     )
 
     # Check the dimensions
-    assert da.dims == ("time", "lev", "lat", "lon")
+    assert da.dims == ("time", "plev", "lat", "lon")
 
     # Check the sizes
     assert da.sizes == {
         "time": times.size,
-        "lev": levs.size,
+        "plev": levs.size,
         "lat": lats.size,
         "lon": lons.size,
     }
@@ -262,10 +262,10 @@ def test_create_dummy_geo_field_invalid_arguments():
     lats = lat_coord(10)
 
     with pytest.raises(AttributeError):
-        create_dummy_geo_field(invalid_dataarray, lats)
+        create_dummy_geo_field(invalid_dataarray, lats)  # type: ignore
 
     with pytest.raises(AttributeError):
-        create_dummy_geo_field(lons, invalid_dataarray)
+        create_dummy_geo_field(lons, invalid_dataarray)  # type: ignore
 
 
 def test_create_dummy_geo_dataset():
@@ -274,10 +274,10 @@ def test_create_dummy_geo_dataset():
     lats = lat_coord(10)
     levs = plev_coord(5)
     times = time_coord()
-    field_names = ["temp", "precip", "wind"]
+    field_names = ["temp", "humidity", "wind"]
     field_attrs = {
         "temp": {"units": "K"},
-        "precip": {"units": "mm/hr"},
+        "humidity": {"units": "kg / kg"},
         "wind": {"units": "m/s"},
     }
 
@@ -291,12 +291,12 @@ def test_create_dummy_geo_dataset():
         assert ds[name].attrs == field_attrs[name]
 
         # Check the dimensions
-        assert ds[name].dims == ("time", "lev", "lat", "lon")
+        assert ds[name].dims == ("time", "plev", "lat", "lon")
 
         # Check the sizes
         assert ds[name].sizes == {
             "time": times.size,
-            "lev": levs.size,
+            "plev": levs.size,
             "lat": lats.size,
             "lon": lons.size,
         }
@@ -310,10 +310,10 @@ def test_create_dummy_geo_dataset_invalid_arguments():
     field_names = ["temp"]
 
     with pytest.raises(AttributeError):
-        create_dummy_geo_dataset(field_names, invalid_dataarray, lats)
+        create_dummy_geo_dataset(field_names, invalid_dataarray, lats)  # type: ignore
 
     with pytest.raises(AttributeError):
-        create_dummy_geo_dataset(field_names, lons, invalid_dataarray)
+        create_dummy_geo_dataset(field_names, lons, invalid_dataarray)  # type: ignore
 
     mismatched_field_attrs = {"temp": {"units": "K"}, "missing_field": {}}
     with pytest.raises(ValueError) as e:
diff --git a/tests/test_qglogp.py b/tests/test_qglogp.py
index f86e83f..6b77da7 100644
--- a/tests/test_qglogp.py
+++ b/tests/test_qglogp.py
@@ -10,13 +10,14 @@
     plumb_wave_activity_flux,
 )
 from pyzome.mock_data import (
-    create_dummy_geo_dataset, 
-    create_dummy_geo_field, 
-    lon_coord, 
-    lat_coord, 
-    plev_coord
+    create_dummy_geo_dataset,
+    create_dummy_geo_field,
+    lon_coord,
+    lat_coord,
+    plev_coord,
 )
 from pyzome.constants import SCALE_HEIGHT, PREF
+from pyzome.exceptions import CoordinateError
 
 
 def test_add_logp_altitude():
@@ -25,7 +26,7 @@ def test_add_logp_altitude():
         lons=lon_coord(2.5),
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
-        name = "T",
+        name="T",
         attrs={"units": "K"},
     )
 
@@ -33,7 +34,9 @@ def test_add_logp_altitude():
     assert da.z.attrs["units"] == "m"
     assert da.z.attrs["long_name"] == "log-pressure altitude"
     assert da.z.attrs["note"] == "added by pyzome"
-    np.testing.assert_allclose(da.z.values, -SCALE_HEIGHT*np.log(da.lev.values/PREF))
+    np.testing.assert_allclose(
+        da.z.values, -SCALE_HEIGHT * np.log(da.plev.values / PREF)
+    )
 
 
 def test_buoyancy_frequency_squared():
@@ -42,7 +45,7 @@ def test_buoyancy_frequency_squared():
         lons=None,
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
-        name = "T",
+        name="T",
         attrs={"units": "K"},
     )
     T = add_logp_altitude(T)
@@ -55,7 +58,7 @@ def test_buoyancy_frequency_squared():
 def test_merid_grad_qgpv():
     """Test that merid_grad_qgpv computation works from start to finish"""
     ds = create_dummy_geo_dataset(
-        field_names = ["u", "N2"],
+        field_names=["u", "N2"],
         lons=None,
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
@@ -72,29 +75,35 @@ def test_refractive_index():
     """Test that refractive_index computation works from start to finish"""
 
     ds = create_dummy_geo_dataset(
-        field_names = ["u", "q_phi", "N2"],
+        field_names=["u", "q_phi", "N2"],
         lons=None,
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
-        field_attrs={"u": {"units": "m s-1"}, "q_phi": {"units": "s-1"}, "N2": {"units": "s-2"}},
+        field_attrs={
+            "u": {"units": "m s-1"},
+            "q_phi": {"units": "s-1"},
+            "N2": {"units": "s-2"},
+        },
     )
     ds = add_logp_altitude(ds)
 
     ri2 = refractive_index(ds.u, ds.q_phi, np.abs(ds.N2), 1)
-    term1, term2, term3  = refractive_index(ds.u, ds.q_phi, np.abs(ds.N2), 1, terms=True)
+    term1, term2, term3 = refractive_index(ds.u, ds.q_phi, np.abs(ds.N2), 1, terms=True)
     np.testing.assert_allclose(ri2.values, (term1 + term2 + term3).values)
 
     # Test with alternate buoyancy frequency term
     ri2 = refractive_index(ds.u, ds.q_phi, np.abs(ds.N2), 1, cd_approx_term=True)
-    term1, term2, term3 = refractive_index(ds.u, ds.q_phi, np.abs(ds.N2), 1, terms=True, cd_approx_term=True)
+    term1, term2, term3 = refractive_index(
+        ds.u, ds.q_phi, np.abs(ds.N2), 1, terms=True, cd_approx_term=True
+    )
     np.testing.assert_allclose(ri2.values, (term1 + term2 + term3).values)
 
 
 def test_plumb_wave_activity_flux():
     """Test that plumb_wave_activity_flux computation works from start to finish"""
-    
+
     ds = create_dummy_geo_dataset(
-        field_names = ["psi", "N2"],
+        field_names=["psi", "N2"],
         lons=lon_coord(2.5),
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
@@ -102,7 +111,7 @@ def test_plumb_wave_activity_flux():
     )
     ds = add_logp_altitude(ds)
 
-    combos = [["x","y","z"], ["x","z"], ["y","z"], ["x"], ["y"], ["z"]]
+    combos = [["x", "y", "z"], ["x", "z"], ["y", "z"], ["x"], ["y"], ["z"]]
     for components in combos:
         F = plumb_wave_activity_flux(ds.psi, ds.N2, components=components)
         assert len(F) == len(components)
@@ -110,7 +119,7 @@ def test_plumb_wave_activity_flux():
 
 def test_plumb_wave_activity_flux_no_lats():
     ds = create_dummy_geo_dataset(
-        field_names = ["psi", "N2"],
+        field_names=["psi", "N2"],
         lons=lon_coord(2.5),
         lats=None,
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
@@ -118,14 +127,14 @@ def test_plumb_wave_activity_flux_no_lats():
     )
     ds = add_logp_altitude(ds)
 
-    with pytest.raises(ValueError) as e:
-        F = plumb_wave_activity_flux(ds.psi, ds.N2, components=["y"])
-    assert "A latitude coordinate must be available" in str(e.value)
+    with pytest.raises(CoordinateError) as e:
+        _ = plumb_wave_activity_flux(ds.psi, ds.N2, components=["y"])
+    assert "Unable to match any of the coordinates in dat for lat" in str(e.value)
 
 
 def test_plumb_wave_activity_flux_no_lons():
     ds = create_dummy_geo_dataset(
-        field_names = ["psi", "N2"],
+        field_names=["psi", "N2"],
         lons=None,
         lats=lat_coord(2.5),
         levs=plev_coord(3, left_lim_exponent=5, right_lim_exponent=2, units="Pa"),
@@ -133,6 +142,6 @@ def test_plumb_wave_activity_flux_no_lons():
     )
     ds = add_logp_altitude(ds)
 
-    with pytest.raises(ValueError) as e:
-        F = plumb_wave_activity_flux(ds.psi, ds.N2, components=["x", "z"])
-    assert "A longitude coordinate must be available" in str(e.value)
+    with pytest.raises(CoordinateError) as e:
+        _ = plumb_wave_activity_flux(ds.psi, ds.N2, components=["x", "z"])
+    assert "Unable to match any of the coordinates in dat for lon" in str(e.value)
diff --git a/tests/test_recipes_zmd.py b/tests/test_recipes_zmd.py
index 39faaaf..72f497b 100644
--- a/tests/test_recipes_zmd.py
+++ b/tests/test_recipes_zmd.py
@@ -86,5 +86,5 @@ def test_create_zonal_mean_dataset_fails_no_valid_fields():
     )
 
     with pytest.raises(ValueError) as e:
-        zmds = create_zonal_mean_dataset(ds)
+        _ = create_zonal_mean_dataset(ds)
     assert "No valid fields found in provided dataset" in str(e.value)
diff --git a/tests/test_zonal_waves.py b/tests/test_zonal_waves.py
index 50ba1a0..1c1c473 100644
--- a/tests/test_zonal_waves.py
+++ b/tests/test_zonal_waves.py
@@ -73,7 +73,7 @@ def test_zonal_wave_coeffs_unknown_fftpkg():
     """Test that zonal wave coeffs fails with unknown fftpkg"""
     da = create_dummy_geo_field(lon_coord(10), lat_coord(10))
     with pytest.raises(ValueError) as e:
-        da_zonal_wave_coeffs = zonal_wave_coeffs(da, fftpkg="unknown")
+        _ = zonal_wave_coeffs(da, fftpkg="unknown")
     assert "fftpkg must be 'scipy' or 'xrft'" in str(e.value)
 
 
@@ -100,7 +100,7 @@ def test_inflate_zonal_wave_coeffs_missing_nlons():
     _ = da_fc.attrs.pop("nlons")
 
     with pytest.raises(KeyError) as e:
-        da_fc_inflated = inflate_zonal_wave_coeffs(da_fc)
+        _ = inflate_zonal_wave_coeffs(da_fc)
     assert "input DataArray must have an 'nlons' attribute" in str(e.value)
 
 
@@ -122,7 +122,7 @@ def test_inflate_zonal_wave_coeffs_unexpected_wavenums():
     da_fc = da_fc.assign_coords({"zonal_wavenum": [1, 2, 103, 104]})
 
     with pytest.raises(ValueError) as e:
-        da_fc_inflated = inflate_zonal_wave_coeffs(da_fc)
+        _ = inflate_zonal_wave_coeffs(da_fc)
     assert (
         "input DataArray wavenumbers are not a subset of the expected wavenumbers based on the 'nlons' attribute"
         in str(e.value)
@@ -160,7 +160,7 @@ def test_zonal_wave_ampl_phase_missing_nlons():
 
     with pytest.raises(KeyError) as e:
         _ = da_fc.attrs.pop("nlons")
-        ampl, phase = zonal_wave_ampl_phase(da_fc)
+        _, _ = zonal_wave_ampl_phase(da_fc)
     assert "input DataArray must have an 'nlons' attribute" in str(e.value)
 
 
@@ -246,12 +246,12 @@ def test_filter_by_zonal_wave_truncation_no_attrs():
 
     with pytest.raises(KeyError) as e:
         _ = da_fc.attrs.pop("lon0")
-        da_filt = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3])
+        _ = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3])
     assert "input DataArray must have a 'lon0' attribute" in str(e.value)
 
     with pytest.raises(KeyError) as e:
         _ = da_fc.attrs.pop("nlons")
-        da_filt = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3])
+        _ = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3])
     assert "input DataArray must have an 'nlons' attribute" in str(e.value)
 
 
@@ -276,8 +276,8 @@ def test_filter_by_zonal_wave_truncation_fftpkg_fails():
     da_fc = zonal_wave_coeffs(da, fftpkg="scipy")
 
     with pytest.raises(ValueError) as e:
-        da_filt = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3], fftpkg="foo")
-    assert "Invalid fftpkg: must be" in str(e.value)
+        _ = filter_by_zonal_wave_truncation(da_fc, [1, 2, 3], fftpkg="foo")
+    assert "fftpkg must be 'scipy' or 'xrft'" in str(e.value)
 
 
 def test_zonal_wave_contributions():
@@ -353,8 +353,7 @@ def test_zonal_wave_covariance_no_nlons():
     with pytest.raises(ValueError) as e:
         _ = da_fc1.attrs.pop("nlons")
         _ = da_fc2.attrs.pop("nlons")
-
-        zon_cov = zonal_wave_covariance(da_fc1, da_fc2)
+        _ = zonal_wave_covariance(da_fc1, da_fc2)
     assert (
         "nlons must either be provided as a kwarg or be in the attrs of fc1 or fc2"
         in str(e.value)