Zonal Wavenumber Decompositions (pyzome.zonal_waves)

These functions provide utilities for decomposing fields as a function of zonal wavenumber. The primary function of interest is pyzome.zonal_wave_coeffs(), which performs a rFFT operation across the longitude dimension of the input, returning the complex Fourier coefficients. 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.

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.

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

Functions

pyzome.zonal_waves.zonal_wave_coeffs(dat: DataArray, *, waves: Sequence[int] | None = None, fftpkg: str = 'scipy', lon_coord: str = '') → DataArray

Compute the Fourier transform of data along the zonal direction.

This function is imported at the top level of the package by default.

Parameters:

  • 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 for all.

  • 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.

  • lon_coord (str, optional) – The coordinate name of the longitude dimension. If given an empty string (the default), the function tries to infer which coordinate corresponds to the longitude

Returns:

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 and lon0 that specify the number of longitudes and the starting longitude of the input data.

Return type:

xarray.DataArray

pyzome.zonal_waves.inflate_zonal_wave_coeffs(fc_subset: DataArray, wave_coord: str = '') → DataArray

Inflates an array of zonal wavenumber Fourier coefficients to its expected full-spectrum size. The full size is determined by the “nlons” attribute of the input DataArray.

Inflating is done by inserting zeros in the missing wavenumber bins. This can be useful for keeping only a subset of wavenumbers while still allowing for taking an inverse FFT for filtering or determining the contributions of individual wavenumbers to the full field.

Parameters:

  • fc_subset (xarray.DataArray) – Fourier coefficients as a function of zonal wavenumber, 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 corresponds to the wavenumbers

Returns:

Fourier coefficients as a function of zonal wavenumber, with the spectrum size determined by the “nlons” attribute of the input DataArray.

Return type:

xarray.DataArray

pyzome.zonal_waves.zonal_wave_ampl_phase(fc: DataArray, phase_deg: bool = False, wave_coord: str = '') → tuple[DataArray, DataArray]

Calculates the amplitudes and relative phases of waves in the zonal direction.

Parameters:

  • fc (xarray.DataArray) – Fourier coefficients as a function of zonal wavenumber, as returned by zonal_wave_coeffs.

  • phase_deg (boolean, optional) – Whether to return the relative phases in radians or degrees.

  • 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 corresponds to the wavenumbers

Returns:

Contains the amplitudes and phases for the input field zonal waves.

Return type:

tuple of two xarray.DataArray

See also

zonal_wave_coeffs

pyzome.zonal_waves.filter_by_zonal_wave_truncation(fc: DataArray, waves: Sequence[int], fftpkg: str = 'scipy', wave_coord: str = '', lons: DataArray | None = None) → DataArray

Filters a field by truncating the zonal wavenumbers. This is done by taking an inverse rFFT of the Fourier coefficients, with the unwanted wavenumbers set to zero.

Parameters:

  • fc (xarray.DataArray) – Fourier coefficients as a function of zonal wavenumber, as returned by zonal_wave_coeffs.

  • waves (iterable of int or slice) – The wavenumbers to keep. If a slice is given, it must be a slice of integers. If an iterable is given, it must be an iterable of integers.

  • fftpkg ({'scipy', 'numpy'}, optional) – The FFT package to use. Defaults to ‘scipy’.

  • 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 corresponds to the wavenumbers

  • 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:

The field filtered to retain only the given zonal wavenumbers.

Return type:

xarray.DataArray

pyzome.zonal_waves.zonal_wave_contributions(fc: DataArray, waves: Sequence[int], fftpkg: str = 'scipy', wave_coord: str = '', lons: DataArray | None = None) → DataArray

Computes the individual contributions of each zonal wavenumber to the input field. This is done by taking an inverse FFT of the Fourier coefficients, with all but one wavenumber set to zero.

Parameters:

  • fc (xarray.DataArray) – Fourier coefficients as a function of zonal wavenumber, as returned by zonal_wave_coeffs.

  • waves (iterable of int, optional) – The zonal wavenumbers to maintain in the output.

  • 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.

  • 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 corresponds to the wavenumbers

  • 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:

The individual contributions of each wavenumber to the original input field.

Return type:

xarray.DataArray

See also

zonal_wave_coeffs, filter_by_zonal_wave_truncation

pyzome.zonal_waves.zonal_wave_covariance(fc1: DataArray, fc2: DataArray, wave_coord: str = '', nlons: int | None = None) → DataArray

Calculates the covariance of two fields partititioned into zonal wavenumbers.

Parameters:

  • fc1 (xarray.DataArray) – The zonal Fourier coefficients of the first field, as returned by 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 corresponds to the wavenumbers.

  • nlons (int, optional) – The number of longitudes in the original data. If not provided, the function will try to find this in the attrs of fc1 or fc2.

Return type:

xarray.DataArray

See also

zonal_wave_coeffs