Transforms

Spectral and signal transforms on xarray data — Fourier / wavelet / cosine spectra, rotary and cross spectra, oceanographic spectral fluxes, matrix decompositions (PCA / EOF / ICA / NMF / KMeans), and binary-mask morphology. The numpy compute cores are jaxtyped; the xarray wrappers carry coordinate and attribute handling.

Spectra

PowerSpectrum

PowerSpectrum(variable: str, dim: str | Sequence[str], *, isotropic: bool = False, **kwargs: Any)

Bases: Operator

Power spectrum of ds[variable]. Set isotropic=True for the radial (2-D) variant.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    variable: str,
    dim: str | Sequence[str],
    *,
    isotropic: bool = False,
    **kwargs: Any,
) -> None:
    self.variable = variable
    self.dim = dim
    self.isotropic = isotropic
    self.kwargs = dict(kwargs)

CrossSpectrum

CrossSpectrum(var_a: str, var_b: str, dim: str | Sequence[str], **kwargs: Any)

Bases: Operator

Cross-power spectrum of ds[var_a] and ds[var_b].

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    var_a: str,
    var_b: str,
    dim: str | Sequence[str],
    **kwargs: Any,
) -> None:
    self.var_a = var_a
    self.var_b = var_b
    self.dim = dim
    self.kwargs = dict(kwargs)

Coherence

Coherence(var_a: str, var_b: str, dim: str | Sequence[str], **kwargs: Any)

Bases: Operator

Magnitude-squared coherence of two variables.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    var_a: str,
    var_b: str,
    dim: str | Sequence[str],
    **kwargs: Any,
) -> None:
    self.var_a = var_a
    self.var_b = var_b
    self.dim = dim
    self.kwargs = dict(kwargs)

RotarySpectrum

RotarySpectrum(u: str, v: str, dim: str, *, avg_dims: str | Sequence[str] | None = None)

Bases: Operator

Rotary power spectrum from ds[u] / ds[v] along dim.

The operator selects the named velocity variables from the input Dataset and calls :func:rotary_spectrum on the resulting DataArrays.

Parameters:

Name	Type	Description	Default
u	str	Name of the zonal velocity variable in the input Dataset.	required
v	str	Name of the meridional velocity variable in the input Dataset.	required
dim	str	Dimension to Fourier transform.	required
avg_dims	str | Sequence[str] | None	Optional dimension or dimensions to average in each output.	None

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    u: str,
    v: str,
    dim: str,
    *,
    avg_dims: str | Sequence[str] | None = None,
) -> None:
    self.u = u
    self.v = v
    self.dim = dim
    self.avg_dims = avg_dims

STFT

STFT(variable: str, dim: str, *, window_size: int, hop: int | None = None, window: str = 'tukey', detrend: str | None = 'linear')

Bases: Operator

Short-time Fourier transform of ds[variable] along dim.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    variable: str,
    dim: str,
    *,
    window_size: int,
    hop: int | None = None,
    window: str = "tukey",
    detrend: str | None = "linear",
) -> None:
    self.variable = variable
    self.dim = dim
    self.window_size = window_size
    self.hop = hop
    self.window = window
    self.detrend = detrend

DCT

DCT(variable: str, dim: str, *, type: int = 2, norm: str | None = 'ortho')

Bases: Operator

Discrete Cosine Transform of ds[variable] along dim.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    variable: str,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> None:
    self.variable = variable
    self.dim = dim
    self.type = type
    self.norm = norm

DST

DST(variable: str, dim: str, *, type: int = 2, norm: str | None = 'ortho')

Bases: Operator

Discrete Sine Transform of ds[variable] along dim.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    variable: str,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> None:
    self.variable = variable
    self.dim = dim
    self.type = type
    self.norm = norm

CWT

CWT(variable: str, dim: str, *, scales: Sequence[float], wavelet: str = 'morl', sampling_period: float = 1.0)

Bases: Operator

Continuous Wavelet Transform of ds[variable] along dim.

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    variable: str,
    dim: str,
    *,
    scales: Sequence[float],
    wavelet: str = "morl",
    sampling_period: float = 1.0,
) -> None:
    self.variable = variable
    self.dim = dim
    self.scales = list(scales)
    self.wavelet = wavelet
    self.sampling_period = sampling_period

Spectral flux

KESpectralFlux

KESpectralFlux(u: str, v: str, dim: Sequence[str], *, window: str | None = 'tukey', detrend: str | None = 'linear', avg_dims: str | Sequence[str] | None = None, return_2d: bool = False)

Bases: Operator

Kinetic-energy spectral flux from ds[u] and ds[v].

Parameters:

Name	Type	Description	Default
u	str	Name of the zonal velocity variable in the input Dataset.	required
v	str	Name of the meridional velocity variable in the input Dataset.	required
dim	Sequence[str]	Two spatial dimensions to Fourier transform.	required
window	str | None	Optional Fourier window.	'tukey'
detrend	str | None	Optional Fourier detrending mode.	'linear'
avg_dims	str | Sequence[str] | None	Optional non-spectral dimensions to average.	None
return_2d	bool	If True, include transfer_2d in the output Dataset.	False

Returns:

Type	Description
	Dataset with transfer and flux variables, plus transfer_2d
	when requested.

Examples:

>>> op = KESpectralFlux("u", "v", ("x", "y"), avg_dims="time")
>>> flux_ds = op(ds)

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    u: str,
    v: str,
    dim: Sequence[str],
    *,
    window: str | None = "tukey",
    detrend: str | None = "linear",
    avg_dims: str | Sequence[str] | None = None,
    return_2d: bool = False,
) -> None:
    self.u = u
    self.v = v
    self.dim = dim
    self.window = window
    self.detrend = detrend
    self.avg_dims = avg_dims
    self.return_2d = return_2d

EnstrophySpectralFlux

EnstrophySpectralFlux(u: str, v: str, dim: Sequence[str], *, window: str | None = 'tukey', detrend: str | None = 'linear', avg_dims: str | Sequence[str] | None = None, return_2d: bool = False)

Bases: Operator

Enstrophy spectral flux from ds[u] and ds[v].

Parameters:

Name	Type	Description	Default
u	str	Name of the zonal velocity variable in the input Dataset.	required
v	str	Name of the meridional velocity variable in the input Dataset.	required
dim	Sequence[str]	Two spatial dimensions to Fourier transform.	required
window	str | None	Optional Fourier window.	'tukey'
detrend	str | None	Optional Fourier detrending mode.	'linear'
avg_dims	str | Sequence[str] | None	Optional non-spectral dimensions to average.	None
return_2d	bool	If True, include transfer_2d in the output Dataset.	False

Returns:

Type	Description
	Dataset with vorticity-based enstrophy transfer and flux
	variables, plus transfer_2d when requested.

Examples:

>>> op = EnstrophySpectralFlux("u", "v", ("x", "y"))
>>> flux_ds = op(ds)

Source code in packages/xrtoolz-ops/src/xrops/transforms/operators.py

def __init__(
    self,
    u: str,
    v: str,
    dim: Sequence[str],
    *,
    window: str | None = "tukey",
    detrend: str | None = "linear",
    avg_dims: str | Sequence[str] | None = None,
    return_2d: bool = False,
) -> None:
    self.u = u
    self.v = v
    self.dim = dim
    self.window = window
    self.detrend = detrend
    self.avg_dims = avg_dims
    self.return_2d = return_2d

Spectral functions

These pure functions back the operators above; each takes xr.DataArray/xr.Dataset and a dim: argument.

power_spectrum

power_spectrum(da: DataArray, dim: str | Sequence[str], *, isotropic: bool = False, **kwargs: Any) -> DataArray

Power spectral density of da along dim.

Parameters:

Name	Type	Description	Default
da	DataArray	Input field.	required
dim	str | Sequence[str]	Single dim or sequence of dims to transform. Pass a single string for 1-D spectra; a tuple/list for multi-D (space-time) spectra.	required
isotropic	bool	If True, compute the radially-averaged (isotropic) power spectrum. Requires exactly two spatial dims.	False
**kwargs	Any	Forwarded to :func:xrft.power_spectrum / :func:xrft.isotropic_power_spectrum. Sensible defaults (Tukey window, linear detrend, density scaling) fill in keys not provided.	{}

Returns:

Type	Description
DataArray	DataArray of the spectrum, named f"{da.name}_psd" (or
DataArray	f"{da.name}_iso_psd" if isotropic=True).

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def power_spectrum(
    da: xr.DataArray,
    dim: str | Sequence[str],
    *,
    isotropic: bool = False,
    **kwargs: Any,
) -> xr.DataArray:
    """Power spectral density of ``da`` along ``dim``.

    Args:
        da: Input field.
        dim: Single dim or sequence of dims to transform. Pass a single
            string for 1-D spectra; a tuple/list for multi-D
            (space-time) spectra.
        isotropic: If ``True``, compute the radially-averaged (isotropic)
            power spectrum. Requires exactly two spatial dims.
        **kwargs: Forwarded to :func:`xrft.power_spectrum` /
            :func:`xrft.isotropic_power_spectrum`. Sensible defaults
            (Tukey window, linear detrend, density scaling) fill in
            keys not provided.

    Returns:
        DataArray of the spectrum, named ``f"{da.name}_psd"`` (or
        ``f"{da.name}_iso_psd"`` if ``isotropic=True``).
    """
    dims = [dim] if isinstance(dim, str) else list(dim)
    opts = {**_DEFAULT_PSD_KWARGS, **kwargs}
    if isotropic:
        if len(dims) != 2:
            raise ValueError(
                f"isotropic_power_spectrum requires exactly 2 spatial dims; "
                f"got dim={dims}."
            )
        out = xrft.isotropic_power_spectrum(da, dim=dims, **opts)
        out.name = _output_name(da, "iso_psd")
        return out
    out = xrft.power_spectrum(da, dim=dims, **opts)
    out.name = _output_name(da, "psd")
    return out

