diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 42fcb479f..0be500d9f 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -463,6 +463,7 @@ def get_patch_name(self, *args, **kwargs) -> str: 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 aa44c7985..a8d279b0d 100644 --- a/dascore/proc/__init__.py +++ b/dascore/proc/__init__.py @@ -18,4 +18,5 @@ from .hampel import hampel_filter from .wiener import wiener_filter from .align import align_to_coord +from .adaptive_spectral_filter import adaptive_spectral_filter from .inventory import enrich diff --git a/dascore/proc/_adaptive_spectral_filter_numba.py b/dascore/proc/_adaptive_spectral_filter_numba.py new file mode 100644 index 000000000..b828d0fe9 --- /dev/null +++ b/dascore/proc/_adaptive_spectral_filter_numba.py @@ -0,0 +1,117 @@ +""" +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 + +import numpy as np + +from dascore.proc.adaptive_spectral_filter import ( + _finalize_output, + _prepare_work_arrays, + _validate_filter_inputs, +) +from dascore.utils.jit import maybe_numba_jit + + +# 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, + window0: int, + window1: int, + stride0: int, + stride1: int, + n_tiles0: int, + n_tiles1: int, + parity0: int, + parity1: int, + exponent: float, + normalize_power: bool, +) -> None: + """ + 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: + 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] + + +_NUMBA_ENGINE_AVAILABLE = _filter_tile_group.jit_available + + +def _adaptive_spectral_filter_numba( + data: np.ndarray, + *, + window_size: tuple[int, int], + overlap: tuple[int, int], + exponent: float = 0.8, + normalize_power: bool = False, +) -> np.ndarray: + """ + Filter a 2D array with the optional Numba/rocket-fft implementation. + + 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) + ) + 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): + _filter_tile_group( + padded, + filtered, + taper, + *window_size, + *stride, + *n_tiles, + parity0, + parity1, + float(exponent), + bool(normalize_power), + ) + 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 new file mode 100644 index 000000000..5fd918120 --- /dev/null +++ b/dascore/proc/adaptive_spectral_filter.py @@ -0,0 +1,542 @@ +""" +Adaptive spectral filtering for DASCore patches. + +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 + +from collections.abc import Callable, Mapping +from itertools import product +from math import prod +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, + PatchCoordinateError, +) +from dascore.utils.misc import is_power_of_two +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__ = ("AdaptiveSpectralFilter", "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 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) + + +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) + _check_exponent(exponent) + for axis, (window, axis_overlap) in enumerate(zip(window_size, overlap)): + _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( + data: np.ndarray, + *, + window_size: tuple[int, ...], + overlap: tuple[int, ...], +) -> 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 = tuple( + slice(step, length + step) for length, step in zip(working.shape, stride) + ) + padded[inner] = working + n_tiles = tuple(pad_len // step for pad_len, step in zip(padded.shape, stride)) + return padded, taper, stride, n_tiles + + +def _finalize_output( + filtered: np.ndarray, + shape: tuple[int, ...], + dtype: np.dtype, + stride: tuple[int, ...], +) -> np.ndarray: + """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)) + 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 + + +def _extract_tiles( + padded: np.ndarray, + window_size: tuple[int, ...], + stride: tuple[int, ...], + n_tiles: tuple[int, ...], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """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) + 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( + out: np.ndarray, + tiles: np.ndarray, + taper: np.ndarray, + begins: np.ndarray, + sizes: np.ndarray, +) -> None: + """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)) + ): + 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.8, + 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) + ) + padded, taper, stride, n_tiles = _prepare_work_arrays( + data, window_size=window_size, overlap=overlap + ) + 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) + 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 + ) + filtered = np.zeros_like(padded) + _overlap_add_tiles(filtered, tiles, taper, begins, sizes) + return _finalize_output(filtered, data.shape, data.dtype, stride) + + +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, + sample_dims: frozenset[str] = frozenset(), +) -> tuple[int, ...]: + """Convert per-dimension values in samples or coordinate units to sample counts.""" + out: list[int] = [] + for dim, value in values.items(): + if samples or dim in sample_dims: + count = int(value) + else: + 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" + msg = f"{name} for dimension {dim!r} must be {requirement}." + raise ParameterError(msg) + out.append(count) + return tuple(out) + + +def _normalize_overlap( + overlap: Any, + dims: tuple[str, ...], + windows: tuple[int, ...], +) -> tuple[dict[str, Any], frozenset[str]]: + """Return per-dimension overlap values and internally defaulted dimensions.""" + # 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): + extra = set(overlap) - set(dims) + if extra: + names = sorted(map(str, extra)) + msg = f"overlap contains dimensions not being filtered: {names}" + raise ParameterError(msg) + return defaults | dict(overlap), frozenset(set(dims) - set(overlap)) + # Uncoerced: the value the caller gave is read in coordinate units + # unless samples says otherwise, and int() of a timedelta64 or a + # fractional second is not the overlap they asked for. Only the + # defaults above are sample counts, which is what the returned set says. + return dict.fromkeys(dims, overlap), frozenset() + + +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 + 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) + # 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." + ) + raise MissingOptionalDependencyError(msg) + return _adaptive_spectral_filter_scipy + + +@patch_function() +def adaptive_spectral_filter( + patch: PatchType, + *, + overlap: Any = None, + exponent: float = 0.8, + 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 - 1`` samples, the largest overlap allowed. + exponent + 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. 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 + 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 + 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. + 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 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. + + Examples + -------- + >>> import dascore as dc + >>> 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 + ... ) + >>> # 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 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. + """ + return AdaptiveSpectralFilter( + overlap=overlap, + exponent=exponent, + normalize_power=normalize_power, + samples=samples, + 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, + ) + _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/dascore/utils/jit.py b/dascore/utils/jit.py index 2f91a65b5..211aea779 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 @@ -20,17 +21,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. @@ -67,7 +79,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 # noqa: PLC0415 @@ -75,7 +92,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 @@ -83,21 +110,22 @@ 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"{_get_install_message('numba')}" + f"{module_names} but it is not installed. " + f"{_get_install_message(missing_modules)}" ) 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) @@ -108,9 +136,13 @@ def decorated(*args, **kwargs): # only when the import failed. jit = numba.jit # ty: ignore[unresolved-attribute] out_func = jit(**compiler_kwargs)(func) - # Make the original func accessible via .func; function objects accept - # new attributes even though their declared type does not. + # Make the original func accessible via .func, and say which deps the + # jit wanted; function objects accept new attributes even though their + # declared type does not. + missing = tuple(missing_modules) out_func.func = func # ty: ignore[invalid-assignment] + out_func.jit_available = has_all_deps # ty: ignore[invalid-assignment] + out_func.missing_jit_deps = missing # ty: ignore[invalid-assignment] return out_func return _wrapper diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 0e0e74719..495148e87 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -1402,6 +1402,23 @@ def tukey_fence(data, fence_multiplier=1.5) -> np.ndarray: 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 + ---------- + value + The value to test. + """ + return value > 0 and (value & (value - 1) == 0) + + def is_strictly_monotonic(values, increasing: bool | None = None) -> bool: """ Return True if a 1D sequence is strictly monotonic. diff --git a/dascore/utils/signal.py b/dascore/utils/signal.py index 7b2c05803..9c72bbee2 100644 --- a/dascore/utils/signal.py +++ b/dascore/utils/signal.py @@ -2,6 +2,12 @@ Utilities for signal processing. """ +from __future__ import annotations + +from functools import lru_cache + +import numpy as np + from dascore.exceptions import ParameterError from dascore.utils.imports import lazy_import @@ -33,3 +39,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/recipes/adaptive_spectral_filter.qmd b/docs/recipes/adaptive_spectral_filter.qmd new file mode 100644 index 000000000..0f6a308e0 --- /dev/null +++ b/docs/recipes/adaptive_spectral_filter.qmd @@ -0,0 +1,59 @@ +--- +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 favours over everything spread thinly across the spectrum. + +```{python} +import numpy as np +import matplotlib.pyplot as plt + +import dascore as dc + +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): + """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() + + +fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True) +show(axes, [patch, filtered], ["band-passed", "filtered"]) +``` + +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. + +```{python} +exponents = [0.2, 0.5, 0.8, 1.2] +fig, axes = plt.subplots(1, 4, figsize=(20, 5), sharey=True) +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], +) +``` + +## 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 the arrivals; a band-pass above 50 Hz before filtering is what leaves the event alone. diff --git a/docs/references.bib b/docs/references.bib index 45cb0765c..4bd509a43 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -68,6 +68,17 @@ @article{schimmel1997noise 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}, + volume={231}, + number={2}, + pages={944--949}, + year={2022}, + doi={10.1093/gji/ggac229} +} + @article{langet2014, TITLE = {{Continuous Kurtosis-Based Migration for Seismic Event Detection and Location, with Application to Piton de la Fournaise Volcano, La R{\'e}union}}, AUTHOR = {Langet, Nad{\`e}ge and Maggi, Alessia and Michelini, Alberto and Brenguier, Florent}, diff --git a/pyproject.toml b/pyproject.toml index 684275c5a..7b6f16035 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,6 +77,7 @@ extras = [ "findiff", "obspy", "numba", + "rocket-fft", "segyio", "bottleneck", "pymseed[numpy]", 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 new file mode 100644 index 000000000..459c96dcf --- /dev/null +++ b/tests/test_proc/test_adaptive_spectral_filter.py @@ -0,0 +1,860 @@ +"""Tests for adaptive spectral filtering.""" + +from __future__ import annotations + +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 ( + AdaptiveSpectralFilter, + _adaptive_spectral_filter_scipy, + _get_engine, + _validate_window_and_overlap, +) +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, ...], + *, + 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_fractional_scalar_overlap_reads_coordinate_units(self) -> None: + """One overlap for every dimension is a unit value like a mapping is.""" + patch = _patch((64,), ("distance",), distance_step=0.5) + + scalar = patch.adaptive_spectral_filter( + distance=8.0, overlap=3.5, samples=False, engine="scipy" + ) + by_samples = patch.adaptive_spectral_filter( + distance=16, overlap=7, samples=True, engine="scipy" + ) + + np.testing.assert_allclose(scalar.data, by_samples.data, rtol=1e-5, atol=1e-5) + + def test_scalar_overlap_may_be_time_like(self) -> None: + """A timedelta overlap survives to the coordinate, rather than int().""" + patch = _patch((64,), ("time",), time_step=np.timedelta64(4, "ms")) + + by_units = patch.adaptive_spectral_filter( + time=np.timedelta64(64, "ms"), + overlap=np.timedelta64(28, "ms"), + samples=False, + engine="scipy", + ) + by_samples = patch.adaptive_spectral_filter( + time=16, overlap=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": 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_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": 7, "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=7, + 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: value // 2 - 1 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_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) + + 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_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_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 = _numba_engine() + 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 = _numba_engine() + + assert _get_engine("auto", 2) is numba_mod._adaptive_spectral_filter_numba + + def test_kernel_runs_in_python(self) -> None: + """The tile kernel gives SciPy's answer when run uncompiled.""" + numba_mod = _numba_engine() + 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) + 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_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.zeros((32, 32), dtype=np.float32) + + 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) diff --git a/tests/test_utils/test_jit.py b/tests/test_utils/test_jit.py index 10f0e7adc..5f86ec25a 100644 --- a/tests/test_utils/test_jit.py +++ b/tests/test_utils/test_jit.py @@ -48,6 +48,54 @@ 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 is not one of the missing ones.""" + + @maybe_numba_jit(deps="json") + def _jit_test_func(ar): + return ar + + 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 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 "rocket_fft" in _jit_test_func.missing_jit_deps + 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 "dascore_not_a_real_module" in _jit_test_func.missing_jit_deps + def test_example(self): """Test docstring examples.""" pytest.importorskip("numba") diff --git a/tests/test_workflow/_calls.py b/tests/test_workflow/_calls.py index 5764b2bd5..77ec4a187 100644 --- a/tests/test_workflow/_calls.py +++ b/tests/test_workflow/_calls.py @@ -205,6 +205,7 @@ def _inventory(): ("slope_filter", "default", (Lazy(_slope_filter),), {}), ("wiener_filter", "default", (), {"time": 5, "samples": True}), ("hampel_filter", "default", (), {"time": 5, "samples": True}), + ("adaptive_spectral_filter", "default", (), {"time": 32, "samples": True}), ("select", "default", (), {"distance": (10, 40)}), ("unselect", "default", (), {"distance": (10, 40)}), ("order", "default", (), {"distance": (30, 10, 20), "samples": True}), diff --git a/tests/test_workflow/test_patch_op.py b/tests/test_workflow/test_patch_op.py index cd7bd93e5..354817647 100644 --- a/tests/test_workflow/test_patch_op.py +++ b/tests/test_workflow/test_patch_op.py @@ -594,29 +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", - "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."""