Advanced Analysis Methods

Advanced analysis methods for weather and air quality, including temporal resampling, climatology, and Kolmogorov-Zurbenko (KZ) filters.

Advanced analysis methods for weather and air quality (Aero Protocol Compliant).

anomalies(data, freq='month', dim='time')

Compute anomalies by subtracting the climatology (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data with a time-like coordinate. freq : str, optional Climatology frequency ('season', 'month', 'dayofyear', 'hour'). Default is 'month'. dim : str, optional Dimension along which to compute the anomalies. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] Anomalies (data - climatology).

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=366*2, freq="D") da = xr.DataArray(np.random.rand(732), coords={"time": times}, dims="time") monthly_anom = anomalies(da, freq="month")

Source code in src/monet_stats/analysis.py

def anomalies(
    data: Union[xr.DataArray, xr.Dataset],
    freq: str = "month",
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute anomalies by subtracting the climatology (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with a time-like coordinate.
    freq : str, optional
        Climatology frequency ('season', 'month', 'dayofyear', 'hour').
        Default is 'month'.
    dim : str, optional
        Dimension along which to compute the anomalies. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Anomalies (data - climatology).

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=366*2, freq="D")
    >>> da = xr.DataArray(np.random.rand(732), coords={"time": times}, dims="time")
    >>> monthly_anom = anomalies(da, freq="month")
    """
    group = f"{dim}.{freq}"
    # Compute climatology (this reduces the 'dim' dimension)
    climo = climatology(data, freq=freq, method="mean", dim=dim)

    # Subtract climatology using groupby broadcasting
    # data.groupby(group) - climo aligns the 'freq' coordinate automatically
    res = data.groupby(group) - climo

    return _update_history(res, f"Anomalies ({freq})")

climatology(data, freq='season', method='mean', dim='time')

Compute climatological statistics (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data with a time-like coordinate. freq : str, optional Climatology frequency ('season', 'month', 'dayofyear', 'hour'). Default is 'season'. method : str, optional Statistical method to apply ('mean', 'std', 'min', 'max', 'median'). Default is 'mean'. dim : str, optional Dimension along which to compute climatology. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] Climatological statistics.

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=365*2, freq="D") da = xr.DataArray(np.random.rand(730), coords={"time": times}, dims="time") seasonal_climo = climatology(da, freq="season", method="mean")

Source code in src/monet_stats/analysis.py

def climatology(
    data: Union[xr.DataArray, xr.Dataset],
    freq: str = "season",
    method: str = "mean",
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute climatological statistics (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with a time-like coordinate.
    freq : str, optional
        Climatology frequency ('season', 'month', 'dayofyear', 'hour').
        Default is 'season'.
    method : str, optional
        Statistical method to apply ('mean', 'std', 'min', 'max', 'median').
        Default is 'mean'.
    dim : str, optional
        Dimension along which to compute climatology. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Climatological statistics.

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=365*2, freq="D")
    >>> da = xr.DataArray(np.random.rand(730), coords={"time": times}, dims="time")
    >>> seasonal_climo = climatology(da, freq="season", method="mean")
    """
    group = f"{dim}.{freq}"
    # Use native xarray groupby methods for better performance and Dask integration
    grouped = data.groupby(group)
    if hasattr(grouped, method):
        result = getattr(grouped, method)(dim=dim)
    else:
        # Fallback for methods not directly implemented on GroupBy
        result = grouped.reduce(getattr(np, f"nan{method}" if "nan" not in method else method), dim=dim)

    return _update_history(result, f"Climatology ({freq}) using {method}")

detrend(data, method='linear', dim='time')

Remove trend from data (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. method : str, optional Detrending method ('linear', 'constant'). - 'linear': least-squares linear detrend. - 'constant': subtract mean. Default is 'linear'. dim : str, optional Dimension along which to detrend. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] Detrended data.

Examples

import xarray as xr import numpy as np da = xr.DataArray(np.arange(10) + np.random.randn(10), dims="time") detrended = detrend(da, method="linear")

Source code in src/monet_stats/analysis.py

def detrend(
    data: Union[xr.DataArray, xr.Dataset],
    method: str = "linear",
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Remove trend from data (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data.
    method : str, optional
        Detrending method ('linear', 'constant').
        - 'linear': least-squares linear detrend.
        - 'constant': subtract mean.
        Default is 'linear'.
    dim : str, optional
        Dimension along which to detrend. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Detrended data.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> da = xr.DataArray(np.arange(10) + np.random.randn(10), dims="time")
    >>> detrended = detrend(da, method="linear")
    """
    if method == "constant":
        res = data - data.mean(dim=dim)
        return _update_history(res, "Detrended (constant)")

    if method == "linear":
        from scipy.signal import detrend as scipy_detrend

        if isinstance(data, xr.Dataset):
            res = data.map(detrend, method=method, dim=dim)
            return _update_history(res, "Detrended (linear)")

        # Core dimensions for apply_ufunc must be a single chunk if using dask
        is_lazy = False
        if hasattr(data.data, "chunks"):
            is_lazy = True

        if is_lazy:
            data = data.chunk({dim: -1})

        res = xr.apply_ufunc(
            scipy_detrend,
            data,
            input_core_dims=[[dim]],
            output_core_dims=[[dim]],
            kwargs={"axis": -1},
            dask="parallelized",
            output_dtypes=[data.dtype],
            keep_attrs=True,
        )
        return _update_history(res, "Detrended (linear)")

    raise ValueError(f"Unknown detrending method: {method}")

diurnal_cycle(data, method='mean', dim='time')

Compute the diurnal cycle (average hourly profile) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data with a time-like coordinate. method : str, optional Statistical method to apply ('mean', 'median', 'std'). Default is 'mean'. dim : str, optional Dimension along which to compute the cycle. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] Diurnal cycle (24 values, one for each hour).

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=24*10, freq="h") da = xr.DataArray(np.random.rand(240), coords={"time": times}, dims="time") cycle = diurnal_cycle(da, method="mean")

Source code in src/monet_stats/analysis.py

def diurnal_cycle(
    data: Union[xr.DataArray, xr.Dataset],
    method: str = "mean",
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute the diurnal cycle (average hourly profile) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with a time-like coordinate.
    method : str, optional
        Statistical method to apply ('mean', 'median', 'std'). Default is 'mean'.
    dim : str, optional
        Dimension along which to compute the cycle. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Diurnal cycle (24 values, one for each hour).

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=24*10, freq="h")
    >>> da = xr.DataArray(np.random.rand(240), coords={"time": times}, dims="time")
    >>> cycle = diurnal_cycle(da, method="mean")
    """
    return climatology(data, freq="hour", method=method, dim=dim)

exceedance_count(data, threshold, dim='time', axis=None)

Count exceedances of a threshold (Aero Protocol).

Parameters

data : xarray.DataArray, xarray.Dataset, or numpy.ndarray Input data. threshold : float Value above which an exceedance is counted. dim : str, optional Dimension along which to count exceedances (xarray only). Default is 'time'. axis : int, optional Axis along which to count exceedances (numpy only). Default is None (all).

Returns

Union[xr.DataArray, xr.Dataset, np.ndarray] Number of exceedances.

Examples

import xarray as xr da = xr.DataArray([1, 5, 2, 6, 3]) exceedance_count(da, threshold=4) array(2)

Source code in src/monet_stats/analysis.py

def exceedance_count(
    data: Union[xr.DataArray, xr.Dataset, np.ndarray],
    threshold: float,
    dim: str = "time",
    axis: Optional[int] = None,
) -> Union[xr.DataArray, xr.Dataset, np.ndarray]:
    """
    Count exceedances of a threshold (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray, xarray.Dataset, or numpy.ndarray
        Input data.
    threshold : float
        Value above which an exceedance is counted.
    dim : str, optional
        Dimension along which to count exceedances (xarray only). Default is 'time'.
    axis : int, optional
        Axis along which to count exceedances (numpy only). Default is None (all).

    Returns
    -------
    Union[xr.DataArray, xr.Dataset, np.ndarray]
        Number of exceedances.

    Examples
    --------
    >>> import xarray as xr
    >>> da = xr.DataArray([1, 5, 2, 6, 3])
    >>> exceedance_count(da, threshold=4)
    <xarray.DataArray ()>
    array(2)
    """
    if isinstance(data, (xr.DataArray, xr.Dataset)):
        res = (data > threshold).sum(dim=dim)
        return _update_history(res, f"Exceedance count (threshold={threshold})")

    res = np.sum(data > threshold, axis=axis)
    return res

fft_analysis(data, dim='time', output='psd')

Perform Fast Fourier Transform (FFT) analysis (Aero Protocol).

Parameters

data : xarray.DataArray Input data. dim : str, optional Dimension along which to perform FFT. Default is 'time'. output : str, optional Type of output to return: - 'psd': Power Spectral Density (magnitude squared of FFT). - 'magnitude': Magnitude of FFT. - 'complex': Complex FFT results. Default is 'psd'.

Returns

xarray.DataArray FFT results. The coordinate for 'dim' is replaced by frequency indices.

Examples

import xarray as xr import numpy as np t = np.linspace(0, 10, 100) signal = np.sin(2 * np.pi * 1.5 * t) # 1.5 Hz signal da = xr.DataArray(signal, coords={"time": t}, dims="time") psd = fft_analysis(da, dim="time", output="psd")

Source code in src/monet_stats/analysis.py

def fft_analysis(
    data: xr.DataArray,
    dim: str = "time",
    output: str = "psd",
) -> xr.DataArray:
    """
    Perform Fast Fourier Transform (FFT) analysis (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray
        Input data.
    dim : str, optional
        Dimension along which to perform FFT. Default is 'time'.
    output : str, optional
        Type of output to return:
        - 'psd': Power Spectral Density (magnitude squared of FFT).
        - 'magnitude': Magnitude of FFT.
        - 'complex': Complex FFT results.
        Default is 'psd'.

    Returns
    -------
    xarray.DataArray
        FFT results. The coordinate for 'dim' is replaced by frequency indices.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> t = np.linspace(0, 10, 100)
    >>> signal = np.sin(2 * np.pi * 1.5 * t)  # 1.5 Hz signal
    >>> da = xr.DataArray(signal, coords={"time": t}, dims="time")
    >>> psd = fft_analysis(da, dim="time", output="psd")
    """

    def _fft_wrapper(x: np.ndarray) -> np.ndarray:
        return np.fft.fft(x, axis=-1)

    # Core dimensions for apply_ufunc must be a single chunk if using dask
    is_lazy = False
    if hasattr(data.data, "chunks"):
        is_lazy = True

    if is_lazy:
        data = data.chunk({dim: -1})

    res_complex = xr.apply_ufunc(
        _fft_wrapper,
        data,
        input_core_dims=[[dim]],
        output_core_dims=[[dim]],
        dask="parallelized",
        output_dtypes=[np.complex128],
    )

    if output == "complex":
        res = res_complex
    elif output == "magnitude":
        res = np.abs(res_complex)
    elif output == "psd":
        res = np.abs(res_complex) ** 2
    else:
        raise ValueError(f"Unknown output type: {output}")

    # Update coordinate to frequency index
    n = data.sizes[dim]
    res = res.assign_coords({dim: np.arange(n)})

    return _update_history(res, f"FFT analysis (output={output})")

kz_filter(data, m, k, dim='time', axis=-1)

Kolmogorov-Zurbenko (KZ) filter (Aero Protocol).

The KZ filter is a low-pass filter implemented as k iterations of a moving average of window size m.

Parameters

data : xarray.DataArray, xarray.Dataset, or numpy.ndarray Input data. m : int Window size for the moving average (must be an odd integer for symmetry). k : int Number of iterations. dim : str, optional Dimension along which to apply the filter (xarray only). Default is 'time'. axis : int, optional Axis along which to apply the filter (numpy only). Default is -1.

Returns

Union[xr.DataArray, xr.Dataset, np.ndarray] Filtered data.

Notes

The KZ filter is widely used in air quality analysis to separate different time scales in a time series (e.g., seasonal, long-term, and short-term).

Examples

import numpy as np x = np.random.rand(100) filtered = kz_filter(x, m=5, k=3)

Source code in src/monet_stats/analysis.py

def kz_filter(
    data: Union[xr.DataArray, xr.Dataset, np.ndarray],
    m: int,
    k: int,
    dim: str = "time",
    axis: int = -1,
) -> Union[xr.DataArray, xr.Dataset, np.ndarray]:
    """
    Kolmogorov-Zurbenko (KZ) filter (Aero Protocol).

    The KZ filter is a low-pass filter implemented as k iterations of a moving
    average of window size m.

    Parameters
    ----------
    data : xarray.DataArray, xarray.Dataset, or numpy.ndarray
        Input data.
    m : int
        Window size for the moving average (must be an odd integer for symmetry).
    k : int
        Number of iterations.
    dim : str, optional
        Dimension along which to apply the filter (xarray only). Default is 'time'.
    axis : int, optional
        Axis along which to apply the filter (numpy only). Default is -1.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset, np.ndarray]
        Filtered data.

    Notes
    -----
    The KZ filter is widely used in air quality analysis to separate different
    time scales in a time series (e.g., seasonal, long-term, and short-term).

    Examples
    --------
    >>> import numpy as np
    >>> x = np.random.rand(100)
    >>> filtered = kz_filter(x, m=5, k=3)
    """
    if m % 2 == 0:
        warnings.warn("KZ filter window size m should ideally be odd for symmetry.", stacklevel=2)

    # A KZ filter (m, k) is equivalent to a convolution with a kernel
    # that is the k-fold convolution of a boxcar of width m.
    boxcar = np.ones(m) / m
    kernel = boxcar
    for _ in range(k - 1):
        kernel = np.convolve(kernel, boxcar)

    if not isinstance(data, (xr.DataArray, xr.Dataset)):
        # Fast path for NumPy using single convolution
        res_np = np.asanyarray(data, dtype=float)
        # Use scipy.ndimage.convolve1d for multi-dimensional support and speed.
        # mode='constant', cval=np.nan ensures that edges where the kernel
        # overlaps with the boundary become NaN, matching Xarray's rolling behavior.
        return convolve1d(res_np, kernel, axis=axis, mode="constant", cval=np.nan)

    if isinstance(data, xr.Dataset):
        # Apply to each variable in the dataset
        res = data.map(kz_filter, m=m, k=k, dim=dim)
        return _update_history(res, f"KZ filter (m={m}, k={k})")

    # Xarray DataArray path (handles Dask natively via apply_ufunc)
    def _kz_wrapper(x: np.ndarray, kernel: np.ndarray) -> np.ndarray:
        return convolve1d(x, kernel, axis=-1, mode="constant", cval=np.nan)

    # Ensure dimension is in a single chunk for Dask-backed arrays
    is_lazy = False
    if isinstance(data, xr.DataArray):
        if hasattr(data.data, "chunks"):
            is_lazy = True
    elif isinstance(data, xr.Dataset):
        if any(hasattr(data[v].data, "chunks") for v in data.data_vars):
            is_lazy = True

    if is_lazy:
        data = data.chunk({dim: -1})

    res = xr.apply_ufunc(
        _kz_wrapper,
        data,
        input_core_dims=[[dim]],
        output_core_dims=[[dim]],
        kwargs={"kernel": kernel},
        dask="parallelized",
        output_dtypes=[float],
        keep_attrs=True,
    )

    return _update_history(res, f"KZ filter (m={m}, k={k})")

mda1(data, dim='time')

Compute Maximum Daily 1-hour Average (MDA1) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. Must have hourly frequency. dim : str, optional Dimension along which to compute. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] MDA1 values (one per day).

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=24*5, freq="h") da = xr.DataArray(np.random.rand(120), coords={"time": times}, dims="time") mda1_vals = mda1(da)

Source code in src/monet_stats/analysis.py

def mda1(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute Maximum Daily 1-hour Average (MDA1) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data. Must have hourly frequency.
    dim : str, optional
        Dimension along which to compute. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        MDA1 values (one per day).

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=24*5, freq="h")
    >>> da = xr.DataArray(np.random.rand(120), coords={"time": times}, dims="time")
    >>> mda1_vals = mda1(da)
    """
    res = data.resample({dim: "D"}).max()
    return _update_history(res, "MDA1 (Maximum Daily 1-hour Average)")

mda8(data, dim='time', min_periods=6, center=False)

Compute Maximum Daily 8-hour Average (MDA8) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. Must have hourly frequency. dim : str, optional Dimension along which to compute. Default is 'time'. min_periods : int, optional Minimum number of observations for the 8-hour rolling mean. Default is 6. center : bool, optional Whether to center the 8-hour rolling window. Regulatory MDA8 (e.g., EPA) typically uses a non-centered (trailing) window. Default is False.

Returns

Union[xr.DataArray, xr.Dataset] MDA8 values (one per day).

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=24*5, freq="h") da = xr.DataArray(np.random.rand(120), coords={"time": times}, dims="time") ozone_mda8 = mda8(da)

Source code in src/monet_stats/analysis.py

def mda8(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
    min_periods: int = 6,
    center: bool = False,
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute Maximum Daily 8-hour Average (MDA8) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data. Must have hourly frequency.
    dim : str, optional
        Dimension along which to compute. Default is 'time'.
    min_periods : int, optional
        Minimum number of observations for the 8-hour rolling mean. Default is 6.
    center : bool, optional
        Whether to center the 8-hour rolling window. Regulatory MDA8 (e.g., EPA)
        typically uses a non-centered (trailing) window. Default is False.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        MDA8 values (one per day).

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=24*5, freq="h")
    >>> da = xr.DataArray(np.random.rand(120), coords={"time": times}, dims="time")
    >>> ozone_mda8 = mda8(da)
    """
    rolling_8h = rolling_mean_8h(data, dim=dim, min_periods=min_periods, center=center)
    # Group by day and take maximum
    res = rolling_8h.resample({dim: "D"}).max()
    return _update_history(res, "MDA8 (Maximum Daily 8-hour Average)")

monthly_climatology(data, dim='time', method='mean')

Compute monthly climatology (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data with a time-like coordinate. dim : str, optional Dimension along which to compute the climatology. Default is 'time'. method : str, optional Statistical method to apply ('mean', 'std', 'min', 'max', 'median'). Default is 'mean'.

Returns

Union[xr.DataArray, xr.Dataset] Monthly climatology (12 values, one for each month).

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=366*2, freq="D") da = xr.DataArray(np.random.rand(732), coords={"time": times}, dims="time") m_climo = monthly_climatology(da)

Source code in src/monet_stats/analysis.py

def monthly_climatology(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
    method: str = "mean",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute monthly climatology (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with a time-like coordinate.
    dim : str, optional
        Dimension along which to compute the climatology. Default is 'time'.
    method : str, optional
        Statistical method to apply ('mean', 'std', 'min', 'max', 'median').
        Default is 'mean'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Monthly climatology (12 values, one for each month).

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=366*2, freq="D")
    >>> da = xr.DataArray(np.random.rand(732), coords={"time": times}, dims="time")
    >>> m_climo = monthly_climatology(da)
    """
    # Shortcut to the general climatology function with freq='month'
    res = climatology(data, freq="month", method=method, dim=dim)
    return _update_history(res, f"Monthly climatology using {method}")

peak_timing(data, dim='time')

Identify the coordinate value of the maximum (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. dim : str, optional Dimension along which to find the peak. Default is 'time'.

Returns

Union[xr.DataArray, xr.Dataset] Coordinate values where the maximum occurs.

Examples

import xarray as xr import pandas as pd times = pd.date_range("2020-01-01", periods=24, freq="h") da = xr.DataArray(np.random.rand(24), coords={"time": times}, dims="time") peak_hour = peak_timing(da, dim="time")

Source code in src/monet_stats/analysis.py

def peak_timing(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Identify the coordinate value of the maximum (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data.
    dim : str, optional
        Dimension along which to find the peak. Default is 'time'.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Coordinate values where the maximum occurs.

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> times = pd.date_range("2020-01-01", periods=24, freq="h")
    >>> da = xr.DataArray(np.random.rand(24), coords={"time": times}, dims="time")
    >>> peak_hour = peak_timing(da, dim="time")
    """
    # Ensure dimension is in a single chunk for Dask-backed arrays
    is_lazy = False
    if isinstance(data, xr.DataArray):
        if hasattr(data.data, "chunks"):
            is_lazy = True
    elif isinstance(data, xr.Dataset):
        if any(hasattr(data[v].data, "chunks") for v in data.data_vars):
            is_lazy = True

    if is_lazy:
        data = data.chunk({dim: -1})

    # idxmax returns the coordinate of the maximum
    res = data.idxmax(dim=dim)
    return _update_history(res, f"Peak timing along {dim}")

percentile(data, q, dim='time', axis=None, **kwargs)

Compute percentiles (Aero Protocol).

Parameters

data : xarray.DataArray, xarray.Dataset, or numpy.ndarray Input data. q : float or list of float Percentile(s) to compute (0-100). dim : str, optional Dimension(s) over which to compute percentiles (xarray only). Default is 'time'. axis : int, optional Axis over which to compute percentiles (numpy only). Default is None. **kwargs : Any Additional keyword arguments passed to xarray.quantile or np.percentile.

Returns

Union[xr.DataArray, xr.Dataset, np.ndarray] Computed percentiles.

Source code in src/monet_stats/analysis.py

def percentile(
    data: Union[xr.DataArray, xr.Dataset, np.ndarray],
    q: Union[float, list, np.ndarray],
    dim: str = "time",
    axis: Optional[int] = None,
    **kwargs: Any,
) -> Union[xr.DataArray, xr.Dataset, np.ndarray]:
    """
    Compute percentiles (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray, xarray.Dataset, or numpy.ndarray
        Input data.
    q : float or list of float
        Percentile(s) to compute (0-100).
    dim : str, optional
        Dimension(s) over which to compute percentiles (xarray only). Default is 'time'.
    axis : int, optional
        Axis over which to compute percentiles (numpy only). Default is None.
    **kwargs : Any
        Additional keyword arguments passed to xarray.quantile or np.percentile.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset, np.ndarray]
        Computed percentiles.
    """
    if isinstance(data, (xr.DataArray, xr.Dataset)):
        # Ensure dimension is in a single chunk for Dask-backed arrays
        is_lazy = False
        if isinstance(data, xr.DataArray):
            if hasattr(data.data, "chunks"):
                is_lazy = True
        elif isinstance(data, xr.Dataset):
            if any(hasattr(data[v].data, "chunks") for v in data.data_vars):
                is_lazy = True

        if is_lazy:
            data = data.chunk({dim: -1})

        # xarray uses 0-1 for quantile, so divide by 100
        res = data.quantile(np.asanyarray(q) / 100.0, dim=dim, **kwargs)
        return _update_history(res, f"Percentile (q={q})")

    return np.percentile(data, q, axis=axis, **kwargs)

power_spectrum(data, dim='time', fs=1.0, window='hann', nperseg=None, **kwargs)

Compute power spectrum using Welch's method (Aero Protocol).

Welch's method computes an estimate of the power spectral density by dividing the data into overlapping segments, computing a periodogram for each segment and averaging the results.

Parameters

data : xarray.DataArray Input data. dim : str, optional Dimension along which to compute the spectrum. Default is 'time'. fs : float, optional Sampling frequency. Default is 1.0. window : str, optional Desired window to use. Default is 'hann'. nperseg : int, optional Length of each segment. Default is None (256). **kwargs : Any Additional keyword arguments passed to scipy.signal.welch.

Returns

xarray.DataArray Power spectral density. The 'dim' dimension is replaced by 'frequency'.

Examples

import xarray as xr import numpy as np t = np.linspace(0, 100, 1000) signal = np.sin(2 * np.pi * 0.1 * t) + np.random.randn(1000) * 2 da = xr.DataArray(signal, coords={"time": t}, dims="time") psd = power_spectrum(da, dim="time", fs=10.0)

Source code in src/monet_stats/analysis.py

def power_spectrum(
    data: xr.DataArray,
    dim: str = "time",
    fs: float = 1.0,
    window: str = "hann",
    nperseg: Optional[int] = None,
    **kwargs: Any,
) -> xr.DataArray:
    """
    Compute power spectrum using Welch's method (Aero Protocol).

    Welch's method computes an estimate of the power spectral density by
    dividing the data into overlapping segments, computing a periodogram for
    each segment and averaging the results.

    Parameters
    ----------
    data : xarray.DataArray
        Input data.
    dim : str, optional
        Dimension along which to compute the spectrum. Default is 'time'.
    fs : float, optional
        Sampling frequency. Default is 1.0.
    window : str, optional
        Desired window to use. Default is 'hann'.
    nperseg : int, optional
        Length of each segment. Default is None (256).
    **kwargs : Any
        Additional keyword arguments passed to scipy.signal.welch.

    Returns
    -------
    xarray.DataArray
        Power spectral density. The 'dim' dimension is replaced by 'frequency'.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> t = np.linspace(0, 100, 1000)
    >>> signal = np.sin(2 * np.pi * 0.1 * t) + np.random.randn(1000) * 2
    >>> da = xr.DataArray(signal, coords={"time": t}, dims="time")
    >>> psd = power_spectrum(da, dim="time", fs=10.0)
    """
    from scipy.signal import welch

    # Core dimensions for apply_ufunc must be a single chunk if using dask
    is_lazy = False
    if hasattr(data.data, "chunks"):
        is_lazy = True

    if is_lazy:
        data = data.chunk({dim: -1})

    def _welch_wrapper(x: np.ndarray, fs: float, window: str, nperseg: int, **kwargs: Any) -> np.ndarray:
        f, psd = welch(x, fs=fs, window=window, nperseg=nperseg, axis=-1, **kwargs)
        return psd

    # Get the number of frequency bins to set output_core_dims size
    # For real FFT, it's nperseg // 2 + 1
    if nperseg is None:
        nperseg = min(data.sizes[dim], 256)

    n_freq = nperseg // 2 + 1

    res = xr.apply_ufunc(
        _welch_wrapper,
        data,
        input_core_dims=[[dim]],
        output_core_dims=[["frequency"]],
        kwargs={"fs": fs, "window": window, "nperseg": nperseg, **kwargs},
        dask="parallelized",
        output_dtypes=[data.dtype],
        dask_gufunc_kwargs={"output_sizes": {"frequency": n_freq}},
    )

    # Assign frequency coordinates
    freqs = np.fft.rfftfreq(nperseg, d=1.0 / fs)
    res = res.assign_coords(frequency=freqs)

    return _update_history(res, "Power spectrum (Welch method)")

resample_data(data, freq='MS', method='mean', dim='time', **kwargs)

Resample data to a new temporal frequency (Aero Protocol).

Parameters

data : xarray.DataArray, xarray.Dataset, pandas.Series, or pandas.DataFrame Input data with a time-like index or coordinate. freq : str, optional Resampling frequency (e.g., 'MS' for monthly start, 'W' for weekly, 'D' for daily). Default is 'MS'. method : str, optional Statistical method to apply ('mean', 'sum', 'min', 'max', 'std', 'median'). Default is 'mean'. dim : str, optional Dimension along which to resample (xarray only). Default is 'time'. **kwargs : Any Additional keyword arguments passed to the resample method.

Returns

Union[xr.DataArray, xr.Dataset, pd.Series, pd.DataFrame] Resampled data.

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=100, freq="D") da = xr.DataArray(np.random.rand(100), coords={"time": times}, dims="time") monthly_mean = resample_data(da, freq="MS", method="mean")

Source code in src/monet_stats/analysis.py

def resample_data(
    data: Union[xr.DataArray, xr.Dataset, pd.Series, pd.DataFrame],
    freq: str = "MS",
    method: str = "mean",
    dim: str = "time",
    **kwargs: Any,
) -> Union[xr.DataArray, xr.Dataset, pd.Series, pd.DataFrame]:
    """
    Resample data to a new temporal frequency (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray, xarray.Dataset, pandas.Series, or pandas.DataFrame
        Input data with a time-like index or coordinate.
    freq : str, optional
        Resampling frequency (e.g., 'MS' for monthly start, 'W' for weekly, 'D' for daily).
        Default is 'MS'.
    method : str, optional
        Statistical method to apply ('mean', 'sum', 'min', 'max', 'std', 'median').
        Default is 'mean'.
    dim : str, optional
        Dimension along which to resample (xarray only). Default is 'time'.
    **kwargs : Any
        Additional keyword arguments passed to the resample method.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset, pd.Series, pd.DataFrame]
        Resampled data.

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=100, freq="D")
    >>> da = xr.DataArray(np.random.rand(100), coords={"time": times}, dims="time")
    >>> monthly_mean = resample_data(da, freq="MS", method="mean")
    """
    if isinstance(data, (xr.DataArray, xr.Dataset)):
        resampled = data.resample({dim: freq}, **kwargs)
        result = getattr(resampled, method)()
        return _update_history(result, f"Resampled to {freq} using {method}")
    elif isinstance(data, (pd.Series, pd.DataFrame)):
        result = data.resample(freq, **kwargs).agg(method)
        return result
    else:
        raise TypeError("data must be an xarray or pandas object with a time index")

rolling_mean_24h(data, dim='time', min_periods=18, center=True)

Compute rolling 24-hour mean (commonly for PM2.5) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. Must have hourly frequency. dim : str, optional Dimension along which to compute the mean. Default is 'time'. min_periods : int, optional Minimum number of observations in window required to have a value. Default is 18 (75% of 24 hours). center : bool, optional If True, set the labels at the center of the window. Default is True.

Returns

Union[xr.DataArray, xr.Dataset] Rolling 24-hour mean.

Source code in src/monet_stats/analysis.py

def rolling_mean_24h(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
    min_periods: int = 18,
    center: bool = True,
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute rolling 24-hour mean (commonly for PM2.5) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data. Must have hourly frequency.
    dim : str, optional
        Dimension along which to compute the mean. Default is 'time'.
    min_periods : int, optional
        Minimum number of observations in window required to have a value.
        Default is 18 (75% of 24 hours).
    center : bool, optional
        If True, set the labels at the center of the window. Default is True.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Rolling 24-hour mean.
    """
    res = data.rolling({dim: 24}, min_periods=min_periods, center=center).mean()
    return _update_history(res, "Rolling 24-hour mean")

rolling_mean_8h(data, dim='time', min_periods=6, center=True)

Compute rolling 8-hour mean (commonly for Ozone) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data. Must have hourly frequency. dim : str, optional Dimension along which to compute the mean. Default is 'time'. min_periods : int, optional Minimum number of observations in window required to have a value. Default is 6 (75% of 8 hours). center : bool, optional If True, set the labels at the center of the window. Default is True.

Returns

Union[xr.DataArray, xr.Dataset] Rolling 8-hour mean.

Source code in src/monet_stats/analysis.py

def rolling_mean_8h(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
    min_periods: int = 6,
    center: bool = True,
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute rolling 8-hour mean (commonly for Ozone) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data. Must have hourly frequency.
    dim : str, optional
        Dimension along which to compute the mean. Default is 'time'.
    min_periods : int, optional
        Minimum number of observations in window required to have a value.
        Default is 6 (75% of 8 hours).
    center : bool, optional
        If True, set the labels at the center of the window. Default is True.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Rolling 8-hour mean.
    """
    res = data.rolling({dim: 8}, min_periods=min_periods, center=center).mean()
    return _update_history(res, "Rolling 8-hour mean")

seasonal_mean(data, dim='time', weighted=True)

Compute seasonal mean (DJF, MAM, JJA, SON) (Aero Protocol).

Parameters

data : xarray.DataArray or xarray.Dataset Input data with a time-like coordinate. dim : str, optional Dimension along which to compute the seasonal mean. Default is 'time'. weighted : bool, optional If True, weight the mean by the number of days in each month for improved scientific accuracy. Default is True.

Returns

Union[xr.DataArray, xr.Dataset] Seasonal means.

Examples

import xarray as xr import pandas as pd import numpy as np times = pd.date_range("2020-01-01", periods=366, freq="D") da = xr.DataArray(np.random.rand(366), coords={"time": times}, dims="time") s_mean = seasonal_mean(da)

Source code in src/monet_stats/analysis.py

def seasonal_mean(
    data: Union[xr.DataArray, xr.Dataset],
    dim: str = "time",
    weighted: bool = True,
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute seasonal mean (DJF, MAM, JJA, SON) (Aero Protocol).

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with a time-like coordinate.
    dim : str, optional
        Dimension along which to compute the seasonal mean. Default is 'time'.
    weighted : bool, optional
        If True, weight the mean by the number of days in each month for
        improved scientific accuracy. Default is True.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Seasonal means.

    Examples
    --------
    >>> import xarray as xr
    >>> import pandas as pd
    >>> import numpy as np
    >>> times = pd.date_range("2020-01-01", periods=366, freq="D")
    >>> da = xr.DataArray(np.random.rand(366), coords={"time": times}, dims="time")
    >>> s_mean = seasonal_mean(da)
    """
    if not weighted:
        return climatology(data, freq="season", method="mean", dim=dim)

    # To handle both daily and monthly data correctly, we first reduce to monthly
    # means (which are already representative of their months) and then apply
    # seasonal weighting based on month length.
    monthly_data = data.resample({dim: "MS"}).mean(dim=dim)

    def _weighted_seasonal_mean(group: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Dataset]:
        """Helper to apply weighted mean to each seasonal group."""
        month_length = group[dim].dt.days_in_month
        return group.weighted(month_length.fillna(0)).mean(dim=dim)

    # Use xarray.weighted for robust NaN handling
    weighted_data = monthly_data.groupby(f"{dim}.season").map(_weighted_seasonal_mean)

    return _update_history(weighted_data, "Seasonal mean (weighted by days in month)")

weighted_spatial_mean(data, lat_dim='lat', lon_dim='lon', weights=None)

Compute area-weighted spatial mean (Aero Protocol).

Supports automatic detection of 'cell_area', custom weights, or falls back to cosine-latitude weighting for regular grids.

Parameters

data : xarray.DataArray or xarray.Dataset Input data with spatial coordinates. lat_dim : str, optional Name of the latitude dimension. Default is 'lat'. lon_dim : str, optional Name of the longitude dimension. Default is 'lon'. weights : xarray.DataArray or numpy.ndarray, optional Custom weights for the mean. If None, it tries to find 'cell_area' in the data or computes cos(lat) weights.

Returns

Union[xr.DataArray, xr.Dataset] Area-weighted spatial mean.

Notes

Aero Protocol: Targets high performance via xarray.weighted and handles Dask-backed arrays lazily.

Examples

import xarray as xr import numpy as np lats = np.arange(-90, 91, 1) lons = np.arange(-180, 181, 1) da = xr.DataArray(np.ones((len(lats), len(lons))), ... coords={"lat": lats, "lon": lons}, ... dims=("lat", "lon")) spatial_mean = weighted_spatial_mean(da)

Source code in src/monet_stats/analysis.py

def weighted_spatial_mean(
    data: Union[xr.DataArray, xr.Dataset],
    lat_dim: str = "lat",
    lon_dim: str = "lon",
    weights: Optional[Union[xr.DataArray, np.ndarray]] = None,
) -> Union[xr.DataArray, xr.Dataset]:
    """
    Compute area-weighted spatial mean (Aero Protocol).

    Supports automatic detection of 'cell_area', custom weights, or falls
    back to cosine-latitude weighting for regular grids.

    Parameters
    ----------
    data : xarray.DataArray or xarray.Dataset
        Input data with spatial coordinates.
    lat_dim : str, optional
        Name of the latitude dimension. Default is 'lat'.
    lon_dim : str, optional
        Name of the longitude dimension. Default is 'lon'.
    weights : xarray.DataArray or numpy.ndarray, optional
        Custom weights for the mean. If None, it tries to find 'cell_area'
        in the data or computes cos(lat) weights.

    Returns
    -------
    Union[xr.DataArray, xr.Dataset]
        Area-weighted spatial mean.

    Notes
    -----
    Aero Protocol: Targets high performance via xarray.weighted and handles
    Dask-backed arrays lazily.

    Examples
    --------
    >>> import xarray as xr
    >>> import numpy as np
    >>> lats = np.arange(-90, 91, 1)
    >>> lons = np.arange(-180, 181, 1)
    >>> da = xr.DataArray(np.ones((len(lats), len(lons))),
    ...                   coords={"lat": lats, "lon": lons},
    ...                   dims=("lat", "lon"))
    >>> spatial_mean = weighted_spatial_mean(da)
    """
    # 2. Automated Weight Selection
    if weights is None:
        if "cell_area" in data.coords:
            weights = data.coords["cell_area"]
        elif isinstance(data, xr.Dataset) and "cell_area" in data.data_vars:
            weights = data["cell_area"]
        elif lat_dim in data.coords or lat_dim in data.dims:
            # Fall back to cosine of latitude
            weights = np.cos(np.deg2rad(data[lat_dim]))
        else:
            warnings.warn(
                f"No weights found and '{lat_dim}' not in coordinates. "
                "Falling back to equal weights (unweighted mean).",
                UserWarning,
                stacklevel=2,
            )
            weights = xr.DataArray(1.0)

    # 3. Robust Weight Broadcasting
    # If weights is NumPy, wrap it intelligently based on detected spatial dimensions
    if not isinstance(weights, xr.DataArray):
        weights_arr = np.asanyarray(weights)
        spatial_dims = [d for d in [lat_dim, lon_dim] if d in data.dims]

        if not spatial_dims:
            # Fallback for non-spatial data or generic weighting
            w_ndim = weights_arr.ndim
            if w_ndim == 0:
                weights = xr.DataArray(weights_arr)
            else:
                # Original behavior as last resort, but warned
                w_dims = data.dims[-w_ndim:]
                weights = xr.DataArray(weights_arr, dims=w_dims)
        else:
            # Try to match spatial_dims to weights shape
            w_shape = weights_arr.shape
            if len(w_shape) == len(spatial_dims):
                # Check for direct match or transposed match
                data_spatial_shape = tuple(data.sizes[d] for d in spatial_dims)
                if w_shape == data_spatial_shape:
                    weights = xr.DataArray(weights_arr, dims=spatial_dims)
                elif w_shape == data_spatial_shape[::-1]:
                    weights = xr.DataArray(weights_arr, dims=spatial_dims[::-1])
                else:
                    # Best guess if sizes don't match exactly (will fail later during alignment)
                    weights = xr.DataArray(weights_arr, dims=spatial_dims)
            elif len(w_shape) == 1 and len(spatial_dims) > 0:
                # Handle case where 1D weights are provided for one of the spatial dims
                for d in spatial_dims:
                    if w_shape[0] == data.sizes[d]:
                        weights = xr.DataArray(weights_arr, dims=[d])
                        break
                else:
                    weights = xr.DataArray(weights_arr, dims=[spatial_dims[0]])
            else:
                weights = xr.DataArray(weights_arr)

    # 4. Optimized Reduction
    # Capture name before overwriting for history tracking
    weights_name = str(getattr(weights, "name", ""))

    # Determine reduction dimensions (re-use spatial_dims if already calculated)
    if "spatial_dims" not in locals():
        reduction_dims = [d for d in [lat_dim, lon_dim] if d in data.dims]
    else:
        reduction_dims = spatial_dims

    if not reduction_dims:
        # If specified dims are not present, reduce over all shared dims with weights
        reduction_dims = [d for d in data.dims if d in weights.dims]

    weights.name = "weights"
    weighted_data = data.weighted(weights)

    if not reduction_dims:
        # If still no reduction dims, default to all dimensions
        res = weighted_data.mean()
    else:
        res = weighted_data.mean(dim=reduction_dims)

    # Update history with info about weights used
    w_info = "area-weighted" if "cell_area" in weights_name else "weighted"
    return _update_history(res, f"Weighted spatial mean ({w_info})")