isotropic_power_spectrum

isotropic_power_spectrum(da: DataArray, dim: Sequence[str], **kwargs: Any) -> DataArray

Radially-averaged 2-D power spectrum. Convenience alias of :func:power_spectrum with isotropic=True.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def isotropic_power_spectrum(
    da: xr.DataArray,
    dim: Sequence[str],
    **kwargs: Any,
) -> xr.DataArray:
    """Radially-averaged 2-D power spectrum. Convenience alias of
    :func:`power_spectrum` with ``isotropic=True``."""
    return power_spectrum(da, dim=dim, isotropic=True, **kwargs)

cross_spectrum

cross_spectrum(da_a: DataArray, da_b: DataArray, dim: str | Sequence[str], **kwargs: Any) -> DataArray

Cross-power spectrum of da_a and da_b along dim.

Returns:

Type	Description
DataArray	DataArray named f"{a.name}_{b.name}_csd".

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def cross_spectrum(
    da_a: xr.DataArray,
    da_b: xr.DataArray,
    dim: str | Sequence[str],
    **kwargs: Any,
) -> xr.DataArray:
    """Cross-power spectrum of ``da_a`` and ``da_b`` along ``dim``.

    Returns:
        DataArray named ``f"{a.name}_{b.name}_csd"``.
    """
    dims = [dim] if isinstance(dim, str) else list(dim)
    opts = {**_DEFAULT_PSD_KWARGS, **kwargs}
    out = xrft.cross_spectrum(da_a, da_b, dim=dims, **opts)
    a_name = da_a.name if da_a.name is not None else "a"
    b_name = da_b.name if da_b.name is not None else "b"
    out.name = f"{a_name}_{b_name}_csd"
    return out

coherence

coherence(da_a: DataArray, da_b: DataArray, dim: str | Sequence[str], **kwargs: Any) -> DataArray

Magnitude-squared coherence |S_ab|² / (S_aa · S_bb).

Bounded in [0, 1]: 1 means the two signals share identical phase-aligned content at that frequency. Computed from the :func:xrft.cross_spectrum and individual auto-spectra under the same averaging window, so the three quantities are commensurable.

Returns:

Type	Description
DataArray	DataArray named f"{a.name}_{b.name}_coh".

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def coherence(
    da_a: xr.DataArray,
    da_b: xr.DataArray,
    dim: str | Sequence[str],
    **kwargs: Any,
) -> xr.DataArray:
    """Magnitude-squared coherence ``|S_ab|² / (S_aa · S_bb)``.

    Bounded in ``[0, 1]``: 1 means the two signals share identical
    phase-aligned content at that frequency. Computed from the
    :func:`xrft.cross_spectrum` and individual auto-spectra under the
    same averaging window, so the three quantities are commensurable.

    Returns:
        DataArray named ``f"{a.name}_{b.name}_coh"``.
    """
    dims = [dim] if isinstance(dim, str) else list(dim)
    opts = {**_DEFAULT_PSD_KWARGS, **kwargs}
    s_aa = xrft.power_spectrum(da_a, dim=dims, **opts)
    s_bb = xrft.power_spectrum(da_b, dim=dims, **opts)
    s_ab = xrft.cross_spectrum(da_a, da_b, dim=dims, **opts)
    coh = (np.abs(s_ab) ** 2) / (s_aa * s_bb)
    a_name = da_a.name if da_a.name is not None else "a"
    b_name = da_b.name if da_b.name is not None else "b"
    coh.name = f"{a_name}_{b_name}_coh"
    return coh

rotary_spectrum

rotary_spectrum(u: DataArray, v: DataArray, *, dim: str, avg_dims: str | Sequence[str] | None = None) -> Dataset

Rotary power spectrum from horizontal velocity components.

Computes the two-sided FFT of the complex velocity u + i v, splits positive and negative wavenumbers into counter-clockwise and clockwise components, then folds both onto a shared positive wavenumber axis. The unnormalized FFT is converted to a density with dx / n scaling so integrating psd_cw + psd_ccw over wavenumber recovers velocity variance. Polarization is NaN where total rotary power is negligible.

Parameters:

Name	Type	Description	Default
u	DataArray	Zonal/eastward velocity DataArray.	required
v	DataArray	Meridional/northward velocity DataArray on the same grid as u.	required
dim	str	Dimension to Fourier transform.	required
avg_dims	str | Sequence[str] | None	Optional dimension or dimensions to average in each output.	None

Returns:

Type	Description
Dataset	Dataset with psd_ccw, psd_cw, and polarization where
Dataset	polarization = (psd_cw - psd_ccw) / (psd_cw + psd_ccw).

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def rotary_spectrum(
    u: xr.DataArray,
    v: xr.DataArray,
    *,
    dim: str,
    avg_dims: str | Sequence[str] | None = None,
) -> xr.Dataset:
    """Rotary power spectrum from horizontal velocity components.

    Computes the two-sided FFT of the complex velocity ``u + i v``, splits
    positive and negative wavenumbers into counter-clockwise and clockwise
    components, then folds both onto a shared positive ``wavenumber`` axis.
    The unnormalized FFT is converted to a density with ``dx / n`` scaling so
    integrating ``psd_cw + psd_ccw`` over wavenumber recovers velocity variance.
    Polarization is NaN where total rotary power is negligible.

    Args:
        u: Zonal/eastward velocity DataArray.
        v: Meridional/northward velocity DataArray on the same grid as ``u``.
        dim: Dimension to Fourier transform.
        avg_dims: Optional dimension or dimensions to average in each output.

    Returns:
        Dataset with ``psd_ccw``, ``psd_cw``, and ``polarization`` where
        ``polarization = (psd_cw - psd_ccw) / (psd_cw + psd_ccw)``.
    """
    if dim not in u.dims or dim not in v.dims:
        raise ValueError(
            f"dim={dim!r} must be present on both u (dims={u.dims}) and "
            f"v (dims={v.dims})."
        )

    w = u + 1j * v
    # Rectangular window + mean-only detrending keeps the dx / n density scaling
    # Parseval-consistent for variance.
    spec = xrft.fft(w, dim=[dim], window=None, detrend="constant", true_amplitude=False)
    freq_dim = f"freq_{dim}"
    # DataArrays can legally have dims without explicit coord variables;
    # fall back to unit spacing in that case rather than KeyError-ing. The
    # density scaling tolerates non-uniform / singleton coords too.
    density_scale = (
        coord_spacing(
            u,
            dim,
            require_coord=False,
            require_min_samples=False,
            require_uniform=False,
        )
        / u.sizes[dim]
    )
    psd = (np.abs(spec) ** 2) * density_scale

    psd_ccw = psd.where(psd[freq_dim] > 0.0, drop=True).rename({freq_dim: "wavenumber"})
    neg = psd.where(psd[freq_dim] < 0.0, drop=True)
    psd_cw = (
        neg.assign_coords({freq_dim: np.abs(neg[freq_dim])})
        .sortby(freq_dim)
        .rename({freq_dim: "wavenumber"})
    )
    # Union the two folds and fill missing bins with 0 so an even-length
    # FFT's lone -Nyquist bin survives, keeping psd_cw + psd_ccw
    # Parseval-consistent with the input variance.
    psd_ccw, psd_cw = xr.align(psd_ccw, psd_cw, join="outer", fill_value=0.0)
    psd_ccw.name = "psd_ccw"
    psd_cw.name = "psd_cw"

    denom = psd_cw + psd_ccw
    polarization = ((psd_cw - psd_ccw) / denom).where(
        denom > _polarization_epsilon(denom)
    )
    polarization.name = "polarization"
    out = xr.Dataset(
        {"psd_ccw": psd_ccw, "psd_cw": psd_cw, "polarization": polarization}
    )
    if avg_dims is not None:
        dims = [avg_dims] if isinstance(avg_dims, str) else list(avg_dims)
        out = out.mean(dim=dims)
    return out

compensated_spectrum

compensated_spectrum(psd: DataArray, *, wavenumber_dim: str = 'freq_r', exponent: float) -> DataArray

Return psd * k**exponent for inertial-range compensation.

Parameters:

Name	Type	Description	Default
psd	DataArray	Spectrum to compensate.	required
wavenumber_dim	str	Name of the wavenumber dimension.	'freq_r'
exponent	float	Power-law exponent, for example 5 / 3 for Kolmogorov scaling.	required

Returns:

