From 6db1dd6db9e3c2c28171ac4de4814863f9cdc7f5 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Sat, 16 May 2026 21:25:33 +0200 Subject: [PATCH 1/8] add adaptive spectral filter --- dascore/core/patch.py | 1 + dascore/proc/__init__.py | 1 + .../proc/_adaptive_spectral_filter_numba.py | 317 +++++++ dascore/proc/adaptive_spectral_filter.py | 487 +++++++++++ dascore/utils/signal.py | 86 ++ docs/references.bib | 9 + pyproject.toml | 1 + .../test_adaptive_spectral_filter.py | 783 ++++++++++++++++++ tests/test_proc/test_taper.py | 2 +- 9 files changed, 1686 insertions(+), 1 deletion(-) create mode 100644 dascore/proc/_adaptive_spectral_filter_numba.py create mode 100644 dascore/proc/adaptive_spectral_filter.py create mode 100644 tests/test_proc/test_adaptive_spectral_filter.py diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 9ea9bd3a2..b301f9ca1 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -416,6 +416,7 @@ def iselect(self, *args, **kwargs): gaussian_filter = dascore.proc.gaussian_filter slope_filter = dascore.proc.slope_filter wiener_filter = dascore.proc.wiener_filter + adaptive_spectral_filter = dascore.proc.adaptive_spectral_filter abs = dascore.proc.abs conj = dascore.proc.conj real = dascore.proc.real diff --git a/dascore/proc/__init__.py b/dascore/proc/__init__.py index 40fc327db..cc292d4ff 100644 --- a/dascore/proc/__init__.py +++ b/dascore/proc/__init__.py @@ -18,3 +18,4 @@ from .hampel import hampel_filter from .wiener import wiener_filter from .align import align_to_coord +from .adaptive_spectral_filter import adaptive_spectral_filter diff --git a/dascore/proc/_adaptive_spectral_filter_numba.py b/dascore/proc/_adaptive_spectral_filter_numba.py new file mode 100644 index 000000000..30288f17a --- /dev/null +++ b/dascore/proc/_adaptive_spectral_filter_numba.py @@ -0,0 +1,317 @@ +"""Optional Numba/rocket-fft engine for adaptive spectral filtering.""" + +from __future__ import annotations + +import numba as nb +import numpy as np +import rocket_fft # noqa: F401 # registers FFT overloads with numba. + +from dascore.proc.adaptive_spectral_filter import ( + _finalize_output, + _prepare_work_arrays, + _validate_filter_inputs, +) + + +def _tile_indices_from_parity_index( + ind: int, + count1: int, + parity0: int, + parity1: int, +) -> tuple[int, int]: + """Map a flattened parity-local index back to full-grid tile indices.""" + ix = parity0 + 2 * (ind // count1) + iy = parity1 + 2 * (ind % count1) + return ix, iy + + +def _tile_bounds( + ix: int, + iy: int, + wx: int, + wy: int, + stride0: int, + stride1: int, + shape0: int, + shape1: int, +) -> tuple[int, int, int, int]: + """Return padded-array origin and valid tile shape for one window.""" + beg0 = ix * stride0 + beg1 = iy * stride1 + end0 = min(beg0 + wx, shape0) + end1 = min(beg1 + wy, shape1) + return beg0, beg1, end0 - beg0, end1 - beg1 + + +def _copy_padded_tile( + padded: np.ndarray, + tile: np.ndarray, + beg0: int, + beg1: int, + n0: int, + n1: int, +) -> None: + """Copy the valid padded-array region into a fixed-shape zeroed tile.""" + for i in range(n0): + for j in range(n1): + tile[i, j] = padded[beg0 + i, beg1 + j] + + +def _complex_power(value: complex) -> np.float32: + """Return the magnitude of a complex FFT coefficient as ``float32``.""" + return np.float32((value.real * value.real + value.imag * value.imag) ** 0.5) + + +def _max_spectral_power(spec: np.ndarray) -> np.float32: + """Return the maximum spectral magnitude in one tile.""" + max_power = np.float32(0.0) + for i in range(spec.shape[0]): + for j in range(spec.shape[1]): + power = _complex_power(spec[i, j]) + if power > max_power: + max_power = power + return max_power + + +def _apply_spectral_weight( + spec: np.ndarray, + exponent: float, + normalize_power: bool, +) -> None: + """Apply adaptive magnitude weighting to one tile spectrum in place.""" + max_power = np.float32(0.0) + if normalize_power: + max_power = _max_spectral_power(spec) + + for i in range(spec.shape[0]): + for j in range(spec.shape[1]): + power = _complex_power(spec[i, j]) + if normalize_power: + if max_power != 0.0: + power = power / max_power + else: + power = np.float32(0.0) + weight = np.float32(power**exponent) + spec[i, j] *= weight + + +def _overlap_add_tile( + filtered: np.ndarray, + tile: np.ndarray, + taper: np.ndarray, + beg0: int, + beg1: int, + n0: int, + n1: int, +) -> None: + """Accumulate the valid region of one filtered tile into the output.""" + for i in range(n0): + for j in range(n1): + filtered[beg0 + i, beg1 + j] += tile[i, j] * taper[i, j] + + +_tile_indices_from_parity_index_numba = nb.njit(cache=True, inline="always")( + _tile_indices_from_parity_index +) +_tile_bounds_numba = nb.njit(cache=True, inline="always")(_tile_bounds) +_copy_padded_tile_numba = nb.njit(cache=True, inline="always")(_copy_padded_tile) +_complex_power_numba = nb.njit(cache=True, inline="always")(_complex_power) + + +def _max_spectral_power_numba_impl(spec: np.ndarray) -> np.float32: + """Return the maximum spectral magnitude using compiled helpers.""" + max_power = np.float32(0.0) + for i in range(spec.shape[0]): + for j in range(spec.shape[1]): + power = _complex_power_numba(spec[i, j]) + if power > max_power: + max_power = power + return max_power + + +def _apply_spectral_weight_numba_impl( + spec: np.ndarray, + exponent: float, + normalize_power: bool, +) -> None: + """Apply adaptive magnitude weighting using compiled helpers.""" + max_power = np.float32(0.0) + if normalize_power: + max_power = _max_spectral_power_numba(spec) + + for i in range(spec.shape[0]): + for j in range(spec.shape[1]): + power = _complex_power_numba(spec[i, j]) + if normalize_power: + if max_power != 0.0: + power = power / max_power + else: + power = np.float32(0.0) + weight = np.float32(power**exponent) + spec[i, j] *= weight + + +_max_spectral_power_numba = nb.njit(cache=True, inline="always")( + _max_spectral_power_numba_impl +) +_apply_spectral_weight_numba = nb.njit(cache=True, inline="always")( + _apply_spectral_weight_numba_impl +) +_overlap_add_tile_numba = nb.njit(cache=True, inline="always")(_overlap_add_tile) + + +def _process_tile_group_python( + padded: np.ndarray, + filtered: np.ndarray, + taper: np.ndarray, + wx: int, + wy: int, + stride0: int, + stride1: int, + nx: int, + ny: int, + parity0: int, + parity1: int, + exponent: float, + normalize_power: bool, +) -> None: + """Process one non-overlapping tile parity group in pure Python.""" + count0 = (nx - parity0 + 1) // 2 + count1 = (ny - parity1 + 1) // 2 + count = count0 * count1 + for ind in range(count): + ix, iy = _tile_indices_from_parity_index(ind, count1, parity0, parity1) + beg0, beg1, n0, n1 = _tile_bounds( + ix, iy, wx, wy, stride0, stride1, padded.shape[0], padded.shape[1] + ) + + tile = np.zeros((wx, wy), dtype=np.float32) + _copy_padded_tile(padded, tile, beg0, beg1, n0, n1) + + spec = np.fft.rfft2(tile) + if exponent != 0.0: + _apply_spectral_weight(spec, exponent, normalize_power) + + tile = np.fft.irfft2(spec, s=(wx, wy)) + _overlap_add_tile(filtered, tile, taper, beg0, beg1, n0, n1) + + +def _process_tile_group_numba_impl( + padded: np.ndarray, + filtered: np.ndarray, + taper: np.ndarray, + wx: int, + wy: int, + stride0: int, + stride1: int, + nx: int, + ny: int, + parity0: int, + parity1: int, + exponent: float, + normalize_power: bool, +) -> None: + """Process one non-overlapping tile parity group with compiled helpers.""" + count0 = (nx - parity0 + 1) // 2 + count1 = (ny - parity1 + 1) // 2 + count = count0 * count1 + for ind in nb.prange(count): # type: ignore[not-iterable] + ix, iy = _tile_indices_from_parity_index_numba(ind, count1, parity0, parity1) + beg0, beg1, n0, n1 = _tile_bounds_numba( + ix, iy, wx, wy, stride0, stride1, padded.shape[0], padded.shape[1] + ) + + tile = np.zeros((wx, wy), dtype=np.float32) + _copy_padded_tile_numba(padded, tile, beg0, beg1, n0, n1) + + spec = np.fft.rfft2(tile) + if exponent != 0.0: + _apply_spectral_weight_numba(spec, exponent, normalize_power) + + tile = np.fft.irfft2(spec, s=(wx, wy)) + _overlap_add_tile_numba(filtered, tile, taper, beg0, beg1, n0, n1) + + +# fastmath is intentional here: the weighting is approximate and tests allow +# small SciPy/Numba differences from parallel floating-point evaluation. +_process_tile_group_numba = nb.njit(cache=True, fastmath=True, parallel=True)( + _process_tile_group_numba_impl +) + + +def _adaptive_spectral_filter_numba( + data: np.ndarray, + *, + window_size: tuple[int, int], + overlap: tuple[int, int], + exponent: float = 0.3, + normalize_power: bool = False, +) -> np.ndarray: + """ + Filter a 2D array with the optional Numba/rocket-fft implementation. + + Parameters + ---------- + data + Two-dimensional input array. The filter computes in ``float32``. + window_size + Two power-of-two window lengths, one per array axis. Values must be + greater than 4. + overlap + Number of samples each neighboring window overlaps on each axis. Each + value must be non-negative and smaller than half the matching window. + exponent + Spectral magnitude exponent used as the adaptive weighting power. ``0`` + leaves the spectrum unweighted before overlap-add reconstruction. + normalize_power + If ``True``, normalize each tile's spectral magnitudes by that tile's + maximum magnitude before applying ``exponent``. + + Returns + ------- + numpy.ndarray + The filtered array with the same shape as ``data``. Floating input + dtypes are restored; non-floating inputs return ``float32`` output. + + Raises + ------ + ValueError + If ``data`` is not two-dimensional, ``exponent`` is not finite, + ``window_size`` and ``overlap`` do not contain exactly two integer + values, any window size is not a power of two greater than 4, or any + overlap is negative or at least half the matching window size. + + Notes + ----- + This implementation uses Numba-compiled loops and rocket-fft-backed NumPy + FFT calls. It is selected by + :func:`dascore.proc.adaptive_spectral_filter.adaptive_spectral_filter` for + two selected dimensions when ``engine="numba"`` or when ``engine="auto"`` + and optional dependencies are installed. + """ + data = np.asarray(data) + _validate_filter_inputs( + data, window_size=window_size, overlap=overlap, exponent=float(exponent) + ) + wx, wy = window_size + working, original_dtype, stride, taper, padded, filtered, n_tiles = ( + _prepare_work_arrays(data, window_size=window_size, overlap=overlap) + ) + for parity0 in range(2): + for parity1 in range(2): + _process_tile_group_numba( + padded, + filtered, + taper, + wx, + wy, + stride[0], + stride[1], + n_tiles[0], + n_tiles[1], + parity0, + parity1, + float(exponent), + bool(normalize_power), + ) + return _finalize_output(filtered, working, original_dtype, stride) diff --git a/dascore/proc/adaptive_spectral_filter.py b/dascore/proc/adaptive_spectral_filter.py new file mode 100644 index 000000000..73f5ee71d --- /dev/null +++ b/dascore/proc/adaptive_spectral_filter.py @@ -0,0 +1,487 @@ +""" +Adaptive spectral filtering for DASCore patches. + +The adaptive spectral filter suppresses incoherent energy by processing a patch +in overlapping windows along one or two selected dimensions. Each window is +transformed to the spectral domain, weighted by a power of its spectral +magnitude, transformed back to the original domain, and accumulated with +tapered overlap-add reconstruction. + +With one selected dimension, this is an adaptive frequency-domain normalization +applied independently to every trace over the remaining patch dimensions. With +two selected dimensions, this is the adaptive frequency-wavenumber filter +described by @isken2022denoising and exposed by Pyrocko +[Lightguide](https://github.com/pyrocko/lightguide). Coherent plane-wave energy +tends to concentrate in the frequency-wavenumber spectrum, so the weighting +emphasizes locally coherent arrivals relative to diffuse or randomly +distributed energy. + +This module exposes a single public patch method, +:func:`adaptive_spectral_filter`. The public function resolves one or two +DASCore dimensions, converts window and overlap values to sample counts, moves +those dimensions to the array tail, and processes every remaining leading index +as an independent batch. The lower-level SciPy and Numba implementations are +private because they operate on raw arrays and do not perform DASCore +coordinate handling. + +The SciPy engine handles one- and two-dimensional selected windows using +``rfftn``/``irfftn``. The optional Numba/rocket-fft engine currently handles +the two-dimensional case only, using parity-separated tile groups so neighboring +writes do not overlap within each parallel loop. Both engines share validation, +padding, tapering, and dtype-restoration logic so two-dimensional outputs remain +directly comparable. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from itertools import product +from math import prod +from typing import Any, Literal + +import numpy as np +from scipy import fft as sp_fft + +from dascore.constants import PatchType +from dascore.exceptions import MissingOptionalDependencyError, ParameterError +from dascore.utils.patch import get_dim_axis_value, patch_function +from dascore.utils.signal import _triangular_taper + +_AdaptiveSpectralEngine = Literal["auto", "numba", "scipy"] +__all__ = ("adaptive_spectral_filter",) + + +def _is_power_of_two(value: int) -> bool: + """Return ``True`` when *value* is a positive power of two.""" + return value > 0 and (value & (value - 1) == 0) + + +def _validate_filter_inputs( + data: np.ndarray, + *, + window_size: tuple[int, ...], + overlap: tuple[int, ...], + exponent: float, +) -> None: + """Validate direct array-filter inputs before entering FFT kernels.""" + if data.ndim not in {1, 2}: + msg = ( + f"adaptive spectral array filters require 1D or 2D input; got {data.ndim}D." + ) + raise ValueError(msg) + if len(window_size) != data.ndim or len(overlap) != data.ndim: + msg = "window_size and overlap must match the input dimensionality." + raise ValueError(msg) + if not np.isfinite(exponent): + msg = "exponent must be finite." + raise ValueError(msg) + + for axis, (window, axis_overlap) in enumerate(zip(window_size, overlap)): + if not isinstance(window, (int, np.integer)): + msg = f"window_size[{axis}] must be an integer; got {window!r}." + raise ValueError(msg) + if not isinstance(axis_overlap, (int, np.integer)): + msg = f"overlap[{axis}] must be an integer; got {axis_overlap!r}." + raise ValueError(msg) + + window = int(window) + axis_overlap = int(axis_overlap) + if window <= 4 or not _is_power_of_two(window): + msg = ( + f"window_size[{axis}] must be a power of two greater than 4; " + f"got {window!r}." + ) + raise ValueError(msg) + if axis_overlap < 0: + msg = f"overlap[{axis}] must be non-negative; got {axis_overlap!r}." + raise ValueError(msg) + if axis_overlap >= window / 2: + msg = f"overlap[{axis}] is too large; maximum is {window // 2 - 1} samples." + raise ValueError(msg) + + +def _prepare_work_arrays( + data: np.ndarray, + *, + window_size: tuple[int, ...], + overlap: tuple[int, ...], +) -> tuple[ + np.ndarray, + np.dtype, + tuple[int, ...], + np.ndarray, + np.ndarray, + np.ndarray, + tuple[int, ...], +]: + """Prepare ``float32`` padded arrays shared by filter implementations.""" + data = np.asarray(data) + original_dtype = data.dtype + working = np.ascontiguousarray(data, dtype=np.float32) + stride = tuple(win - over for win, over in zip(window_size, overlap)) + plateau = tuple(win - 2 * over for win, over in zip(window_size, overlap)) + taper = _triangular_taper(window_size, plateau) + + padded_shape = tuple( + length + 2 * step for length, step in zip(working.shape, stride) + ) + padded = np.zeros(padded_shape, dtype=np.float32) + inner_slices = tuple( + slice(step, length + step) for length, step in zip(working.shape, stride) + ) + padded[inner_slices] = working + filtered = np.zeros_like(padded) + n_tiles = tuple(pad_len // step for pad_len, step in zip(padded.shape, stride)) + return working, original_dtype, stride, taper, padded, filtered, n_tiles + + +def _finalize_output( + filtered: np.ndarray, + working: np.ndarray, + original_dtype: np.dtype, + stride: tuple[int, ...], +) -> np.ndarray: + """Crop padded output and restore floating dtypes where possible.""" + slices = tuple( + slice(step, length + step) for length, step in zip(working.shape, stride) + ) + out = filtered[slices] + if np.issubdtype(original_dtype, np.floating): + return out.astype(original_dtype, copy=False) + return out + + +def _extract_tiles_python( + padded: np.ndarray, + window_size: tuple[int, ...], + stride: tuple[int, ...], + n_tiles: tuple[int, ...], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Extract padded windows into a dense tile stack for batched SciPy FFTs.""" + ndim = len(window_size) + tiles = np.zeros((prod(n_tiles), *window_size), dtype=np.float32) + begins = np.zeros((*n_tiles, ndim), dtype=np.int64) + sizes = np.zeros((*n_tiles, ndim), dtype=np.int64) + for tile_index, tile_inds in enumerate(product(*(range(num) for num in n_tiles))): + beg = tuple(ind * step for ind, step in zip(tile_inds, stride)) + end = tuple( + min(start + win, size) + for start, win, size in zip(beg, window_size, padded.shape) + ) + valid_shape = tuple(stop - start for start, stop in zip(beg, end)) + begins[tile_inds] = beg + sizes[tile_inds] = valid_shape + data_slices = tuple(slice(start, stop) for start, stop in zip(beg, end)) + tile_slices = tuple(slice(0, size) for size in valid_shape) + tiles[(tile_index, *tile_slices)] = padded[data_slices] + return tiles, begins, sizes + + +def _overlap_add_tiles_python( + out: np.ndarray, + tiles: np.ndarray, + taper: np.ndarray, + begins: np.ndarray, + sizes: np.ndarray, +) -> None: + """Apply tapered overlap-add reconstruction from a dense tile stack.""" + grid_shape = begins.shape[:-1] + for tile_index, tile_inds in enumerate( + product(*(range(num) for num in grid_shape)) + ): + beg = tuple(begins[tile_inds]) + valid_shape = tuple(sizes[tile_inds]) + out_slices = tuple( + slice(start, start + size) for start, size in zip(beg, valid_shape) + ) + tile_slices = tuple(slice(0, size) for size in valid_shape) + out[out_slices] += tiles[(tile_index, *tile_slices)] * taper[tile_slices] + + +def _adaptive_spectral_filter_scipy( + data: np.ndarray, + *, + window_size: tuple[int, ...], + overlap: tuple[int, ...], + exponent: float = 0.3, + normalize_power: bool = False, +) -> np.ndarray: + """ + Filter a 1D or 2D array with the SciPy adaptive spectral implementation. + + Parameters + ---------- + data + One- or two-dimensional input array. The filter computes in ``float32``. + window_size + Power-of-two window lengths, one per array axis. Values must be greater + than 4. + overlap + Number of samples each neighboring window overlaps on each axis. Values + must be non-negative and smaller than half the matching window. + exponent + Spectral magnitude exponent used as the adaptive weighting power. ``0`` + leaves the spectrum unweighted before overlap-add reconstruction. + normalize_power + If ``True``, normalize each tile's spectral magnitudes by that tile's + maximum magnitude before applying ``exponent``. + + Returns + ------- + numpy.ndarray + The filtered array with the same shape as ``data``. Floating input + dtypes are restored; non-floating inputs return ``float32`` output. + + Raises + ------ + ValueError + If ``data`` is not one- or two-dimensional, ``exponent`` is not finite, + ``window_size`` and ``overlap`` do not match ``data.ndim``, any window + size is not a power of two greater than 4, or any overlap is negative or + at least half the matching window size. + """ + data = np.asarray(data) + _validate_filter_inputs( + data, window_size=window_size, overlap=overlap, exponent=float(exponent) + ) + working, original_dtype, stride, taper, padded, filtered, n_tiles = ( + _prepare_work_arrays(data, window_size=window_size, overlap=overlap) + ) + tiles, begins, sizes = _extract_tiles_python(padded, window_size, stride, n_tiles) + axes = tuple(range(-data.ndim, 0)) + + spec = sp_fft.rfftn(tiles, s=window_size, axes=axes, workers=-1) + if exponent != 0.0: + power = np.abs(spec).astype(np.float32, copy=False) + if normalize_power: + max_power = power.max(axis=axes, keepdims=True) + power = np.divide( + power, max_power, out=np.zeros_like(power), where=max_power != 0 + ) + spec *= power**exponent + tiles = sp_fft.irfftn(spec, s=window_size, axes=axes, workers=-1).astype( + np.float32, copy=False + ) + _overlap_add_tiles_python(filtered, tiles, taper, begins, sizes) + return _finalize_output(filtered, working, original_dtype, stride) + + +def _get_dim_axis_values(patch: PatchType, kwargs: Mapping[str, Any]): + """Resolve DASCore dimension keyword arguments into dim/axis values.""" + if len(kwargs) not in {1, 2}: + msg = ( + "adaptive_spectral_filter requires one or two dimension window kwargs, " + "e.g. patch.adaptive_spectral_filter(time=32, samples=True)." + ) + raise ParameterError(msg) + return get_dim_axis_value(patch, kwargs=dict(kwargs), allow_multiple=True) + + +def _dim_values_to_samples( + patch: PatchType, + dim_axis_values, + *, + samples: bool, + name: str, + force_sample_dims: frozenset[str] = frozenset(), +) -> tuple[int, ...]: + """Convert DASCore dimension values from units or samples into sample counts.""" + out: list[int] = [] + for dim, _, value in dim_axis_values: + if samples or dim in force_sample_dims: + count = int(value) + else: + coord = patch.get_coord(dim, require_evenly_sampled=True) + count = coord.get_sample_count(value, samples=False) + invalid = count < 0 if name == "overlap" else count <= 0 + if invalid: + requirement = "non-negative" if name == "overlap" else "positive" + msg = f"{name} for dimension {dim!r} must be {requirement}." + raise ParameterError(msg) + out.append(count) + return tuple(out) + + +def _normalize_overlap( + overlap: int | Mapping[str, Any] | None, + dims: tuple[str, ...], + windows: tuple[int, ...], +) -> tuple[dict[str, Any], frozenset[str]]: + """Return per-dimension overlap values and internally defaulted dimensions.""" + defaults = {dim: max(window // 2 - 2, 0) for dim, window in zip(dims, windows)} + if overlap is None: + return defaults, frozenset(dims) + if isinstance(overlap, Mapping): + extra = set(overlap) - set(dims) + if extra: + msg = f"overlap contains dimensions not being filtered: {sorted(extra)}" + raise ParameterError(msg) + return defaults | dict(overlap), frozenset(set(dims) - set(overlap)) + return {dim: int(overlap) for dim in dims}, frozenset() + + +def _validate_window_and_overlap( + dims: tuple[str, ...], + windows: tuple[int, ...], + overlaps: tuple[int, ...], + exponent: float, +) -> None: + """Validate public DASCore window and overlap settings.""" + if not np.isfinite(exponent): + msg = "exponent must be finite." + raise ParameterError(msg) + for dim, window, overlap in zip(dims, windows, overlaps): + if window <= 4 or not _is_power_of_two(window): + msg = f"window size for {dim!r} must be a power of two and > 4." + raise ParameterError(msg) + if overlap < 0: + msg = f"overlap for {dim!r} must be non-negative." + raise ParameterError(msg) + if overlap >= window / 2: + msg = ( + f"overlap for {dim!r} is too large. Maximum overlap is " + f"{window // 2 - 1} samples." + ) + raise ParameterError(msg) + + +def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable: + """Return the requested adaptive spectral array filter implementation.""" + if engine == "scipy" or (engine == "auto" and selected_ndim == 1): + return _adaptive_spectral_filter_scipy + if engine not in {"auto", "numba"}: + msg = "engine must be one of 'auto', 'numba', or 'scipy'." + raise ParameterError(msg) + if selected_ndim != 2: + msg = "engine='numba' currently supports exactly two selected dimensions." + raise ParameterError(msg) + try: + from dascore.proc._adaptive_spectral_filter_numba import ( + _adaptive_spectral_filter_numba, + ) + except ImportError as exc: + if engine == "numba": + msg = ( + "engine='numba' requires optional dependencies numba and " + "rocket-fft to be installed." + ) + raise MissingOptionalDependencyError(msg) from exc + return _adaptive_spectral_filter_scipy + return _adaptive_spectral_filter_numba + + +@patch_function() +def adaptive_spectral_filter( + patch: PatchType, + *, + overlap: int | Mapping[str, Any] | None = None, + exponent: float = 0.3, + normalize_power: bool = False, + samples: bool = False, + engine: _AdaptiveSpectralEngine = "auto", + **kwargs: Any, +) -> PatchType: + """ + Apply adaptive spectral filtering over one or two patch dimensions. + + Parameters + ---------- + patch + DASCore patch whose data should be filtered. + overlap + Window overlap in samples when ``samples=True`` or in coordinate units + otherwise. A single value applies to all selected dimensions; a mapping + can specify dimensions independently. When omitted, each dimension + defaults to ``window // 2 - 2`` samples. + exponent + Spectral magnitude exponent used as the adaptive weighting power. ``0`` + leaves the spectrum unweighted before overlap-add reconstruction. + normalize_power + If ``True``, normalize each tile's spectral magnitudes by that tile's + maximum magnitude before applying ``exponent``. + samples + If ``True``, dimension kwargs and overlap values are interpreted as + sample counts. If ``False``, values are converted through evenly sampled + patch coordinates. + engine + ``"auto"`` uses SciPy for one selected dimension and the optional + Numba/rocket-fft implementation for two selected dimensions when + available. ``"numba"`` requires two selected dimensions and the optional + fast engine. ``"scipy"`` always uses the SciPy FFT implementation. + **kwargs + One or two dimension names and their window sizes, such as ``time=32`` + or ``time=32, distance=32``. + + Returns + ------- + Patch + A new patch with filtered data and original dimensions and coordinates. + + Raises + ------ + ParameterError + If one or two dimensions are not selected, if selected window or overlap + values are invalid, if ``exponent`` is not finite, or if an invalid + engine name is requested. + MissingOptionalDependencyError + If ``engine="numba"`` is requested for two selected dimensions but the + optional fast-engine dependencies are not installed. + + Examples + -------- + >>> import dascore as dc + >>> patch = dc.get_example_patch() + >>> filtered_1d = patch.adaptive_spectral_filter(time=32, samples=True) + >>> filtered_2d = patch.adaptive_spectral_filter( + ... time=32, distance=32, samples=True + ... ) + >>> filtered_1d.shape == filtered_2d.shape == patch.shape + True + + Notes + ----- + With two selected dimensions, this method is equivalent to the adaptive + frequency-wavenumber (f-k) filter described in @isken2022denoising and + follows the behavior exposed by Pyrocko + [Lightguide](https://github.com/pyrocko/lightguide). + """ + dim_axis_values = _get_dim_axis_values(patch, kwargs) + dims = tuple(x.dim for x in dim_axis_values) + axes = tuple(x.axis for x in dim_axis_values) + windows = _dim_values_to_samples( + patch, dim_axis_values, samples=samples, name="window" + ) + overlap_values, default_overlap_dims = _normalize_overlap(overlap, dims, windows) + overlap_dim_axis_values = get_dim_axis_value( + patch, kwargs=overlap_values, allow_multiple=True + ) + overlaps = _dim_values_to_samples( + patch, + samples=samples, + dim_axis_values=overlap_dim_axis_values, + name="overlap", + force_sample_dims=default_overlap_dims, + ) + _validate_window_and_overlap(dims, windows, overlaps, float(exponent)) + + data = np.asarray(patch.data) + selected_ndim = len(axes) + moved = np.moveaxis(data, axes, tuple(range(-selected_ndim, 0))) + batch_shape = moved.shape[:-selected_ndim] + selected_shape = moved.shape[-selected_ndim:] + working = moved.reshape((-1, *selected_shape)) + filtered = np.empty_like(working, dtype=np.float32) + engine_func = _get_engine(engine, selected_ndim) + for ind, array in enumerate(working): + filtered[ind] = engine_func( + array, + window_size=windows, + overlap=overlaps, + exponent=float(exponent), + normalize_power=bool(normalize_power), + ) + filtered = filtered.reshape((*batch_shape, *selected_shape)) + filtered = np.moveaxis(filtered, tuple(range(-selected_ndim, 0)), axes) + if np.issubdtype(data.dtype, np.floating): + filtered = filtered.astype(data.dtype, copy=False) + return patch.update(data=filtered) diff --git a/dascore/utils/signal.py b/dascore/utils/signal.py index 0d2bbddfe..a2837b20e 100644 --- a/dascore/utils/signal.py +++ b/dascore/utils/signal.py @@ -2,6 +2,9 @@ Utilities for signal processing. """ +from functools import lru_cache + +import numpy as np from scipy.signal import windows from dascore.exceptions import ParameterError @@ -34,3 +37,86 @@ def _get_window_function(window_type): raise ParameterError(msg) func = WINDOW_FUNCTIONS[window_type] return func + + +def _triangular_taper_1d(size: int, plateau: int) -> np.ndarray: + """Return a one-dimensional triangular plateau taper.""" + ramp_size = (size - plateau) // 2 + taper = np.ones(size, dtype=np.float32) + if ramp_size: + ramp = _get_window_function("triang")(ramp_size * 2 + 1)[:ramp_size] + taper[:ramp_size] = ramp + taper[size - ramp_size :] = ramp[::-1] + return taper + + +@lru_cache(maxsize=64) +def _cached_triangular_taper( + window_size: tuple[int, ...], plateau: tuple[int, ...] +) -> np.ndarray: + """Build the cached taper array used by :func:`_triangular_taper`.""" + if len(window_size) != len(plateau): + msg = "window_size and plateau must have the same length." + raise ValueError(msg) + if len(window_size) not in {1, 2}: + msg = "Only one- and two-dimensional tapers are supported." + raise ValueError(msg) + if any(plat > win for win, plat in zip(window_size, plateau)): + msg = "Plateau cannot be larger than window size." + raise ValueError(msg) + if any(plat < 0 for plat in plateau): + msg = "Plateau sizes must be non-negative." + raise ValueError(msg) + if any(win % 2 for win in window_size): + msg = "Window sizes must be even." + raise ValueError(msg) + tapers = [ + _triangular_taper_1d(win, plat) for win, plat in zip(window_size, plateau) + ] + if len(tapers) == 1: + return tapers[0].astype(np.float32) + return (tapers[0][:, None] * tapers[1][None, :]).astype(np.float32) + + +def _triangular_taper( + window_size: tuple[int, ...], plateau: tuple[int, ...] +) -> np.ndarray: + """ + Return a one- or two-dimensional triangular plateau taper. + + Parameters + ---------- + window_size + Number of samples in the window dimensions. Values must be even. + plateau + Number of central samples with unit weight in each dimension. Values + must be non-negative and no larger than the corresponding window size. + + Returns + ------- + numpy.ndarray + A ``float32`` array with shape ``window_size``. The returned array is a + copy, so callers can mutate it without corrupting the internal cache. + + Raises + ------ + ValueError + If the inputs have different lengths, if the taper dimensionality is + not one or two, if any plateau is negative, if any plateau is greater + than the corresponding window size, or if any window size is odd. + + Notes + ----- + The taper is separable in 2D. Each axis contains a central unit-weight + plateau and triangular ramps on both sides generated from DASCore's + registered ``"triang"`` window function. When plateau equals window size, + the result is all ones along that axis. + + Examples + -------- + >>> from dascore.utils.signal import _triangular_taper + >>> taper = _triangular_taper((8, 8), (2, 2)) + >>> taper.shape + (8, 8) + """ + return _cached_triangular_taper(window_size, plateau).copy() diff --git a/docs/references.bib b/docs/references.bib index b633bf422..2a4428042 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -67,3 +67,12 @@ @article{schimmel1997noise year={1997}, publisher={Blackwell Publishing Ltd Oxford, UK} } + +@article{isken2022denoising, + title={De-noising distributed acoustic sensing data using an adaptive frequency-wavenumber filter}, + author={Isken, Marius Paul and Vasyura-Bathke, Hannes and Dahm, Torsten and Heimann, Sebastian}, + journal={Geophysical Journal International}, + pages={ggac229}, + year={2022}, + doi={10.1093/gji/ggac229} +} diff --git a/pyproject.toml b/pyproject.toml index 764a87bf8..9c3eb2fef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,7 @@ extras = [ "findiff", "obspy", "numba", + "rocket-fft", "segyio", "bottleneck", ] diff --git a/tests/test_proc/test_adaptive_spectral_filter.py b/tests/test_proc/test_adaptive_spectral_filter.py new file mode 100644 index 000000000..b91948bd6 --- /dev/null +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -0,0 +1,783 @@ +"""Tests for adaptive spectral filtering.""" + +from __future__ import annotations + +import builtins +from typing import Any + +import numpy as np +import pytest + +import dascore as dc +from dascore.exceptions import ( + CoordError, + MissingOptionalDependencyError, + ParameterError, + PatchCoordinateError, +) +from dascore.proc import adaptive_spectral_filter as adaptive_spectral_filter_func +from dascore.proc.adaptive_spectral_filter import ( + _adaptive_spectral_filter_scipy, + _get_engine, + _validate_window_and_overlap, +) +from dascore.utils.signal import _triangular_taper + + +def _patch( + shape: tuple[int, ...], + dims: tuple[str, ...], + *, + dtype=np.float32, + time_step=np.timedelta64(4, "ms"), + distance_step=1.0, +) -> dc.Patch: + """Return a deterministic patch for adaptive spectral tests.""" + rng = np.random.default_rng(20260508) + data = rng.normal(size=shape).astype(dtype) + coords = {} + for dim, length in zip(dims, shape, strict=True): + if dim == "time": + coords[dim] = np.datetime64("2020-01-01") + np.arange(length) * time_step + elif dim == "distance": + coords[dim] = np.arange(length, dtype=float) * distance_step + else: + coords[dim] = np.arange(length, dtype=float) + return dc.Patch(data=data, coords=coords, dims=dims) + + +class TestAdaptiveSpectralFilter: + """Tests for the adaptive spectral patch method.""" + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) + def test_dtype_shape_dims_and_coords_preserved(self, dtype) -> None: + """Adaptive spectral should preserve patch structure and floating dtype.""" + patch = _patch((64, 64), ("distance", "time"), dtype=dtype) + + out = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap={"distance": 7, "time": 7}, + samples=True, + engine="scipy", + ) + + assert np.asarray(out.data).dtype == np.asarray(patch.data).dtype + assert out.shape == patch.shape + assert out.dims == patch.dims + assert out.coords == patch.coords + assert out.attrs.history[-1].startswith("adaptive_spectral_filter") + + def test_time_distance_reversed_dims_are_supported(self) -> None: + """Selected dimensions need not be in a fixed order.""" + patch = _patch((64, 80), ("time", "distance"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + time=16, + distance=32, + overlap={"time": 7, "distance": 14}, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + assert out.dims == ("time", "distance") + assert np.isfinite(out.data).all() + + def test_arbitrary_2d_dimension_names_are_supported(self) -> None: + """Adaptive spectral should work with any two patch dimensions.""" + patch = _patch((64, 64), ("channel", "sample"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + channel=16, + sample=16, + overlap={"channel": 7, "sample": 7}, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + assert out.dims == ("channel", "sample") + + def test_requires_explicit_dimension_kwargs(self) -> None: + """At least one dimension window is required.""" + patch = _patch((64, 64), ("channel", "sample"), dtype=np.float32) + + with pytest.raises(ParameterError, match="one or two dimension window kwargs"): + patch.adaptive_spectral_filter(samples=True, engine="scipy") + + def test_rejects_non_positive_window(self) -> None: + """Window sizes must resolve to positive sample counts.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(ParameterError, match=r"window.*must be positive"): + patch.adaptive_spectral_filter( + distance=0, time=16, samples=True, engine="scipy" + ) + + def test_one_dimension_filter_is_supported(self) -> None: + """A single selected dimension should run the 1D spectral path.""" + patch = _patch((8, 64), ("distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + time=16, + overlap={"time": 7}, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + assert out.dims == patch.dims + assert out.coords == patch.coords + assert np.isfinite(out.data).all() + + def test_one_dimension_filter_batches_other_dims(self) -> None: + """Unselected dimensions should be batches for 1D filtering.""" + patch = _patch((3, 8, 64), ("shot", "distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter(time=16, samples=True, engine="scipy") + expected = np.stack( + [ + dc.Patch( + data=np.asarray(patch.data)[ind], + coords={ + "distance": patch.get_array("distance"), + "time": patch.get_array("time"), + }, + dims=("distance", "time"), + ) + .adaptive_spectral_filter(time=16, samples=True, engine="scipy") + .data + for ind in range(patch.shape[0]) + ] + ) + + np.testing.assert_allclose(out.data, expected, rtol=1e-5, atol=1e-5) + + def test_coordinate_unit_window_and_overlap_conversion(self) -> None: + """Coordinate units and sample counts should resolve identically.""" + patch = _patch( + (64, 64), + ("distance", "time"), + dtype=np.float32, + distance_step=2.0, + time_step=np.timedelta64(4, "ms"), + ) + + by_units = patch.adaptive_spectral_filter( + distance=32.0, + time=np.timedelta64(64, "ms"), + overlap={"distance": 14.0, "time": np.timedelta64(28, "ms")}, + samples=False, + engine="scipy", + ) + by_samples = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap={"distance": 7, "time": 7}, + samples=True, + engine="scipy", + ) + + np.testing.assert_allclose(by_units.data, by_samples.data, rtol=1e-5, atol=1e-5) + + def test_default_overlap_stays_in_samples_when_windows_use_units(self) -> None: + """Computed overlap defaults should not be interpreted as coordinate units.""" + patch = _patch( + (64, 64), + ("distance", "time"), + dtype=np.float32, + distance_step=2.0, + time_step=np.timedelta64(4, "ms"), + ) + + by_units = patch.adaptive_spectral_filter( + distance=32.0, + time=np.timedelta64(64, "ms"), + samples=False, + engine="scipy", + ) + by_samples = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap={"distance": 6, "time": 6}, + samples=True, + engine="scipy", + ) + + np.testing.assert_allclose(by_units.data, by_samples.data, rtol=1e-5, atol=1e-5) + + def test_partial_overlap_defaults_stay_in_samples_with_units(self) -> None: + """Missing overlap mapping entries should stay sample-count defaults.""" + patch = _patch( + (64, 64), + ("distance", "time"), + dtype=np.float32, + distance_step=2.0, + time_step=np.timedelta64(4, "ms"), + ) + + by_units = patch.adaptive_spectral_filter( + distance=32.0, + time=np.timedelta64(64, "ms"), + overlap={"time": np.timedelta64(24, "ms")}, + samples=False, + engine="scipy", + ) + by_samples = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap={"distance": 6, "time": 6}, + samples=True, + engine="scipy", + ) + + np.testing.assert_allclose(by_units.data, by_samples.data, rtol=1e-5, atol=1e-5) + + def test_1d_default_overlap_stays_in_samples_with_units(self) -> None: + """Computed 1D overlap defaults should stay in sample counts.""" + patch = _patch( + (8, 64), + ("distance", "time"), + dtype=np.float32, + time_step=np.timedelta64(4, "ms"), + ) + + by_units = patch.adaptive_spectral_filter( + time=np.timedelta64(64, "ms"), + samples=False, + engine="scipy", + ) + by_samples = patch.adaptive_spectral_filter( + time=16, + overlap=6, + samples=True, + engine="scipy", + ) + + np.testing.assert_allclose(by_units.data, by_samples.data, rtol=1e-5, atol=1e-5) + + @pytest.mark.parametrize( + "shape,dims,kwargs", + [ + ((3, 64, 64), ("shot", "distance", "time"), {"distance": 16, "time": 16}), + ( + (2, 3, 64, 64), + ("component", "shot", "distance", "time"), + {"distance": 16, "time": 16}, + ), + ((64, 3, 64), ("distance", "shot", "time"), {"distance": 16, "time": 16}), + ], + ) + def test_batches_over_non_selected_dimensions( + self, + shape: tuple[int, ...], + dims: tuple[str, ...], + kwargs: dict[str, int], + ) -> None: + """Extra dimensions should be processed as independent 2D batches.""" + patch = _patch(shape, dims, dtype=np.float32) + + out = patch.adaptive_spectral_filter( + **kwargs, + overlap={dim: max(value // 2 - 2, 0) for dim, value in kwargs.items()}, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + assert out.dims == patch.dims + assert out.coords == patch.coords + assert np.isfinite(out.data).all() + + def test_batched_output_matches_independent_2d_calls(self) -> None: + """Batched filtering should match independent 2D patch calls.""" + patch = _patch((4, 32, 32), ("depth", "distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + distance=16, time=16, samples=True, engine="scipy" + ) + expected = np.stack( + [ + dc.Patch( + data=np.asarray(patch.data)[ind], + coords={ + "distance": patch.get_array("distance"), + "time": patch.get_array("time"), + }, + dims=("distance", "time"), + ) + .adaptive_spectral_filter( + distance=16, time=16, samples=True, engine="scipy" + ) + .data + for ind in range(patch.shape[0]) + ] + ) + + np.testing.assert_allclose(out.data, expected, rtol=1e-5, atol=1e-5) + + def test_rejects_unknown_overlap_dimension(self) -> None: + """Overlap mappings may only name selected dimensions.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(ParameterError, match="overlap contains dimensions"): + patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap={"distance": 7, "bad": 7}, + samples=True, + engine="scipy", + ) + + def test_scalar_overlap_applies_to_both_dimensions(self) -> None: + """A scalar overlap should apply to each selected dimension.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap=7, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + + def test_scalar_overlap_applies_to_one_dimension(self) -> None: + """A scalar overlap should also work for 1D filtering.""" + patch = _patch((8, 64), ("distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + time=16, + overlap=7, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + + def test_zero_overlap_is_supported(self) -> None: + """Zero overlap should use an all-ones reconstruction taper.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter( + distance=16, + time=16, + overlap=0, + samples=True, + engine="scipy", + ) + + assert out.shape == patch.shape + assert np.isfinite(out.data).all() + + def test_one_dimension_auto_uses_scipy(self) -> None: + """Auto mode should use SciPy for 1D filtering.""" + patch = _patch((8, 64), ("distance", "time"), dtype=np.float32) + + out = patch.adaptive_spectral_filter(time=16, samples=True, engine="auto") + + assert out.shape == patch.shape + + def test_numba_rejects_one_dimension(self) -> None: + """The optional numba engine is intentionally 2D-only.""" + patch = _patch((8, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(ParameterError, match="two selected dimensions"): + patch.adaptive_spectral_filter(time=16, samples=True, engine="numba") + + @pytest.mark.parametrize( + "kwargs,match", + [ + ({"exponent": np.nan}, "exponent must be finite"), + ({"distance": 15}, "power of two"), + ({"overlap": {"distance": 8}}, "too large"), + ({"overlap": {"distance": -1}}, "non-negative"), + ], + ) + def test_patch_validation_branches( + self, kwargs: dict[str, Any], match: str + ) -> None: + """Patch-level validation should raise ParameterError.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + call_kwargs = {"distance": 16, "time": 16, "samples": True, "engine": "scipy"} + call_kwargs.update(kwargs) + + with pytest.raises(ParameterError, match=match): + patch.adaptive_spectral_filter(**call_kwargs) + + def test_rejects_missing_dimension_kwarg(self) -> None: + """Unknown dimensions should raise the normal patch coordinate error.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(PatchCoordinateError, match="not found"): + patch.adaptive_spectral_filter( + distance=16, missing=16, samples=True, engine="scipy" + ) + + def test_uneven_coordinate_conversion_raises_when_not_samples(self) -> None: + """Coordinate-unit windows require evenly sampled coordinates.""" + patch = dc.get_example_patch("wacky_dim_coords_patch") + + with pytest.raises(CoordError): + patch.adaptive_spectral_filter( + distance=16, time=16, samples=False, engine="scipy" + ) + + def test_nan_values_remain_supported(self) -> None: + """NaNs may propagate but should not produce infinities.""" + patch = dc.get_example_patch("patch_with_null", shape=(64, 64)) + + out = patch.adaptive_spectral_filter( + distance=16, time=16, samples=True, engine="scipy" + ) + out_data = np.asarray(out.data) + + assert out.shape == patch.shape + assert np.isnan(out_data).any() + assert not np.isinf(out_data).any() + + def test_invalid_engine_raises(self) -> None: + """Engine values should be constrained.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(ParameterError, match="engine"): + patch.adaptive_spectral_filter( + distance=16, + time=16, + samples=True, + engine="bad", # type: ignore[arg-type] + ) + + def test_proc_export_is_function(self) -> None: + """The processing module should expose the direct patch function.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + out = adaptive_spectral_filter_func( + patch, distance=16, time=16, samples=True, engine="scipy" + ) + + assert out.shape == patch.shape + + +class TestAdaptiveSpectralCore: + """Tests for plain array adaptive spectral helpers.""" + + def test_triangular_taper_values(self) -> None: + """The shared taper should match the overlap-add ramp geometry.""" + taper = _triangular_taper((8, 8), (2, 2)) + expected_1d = np.array([0.25, 0.5, 0.75, 1.0, 1.0, 0.75, 0.5, 0.25]) + + np.testing.assert_allclose(taper, expected_1d[:, None] * expected_1d[None, :]) + assert taper.dtype == np.float32 + + def test_triangular_taper_all_ones_when_plateau_matches_window(self) -> None: + """A full-window plateau should support zero-overlap reconstruction.""" + taper = _triangular_taper((8, 8), (8, 8)) + + np.testing.assert_array_equal(taper, np.ones((8, 8), dtype=np.float32)) + + def test_triangular_taper_one_dimensional(self) -> None: + """The shared taper should also support 1D reconstruction.""" + taper = _triangular_taper((8,), (2,)) + expected = np.array([0.25, 0.5, 0.75, 1.0, 1.0, 0.75, 0.5, 0.25]) + + np.testing.assert_allclose(taper, expected) + assert taper.dtype == np.float32 + + def test_triangular_taper_cache_is_not_mutated_by_callers(self) -> None: + """Callers should receive a copy of the cached taper.""" + taper = _triangular_taper((16, 16), (2, 2)) + expected = taper.copy() + + taper[...] = -1.0 + actual = _triangular_taper((16, 16), (2, 2)) + + np.testing.assert_array_equal(actual, expected) + + @pytest.mark.parametrize( + "window_size,plateau,match", + [ + ((16, 16), (17, 2), "Plateau cannot"), + ((16, 16), (-1, 2), "non-negative"), + ((15, 16), (2, 2), "Window sizes must be even"), + ((16, 16), (2,), "same length"), + ((16, 16, 16), (2, 2, 2), "one- and two-dimensional"), + ], + ) + def test_triangular_taper_rejects_invalid_geometry( + self, + window_size: tuple[int, int], + plateau: tuple[int, int], + match: str, + ) -> None: + """Invalid taper geometry should raise.""" + with pytest.raises(ValueError, match=match): + _triangular_taper(window_size, plateau) + + @pytest.mark.parametrize( + "window_size,overlap,match", + [ + ((15, 16), (7, 7), "power of two"), + ((4, 16), (1, 7), "greater than 4"), + ((16, 16), (-1, 7), "non-negative"), + ((16, 16), (8, 7), "too large"), + ((16.0, 16), (7, 7), "must be an integer"), + ((16, 16), (7.0, 7), "must be an integer"), + ((16,), (7, 7), "match the input dimensionality"), + ], + ) + def test_core_rejects_invalid_window_and_overlap( + self, + window_size: tuple[Any, Any], + overlap: tuple[Any, Any], + match: str, + ) -> None: + """Direct array API should validate window geometry.""" + data = np.ones((32, 32), dtype=np.float32) + + with pytest.raises(ValueError, match=match): + _adaptive_spectral_filter_scipy( + data, window_size=window_size, overlap=overlap + ) + + def test_core_rejects_non_1d_or_2d_input(self) -> None: + """Direct array API is 1D or 2D only.""" + data = np.ones((2, 16, 16), dtype=np.float32) + + with pytest.raises(ValueError, match="1D or 2D input"): + _adaptive_spectral_filter_scipy(data, window_size=(16, 16), overlap=(7, 7)) + + def test_core_rejects_non_finite_exponent(self) -> None: + """Exponent must be finite.""" + data = np.ones((32, 32), dtype=np.float32) + + with pytest.raises(ValueError, match="exponent must be finite"): + _adaptive_spectral_filter_scipy( + data, window_size=(16, 16), overlap=(7, 7), exponent=np.nan + ) + + def test_direct_array_api_returns_float32_for_integer_inputs(self) -> None: + """Non-floating array inputs should return float32 outputs.""" + data = np.ones((32, 32), dtype=np.int16) + + out = _adaptive_spectral_filter_scipy( + data, window_size=(16, 16), overlap=(7, 7) + ) + + assert out.dtype == np.float32 + + def test_direct_array_api_normalizes_power(self) -> None: + """The SciPy path should run power normalization.""" + data = np.ones((32, 32), dtype=np.float32) + + out = _adaptive_spectral_filter_scipy( + data, + window_size=(16, 16), + overlap=(7, 7), + exponent=0.5, + normalize_power=True, + ) + + assert out.shape == data.shape + assert np.isfinite(out).all() + + def test_direct_array_api_filters_one_dimensional_data(self) -> None: + """The SciPy helper should support 1D arrays.""" + data = np.ones(32, dtype=np.float32) + + out = _adaptive_spectral_filter_scipy( + data, + window_size=(16,), + overlap=(7,), + exponent=0.5, + normalize_power=True, + ) + + assert out.shape == data.shape + assert np.isfinite(out).all() + + def test_direct_array_api_supports_zero_overlap(self) -> None: + """The direct SciPy helper should accept non-overlapping windows.""" + data = np.ones((32, 32), dtype=np.float32) + + out = _adaptive_spectral_filter_scipy( + data, window_size=(16, 16), overlap=(0, 0) + ) + + assert out.shape == data.shape + assert np.isfinite(out).all() + + def test_auto_engine_falls_back_when_numba_missing(self, monkeypatch) -> None: + """Auto engine should fall back to SciPy when optional deps are absent.""" + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "dascore.proc._adaptive_spectral_filter_numba": + raise ImportError("simulated missing numba engine") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert _get_engine("auto", 2) is _adaptive_spectral_filter_scipy + + def test_numba_engine_raises_when_missing(self, monkeypatch) -> None: + """Explicit numba engine should raise when optional deps are absent.""" + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "dascore.proc._adaptive_spectral_filter_numba": + raise ImportError("simulated missing numba engine") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + with pytest.raises(MissingOptionalDependencyError, match="engine='numba'"): + _get_engine("numba", 2) + + def test_get_engine_uses_scipy_for_one_dimensional_auto(self) -> None: + """Auto mode should use SciPy when one dimension is selected.""" + assert _get_engine("auto", 1) is _adaptive_spectral_filter_scipy + + def test_get_engine_rejects_numba_for_one_dimension(self) -> None: + """Numba is intentionally unavailable for the 1D helper path.""" + with pytest.raises(ParameterError, match="two selected dimensions"): + _get_engine("numba", 1) + + def test_private_window_overlap_validator_rejects_negative_overlap(self) -> None: + """The private window validator should guard negative overlaps.""" + with pytest.raises(ParameterError, match="non-negative"): + _validate_window_and_overlap(("distance", "time"), (16, 16), (-1, 7), 0.3) + + @pytest.mark.parametrize("exponent", [0.0, 0.3]) + def test_numba_and_scipy_match_when_numba_available(self, exponent) -> None: + """The optional Numba 2D path should match the SciPy implementation.""" + numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + rng = np.random.default_rng(20260511) + data = rng.normal(size=(32, 32)).astype(np.float32) + + numba = numba_mod._adaptive_spectral_filter_numba( + data, + window_size=(16, 16), + overlap=(7, 7), + exponent=exponent, + normalize_power=True, + ) + scipy = _adaptive_spectral_filter_scipy( + data, + window_size=(16, 16), + overlap=(7, 7), + exponent=exponent, + normalize_power=True, + ) + + np.testing.assert_allclose(scipy, numba, rtol=1e-5, atol=1e-5) + + def test_auto_engine_uses_numba_when_available(self) -> None: + """Auto engine should use the Numba 2D path when optional deps import.""" + numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + + assert _get_engine("auto", 2) is numba_mod._adaptive_spectral_filter_numba + + def test_numba_private_helpers_run_in_python(self) -> None: + """The fast-engine helpers should be directly testable in Python.""" + numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + padded = np.arange(16, dtype=np.float32).reshape(4, 4) + tile = np.zeros((2, 2), dtype=np.float32) + + assert numba_mod._tile_indices_from_parity_index(3, 2, 1, 0) == (3, 2) + assert numba_mod._tile_bounds(1, 1, 2, 2, 1, 1, 4, 4) == (1, 1, 2, 2) + numba_mod._copy_padded_tile(padded, tile, 1, 1, 2, 2) + np.testing.assert_array_equal(tile, padded[1:3, 1:3]) + assert numba_mod._complex_power(3 + 4j) == np.float32(5.0) + + spec = np.array([[3 + 4j, 0j]], dtype=np.complex64) + assert numba_mod._max_spectral_power(spec) == np.float32(5.0) + assert numba_mod._max_spectral_power_numba_impl(spec) == np.float32(5.0) + weighted = spec.copy() + numba_mod._apply_spectral_weight(weighted, 1.0, False) + np.testing.assert_allclose(weighted[0, 0], spec[0, 0] * 5.0) + + weighted = spec.copy() + numba_mod._apply_spectral_weight(weighted, 0.3, True) + assert np.isfinite(weighted).all() + + weighted = spec.copy() + numba_mod._apply_spectral_weight_numba_impl(weighted, 0.3, True) + assert np.isfinite(weighted).all() + + weighted = spec.copy() + numba_mod._apply_spectral_weight_numba_impl(weighted, 1.0, False) + np.testing.assert_allclose(weighted[0, 0], spec[0, 0] * 5.0) + + zeros = np.array([[0j]], dtype=np.complex64) + numba_mod._apply_spectral_weight(zeros, 0.3, True) + assert zeros[0, 0] == 0j + + zeros = np.array([[0j]], dtype=np.complex64) + numba_mod._apply_spectral_weight_numba_impl(zeros, 0.3, True) + assert zeros[0, 0] == 0j + + filtered = np.zeros_like(padded) + taper = np.ones((2, 2), dtype=np.float32) + numba_mod._overlap_add_tile(filtered, tile, taper, 1, 1, 2, 2) + np.testing.assert_array_equal(filtered[1:3, 1:3], tile) + + def test_numba_private_tile_group_runs_in_python(self) -> None: + """The tile group algorithm should run without JIT for coverage.""" + numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + data = np.ones((8, 8), dtype=np.float32) + working, _, stride, taper, padded, filtered, n_tiles = ( + numba_mod._prepare_work_arrays(data, window_size=(8, 8), overlap=(3, 3)) + ) + + numba_mod._process_tile_group_python( + padded, + filtered, + taper, + 8, + 8, + stride[0], + stride[1], + n_tiles[0], + n_tiles[1], + 0, + 0, + 0.0, + False, + ) + numba_mod._process_tile_group_python( + padded, + filtered, + taper, + 8, + 8, + stride[0], + stride[1], + n_tiles[0], + n_tiles[1], + 0, + 0, + 0.5, + True, + ) + numba_mod._process_tile_group_numba_impl( + padded, + filtered, + taper, + 8, + 8, + stride[0], + stride[1], + n_tiles[0], + n_tiles[1], + 0, + 0, + 0.5, + True, + ) + out = numba_mod._finalize_output(filtered, working, data.dtype, stride) + + assert out.shape == data.shape + assert np.isfinite(out).all() diff --git a/tests/test_proc/test_taper.py b/tests/test_proc/test_taper.py index 62ff61527..051335aea 100644 --- a/tests/test_proc/test_taper.py +++ b/tests/test_proc/test_taper.py @@ -245,7 +245,7 @@ def test_poorly_shaped_sequence_raises(self, random_patch): def test_bad_use_of_none(self, random_patch): """Ensure bad use of None raises.""" - with pytest.raises(ParameterError, match="Cannot use ... or None"): + with pytest.raises(ParameterError, match=r"Cannot use \.\.\. or None"): random_patch.taper_range(time=(1, None), relative=True) def test_use_none(self, random_patch): From fe7751e65644fea9ca1a53a5ec38b505adab2b45 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Wed, 20 May 2026 15:00:55 +0200 Subject: [PATCH 2/8] update --- .../proc/_adaptive_spectral_filter_numba.py | 93 +++++++++++++------ dascore/proc/adaptive_spectral_filter.py | 23 +++-- dascore/utils/jit.py | 46 +++++++-- dascore/utils/misc.py | 17 ++++ dascore/utils/signal.py | 2 + 5 files changed, 134 insertions(+), 47 deletions(-) diff --git a/dascore/proc/_adaptive_spectral_filter_numba.py b/dascore/proc/_adaptive_spectral_filter_numba.py index 30288f17a..c071c218b 100644 --- a/dascore/proc/_adaptive_spectral_filter_numba.py +++ b/dascore/proc/_adaptive_spectral_filter_numba.py @@ -2,15 +2,16 @@ from __future__ import annotations -import numba as nb import numpy as np -import rocket_fft # noqa: F401 # registers FFT overloads with numba. from dascore.proc.adaptive_spectral_filter import ( _finalize_output, _prepare_work_arrays, _validate_filter_inputs, ) +from dascore.utils.jit import maybe_numba_jit + +_JIT_DEPS = ("rocket_fft",) def _tile_indices_from_parity_index( @@ -20,14 +21,14 @@ def _tile_indices_from_parity_index( parity1: int, ) -> tuple[int, int]: """Map a flattened parity-local index back to full-grid tile indices.""" - ix = parity0 + 2 * (ind // count1) - iy = parity1 + 2 * (ind % count1) - return ix, iy + x_index = parity0 + 2 * (ind // count1) + y_index = parity1 + 2 * (ind % count1) + return x_index, y_index def _tile_bounds( - ix: int, - iy: int, + x_index: int, + y_index: int, wx: int, wy: int, stride0: int, @@ -36,8 +37,8 @@ def _tile_bounds( shape1: int, ) -> tuple[int, int, int, int]: """Return padded-array origin and valid tile shape for one window.""" - beg0 = ix * stride0 - beg1 = iy * stride1 + beg0 = x_index * stride0 + beg1 = y_index * stride1 end0 = min(beg0 + wx, shape0) end1 = min(beg1 + wy, shape1) return beg0, beg1, end0 - beg0, end1 - beg1 @@ -110,12 +111,18 @@ def _overlap_add_tile( filtered[beg0 + i, beg1 + j] += tile[i, j] * taper[i, j] -_tile_indices_from_parity_index_numba = nb.njit(cache=True, inline="always")( - _tile_indices_from_parity_index -) -_tile_bounds_numba = nb.njit(cache=True, inline="always")(_tile_bounds) -_copy_padded_tile_numba = nb.njit(cache=True, inline="always")(_copy_padded_tile) -_complex_power_numba = nb.njit(cache=True, inline="always")(_complex_power) +_tile_indices_from_parity_index_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_tile_indices_from_parity_index) +_tile_bounds_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_tile_bounds) +_copy_padded_tile_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_copy_padded_tile) +_complex_power_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_complex_power) def _max_spectral_power_numba_impl(spec: np.ndarray) -> np.float32: @@ -151,13 +158,15 @@ def _apply_spectral_weight_numba_impl( spec[i, j] *= weight -_max_spectral_power_numba = nb.njit(cache=True, inline="always")( - _max_spectral_power_numba_impl -) -_apply_spectral_weight_numba = nb.njit(cache=True, inline="always")( - _apply_spectral_weight_numba_impl -) -_overlap_add_tile_numba = nb.njit(cache=True, inline="always")(_overlap_add_tile) +_max_spectral_power_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_max_spectral_power_numba_impl) +_apply_spectral_weight_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_apply_spectral_weight_numba_impl) +_overlap_add_tile_numba = maybe_numba_jit( + required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" +)(_overlap_add_tile) def _process_tile_group_python( @@ -180,9 +189,18 @@ def _process_tile_group_python( count1 = (ny - parity1 + 1) // 2 count = count0 * count1 for ind in range(count): - ix, iy = _tile_indices_from_parity_index(ind, count1, parity0, parity1) + x_index, y_index = _tile_indices_from_parity_index( + ind, count1, parity0, parity1 + ) beg0, beg1, n0, n1 = _tile_bounds( - ix, iy, wx, wy, stride0, stride1, padded.shape[0], padded.shape[1] + x_index, + y_index, + wx, + wy, + stride0, + stride1, + padded.shape[0], + padded.shape[1], ) tile = np.zeros((wx, wy), dtype=np.float32) @@ -215,10 +233,19 @@ def _process_tile_group_numba_impl( count0 = (nx - parity0 + 1) // 2 count1 = (ny - parity1 + 1) // 2 count = count0 * count1 - for ind in nb.prange(count): # type: ignore[not-iterable] - ix, iy = _tile_indices_from_parity_index_numba(ind, count1, parity0, parity1) + for ind in numba.prange(count): # noqa: F821 + x_index, y_index = _tile_indices_from_parity_index_numba( + ind, count1, parity0, parity1 + ) beg0, beg1, n0, n1 = _tile_bounds_numba( - ix, iy, wx, wy, stride0, stride1, padded.shape[0], padded.shape[1] + x_index, + y_index, + wx, + wy, + stride0, + stride1, + padded.shape[0], + padded.shape[1], ) tile = np.zeros((wx, wy), dtype=np.float32) @@ -234,9 +261,15 @@ def _process_tile_group_numba_impl( # fastmath is intentional here: the weighting is approximate and tests allow # small SciPy/Numba differences from parallel floating-point evaluation. -_process_tile_group_numba = nb.njit(cache=True, fastmath=True, parallel=True)( - _process_tile_group_numba_impl -) +_process_tile_group_numba = maybe_numba_jit( + required=True, + deps=_JIT_DEPS, + nopython=True, + cache=True, + fastmath=True, + parallel=True, +)(_process_tile_group_numba_impl) +_NUMBA_ENGINE_AVAILABLE = _process_tile_group_numba.jit_available def _adaptive_spectral_filter_numba( diff --git a/dascore/proc/adaptive_spectral_filter.py b/dascore/proc/adaptive_spectral_filter.py index 73f5ee71d..b1745d379 100644 --- a/dascore/proc/adaptive_spectral_filter.py +++ b/dascore/proc/adaptive_spectral_filter.py @@ -44,6 +44,7 @@ from dascore.constants import PatchType from dascore.exceptions import MissingOptionalDependencyError, ParameterError +from dascore.utils.misc import is_power_of_two from dascore.utils.patch import get_dim_axis_value, patch_function from dascore.utils.signal import _triangular_taper @@ -51,11 +52,6 @@ __all__ = ("adaptive_spectral_filter",) -def _is_power_of_two(value: int) -> bool: - """Return ``True`` when *value* is a positive power of two.""" - return value > 0 and (value & (value - 1) == 0) - - def _validate_filter_inputs( data: np.ndarray, *, @@ -77,16 +73,16 @@ def _validate_filter_inputs( raise ValueError(msg) for axis, (window, axis_overlap) in enumerate(zip(window_size, overlap)): - if not isinstance(window, (int, np.integer)): + if not isinstance(window, int | np.integer): msg = f"window_size[{axis}] must be an integer; got {window!r}." raise ValueError(msg) - if not isinstance(axis_overlap, (int, np.integer)): + if not isinstance(axis_overlap, int | np.integer): msg = f"overlap[{axis}] must be an integer; got {axis_overlap!r}." raise ValueError(msg) window = int(window) axis_overlap = int(axis_overlap) - if window <= 4 or not _is_power_of_two(window): + if window <= 4 or not is_power_of_two(window): msg = ( f"window_size[{axis}] must be a power of two greater than 4; " f"got {window!r}." @@ -331,7 +327,7 @@ def _validate_window_and_overlap( msg = "exponent must be finite." raise ParameterError(msg) for dim, window, overlap in zip(dims, windows, overlaps): - if window <= 4 or not _is_power_of_two(window): + if window <= 4 or not is_power_of_two(window): msg = f"window size for {dim!r} must be a power of two and > 4." raise ParameterError(msg) if overlap < 0: @@ -357,6 +353,7 @@ def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable raise ParameterError(msg) try: from dascore.proc._adaptive_spectral_filter_numba import ( + _NUMBA_ENGINE_AVAILABLE, _adaptive_spectral_filter_numba, ) except ImportError as exc: @@ -367,6 +364,14 @@ def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable ) raise MissingOptionalDependencyError(msg) from exc return _adaptive_spectral_filter_scipy + if not _NUMBA_ENGINE_AVAILABLE: + if engine == "numba": + msg = ( + "engine='numba' requires optional dependencies numba and " + "rocket-fft to be installed." + ) + raise MissingOptionalDependencyError(msg) + return _adaptive_spectral_filter_scipy return _adaptive_spectral_filter_numba diff --git a/dascore/utils/jit.py b/dascore/utils/jit.py index b47a5f24c..c46af8252 100644 --- a/dascore/utils/jit.py +++ b/dascore/utils/jit.py @@ -4,6 +4,7 @@ from __future__ import annotations +import importlib import warnings from functools import wraps @@ -18,17 +19,28 @@ def prange(self, count): yield from range(count) -def maybe_numba_jit(required=False, _missing_numba=False, **compiler_kwargs): +def maybe_numba_jit( + required=False, + deps: str | tuple[str, ...] = (), + _missing_numba=False, + _missing_deps: str | tuple[str, ...] = (), + **compiler_kwargs, +): """ Use numba to apply JIT compilation to the decorated function. Parameters ---------- required - If True an ImportError is raised if the wrapped function is called - and the compiler module is not installed. If False, issue a warning. + If True an ImportError is raised if the wrapped function is called and + the compiler module or required dependencies are not installed. If + False, issue a warning. + deps + Extra importable modules required before compilation can occur. _missing_numba If true, simulate missing the numba package. Only used for testing. + _missing_deps + Extra dependencies to treat as missing. Only used for testing. **compiler_kwargs Keyword arguments passed to the compiler function. @@ -65,7 +77,12 @@ def maybe_numba_jit(required=False, _missing_numba=False, **compiler_kwargs): if callable(required): return maybe_numba_jit()(required) - has_numba = True + deps = (deps,) if isinstance(deps, str) else tuple(deps) + _missing_deps = ( + (_missing_deps,) if isinstance(_missing_deps, str) else tuple(_missing_deps) + ) + + missing_modules = [] try: import numba @@ -73,7 +90,17 @@ def maybe_numba_jit(required=False, _missing_numba=False, **compiler_kwargs): raise ImportError("Simulating missing numba.") except ImportError: numba = _DummyNumba() - has_numba = False + missing_modules.append("numba") + + for module_name in deps: + try: + if module_name in _missing_deps: + raise ImportError(f"Simulating missing {module_name}.") + importlib.import_module(module_name) + except ImportError: + missing_modules.append(module_name) + + has_all_deps = not missing_modules def _wrapper(func): # Add numba to the functions global namespace so that it can be used @@ -81,20 +108,21 @@ def _wrapper(func): globs = getattr(func, "__globals__", {}) globs["numba"] = numba - if not has_numba: + if not has_all_deps: @wraps(func) def decorated(*args, **kwargs): + module_names = ", ".join(missing_modules) if required: msg = ( f"{func.__name__} requires python module " - f"numba but it is not installed. " + f"{module_names} but it is not installed. " ) raise ImportError(msg) else: msg = ( f"{func.__name__} can be compiled to improve performance. " - f"Please install numba to enable JIT." + f"Please install {module_names} to enable JIT." ) warnings.warn(msg, UserWarning) return func(*args, **kwargs) @@ -103,6 +131,8 @@ def decorated(*args, **kwargs): else: out_func = numba.jit(**compiler_kwargs)(func) out_func.func = func # make original func accessible via .func + out_func.jit_available = has_all_deps + out_func.missing_jit_deps = tuple(missing_modules) return out_func return _wrapper diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 1b4b37391..921f1bde3 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -910,3 +910,20 @@ def tukey_fence(data, fence_multiplier=1.5) -> np.ndarray: q_upper = np.nanmin([q3 + diff * fence_multiplier, dmax]) lower_and_top = np.asarray([q_lower, q_upper]) return lower_and_top + + +def is_power_of_two(value: int) -> bool: + """ + Return ``True`` when *value* is a positive power of two. + + Powers of two have exactly one set bit in their binary representation. + Subtracting one flips that bit and all lower bits, so ``value & (value - 1)`` + is zero only when ``value`` had a single set bit. The ``value > 0`` check + excludes zero and negative values. + + Parameters + ---------- + values + The value to test. + """ + return value > 0 and (value & (value - 1) == 0) diff --git a/dascore/utils/signal.py b/dascore/utils/signal.py index a2837b20e..5bcbee24f 100644 --- a/dascore/utils/signal.py +++ b/dascore/utils/signal.py @@ -2,6 +2,8 @@ Utilities for signal processing. """ +from __future__ import annotations + from functools import lru_cache import numpy as np From ba53039e59097ab6f2c35fb23248f7f2634214b0 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 15:41:37 +0200 Subject: [PATCH 3/8] Skip the numba engine tests where numba is not installed The min-deps and free-threaded cells have neither numba nor rocket-fft. The tests guarded with importorskip on the engine module, but that module imports either way -- which is exactly what lets engine='auto' fall back -- so nothing skipped and five tests failed on the ImportError raised when the jit wrapper was called. _NUMBA_ENGINE_AVAILABLE is the flag that actually answers the question. --- .../test_adaptive_spectral_filter.py | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/test_proc/test_adaptive_spectral_filter.py b/tests/test_proc/test_adaptive_spectral_filter.py index bab1abd71..cff6b0809 100644 --- a/tests/test_proc/test_adaptive_spectral_filter.py +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -24,6 +24,21 @@ from dascore.utils.signal import _triangular_taper +def _numba_engine(): + """ + Return the optional numba engine module, or skip the test. + + The module imports whether or not numba is installed -- that is what + lets `engine="auto"` fall back -- so importing it proves nothing. Only + the flag says whether its functions can actually run. + """ + name = "dascore.proc._adaptive_spectral_filter_numba" + module = pytest.importorskip(name) + if not module._NUMBA_ENGINE_AVAILABLE: + pytest.skip("numba and rocket-fft are not installed") + return module + + def _patch( shape: tuple[int, ...], dims: tuple[str, ...], @@ -682,7 +697,7 @@ def test_private_window_overlap_validator_rejects_negative_overlap(self) -> None @pytest.mark.parametrize("exponent", [0.0, 0.3]) def test_numba_and_scipy_match_when_numba_available(self, exponent) -> None: """The optional Numba 2D path should match the SciPy implementation.""" - numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + numba_mod = _numba_engine() rng = np.random.default_rng(20260511) data = rng.normal(size=(32, 32)).astype(np.float32) @@ -705,13 +720,13 @@ def test_numba_and_scipy_match_when_numba_available(self, exponent) -> None: def test_auto_engine_uses_numba_when_available(self) -> None: """Auto engine should use the Numba 2D path when optional deps import.""" - numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + numba_mod = _numba_engine() assert _get_engine("auto", 2) is numba_mod._adaptive_spectral_filter_numba def test_numba_private_helpers_run_in_python(self) -> None: """The fast-engine helpers should be directly testable in Python.""" - numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + numba_mod = _numba_engine() padded = np.arange(16, dtype=np.float32).reshape(4, 4) tile = np.zeros((2, 2), dtype=np.float32) @@ -755,7 +770,7 @@ def test_numba_private_helpers_run_in_python(self) -> None: def test_numba_private_tile_group_runs_in_python(self) -> None: """The tile group algorithm should run without JIT for coverage.""" - numba_mod = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + numba_mod = _numba_engine() data = np.ones((8, 8), dtype=np.float32) working, _, stride, taper, padded, filtered, n_tiles = ( numba_mod._prepare_work_arrays(data, window_size=(8, 8), overlap=(3, 3)) From 86cff5acc0ba6ddeab6b71f4f1de4f842a97f9a4 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 15:57:49 +0200 Subject: [PATCH 4/8] Test the two things this branch added and never covered The coverage gate found both. maybe_numba_jit grew deps and _missing_deps here, so the loop that imports them, and the message naming every missing module, had no test. And _get_engine's fallback on _NUMBA_ENGINE_AVAILABLE is separate from its fallback on a failed import: the engine module imports whether or not numba is installed, which is the whole point, so only the flag says the jit can actually run. --- .../test_adaptive_spectral_filter.py | 15 ++++++ tests/test_utils/test_jit.py | 46 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/tests/test_proc/test_adaptive_spectral_filter.py b/tests/test_proc/test_adaptive_spectral_filter.py index cff6b0809..d4d5dcb56 100644 --- a/tests/test_proc/test_adaptive_spectral_filter.py +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -666,6 +666,21 @@ def fake_import(name, *args, **kwargs): assert _get_engine("auto", 2) is _adaptive_spectral_filter_scipy + def test_auto_engine_falls_back_when_deps_are_absent(self, monkeypatch) -> None: + """The engine module imports without numba; the flag says it cannot run.""" + module = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + monkeypatch.setattr(module, "_NUMBA_ENGINE_AVAILABLE", False) + + assert _get_engine("auto", 2) is _adaptive_spectral_filter_scipy + + def test_numba_engine_raises_when_deps_are_absent(self, monkeypatch) -> None: + """Asking for it by name says which dependencies are wanted.""" + module = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") + monkeypatch.setattr(module, "_NUMBA_ENGINE_AVAILABLE", False) + + with pytest.raises(MissingOptionalDependencyError, match="engine='numba'"): + _get_engine("numba", 2) + def test_numba_engine_raises_when_missing(self, monkeypatch) -> None: """Explicit numba engine should raise when optional deps are absent.""" real_import = builtins.__import__ diff --git a/tests/test_utils/test_jit.py b/tests/test_utils/test_jit.py index 10f0e7adc..dad5c1bfa 100644 --- a/tests/test_utils/test_jit.py +++ b/tests/test_utils/test_jit.py @@ -48,6 +48,52 @@ def _jit_test_func(ar): with pytest.raises(ImportError, match=match): _jit_test_func(np.array([1, 2, 3])) + def test_extra_dep_present(self): + """A dep which imports leaves the jit available.""" + + @maybe_numba_jit(deps="json") + def _jit_test_func(ar): + return ar + + assert _jit_test_func.missing_jit_deps == () + + def test_missing_extra_dep_warns(self): + """A dep which does not import is named, numba being fine.""" + + @maybe_numba_jit(deps="rocket_fft", _missing_deps="rocket_fft") + def _jit_test_func(ar): + return ar + + assert not _jit_test_func.jit_available + assert _jit_test_func.missing_jit_deps == ("rocket_fft",) + with pytest.warns(UserWarning, match="rocket_fft"): + _jit_test_func(np.array([1, 2, 3])) + + def test_missing_extra_dep_raises(self): + """Every missing module is named, and said how to install.""" + + @maybe_numba_jit( + required=True, + deps=("rocket_fft",), + _missing_numba=True, + _missing_deps=("rocket_fft",), + ) + def _jit_test_func(ar): + return ar + + with pytest.raises(ImportError, match="numba, rocket_fft") as info: + _jit_test_func(np.array([1, 2, 3])) + assert "pip install numba rocket_fft" in str(info.value) + + def test_absent_dep_is_found_without_simulating_it(self): + """A module which really is not installed is reported as missing.""" + + @maybe_numba_jit(deps="dascore_not_a_real_module") + def _jit_test_func(ar): + return ar + + assert _jit_test_func.missing_jit_deps == ("dascore_not_a_real_module",) + def test_example(self): """Test docstring examples.""" pytest.importorskip("numba") From 9936c5e64db7afa11b37a14560e345c3ab017ba1 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 16:01:22 +0200 Subject: [PATCH 5/8] Say nothing about numba in the extra-dep tests The min-deps cells do not install numba, so missing_jit_deps holds it there and not here; asserting the whole tuple only passed on machines that happened to have it. The tests ask about the dep each one is for. --- tests/test_utils/test_jit.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_utils/test_jit.py b/tests/test_utils/test_jit.py index dad5c1bfa..5f86ec25a 100644 --- a/tests/test_utils/test_jit.py +++ b/tests/test_utils/test_jit.py @@ -48,24 +48,26 @@ def _jit_test_func(ar): with pytest.raises(ImportError, match=match): _jit_test_func(np.array([1, 2, 3])) + # These say nothing about numba itself: the min-deps cells do not + # install it, so it is in missing_jit_deps there and not here. def test_extra_dep_present(self): - """A dep which imports leaves the jit available.""" + """A dep which imports is not one of the missing ones.""" @maybe_numba_jit(deps="json") def _jit_test_func(ar): return ar - assert _jit_test_func.missing_jit_deps == () + assert "json" not in _jit_test_func.missing_jit_deps def test_missing_extra_dep_warns(self): - """A dep which does not import is named, numba being fine.""" + """A dep which does not import is named and warned about.""" @maybe_numba_jit(deps="rocket_fft", _missing_deps="rocket_fft") def _jit_test_func(ar): return ar assert not _jit_test_func.jit_available - assert _jit_test_func.missing_jit_deps == ("rocket_fft",) + assert "rocket_fft" in _jit_test_func.missing_jit_deps with pytest.warns(UserWarning, match="rocket_fft"): _jit_test_func(np.array([1, 2, 3])) @@ -92,7 +94,7 @@ def test_absent_dep_is_found_without_simulating_it(self): def _jit_test_func(ar): return ar - assert _jit_test_func.missing_jit_deps == ("dascore_not_a_real_module",) + assert "dascore_not_a_real_module" in _jit_test_func.missing_jit_deps def test_example(self): """Test docstring examples.""" From c0e10d53b4710e955377604486090764cc449943 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 20:10:42 +0200 Subject: [PATCH 6/8] Show the adaptive spectral filter working, and make it one kernel The default exponent of 0.3 recovered a fraction of what Lightguide's 0.8 does on the same data (3.6 dB against 9.9 dB of SNR gain on a synthetic event), and nothing in the suite asked whether the filter helped at all. The defaults now match the reference implementation, whose output the SciPy engine reproduces to 1e-7, and a TestEfficacy class asserts the filter recovers buried arrivals, that two dimensions beat one, and that the exponent is the strength. The numba engine carried every helper twice, once for Python and once for compilation, plus a third copy of the tile loop, all kept alive by coverage tests of the copies. It is now a single kernel which runs uncompiled through .func for the same coverage, at the same speed. The two validators share one set of rules, the import-failure branch a module designed to import without numba could never take is gone, and the engine helpers hand back four values instead of seven. A recipe shows the filter on a synthetic event, both example events, an exponent sweep, and the two things it does not do: keep amplitude, and remove coherent noise. --- .../proc/_adaptive_spectral_filter_numba.py | 355 +++--------------- dascore/proc/adaptive_spectral_filter.py | 309 +++++++-------- docs/recipes/adaptive_spectral_filter.qmd | 138 +++++++ scripts/_templates/_quarto.yml | 1 + .../test_adaptive_spectral_filter.py | 245 ++++++------ 5 files changed, 460 insertions(+), 588 deletions(-) create mode 100644 docs/recipes/adaptive_spectral_filter.qmd diff --git a/dascore/proc/_adaptive_spectral_filter_numba.py b/dascore/proc/_adaptive_spectral_filter_numba.py index 398f04fc4..b828d0fe9 100644 --- a/dascore/proc/_adaptive_spectral_filter_numba.py +++ b/dascore/proc/_adaptive_spectral_filter_numba.py @@ -1,4 +1,9 @@ -"""Optional Numba/rocket-fft engine for adaptive spectral filtering.""" +""" +Optional Numba/rocket-fft engine for the two-dimensional adaptive spectral filter. + +The module imports whether or not numba and rocket-fft are installed; only +``_NUMBA_ENGINE_AVAILABLE`` says whether the kernel can actually compile. +""" from __future__ import annotations @@ -11,265 +16,62 @@ ) from dascore.utils.jit import maybe_numba_jit -_JIT_DEPS = ("rocket_fft",) - - -def _tile_indices_from_parity_index( - ind: int, - count1: int, - parity0: int, - parity1: int, -) -> tuple[int, int]: - """Map a flattened parity-local index back to full-grid tile indices.""" - x_index = parity0 + 2 * (ind // count1) - y_index = parity1 + 2 * (ind % count1) - return x_index, y_index - - -def _tile_bounds( - x_index: int, - y_index: int, - wx: int, - wy: int, - stride0: int, - stride1: int, - shape0: int, - shape1: int, -) -> tuple[int, int, int, int]: - """Return padded-array origin and valid tile shape for one window.""" - beg0 = x_index * stride0 - beg1 = y_index * stride1 - end0 = min(beg0 + wx, shape0) - end1 = min(beg1 + wy, shape1) - return beg0, beg1, end0 - beg0, end1 - beg1 - - -def _copy_padded_tile( - padded: np.ndarray, - tile: np.ndarray, - beg0: int, - beg1: int, - n0: int, - n1: int, -) -> None: - """Copy the valid padded-array region into a fixed-shape zeroed tile.""" - for i in range(n0): - for j in range(n1): - tile[i, j] = padded[beg0 + i, beg1 + j] - - -def _complex_power(value: complex) -> np.float32: - """Return the magnitude of a complex FFT coefficient as ``float32``.""" - return np.float32((value.real * value.real + value.imag * value.imag) ** 0.5) - - -def _max_spectral_power(spec: np.ndarray) -> np.float32: - """Return the maximum spectral magnitude in one tile.""" - max_power = np.float32(0.0) - for i in range(spec.shape[0]): - for j in range(spec.shape[1]): - power = _complex_power(spec[i, j]) - if power > max_power: - max_power = power - return max_power - - -def _apply_spectral_weight( - spec: np.ndarray, - exponent: float, - normalize_power: bool, -) -> None: - """Apply adaptive magnitude weighting to one tile spectrum in place.""" - max_power = np.float32(0.0) - if normalize_power: - max_power = _max_spectral_power(spec) - - for i in range(spec.shape[0]): - for j in range(spec.shape[1]): - power = _complex_power(spec[i, j]) - if normalize_power: - if max_power != 0.0: - power = power / max_power - else: - power = np.float32(0.0) - weight = np.float32(power**exponent) - spec[i, j] *= weight - - -def _overlap_add_tile( - filtered: np.ndarray, - tile: np.ndarray, - taper: np.ndarray, - beg0: int, - beg1: int, - n0: int, - n1: int, -) -> None: - """Accumulate the valid region of one filtered tile into the output.""" - for i in range(n0): - for j in range(n1): - filtered[beg0 + i, beg1 + j] += tile[i, j] * taper[i, j] - - -_tile_indices_from_parity_index_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_tile_indices_from_parity_index) -_tile_bounds_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_tile_bounds) -_copy_padded_tile_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_copy_padded_tile) -_complex_power_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_complex_power) - - -def _max_spectral_power_numba_impl(spec: np.ndarray) -> np.float32: - """Return the maximum spectral magnitude using compiled helpers.""" - max_power = np.float32(0.0) - for i in range(spec.shape[0]): - for j in range(spec.shape[1]): - power = _complex_power_numba(spec[i, j]) - if power > max_power: - max_power = power - return max_power - - -def _apply_spectral_weight_numba_impl( - spec: np.ndarray, - exponent: float, - normalize_power: bool, -) -> None: - """Apply adaptive magnitude weighting using compiled helpers.""" - max_power = np.float32(0.0) - if normalize_power: - max_power = _max_spectral_power_numba(spec) - - for i in range(spec.shape[0]): - for j in range(spec.shape[1]): - power = _complex_power_numba(spec[i, j]) - if normalize_power: - if max_power != 0.0: - power = power / max_power - else: - power = np.float32(0.0) - weight = np.float32(power**exponent) - spec[i, j] *= weight - -_max_spectral_power_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_max_spectral_power_numba_impl) -_apply_spectral_weight_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_apply_spectral_weight_numba_impl) -_overlap_add_tile_numba = maybe_numba_jit( - required=True, deps=_JIT_DEPS, nopython=True, cache=True, inline="always" -)(_overlap_add_tile) - - -def _process_tile_group_python( - padded: np.ndarray, - filtered: np.ndarray, - taper: np.ndarray, - wx: int, - wy: int, - stride0: int, - stride1: int, - nx: int, - ny: int, - parity0: int, - parity1: int, - exponent: float, - normalize_power: bool, -) -> None: - """Process one non-overlapping tile parity group in pure Python.""" - count0 = (nx - parity0 + 1) // 2 - count1 = (ny - parity1 + 1) // 2 - count = count0 * count1 - for ind in range(count): - x_index, y_index = _tile_indices_from_parity_index( - ind, count1, parity0, parity1 - ) - beg0, beg1, n0, n1 = _tile_bounds( - x_index, - y_index, - wx, - wy, - stride0, - stride1, - padded.shape[0], - padded.shape[1], - ) - - tile = np.zeros((wx, wy), dtype=np.float32) - _copy_padded_tile(padded, tile, beg0, beg1, n0, n1) - - spec = np.fft.rfft2(tile) - if exponent != 0.0: - _apply_spectral_weight(spec, exponent, normalize_power) - - tile = np.fft.irfft2(spec, s=(wx, wy)) - _overlap_add_tile(filtered, tile, taper, beg0, beg1, n0, n1) - - -def _process_tile_group_numba_impl( +# fastmath is intentional: the weighting is approximate, and tests allow small +# SciPy/Numba differences from parallel floating-point evaluation. +@maybe_numba_jit( + required=True, + deps="rocket_fft", + nopython=True, + cache=True, + fastmath=True, + parallel=True, +) +def _filter_tile_group( padded: np.ndarray, filtered: np.ndarray, taper: np.ndarray, - wx: int, - wy: int, + window0: int, + window1: int, stride0: int, stride1: int, - nx: int, - ny: int, + n_tiles0: int, + n_tiles1: int, parity0: int, parity1: int, exponent: float, normalize_power: bool, ) -> None: - """Process one non-overlapping tile parity group with compiled helpers.""" - count0 = (nx - parity0 + 1) // 2 - count1 = (ny - parity1 + 1) // 2 - count = count0 * count1 - for ind in numba.prange(count): # noqa: F821 # ty: ignore[unresolved-reference] - x_index, y_index = _tile_indices_from_parity_index_numba( - ind, count1, parity0, parity1 - ) - beg0, beg1, n0, n1 = _tile_bounds_numba( - x_index, - y_index, - wx, - wy, - stride0, - stride1, - padded.shape[0], - padded.shape[1], - ) - - tile = np.zeros((wx, wy), dtype=np.float32) - _copy_padded_tile_numba(padded, tile, beg0, beg1, n0, n1) + """ + Filter every tile whose grid indices share a parity, adding into filtered. + Same-parity tiles start two strides apart, and an overlap under half the + window keeps a window shorter than two strides, so no two iterations of + the parallel loop write to the same output sample. + """ + count0 = (n_tiles0 - parity0 + 1) // 2 + count1 = (n_tiles1 - parity1 + 1) // 2 + for ind in numba.prange(count0 * count1): # noqa: F821 # ty: ignore[unresolved-reference] + beg0 = (parity0 + 2 * (ind // count1)) * stride0 + beg1 = (parity1 + 2 * (ind % count1)) * stride1 + n0 = min(window0, padded.shape[0] - beg0) + n1 = min(window1, padded.shape[1] - beg1) + tile = np.zeros((window0, window1), dtype=np.float32) + tile[:n0, :n1] = padded[beg0 : beg0 + n0, beg1 : beg1 + n1] spec = np.fft.rfft2(tile) if exponent != 0.0: - _apply_spectral_weight_numba(spec, exponent, normalize_power) - - tile = np.fft.irfft2(spec, s=(wx, wy)) - _overlap_add_tile_numba(filtered, tile, taper, beg0, beg1, n0, n1) + power = np.abs(spec) + if normalize_power: + # A silent tile weights to zero rather than dividing by it. + max_power = power.max() + power = power / max_power if max_power > 0.0 else power + # float32 so the weighted spectrum keeps rfft2's complex64 type. + spec = spec * (power**exponent).astype(np.float32) + tile = np.fft.irfft2(spec, s=(window0, window1)) + filtered[beg0 : beg0 + n0, beg1 : beg1 + n1] += tile[:n0, :n1] * taper[:n0, :n1] -# fastmath is intentional here: the weighting is approximate and tests allow -# small SciPy/Numba differences from parallel floating-point evaluation. -_process_tile_group_numba = maybe_numba_jit( - required=True, - deps=_JIT_DEPS, - nopython=True, - cache=True, - fastmath=True, - parallel=True, -)(_process_tile_group_numba_impl) -_NUMBA_ENGINE_AVAILABLE = _process_tile_group_numba.jit_available +_NUMBA_ENGINE_AVAILABLE = _filter_tile_group.jit_available def _adaptive_spectral_filter_numba( @@ -277,74 +79,39 @@ def _adaptive_spectral_filter_numba( *, window_size: tuple[int, int], overlap: tuple[int, int], - exponent: float = 0.3, + exponent: float = 0.8, normalize_power: bool = False, ) -> np.ndarray: """ Filter a 2D array with the optional Numba/rocket-fft implementation. - Parameters - ---------- - data - Two-dimensional input array. The filter computes in ``float32``. - window_size - Two power-of-two window lengths, one per array axis. Values must be - greater than 4. - overlap - Number of samples each neighboring window overlaps on each axis. Each - value must be non-negative and smaller than half the matching window. - exponent - Spectral magnitude exponent used as the adaptive weighting power. ``0`` - leaves the spectrum unweighted before overlap-add reconstruction. - normalize_power - If ``True``, normalize each tile's spectral magnitudes by that tile's - maximum magnitude before applying ``exponent``. - - Returns - ------- - numpy.ndarray - The filtered array with the same shape as ``data``. Floating input - dtypes are restored; non-floating inputs return ``float32`` output. - - Raises - ------ - ValueError - If ``data`` is not two-dimensional, ``exponent`` is not finite, - ``window_size`` and ``overlap`` do not contain exactly two integer - values, any window size is not a power of two greater than 4, or any - overlap is negative or at least half the matching window size. - - Notes - ----- - This implementation uses Numba-compiled loops and rocket-fft-backed NumPy - FFT calls. It is selected by - :func:`dascore.proc.adaptive_spectral_filter.adaptive_spectral_filter` for - two selected dimensions when ``engine="numba"`` or when ``engine="auto"`` - and optional dependencies are installed. + Takes the arguments :func:`_adaptive_spectral_filter_scipy` takes and + returns what it returns, to within floating-point evaluation order; the + two-dimensional restriction is the only difference. """ data = np.asarray(data) _validate_filter_inputs( data, window_size=window_size, overlap=overlap, exponent=float(exponent) ) - wx, wy = window_size - working, original_dtype, stride, taper, padded, filtered, n_tiles = ( - _prepare_work_arrays(data, window_size=window_size, overlap=overlap) + if data.ndim != 2: + msg = "The numba engine filters two-dimensional arrays only." + raise ValueError(msg) + padded, taper, stride, n_tiles = _prepare_work_arrays( + data, window_size=window_size, overlap=overlap ) + filtered = np.zeros_like(padded) for parity0 in range(2): for parity1 in range(2): - _process_tile_group_numba( + _filter_tile_group( padded, filtered, taper, - wx, - wy, - stride[0], - stride[1], - n_tiles[0], - n_tiles[1], + *window_size, + *stride, + *n_tiles, parity0, parity1, float(exponent), bool(normalize_power), ) - return _finalize_output(filtered, working, original_dtype, stride) + return _finalize_output(filtered, data.shape, data.dtype, stride) diff --git a/dascore/proc/adaptive_spectral_filter.py b/dascore/proc/adaptive_spectral_filter.py index 071fc6d60..ec2239665 100644 --- a/dascore/proc/adaptive_spectral_filter.py +++ b/dascore/proc/adaptive_spectral_filter.py @@ -1,35 +1,25 @@ """ Adaptive spectral filtering for DASCore patches. -The adaptive spectral filter suppresses incoherent energy by processing a patch -in overlapping windows along one or two selected dimensions. Each window is -transformed to the spectral domain, weighted by a power of its spectral -magnitude, transformed back to the original domain, and accumulated with -tapered overlap-add reconstruction. - -With one selected dimension, this is an adaptive frequency-domain normalization -applied independently to every trace over the remaining patch dimensions. With -two selected dimensions, this is the adaptive frequency-wavenumber filter -described by @isken2022denoising and exposed by Pyrocko -[Lightguide](https://github.com/pyrocko/lightguide). Coherent plane-wave energy -tends to concentrate in the frequency-wavenumber spectrum, so the weighting -emphasizes locally coherent arrivals relative to diffuse or randomly -distributed energy. - -This module exposes a single public patch method, -:func:`adaptive_spectral_filter`. The public function resolves one or two -DASCore dimensions, converts window and overlap values to sample counts, moves -those dimensions to the array tail, and processes every remaining leading index -as an independent batch. The lower-level SciPy and Numba implementations are -private because they operate on raw arrays and do not perform DASCore -coordinate handling. - -The SciPy engine handles one- and two-dimensional selected windows using -``rfftn``/``irfftn``. The optional Numba/rocket-fft engine currently handles -the two-dimensional case only, using parity-separated tile groups so neighboring -writes do not overlap within each parallel loop. Both engines share validation, -padding, tapering, and dtype-restoration logic so two-dimensional outputs remain -directly comparable. +The filter walks a patch in overlapping windows along one or two dimensions. +Each window is transformed to the spectral domain, every coefficient is +weighted by a power of its own magnitude, and the window is transformed back +and added into the output under a tapered overlap-add. Energy which is +coherent within a window concentrates in a few large coefficients, which the +weighting keeps; energy spread across the spectrum is suppressed. + +Over two dimensions this is the adaptive frequency-wavenumber filter of +@isken2022denoising as implemented by Pyrocko +[Lightguide](https://github.com/pyrocko/lightguide), which the SciPy engine +here matches to floating-point precision. Over one dimension it is the same +weighting applied to each trace on its own. + +The public patch function converts dimension names and coordinate units to +axes and sample counts and batches over every dimension not selected. The +engines work on raw one- or two-dimensional arrays; the optional +Numba/rocket-fft engine in `_adaptive_spectral_filter_numba` handles the +two-dimensional case and shares this module's validation, padding, and taper +so the two produce the same output. """ from __future__ import annotations @@ -52,6 +42,34 @@ __all__ = ("adaptive_spectral_filter",) +def _check_window(window: Any, overlap: Any, label: str) -> None: + """Raise ValueError unless a window and its overlap can tile an axis.""" + if not isinstance(window, int | np.integer): + msg = f"window for {label} must be an integer; got {window!r}." + raise ValueError(msg) + if not isinstance(overlap, int | np.integer): + msg = f"overlap for {label} must be an integer; got {overlap!r}." + raise ValueError(msg) + if window <= 4 or not is_power_of_two(window): + msg = ( + f"window for {label} must be a power of two greater than 4; got {window!r}." + ) + raise ValueError(msg) + if overlap < 0: + msg = f"overlap for {label} must be non-negative; got {overlap!r}." + raise ValueError(msg) + if overlap >= window / 2: + msg = f"overlap for {label} is too large; maximum is {window // 2 - 1} samples." + raise ValueError(msg) + + +def _check_exponent(exponent: float) -> None: + """Raise ValueError unless the exponent is finite.""" + if not np.isfinite(exponent): + msg = "exponent must be finite." + raise ValueError(msg) + + def _validate_filter_inputs( data: np.ndarray, *, @@ -68,32 +86,24 @@ def _validate_filter_inputs( if len(window_size) != data.ndim or len(overlap) != data.ndim: msg = "window_size and overlap must match the input dimensionality." raise ValueError(msg) - if not np.isfinite(exponent): - msg = "exponent must be finite." - raise ValueError(msg) - + _check_exponent(exponent) for axis, (window, axis_overlap) in enumerate(zip(window_size, overlap)): - if not isinstance(window, int | np.integer): - msg = f"window_size[{axis}] must be an integer; got {window!r}." - raise ValueError(msg) - if not isinstance(axis_overlap, int | np.integer): - msg = f"overlap[{axis}] must be an integer; got {axis_overlap!r}." - raise ValueError(msg) - - window = int(window) - axis_overlap = int(axis_overlap) - if window <= 4 or not is_power_of_two(window): - msg = ( - f"window_size[{axis}] must be a power of two greater than 4; " - f"got {window!r}." - ) - raise ValueError(msg) - if axis_overlap < 0: - msg = f"overlap[{axis}] must be non-negative; got {axis_overlap!r}." - raise ValueError(msg) - if axis_overlap >= window / 2: - msg = f"overlap[{axis}] is too large; maximum is {window // 2 - 1} samples." - raise ValueError(msg) + _check_window(window, axis_overlap, f"axis {axis}") + + +def _validate_window_and_overlap( + dims: tuple[str, ...], + windows: tuple[int, ...], + overlaps: tuple[int, ...], + exponent: float, +) -> None: + """Validate the patch-level settings, naming dimensions rather than axes.""" + try: + _check_exponent(exponent) + for dim, window, overlap in zip(dims, windows, overlaps): + _check_window(window, overlap, repr(dim)) + except ValueError as exc: + raise ParameterError(str(exc)) from exc def _prepare_work_arrays( @@ -101,59 +111,50 @@ def _prepare_work_arrays( *, window_size: tuple[int, ...], overlap: tuple[int, ...], -) -> tuple[ - np.ndarray, - np.dtype, - tuple[int, ...], - np.ndarray, - np.ndarray, - np.ndarray, - tuple[int, ...], -]: - """Prepare ``float32`` padded arrays shared by filter implementations.""" - data = np.asarray(data) - original_dtype = data.dtype +) -> tuple[np.ndarray, np.ndarray, tuple[int, ...], tuple[int, ...]]: + """ + Return the padded float32 input, the taper, the stride, and the tile grid. + + The input is padded by one stride of zeros on every side so the tiles + which straddle its edges see a full taper ramp. + """ working = np.ascontiguousarray(data, dtype=np.float32) stride = tuple(win - over for win, over in zip(window_size, overlap)) plateau = tuple(win - 2 * over for win, over in zip(window_size, overlap)) taper = _triangular_taper(window_size, plateau) - padded_shape = tuple( length + 2 * step for length, step in zip(working.shape, stride) ) padded = np.zeros(padded_shape, dtype=np.float32) - inner_slices = tuple( + inner = tuple( slice(step, length + step) for length, step in zip(working.shape, stride) ) - padded[inner_slices] = working - filtered = np.zeros_like(padded) + padded[inner] = working n_tiles = tuple(pad_len // step for pad_len, step in zip(padded.shape, stride)) - return working, original_dtype, stride, taper, padded, filtered, n_tiles + return padded, taper, stride, n_tiles def _finalize_output( filtered: np.ndarray, - working: np.ndarray, - original_dtype: np.dtype, + shape: tuple[int, ...], + dtype: np.dtype, stride: tuple[int, ...], ) -> np.ndarray: - """Crop padded output and restore floating dtypes where possible.""" - slices = tuple( - slice(step, length + step) for length, step in zip(working.shape, stride) - ) - out = filtered[slices] - if np.issubdtype(original_dtype, np.floating): - return out.astype(original_dtype, copy=False) + """Crop the padding away and restore a floating input dtype.""" + inner = tuple(slice(step, length + step) for length, step in zip(shape, stride)) + out = filtered[inner] + if np.issubdtype(dtype, np.floating): + return out.astype(dtype, copy=False) return out -def _extract_tiles_python( +def _extract_tiles( padded: np.ndarray, window_size: tuple[int, ...], stride: tuple[int, ...], n_tiles: tuple[int, ...], ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Extract padded windows into a dense tile stack for batched SciPy FFTs.""" + """Copy every window into a dense tile stack for batched FFTs.""" ndim = len(window_size) tiles = np.zeros((prod(n_tiles), *window_size), dtype=np.float32) begins = np.zeros((*n_tiles, ndim), dtype=np.int64) @@ -173,14 +174,14 @@ def _extract_tiles_python( return tiles, begins, sizes -def _overlap_add_tiles_python( +def _overlap_add_tiles( out: np.ndarray, tiles: np.ndarray, taper: np.ndarray, begins: np.ndarray, sizes: np.ndarray, ) -> None: - """Apply tapered overlap-add reconstruction from a dense tile stack.""" + """Add every tapered tile back into the padded output.""" grid_shape = begins.shape[:-1] for tile_index, tile_inds in enumerate( product(*(range(num) for num in grid_shape)) @@ -199,7 +200,7 @@ def _adaptive_spectral_filter_scipy( *, window_size: tuple[int, ...], overlap: tuple[int, ...], - exponent: float = 0.3, + exponent: float = 0.8, normalize_power: bool = False, ) -> np.ndarray: """ @@ -240,10 +241,10 @@ def _adaptive_spectral_filter_scipy( _validate_filter_inputs( data, window_size=window_size, overlap=overlap, exponent=float(exponent) ) - working, original_dtype, stride, taper, padded, filtered, n_tiles = ( - _prepare_work_arrays(data, window_size=window_size, overlap=overlap) + padded, taper, stride, n_tiles = _prepare_work_arrays( + data, window_size=window_size, overlap=overlap ) - tiles, begins, sizes = _extract_tiles_python(padded, window_size, stride, n_tiles) + tiles, begins, sizes = _extract_tiles(padded, window_size, stride, n_tiles) axes = tuple(range(-data.ndim, 0)) spec = sp_fft.rfftn(tiles, s=window_size, axes=axes, workers=-1) @@ -258,19 +259,9 @@ def _adaptive_spectral_filter_scipy( tiles = sp_fft.irfftn(spec, s=window_size, axes=axes, workers=-1).astype( np.float32, copy=False ) - _overlap_add_tiles_python(filtered, tiles, taper, begins, sizes) - return _finalize_output(filtered, working, original_dtype, stride) - - -def _get_dim_axis_values(patch: PatchType, kwargs: Mapping[str, Any]): - """Resolve DASCore dimension keyword arguments into dim/axis values.""" - if len(kwargs) not in {1, 2}: - msg = ( - "adaptive_spectral_filter requires one or two dimension window kwargs, " - "e.g. patch.adaptive_spectral_filter(time=32, samples=True)." - ) - raise ParameterError(msg) - return get_dim_axis_value(patch, kwargs=dict(kwargs), allow_multiple=True) + filtered = np.zeros_like(padded) + _overlap_add_tiles(filtered, tiles, taper, begins, sizes) + return _finalize_output(filtered, data.shape, data.dtype, stride) def _dim_values_to_samples( @@ -304,7 +295,8 @@ def _normalize_overlap( windows: tuple[int, ...], ) -> tuple[dict[str, Any], frozenset[str]]: """Return per-dimension overlap values and internally defaulted dimensions.""" - defaults = {dim: max(window // 2 - 2, 0) for dim, window in zip(dims, windows)} + # The largest overlap the window allows, which is what Lightguide uses. + defaults = {dim: window // 2 - 1 for dim, window in zip(dims, windows)} if overlap is None: return defaults, frozenset(dims) if isinstance(overlap, Mapping): @@ -321,31 +313,6 @@ def _normalize_overlap( return dict.fromkeys(dims, overlap), frozenset() -def _validate_window_and_overlap( - dims: tuple[str, ...], - windows: tuple[int, ...], - overlaps: tuple[int, ...], - exponent: float, -) -> None: - """Validate public DASCore window and overlap settings.""" - if not np.isfinite(exponent): - msg = "exponent must be finite." - raise ParameterError(msg) - for dim, window, overlap in zip(dims, windows, overlaps): - if window <= 4 or not is_power_of_two(window): - msg = f"window size for {dim!r} must be a power of two and > 4." - raise ParameterError(msg) - if overlap < 0: - msg = f"overlap for {dim!r} must be non-negative." - raise ParameterError(msg) - if overlap >= window / 2: - msg = ( - f"overlap for {dim!r} is too large. Maximum overlap is " - f"{window // 2 - 1} samples." - ) - raise ParameterError(msg) - - def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable: """Return the requested adaptive spectral array filter implementation.""" if engine == "scipy" or (engine == "auto" and selected_ndim == 1): @@ -356,30 +323,22 @@ def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable if selected_ndim != 2: msg = "engine='numba' currently supports exactly two selected dimensions." raise ParameterError(msg) - try: - # Deferred: the numba engine is optional, and importing it eagerly - # would make numba and rocket-fft required to import dascore. - from dascore.proc._adaptive_spectral_filter_numba import ( # noqa: PLC0415 - _NUMBA_ENGINE_AVAILABLE, - _adaptive_spectral_filter_numba, + # Deferred: the numba engine is optional, and importing it eagerly + # would compile it whenever dascore is imported. + from dascore.proc._adaptive_spectral_filter_numba import ( # noqa: PLC0415 + _NUMBA_ENGINE_AVAILABLE, + _adaptive_spectral_filter_numba, + ) + + if _NUMBA_ENGINE_AVAILABLE: + return _adaptive_spectral_filter_numba + if engine == "numba": + msg = ( + "engine='numba' requires optional dependencies numba and " + "rocket-fft to be installed." ) - except ImportError as exc: - if engine == "numba": - msg = ( - "engine='numba' requires optional dependencies numba and " - "rocket-fft to be installed." - ) - raise MissingOptionalDependencyError(msg) from exc - return _adaptive_spectral_filter_scipy - if not _NUMBA_ENGINE_AVAILABLE: - if engine == "numba": - msg = ( - "engine='numba' requires optional dependencies numba and " - "rocket-fft to be installed." - ) - raise MissingOptionalDependencyError(msg) - return _adaptive_spectral_filter_scipy - return _adaptive_spectral_filter_numba + raise MissingOptionalDependencyError(msg) + return _adaptive_spectral_filter_scipy @patch_function() @@ -387,7 +346,7 @@ def adaptive_spectral_filter( patch: PatchType, *, overlap: Any = None, - exponent: float = 0.3, + exponent: float = 0.8, normalize_power: bool = False, samples: bool = False, engine: _AdaptiveSpectralEngine = "auto", @@ -404,13 +363,17 @@ def adaptive_spectral_filter( Window overlap in samples when ``samples=True`` or in coordinate units otherwise. A single value applies to all selected dimensions; a mapping can specify dimensions independently. When omitted, each dimension - defaults to ``window // 2 - 2`` samples. + defaults to ``window // 2 - 1`` samples, the largest overlap allowed. exponent - Spectral magnitude exponent used as the adaptive weighting power. ``0`` - leaves the spectrum unweighted before overlap-add reconstruction. + Spectral magnitude exponent used as the adaptive weighting power. + Larger values suppress incoherent energy harder; ``0`` leaves the + spectrum unweighted, and values above 1 begin to remove weak coherent + arrivals along with the noise. normalize_power If ``True``, normalize each tile's spectral magnitudes by that tile's - maximum magnitude before applying ``exponent``. + maximum magnitude before applying ``exponent``. This keeps the + amplitude of every window near its input level, at the cost of + suppressing much less noise in windows which hold no signal. samples If ``True``, dimension kwargs and overlap values are interpreted as sample counts. If ``False``, values are converted through evenly sampled @@ -442,22 +405,36 @@ def adaptive_spectral_filter( Examples -------- >>> import dascore as dc - >>> patch = dc.get_example_patch() - >>> filtered_1d = patch.adaptive_spectral_filter(time=32, samples=True) - >>> filtered_2d = patch.adaptive_spectral_filter( - ... time=32, distance=32, samples=True + >>> patch = dc.get_example_patch("example_event_2").pass_filter(time=(1, 300)) + >>> # Suppress energy which is not coherent across both time and distance, + >>> # in windows of 16 samples along each. + >>> filtered = patch.adaptive_spectral_filter( + ... time=16, distance=16, samples=True ... ) - >>> filtered_1d.shape == filtered_2d.shape == patch.shape - True + >>> # Or weight each trace's spectrum on its own. + >>> per_trace = patch.adaptive_spectral_filter(time=32, samples=True) Notes ----- - With two selected dimensions, this method is equivalent to the adaptive - frequency-wavenumber (f-k) filter described in @isken2022denoising and - follows the behavior exposed by Pyrocko - [Lightguide](https://github.com/pyrocko/lightguide). + - With two selected dimensions this is the adaptive frequency-wavenumber + (AFK) filter of @isken2022denoising, and matches Pyrocko + [Lightguide](https://github.com/pyrocko/lightguide)'s `afk_filter`, + whose defaults are ``window_size=16, overlap=7, exponent=0.8``. + - The filter is not amplitude preserving. Each coefficient is scaled by + its own magnitude to the power of ``exponent``, so the output's units + are not the input's and its amplitudes grow with the input's; compare + arrivals within one output rather than across inputs. + - Windows must be powers of two greater than 4 samples. A window should + hold a few cycles of the arrivals to keep and be short against the + distance over which their moveout changes. """ - dim_axis_values = _get_dim_axis_values(patch, kwargs) + if len(kwargs) not in {1, 2}: + msg = ( + "adaptive_spectral_filter requires one or two dimension window kwargs, " + "e.g. patch.adaptive_spectral_filter(time=32, samples=True)." + ) + raise ParameterError(msg) + dim_axis_values = get_dim_axis_value(patch, kwargs=kwargs, allow_multiple=True) dims = tuple(x.dim for x in dim_axis_values) axes = tuple(x.axis for x in dim_axis_values) windows = _dim_values_to_samples( diff --git a/docs/recipes/adaptive_spectral_filter.qmd b/docs/recipes/adaptive_spectral_filter.qmd new file mode 100644 index 000000000..7a58f88d1 --- /dev/null +++ b/docs/recipes/adaptive_spectral_filter.qmd @@ -0,0 +1,138 @@ +--- +title: Adaptive Spectral (AFK) Filtering +execute: + warning: False +--- + +[`Patch.adaptive_spectral_filter`](`dascore.Patch.adaptive_spectral_filter`) suppresses energy which is not coherent within a small window of the data. Over time and distance together it is the adaptive frequency-wavenumber (AFK) filter of @isken2022denoising, as implemented in Pyrocko's [Lightguide](https://github.com/pyrocko/lightguide), and it is a good first tool for pulling arrivals out of DAS data whose noise is not organized along any particular moveout. + +The filter walks the patch in overlapping windows. Each window is Fourier transformed, every coefficient is multiplied by its own magnitude raised to `exponent`, and the window is transformed back and blended into the output. An arrival which is coherent across the window concentrates its energy in a few large coefficients, which the weighting keeps; noise spread across the spectrum is suppressed everywhere. + +## A synthetic event + +Two linear-moveout wavelets and one hyperbola, buried in white noise at half their amplitude: + +```{python} +import numpy as np +import matplotlib.pyplot as plt + +import dascore as dc + +rng = np.random.default_rng(0) +time = np.arange(1024) * 0.002 +distance = np.arange(256) * 2.0 + + +def ricker(t, frequency=25.0): + width = (np.pi * frequency * t) ** 2 + return (1 - 2 * width) * np.exp(-width) + + +clean = np.zeros((256, 1024), dtype=np.float32) +for start, velocity in [(0.4, 1500.0), (0.9, -3000.0), (1.3, 800.0)]: + clean += ricker(time[None, :] - (start + distance[:, None] / velocity)) +clean += 0.8 * ricker( + time[None, :] - np.sqrt(1.6**2 + (distance[:, None] - 250) ** 2 / 2000**2) +) +coords = {"distance": distance, "time": dc.to_datetime64(time)} +clean_patch = dc.Patch(data=clean, coords=coords, dims=("distance", "time")) +noisy_patch = clean_patch.new( + data=clean + rng.normal(0, 0.5, clean.shape).astype(np.float32) +) + + +def show(axes, patches, titles): + """Draw each patch on its own colour scale; the filter does not keep amplitude.""" + for ax, patch, title in zip(axes, patches, titles): + scale = np.percentile(np.abs(patch.data), 99) + patch.viz.waterfall(ax=ax, scale=scale, scale_type="absolute", cmap="bwr") + ax.set_title(title) + axes[0].figure.tight_layout() + + +filtered = noisy_patch.adaptive_spectral_filter(time=16, distance=16, samples=True) + +fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) +show(axes, [clean_patch, noisy_patch, filtered], ["clean", "noisy", "filtered"]) +``` + +Only the window sizes were given; `overlap` defaults to the largest the window allows and `exponent` to 0.8, the settings Lightguide uses. + +## A real event + +```{python} +patch = dc.get_example_patch("example_event_2").pass_filter(time=(1, 300)) +filtered = patch.adaptive_spectral_filter(time=16, distance=16, samples=True) + +fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) +show(axes, [patch, filtered], ["band-passed", "filtered"]) +``` + +## Choosing the parameters + +`exponent` is the filter's strength. At 0 the data pass through unchanged; near 1 the weighting is the magnitude itself, and above 1 weak but coherent arrivals start to go with the noise. The default of 0.8 is a good starting point. + +```{python} +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +exponents = [0.2, 0.5, 0.8, 1.2] +show( + axes, + [ + patch.adaptive_spectral_filter( + time=16, distance=16, samples=True, exponent=exponent + ) + for exponent in exponents + ], + [f"exponent={exponent}" for exponent in exponents], +) +``` + +The windows must be powers of two greater than 4 samples. A window should hold a few cycles of the arrivals to keep and stay short against the distance over which their moveout changes; 16 samples along each dimension serves most data, and 32 sharpens straight, long arrivals. Windows can also be given in coordinate units: + +```{python} +from dascore.units import s, m + +filtered = patch.adaptive_spectral_filter(time=1.6e-3 * s, distance=16 * m) +``` + +Selecting one dimension weights each trace's spectrum on its own. That still favours the strong arrivals, but a single trace carries no information about coherence across the fiber, so noise which shares an arrival's frequencies stays: + +```{python} +fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) +show( + axes, + [ + patch.adaptive_spectral_filter(time=16, samples=True), + patch.adaptive_spectral_filter(time=16, distance=16, samples=True), + ], + ["time only", "time and distance"], +) +``` + +## What it does not do + +The filter is not amplitude preserving: every coefficient is scaled by a power of its own magnitude, so the output's units are not the input's and its amplitudes grow with the input's. Compare arrivals within one filtered patch, not across patches, and do not read the output as strain. + +Nor does it remove coherent noise. Energy which is organized within a window is kept whatever its origin, so per-channel offsets, instrument ringing, or a strong surface wave must be removed first. The first example event carries low-frequency per-channel striping, which the filter keeps as faithfully as it keeps the arrivals; band-passing it away beforehand is what leaves the event alone: + +```{python} +event = dc.get_example_patch("example_event_1").detrend("time").taper(time=0.05) +striped = event.pass_filter(time=(1, 300)) +unstriped = event.pass_filter(time=(50, 1000)) + +fig, axes = plt.subplots(2, 2, figsize=(12, 10), sharex=True, sharey=True) +show( + axes.ravel(), + [ + striped, + striped.adaptive_spectral_filter(time=16, distance=16, samples=True), + unstriped, + unstriped.adaptive_spectral_filter(time=16, distance=16, samples=True), + ], + ["1-300 Hz", "filtered", "50-1000 Hz", "filtered"], +) +``` + +## Performance + +Two-dimensional filtering runs on an optional compiled engine when `numba` and `rocket-fft` are installed, which is about ten times faster than the SciPy path on large patches; `engine="scipy"` forces the reference implementation. diff --git a/scripts/_templates/_quarto.yml b/scripts/_templates/_quarto.yml index 7f599d918..891081fcf 100644 --- a/scripts/_templates/_quarto.yml +++ b/scripts/_templates/_quarto.yml @@ -168,6 +168,7 @@ website: - recipes/correlate.qmd - recipes/edge_effects.qmd - recipes/fk.qmd + - recipes/adaptive_spectral_filter.qmd - recipes/real_time_proc.qmd - recipes/parallelization.qmd - recipes/low_freq_proc.qmd diff --git a/tests/test_proc/test_adaptive_spectral_filter.py b/tests/test_proc/test_adaptive_spectral_filter.py index d4d5dcb56..07bafe26e 100644 --- a/tests/test_proc/test_adaptive_spectral_filter.py +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -2,7 +2,6 @@ from __future__ import annotations -import builtins from typing import Any import numpy as np @@ -244,7 +243,7 @@ def test_default_overlap_stays_in_samples_when_windows_use_units(self) -> None: by_samples = patch.adaptive_spectral_filter( distance=16, time=16, - overlap={"distance": 6, "time": 6}, + overlap={"distance": 7, "time": 7}, samples=True, engine="scipy", ) @@ -271,7 +270,7 @@ def test_partial_overlap_defaults_stay_in_samples_with_units(self) -> None: by_samples = patch.adaptive_spectral_filter( distance=16, time=16, - overlap={"distance": 6, "time": 6}, + overlap={"distance": 7, "time": 6}, samples=True, engine="scipy", ) @@ -294,7 +293,7 @@ def test_1d_default_overlap_stays_in_samples_with_units(self) -> None: ) by_samples = patch.adaptive_spectral_filter( time=16, - overlap=6, + overlap=7, samples=True, engine="scipy", ) @@ -324,7 +323,7 @@ def test_batches_over_non_selected_dimensions( out = patch.adaptive_spectral_filter( **kwargs, - overlap={dim: max(value // 2 - 2, 0) for dim, value in kwargs.items()}, + overlap={dim: value // 2 - 1 for dim, value in kwargs.items()}, samples=True, engine="scipy", ) @@ -653,19 +652,6 @@ def test_direct_array_api_supports_zero_overlap(self) -> None: assert out.shape == data.shape assert np.isfinite(out).all() - def test_auto_engine_falls_back_when_numba_missing(self, monkeypatch) -> None: - """Auto engine should fall back to SciPy when optional deps are absent.""" - real_import = builtins.__import__ - - def fake_import(name, *args, **kwargs): - if name == "dascore.proc._adaptive_spectral_filter_numba": - raise ImportError("simulated missing numba engine") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert _get_engine("auto", 2) is _adaptive_spectral_filter_scipy - def test_auto_engine_falls_back_when_deps_are_absent(self, monkeypatch) -> None: """The engine module imports without numba; the flag says it cannot run.""" module = pytest.importorskip("dascore.proc._adaptive_spectral_filter_numba") @@ -681,20 +667,6 @@ def test_numba_engine_raises_when_deps_are_absent(self, monkeypatch) -> None: with pytest.raises(MissingOptionalDependencyError, match="engine='numba'"): _get_engine("numba", 2) - def test_numba_engine_raises_when_missing(self, monkeypatch) -> None: - """Explicit numba engine should raise when optional deps are absent.""" - real_import = builtins.__import__ - - def fake_import(name, *args, **kwargs): - if name == "dascore.proc._adaptive_spectral_filter_numba": - raise ImportError("simulated missing numba engine") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - with pytest.raises(MissingOptionalDependencyError, match="engine='numba'"): - _get_engine("numba", 2) - def test_get_engine_uses_scipy_for_one_dimensional_auto(self) -> None: """Auto mode should use SciPy when one dimension is selected.""" assert _get_engine("auto", 1) is _adaptive_spectral_filter_scipy @@ -739,104 +711,121 @@ def test_auto_engine_uses_numba_when_available(self) -> None: assert _get_engine("auto", 2) is numba_mod._adaptive_spectral_filter_numba - def test_numba_private_helpers_run_in_python(self) -> None: - """The fast-engine helpers should be directly testable in Python.""" + def test_kernel_runs_in_python(self) -> None: + """The tile kernel gives SciPy's answer when run uncompiled.""" numba_mod = _numba_engine() - padded = np.arange(16, dtype=np.float32).reshape(4, 4) - tile = np.zeros((2, 2), dtype=np.float32) - - assert numba_mod._tile_indices_from_parity_index(3, 2, 1, 0) == (3, 2) - assert numba_mod._tile_bounds(1, 1, 2, 2, 1, 1, 4, 4) == (1, 1, 2, 2) - numba_mod._copy_padded_tile(padded, tile, 1, 1, 2, 2) - np.testing.assert_array_equal(tile, padded[1:3, 1:3]) - assert numba_mod._complex_power(3 + 4j) == np.float32(5.0) - - spec = np.array([[3 + 4j, 0j]], dtype=np.complex64) - assert numba_mod._max_spectral_power(spec) == np.float32(5.0) - assert numba_mod._max_spectral_power_numba_impl(spec) == np.float32(5.0) - weighted = spec.copy() - numba_mod._apply_spectral_weight(weighted, 1.0, False) - np.testing.assert_allclose(weighted[0, 0], spec[0, 0] * 5.0) - - weighted = spec.copy() - numba_mod._apply_spectral_weight(weighted, 0.3, True) - assert np.isfinite(weighted).all() - - weighted = spec.copy() - numba_mod._apply_spectral_weight_numba_impl(weighted, 0.3, True) - assert np.isfinite(weighted).all() - - weighted = spec.copy() - numba_mod._apply_spectral_weight_numba_impl(weighted, 1.0, False) - np.testing.assert_allclose(weighted[0, 0], spec[0, 0] * 5.0) - - zeros = np.array([[0j]], dtype=np.complex64) - numba_mod._apply_spectral_weight(zeros, 0.3, True) - assert zeros[0, 0] == 0j - - zeros = np.array([[0j]], dtype=np.complex64) - numba_mod._apply_spectral_weight_numba_impl(zeros, 0.3, True) - assert zeros[0, 0] == 0j - + rng = np.random.default_rng(20260511) + data = rng.normal(size=(24, 40)).astype(np.float32) + kwargs = dict(window_size=(8, 16), overlap=(3, 7)) + padded, taper, stride, n_tiles = numba_mod._prepare_work_arrays(data, **kwargs) filtered = np.zeros_like(padded) - taper = np.ones((2, 2), dtype=np.float32) - numba_mod._overlap_add_tile(filtered, tile, taper, 1, 1, 2, 2) - np.testing.assert_array_equal(filtered[1:3, 1:3], tile) + for parity0, parity1 in [(0, 0), (0, 1), (1, 0), (1, 1)]: + numba_mod._filter_tile_group.func( + padded, + filtered, + taper, + 8, + 16, + *stride, + *n_tiles, + parity0, + parity1, + 0.8, + True, + ) + out = numba_mod._finalize_output(filtered, data.shape, data.dtype, stride) + expected = _adaptive_spectral_filter_scipy( + data, exponent=0.8, normalize_power=True, **kwargs + ) + + np.testing.assert_allclose(out, expected, rtol=1e-5, atol=1e-5) - def test_numba_private_tile_group_runs_in_python(self) -> None: - """The tile group algorithm should run without JIT for coverage.""" + def test_silent_tile_normalizes_to_zero(self) -> None: + """A tile with no energy has no maximum to divide by, and stays silent.""" numba_mod = _numba_engine() - data = np.ones((8, 8), dtype=np.float32) - working, _, stride, taper, padded, filtered, n_tiles = ( - numba_mod._prepare_work_arrays(data, window_size=(8, 8), overlap=(3, 3)) - ) - - numba_mod._process_tile_group_python( - padded, - filtered, - taper, - 8, - 8, - stride[0], - stride[1], - n_tiles[0], - n_tiles[1], - 0, - 0, - 0.0, - False, - ) - numba_mod._process_tile_group_python( - padded, - filtered, - taper, - 8, - 8, - stride[0], - stride[1], - n_tiles[0], - n_tiles[1], - 0, - 0, - 0.5, - True, - ) - numba_mod._process_tile_group_numba_impl( - padded, - filtered, - taper, - 8, - 8, - stride[0], - stride[1], - n_tiles[0], - n_tiles[1], - 0, - 0, - 0.5, - True, - ) - out = numba_mod._finalize_output(filtered, working, data.dtype, stride) + data = np.zeros((32, 32), dtype=np.float32) - assert out.shape == data.shape - assert np.isfinite(out).all() + out = numba_mod._adaptive_spectral_filter_numba( + data, window_size=(16, 16), overlap=(7, 7), normalize_power=True + ) + + assert not out.any() + + def test_numba_engine_is_two_dimensional(self) -> None: + """The kernel's parity trick is written for two axes.""" + numba_mod = _numba_engine() + + with pytest.raises(ValueError, match="two-dimensional"): + numba_mod._adaptive_spectral_filter_numba( + np.zeros(64, dtype=np.float32), window_size=(16,), overlap=(7,) + ) + + +class TestEfficacy: + """The filter should recover coherent arrivals from noise, not just run.""" + + @pytest.fixture(scope="class") + def clean_and_noisy(self) -> tuple[np.ndarray, np.ndarray]: + """Two linear-moveout wavelets, and the same under white noise.""" + rng = np.random.default_rng(20260827) + time = np.arange(512) * 0.002 + distance = np.arange(96) * 2.0 + clean = np.zeros((96, 512), dtype=np.float32) + for start, velocity in [(0.2, 1500.0), (0.6, -2500.0)]: + arrival = time[None, :] - (start + distance[:, None] / velocity) + width = (np.pi * 25.0 * arrival) ** 2 + clean += (1 - 2 * width) * np.exp(-width) + noisy = clean + rng.normal(0, 0.5, clean.shape).astype(np.float32) + return clean, noisy + + @staticmethod + def _correlation(a: np.ndarray, b: np.ndarray) -> float: + """Correlation with the clean signal; the filter does not keep amplitude.""" + a, b = a - a.mean(), b - b.mean() + return float((a * b).sum() / np.sqrt((a * a).sum() * (b * b).sum())) + + def test_two_dimensional_filter_recovers_arrivals(self, clean_and_noisy) -> None: + """Filtering over both dimensions brings the data much closer to clean.""" + clean, noisy = clean_and_noisy + out = _adaptive_spectral_filter_scipy( + noisy, window_size=(16, 16), overlap=(7, 7), exponent=0.8 + ) + + before = self._correlation(noisy, clean) + after = self._correlation(out, clean) + assert before < 0.4 + assert after > 0.65 + + def test_two_dimensions_beat_one(self, clean_and_noisy) -> None: + """Coherence across distance is what a single trace cannot see.""" + clean, noisy = clean_and_noisy + both = _adaptive_spectral_filter_scipy( + noisy, window_size=(16, 16), overlap=(7, 7), exponent=0.8 + ) + per_trace = np.stack( + [ + _adaptive_spectral_filter_scipy( + trace, window_size=(16,), overlap=(7,), exponent=0.8 + ) + for trace in noisy + ] + ) + + assert ( + self._correlation(both, clean) > self._correlation(per_trace, clean) + 0.2 + ) + + def test_larger_exponent_suppresses_more(self, clean_and_noisy) -> None: + """The exponent is the filter's strength.""" + clean, noisy = clean_and_noisy + scores = [ + self._correlation( + _adaptive_spectral_filter_scipy( + noisy, window_size=(16, 16), overlap=(7, 7), exponent=exponent + ), + clean, + ) + for exponent in (0.0, 0.3, 0.8) + ] + + assert scores == sorted(scores) From ba5768da19d58534aa1fb4da919e051e6931e499 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 20:28:18 +0200 Subject: [PATCH 7/8] Make the adaptive spectral filter a PatchProcessor The operation now has the kernel seam: AdaptiveSpectralFilter holds the call as the patch function takes it, `geometry` turns the windows and overlaps into sample counts once the coordinates are known, and `kernel` is the numpy implementation a `register_kernel` for another backend would replace. `fn.op(...)` returns the processor, and both routes give the same patch. Two things the Codex review found: a negative exponent turned every silent tile into NaN, and is refused; float16 input overflowed at the new default, and comes back as float32. The recipe is cut to the real event, the exponent sweep, and what the filter does not do. --- dascore/proc/adaptive_spectral_filter.py | 200 ++++++++++++------ docs/recipes/adaptive_spectral_filter.qmd | 97 +-------- .../test_adaptive_spectral_filter.py | 29 +++ tests/test_workflow/test_patch_op.py | 1 + 4 files changed, 172 insertions(+), 155 deletions(-) diff --git a/dascore/proc/adaptive_spectral_filter.py b/dascore/proc/adaptive_spectral_filter.py index ec2239665..5fd918120 100644 --- a/dascore/proc/adaptive_spectral_filter.py +++ b/dascore/proc/adaptive_spectral_filter.py @@ -27,19 +27,26 @@ from collections.abc import Callable, Mapping from itertools import product from math import prod -from typing import Any, Literal +from typing import Any, Literal, NamedTuple import numpy as np +from pydantic import ConfigDict from scipy import fft as sp_fft from dascore.constants import PatchType -from dascore.exceptions import MissingOptionalDependencyError, ParameterError +from dascore.exceptions import ( + MissingOptionalDependencyError, + ParameterError, + PatchCoordinateError, +) from dascore.utils.misc import is_power_of_two -from dascore.utils.patch import get_dim_axis_value, patch_function +from dascore.utils.patch import patch_function from dascore.utils.signal import _triangular_taper +from dascore.workflow.meta import PatchMeta +from dascore.workflow.processor import PatchProcessor, register_implementation _AdaptiveSpectralEngine = Literal["auto", "numba", "scipy"] -__all__ = ("adaptive_spectral_filter",) +__all__ = ("AdaptiveSpectralFilter", "adaptive_spectral_filter") def _check_window(window: Any, overlap: Any, label: str) -> None: @@ -64,9 +71,11 @@ def _check_window(window: Any, overlap: Any, label: str) -> None: def _check_exponent(exponent: float) -> None: - """Raise ValueError unless the exponent is finite.""" - if not np.isfinite(exponent): - msg = "exponent must be finite." + """Raise ValueError unless the exponent is finite and non-negative.""" + # Negative: a silent coefficient would be raised to a negative power, + # and zero times infinity is the NaN every sample of the tile becomes. + if not np.isfinite(exponent) or exponent < 0: + msg = f"exponent must be finite and non-negative; got {exponent!r}." raise ValueError(msg) @@ -140,10 +149,20 @@ def _finalize_output( dtype: np.dtype, stride: tuple[int, ...], ) -> np.ndarray: - """Crop the padding away and restore a floating input dtype.""" + """Crop the padding away and restore the input dtype where that is safe.""" inner = tuple(slice(step, length + step) for length, step in zip(shape, stride)) - out = filtered[inner] - if np.issubdtype(dtype, np.floating): + return _restore_dtype(filtered[inner], dtype) + + +def _restore_dtype(out: np.ndarray, dtype: np.dtype) -> np.ndarray: + """ + Return float32 output in the input's dtype, if that is a wide enough float. + + The output grows as the input to the power of one plus the exponent, + which overflows float16 at the default exponent for inputs of a few + hundred; those, and integers, come back as float32. + """ + if np.issubdtype(dtype, np.floating) and np.dtype(dtype).itemsize >= 4: return out.astype(dtype, copy=False) return out @@ -264,22 +283,38 @@ def _adaptive_spectral_filter_scipy( return _finalize_output(filtered, data.shape, data.dtype, stride) -def _dim_values_to_samples( - patch: PatchType, - dim_axis_values, +class _Geometry(NamedTuple): + """The selected axes and their windows and overlaps, all in samples.""" + + axes: tuple[int, ...] + windows: tuple[int, ...] + overlaps: tuple[int, ...] + + +def _axes(meta: PatchMeta, dims: tuple[str, ...]) -> tuple[int, ...]: + """Return the axis of each dimension named, refusing one the patch lacks.""" + for dim in dims: + if dim not in meta.dims: + msg = f"Dimension {dim!r} not found in patch dimensions {meta.dims}." + raise PatchCoordinateError(msg) + return tuple(meta.get_axis(dim) for dim in dims) + + +def _sample_counts( + meta: PatchMeta, + values: Mapping[str, Any], *, samples: bool, name: str, - force_sample_dims: frozenset[str] = frozenset(), + sample_dims: frozenset[str] = frozenset(), ) -> tuple[int, ...]: - """Convert DASCore dimension values from units or samples into sample counts.""" + """Convert per-dimension values in samples or coordinate units to sample counts.""" out: list[int] = [] - for dim, _, value in dim_axis_values: - if samples or dim in force_sample_dims: + for dim, value in values.items(): + if samples or dim in sample_dims: count = int(value) else: - coord = patch.get_coord(dim, require_evenly_sampled=True) - count = coord.get_sample_count(value, samples=False) + count = meta.coords.get_coord(dim).get_sample_count(value) invalid = count < 0 if name == "overlap" else count <= 0 if invalid: requirement = "non-negative" if name == "overlap" else "positive" @@ -313,7 +348,7 @@ def _normalize_overlap( return dict.fromkeys(dims, overlap), frozenset() -def _get_engine(engine: _AdaptiveSpectralEngine, selected_ndim: int) -> Callable: +def _get_engine(engine: str, selected_ndim: int) -> Callable: """Return the requested adaptive spectral array filter implementation.""" if engine == "scipy" or (engine == "auto" and selected_ndim == 1): return _adaptive_spectral_filter_scipy @@ -368,7 +403,7 @@ def adaptive_spectral_filter( Spectral magnitude exponent used as the adaptive weighting power. Larger values suppress incoherent energy harder; ``0`` leaves the spectrum unweighted, and values above 1 begin to remove weak coherent - arrivals along with the noise. + arrivals along with the noise. Must be non-negative. normalize_power If ``True``, normalize each tile's spectral magnitudes by that tile's maximum magnitude before applying ``exponent``. This keeps the @@ -391,13 +426,14 @@ def adaptive_spectral_filter( ------- Patch A new patch with filtered data and original dimensions and coordinates. + The data are ``float32``, or ``float64`` for ``float64`` input. Raises ------ ParameterError If one or two dimensions are not selected, if selected window or overlap - values are invalid, if ``exponent`` is not finite, or if an invalid - engine name is requested. + values are invalid, if ``exponent`` is not finite and non-negative, or + if an invalid engine name is requested. MissingOptionalDependencyError If ``engine="numba"`` is requested for two selected dimensions but the optional fast-engine dependencies are not installed. @@ -428,49 +464,79 @@ def adaptive_spectral_filter( hold a few cycles of the arrivals to keep and be short against the distance over which their moveout changes. """ - if len(kwargs) not in {1, 2}: - msg = ( - "adaptive_spectral_filter requires one or two dimension window kwargs, " - "e.g. patch.adaptive_spectral_filter(time=32, samples=True)." - ) - raise ParameterError(msg) - dim_axis_values = get_dim_axis_value(patch, kwargs=kwargs, allow_multiple=True) - dims = tuple(x.dim for x in dim_axis_values) - axes = tuple(x.axis for x in dim_axis_values) - windows = _dim_values_to_samples( - patch, dim_axis_values, samples=samples, name="window" - ) - overlap_values, default_overlap_dims = _normalize_overlap(overlap, dims, windows) - overlap_dim_axis_values = get_dim_axis_value( - patch, kwargs=overlap_values, allow_multiple=True - ) - overlaps = _dim_values_to_samples( - patch, + return AdaptiveSpectralFilter( + overlap=overlap, + exponent=exponent, + normalize_power=normalize_power, samples=samples, - dim_axis_values=overlap_dim_axis_values, - name="overlap", - force_sample_dims=default_overlap_dims, - ) - _validate_window_and_overlap(dims, windows, overlaps, float(exponent)) - - data = np.asarray(patch.data) - selected_ndim = len(axes) - moved = np.moveaxis(data, axes, tuple(range(-selected_ndim, 0))) - batch_shape = moved.shape[:-selected_ndim] - selected_shape = moved.shape[-selected_ndim:] - working = moved.reshape((-1, *selected_shape)) - filtered = np.empty_like(working, dtype=np.float32) - engine_func = _get_engine(engine, selected_ndim) - for ind, array in enumerate(working): - filtered[ind] = engine_func( - array, - window_size=windows, - overlap=overlaps, - exponent=float(exponent), - normalize_power=bool(normalize_power), + engine=engine, + **kwargs, + )._apply(patch) + + +class AdaptiveSpectralFilter(PatchProcessor): + """ + Weight every window's spectrum by a power of its own magnitude. + + The dimensions to filter arrive as extras carrying their window sizes, + as the patch function takes them. Windows and overlaps may be given in + coordinate units, so they are only sample counts once the coordinates + are known: `geometry` is that conversion, and where a kernel for any + backend starts. + """ + + model_config = ConfigDict(extra="allow", frozen=True) + + overlap: Any = None + exponent: float = 0.8 + normalize_power: bool = False + samples: bool = False + # A str rather than the Literal, so a wrong name is refused by + # `_get_engine` as a ParameterError like every other bad argument. + engine: str = "auto" + + def geometry(self, meta: PatchMeta) -> _Geometry: + """Return the selected axes and their windows and overlaps, in samples.""" + selected = self.model_extra or {} + if len(selected) not in {1, 2}: + msg = ( + "adaptive_spectral_filter requires one or two dimension window " + "kwargs, e.g. patch.adaptive_spectral_filter(time=32, samples=True)." + ) + raise ParameterError(msg) + dims = tuple(selected) + axes = _axes(meta, dims) + windows = _sample_counts(meta, selected, samples=self.samples, name="window") + overlap_values, sample_dims = _normalize_overlap(self.overlap, dims, windows) + overlaps = _sample_counts( + meta, + overlap_values, + samples=self.samples, + name="overlap", + sample_dims=sample_dims, ) - filtered = filtered.reshape((*batch_shape, *selected_shape)) - filtered = np.moveaxis(filtered, tuple(range(-selected_ndim, 0)), axes) - if np.issubdtype(data.dtype, np.floating): - filtered = filtered.astype(data.dtype, copy=False) - return patch.update(data=filtered) + _validate_window_and_overlap(dims, windows, overlaps, float(self.exponent)) + return _Geometry(axes, windows, overlaps) + + def kernel(self, data, meta, out_meta): + """Filter every batch over the selected axes and stack the results.""" + axes, windows, overlaps = self.geometry(meta) + engine = _get_engine(self.engine, len(axes)) + data = np.asarray(data) + tail = tuple(range(-len(axes), 0)) + moved = np.moveaxis(data, axes, tail) + working = moved.reshape((-1, *moved.shape[-len(axes) :])) + filtered = np.empty_like(working, dtype=np.float32) + for ind, array in enumerate(working): + filtered[ind] = engine( + array, + window_size=windows, + overlap=overlaps, + exponent=float(self.exponent), + normalize_power=bool(self.normalize_power), + ) + filtered = np.moveaxis(filtered.reshape(moved.shape), tail, axes) + return _restore_dtype(filtered, data.dtype) + + +register_implementation("adaptive_spectral_filter", AdaptiveSpectralFilter) diff --git a/docs/recipes/adaptive_spectral_filter.qmd b/docs/recipes/adaptive_spectral_filter.qmd index 7a58f88d1..0f6a308e0 100644 --- a/docs/recipes/adaptive_spectral_filter.qmd +++ b/docs/recipes/adaptive_spectral_filter.qmd @@ -6,11 +6,7 @@ execute: [`Patch.adaptive_spectral_filter`](`dascore.Patch.adaptive_spectral_filter`) suppresses energy which is not coherent within a small window of the data. Over time and distance together it is the adaptive frequency-wavenumber (AFK) filter of @isken2022denoising, as implemented in Pyrocko's [Lightguide](https://github.com/pyrocko/lightguide), and it is a good first tool for pulling arrivals out of DAS data whose noise is not organized along any particular moveout. -The filter walks the patch in overlapping windows. Each window is Fourier transformed, every coefficient is multiplied by its own magnitude raised to `exponent`, and the window is transformed back and blended into the output. An arrival which is coherent across the window concentrates its energy in a few large coefficients, which the weighting keeps; noise spread across the spectrum is suppressed everywhere. - -## A synthetic event - -Two linear-moveout wavelets and one hyperbola, buried in white noise at half their amplitude: +The filter walks the patch in overlapping windows. Each window is Fourier transformed, every coefficient is multiplied by its own magnitude raised to `exponent`, and the window is transformed back and blended into the output. An arrival which is coherent across the window concentrates its energy in a few large coefficients, which the weighting favours over everything spread thinly across the spectrum. ```{python} import numpy as np @@ -18,27 +14,8 @@ import matplotlib.pyplot as plt import dascore as dc -rng = np.random.default_rng(0) -time = np.arange(1024) * 0.002 -distance = np.arange(256) * 2.0 - - -def ricker(t, frequency=25.0): - width = (np.pi * frequency * t) ** 2 - return (1 - 2 * width) * np.exp(-width) - - -clean = np.zeros((256, 1024), dtype=np.float32) -for start, velocity in [(0.4, 1500.0), (0.9, -3000.0), (1.3, 800.0)]: - clean += ricker(time[None, :] - (start + distance[:, None] / velocity)) -clean += 0.8 * ricker( - time[None, :] - np.sqrt(1.6**2 + (distance[:, None] - 250) ** 2 / 2000**2) -) -coords = {"distance": distance, "time": dc.to_datetime64(time)} -clean_patch = dc.Patch(data=clean, coords=coords, dims=("distance", "time")) -noisy_patch = clean_patch.new( - data=clean + rng.normal(0, 0.5, clean.shape).astype(np.float32) -) +patch = dc.get_example_patch("example_event_2").pass_filter(time=(1, 300)) +filtered = patch.adaptive_spectral_filter(time=16, distance=16, samples=True) def show(axes, patches, titles): @@ -50,31 +27,19 @@ def show(axes, patches, titles): axes[0].figure.tight_layout() -filtered = noisy_patch.adaptive_spectral_filter(time=16, distance=16, samples=True) - -fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True) -show(axes, [clean_patch, noisy_patch, filtered], ["clean", "noisy", "filtered"]) -``` - -Only the window sizes were given; `overlap` defaults to the largest the window allows and `exponent` to 0.8, the settings Lightguide uses. - -## A real event - -```{python} -patch = dc.get_example_patch("example_event_2").pass_filter(time=(1, 300)) -filtered = patch.adaptive_spectral_filter(time=16, distance=16, samples=True) - fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) show(axes, [patch, filtered], ["band-passed", "filtered"]) ``` -## Choosing the parameters +Only the window sizes were given; `overlap` defaults to the largest the window allows and `exponent` to 0.8, the settings Lightguide uses. Windows must be powers of two greater than 4 samples and can also be given in coordinate units, as in `patch.adaptive_spectral_filter(time=1.6e-3 * s, distance=16 * m)`. Selecting a single dimension weights each trace's spectrum on its own, which favours the strong arrivals but cannot see coherence across the fiber. + +## Choosing the exponent -`exponent` is the filter's strength. At 0 the data pass through unchanged; near 1 the weighting is the magnitude itself, and above 1 weak but coherent arrivals start to go with the noise. The default of 0.8 is a good starting point. +`exponent` is the filter's strength. At 0 the data pass through unchanged; near 1 the weighting is the magnitude itself, and above 1 weak but coherent arrivals start to go with the noise. ```{python} -fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) exponents = [0.2, 0.5, 0.8, 1.2] +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) show( axes, [ @@ -87,52 +52,8 @@ show( ) ``` -The windows must be powers of two greater than 4 samples. A window should hold a few cycles of the arrivals to keep and stay short against the distance over which their moveout changes; 16 samples along each dimension serves most data, and 32 sharpens straight, long arrivals. Windows can also be given in coordinate units: - -```{python} -from dascore.units import s, m - -filtered = patch.adaptive_spectral_filter(time=1.6e-3 * s, distance=16 * m) -``` - -Selecting one dimension weights each trace's spectrum on its own. That still favours the strong arrivals, but a single trace carries no information about coherence across the fiber, so noise which shares an arrival's frequencies stays: - -```{python} -fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) -show( - axes, - [ - patch.adaptive_spectral_filter(time=16, samples=True), - patch.adaptive_spectral_filter(time=16, distance=16, samples=True), - ], - ["time only", "time and distance"], -) -``` - ## What it does not do The filter is not amplitude preserving: every coefficient is scaled by a power of its own magnitude, so the output's units are not the input's and its amplitudes grow with the input's. Compare arrivals within one filtered patch, not across patches, and do not read the output as strain. -Nor does it remove coherent noise. Energy which is organized within a window is kept whatever its origin, so per-channel offsets, instrument ringing, or a strong surface wave must be removed first. The first example event carries low-frequency per-channel striping, which the filter keeps as faithfully as it keeps the arrivals; band-passing it away beforehand is what leaves the event alone: - -```{python} -event = dc.get_example_patch("example_event_1").detrend("time").taper(time=0.05) -striped = event.pass_filter(time=(1, 300)) -unstriped = event.pass_filter(time=(50, 1000)) - -fig, axes = plt.subplots(2, 2, figsize=(12, 10), sharex=True, sharey=True) -show( - axes.ravel(), - [ - striped, - striped.adaptive_spectral_filter(time=16, distance=16, samples=True), - unstriped, - unstriped.adaptive_spectral_filter(time=16, distance=16, samples=True), - ], - ["1-300 Hz", "filtered", "50-1000 Hz", "filtered"], -) -``` - -## Performance - -Two-dimensional filtering runs on an optional compiled engine when `numba` and `rocket-fft` are installed, which is about ten times faster than the SciPy path on large patches; `engine="scipy"` forces the reference implementation. +Nor does it remove coherent noise. Energy which is organized within a window is kept whatever its origin, so per-channel offsets, instrument ringing, or a strong surface wave must be removed first. The first example event carries low-frequency per-channel striping which the filter keeps as faithfully as the arrivals; a band-pass above 50 Hz before filtering is what leaves the event alone. diff --git a/tests/test_proc/test_adaptive_spectral_filter.py b/tests/test_proc/test_adaptive_spectral_filter.py index 07bafe26e..459c96dcf 100644 --- a/tests/test_proc/test_adaptive_spectral_filter.py +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -16,6 +16,7 @@ ) from dascore.proc import adaptive_spectral_filter as adaptive_spectral_filter_func from dascore.proc.adaptive_spectral_filter import ( + AdaptiveSpectralFilter, _adaptive_spectral_filter_scipy, _get_engine, _validate_window_and_overlap, @@ -493,6 +494,34 @@ def test_invalid_engine_raises(self) -> None: engine="bad", # type: ignore[arg-type] ) + def test_negative_exponent_is_refused(self) -> None: + """A negative power of a silent coefficient is zero times infinity.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + + with pytest.raises(ParameterError, match="non-negative"): + patch.adaptive_spectral_filter( + time=16, distance=16, samples=True, exponent=-0.5 + ) + + def test_float16_comes_back_as_float32(self) -> None: + """Output grows as input to the 1.8, which float16 cannot hold.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + patch = patch.update(data=(patch.data * 100).astype(np.float16)) + + out = patch.adaptive_spectral_filter(time=16, distance=16, samples=True) + + assert out.data.dtype == np.float32 + assert np.isfinite(out.data).all() + + def test_op_is_the_processor(self) -> None: + """The seam: the call names the processor, and the two routes agree.""" + patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) + op = adaptive_spectral_filter_func.op(time=16, distance=16, samples=True) + + assert isinstance(op, AdaptiveSpectralFilter) + expected = patch.adaptive_spectral_filter(time=16, distance=16, samples=True) + assert op(patch).equals(expected) + def test_proc_export_is_function(self) -> None: """The processing module should expose the direct patch function.""" patch = _patch((64, 64), ("distance", "time"), dtype=np.float32) diff --git a/tests/test_workflow/test_patch_op.py b/tests/test_workflow/test_patch_op.py index cd7bd93e5..968a76190 100644 --- a/tests/test_workflow/test_patch_op.py +++ b/tests/test_workflow/test_patch_op.py @@ -608,6 +608,7 @@ def test_the_registry_gains_one_tag(self): "PatchOp", "PatchProcessor", "Abs", + "AdaptiveSpectralFilter", "Conj", "Demean", "Imag", From 823171965a27b80ba4300b6783c852af85f007cf Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 27 Aug 2026 20:31:56 +0200 Subject: [PATCH 8/8] Derive the processor table instead of spelling it out A test which lists every registered processor by name has to be edited to add one, and so notices nothing; what it meant is that a processor class exists only as the implementation of a patch function, with a seam. Both halves can be asked of the registry. --- tests/test_workflow/test_patch_op.py | 36 +++++++++++----------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/tests/test_workflow/test_patch_op.py b/tests/test_workflow/test_patch_op.py index 968a76190..354817647 100644 --- a/tests/test_workflow/test_patch_op.py +++ b/tests/test_workflow/test_patch_op.py @@ -594,30 +594,22 @@ def test_the_registry_gains_one_tag(self): One class for every patch function, not one class each. `PatchOp` is the whole cost of naming any of them in a document. + The module argues against a class per patch function, so a + processor class is allowed only where it is the registered + implementation of one -- and then only because it gives the + operation a seam a whole function does not have. """ - tags = { - tag - for tag, cls in registered_models().items() - if isinstance(cls, type) and issubclass(cls, (PatchOp, PatchProcessor)) - } - # Spelled out rather than counted, so that a class added without - # a reason to is noticed. The module argues against a class per - # patch function; these are the exceptions it names -- the ones - # wanting a kernel seam. - assert tags == { - "PatchOp", - "PatchProcessor", - "Abs", - "AdaptiveSpectralFilter", - "Conj", - "Demean", - "Imag", - "Normalize", - "Real", - "RenameCoords", - "Standardize", - "Transpose", + processors = { + cls + for cls in registered_models().values() + if isinstance(cls, type) + and issubclass(cls, PatchProcessor) + and cls is not PatchProcessor } + assert processors == set(processor_module._IMPLEMENTATIONS.values()) + seams = ("derive_meta", "plan_kernel", "kernel", "reconcile") + for cls in processors: + assert any(seam in cls.__dict__ for seam in seams), cls.__name__ def test_the_document_names_the_operation(self): """The name and the arguments are what a document holds."""