xarray accessors (pyzome.accessor)

pyzome provides a set of xarray 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

class pyzome.accessor.PyzomeAccessor(xarray_obj)

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()
coord_map(coord: str)

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:

The actual coordinate name in the xarray object

Return type:

str

Raises:

KeyError – If the underlying coordinate name is not found in the xarray object

property lon

The longitude coordinate of the xarray object, if it exists/was identified

property lat

The latitude coordinate of the xarray object, if it exists/was identified

property plev

The pressure level coordinate of the xarray object, if it exists/was identified

property zonal_wavenum

The zonal wavenumber coordinate of the xarray object, if it exists/was identified

add_logp_altitude(H: float = 7000, p0: float = 100000.0)

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.

zonal_mean(strict: bool = True)

Computes the zonal mean of the xarray object, if a longitude coordinate was identifed.

meridional_mean(lat1: float, lat2: float, strict: bool = True)

Computes the meridional mean of the xarray object, if a latitude coordinate was identifed.

PyzomeDataArrayAccessor

class pyzome.accessor.PyzomeDataArrayAccessor(xarray_obj: DataArray)

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])
zonal_wave_coeffs(waves: Sequence[int] | None = None, fftpkg: str = 'scipy')

Computes the zonal wave coefficients of a DataArray, if a longitude coordinate was identifed.

inflate_zonal_wave_coeffs()

Inflates the zonal wave coefficients of a DataArray to the full spectrum expected for inverse transforms.

filter_by_zonal_wave_truncation(waves: Sequence[int], fftpkg: str = 'scipy', lons: DataArray | None = None)

Filters the input DataArray by truncating to include only the specified zonal wavenumbers.

zonal_wave_contributions(waves: Sequence[int], fftpkg: str = 'scipy', lons: DataArray | None = None)

Computes the individual contributions of each given zonal wavenumber to the input DataArray.