Type	Description
DataArray	DataArray named f"{psd.name}_compensated" that is flat where
DataArray	psd ∝ k^{-exponent}.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def compensated_spectrum(
    psd: xr.DataArray,
    *,
    wavenumber_dim: str = "freq_r",
    exponent: float,
) -> xr.DataArray:
    """Return ``psd * k**exponent`` for inertial-range compensation.

    Args:
        psd: Spectrum to compensate.
        wavenumber_dim: Name of the wavenumber dimension.
        exponent: Power-law exponent, for example ``5 / 3`` for Kolmogorov
            scaling.

    Returns:
        DataArray named ``f"{psd.name}_compensated"`` that is flat where
        ``psd ∝ k^{-exponent}``.
    """
    out = psd * psd[wavenumber_dim] ** exponent
    out.name = _output_name(psd, "compensated", fallback="spectrum")
    return out

stft

stft(da: DataArray, dim: str, *, window_size: int, hop: int | None = None, window: str = 'tukey', detrend: str | None = 'linear') -> DataArray

Short-time Fourier transform along dim.

Splits the signal into overlapping windows of length window_size spaced by hop samples (default window_size // 2), computes a Fourier transform per window, and stacks the result on a new segment axis whose coordinate is the centre time of each window.

Parameters:

Name	Type	Description	Default
da	DataArray	1-D signal (along dim) plus arbitrary extra dims that are broadcast over.	required
dim	str	Time / sample dimension to slide the window along.	required
window_size	int	Length (in samples) of each STFT window.	required
hop	int | None	Stride between window starts. Defaults to window_size // 2 (50 % overlap).	None
window	str	Window name forwarded to :func:xrft.power_spectrum's internal window factory (Tukey by default).	'tukey'
detrend	str | None	Per-window detrending. None to disable.	'linear'

Returns:

Type	Description
DataArray	DataArray with dims (*outer_dims, segment, freq_<dim>),
DataArray	named f"{da.name}_stft".

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def stft(
    da: xr.DataArray,
    dim: str,
    *,
    window_size: int,
    hop: int | None = None,
    window: str = "tukey",
    detrend: str | None = "linear",
) -> xr.DataArray:
    """Short-time Fourier transform along ``dim``.

    Splits the signal into overlapping windows of length ``window_size``
    spaced by ``hop`` samples (default ``window_size // 2``), computes a
    Fourier transform per window, and stacks the result on a new
    ``segment`` axis whose coordinate is the centre time of each
    window.

    Args:
        da: 1-D signal (along ``dim``) plus arbitrary extra dims that
            are broadcast over.
        dim: Time / sample dimension to slide the window along.
        window_size: Length (in samples) of each STFT window.
        hop: Stride between window starts. Defaults to
            ``window_size // 2`` (50 % overlap).
        window: Window name forwarded to :func:`xrft.power_spectrum`'s
            internal window factory (Tukey by default).
        detrend: Per-window detrending. ``None`` to disable.

    Returns:
        DataArray with dims ``(*outer_dims, segment, freq_<dim>)``,
        named ``f"{da.name}_stft"``.
    """
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    if window_size < 2:
        raise ValueError(f"window_size must be >= 2; got {window_size}.")
    if hop is None:
        hop = window_size // 2 if window_size > 1 else 1
    if hop < 1:
        raise ValueError(f"hop must be >= 1; got {hop}.")

    n_total = da.sizes[dim]
    if window_size > n_total:
        raise ValueError(f"window_size={window_size} exceeds {dim!r} length {n_total}.")

    starts = np.arange(0, n_total - window_size + 1, hop, dtype=int)
    if starts.size == 0:
        raise ValueError(
            f"No STFT windows fit: window_size={window_size}, hop={hop}, n={n_total}."
        )

    # Use FFT (complex) per segment for true STFT, not power spectrum,
    # so callers can square / phase as they see fit.
    segments = []
    coord = da[dim] if dim in da.coords else None
    centres = []
    for s in starts:
        segment = da.isel({dim: slice(int(s), int(s) + window_size)})
        spec = xrft.fft(
            segment, dim=[dim], window=window, detrend=detrend, true_amplitude=False
        )
        segments.append(spec)
        if coord is not None:
            centres.append(coord.isel({dim: int(s) + window_size // 2}).values.item())
        else:
            centres.append(int(s) + window_size // 2)

    out = xr.concat(segments, dim="segment")
    if centres:
        out = out.assign_coords(segment=("segment", np.asarray(centres)))
    out.name = _output_name(da, "stft")
    return out

dct

dct(da: DataArray, dim: str, *, type: int = 2, norm: str | None = 'ortho') -> DataArray

Discrete Cosine Transform along dim.

Parameters:

Name	Type	Description	Default
da	DataArray	Input field.	required
dim	str	Dimension to transform.	required
type	int	DCT type 1, 2, 3, 4 (default 2 — the DCT-II used by JPEG and most signal-processing applications).	2
norm	str | None	"ortho" (default) makes the transform unitary so that idct(dct(x)) == x. Pass None for the unnormalised scipy default.	'ortho'

Returns:

Type	Description
DataArray	DataArray of the same shape, named f"{name}_dct".

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/dct.py

def dct(
    da: xr.DataArray,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> xr.DataArray:
    """Discrete Cosine Transform along ``dim``.

    Args:
        da: Input field.
        dim: Dimension to transform.
        type: DCT type ``1, 2, 3, 4`` (default ``2`` — the DCT-II used
            by JPEG and most signal-processing applications).
        norm: ``"ortho"`` (default) makes the transform unitary so that
            ``idct(dct(x)) == x``. Pass ``None`` for the unnormalised
            scipy default.

    Returns:
        DataArray of the same shape, named ``f"{name}_dct"``.
    """
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    axis = da.get_axis_num(dim)
    raw = _scipy_dct(da.values, type=type, axis=axis, norm=norm)
    return xr.DataArray(
        raw,
        dims=da.dims,
        coords={k: da.coords[k] for k in da.coords},
        attrs=dict(da.attrs),
        name=_output_name(da, "dct"),
    )

idct

idct(da: DataArray, dim: str, *, type: int = 2, norm: str | None = 'ortho') -> DataArray

Inverse DCT along dim. See :func:dct for arguments.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/dct.py

def idct(
    da: xr.DataArray,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> xr.DataArray:
    """Inverse DCT along ``dim``. See :func:`dct` for arguments."""
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    axis = da.get_axis_num(dim)
    raw = _scipy_idct(da.values, type=type, axis=axis, norm=norm)
    return xr.DataArray(
        raw,
        dims=da.dims,
        coords={k: da.coords[k] for k in da.coords},
        attrs=dict(da.attrs),
        name=_output_name(da, "idct"),
    )

dst

dst(da: DataArray, dim: str, *, type: int = 2, norm: str | None = 'ortho') -> DataArray

Discrete Sine Transform along dim. See :func:dct.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/dct.py

def dst(
    da: xr.DataArray,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> xr.DataArray:
    """Discrete Sine Transform along ``dim``. See :func:`dct`."""
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    axis = da.get_axis_num(dim)
    raw = _scipy_dst(da.values, type=type, axis=axis, norm=norm)
    return xr.DataArray(
        raw,
        dims=da.dims,
        coords={k: da.coords[k] for k in da.coords},
        attrs=dict(da.attrs),
        name=_output_name(da, "dst"),
    )

idst

idst(da: DataArray, dim: str, *, type: int = 2, norm: str | None = 'ortho') -> DataArray

Inverse DST along dim. See :func:dct.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/dct.py

def idst(
    da: xr.DataArray,
    dim: str,
    *,
    type: int = 2,
    norm: str | None = "ortho",
) -> xr.DataArray:
    """Inverse DST along ``dim``. See :func:`dct`."""
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    axis = da.get_axis_num(dim)
    raw = _scipy_idst(da.values, type=type, axis=axis, norm=norm)
    return xr.DataArray(
        raw,
        dims=da.dims,
        coords={k: da.coords[k] for k in da.coords},
        attrs=dict(da.attrs),
        name=_output_name(da, "idst"),
    )

cwt

cwt(da: DataArray, dim: str, *, scales: Sequence[float] | ndarray, wavelet: str = 'morl', sampling_period: float = 1.0) -> DataArray

Continuous Wavelet Transform along dim.

Parameters:

Name	Type	Description	Default
da	DataArray	Input field.	required
dim	str	Dimension to transform.	required
scales	Sequence[float] | ndarray	Wavelet scales (positive floats). Larger scale → lower frequency content. See :func:pywt.scale2frequency.	required
wavelet	str	Wavelet name (e.g. "morl" for Morlet, "mexh", "cmor", "gaus1").	'morl'
sampling_period	float	Sample spacing along dim in seconds (or equivalent), used to convert scales to frequencies.	1.0

Returns:

Type	Description
DataArray	Complex-valued DataArray with dims (*outer_dims, scale, dim),
DataArray	named f"{name}_cwt". A frequency 1-D coord is attached
DataArray	along the new scale axis.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/wavelet.py

def cwt(
    da: xr.DataArray,
    dim: str,
    *,
    scales: Sequence[float] | np.ndarray,
    wavelet: str = "morl",
    sampling_period: float = 1.0,
) -> xr.DataArray:
    """Continuous Wavelet Transform along ``dim``.

    Args:
        da: Input field.
        dim: Dimension to transform.
        scales: Wavelet scales (positive floats). Larger scale → lower
            frequency content. See :func:`pywt.scale2frequency`.
        wavelet: Wavelet name (e.g. ``"morl"`` for Morlet,
            ``"mexh"``, ``"cmor"``, ``"gaus1"``).
        sampling_period: Sample spacing along ``dim`` in seconds (or
            equivalent), used to convert scales to frequencies.

    Returns:
        Complex-valued DataArray with dims ``(*outer_dims, scale, dim)``,
        named ``f"{name}_cwt"``. A ``frequency`` 1-D coord is attached
        along the new ``scale`` axis.
    """
    pywt = _import_pywt()
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")
    scales_arr = np.asarray(scales, dtype=float)
    if scales_arr.ndim != 1 or scales_arr.size == 0:
        raise ValueError("scales must be a non-empty 1-D sequence of positive floats.")
    if (scales_arr <= 0).any():
        raise ValueError("scales must be strictly positive.")

    axis = da.get_axis_num(dim)
    coeffs, freqs = pywt.cwt(
        da.values, scales_arr, wavelet, sampling_period=sampling_period, axis=axis
    )

    new_dims = ("scale", *da.dims)
    coords: dict = {k: da.coords[k] for k in da.coords}
    coords["scale"] = scales_arr
    coords["frequency"] = ("scale", np.asarray(freqs))

    return xr.DataArray(
        coeffs,
        dims=new_dims,
        coords=coords,
        attrs=dict(da.attrs),
        name=_output_name(da, "cwt"),
    )

dwt

dwt(da: DataArray, dim: str, *, wavelet: str = 'db4', level: int | None = None, mode: str = 'symmetric') -> dict[str, DataArray]

Multi-level discrete wavelet decomposition along dim.

Returns a dict {"approx": ..., "detail_1": ..., ..., "detail_L": ...} where L = level (or the maximum allowed by the signal length if level is None). Each entry is a DataArray with the transformed axis renamed (dim is replaced by f"{dim}_dwt_<k>") since the level-wise outputs have different lengths.

Parameters:

Name	Type	Description	Default
da	DataArray	Input field.	required
dim	str	Dimension to decompose.	required
wavelet	str	Wavelet name ("db4", "sym5", "haar"…).	'db4'
level	int | None	Decomposition depth. None uses :func:pywt.dwt_max_level.	None
mode	str	Signal-extension mode at boundaries.	'symmetric'

Returns:

Type	Description
dict[str, DataArray]	Dictionary of approximation and detail coefficients as
dict[str, DataArray]	DataArrays.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/wavelet.py

def dwt(
    da: xr.DataArray,
    dim: str,
    *,
    wavelet: str = "db4",
    level: int | None = None,
    mode: str = "symmetric",
) -> dict[str, xr.DataArray]:
    """Multi-level discrete wavelet decomposition along ``dim``.

    Returns a dict ``{"approx": ..., "detail_1": ..., ..., "detail_L": ...}``
    where ``L = level`` (or the maximum allowed by the signal length if
    ``level`` is ``None``). Each entry is a DataArray with the
    transformed axis renamed (``dim`` is replaced by ``f"{dim}_dwt_<k>"``)
    since the level-wise outputs have different lengths.

    Args:
        da: Input field.
        dim: Dimension to decompose.
        wavelet: Wavelet name (``"db4"``, ``"sym5"``, ``"haar"``…).
        level: Decomposition depth. ``None`` uses
            :func:`pywt.dwt_max_level`.
        mode: Signal-extension mode at boundaries.

    Returns:
        Dictionary of approximation and detail coefficients as
        DataArrays.
    """
    pywt = _import_pywt()
    if dim not in da.dims:
        raise ValueError(f"dim={dim!r} not present on DataArray with dims={da.dims}.")

    axis = da.get_axis_num(dim)
    n = da.sizes[dim]
    max_level = pywt.dwt_max_level(n, pywt.Wavelet(wavelet).dec_len)
    eff_level = max_level if level is None else int(level)
    if eff_level < 1 or eff_level > max_level:
        raise ValueError(
            f"level must be in [1, {max_level}] for n={n} and wavelet={wavelet!r}; "
            f"got {eff_level}."
        )

    coeffs = pywt.wavedec(
        da.values, wavelet=wavelet, level=eff_level, mode=mode, axis=axis
    )
    # ``coeffs[0]`` is the approximation; ``coeffs[1:]`` are details
    # ordered coarsest→finest.
    out: dict[str, xr.DataArray] = {}
    other_dims = tuple(d for d in da.dims if d != dim)
    other_coords = {k: da.coords[k] for k in da.coords if k in other_dims}

    def _wrap(arr: np.ndarray, label: str) -> xr.DataArray:
        new_dim = f"{dim}_dwt_{label}"
        new_dims = tuple(new_dim if d == dim else d for d in da.dims)
        return xr.DataArray(
            arr,
            dims=new_dims,
            coords=other_coords,
            attrs=dict(da.attrs),
            name=_output_name(da, f"dwt_{label}"),
        )

    out["approx"] = _wrap(coeffs[0], "approx")
    for k, det in enumerate(coeffs[1:], start=1):
        out[f"detail_{k}"] = _wrap(det, f"detail_{k}")
    return out

drop_negative_frequencies

drop_negative_frequencies(da: T, dims: Sequence[str], *, drop: bool = True) -> T

Mean along dims after restricting every other freq axis to its positive half.

Used to collapse a multi-D spectrum onto a subset of frequency axes without contaminating the average with the zero / negative halves of the remaining axes. This is the renamed and clarified form of the historical conditional_average.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def drop_negative_frequencies[T: xr.DataArray | xr.Dataset](
    da: T,
    dims: Sequence[str],
    *,
    drop: bool = True,
) -> T:
    """Mean along ``dims`` after restricting every other freq axis to its
    positive half.

    Used to collapse a multi-D spectrum onto a subset of frequency axes
    without contaminating the average with the zero / negative halves
    of the remaining axes. This is the renamed and clarified form of
    the historical ``conditional_average``.
    """
    dims = list(dims)
    remaining = [d for d in da.dims if d not in dims]
    if not remaining:
        return da.mean(dim=dims)

    cond = da[remaining[0]] > 0.0
    for d in remaining[1:]:
        cond = cond & (da[d] > 0.0)
    return da.mean(dim=dims).where(cond, drop=drop)

to_phase

to_phase(da: DataArray, *, time_dim: str, period: float, n_bins: int, epoch: float = 0.0) -> DataArray

Fold time_dim onto a phase axis by binning + averaging.

The time coordinate must be numeric in the same units as period. Phase is computed as ((t - epoch) / period) mod 1 and binned into n_bins evenly-spaced bins on [0, 1). Output carries a "phase" dim with coordinate values at bin centers.

The Layer-1 ToPhase operator handles Dataset selection / per-variable looping.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/coord_remap.py

def to_phase(
    da: xr.DataArray,
    *,
    time_dim: str,
    period: float,
    n_bins: int,
    epoch: float = 0.0,
) -> xr.DataArray:
    """Fold ``time_dim`` onto a phase axis by binning + averaging.

    The time coordinate must be numeric in the same units as ``period``.
    Phase is computed as ``((t - epoch) / period) mod 1`` and binned
    into ``n_bins`` evenly-spaced bins on ``[0, 1)``. Output carries
    a ``"phase"`` dim with coordinate values at bin centers.

    The Layer-1 ``ToPhase`` operator handles Dataset selection /
    per-variable looping.
    """
    if time_dim not in da.dims:
        raise ValueError(
            f"time_dim {time_dim!r} not in DataArray dims {tuple(da.dims)}"
        )
    if time_dim not in da.coords:
        raise ValueError(
            f"DataArray must carry a coordinate named {time_dim!r} to compute phase"
        )
    if period <= 0:
        raise ValueError(f"period must be > 0, got {period}")
    if n_bins < 1:
        raise ValueError(f"n_bins must be >= 1, got {n_bins}")
    if not np.issubdtype(da.dtype, np.number):
        raise TypeError(
            f"to_phase requires numeric data; got dtype {da.dtype}. "
            "Drop or convert the variable before calling to_phase."
        )

    t = np.asarray(da[time_dim].values, dtype=float)
    finite_t = _finite_mask(t)
    phase = np.where(finite_t, ((t - epoch) / period) % 1.0, 0.0)
    edges = np.linspace(0.0, 1.0, n_bins + 1)
    centers = 0.5 * (edges[:-1] + edges[1:])
    bin_idx = np.clip(np.searchsorted(edges, phase, side="right") - 1, 0, n_bins - 1)

    axis = da.get_axis_num(time_dim)
    moved = np.moveaxis(da.values, axis, 0)
    flat = moved.reshape(moved.shape[0], -1)
    # Use a complex accumulator if the data is complex so the imaginary
    # part isn't dropped (P1 review).
    is_complex = np.iscomplexobj(flat)
    acc_dtype = np.complex128 if is_complex else float
    # cast: with dtype resolved at runtime (complex128 or float), the
    # specialised ndarray type numpy's stubs infer would reject the row
    # assignments below.
    sums = cast(np.ndarray, np.zeros((n_bins, flat.shape[1]), dtype=acc_dtype))
    counts = np.zeros((n_bins, flat.shape[1]), dtype=float)
    # A sample is valid only if its time coord is finite AND every data
    # value is finite. Excluding NaN-time rows keeps stale rows out of
    # the phase means (P2 review). ``np.isfinite`` handles integer
    # dtypes (treated as all-finite) where ``np.isnan`` would raise.
    valid_value = (
        np.isfinite(flat)
        if not is_complex
        else (np.isfinite(flat.real) & np.isfinite(flat.imag))
    )
    valid = finite_t[:, None] & valid_value
    for b in range(n_bins):
        m = (bin_idx == b) & finite_t
        if not m.any():
            continue
        sub = flat[m]
        sub_valid = valid[m]
        zero = acc_dtype(0)
        sums[b] = np.where(sub_valid, sub, zero).sum(axis=0)
        counts[b] = sub_valid.sum(axis=0)
    with np.errstate(invalid="ignore"):
        nan_fill = (np.nan + 0j) if is_complex else np.nan
        mean = np.where(counts > 0, sums / np.where(counts > 0, counts, 1.0), nan_fill)
    out_shape = (n_bins, *moved.shape[1:])
    new_values = mean.reshape(out_shape)
    new_values = np.moveaxis(new_values, 0, axis)
    new_dims = tuple("phase" if d == time_dim else d for d in da.dims)
    new_coords = {
        cname: c
        for cname, c in da.coords.items()
        if time_dim not in c.dims and cname != time_dim
    }
    new_coords["phase"] = xr.DataArray(centers, dims=("phase",), name="phase")
    return xr.DataArray(
        new_values,
        dims=new_dims,
        coords=new_coords,
        attrs=dict(da.attrs),
        name=da.name,
    )

fit_spectral_slope

fit_spectral_slope(psd: DataArray, *, wavenumber_dim: str = 'freq_r', k_min: float, k_max: float) -> tuple[float, float]

Fit log(psd) = slope * log(k) + intercept over [k_min, k_max].

Parameters:

Name	Type	Description	Default
psd	DataArray	One-dimensional positive spectrum.	required
wavenumber_dim	str	Name of the wavenumber dimension.	'freq_r'
k_min	float	Lower inclusive wavenumber bound.	required
k_max	float	Upper inclusive wavenumber bound.	required

Returns:

Type	Description
tuple[float, float]	(slope, intercept) from numpy.polyfit in log-log space.

Raises:

Type	Description
ValueError	If fewer than two positive finite samples remain in the requested fit window.

Notes

Use the fitted slope to identify inertial-range power laws such as -5/3 inverse-cascade or -3 enstrophy-cascade scaling.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def fit_spectral_slope(
    psd: xr.DataArray,
    *,
    wavenumber_dim: str = "freq_r",
    k_min: float,
    k_max: float,
) -> tuple[float, float]:
    """Fit ``log(psd) = slope * log(k) + intercept`` over ``[k_min, k_max]``.

    Args:
        psd: One-dimensional positive spectrum.
        wavenumber_dim: Name of the wavenumber dimension.
        k_min: Lower inclusive wavenumber bound.
        k_max: Upper inclusive wavenumber bound.

    Returns:
        ``(slope, intercept)`` from ``numpy.polyfit`` in log-log space.

    Raises:
        ValueError: If fewer than two positive finite samples remain in the
            requested fit window.

    Notes:
        Use the fitted slope to identify inertial-range power laws such as
        ``-5/3`` inverse-cascade or ``-3`` enstrophy-cascade scaling.
    """
    k = psd[wavenumber_dim]
    subset = psd.where((k >= k_min) & (k <= k_max) & (k > 0) & (psd > 0), drop=True)
    x = np.log(subset[wavenumber_dim].values)
    y = np.log(subset.values)
    finite = np.isfinite(x) & np.isfinite(y)
    if finite.sum() < 2:
        raise ValueError(
            f"At least two positive finite spectral samples required in "
            f"[{k_min}, {k_max}]; found {finite.sum()}."
        )
    slope, intercept = np.polyfit(x[finite], y[finite], 1)
    return float(slope), float(intercept)

integral_scale

integral_scale(psd: DataArray, *, wavenumber_dim: str = 'freq_r', moment: int = 1) -> DataArray

Energy-weighted integral scale or Taylor microscale.

Parameters:

Name	Type	Description	Default
psd	DataArray	One-dimensional spectrum with a wavenumber coordinate.	required
wavenumber_dim	str	Name of the wavenumber dimension.	'freq_r'
moment	int	1 for the integral scale, 2 for the Taylor microscale.	1

Returns:

Type	Description
DataArray	DataArray with wavenumber_dim removed.

Notes

moment=1 returns ∫ψ dk / ∫kψ dk. moment=2 returns the Taylor microscale λ = sqrt(∫ψ dk / ∫k²ψ dk). Integrals use xarray's coordinate-aware trapezoidal rule, so non-uniform wavenumber grids (e.g. radially binned spectra) are handled correctly.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def integral_scale(
    psd: xr.DataArray,
    *,
    wavenumber_dim: str = "freq_r",
    moment: int = 1,
) -> xr.DataArray:
    """Energy-weighted integral scale or Taylor microscale.

    Args:
        psd: One-dimensional spectrum with a wavenumber coordinate.
        wavenumber_dim: Name of the wavenumber dimension.
        moment: ``1`` for the integral scale, ``2`` for the Taylor microscale.

    Returns:
        DataArray with ``wavenumber_dim`` removed.

    Notes:
        ``moment=1`` returns ``∫ψ dk / ∫kψ dk``. ``moment=2`` returns the
        Taylor microscale ``λ = sqrt(∫ψ dk / ∫k²ψ dk)``. Integrals use
        xarray's coordinate-aware trapezoidal rule, so non-uniform
        wavenumber grids (e.g. radially binned spectra) are handled
        correctly.
    """
    if moment not in (1, 2):
        raise ValueError(
            f"moment must be 1 (integral scale) or 2 (Taylor microscale); got {moment}."
        )
    k = psd[wavenumber_dim]
    numerator = psd.integrate(coord=wavenumber_dim)
    denominator = (psd * k**moment).integrate(coord=wavenumber_dim)
    out = numerator / denominator
    if moment == 2:
        out = out**0.5
    out.name = _output_name(psd, f"moment{moment}_scale", fallback="spectrum")
    return out

ke_spectral_flux

ke_spectral_flux(u: DataArray, v: DataArray, *, dim: Sequence[str], window: str | None = 'tukey', detrend: str | None = 'linear', avg_dims: str | Sequence[str] | None = None, return_2d: bool = False) -> Dataset

Kinetic-energy spectral flux Π(k).

Parameters:

Name	Type	Description	Default
u	DataArray	Zonal velocity component with both spatial dimensions in dim.	required
v	DataArray	Meridional velocity component on the same grid as u.	required
dim	Sequence[str]	Two spatial dimensions to Fourier transform, for example ("x", "y").	required
window	str | None	Optional window forwarded to :func:xrft.fft.	'tukey'
detrend	str | None	Optional detrending forwarded to :func:xrft.fft.	'linear'
avg_dims	str | Sequence[str] | None	Optional non-spectral dimensions to average before radial integration, for example "time".	None
return_2d	bool	If True, include the unbinned two-dimensional transfer field as transfer_2d.	False

Returns:

Type	Description
Dataset	Dataset with transfer as radially summed KE transfer T(k) and
Dataset	flux as Π(k) = Σ_{k' >= k} T(k'). transfer_2d is included
Dataset	when requested.

Notes

Positive flux denotes downscale kinetic-energy transfer. The nonlinear advection terms are formed in physical space from Fourier-space gradients, transformed back to spectral space, radially binned by freq_r, and accumulated from high to low wavenumber.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def ke_spectral_flux(
    u: xr.DataArray,
    v: xr.DataArray,
    *,
    dim: Sequence[str],
    window: str | None = "tukey",
    detrend: str | None = "linear",
    avg_dims: str | Sequence[str] | None = None,
    return_2d: bool = False,
) -> xr.Dataset:
    """Kinetic-energy spectral flux ``Π(k)``.

    Args:
        u: Zonal velocity component with both spatial dimensions in ``dim``.
        v: Meridional velocity component on the same grid as ``u``.
        dim: Two spatial dimensions to Fourier transform, for example
            ``("x", "y")``.
        window: Optional window forwarded to :func:`xrft.fft`.
        detrend: Optional detrending forwarded to :func:`xrft.fft`.
        avg_dims: Optional non-spectral dimensions to average before radial
            integration, for example ``"time"``.
        return_2d: If ``True``, include the unbinned two-dimensional transfer
            field as ``transfer_2d``.

    Returns:
        Dataset with ``transfer`` as radially summed KE transfer ``T(k)`` and
        ``flux`` as ``Π(k) = Σ_{k' >= k} T(k')``. ``transfer_2d`` is included
        when requested.

    Notes:
        Positive ``flux`` denotes downscale kinetic-energy transfer. The
        nonlinear advection terms are formed in physical space from
        Fourier-space gradients, transformed back to spectral space, radially
        binned by ``freq_r``, and accumulated from high to low wavenumber.
    """
    dims = _validate_spatial_dims(dim)
    freq_dims = [f"freq_{name}" for name in dims]
    u_hat, v_hat, du_dx, du_dy, dv_dx, dv_dy = _fourier_uv_gradients(
        u, v, dims, window=window, detrend=detrend
    )
    # Reconstruct physical-space velocities from the (possibly windowed and
    # detrended) spectra so the advection product mixes fields with matching
    # preprocessing — using raw u/v here would bias the transfer when
    # window/detrend are non-None (which they are by default).
    u_proc = _ifft2(u_hat, freq_dims, u)
    v_proc = _ifft2(v_hat, freq_dims, v)
    phi_u = u_proc * du_dx + v_proc * du_dy
    phi_v = u_proc * dv_dx + v_proc * dv_dy
    phi_u_hat = _fft2(phi_u, dims, window=None, detrend=None)
    phi_v_hat = _fft2(phi_v, dims, window=None, detrend=None)
    norm = float(np.prod([u.sizes[name] for name in dims]) ** 2)
    transfer_2d = (
        -np.real(np.conj(u_hat) * phi_u_hat + np.conj(v_hat) * phi_v_hat) / norm
    )
    transfer_2d.name = "transfer_2d"
    return _spectral_flux_dataset(
        transfer_2d, freq_dims, avg_dims=_as_list(avg_dims), return_2d=return_2d
    )

enstrophy_spectral_flux

enstrophy_spectral_flux(u: DataArray, v: DataArray, *, dim: Sequence[str], window: str | None = 'tukey', detrend: str | None = 'linear', avg_dims: str | Sequence[str] | None = None, return_2d: bool = False) -> Dataset

Enstrophy spectral flux Π_Z(k).

Parameters:

Name	Type	Description	Default
u	DataArray	Zonal velocity component with both spatial dimensions in dim.	required
v	DataArray	Meridional velocity component on the same grid as u.	required
dim	Sequence[str]	Two spatial dimensions to Fourier transform.	required
window	str | None	Optional window forwarded to :func:xrft.fft.	'tukey'
detrend	str | None	Optional detrending forwarded to :func:xrft.fft.	'linear'
avg_dims	str | Sequence[str] | None	Optional non-spectral dimensions to average before radial integration.	None
return_2d	bool	If True, include the unbinned two-dimensional transfer field as transfer_2d.	False

Returns:

Type	Description
Dataset	Dataset with radially summed enstrophy transfer and cumulative
Dataset	flux. transfer_2d is included when requested.

Notes

Vorticity is computed spectrally as ζ̂ = 2πi(k_x v̂ - k_y û). Positive flux denotes downscale enstrophy transfer, accumulated from high to low radial wavenumber with the same convention as :func:ke_spectral_flux.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/fourier.py

def enstrophy_spectral_flux(
    u: xr.DataArray,
    v: xr.DataArray,
    *,
    dim: Sequence[str],
    window: str | None = "tukey",
    detrend: str | None = "linear",
    avg_dims: str | Sequence[str] | None = None,
    return_2d: bool = False,
) -> xr.Dataset:
    """Enstrophy spectral flux ``Π_Z(k)``.

    Args:
        u: Zonal velocity component with both spatial dimensions in ``dim``.
        v: Meridional velocity component on the same grid as ``u``.
        dim: Two spatial dimensions to Fourier transform.
        window: Optional window forwarded to :func:`xrft.fft`.
        detrend: Optional detrending forwarded to :func:`xrft.fft`.
        avg_dims: Optional non-spectral dimensions to average before radial
            integration.
        return_2d: If ``True``, include the unbinned two-dimensional transfer
            field as ``transfer_2d``.

    Returns:
        Dataset with radially summed enstrophy ``transfer`` and cumulative
        ``flux``. ``transfer_2d`` is included when requested.

    Notes:
        Vorticity is computed spectrally as ``ζ̂ = 2πi(k_x v̂ - k_y û)``.
        Positive ``flux`` denotes downscale enstrophy transfer, accumulated
        from high to low radial wavenumber with the same convention as
        :func:`ke_spectral_flux`.
    """
    dims = _validate_spatial_dims(dim)
    freq_dims = [f"freq_{name}" for name in dims]
    u_hat, v_hat, _du_dx, _du_dy, _dv_dx, _dv_dy = _fourier_uv_gradients(
        u, v, dims, window=window, detrend=detrend
    )
    kx = v_hat[freq_dims[0]]
    ky = u_hat[freq_dims[1]]
    zeta_hat = 2.0j * np.pi * (kx * v_hat - ky * u_hat)
    dzeta_dx = _gradient_from_hat(zeta_hat, freq_dims[0], freq_dims, u)
    dzeta_dy = _gradient_from_hat(zeta_hat, freq_dims[1], freq_dims, u)
    # Match the windowed/detrended preprocessing of zeta_hat in the
    # advection product — see ke_spectral_flux for the same fix.
    u_proc = _ifft2(u_hat, freq_dims, u)
    v_proc = _ifft2(v_hat, freq_dims, v)
    adv_zeta = u_proc * dzeta_dx + v_proc * dzeta_dy
    adv_zeta_hat = _fft2(adv_zeta, dims, window=None, detrend=None)
    norm = float(np.prod([u.sizes[name] for name in dims]) ** 2)
    transfer_2d = -np.real(np.conj(zeta_hat) * adv_zeta_hat) / norm
    transfer_2d.name = "transfer_2d"
    return _spectral_flux_dataset(
        transfer_2d, freq_dims, avg_dims=_as_list(avg_dims), return_2d=return_2d
    )

Matrix decompositions

These return xrops.learn.XarrayEstimator instances configured for the named decomposition (see scikit-learn on xarray).

pca

pca(sample_dims: Hashable | Sequence[Hashable], n_components: int | float | str | None = None, *, whiten: bool = False, random_state: int | None = None, **kwargs: Any) -> XarrayEstimator

Principal Component Analysis on N-D xarray inputs.

Returns:

Type	Description
XarrayEstimator	An unfitted :class:XarrayEstimator wrapping
XarrayEstimator	class:sklearn.decomposition.PCA. Call .fit(da) /
XarrayEstimator	.transform(da) / .fit_transform(da) /
XarrayEstimator	.inverse_transform(scores) as you would on the raw sklearn
XarrayEstimator	object — inputs and outputs are :class:xr.DataArray.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/decompose.py

def pca(
    sample_dims: Hashable | Sequence[Hashable],
    n_components: int | float | str | None = None,
    *,
    whiten: bool = False,
    random_state: int | None = None,
    **kwargs: Any,
) -> XarrayEstimator:
    """Principal Component Analysis on N-D xarray inputs.

    Returns:
        An unfitted :class:`XarrayEstimator` wrapping
        :class:`sklearn.decomposition.PCA`. Call ``.fit(da)`` /
        ``.transform(da)`` / ``.fit_transform(da)`` /
        ``.inverse_transform(scores)`` as you would on the raw sklearn
        object — inputs and outputs are :class:`xr.DataArray`.
    """
    return XarrayEstimator(
        _sklearn("decomposition", "PCA")(
            n_components=n_components,
            whiten=whiten,
            random_state=random_state,
            **kwargs,
        ),
        sample_dims=sample_dims,
        feature_dim="component",
    )

eof

eof(sample_dims: Hashable | Sequence[Hashable], n_modes: int | None = None, *, whiten: bool = False, random_state: int | None = None, **kwargs: Any) -> XarrayEstimator

Empirical Orthogonal Functions — geophysical PCA.

Identical numerics to :func:pca, but the new feature axis is named "mode" rather than "component" to match the geosciences convention.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/decompose.py

def eof(
    sample_dims: Hashable | Sequence[Hashable],
    n_modes: int | None = None,
    *,
    whiten: bool = False,
    random_state: int | None = None,
    **kwargs: Any,
) -> XarrayEstimator:
    """Empirical Orthogonal Functions — geophysical PCA.

    Identical numerics to :func:`pca`, but the new feature axis is
    named ``"mode"`` rather than ``"component"`` to match the
    geosciences convention.
    """
    return XarrayEstimator(
        _sklearn("decomposition", "PCA")(
            n_components=n_modes,
            whiten=whiten,
            random_state=random_state,
            **kwargs,
        ),
        sample_dims=sample_dims,
        feature_dim="mode",
    )

ica

ica(sample_dims: Hashable | Sequence[Hashable], n_components: int | None = None, *, random_state: int | None = None, **kwargs: Any) -> XarrayEstimator

Independent Component Analysis (FastICA).

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/decompose.py

def ica(
    sample_dims: Hashable | Sequence[Hashable],
    n_components: int | None = None,
    *,
    random_state: int | None = None,
    **kwargs: Any,
) -> XarrayEstimator:
    """Independent Component Analysis (FastICA)."""
    return XarrayEstimator(
        _sklearn("decomposition", "FastICA")(
            n_components=n_components, random_state=random_state, **kwargs
        ),
        sample_dims=sample_dims,
        feature_dim="component",
    )

nmf

nmf(sample_dims: Hashable | Sequence[Hashable], n_components: int | None = None, *, init: str | None = None, random_state: int | None = None, **kwargs: Any) -> XarrayEstimator

Non-negative Matrix Factorisation. Input must be non-negative.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/decompose.py

def nmf(
    sample_dims: Hashable | Sequence[Hashable],
    n_components: int | None = None,
    *,
    init: str | None = None,
    random_state: int | None = None,
    **kwargs: Any,
) -> XarrayEstimator:
    """Non-negative Matrix Factorisation. Input must be non-negative."""
    return XarrayEstimator(
        _sklearn("decomposition", "NMF")(
            n_components=n_components,
            init=init,
            random_state=random_state,
            **kwargs,
        ),
        sample_dims=sample_dims,
        feature_dim="component",
    )

kmeans

kmeans(sample_dims: Hashable | Sequence[Hashable], n_clusters: int = 8, *, n_init: int | str = 'auto', random_state: int | None = None, **kwargs: Any) -> XarrayEstimator

K-Means clustering. predict(da) returns 1-D labels along sample_dims.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/decompose.py

def kmeans(
    sample_dims: Hashable | Sequence[Hashable],
    n_clusters: int = 8,
    *,
    n_init: int | str = "auto",
    random_state: int | None = None,
    **kwargs: Any,
) -> XarrayEstimator:
    """K-Means clustering. ``predict(da)`` returns 1-D labels along
    ``sample_dims``."""
    return XarrayEstimator(
        _sklearn("cluster", "KMeans")(
            n_clusters=n_clusters,
            n_init=n_init,
            random_state=random_state,
            **kwargs,
        ),
        sample_dims=sample_dims,
        feature_dim="component",
    )

Mask morphology

binary_opening_2d

binary_opening_2d(mask: DataArray, *, footprint: Footprint = 1, lon: str = 'lon', lat: str = 'lat') -> DataArray

Morphologically open a 2-D boolean mask slice-by-slice.

Parameters:

Name	Type	Description	Default
mask	DataArray	Boolean mask with lat and lon dimensions.	required
footprint	Footprint	Structuring element specification.	1
lon	str	Longitude dimension name.	'lon'
lat	str	Latitude dimension name.	'lat'

Returns:

Type	Description
DataArray	Same-shaped boolean mask after erosion followed by dilation.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/morphology.py

def binary_opening_2d(
    mask: xr.DataArray,
    *,
    footprint: Footprint = 1,
    lon: str = "lon",
    lat: str = "lat",
) -> xr.DataArray:
    """Morphologically open a 2-D boolean mask slice-by-slice.

    Args:
        mask: Boolean mask with ``lat`` and ``lon`` dimensions.
        footprint: Structuring element specification.
        lon: Longitude dimension name.
        lat: Latitude dimension name.

    Returns:
        Same-shaped boolean mask after erosion followed by dilation.
    """
    morph = _require_skimage()
    fp = _resolve_footprint(footprint)
    return _wrap2d(
        mask,
        lon=lon,
        lat=lat,
        fn=lambda m: morph.opening(m, footprint=fp),
    )

binary_closing_2d

binary_closing_2d(mask: DataArray, *, footprint: Footprint = 1, lon: str = 'lon', lat: str = 'lat') -> DataArray

Morphologically close a 2-D boolean mask slice-by-slice.

Parameters:

Name	Type	Description	Default
mask	DataArray	Boolean mask with lat and lon dimensions.	required
footprint	Footprint	Structuring element specification.	1
lon	str	Longitude dimension name.	'lon'
lat	str	Latitude dimension name.	'lat'

Returns:

Type	Description
DataArray	Same-shaped boolean mask after dilation followed by erosion.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/morphology.py

def binary_closing_2d(
    mask: xr.DataArray,
    *,
    footprint: Footprint = 1,
    lon: str = "lon",
    lat: str = "lat",
) -> xr.DataArray:
    """Morphologically close a 2-D boolean mask slice-by-slice.

    Args:
        mask: Boolean mask with ``lat`` and ``lon`` dimensions.
        footprint: Structuring element specification.
        lon: Longitude dimension name.
        lat: Latitude dimension name.

    Returns:
        Same-shaped boolean mask after dilation followed by erosion.
    """
    morph = _require_skimage()
    fp = _resolve_footprint(footprint)
    return _wrap2d(
        mask,
        lon=lon,
        lat=lat,
        fn=lambda m: morph.closing(m, footprint=fp),
    )

remove_small_holes_2d

remove_small_holes_2d(mask: DataArray, *, area: int = 4, lon: str = 'lon', lat: str = 'lat') -> DataArray

Flip False connected components smaller than area to True.

True means "masked / needs filling" in xrops interpolation workflows, so this removes small unmasked islands inside larger gaps.

Parameters:

Name	Type	Description	Default
mask	DataArray	Boolean mask with lat and lon dimensions.	required
area	int	False regions smaller than this pixel count are filled.	4
lon	str	Longitude dimension name.	'lon'
lat	str	Latitude dimension name.	'lat'

Returns:

Type	Description
DataArray	Same-shaped boolean mask with small False holes filled.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/morphology.py

def remove_small_holes_2d(
    mask: xr.DataArray,
    *,
    area: int = 4,
    lon: str = "lon",
    lat: str = "lat",
) -> xr.DataArray:
    """Flip False connected components smaller than ``area`` to True.

    ``True`` means "masked / needs filling" in xrops interpolation
    workflows, so this removes small unmasked islands inside larger gaps.

    Args:
        mask: Boolean mask with ``lat`` and ``lon`` dimensions.
        area: False regions smaller than this pixel count are filled.
        lon: Longitude dimension name.
        lat: Latitude dimension name.

    Returns:
        Same-shaped boolean mask with small False holes filled.
    """
    _validate_area(area)
    return _wrap2d(
        mask,
        lon=lon,
        lat=lat,
        fn=lambda m: _remove_small_holes(m, area=area),
    )

remove_small_objects_2d

remove_small_objects_2d(mask: DataArray, *, area: int = 4, lon: str = 'lon', lat: str = 'lat') -> DataArray

Flip True connected components smaller than area to False.

True means "masked / needs filling" in xrops interpolation workflows, so this removes isolated masked specks.

Parameters:

Name	Type	Description	Default
mask	DataArray	Boolean mask with lat and lon dimensions.	required
area	int	True regions smaller than this pixel count are dropped.	4
lon	str	Longitude dimension name.	'lon'
lat	str	Latitude dimension name.	'lat'

Returns:

Type	Description
DataArray	Same-shaped boolean mask with small True objects removed.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/morphology.py

def remove_small_objects_2d(
    mask: xr.DataArray,
    *,
    area: int = 4,
    lon: str = "lon",
    lat: str = "lat",
) -> xr.DataArray:
    """Flip True connected components smaller than ``area`` to False.

    ``True`` means "masked / needs filling" in xrops interpolation
    workflows, so this removes isolated masked specks.

    Args:
        mask: Boolean mask with ``lat`` and ``lon`` dimensions.
        area: True regions smaller than this pixel count are dropped.
        lon: Longitude dimension name.
        lat: Latitude dimension name.

    Returns:
        Same-shaped boolean mask with small True objects removed.
    """
    _validate_area(area)
    return _wrap2d(
        mask,
        lon=lon,
        lat=lat,
        fn=lambda m: _remove_small_objects(m, area=area),
    )

clean_mask

clean_mask(mask: DataArray, *, fill_holes_area: int | None = 4, drop_objects_area: int | None = None, closing_footprint: Footprint | None = None, opening_footprint: Footprint | None = None, lon: str = 'lon', lat: str = 'lat') -> DataArray

Run the common mask-despeckling pipeline.

Steps run in fixed order when their corresponding keyword is not None: remove small holes, remove small objects, binary closing, then binary opening.

Parameters:

Name	Type	Description	Default
mask	DataArray	Boolean mask with lat and lon dimensions.	required
fill_holes_area	int | None	Optional hole area threshold.	4
drop_objects_area	int | None	Optional object area threshold.	None
closing_footprint	Footprint | None	Optional binary-closing footprint.	None
opening_footprint	Footprint | None	Optional binary-opening footprint.	None
lon	str	Longitude dimension name.	'lon'
lat	str	Latitude dimension name.	'lat'

Returns:

Type	Description
DataArray	Same-shaped cleaned boolean mask.

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/morphology.py

def clean_mask(
    mask: xr.DataArray,
    *,
    fill_holes_area: int | None = 4,
    drop_objects_area: int | None = None,
    closing_footprint: Footprint | None = None,
    opening_footprint: Footprint | None = None,
    lon: str = "lon",
    lat: str = "lat",
) -> xr.DataArray:
    """Run the common mask-despeckling pipeline.

    Steps run in fixed order when their corresponding keyword is not
    ``None``: remove small holes, remove small objects, binary closing,
    then binary opening.

    Args:
        mask: Boolean mask with ``lat`` and ``lon`` dimensions.
        fill_holes_area: Optional hole area threshold.
        drop_objects_area: Optional object area threshold.
        closing_footprint: Optional binary-closing footprint.
        opening_footprint: Optional binary-opening footprint.
        lon: Longitude dimension name.
        lat: Latitude dimension name.

    Returns:
        Same-shaped cleaned boolean mask.
    """
    out = mask
    if fill_holes_area is not None:
        out = remove_small_holes_2d(out, area=fill_holes_area, lon=lon, lat=lat)
    if drop_objects_area is not None:
        out = remove_small_objects_2d(out, area=drop_objects_area, lon=lon, lat=lat)
    if closing_footprint is not None:
        out = binary_closing_2d(out, footprint=closing_footprint, lon=lon, lat=lat)
    if opening_footprint is not None:
        out = binary_opening_2d(out, footprint=opening_footprint, lon=lon, lat=lat)
    return out

Axis remapping

remap_axis interpolates one dimension of a DataArray onto a new axis of target values. By default the source axis is the 1-D dimension coordinate and targets outside its range become NaN. Pass source_coords= (a coordinate name or a broadcastable DataArray) when the source levels vary per column — terrain-following or hybrid vertical grids — and extrapolate="nearest" to hold each column's end value beyond its own range. See the Vertical & axis remapping section of the interpolation API for the operator-level presets.

remap_axis

remap_axis(da: DataArray, *, source_dim: str, target_coords: DataArray | ndarray, target_name: str | None = None, method: str = 'linear', source_coords: str | DataArray | None = None, extrapolate: Extrapolate = 'nan') -> DataArray

Remap da along source_dim onto target_coords.

If target_coords is a :class:xr.DataArray and target_name is None, the new dim name is taken from target_coords.name; otherwise target_name (or source_dim) is used.

Every coordinate that depends on source_dim is dropped from the output; the new target coordinate is attached under the new name. The Layer-1 RemapAxis operator handles Dataset selection / per-variable looping.

Parameters:

Name	Type	Description	Default
da	DataArray	Numeric DataArray carrying source_dim.	required
source_dim	str	Name of the dimension being replaced.	required
target_coords	DataArray | ndarray	1-D target axis values.	required
target_name	str | None	Explicit name for the new dimension.	None
method	str	"linear" or "nearest".	'linear'
source_coords	str | DataArray | None	Source axis values. None (default) uses the 1-D dimension coordinate da[source_dim]. A coordinate name on da or a DataArray broadcastable against da supplies per-column levels — it must carry source_dim and may vary along any of the other dims (e.g. z_agl(time, level, y, x) on a terrain-following grid). Every column must be strictly monotone along source_dim.	None
extrapolate	Extrapolate	"nan" (default) — targets outside a column's range become NaN; "nearest" — they hold that column's end value.	'nan'

Returns:

Type	Description
DataArray	da with source_dim replaced by the target axis. When
DataArray	source_coords is given, chunked (dask) inputs stay lazy as
DataArray	long as source_dim sits in a single chunk.

Raises:

Type	Description
ValueError	If source_dim is missing, source_coords does not carry it, a column is not monotone, or extrapolate / method is unknown.
TypeError	If da is not numeric.

Example

Remap a model-level field with a 4-D height-above-ground coordinate onto fixed heights, holding the lowest / highest model level beyond each column's range:

>>> import numpy as np, xarray as xr
>>> from xrops.transforms import remap_axis
>>> level = np.arange(4)
>>> z_agl = 20.0 * (level[None, :, None, None] + 1) + np.zeros((2, 4, 3, 3))
>>> z_agl[..., 0, 0] += 15.0  # one column starts higher up
>>> t = xr.DataArray(
...     300.0 - 0.01 * z_agl,
...     dims=("time", "level", "y", "x"),
...     coords={"z_agl": (("time", "level", "y", "x"), z_agl)},
... )
>>> out = remap_axis(
...     t,
...     source_dim="level",
...     target_coords=np.array([10.0, 50.0, 90.0]),
...     target_name="height",
...     source_coords="z_agl",
...     extrapolate="nearest",
... )
>>> out.dims, out.sizes["height"]
(('time', 'height', 'y', 'x'), 3)
>>> float(out.isel(time=0, height=1, y=1, x=1))  # in range: linear
299.5

Source code in packages/xrtoolz-ops/src/xrops/transforms/_src/coord_remap.py

def remap_axis(
    da: xr.DataArray,
    *,
    source_dim: str,
    target_coords: xr.DataArray | np.ndarray,
    target_name: str | None = None,
    method: str = "linear",
    source_coords: str | xr.DataArray | None = None,
    extrapolate: Extrapolate = "nan",
) -> xr.DataArray:
    """Remap ``da`` along ``source_dim`` onto ``target_coords``.

    If ``target_coords`` is a :class:`xr.DataArray` and ``target_name``
    is None, the new dim name is taken from ``target_coords.name``;
    otherwise ``target_name`` (or ``source_dim``) is used.

    Every coordinate that depends on ``source_dim`` is dropped from the
    output; the new target coordinate is attached under the new name.
    The Layer-1 ``RemapAxis`` operator handles Dataset selection /
    per-variable looping.

    Args:
        da: Numeric DataArray carrying ``source_dim``.
        source_dim: Name of the dimension being replaced.
        target_coords: 1-D target axis values.
        target_name: Explicit name for the new dimension.
        method: ``"linear"`` or ``"nearest"``.
        source_coords: Source axis *values*. ``None`` (default) uses the
            1-D dimension coordinate ``da[source_dim]``. A coordinate name
            on ``da`` or a DataArray broadcastable against ``da`` supplies
            per-column levels — it must carry ``source_dim`` and may vary
            along any of the other dims (e.g. ``z_agl(time, level, y, x)``
            on a terrain-following grid). Every column must be strictly
            monotone along ``source_dim``.
        extrapolate: ``"nan"`` (default) — targets outside a column's
            range become ``NaN``; ``"nearest"`` — they hold that column's
            end value.

    Returns:
        ``da`` with ``source_dim`` replaced by the target axis. When
        ``source_coords`` is given, chunked (dask) inputs stay lazy as
        long as ``source_dim`` sits in a single chunk.

    Raises:
        ValueError: If ``source_dim`` is missing, ``source_coords`` does
            not carry it, a column is not monotone, or ``extrapolate`` /
            ``method`` is unknown.
        TypeError: If ``da`` is not numeric.

    Example:
        Remap a model-level field with a 4-D height-above-ground
        coordinate onto fixed heights, holding the lowest / highest model
        level beyond each column's range:

        ```pycon
        >>> import numpy as np, xarray as xr
        >>> from xrops.transforms import remap_axis
        >>> level = np.arange(4)
        >>> z_agl = 20.0 * (level[None, :, None, None] + 1) + np.zeros((2, 4, 3, 3))
        >>> z_agl[..., 0, 0] += 15.0  # one column starts higher up
        >>> t = xr.DataArray(
        ...     300.0 - 0.01 * z_agl,
        ...     dims=("time", "level", "y", "x"),
        ...     coords={"z_agl": (("time", "level", "y", "x"), z_agl)},
        ... )
        >>> out = remap_axis(
        ...     t,
        ...     source_dim="level",
        ...     target_coords=np.array([10.0, 50.0, 90.0]),
        ...     target_name="height",
        ...     source_coords="z_agl",
        ...     extrapolate="nearest",
        ... )
        >>> out.dims, out.sizes["height"]
        (('time', 'height', 'y', 'x'), 3)
        >>> float(out.isel(time=0, height=1, y=1, x=1))  # in range: linear
        299.5

        ```
    """
    if source_dim not in da.dims:
        raise ValueError(
            f"source_dim {source_dim!r} not in DataArray dims {tuple(da.dims)}"
        )
    if extrapolate not in _kernels._EXTRAPOLATE_MODES:
        raise ValueError(
            f"unknown extrapolate {extrapolate!r}; expected one of "
            f"{_kernels._EXTRAPOLATE_MODES!r}"
        )
    if source_coords is None and source_dim not in da.coords:
        raise ValueError(
            f"DataArray must carry a coordinate named {source_dim!r} "
            "for the source axis values"
        )
    if not np.issubdtype(da.dtype, np.number):
        raise TypeError(
            f"remap_axis requires numeric data; got dtype {da.dtype}. "
            "Drop or convert the variable before calling remap_axis."
        )

    if isinstance(target_coords, xr.DataArray):
        new_name = target_name or target_coords.name or source_dim
        tgt = np.asarray(target_coords.values, dtype=float)
    else:
        new_name = target_name or source_dim
        tgt = np.asarray(target_coords, dtype=float)

    if source_coords is not None or extrapolate != "nan":
        return _remap_axis_columns(
            da,
            source_dim=source_dim,
            tgt=tgt,
            new_name=str(new_name),
            method=method,
            source_coords=source_coords,
            extrapolate=extrapolate,
        )

    src = np.asarray(da[source_dim].values, dtype=float)

    axis = da.get_axis_num(source_dim)
    new_values = _kernels.remap_axis(
        da.values,
        axis=axis,
        source_coords=src,
        target_coords=tgt,
        method=method,
    )
    new_dims = tuple(new_name if d == source_dim else d for d in da.dims)
    new_coords = {
        cname: c
        for cname, c in da.coords.items()
        if source_dim not in c.dims and cname != source_dim
    }
    new_coords[new_name] = xr.DataArray(tgt, dims=(new_name,), name=new_name)
    return xr.DataArray(
        new_values,
        dims=new_dims,
        coords=new_coords,
        attrs=dict(da.attrs),
        name=da.name,
    )