From 92727cd4f52734fe88952d8b74d4dbe47188f5fd Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 16 Oct 2025 10:07:30 +0100 Subject: [PATCH 01/15] start mute implementation --- dascore/core/patch.py | 1 + dascore/examples.py | 18 +- dascore/proc/__init__.py | 1 + dascore/proc/mute.py | 721 +++++++++++++++++++++++++++++++++++ tests/test_examples.py | 6 + tests/test_proc/test_mute.py | 436 +++++++++++++++++++++ 6 files changed, 1174 insertions(+), 9 deletions(-) create mode 100644 dascore/proc/mute.py create mode 100644 tests/test_proc/test_mute.py diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 483b1c8e5..9f53b21de 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -438,6 +438,7 @@ def iresample(self, *args, **kwargs): standardize = dascore.proc.standardize taper = dascore.proc.taper taper_range = dascore.proc.taper_range + mute = dascore.proc.mute rolling = dascore.proc.rolling whiten = dascore.proc.whiten diff --git a/dascore/examples.py b/dascore/examples.py index 3fb24c7a6..790231ae9 100644 --- a/dascore/examples.py +++ b/dascore/examples.py @@ -441,7 +441,7 @@ def _ricker(time, delay): @register_func(EXAMPLE_PATCHES, key="delta_patch") def delta_patch( - dim="time", + dim: tuple[str, ...] | str = ("time", "distance"), shape=(10, 200), time_min="2020-01-01", time_step=1 / 250, @@ -457,21 +457,21 @@ def delta_patch( Parameters ---------- - dim : str + dim The dimension at the center of which to place the unit value. - Typically ``"time"`` or ``"distance"``. - shape : tuple of int + Typically, ``"time"`` or ``"distance"``. + shape The shape of the data as (distance, time). Defaults to (10, 200). This is used only if no existing ``patch`` is provided. - time_min : str or datetime64 + time_min The start time of the patch. - time_step : float + time_step The time step in seconds between samples. - distance_min : float + distance_min The minimum distance coordinate. - distance_step : float + distance_step The distance step in meters between samples. - patch : dascore.Patch + patch If provided, creates the delta patch based on this existing patch. Default is None. """ diff --git a/dascore/proc/__init__.py b/dascore/proc/__init__.py index 944d5ab04..6172e8007 100644 --- a/dascore/proc/__init__.py +++ b/dascore/proc/__init__.py @@ -12,6 +12,7 @@ from .resample import decimate, interpolate, resample from .rolling import rolling from .taper import taper, taper_range +from .mute import mute from .units import convert_units, set_units, simplify_units from .whiten import whiten from .hampel import hampel_filter diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py new file mode 100644 index 000000000..da5e31ed0 --- /dev/null +++ b/dascore/proc/mute.py @@ -0,0 +1,721 @@ +"""Processing for muting (zeroing) patch data in specified regions.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Literal + +import numpy as np + +from dascore.constants import PatchType +from dascore.exceptions import ParameterError +from dascore.units import Quantity +from dascore.utils.docs import compose_docstring +from dascore.utils.misc import broadcast_for_index +from dascore.utils.patch import patch_function +from dascore.utils.signal import WINDOW_FUNCTIONS, _get_window_function +from dascore.utils.time import to_float + + +parameter_docstring = """ +Parameters +---------- +patch + The patch instance. +mode + - "union": mute data inside the specified region [default] + - "complement": mute data outside the specified region +taper + Taper width at mute boundaries. Can be: + - None: sharp mute (no taper) + - float (0.0-1.0): fraction of dimension range (e.g., 0.05 = 5%) + The number of samples in the mute is held constant across dimensions + by calculating the samples the fraction represents for each dimension + and using the minimum sample count. This helps avoid distorted mutes + based on dimensions with vastly different lengths. + - Quantity with units: absolute value (e.g., 0.02*dc.units.s) + Note: If multiple dimensions specified, must use dict. + - dict: {dim: taper_value} for dimension-specific taper + Values can be floats (fractions) or Quantities (absolute) + - Callable: custom function that receives and modifies envelope array. +window_type + Window function for tapering (only used for non-callable taper). + Options: + {taper_type}. +relative + If True (default), values are relative to coordinate edges. + Positive values are offsets from start, negative from end. +samples + If True, values specified in samples rather than coordinate values. +**kwargs + Dimension specifications as (boundary_1, boundary_2) pairs. + Each boundary can be: + - Scalar: constant value (defines plane perpendicular to dimension) + - Array/list: coordinates defining line/curve/surface + - None or ...: edge of coordinate (min or max) + +""" + + + +@patch_function() +@compose_docstring(taper_type=sorted(WINDOW_FUNCTIONS), params=parameter_docstring) +def mute_envelope( + patch: PatchType, + *, + mode: Literal["union", "complement"] = "union", + taper: float | dict | None = None, + window_type: str = "hann", + relative: bool = True, + samples: bool = False, + **kwargs, +) -> PatchType: + """ + Calculate a patch envelope which applies mute defined by parameters. + + The resulting patch simply needs to be multiplied by the original + patch to achieve muting. This allows for more fine-grained control + of the + + """ + + +@patch_function() +@compose_docstring( + taper_type=sorted(WINDOW_FUNCTIONS), + params=parameter_docstring, +) +def mute( + patch: PatchType, + *, + mode: Literal["union", "complement"] = "union", + taper: float | dict | None = None, + window_type: str = "hann", + relative: bool = True, + samples: bool = False, + **kwargs, +) -> PatchType: + """ + Mute (zero out) patch data in specified regions. + + Each dimension accepts two boundary specifications that define + the mute region. Boundaries can be scalar values (for planes/blocks) + or arrays of values (for lines/curves/surfaces). + + {params} + + Examples + -------- + >>> import dascore as dc + >>> from scipy.ndimage import gaussian_filter + >>> + >>> patch = dc.get_example_patch("ricker_moveout") + >>> + >>> + >>> # Mute first 0.5s (relative to start by default) + >>> muted = patch.mute(time=(0, 0.5)) + >>> + >>> # Mute everything except middle section + >>> kept = patch.mute(time=(0.2, -0.2), mode="complement") + >>> + >>> # Taper with absolute units + >>> muted = patch.mute( + ... time=(0.2, 0.8), + ... taper={'time': 0.02 * dc.units.s}, + ... ) + >>> + >>> # Custom taper function + >>> def custom_taper(envelope): + ... # Apply gaussian filter to envelope to smooth sharp edges. + ... return gaussian_filter(envelope, sigma=2) + >>> + >>> muted = patch.mute(time=(0.2, 0.8), taper=custom_taper) + >>> + >>> # --- LINEAR VELOCITY MUTES --- + >>> + >>> # Classic first break mute: mute early arrivals + >>> # Line from (t=0, d=0) to (t=0.3, d=300) defines velocity=1000 m/s + >>> muted = patch.mute( + ... time=(0, [0, 0.3]), + ... distance=(0, [0, 300]), + ... taper=0.02, + ... relative=False, + ... ) + >>> + >>> # Mute late arrivals: from velocity line to end + >>> muted = patch.mute( + ... time=([0, 0.3], None), + ... distance=([0, 300], None), + ... relative=False, + ... ) + >>> + >>> # Mute wedge between two velocity lines + >>> muted = patch.mute( + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), + ... relative=False, + ... ) + >>> + >>> # Curved mute using multiple control points + >>> muted = patch.mute( + ... time=(0, [0, 0.1, 0.25, 0.35]), + ... distance=(0, [0, 50, 200, 300]), + ... relative=False, + ... ) + >>> + >>> # --- MIXED GEOMETRY --- + >>> + >>> # Linear mute in time-distance, block in other dimension + >>> patch_3d = dc.get_example_patch("nd_patch", dim_count=3) + >>> muted = patch_3d.mute( + ... dim_1=(0, [0, 5]), + ... dim_2=(0, [0, 5]), + ... dim_3=(2, 8), + ... ) + + Notes + ----- + - For linear/planar mutes, all specified dimensions must have the + same number of points to define the geometry. + - Taper is applied perpendicular to the mute boundary. + - Block mutes can be combined with planar mutes across + different dimensions. + - By default, relative=True means values are offsets from coordinate + edges. Use relative=False for absolute coordinate values. + + See Also + -------- + [`Patch.select`](`dascore.Patch.select`) + [`Patch.taper_range`](`dascore.Patch.taper_range`) + """ + # Validate we have dimension specifications + if not kwargs: + msg = "At least one dimension must be specified for muting" + raise ParameterError(msg) + + # Parse boundaries and categorize them + scalar_bounds, array_bounds = _parse_mute_boundaries(patch, kwargs) + + # Validate array boundaries have matching lengths + if array_bounds: + _validate_array_boundaries(array_bounds) + + # Handle coordinate conversions based on mode + if samples: + # Samples mode: values are already indices + # If relative, treat as relative indices (offset from start/end) + # If not relative, treat as absolute indices + if relative: + # Convert relative sample indices to absolute + scalar_bounds = _convert_relative_samples(patch, scalar_bounds) + array_bounds = _convert_relative_samples_arrays(patch, array_bounds) + # else: already absolute sample indices, use directly + else: + # Coordinate mode: values are coordinate values + if relative: + # Convert relative coordinates to absolute + scalar_bounds = _convert_scalar_relative(patch, scalar_bounds) + array_bounds = _convert_array_relative(patch, array_bounds) + # Now convert coordinate values to sample indices + scalar_bounds = _convert_to_samples(patch, scalar_bounds) + array_bounds = _convert_arrays_to_samples(patch, array_bounds) + + # Build mute envelope + envelope = np.ones(patch.shape, dtype=float) + + # Apply block mutes (scalar boundaries) + if scalar_bounds: + block_envelope = _get_block_mute_envelope( + patch, scalar_bounds, taper, window_type + ) + envelope *= block_envelope + + # Apply geometric mutes (array boundaries) - Phase 2 + if array_bounds: + geom_envelope = _get_geometric_mute_envelope( + patch, array_bounds, taper, window_type + ) + envelope *= geom_envelope + + # Apply custom taper function if provided + if callable(taper): + envelope = taper(envelope) + if not isinstance(envelope, np.ndarray): + msg = ( + "Custom taper function must return a numpy array. " + f"Got {type(envelope)}" + ) + raise ParameterError(msg) + if envelope.shape != patch.shape: + msg = ( + f"Custom taper function must return array with shape {patch.shape}. " + f"Got shape {envelope.shape}" + ) + raise ParameterError(msg) + + # Invert envelope for complement mode + if mode == "complement": + envelope = 1.0 - envelope + + # Apply envelope to data + return patch.new(data=patch.data * envelope) + + +def _parse_mute_boundaries(patch, kwargs): + """ + Parse kwargs into boundary specifications. + + Returns + ------- + scalar_bounds : dict + {dim: (min, max)} for dimensions with scalar boundaries + array_bounds : dict + {dim: (array1, array2)} for dimensions with array boundaries + """ + scalar_bounds = {} + array_bounds = {} + + for dim, value in kwargs.items(): + # Validate dimension exists + if dim not in patch.dims: + valid_dims = sorted(patch.dims) + msg = f"Dimension '{dim}' not found. Valid dimensions: {valid_dims}" + raise ParameterError(msg) + + # Must be a tuple of length 2 + if not isinstance(value, tuple) or len(value) != 2: + msg = ( + f"Each dimension must specify 2 boundaries as a tuple. " + f"Got {value} for dimension '{dim}'" + ) + raise ParameterError(msg) + + bound1, bound2 = value + + # Check if boundaries are scalars or arrays + is_scalar_1 = _is_scalar_or_none(bound1) + is_scalar_2 = _is_scalar_or_none(bound2) + + if is_scalar_1 and is_scalar_2: + scalar_bounds[dim] = (bound1, bound2) + else: + # At least one is an array + # Only convert to array if it's actually array-like (not scalar) + if is_scalar_1: + # bound1 is scalar, bound2 is array - this is an error for now + msg = ( + f"Cannot mix scalar and array boundaries for dimension '{dim}'. " + f"Both boundaries must be scalars or both must be arrays with " + f"matching lengths." + ) + raise ParameterError(msg) + arr1 = None if bound1 is None else np.atleast_1d(bound1) + arr2 = None if bound2 is None else np.atleast_1d(bound2) + array_bounds[dim] = (arr1, arr2) + + return scalar_bounds, array_bounds + + +def _is_scalar_or_none(value): + """Check if value is scalar or None (not an array/list).""" + if value is None or value is ...: + return True + return np.isscalar(value) + + +def _validate_array_boundaries(array_bounds): + """Validate that all array boundaries have matching lengths.""" + lengths = {} + for dim, (arr1, arr2) in array_bounds.items(): + len1 = len(arr1) if arr1 is not None else None + len2 = len(arr2) if arr2 is not None else None + if len1 is not None: + lengths[f"{dim}_bound1"] = len1 + if len2 is not None: + lengths[f"{dim}_bound2"] = len2 + + # Check all non-None lengths match + unique_lengths = set(lengths.values()) + if len(unique_lengths) > 1: + msg = ( + "All array boundaries must have the same number of points. " + f"Got varying lengths: {lengths}" + ) + raise ParameterError(msg) + + +def _convert_scalar_relative(patch, scalar_bounds): + """Convert relative scalar values to absolute coordinate values.""" + result = {} + for dim, (val1, val2) in scalar_bounds.items(): + coord = patch.get_coord(dim) + coord_min, coord_max = coord.min(), coord.max() + + # Convert each boundary + abs_val1 = _relative_to_absolute_scalar(val1, coord_min, coord_max) + abs_val2 = _relative_to_absolute_scalar(val2, coord_min, coord_max) + result[dim] = (abs_val1, abs_val2) + + return result + + +def _convert_array_relative(patch, array_bounds): + """Convert relative array values to absolute coordinate values.""" + result = {} + for dim, (arr1, arr2) in array_bounds.items(): + coord = patch.get_coord(dim) + coord_min, coord_max = coord.min(), coord.max() + + # Convert each boundary array + abs_arr1 = _relative_to_absolute_array(arr1, coord_min, coord_max) + abs_arr2 = _relative_to_absolute_array(arr2, coord_min, coord_max) + result[dim] = (abs_arr1, abs_arr2) + + return result + + +def _relative_to_absolute_scalar(value, coord_min, coord_max): + """Convert single relative value to absolute.""" + if value is None or value is ...: + return value + + # Handle datetime/timedelta types + if isinstance(coord_min, np.datetime64): + # Convert value to timedelta if it's a number + if not isinstance(value, (np.timedelta64, np.datetime64)): + value = np.timedelta64(int(value * 1e9), 'ns') + if value >= np.timedelta64(0): + return coord_min + value + else: + return coord_max + value + else: + # Numeric types + if value >= 0: + return coord_min + value + else: + return coord_max + value + + +def _relative_to_absolute_array(arr, coord_min, coord_max): + """Convert array of relative values to absolute.""" + if arr is None: + return None + + # Handle datetime/timedelta types + if isinstance(coord_min, np.datetime64): + # Convert values to timedeltas if they're numbers + if not isinstance(arr[0], (np.timedelta64, np.datetime64)): + arr = np.array([np.timedelta64(int(v * 1e9), 'ns') for v in arr]) + result = np.empty_like(arr, dtype=coord_min.dtype) + positive_mask = arr >= np.timedelta64(0) + result[positive_mask] = coord_min + arr[positive_mask] + result[~positive_mask] = coord_max + arr[~positive_mask] + return result + else: + result = np.empty_like(arr, dtype=float) + positive_mask = arr >= 0 + result[positive_mask] = coord_min + arr[positive_mask] + result[~positive_mask] = coord_max + arr[~positive_mask] + return result + + +def _convert_relative_samples(patch, scalar_bounds): + """Convert relative sample indices to absolute indices.""" + result = {} + for dim, (val1, val2) in scalar_bounds.items(): + coord = patch.get_coord(dim) + length = len(coord) + + # Convert each boundary + abs_val1 = _relative_sample_to_absolute(val1, length) + abs_val2 = _relative_sample_to_absolute(val2, length) + result[dim] = (abs_val1, abs_val2) + + return result + + +def _convert_relative_samples_arrays(patch, array_bounds): + """Convert relative sample indices in arrays to absolute indices.""" + result = {} + for dim, (arr1, arr2) in array_bounds.items(): + coord = patch.get_coord(dim) + length = len(coord) + + # Convert each boundary array + abs_arr1 = _relative_samples_array_to_absolute(arr1, length) + abs_arr2 = _relative_samples_array_to_absolute(arr2, length) + result[dim] = (abs_arr1, abs_arr2) + + return result + + +def _relative_sample_to_absolute(value, length): + """Convert relative sample index to absolute.""" + if value is None or value is ...: + return value + if value >= 0: + return int(value) + else: + return int(length + value) + + +def _relative_samples_array_to_absolute(arr, length): + """Convert array of relative sample indices to absolute.""" + if arr is None: + return None + + result = np.empty_like(arr, dtype=int) + for i, val in enumerate(arr): + if val >= 0: + result[i] = int(val) + else: + result[i] = int(length + val) + return result + + +def _convert_to_samples(patch, scalar_bounds): + """Convert scalar boundaries from coordinate values to sample indices.""" + result = {} + for dim, (val1, val2) in scalar_bounds.items(): + coord = patch.get_coord(dim) + # Convert to indices + idx1 = _value_to_index(coord, val1) if val1 is not None else None + idx2 = _value_to_index(coord, val2) if val2 is not None else None + result[dim] = (idx1, idx2) + return result + + +def _convert_arrays_to_samples(patch, array_bounds): + """Convert array boundaries from coordinate values to sample indices.""" + result = {} + for dim, (arr1, arr2) in array_bounds.items(): + coord = patch.get_coord(dim) + idx_arr1 = _array_to_indices(coord, arr1) if arr1 is not None else None + idx_arr2 = _array_to_indices(coord, arr2) if arr2 is not None else None + result[dim] = (idx_arr1, idx_arr2) + return result + + +def _value_to_index(coord, value): + """Convert coordinate value to index.""" + if value is None or value is ...: + return None + return coord.get_next_index(value) + + +def _array_to_indices(coord, arr): + """Convert array of coordinate values to indices.""" + if arr is None: + return None + return np.array([coord.get_next_index(val) for val in arr]) + + +def _parse_taper_value(taper_val, coord, dim_name): + """ + Parse a taper value into absolute coordinate units. + + Parameters + ---------- + taper_val + Can be float (fraction), Quantity (absolute), or None + coord + The coordinate object for this dimension + dim_name + Name of the dimension (for error messages) + + Returns + ------- + float or timedelta + Taper width in coordinate units, or None if taper_val is None + """ + if taper_val is None: + return None + + # Check if it's a Quantity with units + if isinstance(taper_val, Quantity): + # For time quantities, keep as timedelta; for others convert to float + coord_range = coord.max() - coord.min() + if isinstance(coord_range, np.timedelta64): + # Convert Quantity to timedelta64 + # First convert to seconds, then to nanoseconds for timedelta64 + seconds = taper_val.to('s').magnitude + return np.timedelta64(int(seconds * 1e9), 'ns') + else: + # Convert to float (handles distance, etc.) + return float(taper_val.magnitude) + + # Otherwise it's a float - treat as fraction of dimension range + if not isinstance(taper_val, (int, float)): + msg = ( + f"Taper value for dimension '{dim_name}' must be a float " + f"(fraction) or Quantity (with units). Got {type(taper_val)}" + ) + raise ParameterError(msg) + + # Get dimension range and multiply by fraction + coord_range = coord.max() - coord.min() + if isinstance(coord_range, np.timedelta64): + # For datetime coords, multiply timedelta by fraction + return coord_range * taper_val + else: + return taper_val * coord_range + + +def _validate_taper_dict(taper, scalar_bounds): + """ + Validate that taper dict has valid dimension names. + + Parameters + ---------- + taper : dict + Dictionary of dimension names to taper values + scalar_bounds : dict + Dictionary of scalar boundaries being used + """ + if not isinstance(taper, dict): + return + + for dim in taper: + if dim not in scalar_bounds: + valid_dims = sorted(scalar_bounds.keys()) + msg = ( + f"Taper dimension '{dim}' not found in mute dimensions. " + f"Valid dimensions: {valid_dims}" + ) + raise ParameterError(msg) + + +def _get_block_mute_envelope(patch, scalar_bounds, taper, window_type): + """ + Create envelope for rectangular block mutes. + + Similar to taper_range but creates a mute envelope instead. + """ + # Validate taper if it's a dict + if isinstance(taper, dict): + _validate_taper_dict(taper, scalar_bounds) + + # Check if taper is a Quantity with multiple dimensions + if isinstance(taper, Quantity) and len(scalar_bounds) > 1: + msg = ( + "Cannot use Quantity (with units) for taper when multiple dimensions " + "are specified. Use a dict like: taper={'time': 0.02*dc.units.s, " + "'distance': 10*dc.units.m}" + ) + raise ParameterError(msg) + + envelope = np.ones(patch.shape, dtype=float) + + for dim, (val1, val2) in scalar_bounds.items(): + coord = patch.get_coord(dim) + axis = patch.get_axis(dim) + + # Values are already sample indices at this point + # Handle None values (edges) + if val1 is None or val1 is ...: + val1 = 0 + if val2 is None or val2 is ...: + val2 = len(coord) - 1 + + # Create slice object from indices + slice_obj = slice(int(val1), int(val2) + 1) + + # Create 1D envelope for this dimension + dim_envelope = np.ones(len(coord)) + dim_envelope[slice_obj] = 0 # Zero out mute region + + # Apply taper if specified and not callable + if taper is not None and not callable(taper): + # Get taper value for this dimension + if isinstance(taper, dict): + taper_val = taper.get(dim) + else: + taper_val = taper + + if taper_val is not None: + # Parse taper value (handles fraction vs Quantity) + taper_width = _parse_taper_value(taper_val, coord, dim) + dim_envelope = _apply_taper_1d( + dim_envelope, slice_obj, taper_width, coord, window_type + ) + + # Broadcast to full shape and multiply + indexer = broadcast_for_index( + patch.ndim, axis, value=slice(None), fill=None + ) + envelope *= dim_envelope[indexer] + + return envelope + + +def _apply_taper_1d(envelope, slice_obj, taper_width, coord, window_type): + """Apply taper to 1D envelope at boundaries.""" + func = _get_window_function(window_type) + + # Get indices of mute region + start_idx = slice_obj.start if slice_obj.start is not None else 0 + stop_idx = slice_obj.stop if slice_obj.stop is not None else len(coord) + + # Convert taper width to samples + if hasattr(coord, 'step') and coord.step is not None: + # Handle both timedelta and numeric step/taper_width + if isinstance(taper_width, np.timedelta64): + # Both are timedeltas - can divide directly + taper_samples = int(taper_width / coord.step) + elif isinstance(coord.step, np.timedelta64): + # taper_width is numeric, step is timedelta - shouldn't happen + # Convert timedelta step to float for division + step_float = to_float(coord.step) + taper_samples = int(taper_width / step_float) + else: + # Both numeric + taper_samples = int(taper_width / coord.step) + else: + # For non-evenly sampled coords, estimate + coord_range = coord.max() - coord.min() + if isinstance(taper_width, np.timedelta64) and isinstance(coord_range, np.timedelta64): + # Both timedeltas + taper_samples = int(len(coord) * (taper_width / coord_range)) + elif isinstance(taper_width, np.timedelta64): + # Convert timedelta to float + taper_float = to_float(taper_width) + range_float = to_float(coord_range) + taper_samples = int(len(coord) * taper_float / range_float) + elif isinstance(coord_range, np.timedelta64): + # Convert coord_range to float + range_float = to_float(coord_range) + taper_samples = int(len(coord) * taper_width / range_float) + else: + # Both numeric + taper_samples = int(len(coord) * taper_width / coord_range) + + taper_samples = max(1, min(taper_samples, (stop_idx - start_idx) // 2)) + + # Apply taper at start of mute region + if taper_samples > 0 and start_idx > 0: + window = func(2 * taper_samples)[:taper_samples] + taper_start = max(0, start_idx - taper_samples) + envelope[taper_start:start_idx] *= window + + # Apply taper at end of mute region + if taper_samples > 0 and stop_idx < len(coord): + window = func(2 * taper_samples)[taper_samples:] + taper_end = min(len(coord), stop_idx + taper_samples) + envelope[stop_idx:taper_end] *= window[:taper_end - stop_idx] + + return envelope + + +def _get_geometric_mute_envelope(patch, array_bounds, taper, window_type): + """ + Create envelope for geometric (line/plane/hyperplane) mutes. + + This is Phase 2 - to be implemented for linear/planar mutes. + For now, raise an informative error. + """ + msg = ( + "Geometric mutes (using arrays to define lines/planes) are not yet " + "implemented. Use scalar boundaries for block mutes, or stay tuned " + "for the next release!" + ) + raise NotImplementedError(msg) diff --git a/tests/test_examples.py b/tests/test_examples.py index a964007af..19a5397b9 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -73,6 +73,12 @@ def test_moveout(self): class TestDeltaPatch: """Tests for the delta_patch example.""" + @pytest.mark.parametrize("shape", ((10, 10), (100, 100), (1, 10))) + def test_shape(self, shape): + """Ensure the shape parameter controls the shape of the patch.""" + patch = dc.get_example_patch("delta_patch", shape=shape) + assert patch.shape == shape + @pytest.mark.parametrize("invalid_dim", ["inv_dim", "", None, 123, 1.1]) def test_delta_patch_invalid_dim(self, invalid_dim): """ diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py new file mode 100644 index 000000000..917dd3b28 --- /dev/null +++ b/tests/test_proc/test_mute.py @@ -0,0 +1,436 @@ +"""Tests for mute processing function.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.exceptions import ParameterError + + +@pytest.fixture(scope="session") +def patch_ones(random_patch): + """Return a patch filled with ones.""" + return random_patch.new(data=np.ones_like(random_patch.data)) + + +@pytest.fixture(scope="session") +def ricker_patch(): + """Return ricker moveout patch for velocity mute testing.""" + return dc.get_example_patch("ricker_moveout") + + +class TestMuteBasics: + """Basic tests for mute functionality.""" + + def test_mute_exists(self, random_patch): + """Ensure mute method exists on patch.""" + assert hasattr(random_patch, "mute") + assert callable(random_patch.mute) + + def test_mute_no_kwargs_raises(self, random_patch): + """Mute without dimension specifications should raise.""" + with pytest.raises(ParameterError, match="At least one dimension"): + random_patch.mute() + + def test_invalid_dimension_raises(self, random_patch): + """Mute with invalid dimension should raise.""" + with pytest.raises(ParameterError, match="not found"): + random_patch.mute(invalid_dim=(0, 10)) + + def test_not_tuple_raises(self, random_patch): + """Boundary must be tuple of length 2.""" + with pytest.raises(ParameterError, match="must specify 2 boundaries"): + random_patch.mute(time=5) + + def test_tuple_wrong_length_raises(self, random_patch): + """Tuple must have exactly 2 elements.""" + with pytest.raises(ParameterError, match="must specify 2 boundaries"): + random_patch.mute(time=(1, 2, 3)) + + +class TestBlockMutes1D: + """Test 1D block mutes (single dimension).""" + + def test_mute_start_of_time(self, patch_ones): + """Mute first portion of time dimension.""" + # Use relative=True (default) with numeric offset + muted = patch_ones.mute(time=(0, 0.5)) + + # Check that first samples are zero + assert np.allclose(muted.data[:, 0], 0) + # Check that later samples are still 1 + assert np.allclose(muted.data[:, -1], 1) + + def test_mute_end_of_time(self, patch_ones): + """Mute last portion of time dimension.""" + # Use relative with negative values (from end) + # (-0.5, None) means from 0.5s before end to the end + muted = patch_ones.mute(time=(-0.5, None)) + + # Check that last samples are zero + assert np.allclose(muted.data[:, -1], 0) + # Check that early samples are still 1 + assert np.allclose(muted.data[:, 0], 1) + + def test_mute_middle_of_distance(self, patch_ones): + """Mute middle portion of distance dimension.""" + dist = patch_ones.coords.get_array("distance") + mid_start = dist[len(dist) // 4] + mid_end = dist[3 * len(dist) // 4] + + muted = patch_ones.mute(distance=(mid_start, mid_end), relative=False) + + # Check middle is muted + mid_idx = len(dist) // 2 + assert np.allclose(muted.data[mid_idx, :], 0) + # Check edges are not muted + assert np.allclose(muted.data[0, :], 1) + assert np.allclose(muted.data[-1, :], 1) + + def test_mute_relative_default(self, patch_ones): + """Test that relative=True is default.""" + # Should mute first second (relative to start) + muted = patch_ones.mute(time=(0, 1)) + # First samples should be zero + assert muted.data[:, 0].max() < 0.1 + + def test_mute_relative_negative(self, patch_ones): + """Test relative with negative values (from end).""" + muted = patch_ones.mute(time=(-1, None)) + # Last samples should be zero + assert muted.data[:, -1].max() < 0.1 + + +class TestBlockMutes2D: + """Test 2D block mutes (multiple dimensions).""" + + def test_mute_rectangular_region(self, patch_ones): + """Mute a rectangular region in 2D.""" + time_coords = patch_ones.coords.get_array("time") + dist_coords = patch_ones.coords.get_array("distance") + + t1, t2 = time_coords[10], time_coords[20] + d1, d2 = dist_coords[5], dist_coords[15] + + muted = patch_ones.mute( + time=(t1, t2), + distance=(d1, d2), + relative=False, + ) + + # Check interior is muted + assert np.allclose(muted.data[10, 15], 0) + # Check corners are not muted + assert np.allclose(muted.data[0, 0], 1) + assert np.allclose(muted.data[-1, -1], 1) + + def test_mute_strips(self, patch_ones): + """Mute strips along each dimension.""" + # Mute a time strip + muted_time = patch_ones.mute(time=(0.5, 1.0)) + # All distances should be affected equally + assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) + + # Mute a distance strip + muted_dist = patch_ones.mute(distance=(10, 20)) + # All times should be affected equally + assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) + + +class TestMuteModes: + """Test different mute modes.""" + + def test_mode_union(self, patch_ones): + """Test union mode (default, mutes inside region).""" + muted = patch_ones.mute(time=(2.0, 6.0), mode="union") + # Middle (at ~4s) should be zero + mid_idx = len(patch_ones.coords.get_array("time")) // 2 + assert muted.data[:, mid_idx].max() < 0.1 + # Edges should be one + assert muted.data[:, 0].min() > 0.9 + + def test_mode_complement(self, patch_ones): + """Test complement mode (mutes outside region).""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[10], time_coords[20] + + muted = patch_ones.mute(time=(t1, t2), mode="complement", relative=False) + + # Edges should be zero + assert muted.data[:, 0].max() < 0.1 + assert muted.data[:, -1].max() < 0.1 + # Middle should be one + mid_idx = 15 + assert muted.data[:, mid_idx].min() > 0.9 + + +class TestMuteTaper: + """Test taper application in mutes.""" + + def test_mute_with_taper(self, patch_ones): + """Mute with taper creates gradual transition.""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + # Mute without taper + muted_no_taper = patch_ones.mute( + time=(t1, t2), + relative=False, + taper=None, + ) + + # Mute with taper (5% of dimension range) + muted_with_taper = patch_ones.mute( + time=(t1, t2), + relative=False, + taper=0.05, + ) + + # Without taper should be sharp transition + assert muted_no_taper.data[:, 19].min() > 0.9 + assert muted_no_taper.data[:, 20].max() < 0.1 + + # With taper should have gradual transition + # Values in taper region should be between 0 and 1 + taper_region = muted_with_taper.data[:, 15:20] + assert (taper_region > 0).any() + assert (taper_region < 1).any() + + def test_taper_dict(self, patch_ones): + """Test dimension-specific taper using dict.""" + time_coords = patch_ones.coords.get_array("time") + dist_coords = patch_ones.coords.get_array("distance") + + taper_dict = { + "time": 0.05, # 5% of time range + "distance": 0, # No taper on distance + } + + muted = patch_ones.mute( + time=(time_coords[20], time_coords[40]), + distance=(dist_coords[10], dist_coords[30]), + relative=False, + taper=taper_dict, + ) + + # Should have taper in time, sharp in distance + assert muted.shape == patch_ones.shape + + def test_taper_window_types(self, patch_ones): + """Test different window types for taper.""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + for window_type in ["hann", "hamming", "triang"]: + muted = patch_ones.mute( + time=(t1, t2), + relative=False, + taper=0.05, # 5% of dimension range + window_type=window_type, + ) + assert isinstance(muted, dc.Patch) + assert muted.shape == patch_ones.shape + + def test_taper_with_quantity(self, patch_ones): + """Test taper with Quantity (absolute units).""" + from dascore.units import s + + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + # Use absolute time value with units (single dimension) + muted = patch_ones.mute( + time=(t1, t2), + relative=False, + taper={'time': 0.05 * s}, + ) + + # Should have gradual transition + taper_region = muted.data[:, 15:20] + assert (taper_region > 0).any() + assert (taper_region < 1).any() + + def test_taper_quantity_multiple_dims_raises(self, patch_ones): + """Test that Quantity taper with multiple dims raises error.""" + from dascore.units import s + + time_coords = patch_ones.coords.get_array("time") + dist_coords = patch_ones.coords.get_array("distance") + + # Should raise error if Quantity used without dict + with pytest.raises(ParameterError, match="Cannot use Quantity"): + patch_ones.mute( + time=(time_coords[20], time_coords[40]), + distance=(dist_coords[10], dist_coords[30]), + relative=False, + taper=0.05 * s, + ) + + def test_taper_mixed_quantity_fraction(self, patch_ones): + """Test mixed taper: fraction for one dim, Quantity for another.""" + from dascore.units import s, m + + time_coords = patch_ones.coords.get_array("time") + dist_coords = patch_ones.coords.get_array("distance") + + muted = patch_ones.mute( + time=(time_coords[20], time_coords[40]), + distance=(dist_coords[10], dist_coords[30]), + relative=False, + taper={'time': 0.05, 'distance': 10 * m}, + ) + + assert muted.shape == patch_ones.shape + + def test_taper_custom_function(self, patch_ones): + """Test custom taper function.""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + # Define custom taper that squares the envelope + def custom_taper(envelope): + return envelope ** 2 + + muted = patch_ones.mute( + time=(t1, t2), + relative=False, + taper=custom_taper, + ) + + # Middle should still be zero + assert muted.data[:, 30].max() < 0.1 + # Edges should still be one + assert muted.data[:, 0].min() > 0.9 + + def test_taper_custom_function_invalid_return(self, patch_ones): + """Test that custom taper with invalid return raises error.""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + # Function that returns wrong type + def bad_taper(envelope): + return list(envelope) + + with pytest.raises(ParameterError, match="must return a numpy array"): + patch_ones.mute( + time=(t1, t2), + relative=False, + taper=bad_taper, + ) + + def test_taper_custom_function_wrong_shape(self, patch_ones): + """Test that custom taper with wrong shape raises error.""" + time_coords = patch_ones.coords.get_array("time") + t1, t2 = time_coords[20], time_coords[40] + + # Function that returns wrong shape + def bad_shape_taper(envelope): + return np.ones((10, 10)) + + with pytest.raises(ParameterError, match="must return array with shape"): + patch_ones.mute( + time=(t1, t2), + relative=False, + taper=bad_shape_taper, + ) + + +class TestMuteEdgeCases: + """Test edge cases and boundary conditions.""" + + def test_mute_with_none_boundaries(self, patch_ones): + """Test using None to reference coordinate edges.""" + # Mute from start to middle + time_coords = patch_ones.coords.get_array("time") + mid_time = time_coords[len(time_coords) // 2] + + muted = patch_ones.mute(time=(None, mid_time), relative=False) + + # First half should be zero + assert muted.data[:, 0].max() < 0.1 + # Second half should be one + assert muted.data[:, -1].min() > 0.9 + + def test_mute_entire_dimension(self, patch_ones): + """Mute entire dimension.""" + muted = patch_ones.mute(time=(None, None), relative=False) + # Everything should be zero + assert np.allclose(muted.data, 0) + + def test_mute_zero_width(self, patch_ones): + """Mute with same start and end values.""" + time_coords = patch_ones.coords.get_array("time") + t = time_coords[10] + + muted = patch_ones.mute(time=(t, t), relative=False) + # Should mute just that one sample (or very close) + assert muted.data[:, 10].max() < 0.1 + + +class TestMuteWithSamples: + """Test mute with samples=True.""" + + def test_mute_samples_mode(self, patch_ones): + """Test mute using sample indices.""" + # Mute samples 10 to 20 + muted = patch_ones.mute(time=(10, 20), samples=True) + + # Samples 10-20 should be zero + assert muted.data[:, 15].max() < 0.1 + # Other samples should be one + assert muted.data[:, 0].min() > 0.9 + assert muted.data[:, -1].min() > 0.9 + + +class TestGeometricMutes: + """Test geometric (array-based) mutes - Phase 2.""" + + def test_array_boundaries_not_implemented(self, ricker_patch): + """Array boundaries should raise error for mixing scalar/array.""" + # Currently we don't allow mixing scalars and arrays + with pytest.raises(ParameterError, match="Cannot mix scalar and array"): + ricker_patch.mute( + time=(0, [0, 0.3]), + distance=(0, [0, 300]), + relative=False, + ) + + def test_mismatched_array_lengths_raises(self, random_patch): + """Arrays with different lengths should raise error during parsing.""" + with pytest.raises(ParameterError, match="Cannot mix scalar and array"): + # This currently raises mixing error before length checking + random_patch.mute( + time=(0, [0, 0.3, 0.5]), # scalar and array + distance=(0, [0, 300]), # scalar and array + relative=False, + ) + + +class TestMutePreservesMetadata: + """Test that mute preserves patch metadata.""" + + def test_preserves_coords(self, random_patch): + """Mute should preserve coordinates.""" + muted = random_patch.mute(time=(0, 0.5)) + assert muted.coords == random_patch.coords + + def test_preserves_attrs(self, random_patch): + """Mute should preserve attributes (except history).""" + muted = random_patch.mute(time=(0, 0.5)) + # History will be different due to processing + assert muted.attrs.data_type == random_patch.attrs.data_type + assert muted.attrs.network == random_patch.attrs.network + assert muted.attrs.station == random_patch.attrs.station + + def test_preserves_shape(self, random_patch): + """Mute should preserve shape.""" + muted = random_patch.mute(time=(0, 0.5)) + assert muted.shape == random_patch.shape + + def test_preserves_dtype(self, random_patch): + """Mute should preserve data type.""" + muted = random_patch.mute(time=(0, 0.5)) + # Allow for float conversion + assert muted.dtype in (random_patch.dtype, np.float64, np.float32) From 9226a362073805ca6991daf518663accc09a8f12 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 16 Oct 2025 17:41:46 +0100 Subject: [PATCH 02/15] mute progress --- dascore/core/patch.py | 1 + dascore/proc/basic.py | 28 ++ dascore/proc/mute.py | 820 ++++++++-------------------------- dascore/proc/taper.py | 4 +- dascore/utils/signal.py | 1 + tests/test_proc/test_mute.py | 805 +++++++++++++++++---------------- tests/test_proc/test_taper.py | 35 ++ 7 files changed, 673 insertions(+), 1021 deletions(-) diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 9f53b21de..02bdf1518 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -364,6 +364,7 @@ def T(self): # noqa: N802 apply_ufunc = dascore.utils.array.apply_ufunc get_patch_names = get_patch_names get_axis = dascore.proc.get_axis + full = dascore.proc.full def get_patch_name(self, *args, **kwargs) -> str: """ diff --git a/dascore/proc/basic.py b/dascore/proc/basic.py index 2205e6cb1..ca628d6ab 100644 --- a/dascore/proc/basic.py +++ b/dascore/proc/basic.py @@ -789,3 +789,31 @@ def flip(patch, *dims, flip_coords=True): data = np.flip(patch.data, axis=axes) if dims else patch.data coords = patch.coords.flip(*dims) if flip_coords else patch.coords return patch.new(data=data, coords=coords) + + +@patch_function() +def full(patch, fill_value): + """ + Return an identical patch with the data replaced by fill_value. + + Parameters + ---------- + patch + The patch to fill. + fill_value + The value in the output patch. + + Examples + -------- + >>> import dascore as dc + >>> patch = dc.get_example_patch() + >>> + >>> # Get a patch identical to original but with data array containing + >>> # Only 1s. + >>> one_patch = patch.full(1.0) + >>> + >>> # Same thing, except for 0s. + >>> zero_patch = patch.full(0.0) + """ + array = np.full(patch.data.shape, fill_value) + return patch.update(data=array) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index da5e31ed0..20e5cc25e 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -2,107 +2,179 @@ from __future__ import annotations -from collections.abc import Callable -from typing import Literal +from collections.abc import Mapping import numpy as np from dascore.constants import PatchType from dascore.exceptions import ParameterError -from dascore.units import Quantity -from dascore.utils.docs import compose_docstring -from dascore.utils.misc import broadcast_for_index -from dascore.utils.patch import patch_function -from dascore.utils.signal import WINDOW_FUNCTIONS, _get_window_function -from dascore.utils.time import to_float +from dascore.utils.patch import get_dim_axis_value, patch_function -parameter_docstring = """ -Parameters ----------- -patch - The patch instance. -mode - - "union": mute data inside the specified region [default] - - "complement": mute data outside the specified region -taper - Taper width at mute boundaries. Can be: - - None: sharp mute (no taper) - - float (0.0-1.0): fraction of dimension range (e.g., 0.05 = 5%) - The number of samples in the mute is held constant across dimensions - by calculating the samples the fraction represents for each dimension - and using the minimum sample count. This helps avoid distorted mutes - based on dimensions with vastly different lengths. - - Quantity with units: absolute value (e.g., 0.02*dc.units.s) - Note: If multiple dimensions specified, must use dict. - - dict: {dim: taper_value} for dimension-specific taper - Values can be floats (fractions) or Quantities (absolute) - - Callable: custom function that receives and modifies envelope array. -window_type - Window function for tapering (only used for non-callable taper). - Options: - {taper_type}. -relative - If True (default), values are relative to coordinate edges. - Positive values are offsets from start, negative from end. -samples - If True, values specified in samples rather than coordinate values. -**kwargs - Dimension specifications as (boundary_1, boundary_2) pairs. - Each boundary can be: - - Scalar: constant value (defines plane perpendicular to dimension) - - Array/list: coordinates defining line/curve/surface - - None or ...: edge of coordinate (min or max) - -""" - - - -@patch_function() -@compose_docstring(taper_type=sorted(WINDOW_FUNCTIONS), params=parameter_docstring) -def mute_envelope( - patch: PatchType, - *, - mode: Literal["union", "complement"] = "union", - taper: float | dict | None = None, - window_type: str = "hann", - relative: bool = True, - samples: bool = False, - **kwargs, -) -> PatchType: +def _get_mute_params(patch, kwargs): """ - Calculate a patch envelope which applies mute defined by parameters. - - The resulting patch simply needs to be multiplied by the original - patch to achieve muting. This allows for more fine-grained control - of the + Get (and validate) the muting parameters. + Returns arrays of dims, axes, and values. """ + # Validate we have dimension specifications + if not kwargs or len(kwargs) > 2: + msg = ( + "Mute requires one or two keyword arguments must be provided. " + f"You passed {kwargs}" + ) + raise ParameterError(msg) + dim_ax_vals = get_dim_axis_value(patch, kwargs=kwargs, allow_multiple=True) + dims = [x[0] for x in dim_ax_vals] + axes = np.array([x[1] for x in dim_ax_vals]) + # Handle single dimension Mute. + if len(dims) == 1: + vals = np.array([x[2] for x in dim_ax_vals]) + if vals.size != 2: + msg = "Mute requires two boundaries when using a single dimension." + raise ParameterError(msg) + else: + breakpoint() + print() + return dims, axes, vals + + +def _get_taper_vals(dims, taper, patch, samples): + """Get the taper values ordered by dimension, in samples.""" + # Just broadcast taper to same length as dims. + if taper is None: + return [0] * len(dims) + if not isinstance(taper, Mapping): + vals = [taper] * len(dims) + else: + # Otherwise each dimension's taper must be specified. + if not set(dims) == set(taper): + msg = ( + f"If a taper dictionary is used in Mute, it must have all " + f"the same keys as the dimensions. Kwarg dims are {dims} and" + f"taper keys are {list(taper)}." + ) + raise ParameterError(msg) + vals = [taper[dim] for dim in dims] + out = [] + for dim, val in zip(dims, vals): + coord = patch.get_coord(dim) + if val is None: + out.append(0) + elif samples: + out.append(val) + else: + coord = patch.get_coord(dim) + coord.coord_range(val) + out.append(val) + return np.array(out) + + +def _mute_patch_1d( + patch, + dim, + vals, + taper_samps, + window, + samples, + relative, + invert, +): + """Apply mute to 1D patch using patch.range_taper.""" + coord = patch.get_coord(dim) + + if not samples: + # Get range represented by values. This is a bit ugly... + sel = coord.select(tuple(*vals), relative=relative)[1] + start = sel.start if sel.start is not None else 0 + # stop can be None (open interval) or an exclusive upper bound + if sel.stop is None: + stop = len(coord) + else: + stop = sel.stop + trange = (start, stop) + else: + # vals is a 2D array with shape (1, 2), flatten to get (start, stop) + trange = vals.flatten() if vals.ndim > 1 else vals + + # Handle edge case: if muting to the end with no taper, use 2-value form + # because 4-value form can't express "mute to the very last sample" + max_idx = len(coord) - 1 + if taper_samps == 0 and trange[1] >= len(coord): + taper_dict = {dim: [trange[0], max_idx]} + else: + # Clamp taper boundaries to valid sample indices. + # taper_range uses exclusive upper bounds in 4-value form [a,b,c,d] + # which zeros indices [b, c). + taper_dict = { + dim: [ + max(0, trange[0] - taper_samps), + trange[0], + min(max_idx, trange[1]), + min(max_idx, trange[1] + taper_samps), + ] + } + out = patch.taper_range( + invert=not invert, + samples=True, + window_type=window, + **taper_dict, + ) + return out @patch_function() -@compose_docstring( - taper_type=sorted(WINDOW_FUNCTIONS), - params=parameter_docstring, -) def mute( patch: PatchType, *, - mode: Literal["union", "complement"] = "union", taper: float | dict | None = None, window_type: str = "hann", + invert: bool = False, relative: bool = True, samples: bool = False, **kwargs, ) -> PatchType: """ - Mute (zero out) patch data in specified regions. + Mute (zero out) data in a region specified by one or more lines. Each dimension accepts two boundary specifications that define - the mute region. Boundaries can be scalar values (for planes/blocks) - or arrays of values (for lines/curves/surfaces). + the mute region. Boundaries can be scalar values or arrays of values. - {params} + Parameters + ---------- + patch + The patch instance. + taper + Taper width at mute boundaries. Can be: + - None: sharp mute (no taper) + - float (0.0-1.0): fraction of dimension range (e.g., 0.05 = 5%) + The number of samples in the mute is held constant across dimensions + by calculating the samples the fraction represents for each dimension + and using the minimum sample count. This helps avoid distorted mutes + based on dimensions with vastly different lengths. + - Quantity with units: absolute value (e.g., 0.02*dc.units.s) + Note: If multiple dimensions specified, must use dict. + - dict: (dim: taper_value) for dimension-specific taper + Values can be floats (fractions) or Quantities (absolute) + - Callable: custom function that receives and modifies envelope array. + window_type + Window function for tapering (only used for non-callable taper). + Options: + {sorted(WINDOW_FUNCTIONS)}. + invert + If True, invert the taper such that the values outside the defined region + are set to 0. + relative + If True (default), values are relative to coordinate edges. + Positive values are offsets from start, negative from end. + samples + If True, values specified in samples rather than coordinate values. + **kwargs + Dimension specifications as (boundary_1, boundary_2) pairs. + Each boundary can be: + - Scalar: constant value (defines plane perpendicular to dimension) + - Array/list: coordinates defining line/curve/surface + - None or ...: edge of coordinate (min or max) Examples -------- @@ -124,598 +196,74 @@ def mute( ... taper={'time': 0.02 * dc.units.s}, ... ) >>> - >>> # Custom taper function - >>> def custom_taper(envelope): - ... # Apply gaussian filter to envelope to smooth sharp edges. - ... return gaussian_filter(envelope, sigma=2) - >>> - >>> muted = patch.mute(time=(0.2, 0.8), taper=custom_taper) - >>> - >>> # --- LINEAR VELOCITY MUTES --- - >>> >>> # Classic first break mute: mute early arrivals >>> # Line from (t=0, d=0) to (t=0.3, d=300) defines velocity=1000 m/s >>> muted = patch.mute( ... time=(0, [0, 0.3]), - ... distance=(0, [0, 300]), + ... distance=(None, [0, 300]), ... taper=0.02, - ... relative=False, ... ) >>> >>> # Mute late arrivals: from velocity line to end >>> muted = patch.mute( ... time=([0, 0.3], None), - ... distance=([0, 300], None), - ... relative=False, + ... distance=([0, 300], 0), ... ) >>> >>> # Mute wedge between two velocity lines >>> muted = patch.mute( ... time=([0, 0.375], [0, 0.25]), ... distance=([0, 300], [0, 300]), - ... relative=False, ... ) >>> - >>> # Curved mute using multiple control points + >>> # Mute wedge outside two velocity lines >>> muted = patch.mute( - ... time=(0, [0, 0.1, 0.25, 0.35]), - ... distance=(0, [0, 50, 200, 300]), - ... relative=False, + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), + ... invert=True, ... ) >>> - >>> # --- MIXED GEOMETRY --- - >>> - >>> # Linear mute in time-distance, block in other dimension - >>> patch_3d = dc.get_example_patch("nd_patch", dim_count=3) - >>> muted = patch_3d.mute( - ... dim_1=(0, [0, 5]), - ... dim_2=(0, [0, 5]), - ... dim_3=(2, 8), + >>> # Apply custom tapering + >>> ones = patch.full(1.0) + >>> envelope = ones.mute( + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), ... ) + >>> # Knock down edges with gaussian filter. + >>> smooth = envelope.gaussian_filter(time=5, distance=5, samples=True) + >>> # Then multiply the two patches. + >>> result = patch * smooth Notes ----- - - For linear/planar mutes, all specified dimensions must have the - same number of points to define the geometry. - - Taper is applied perpendicular to the mute boundary. - - Block mutes can be combined with planar mutes across - different dimensions. - - By default, relative=True means values are offsets from coordinate + - By relative=True (the default) means values are offsets from coordinate edges. Use relative=False for absolute coordinate values. + - Currently, mute doesn't support more than 2 dimensions. + + - For more control over tapering, use a patch with one values then apply + custom tapering/smooting before multiplying with the original patch. + See example section for more details. + See Also -------- [`Patch.select`](`dascore.Patch.select`) [`Patch.taper_range`](`dascore.Patch.taper_range`) """ - # Validate we have dimension specifications - if not kwargs: - msg = "At least one dimension must be specified for muting" - raise ParameterError(msg) - - # Parse boundaries and categorize them - scalar_bounds, array_bounds = _parse_mute_boundaries(patch, kwargs) - - # Validate array boundaries have matching lengths - if array_bounds: - _validate_array_boundaries(array_bounds) - - # Handle coordinate conversions based on mode - if samples: - # Samples mode: values are already indices - # If relative, treat as relative indices (offset from start/end) - # If not relative, treat as absolute indices - if relative: - # Convert relative sample indices to absolute - scalar_bounds = _convert_relative_samples(patch, scalar_bounds) - array_bounds = _convert_relative_samples_arrays(patch, array_bounds) - # else: already absolute sample indices, use directly - else: - # Coordinate mode: values are coordinate values - if relative: - # Convert relative coordinates to absolute - scalar_bounds = _convert_scalar_relative(patch, scalar_bounds) - array_bounds = _convert_array_relative(patch, array_bounds) - # Now convert coordinate values to sample indices - scalar_bounds = _convert_to_samples(patch, scalar_bounds) - array_bounds = _convert_arrays_to_samples(patch, array_bounds) - - # Build mute envelope - envelope = np.ones(patch.shape, dtype=float) - - # Apply block mutes (scalar boundaries) - if scalar_bounds: - block_envelope = _get_block_mute_envelope( - patch, scalar_bounds, taper, window_type - ) - envelope *= block_envelope - - # Apply geometric mutes (array boundaries) - Phase 2 - if array_bounds: - geom_envelope = _get_geometric_mute_envelope( - patch, array_bounds, taper, window_type - ) - envelope *= geom_envelope - - # Apply custom taper function if provided - if callable(taper): - envelope = taper(envelope) - if not isinstance(envelope, np.ndarray): - msg = ( - "Custom taper function must return a numpy array. " - f"Got {type(envelope)}" - ) - raise ParameterError(msg) - if envelope.shape != patch.shape: - msg = ( - f"Custom taper function must return array with shape {patch.shape}. " - f"Got shape {envelope.shape}" - ) - raise ParameterError(msg) - - # Invert envelope for complement mode - if mode == "complement": - envelope = 1.0 - envelope - - # Apply envelope to data - return patch.new(data=patch.data * envelope) - - -def _parse_mute_boundaries(patch, kwargs): - """ - Parse kwargs into boundary specifications. - - Returns - ------- - scalar_bounds : dict - {dim: (min, max)} for dimensions with scalar boundaries - array_bounds : dict - {dim: (array1, array2)} for dimensions with array boundaries - """ - scalar_bounds = {} - array_bounds = {} - - for dim, value in kwargs.items(): - # Validate dimension exists - if dim not in patch.dims: - valid_dims = sorted(patch.dims) - msg = f"Dimension '{dim}' not found. Valid dimensions: {valid_dims}" - raise ParameterError(msg) - - # Must be a tuple of length 2 - if not isinstance(value, tuple) or len(value) != 2: - msg = ( - f"Each dimension must specify 2 boundaries as a tuple. " - f"Got {value} for dimension '{dim}'" - ) - raise ParameterError(msg) - - bound1, bound2 = value - - # Check if boundaries are scalars or arrays - is_scalar_1 = _is_scalar_or_none(bound1) - is_scalar_2 = _is_scalar_or_none(bound2) - - if is_scalar_1 and is_scalar_2: - scalar_bounds[dim] = (bound1, bound2) - else: - # At least one is an array - # Only convert to array if it's actually array-like (not scalar) - if is_scalar_1: - # bound1 is scalar, bound2 is array - this is an error for now - msg = ( - f"Cannot mix scalar and array boundaries for dimension '{dim}'. " - f"Both boundaries must be scalars or both must be arrays with " - f"matching lengths." - ) - raise ParameterError(msg) - arr1 = None if bound1 is None else np.atleast_1d(bound1) - arr2 = None if bound2 is None else np.atleast_1d(bound2) - array_bounds[dim] = (arr1, arr2) - - return scalar_bounds, array_bounds - - -def _is_scalar_or_none(value): - """Check if value is scalar or None (not an array/list).""" - if value is None or value is ...: - return True - return np.isscalar(value) - - -def _validate_array_boundaries(array_bounds): - """Validate that all array boundaries have matching lengths.""" - lengths = {} - for dim, (arr1, arr2) in array_bounds.items(): - len1 = len(arr1) if arr1 is not None else None - len2 = len(arr2) if arr2 is not None else None - if len1 is not None: - lengths[f"{dim}_bound1"] = len1 - if len2 is not None: - lengths[f"{dim}_bound2"] = len2 - - # Check all non-None lengths match - unique_lengths = set(lengths.values()) - if len(unique_lengths) > 1: - msg = ( - "All array boundaries must have the same number of points. " - f"Got varying lengths: {lengths}" + dims, axes, values = _get_mute_params(patch, kwargs) + taper_vals = _get_taper_vals(dims, taper, patch, samples) + # Easy path for 1D mute. + if len(dims) == 1: + out = _mute_patch_1d( + patch, + dims[0], + values, + taper_vals[0], + window=window_type if taper is not None else "boxcar", + samples=samples, + relative=relative, + invert=invert, ) - raise ParameterError(msg) - + return out -def _convert_scalar_relative(patch, scalar_bounds): - """Convert relative scalar values to absolute coordinate values.""" - result = {} - for dim, (val1, val2) in scalar_bounds.items(): - coord = patch.get_coord(dim) - coord_min, coord_max = coord.min(), coord.max() - - # Convert each boundary - abs_val1 = _relative_to_absolute_scalar(val1, coord_min, coord_max) - abs_val2 = _relative_to_absolute_scalar(val2, coord_min, coord_max) - result[dim] = (abs_val1, abs_val2) - - return result - - -def _convert_array_relative(patch, array_bounds): - """Convert relative array values to absolute coordinate values.""" - result = {} - for dim, (arr1, arr2) in array_bounds.items(): - coord = patch.get_coord(dim) - coord_min, coord_max = coord.min(), coord.max() - - # Convert each boundary array - abs_arr1 = _relative_to_absolute_array(arr1, coord_min, coord_max) - abs_arr2 = _relative_to_absolute_array(arr2, coord_min, coord_max) - result[dim] = (abs_arr1, abs_arr2) - - return result - - -def _relative_to_absolute_scalar(value, coord_min, coord_max): - """Convert single relative value to absolute.""" - if value is None or value is ...: - return value - - # Handle datetime/timedelta types - if isinstance(coord_min, np.datetime64): - # Convert value to timedelta if it's a number - if not isinstance(value, (np.timedelta64, np.datetime64)): - value = np.timedelta64(int(value * 1e9), 'ns') - if value >= np.timedelta64(0): - return coord_min + value - else: - return coord_max + value - else: - # Numeric types - if value >= 0: - return coord_min + value - else: - return coord_max + value - - -def _relative_to_absolute_array(arr, coord_min, coord_max): - """Convert array of relative values to absolute.""" - if arr is None: - return None - - # Handle datetime/timedelta types - if isinstance(coord_min, np.datetime64): - # Convert values to timedeltas if they're numbers - if not isinstance(arr[0], (np.timedelta64, np.datetime64)): - arr = np.array([np.timedelta64(int(v * 1e9), 'ns') for v in arr]) - result = np.empty_like(arr, dtype=coord_min.dtype) - positive_mask = arr >= np.timedelta64(0) - result[positive_mask] = coord_min + arr[positive_mask] - result[~positive_mask] = coord_max + arr[~positive_mask] - return result - else: - result = np.empty_like(arr, dtype=float) - positive_mask = arr >= 0 - result[positive_mask] = coord_min + arr[positive_mask] - result[~positive_mask] = coord_max + arr[~positive_mask] - return result - - -def _convert_relative_samples(patch, scalar_bounds): - """Convert relative sample indices to absolute indices.""" - result = {} - for dim, (val1, val2) in scalar_bounds.items(): - coord = patch.get_coord(dim) - length = len(coord) - - # Convert each boundary - abs_val1 = _relative_sample_to_absolute(val1, length) - abs_val2 = _relative_sample_to_absolute(val2, length) - result[dim] = (abs_val1, abs_val2) - - return result - - -def _convert_relative_samples_arrays(patch, array_bounds): - """Convert relative sample indices in arrays to absolute indices.""" - result = {} - for dim, (arr1, arr2) in array_bounds.items(): - coord = patch.get_coord(dim) - length = len(coord) - - # Convert each boundary array - abs_arr1 = _relative_samples_array_to_absolute(arr1, length) - abs_arr2 = _relative_samples_array_to_absolute(arr2, length) - result[dim] = (abs_arr1, abs_arr2) - - return result - - -def _relative_sample_to_absolute(value, length): - """Convert relative sample index to absolute.""" - if value is None or value is ...: - return value - if value >= 0: - return int(value) - else: - return int(length + value) - - -def _relative_samples_array_to_absolute(arr, length): - """Convert array of relative sample indices to absolute.""" - if arr is None: - return None - - result = np.empty_like(arr, dtype=int) - for i, val in enumerate(arr): - if val >= 0: - result[i] = int(val) - else: - result[i] = int(length + val) - return result - - -def _convert_to_samples(patch, scalar_bounds): - """Convert scalar boundaries from coordinate values to sample indices.""" - result = {} - for dim, (val1, val2) in scalar_bounds.items(): - coord = patch.get_coord(dim) - # Convert to indices - idx1 = _value_to_index(coord, val1) if val1 is not None else None - idx2 = _value_to_index(coord, val2) if val2 is not None else None - result[dim] = (idx1, idx2) - return result - - -def _convert_arrays_to_samples(patch, array_bounds): - """Convert array boundaries from coordinate values to sample indices.""" - result = {} - for dim, (arr1, arr2) in array_bounds.items(): - coord = patch.get_coord(dim) - idx_arr1 = _array_to_indices(coord, arr1) if arr1 is not None else None - idx_arr2 = _array_to_indices(coord, arr2) if arr2 is not None else None - result[dim] = (idx_arr1, idx_arr2) - return result - - -def _value_to_index(coord, value): - """Convert coordinate value to index.""" - if value is None or value is ...: - return None - return coord.get_next_index(value) - - -def _array_to_indices(coord, arr): - """Convert array of coordinate values to indices.""" - if arr is None: - return None - return np.array([coord.get_next_index(val) for val in arr]) - - -def _parse_taper_value(taper_val, coord, dim_name): - """ - Parse a taper value into absolute coordinate units. - - Parameters - ---------- - taper_val - Can be float (fraction), Quantity (absolute), or None - coord - The coordinate object for this dimension - dim_name - Name of the dimension (for error messages) - - Returns - ------- - float or timedelta - Taper width in coordinate units, or None if taper_val is None - """ - if taper_val is None: - return None - - # Check if it's a Quantity with units - if isinstance(taper_val, Quantity): - # For time quantities, keep as timedelta; for others convert to float - coord_range = coord.max() - coord.min() - if isinstance(coord_range, np.timedelta64): - # Convert Quantity to timedelta64 - # First convert to seconds, then to nanoseconds for timedelta64 - seconds = taper_val.to('s').magnitude - return np.timedelta64(int(seconds * 1e9), 'ns') - else: - # Convert to float (handles distance, etc.) - return float(taper_val.magnitude) - - # Otherwise it's a float - treat as fraction of dimension range - if not isinstance(taper_val, (int, float)): - msg = ( - f"Taper value for dimension '{dim_name}' must be a float " - f"(fraction) or Quantity (with units). Got {type(taper_val)}" - ) - raise ParameterError(msg) - - # Get dimension range and multiply by fraction - coord_range = coord.max() - coord.min() - if isinstance(coord_range, np.timedelta64): - # For datetime coords, multiply timedelta by fraction - return coord_range * taper_val - else: - return taper_val * coord_range - - -def _validate_taper_dict(taper, scalar_bounds): - """ - Validate that taper dict has valid dimension names. - - Parameters - ---------- - taper : dict - Dictionary of dimension names to taper values - scalar_bounds : dict - Dictionary of scalar boundaries being used - """ - if not isinstance(taper, dict): - return - - for dim in taper: - if dim not in scalar_bounds: - valid_dims = sorted(scalar_bounds.keys()) - msg = ( - f"Taper dimension '{dim}' not found in mute dimensions. " - f"Valid dimensions: {valid_dims}" - ) - raise ParameterError(msg) - - -def _get_block_mute_envelope(patch, scalar_bounds, taper, window_type): - """ - Create envelope for rectangular block mutes. - - Similar to taper_range but creates a mute envelope instead. - """ - # Validate taper if it's a dict - if isinstance(taper, dict): - _validate_taper_dict(taper, scalar_bounds) - - # Check if taper is a Quantity with multiple dimensions - if isinstance(taper, Quantity) and len(scalar_bounds) > 1: - msg = ( - "Cannot use Quantity (with units) for taper when multiple dimensions " - "are specified. Use a dict like: taper={'time': 0.02*dc.units.s, " - "'distance': 10*dc.units.m}" - ) - raise ParameterError(msg) - - envelope = np.ones(patch.shape, dtype=float) - - for dim, (val1, val2) in scalar_bounds.items(): - coord = patch.get_coord(dim) - axis = patch.get_axis(dim) - - # Values are already sample indices at this point - # Handle None values (edges) - if val1 is None or val1 is ...: - val1 = 0 - if val2 is None or val2 is ...: - val2 = len(coord) - 1 - - # Create slice object from indices - slice_obj = slice(int(val1), int(val2) + 1) - - # Create 1D envelope for this dimension - dim_envelope = np.ones(len(coord)) - dim_envelope[slice_obj] = 0 # Zero out mute region - - # Apply taper if specified and not callable - if taper is not None and not callable(taper): - # Get taper value for this dimension - if isinstance(taper, dict): - taper_val = taper.get(dim) - else: - taper_val = taper - - if taper_val is not None: - # Parse taper value (handles fraction vs Quantity) - taper_width = _parse_taper_value(taper_val, coord, dim) - dim_envelope = _apply_taper_1d( - dim_envelope, slice_obj, taper_width, coord, window_type - ) - - # Broadcast to full shape and multiply - indexer = broadcast_for_index( - patch.ndim, axis, value=slice(None), fill=None - ) - envelope *= dim_envelope[indexer] - - return envelope - - -def _apply_taper_1d(envelope, slice_obj, taper_width, coord, window_type): - """Apply taper to 1D envelope at boundaries.""" - func = _get_window_function(window_type) - - # Get indices of mute region - start_idx = slice_obj.start if slice_obj.start is not None else 0 - stop_idx = slice_obj.stop if slice_obj.stop is not None else len(coord) - - # Convert taper width to samples - if hasattr(coord, 'step') and coord.step is not None: - # Handle both timedelta and numeric step/taper_width - if isinstance(taper_width, np.timedelta64): - # Both are timedeltas - can divide directly - taper_samples = int(taper_width / coord.step) - elif isinstance(coord.step, np.timedelta64): - # taper_width is numeric, step is timedelta - shouldn't happen - # Convert timedelta step to float for division - step_float = to_float(coord.step) - taper_samples = int(taper_width / step_float) - else: - # Both numeric - taper_samples = int(taper_width / coord.step) - else: - # For non-evenly sampled coords, estimate - coord_range = coord.max() - coord.min() - if isinstance(taper_width, np.timedelta64) and isinstance(coord_range, np.timedelta64): - # Both timedeltas - taper_samples = int(len(coord) * (taper_width / coord_range)) - elif isinstance(taper_width, np.timedelta64): - # Convert timedelta to float - taper_float = to_float(taper_width) - range_float = to_float(coord_range) - taper_samples = int(len(coord) * taper_float / range_float) - elif isinstance(coord_range, np.timedelta64): - # Convert coord_range to float - range_float = to_float(coord_range) - taper_samples = int(len(coord) * taper_width / range_float) - else: - # Both numeric - taper_samples = int(len(coord) * taper_width / coord_range) - - taper_samples = max(1, min(taper_samples, (stop_idx - start_idx) // 2)) - - # Apply taper at start of mute region - if taper_samples > 0 and start_idx > 0: - window = func(2 * taper_samples)[:taper_samples] - taper_start = max(0, start_idx - taper_samples) - envelope[taper_start:start_idx] *= window - - # Apply taper at end of mute region - if taper_samples > 0 and stop_idx < len(coord): - window = func(2 * taper_samples)[taper_samples:] - taper_end = min(len(coord), stop_idx + taper_samples) - envelope[stop_idx:taper_end] *= window[:taper_end - stop_idx] - - return envelope - - -def _get_geometric_mute_envelope(patch, array_bounds, taper, window_type): - """ - Create envelope for geometric (line/plane/hyperplane) mutes. - - This is Phase 2 - to be implemented for linear/planar mutes. - For now, raise an informative error. - """ - msg = ( - "Geometric mutes (using arrays to define lines/planes) are not yet " - "implemented. Use scalar boundaries for block mutes, or stay tuned " - "for the next release!" - ) - raise NotImplementedError(msg) + breakpoint() diff --git a/dascore/proc/taper.py b/dascore/proc/taper.py index 78933f4f7..0ebae3d75 100644 --- a/dascore/proc/taper.py +++ b/dascore/proc/taper.py @@ -153,7 +153,9 @@ def _get_taper_coord_inds(coord, values, relative, samples): # None or ... means min_val in first half of list else max_val out[num] = 0 if (num / len(out)) < 0.5 else len(coord) - 1 else: - out[num] = coord.get_next_index(val, samples=samples, relative=relative) + out[num] = coord.get_next_index( + val, samples=samples, relative=relative, allow_out_of_bounds=True + ) # Always need a len 4 sequence if len(out) == 2: out = [0, *out, len(coord)] diff --git a/dascore/utils/signal.py b/dascore/utils/signal.py index 3a4375329..0d2bbddfe 100644 --- a/dascore/utils/signal.py +++ b/dascore/utils/signal.py @@ -19,6 +19,7 @@ parzen=windows.parzen, triang=windows.triang, ramp=windows.triang, + boxcar=windows.boxcar, ) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 917dd3b28..2affac1d8 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -9,6 +9,34 @@ from dascore.exceptions import ParameterError +def _get_testable_coord_values(coord, relative=False, samples=False): + """Get some values in the coordinate for testing different ranges.""" + start_ind, stop_ind = len(coord) // 4, 3 * len(coord) // 4 + if samples: + return (start_ind, stop_ind) + start, stop = coord[start_ind], coord[stop_ind] + if relative: + start, stop = start - coord.min(), stop - coord.min() + return (start, stop) + + +def _assert_coord_ranges( + patch, + dim, + zero_ranges, + one_ranges, + relative=True, + samples=False, +): + """Assert that the expected values occur in the patch.""" + for zrange in zero_ranges: + sub = patch.select(**{dim: zrange}, relative=relative, samples=samples) + assert np.allclose(sub.data, 0) + for orange in one_ranges: + sub = patch.select(**{dim: orange}, relative=relative, samples=samples) + assert np.allclose(sub.data, 1) + + @pytest.fixture(scope="session") def patch_ones(random_patch): """Return a patch filled with ones.""" @@ -24,413 +52,422 @@ def ricker_patch(): class TestMuteBasics: """Basic tests for mute functionality.""" - def test_mute_exists(self, random_patch): - """Ensure mute method exists on patch.""" - assert hasattr(random_patch, "mute") - assert callable(random_patch.mute) - def test_mute_no_kwargs_raises(self, random_patch): """Mute without dimension specifications should raise.""" - with pytest.raises(ParameterError, match="At least one dimension"): + with pytest.raises(ParameterError, match="one or two keyword"): random_patch.mute() - def test_invalid_dimension_raises(self, random_patch): - """Mute with invalid dimension should raise.""" - with pytest.raises(ParameterError, match="not found"): - random_patch.mute(invalid_dim=(0, 10)) - def test_not_tuple_raises(self, random_patch): """Boundary must be tuple of length 2.""" - with pytest.raises(ParameterError, match="must specify 2 boundaries"): + with pytest.raises(ParameterError, match="two boundaries when using"): random_patch.mute(time=5) def test_tuple_wrong_length_raises(self, random_patch): """Tuple must have exactly 2 elements.""" - with pytest.raises(ParameterError, match="must specify 2 boundaries"): + with pytest.raises(ParameterError, match="two boundaries when using"): random_patch.mute(time=(1, 2, 3)) -class TestBlockMutes1D: +class Test1DMute: """Test 1D block mutes (single dimension).""" - def test_mute_start_of_time(self, patch_ones): + def test_1d_mute_no_taper(self, patch_ones): """Mute first portion of time dimension.""" # Use relative=True (default) with numeric offset - muted = patch_ones.mute(time=(0, 0.5)) - - # Check that first samples are zero - assert np.allclose(muted.data[:, 0], 0) - # Check that later samples are still 1 - assert np.allclose(muted.data[:, -1], 1) - - def test_mute_end_of_time(self, patch_ones): - """Mute last portion of time dimension.""" - # Use relative with negative values (from end) - # (-0.5, None) means from 0.5s before end to the end - muted = patch_ones.mute(time=(-0.5, None)) - - # Check that last samples are zero - assert np.allclose(muted.data[:, -1], 0) - # Check that early samples are still 1 - assert np.allclose(muted.data[:, 0], 1) - - def test_mute_middle_of_distance(self, patch_ones): - """Mute middle portion of distance dimension.""" - dist = patch_ones.coords.get_array("distance") - mid_start = dist[len(dist) // 4] - mid_end = dist[3 * len(dist) // 4] - - muted = patch_ones.mute(distance=(mid_start, mid_end), relative=False) - - # Check middle is muted - mid_idx = len(dist) // 2 - assert np.allclose(muted.data[mid_idx, :], 0) - # Check edges are not muted - assert np.allclose(muted.data[0, :], 1) - assert np.allclose(muted.data[-1, :], 1) - - def test_mute_relative_default(self, patch_ones): - """Test that relative=True is default.""" - # Should mute first second (relative to start) - muted = patch_ones.mute(time=(0, 1)) - # First samples should be zero - assert muted.data[:, 0].max() < 0.1 - - def test_mute_relative_negative(self, patch_ones): - """Test relative with negative values (from end).""" - muted = patch_ones.mute(time=(-1, None)) - # Last samples should be zero - assert muted.data[:, -1].max() < 0.1 - - -class TestBlockMutes2D: - """Test 2D block mutes (multiple dimensions).""" - - def test_mute_rectangular_region(self, patch_ones): - """Mute a rectangular region in 2D.""" - time_coords = patch_ones.coords.get_array("time") - dist_coords = patch_ones.coords.get_array("distance") - - t1, t2 = time_coords[10], time_coords[20] - d1, d2 = dist_coords[5], dist_coords[15] - - muted = patch_ones.mute( - time=(t1, t2), - distance=(d1, d2), - relative=False, - ) - - # Check interior is muted - assert np.allclose(muted.data[10, 15], 0) - # Check corners are not muted - assert np.allclose(muted.data[0, 0], 1) - assert np.allclose(muted.data[-1, -1], 1) - - def test_mute_strips(self, patch_ones): - """Mute strips along each dimension.""" - # Mute a time strip - muted_time = patch_ones.mute(time=(0.5, 1.0)) - # All distances should be affected equally - assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) - - # Mute a distance strip - muted_dist = patch_ones.mute(distance=(10, 20)) - # All times should be affected equally - assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) - - -class TestMuteModes: - """Test different mute modes.""" - - def test_mode_union(self, patch_ones): - """Test union mode (default, mutes inside region).""" - muted = patch_ones.mute(time=(2.0, 6.0), mode="union") - # Middle (at ~4s) should be zero - mid_idx = len(patch_ones.coords.get_array("time")) // 2 - assert muted.data[:, mid_idx].max() < 0.1 - # Edges should be one - assert muted.data[:, 0].min() > 0.9 - - def test_mode_complement(self, patch_ones): - """Test complement mode (mutes outside region).""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[10], time_coords[20] - - muted = patch_ones.mute(time=(t1, t2), mode="complement", relative=False) - - # Edges should be zero - assert muted.data[:, 0].max() < 0.1 - assert muted.data[:, -1].max() < 0.1 - # Middle should be one - mid_idx = 15 - assert muted.data[:, mid_idx].min() > 0.9 - - -class TestMuteTaper: - """Test taper application in mutes.""" - - def test_mute_with_taper(self, patch_ones): - """Mute with taper creates gradual transition.""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - # Mute without taper - muted_no_taper = patch_ones.mute( - time=(t1, t2), - relative=False, - taper=None, + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted = patch_ones.mute(time=(v1, v2), relative=True) + _assert_coord_ranges( + patch=muted, + dim="time", + zero_ranges=[(v1, v2)], + one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], + relative=True, + samples=False, ) - # Mute with taper (5% of dimension range) - muted_with_taper = patch_ones.mute( - time=(t1, t2), - relative=False, - taper=0.05, + def test_mute_open_interval(self, patch_ones): + """Mute using None for interval ends.""" + coord = patch_ones.get_coord("distance") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted1 = patch_ones.mute(distance=(v2, ...), relative=True) + _assert_coord_ranges( + patch=muted1, + dim="distance", + zero_ranges=[(v2, ...)], + one_ranges=[(..., v2 - coord.step)], + relative=True, + samples=False, ) - # Without taper should be sharp transition - assert muted_no_taper.data[:, 19].min() > 0.9 - assert muted_no_taper.data[:, 20].max() < 0.1 - - # With taper should have gradual transition - # Values in taper region should be between 0 and 1 - taper_region = muted_with_taper.data[:, 15:20] - assert (taper_region > 0).any() - assert (taper_region < 1).any() - - def test_taper_dict(self, patch_ones): - """Test dimension-specific taper using dict.""" - time_coords = patch_ones.coords.get_array("time") - dist_coords = patch_ones.coords.get_array("distance") - - taper_dict = { - "time": 0.05, # 5% of time range - "distance": 0, # No taper on distance - } - - muted = patch_ones.mute( - time=(time_coords[20], time_coords[40]), - distance=(dist_coords[10], dist_coords[30]), + def test_mute_absolute(self, patch_ones): + """Test absolute coordinates work.""" + coord = patch_ones.get_coord("distance") + v1, v2 = _get_testable_coord_values(coord, relative=False) + muted1 = patch_ones.mute(distance=(v1, v2), relative=False) + _assert_coord_ranges( + patch=muted1, + dim="distance", + zero_ranges=[(v1, v2)], + one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], relative=False, - taper=taper_dict, + samples=False, ) - # Should have taper in time, sharp in distance - assert muted.shape == patch_ones.shape - - def test_taper_window_types(self, patch_ones): - """Test different window types for taper.""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - for window_type in ["hann", "hamming", "triang"]: - muted = patch_ones.mute( - time=(t1, t2), - relative=False, - taper=0.05, # 5% of dimension range - window_type=window_type, - ) - assert isinstance(muted, dc.Patch) - assert muted.shape == patch_ones.shape - - def test_taper_with_quantity(self, patch_ones): - """Test taper with Quantity (absolute units).""" - from dascore.units import s - - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - # Use absolute time value with units (single dimension) - muted = patch_ones.mute( - time=(t1, t2), - relative=False, - taper={'time': 0.05 * s}, + def test_mute_samples(self, patch_ones): + """Test that sample mute works.""" + coord = patch_ones.get_coord("distance") + v1, v2 = _get_testable_coord_values(coord, samples=True) + muted1 = patch_ones.mute(distance=(v1, v2), samples=True) + _assert_coord_ranges( + patch=muted1, + dim="distance", + zero_ranges=[(v1, v2)], + one_ranges=[(..., v1 - 1), (v2 + 1, ...)], + samples=True, ) - # Should have gradual transition - taper_region = muted.data[:, 15:20] - assert (taper_region > 0).any() - assert (taper_region < 1).any() - - def test_taper_quantity_multiple_dims_raises(self, patch_ones): - """Test that Quantity taper with multiple dims raises error.""" - from dascore.units import s - - time_coords = patch_ones.coords.get_array("time") - dist_coords = patch_ones.coords.get_array("distance") - - # Should raise error if Quantity used without dict - with pytest.raises(ParameterError, match="Cannot use Quantity"): - patch_ones.mute( - time=(time_coords[20], time_coords[40]), - distance=(dist_coords[10], dist_coords[30]), - relative=False, - taper=0.05 * s, - ) - - def test_taper_mixed_quantity_fraction(self, patch_ones): - """Test mixed taper: fraction for one dim, Quantity for another.""" - from dascore.units import s, m - - time_coords = patch_ones.coords.get_array("time") - dist_coords = patch_ones.coords.get_array("distance") - - muted = patch_ones.mute( - time=(time_coords[20], time_coords[40]), - distance=(dist_coords[10], dist_coords[30]), - relative=False, - taper={'time': 0.05, 'distance': 10 * m}, - ) - - assert muted.shape == patch_ones.shape - - def test_taper_custom_function(self, patch_ones): - """Test custom taper function.""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - # Define custom taper that squares the envelope - def custom_taper(envelope): - return envelope ** 2 - - muted = patch_ones.mute( - time=(t1, t2), - relative=False, - taper=custom_taper, - ) - - # Middle should still be zero - assert muted.data[:, 30].max() < 0.1 - # Edges should still be one - assert muted.data[:, 0].min() > 0.9 - - def test_taper_custom_function_invalid_return(self, patch_ones): - """Test that custom taper with invalid return raises error.""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - # Function that returns wrong type - def bad_taper(envelope): - return list(envelope) - - with pytest.raises(ParameterError, match="must return a numpy array"): - patch_ones.mute( - time=(t1, t2), - relative=False, - taper=bad_taper, - ) - - def test_taper_custom_function_wrong_shape(self, patch_ones): - """Test that custom taper with wrong shape raises error.""" - time_coords = patch_ones.coords.get_array("time") - t1, t2 = time_coords[20], time_coords[40] - - # Function that returns wrong shape - def bad_shape_taper(envelope): - return np.ones((10, 10)) - - with pytest.raises(ParameterError, match="must return array with shape"): - patch_ones.mute( - time=(t1, t2), - relative=False, - taper=bad_shape_taper, - ) - - -class TestMuteEdgeCases: - """Test edge cases and boundary conditions.""" - - def test_mute_with_none_boundaries(self, patch_ones): - """Test using None to reference coordinate edges.""" - # Mute from start to middle - time_coords = patch_ones.coords.get_array("time") - mid_time = time_coords[len(time_coords) // 2] - - muted = patch_ones.mute(time=(None, mid_time), relative=False) - - # First half should be zero - assert muted.data[:, 0].max() < 0.1 - # Second half should be one - assert muted.data[:, -1].min() > 0.9 - - def test_mute_entire_dimension(self, patch_ones): - """Mute entire dimension.""" - muted = patch_ones.mute(time=(None, None), relative=False) - # Everything should be zero - assert np.allclose(muted.data, 0) - - def test_mute_zero_width(self, patch_ones): - """Mute with same start and end values.""" - time_coords = patch_ones.coords.get_array("time") - t = time_coords[10] - - muted = patch_ones.mute(time=(t, t), relative=False) - # Should mute just that one sample (or very close) - assert muted.data[:, 10].max() < 0.1 - - -class TestMuteWithSamples: - """Test mute with samples=True.""" - - def test_mute_samples_mode(self, patch_ones): - """Test mute using sample indices.""" - # Mute samples 10 to 20 - muted = patch_ones.mute(time=(10, 20), samples=True) - - # Samples 10-20 should be zero - assert muted.data[:, 15].max() < 0.1 - # Other samples should be one - assert muted.data[:, 0].min() > 0.9 - assert muted.data[:, -1].min() > 0.9 - - -class TestGeometricMutes: - """Test geometric (array-based) mutes - Phase 2.""" - - def test_array_boundaries_not_implemented(self, ricker_patch): - """Array boundaries should raise error for mixing scalar/array.""" - # Currently we don't allow mixing scalars and arrays - with pytest.raises(ParameterError, match="Cannot mix scalar and array"): - ricker_patch.mute( - time=(0, [0, 0.3]), - distance=(0, [0, 300]), - relative=False, - ) - - def test_mismatched_array_lengths_raises(self, random_patch): - """Arrays with different lengths should raise error during parsing.""" - with pytest.raises(ParameterError, match="Cannot mix scalar and array"): - # This currently raises mixing error before length checking - random_patch.mute( - time=(0, [0, 0.3, 0.5]), # scalar and array - distance=(0, [0, 300]), # scalar and array - relative=False, - ) - - -class TestMutePreservesMetadata: - """Test that mute preserves patch metadata.""" - - def test_preserves_coords(self, random_patch): - """Mute should preserve coordinates.""" - muted = random_patch.mute(time=(0, 0.5)) - assert muted.coords == random_patch.coords - - def test_preserves_attrs(self, random_patch): - """Mute should preserve attributes (except history).""" - muted = random_patch.mute(time=(0, 0.5)) - # History will be different due to processing - assert muted.attrs.data_type == random_patch.attrs.data_type - assert muted.attrs.network == random_patch.attrs.network - assert muted.attrs.station == random_patch.attrs.station - - def test_preserves_shape(self, random_patch): - """Mute should preserve shape.""" - muted = random_patch.mute(time=(0, 0.5)) - assert muted.shape == random_patch.shape - - def test_preserves_dtype(self, random_patch): - """Mute should preserve data type.""" - muted = random_patch.mute(time=(0, 0.5)) - # Allow for float conversion - assert muted.dtype in (random_patch.dtype, np.float64, np.float32) + def test_mute_taper(self, patch_ones): + """Test that taper mute works.""" + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted1 = patch_ones.mute(time=(v1, v2), relative=True, taper=0.1) + # With taper, we can't assert exact zeros/ones, just check shape + assert muted1.shape == patch_ones.shape + # Check that middle region has been modified (not all ones) + sub = muted1.select(time=(v1, v2), relative=True) + assert not np.allclose(sub.data, 1) + + +# +# class TestMute2D: +# """Test 2D block mutes (multiple dimensions).""" +# +# def test_mute_rectangular_region(self, patch_ones): +# """Mute a rectangular region in 2D.""" +# time_coords = patch_ones.coords.get_array("time") +# dist_coords = patch_ones.coords.get_array("distance") +# +# t1, t2 = time_coords[10], time_coords[20] +# d1, d2 = dist_coords[5], dist_coords[15] +# +# muted = patch_ones.mute( +# time=(t1, t2), +# distance=(d1, d2), +# relative=False, +# ) +# +# # Check interior is muted +# assert np.allclose(muted.data[10, 15], 0) +# # Check corners are not muted +# assert np.allclose(muted.data[0, 0], 1) +# assert np.allclose(muted.data[-1, -1], 1) +# +# def test_mute_strips(self, patch_ones): +# """Mute strips along each dimension.""" +# # Mute a time strip +# muted_time = patch_ones.mute(time=(0.5, 1.0)) +# # All distances should be affected equally +# assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) +# +# # Mute a distance strip +# muted_dist = patch_ones.mute(distance=(10, 20)) +# # All times should be affected equally +# assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) +# +# +# class TestMuteModes: +# """Test different mute modes.""" +# +# def test_mode_union(self, patch_ones): +# """Test union mode (default, mutes inside region).""" +# muted = patch_ones.mute(time=(2.0, 6.0), mode="union") +# # Middle (at ~4s) should be zero +# mid_idx = len(patch_ones.coords.get_array("time")) // 2 +# assert muted.data[:, mid_idx].max() < 0.1 +# # Edges should be one +# assert muted.data[:, 0].min() > 0.9 +# +# def test_mode_complement(self, patch_ones): +# """Test complement mode (mutes outside region).""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[10], time_coords[20] +# +# muted = patch_ones.mute(time=(t1, t2), mode="complement", relative=False) +# +# # Edges should be zero +# assert muted.data[:, 0].max() < 0.1 +# assert muted.data[:, -1].max() < 0.1 +# # Middle should be one +# mid_idx = 15 +# assert muted.data[:, mid_idx].min() > 0.9 +# +# +# class TestMuteTaper: +# """Test taper application in mutes.""" +# +# def test_mute_with_taper(self, patch_ones): +# """Mute with taper creates gradual transition.""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# # Mute without taper +# muted_no_taper = patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=None, +# ) +# +# # Mute with taper (5% of dimension range) +# muted_with_taper = patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=0.05, +# ) +# +# # Without taper should be sharp transition +# assert muted_no_taper.data[:, 19].min() > 0.9 +# assert muted_no_taper.data[:, 20].max() < 0.1 +# +# # With taper should have gradual transition +# # Values in taper region should be between 0 and 1 +# taper_region = muted_with_taper.data[:, 15:20] +# assert (taper_region > 0).any() +# assert (taper_region < 1).any() +# +# def test_taper_dict(self, patch_ones): +# """Test dimension-specific taper using dict.""" +# time_coords = patch_ones.coords.get_array("time") +# dist_coords = patch_ones.coords.get_array("distance") +# +# taper_dict = { +# "time": 0.05, # 5% of time range +# "distance": 0, # No taper on distance +# } +# +# muted = patch_ones.mute( +# time=(time_coords[20], time_coords[40]), +# distance=(dist_coords[10], dist_coords[30]), +# relative=False, +# taper=taper_dict, +# ) +# +# # Should have taper in time, sharp in distance +# assert muted.shape == patch_ones.shape +# +# def test_taper_window_types(self, patch_ones): +# """Test different window types for taper.""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# for window_type in ["hann", "hamming", "triang"]: +# muted = patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=0.05, # 5% of dimension range +# window_type=window_type, +# ) +# assert isinstance(muted, dc.Patch) +# assert muted.shape == patch_ones.shape +# +# def test_taper_with_quantity(self, patch_ones): +# """Test taper with Quantity (absolute units).""" +# from dascore.units import s +# +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# # Use absolute time value with units (single dimension) +# muted = patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper={'time': 0.05 * s}, +# ) +# +# # Should have gradual transition +# taper_region = muted.data[:, 15:20] +# assert (taper_region > 0).any() +# assert (taper_region < 1).any() +# +# def test_taper_quantity_multiple_dims_raises(self, patch_ones): +# """Test that Quantity taper with multiple dims raises error.""" +# from dascore.units import s +# +# time_coords = patch_ones.coords.get_array("time") +# dist_coords = patch_ones.coords.get_array("distance") +# +# # Should raise error if Quantity used without dict +# with pytest.raises(ParameterError, match="Cannot use Quantity"): +# patch_ones.mute( +# time=(time_coords[20], time_coords[40]), +# distance=(dist_coords[10], dist_coords[30]), +# relative=False, +# taper=0.05 * s, +# ) +# +# def test_taper_mixed_quantity_fraction(self, patch_ones): +# """Test mixed taper: fraction for one dim, Quantity for another.""" +# from dascore.units import s, m +# +# time_coords = patch_ones.coords.get_array("time") +# dist_coords = patch_ones.coords.get_array("distance") +# +# muted = patch_ones.mute( +# time=(time_coords[20], time_coords[40]), +# distance=(dist_coords[10], dist_coords[30]), +# relative=False, +# taper={'time': 0.05, 'distance': 10 * m}, +# ) +# +# assert muted.shape == patch_ones.shape +# +# def test_taper_custom_function(self, patch_ones): +# """Test custom taper function.""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# # Define custom taper that squares the envelope +# def custom_taper(envelope): +# return envelope ** 2 +# +# muted = patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=custom_taper, +# ) +# +# # Middle should still be zero +# assert muted.data[:, 30].max() < 0.1 +# # Edges should still be one +# assert muted.data[:, 0].min() > 0.9 +# +# def test_taper_custom_function_invalid_return(self, patch_ones): +# """Test that custom taper with invalid return raises error.""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# # Function that returns wrong type +# def bad_taper(envelope): +# return list(envelope) +# +# with pytest.raises(ParameterError, match="must return a numpy array"): +# patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=bad_taper, +# ) +# +# def test_taper_custom_function_wrong_shape(self, patch_ones): +# """Test that custom taper with wrong shape raises error.""" +# time_coords = patch_ones.coords.get_array("time") +# t1, t2 = time_coords[20], time_coords[40] +# +# # Function that returns wrong shape +# def bad_shape_taper(envelope): +# return np.ones((10, 10)) +# +# with pytest.raises(ParameterError, match="must return array with shape"): +# patch_ones.mute( +# time=(t1, t2), +# relative=False, +# taper=bad_shape_taper, +# ) +# +# +# class TestMuteEdgeCases: +# """Test edge cases and boundary conditions.""" +# +# def test_mute_with_none_boundaries(self, patch_ones): +# """Test using None to reference coordinate edges.""" +# # Mute from start to middle +# time_coords = patch_ones.coords.get_array("time") +# mid_time = time_coords[len(time_coords) // 2] +# +# muted = patch_ones.mute(time=(None, mid_time), relative=False) +# +# # First half should be zero +# assert muted.data[:, 0].max() < 0.1 +# # Second half should be one +# assert muted.data[:, -1].min() > 0.9 +# +# def test_mute_entire_dimension(self, patch_ones): +# """Mute entire dimension.""" +# muted = patch_ones.mute(time=(None, None), relative=False) +# # Everything should be zero +# assert np.allclose(muted.data, 0) +# +# def test_mute_zero_width(self, patch_ones): +# """Mute with same start and end values.""" +# time_coords = patch_ones.coords.get_array("time") +# t = time_coords[10] +# +# muted = patch_ones.mute(time=(t, t), relative=False) +# # Should mute just that one sample (or very close) +# assert muted.data[:, 10].max() < 0.1 +# +# +# class TestMuteWithSamples: +# """Test mute with samples=True.""" +# +# def test_mute_samples_mode(self, patch_ones): +# """Test mute using sample indices.""" +# # Mute samples 10 to 20 +# muted = patch_ones.mute(time=(10, 20), samples=True) +# +# # Samples 10-20 should be zero +# assert muted.data[:, 15].max() < 0.1 +# # Other samples should be one +# assert muted.data[:, 0].min() > 0.9 +# assert muted.data[:, -1].min() > 0.9 +# +# +# class TestGeometricMutes: +# """Test geometric (array-based) mutes - Phase 2.""" +# +# def test_array_boundaries_not_implemented(self, ricker_patch): +# """Array boundaries should raise error for mixing scalar/array.""" +# # Currently we don't allow mixing scalars and arrays +# with pytest.raises(ParameterError, match="Cannot mix scalar and array"): +# ricker_patch.mute( +# time=(0, [0, 0.3]), +# distance=(0, [0, 300]), +# relative=False, +# ) +# +# def test_mismatched_array_lengths_raises(self, random_patch): +# """Arrays with different lengths should raise error during parsing.""" +# with pytest.raises(ParameterError, match="Cannot mix scalar and array"): +# # This currently raises mixing error before length checking +# random_patch.mute( +# time=(0, [0, 0.3, 0.5]), # scalar and array +# distance=(0, [0, 300]), # scalar and array +# relative=False, +# ) +# +# +# class TestMutePreservesMetadata: +# """Test that mute preserves patch metadata.""" +# +# def test_preserves_coords(self, random_patch): +# """Mute should preserve coordinates.""" +# muted = random_patch.mute(time=(0, 0.5)) +# assert muted.coords == random_patch.coords +# +# def test_preserves_attrs(self, random_patch): +# """Mute should preserve attributes (except history).""" +# muted = random_patch.mute(time=(0, 0.5)) +# # History will be different due to processing +# assert muted.attrs.data_type == random_patch.attrs.data_type +# assert muted.attrs.network == random_patch.attrs.network +# assert muted.attrs.station == random_patch.attrs.station +# +# def test_preserves_shape(self, random_patch): +# """Mute should preserve shape.""" +# muted = random_patch.mute(time=(0, 0.5)) +# assert muted.shape == random_patch.shape +# +# def test_preserves_dtype(self, random_patch): +# """Mute should preserve data type.""" +# muted = random_patch.mute(time=(0, 0.5)) +# # Allow for float conversion +# assert muted.dtype in (random_patch.dtype, np.float64, np.float32) diff --git a/tests/test_proc/test_taper.py b/tests/test_proc/test_taper.py index 80b887634..600d28e0b 100644 --- a/tests/test_proc/test_taper.py +++ b/tests/test_proc/test_taper.py @@ -26,6 +26,8 @@ def patch_ones(random_patch): @pytest.fixture(scope="session", params=sorted(WINDOW_FUNCTIONS)) def time_tapered_patch(request, patch_ones): """Return a tapered trace.""" + if "boxcar" in str(request.param): + pytest.skip("boxcar doesn't actually apply taper.") # first get a patch with all ones for easy testing patch = patch_ones.update(data=np.ones_like(patch_ones.data)) out = taper(patch, time=0.05, window_type=request.param) @@ -274,3 +276,36 @@ def test_non_cosine_window(self, patch_ones): assert np.allclose(out.select(distance=(55, 90)).data, 1) assert np.allclose(out.select(distance=(126, 149)).data, 0) assert np.allclose(out.select(distance=(225, ...)).data, 0) + + def test_boxcar_window(self, patch_ones): + """The boxcar window is actually no taper.""" + # With no invert, all values are 1. + out1 = patch_ones.taper_range( + distance=(10, 12), invert=False, samples=True, window_type="boxcar" + ) + # With invert, all values are 0. + out2 = patch_ones.taper_range( + distance=(10, 12), invert=True, samples=True, window_type="boxcar" + ) + assert np.allclose(out1.data, 1.0) + assert np.allclose(out2.data, 0.0) + + def test_two_value_invert_mutes_to_last_sample(self, patch_ones): + """Ensure 2-value form with invert can mute from middle to the very end.""" + coord = patch_ones.get_coord("distance") + start_idx = len(coord) // 4 # 75 + end_idx = len(coord) - 1 # 299 (last valid index) + + # Use 2-value form with invert to mute from start_idx to end + out = patch_ones.taper_range( + distance=(start_idx, end_idx), + invert=False, + samples=True, + window_type="boxcar", + ) + # Values before start_idx should be 1 (unmuted) + assert np.allclose(out.data[start_idx - 1, :], 1) + + # Values from start_idx to end should be 0 (muted) + assert np.allclose(out.data[start_idx, :], 0) + assert np.allclose(out.data[end_idx, :], 0) From fba872db5dbf4cee6a8f77b0924de789eb4381be Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Fri, 17 Oct 2025 17:11:00 +0100 Subject: [PATCH 03/15] work on mute --- dascore/proc/mute.py | 320 ++++++++++++++----------- tests/test_proc/test_mute.py | 439 +++++++---------------------------- 2 files changed, 265 insertions(+), 494 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 20e5cc25e..1067022cf 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -2,20 +2,103 @@ from __future__ import annotations -from collections.abc import Mapping +import dataclasses +from collections import defaultdict +from collections.abc import Mapping, Sized import numpy as np +from scipy.ndimage import gaussian_filter +import dascore as dc from dascore.constants import PatchType from dascore.exceptions import ParameterError +from dascore.utils.misc import broadcast_for_index from dascore.utils.patch import get_dim_axis_value, patch_function -def _get_mute_params(patch, kwargs): +@dataclasses.dataclass +class _OriginNLines: + """Get a tuple of origin, line1, line2.""" + + origin: np.ndarray + line1: np.ndarray + line2: np.ndarray + + @classmethod + def _check_degenerate(self, line): + """Ensure a line isn't degenerate else raise.""" + if len(np.unique(line, axis=0)) != 2: + msg = f"Line specified for mute {line} is degenerate!" + raise ParameterError(msg) + + @classmethod + def from_point_list(cls, points): + """""" + cls._check_degenerate(points[:2]) + cls._check_degenerate(points[2:]) + + # We need to ensure one of the points is duplicated; this is the orign + unique, idx, counts = np.unique( + points, axis=0, return_index=True, return_counts=True + ) + if len(unique) == len(points): + msg = ( + f"No common point found in {points}. For unambiguous mute " + f"lines must share an origin." + ) + raise ParameterError(msg) + gt_1_counts = counts > 1 + origin = points[idx[gt_1_counts]] + others = points[idx[~gt_1_counts]] + l1, l2 = others + return cls(origin=origin, line1=l1, line2=l2) + + +def _get_2d_mute_lines(vals, patch, dims, relative=True): + """ + Return values which are two lines in absolute coordinate space, + (expressed as floats) + """ + + def _get_coord_float_values(coord, vals, relative): + """Get the coordinate float values relative to start of coords.""" + out = dc.to_float(np.array(vals)) + if not relative: + out = out - dc.to_float(coord.min()) + return out + + out = [] + # Keep track of the axis that need to be filled in (eg None, paired w/ float) + fill_inds = defaultdict(list) + for ind, (dim, row) in enumerate(zip(dims, vals)): + coord = patch.get_coord(dim) + # In the case of single values (eg None, ..., or some other) + # We need to just mark them and come back later. + if not isinstance(row, Sized): + # We need to get the actual value from the coord. + if row is not None and row is not Ellipsis: + row = _get_coord_float_values(coord, vals, relative=relative) + fill_inds[dim].append(row) + out.append(None) + # Otherwise, we just run with it. + vals = _get_coord_float_values(coord, np.array(row), relative=relative) + out.append(vals) + # Now we can ascertain the line intended by None. + if fill_inds: + breakpoint() + # Then put the 4 points together. This gives us a len 4 array with rows + # as points from first line, then points from second. + points = np.stack(out, axis=-1).reshape(-1, 2) + origin_n_lines = _OriginNLines.from_point_list(points) + breakpoint() + + +def _get_mute_params(patch, kwargs, relative=True): """ Get (and validate) the muting parameters. - Returns arrays of dims, axes, and values. + Returns arrays of dims, axes, and values. Values has different meaning + based on the number of dimensions. """ # Validate we have dimension specifications if not kwargs or len(kwargs) > 2: @@ -27,111 +110,79 @@ def _get_mute_params(patch, kwargs): dim_ax_vals = get_dim_axis_value(patch, kwargs=kwargs, allow_multiple=True) dims = [x[0] for x in dim_ax_vals] axes = np.array([x[1] for x in dim_ax_vals]) - # Handle single dimension Mute. + val_list = [x[2] for x in dim_ax_vals] if len(dims) == 1: - vals = np.array([x[2] for x in dim_ax_vals]) - if vals.size != 2: + vals = np.array(val_list) + if vals.size != 2: # Handle single dimension Mute. msg = "Mute requires two boundaries when using a single dimension." raise ParameterError(msg) - else: - breakpoint() - print() + elif len(dims) > 1: # Dealing with lines. . + vals = _get_2d_mute_lines(val_list, patch, dims, relative) return dims, axes, vals -def _get_taper_vals(dims, taper, patch, samples): - """Get the taper values ordered by dimension, in samples.""" - # Just broadcast taper to same length as dims. - if taper is None: - return [0] * len(dims) - if not isinstance(taper, Mapping): - vals = [taper] * len(dims) - else: - # Otherwise each dimension's taper must be specified. - if not set(dims) == set(taper): - msg = ( - f"If a taper dictionary is used in Mute, it must have all " - f"the same keys as the dimensions. Kwarg dims are {dims} and" - f"taper keys are {list(taper)}." - ) - raise ParameterError(msg) - vals = [taper[dim] for dim in dims] - out = [] - for dim, val in zip(dims, vals): - coord = patch.get_coord(dim) - if val is None: - out.append(0) - elif samples: - out.append(val) +def _get_smooth_sigma(dims, smooth, patch): + """Get sigma values, in samples, for the gaussian kernel.""" + + def _broadcast_smooth_to_dims(dims, smooth): + """Broadcast smooth samples to dims length.""" + # First, we need to get the smooth parameters to line up with the dims. + # For a single value, just broadcast to dim length. + if not isinstance(smooth, Mapping): + vals = [smooth] * len(dims) else: + # Otherwise each dimension's smooth must be specified. + if not set(dims) == set(smooth): + msg = ( + f"If a taper dictionary is used in Mute, it must have all " + f"the same keys as the dimensions. Kwarg dims are {dims} and" + f"taper keys are {list(smooth)}." + ) + raise ParameterError(msg) + vals = [smooth[dim] for dim in dims] + return vals + + def _convert_to_samples(smooth, dims, patch): + """Convert the smooth parameter to number of samples.""" + out = [] + for dim, val in zip(dims, smooth): coord = patch.get_coord(dim) - coord.coord_range(val) - out.append(val) - return np.array(out) - - -def _mute_patch_1d( - patch, - dim, - vals, - taper_samps, - window, - samples, - relative, - invert, -): - """Apply mute to 1D patch using patch.range_taper.""" - coord = patch.get_coord(dim) - - if not samples: - # Get range represented by values. This is a bit ugly... - sel = coord.select(tuple(*vals), relative=relative)[1] - start = sel.start if sel.start is not None else 0 - # stop can be None (open interval) or an exclusive upper bound - if sel.stop is None: - stop = len(coord) - else: - stop = sel.stop - trange = (start, stop) - else: - # vals is a 2D array with shape (1, 2), flatten to get (start, stop) - trange = vals.flatten() if vals.ndim > 1 else vals - - # Handle edge case: if muting to the end with no taper, use 2-value form - # because 4-value form can't express "mute to the very last sample" - max_idx = len(coord) - 1 - if taper_samps == 0 and trange[1] >= len(coord): - taper_dict = {dim: [trange[0], max_idx]} - else: - # Clamp taper boundaries to valid sample indices. - # taper_range uses exclusive upper bounds in 4-value form [a,b,c,d] - # which zeros indices [b, c). - taper_dict = { - dim: [ - max(0, trange[0] - taper_samps), - trange[0], - min(max_idx, trange[1]), - min(max_idx, trange[1] + taper_samps), - ] - } - out = patch.taper_range( - invert=not invert, - samples=True, - window_type=window, - **taper_dict, - ) - return out + if val is None: + out.append(0) + elif isinstance(val, int | np.integer): + out.append(val) + elif isinstance(val, float | np.floating): + if not 0 <= val <= 1: + msg = ( + f"Mute's smooth parameter for {dim} must be between 0 " + f"and 1 when using a floating point value." + ) + raise ParameterError(msg) + out.append(int(np.round(len(coord) * val))) + else: # should capture quantities + out.append(coord.get_sample_count(val)) + return out + + smooth_by_dims = _broadcast_smooth_to_dims(dims, smooth) + smooth_ints = _convert_to_samples(smooth_by_dims, dims, patch) + # Now finagle into input for scipy's gaussian smooth. + return smooth_ints + + +def _line_smooth(data, dims, axes, values, patch): + """ + Mute data between two lines. + """ + breakpoint() @patch_function() def mute( patch: PatchType, *, - taper: float | dict | None = None, - window_type: str = "hann", + smooth=None, invert: bool = False, relative: bool = True, - samples: bool = False, **kwargs, ) -> PatchType: """ @@ -144,31 +195,24 @@ def mute( ---------- patch The patch instance. - taper - Taper width at mute boundaries. Can be: - - None: sharp mute (no taper) - - float (0.0-1.0): fraction of dimension range (e.g., 0.05 = 5%) - The number of samples in the mute is held constant across dimensions - by calculating the samples the fraction represents for each dimension - and using the minimum sample count. This helps avoid distorted mutes - based on dimensions with vastly different lengths. - - Quantity with units: absolute value (e.g., 0.02*dc.units.s) - Note: If multiple dimensions specified, must use dict. - - dict: (dim: taper_value) for dimension-specific taper - Values can be floats (fractions) or Quantities (absolute) - - Callable: custom function that receives and modifies envelope array. - window_type - Window function for tapering (only used for non-callable taper). - Options: - {sorted(WINDOW_FUNCTIONS)}. + smooth + Parameter controlling smoothing of the mute evenlope. Defines the sigma + Can be: + - None: sharp mute + - float (0.0-1.0): fraction of dimension range (e.g., 0.01 = 1%) + which is applied independently to each dimension involved in the + mute. + - int: Indicates number of samples for each dimension. + - Quantity with units, indicates values along a single dimension. + Only applicable if a single dimension is specified. + - dict: {dim: taper_value} for dimension-specific smooth + values which can be any of the above. invert If True, invert the taper such that the values outside the defined region are set to 0. relative If True (default), values are relative to coordinate edges. Positive values are offsets from start, negative from end. - samples - If True, values specified in samples rather than coordinate values. **kwargs Dimension specifications as (boundary_1, boundary_2) pairs. Each boundary can be: @@ -181,7 +225,7 @@ def mute( >>> import dascore as dc >>> from scipy.ndimage import gaussian_filter >>> - >>> patch = dc.get_example_patch("ricker_moveout") + >>> patch = dc.get_example_patch().full(1) >>> >>> >>> # Mute first 0.5s (relative to start by default) @@ -190,11 +234,8 @@ def mute( >>> # Mute everything except middle section >>> kept = patch.mute(time=(0.2, -0.2), mode="complement") >>> - >>> # Taper with absolute units - >>> muted = patch.mute( - ... time=(0.2, 0.8), - ... taper={'time': 0.02 * dc.units.s}, - ... ) + >>> # 1D Mute with smoothed absolute units for time. + >>> muted = patch.mute(time=(0.2, 0.8), smooth=0.02 * dc.units.s) >>> >>> # Classic first break mute: mute early arrivals >>> # Line from (t=0, d=0) to (t=0.3, d=300) defines velocity=1000 m/s @@ -229,8 +270,8 @@ def mute( ... time=([0, 0.375], [0, 0.25]), ... distance=([0, 300], [0, 300]), ... ) - >>> # Knock down edges with gaussian filter. - >>> smooth = envelope.gaussian_filter(time=5, distance=5, samples=True) + >>> # Knock down edges with rolling mean along the time dimension. + >>> smooth = envelope.rolling(time=5, samples=True).mean() >>> # Then multiply the two patches. >>> result = patch * smooth @@ -241,29 +282,30 @@ def mute( - Currently, mute doesn't support more than 2 dimensions. - - For more control over tapering, use a patch with one values then apply - custom tapering/smooting before multiplying with the original patch. - See example section for more details. + - For more control over boundary smoothing, use a patch with one values + then apply custom tapering/smooting before multiplying with the + original patch. See example section for more details. See Also -------- - [`Patch.select`](`dascore.Patch.select`) - [`Patch.taper_range`](`dascore.Patch.taper_range`) + - [`Patch.select`](`dascore.Patch.select`) + - [`Patch.taper_range`](`dascore.Patch.taper_range`) + - [`Patch.gaussian_filter`](`dacore.proc.filter.gaussian_filter`) """ - dims, axes, values = _get_mute_params(patch, kwargs) - taper_vals = _get_taper_vals(dims, taper, patch, samples) + dims, axes, values = _get_mute_params(patch, kwargs, relative) + out = np.zeros_like(patch.data) if invert else np.ones_like(patch.data) + fill_val = 1 if invert else 0 # Easy path for 1D mute. if len(dims) == 1: - out = _mute_patch_1d( - patch, - dims[0], - values, - taper_vals[0], - window=window_type if taper is not None else "boxcar", - samples=samples, - relative=relative, - invert=invert, - ) - return out - - breakpoint() + coord = patch.get_coord(dims[0]) + args = (values[0][0], values[0][1]) + _, c_index = coord.select(args, relative=relative) + index = broadcast_for_index(out.ndim, axes[0], c_index, slice(None)) + out[index] = fill_val + else: + out = _line_smooth(out, dims, axes, values, patch) + # Apply smoothing. + if smooth is not None: + sigma = _get_smooth_sigma(dims, smooth, patch) + out = gaussian_filter(out, sigma=sigma, axes=tuple(axes)) + return patch.update(data=patch.data * out) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 2affac1d8..9203c7ffe 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -9,11 +9,9 @@ from dascore.exceptions import ParameterError -def _get_testable_coord_values(coord, relative=False, samples=False): +def _get_testable_coord_values(coord, relative=False): """Get some values in the coordinate for testing different ranges.""" start_ind, stop_ind = len(coord) // 4, 3 * len(coord) // 4 - if samples: - return (start_ind, stop_ind) start, stop = coord[start_ind], coord[stop_ind] if relative: start, stop = start - coord.min(), stop - coord.min() @@ -26,14 +24,13 @@ def _assert_coord_ranges( zero_ranges, one_ranges, relative=True, - samples=False, ): """Assert that the expected values occur in the patch.""" for zrange in zero_ranges: - sub = patch.select(**{dim: zrange}, relative=relative, samples=samples) + sub = patch.select(**{dim: zrange}, relative=relative) assert np.allclose(sub.data, 0) for orange in one_ranges: - sub = patch.select(**{dim: orange}, relative=relative, samples=samples) + sub = patch.select(**{dim: orange}, relative=relative) assert np.allclose(sub.data, 1) @@ -67,6 +64,13 @@ def test_tuple_wrong_length_raises(self, random_patch): with pytest.raises(ParameterError, match="two boundaries when using"): random_patch.mute(time=(1, 2, 3)) + def test_smooth_bad_float(self, random_patch): + """If a floating point value is used, it must be between 0 and 1.""" + with pytest.raises(ParameterError, match="smooth parameter for"): + random_patch.mute(time=(1, 2), smooth=1.1) + with pytest.raises(ParameterError, match="smooth parameter for"): + random_patch.mute(time=(1, 2), smooth=-0.01) + class Test1DMute: """Test 1D block mutes (single dimension).""" @@ -83,7 +87,6 @@ def test_1d_mute_no_taper(self, patch_ones): zero_ranges=[(v1, v2)], one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], relative=True, - samples=False, ) def test_mute_open_interval(self, patch_ones): @@ -97,7 +100,6 @@ def test_mute_open_interval(self, patch_ones): zero_ranges=[(v2, ...)], one_ranges=[(..., v2 - coord.step)], relative=True, - samples=False, ) def test_mute_absolute(self, patch_ones): @@ -111,363 +113,90 @@ def test_mute_absolute(self, patch_ones): zero_ranges=[(v1, v2)], one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], relative=False, - samples=False, ) - def test_mute_samples(self, patch_ones): - """Test that sample mute works.""" - coord = patch_ones.get_coord("distance") - v1, v2 = _get_testable_coord_values(coord, samples=True) - muted1 = patch_ones.mute(distance=(v1, v2), samples=True) - _assert_coord_ranges( - patch=muted1, - dim="distance", - zero_ranges=[(v1, v2)], - one_ranges=[(..., v1 - 1), (v2 + 1, ...)], - samples=True, - ) + # Test smoothing parameters for 1D case. + def test_single_float_smooth(self, patch_ones): + """Test that taper mute works with a single floating point value.""" + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=0.01) + # With taper, we can't assert exact zeros/ones, just check shape + assert muted1.shape == patch_ones.shape + # Check that middle region has been modified (not all ones) + sub = muted1.select(time=(v1, v2), relative=True) + assert not np.allclose(sub.data, 1) - def test_mute_taper(self, patch_ones): - """Test that taper mute works.""" + def test_single_int_smooth(self, patch_ones): + """Test that smooth can be a single integer which means samples.""" coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted1 = patch_ones.mute(time=(v1, v2), relative=True, taper=0.1) + muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=5) # With taper, we can't assert exact zeros/ones, just check shape assert muted1.shape == patch_ones.shape # Check that middle region has been modified (not all ones) sub = muted1.select(time=(v1, v2), relative=True) assert not np.allclose(sub.data, 1) + def test_single_quantity(self, patch_ones): + """A single quantity should work with a specified dimension.""" + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + smooth_val = 0.01 * coord.units + muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=smooth_val) + # With taper, we can't assert exact zeros/ones, just check shape + assert muted1.shape == patch_ones.shape + # Check that middle region has been modified (not all ones) + sub = muted1.select(time=(v1, v2), relative=True) + assert not np.allclose(sub.data, 1) + + +class TestMuteLines: + """Tests for muting between lines.""" + + def test_generate_raise(self, patch_ones): + """A degenerate line (point) should raise.""" + msg = "Line specified" + with pytest.raises(ParameterError, match=msg): + patch_ones.mute( + time=([0, 0], [0, 0.25]), + distance=([0, 0], [0, 300]), + ) + + def test_no_common_point_raises(self, patch_ones): + """Two lines without a common point should raise.""" + msg = "No common point" + with pytest.raises(ParameterError, match=msg): + patch_ones.mute( + time=([10, 4], [0, 0.25]), + distance=([1, 6], [0, 300]), + ) + + def test_mute_lines(self, patch_ones): + """Mute a rectangular region in 2D.""" + breakpoint() + muted = patch_ones.mute( + time=([0, 0.375], [0, 0.25]), + distance=([0, 300], [0, 300]), + ) + + # Check interior is muted + assert np.allclose(muted.data[10, 15], 0) + # Check corners are not muted + assert np.allclose(muted.data[0, 0], 1) + assert np.allclose(muted.data[-1, -1], 1) + + def test_mute_strips(self, patch_ones): + """Mute strips along each dimension.""" + # Mute a time strip + muted_time = patch_ones.mute(time=(0.5, 1.0)) + # All distances should be affected equally + assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) + + # Mute a distance strip + muted_dist = patch_ones.mute(distance=(10, 20)) + # All times should be affected equally + assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) + # -# class TestMute2D: -# """Test 2D block mutes (multiple dimensions).""" -# -# def test_mute_rectangular_region(self, patch_ones): -# """Mute a rectangular region in 2D.""" -# time_coords = patch_ones.coords.get_array("time") -# dist_coords = patch_ones.coords.get_array("distance") -# -# t1, t2 = time_coords[10], time_coords[20] -# d1, d2 = dist_coords[5], dist_coords[15] -# -# muted = patch_ones.mute( -# time=(t1, t2), -# distance=(d1, d2), -# relative=False, -# ) -# -# # Check interior is muted -# assert np.allclose(muted.data[10, 15], 0) -# # Check corners are not muted -# assert np.allclose(muted.data[0, 0], 1) -# assert np.allclose(muted.data[-1, -1], 1) -# -# def test_mute_strips(self, patch_ones): -# """Mute strips along each dimension.""" -# # Mute a time strip -# muted_time = patch_ones.mute(time=(0.5, 1.0)) -# # All distances should be affected equally -# assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) -# -# # Mute a distance strip -# muted_dist = patch_ones.mute(distance=(10, 20)) -# # All times should be affected equally -# assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) -# -# -# class TestMuteModes: -# """Test different mute modes.""" -# -# def test_mode_union(self, patch_ones): -# """Test union mode (default, mutes inside region).""" -# muted = patch_ones.mute(time=(2.0, 6.0), mode="union") -# # Middle (at ~4s) should be zero -# mid_idx = len(patch_ones.coords.get_array("time")) // 2 -# assert muted.data[:, mid_idx].max() < 0.1 -# # Edges should be one -# assert muted.data[:, 0].min() > 0.9 -# -# def test_mode_complement(self, patch_ones): -# """Test complement mode (mutes outside region).""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[10], time_coords[20] -# -# muted = patch_ones.mute(time=(t1, t2), mode="complement", relative=False) -# -# # Edges should be zero -# assert muted.data[:, 0].max() < 0.1 -# assert muted.data[:, -1].max() < 0.1 -# # Middle should be one -# mid_idx = 15 -# assert muted.data[:, mid_idx].min() > 0.9 -# -# -# class TestMuteTaper: -# """Test taper application in mutes.""" -# -# def test_mute_with_taper(self, patch_ones): -# """Mute with taper creates gradual transition.""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# # Mute without taper -# muted_no_taper = patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=None, -# ) -# -# # Mute with taper (5% of dimension range) -# muted_with_taper = patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=0.05, -# ) -# -# # Without taper should be sharp transition -# assert muted_no_taper.data[:, 19].min() > 0.9 -# assert muted_no_taper.data[:, 20].max() < 0.1 -# -# # With taper should have gradual transition -# # Values in taper region should be between 0 and 1 -# taper_region = muted_with_taper.data[:, 15:20] -# assert (taper_region > 0).any() -# assert (taper_region < 1).any() -# -# def test_taper_dict(self, patch_ones): -# """Test dimension-specific taper using dict.""" -# time_coords = patch_ones.coords.get_array("time") -# dist_coords = patch_ones.coords.get_array("distance") -# -# taper_dict = { -# "time": 0.05, # 5% of time range -# "distance": 0, # No taper on distance -# } -# -# muted = patch_ones.mute( -# time=(time_coords[20], time_coords[40]), -# distance=(dist_coords[10], dist_coords[30]), -# relative=False, -# taper=taper_dict, -# ) -# -# # Should have taper in time, sharp in distance -# assert muted.shape == patch_ones.shape -# -# def test_taper_window_types(self, patch_ones): -# """Test different window types for taper.""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# for window_type in ["hann", "hamming", "triang"]: -# muted = patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=0.05, # 5% of dimension range -# window_type=window_type, -# ) -# assert isinstance(muted, dc.Patch) -# assert muted.shape == patch_ones.shape -# -# def test_taper_with_quantity(self, patch_ones): -# """Test taper with Quantity (absolute units).""" -# from dascore.units import s -# -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# # Use absolute time value with units (single dimension) -# muted = patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper={'time': 0.05 * s}, -# ) -# -# # Should have gradual transition -# taper_region = muted.data[:, 15:20] -# assert (taper_region > 0).any() -# assert (taper_region < 1).any() -# -# def test_taper_quantity_multiple_dims_raises(self, patch_ones): -# """Test that Quantity taper with multiple dims raises error.""" -# from dascore.units import s -# -# time_coords = patch_ones.coords.get_array("time") -# dist_coords = patch_ones.coords.get_array("distance") -# -# # Should raise error if Quantity used without dict -# with pytest.raises(ParameterError, match="Cannot use Quantity"): -# patch_ones.mute( -# time=(time_coords[20], time_coords[40]), -# distance=(dist_coords[10], dist_coords[30]), -# relative=False, -# taper=0.05 * s, -# ) -# -# def test_taper_mixed_quantity_fraction(self, patch_ones): -# """Test mixed taper: fraction for one dim, Quantity for another.""" -# from dascore.units import s, m -# -# time_coords = patch_ones.coords.get_array("time") -# dist_coords = patch_ones.coords.get_array("distance") -# -# muted = patch_ones.mute( -# time=(time_coords[20], time_coords[40]), -# distance=(dist_coords[10], dist_coords[30]), -# relative=False, -# taper={'time': 0.05, 'distance': 10 * m}, -# ) -# -# assert muted.shape == patch_ones.shape -# -# def test_taper_custom_function(self, patch_ones): -# """Test custom taper function.""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# # Define custom taper that squares the envelope -# def custom_taper(envelope): -# return envelope ** 2 -# -# muted = patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=custom_taper, -# ) -# -# # Middle should still be zero -# assert muted.data[:, 30].max() < 0.1 -# # Edges should still be one -# assert muted.data[:, 0].min() > 0.9 -# -# def test_taper_custom_function_invalid_return(self, patch_ones): -# """Test that custom taper with invalid return raises error.""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# # Function that returns wrong type -# def bad_taper(envelope): -# return list(envelope) -# -# with pytest.raises(ParameterError, match="must return a numpy array"): -# patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=bad_taper, -# ) -# -# def test_taper_custom_function_wrong_shape(self, patch_ones): -# """Test that custom taper with wrong shape raises error.""" -# time_coords = patch_ones.coords.get_array("time") -# t1, t2 = time_coords[20], time_coords[40] -# -# # Function that returns wrong shape -# def bad_shape_taper(envelope): -# return np.ones((10, 10)) -# -# with pytest.raises(ParameterError, match="must return array with shape"): -# patch_ones.mute( -# time=(t1, t2), -# relative=False, -# taper=bad_shape_taper, -# ) -# -# -# class TestMuteEdgeCases: -# """Test edge cases and boundary conditions.""" -# -# def test_mute_with_none_boundaries(self, patch_ones): -# """Test using None to reference coordinate edges.""" -# # Mute from start to middle -# time_coords = patch_ones.coords.get_array("time") -# mid_time = time_coords[len(time_coords) // 2] -# -# muted = patch_ones.mute(time=(None, mid_time), relative=False) -# -# # First half should be zero -# assert muted.data[:, 0].max() < 0.1 -# # Second half should be one -# assert muted.data[:, -1].min() > 0.9 -# -# def test_mute_entire_dimension(self, patch_ones): -# """Mute entire dimension.""" -# muted = patch_ones.mute(time=(None, None), relative=False) -# # Everything should be zero -# assert np.allclose(muted.data, 0) -# -# def test_mute_zero_width(self, patch_ones): -# """Mute with same start and end values.""" -# time_coords = patch_ones.coords.get_array("time") -# t = time_coords[10] -# -# muted = patch_ones.mute(time=(t, t), relative=False) -# # Should mute just that one sample (or very close) -# assert muted.data[:, 10].max() < 0.1 -# -# -# class TestMuteWithSamples: -# """Test mute with samples=True.""" -# -# def test_mute_samples_mode(self, patch_ones): -# """Test mute using sample indices.""" -# # Mute samples 10 to 20 -# muted = patch_ones.mute(time=(10, 20), samples=True) -# -# # Samples 10-20 should be zero -# assert muted.data[:, 15].max() < 0.1 -# # Other samples should be one -# assert muted.data[:, 0].min() > 0.9 -# assert muted.data[:, -1].min() > 0.9 -# -# -# class TestGeometricMutes: -# """Test geometric (array-based) mutes - Phase 2.""" -# -# def test_array_boundaries_not_implemented(self, ricker_patch): -# """Array boundaries should raise error for mixing scalar/array.""" -# # Currently we don't allow mixing scalars and arrays -# with pytest.raises(ParameterError, match="Cannot mix scalar and array"): -# ricker_patch.mute( -# time=(0, [0, 0.3]), -# distance=(0, [0, 300]), -# relative=False, -# ) -# -# def test_mismatched_array_lengths_raises(self, random_patch): -# """Arrays with different lengths should raise error during parsing.""" -# with pytest.raises(ParameterError, match="Cannot mix scalar and array"): -# # This currently raises mixing error before length checking -# random_patch.mute( -# time=(0, [0, 0.3, 0.5]), # scalar and array -# distance=(0, [0, 300]), # scalar and array -# relative=False, -# ) -# -# -# class TestMutePreservesMetadata: -# """Test that mute preserves patch metadata.""" -# -# def test_preserves_coords(self, random_patch): -# """Mute should preserve coordinates.""" -# muted = random_patch.mute(time=(0, 0.5)) -# assert muted.coords == random_patch.coords -# -# def test_preserves_attrs(self, random_patch): -# """Mute should preserve attributes (except history).""" -# muted = random_patch.mute(time=(0, 0.5)) -# # History will be different due to processing -# assert muted.attrs.data_type == random_patch.attrs.data_type -# assert muted.attrs.network == random_patch.attrs.network -# assert muted.attrs.station == random_patch.attrs.station -# -# def test_preserves_shape(self, random_patch): -# """Mute should preserve shape.""" -# muted = random_patch.mute(time=(0, 0.5)) -# assert muted.shape == random_patch.shape -# -# def test_preserves_dtype(self, random_patch): -# """Mute should preserve data type.""" -# muted = random_patch.mute(time=(0, 0.5)) -# # Allow for float conversion -# assert muted.dtype in (random_patch.dtype, np.float64, np.float32) From 3d732be59539f7822c1c5b2a0d1258bbe318e854 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Mon, 20 Oct 2025 17:53:58 +0100 Subject: [PATCH 04/15] more work on mute --- dascore/proc/mute.py | 403 +++++++++++++++++++++++------------ dascore/utils/misc.py | 61 ++++++ tests/test_proc/test_mute.py | 59 +++-- 3 files changed, 368 insertions(+), 155 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 1067022cf..acacb645c 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -2,98 +2,284 @@ from __future__ import annotations -import dataclasses from collections import defaultdict from collections.abc import Mapping, Sized import numpy as np +from numpy.typing import NDArray from scipy.ndimage import gaussian_filter import dascore as dc from dascore.constants import PatchType from dascore.exceptions import ParameterError -from dascore.utils.misc import broadcast_for_index +from dascore.utils.misc import ( + get_2d_line_intersection, + vectors_same_direction, +) +from dascore.utils.models import DascoreBaseModel from dascore.utils.patch import get_dim_axis_value, patch_function -@dataclasses.dataclass -class _OriginNLines: - """Get a tuple of origin, line1, line2.""" +class _MuteGeometry(DascoreBaseModel): + """ + Parent class for Mute Geometry. + """ + + dims: tuple[str, ...] + axes: tuple[int, ...] + relative: bool = True + + @classmethod + def from_params(cls, vals, dims, axes, patch, relative): + """Initialize Mute Geometry from input parameters.""" + + def _mask_array( + self, + array: NDArray, + ): + pass + + def _apply_smoothing(self, array, smooth, patch): + """Apply smoothing to the array.""" + sigma = self._get_smooth_sigma(smooth, patch) + return gaussian_filter(array, sigma=sigma, axes=tuple(self.axes)) + + def _get_smooth_sigma(self, smooth, patch): + """Get sigma values, in samples, for the gaussian kernel.""" + + def _broadcast_smooth_to_dims(dims, smooth): + """Broadcast smooth samples to dims length.""" + # First, we need to get the smooth parameters to line up with the dims. + # For a single value, just broadcast to dim length. + if not isinstance(smooth, Mapping): + vals = [smooth] * len(dims) + else: + # Otherwise each dimension's smooth must be specified. + if not set(dims) == set(smooth): + msg = ( + f"If a taper dictionary is used in Mute, it must have all " + f"the same keys as the dimensions. Kwarg dims are {dims} and" + f"taper keys are {list(smooth)}." + ) + raise ParameterError(msg) + vals = [smooth[dim] for dim in dims] + return vals + + def _convert_to_samples(smooth, dims, patch): + """Convert the smooth parameter to number of samples.""" + out = [] + for dim, val in zip(dims, smooth): + coord = patch.get_coord(dim) + if val is None: + out.append(0) + elif isinstance(val, int | np.integer): + out.append(val) + elif isinstance(val, float | np.floating): + if not 0 <= val <= 1: + msg = ( + f"Mute's smooth parameter for {dim} must be between 0 " + f"and 1 when using a floating point value." + ) + raise ParameterError(msg) + out.append(int(np.round(len(coord) * val))) + else: # should capture quantities + out.append(coord.get_sample_count(val)) + return out + + smooth_by_dims = _broadcast_smooth_to_dims(self.dims, smooth) + smooth_ints = _convert_to_samples(smooth_by_dims, self.dims, patch) + # Now finagle into input for scipy's gaussian smooth. + return smooth_ints + + +class _MuteGeometry1D(_MuteGeometry): + """ + Private container to manage 1D Mute Geometry. + """ - origin: np.ndarray - line1: np.ndarray - line2: np.ndarray + lims: tuple @classmethod - def _check_degenerate(self, line): - """Ensure a line isn't degenerate else raise.""" + def from_params(cls, vals, dims, axes, patch, relative): + """Initialize Mute Geometry from input parameters.""" + vals = np.array(vals) + if vals.size != 2: # Handle single dimension Mute. + msg = "Mute requires two boundaries when using a single dimension." + raise ParameterError(msg) + lims = (vals[0][0], vals[0][1]) + return cls(dims=dims, axes=axes, relative=relative, lims=lims) + + def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: + coord = patch.get_coord(self.dims[0]) + _, c_index = coord.select(self.lims, relative=self.relative) + index = [slice(None)] * array.ndim + index[self.axes[0]] = c_index + array[tuple(index)] = fill_value + return array + + +class _MuteGeometry2D(_MuteGeometry): + """ + Private container to help manage the geometry of the slope filter. + + Generally this is initialized with the `from_point_list` class method. + + Parameters + ---------- + origin + The shared origin for the two points. + line1 + A numpy array of the un-normalized first line. + line2 + A numpy array of the un-normalized second line. + norm + The corresponding coordinate range for each dimension (as float) + line1_norm + The first line divided by the norm. + line2_norm + The second line divided by the norm. + + The mute region is selected by finding the cross product between each + normalized (scaled) line + """ + + origin: NDArray[np.floating] + norm: NDArray[np.floating] + line1_norm: NDArray[np.floating] + line2_norm: NDArray[np.floating] + + @staticmethod + def _check_degenerate(line): + """Ensure a line isn't really a point else raise.""" if len(np.unique(line, axis=0)) != 2: msg = f"Line specified for mute {line} is degenerate!" raise ParameterError(msg) @classmethod - def from_point_list(cls, points): - """""" + def _params_from_point_list(cls, points, patch, dims): + """Get the parameters from a list of points.""" cls._check_degenerate(points[:2]) cls._check_degenerate(points[2:]) - # We need to ensure one of the points is duplicated; this is the orign - unique, idx, counts = np.unique( - points, axis=0, return_index=True, return_counts=True - ) - if len(unique) == len(points): - msg = ( - f"No common point found in {points}. For unambiguous mute " - f"lines must share an origin." - ) + origin = get_2d_line_intersection(*points) + norm = np.array([dc.to_float(patch.get_coord(x).coord_range()) for x in dims]) + # Get vectors. Need to normalize in coord space and by l2. + v1 = ((points[1] - points[0]) - origin) / norm + v2 = ((points[3] - points[2]) - origin) / norm + v1_norm = v1 / np.linalg.norm(v1) + v2_norm = v2 / np.linalg.norm(v2) + # Ensure vectors are pointing in the same direction, otherwise the + # mute area is ambiguous. + if not vectors_same_direction(v1, v2): + msg = "Mute vectors must point in the same direction." raise ParameterError(msg) - gt_1_counts = counts > 1 - origin = points[idx[gt_1_counts]] - others = points[idx[~gt_1_counts]] - l1, l2 = others - return cls(origin=origin, line1=l1, line2=l2) + out = dict( + origin=origin, + norm=norm, + line1_norm=v1_norm, + line2_norm=v2_norm, + ) + return out + @classmethod + def from_params(cls, vals, dims, axes, patch, relative=True): + """ + Return values which are two lines in absolute coordinate space, + (expressed as floats) + """ + + def _get_coord_float_values(coord, vals, relative): + """Get the coordinate float values relative to start of coords.""" + out = dc.to_float(np.array(vals)) + if not relative: + out = out - dc.to_float(coord.min()) + return out -def _get_2d_mute_lines(vals, patch, dims, relative=True): - """ - Return values which are two lines in absolute coordinate space, - (expressed as floats) - """ + out = [] + # Keep track of the axis that need to be filled in (eg None, paired w/ float) + fill_inds = defaultdict(list) + for ind, (dim, row) in enumerate(zip(dims, vals)): + coord = patch.get_coord(dim) + # In the case of single values (eg None, ..., or some other) + # We need to just mark them and come back later. + if not isinstance(row, Sized): + # We need to get the actual value from the coord. + if row is not None and row is not Ellipsis: + row = _get_coord_float_values(coord, vals, relative=relative) + fill_inds[dim].append(row) + out.append(None) + # Otherwise, we just run with it. + vals = _get_coord_float_values(coord, np.array(row), relative=relative) + out.append(vals) + # Now we can ascertain the line intended by None. + if fill_inds: + raise NotImplementedError("Working on it.") + # Then put the 4 points together. This gives us a len 4 array with rows + # as points from first line, then points from second. + points = np.stack(out, axis=-1).reshape(-1, 2) + line_params = cls._params_from_point_list(points, patch, dims) + kwargs = dict(dims=dims, axes=axes, relative=relative) | line_params + return cls(**kwargs) + + def _get_normalized_array_coord(self, array, patch): + """ + Get an array that matches the dimensionality of envelope but of its + normalized, relative coordinates. + """ + out = [] + for dim in self.dims: + ax = patch.get_axis(dim) + coord = patch.get_coord(dim) - def _get_coord_float_values(coord, vals, relative): - """Get the coordinate float values relative to start of coords.""" - out = dc.to_float(np.array(vals)) - if not relative: - out = out - dc.to_float(coord.min()) + # Get the coordinate values normalized to coord range. + coord_vals = dc.to_float(coord.values) + coord_range = dc.to_float(coord.coord_range()) + if self.relative: + coord_vals = coord_vals - dc.to_float(coord.min()) + + # First get values in coord. + norm_vals = (coord_vals - self.origin[ax]) / dc.to_float(coord_range) + # We need to transform coordinate values to values between 0 and 1. + # with the same dimensionality as array. + coord_inds = [None] * array.ndim + coord_inds[ax] = slice(None, len(coord)) + norms = norm_vals[tuple(coord_inds)] + + # Next, we set those values on an array with the same shape as array. + inds = [slice(None)] * array.ndim + inds[ax] = slice(None, len(coord)) + carray = np.empty_like(array) + carray[tuple(inds)] = norms + out.append(carray) + # Return new axis as -1 so it will broadcast with lines. + out = np.stack(out, axis=-1) return out - out = [] - # Keep track of the axis that need to be filled in (eg None, paired w/ float) - fill_inds = defaultdict(list) - for ind, (dim, row) in enumerate(zip(dims, vals)): - coord = patch.get_coord(dim) - # In the case of single values (eg None, ..., or some other) - # We need to just mark them and come back later. - if not isinstance(row, Sized): - # We need to get the actual value from the coord. - if row is not None and row is not Ellipsis: - row = _get_coord_float_values(coord, vals, relative=relative) - fill_inds[dim].append(row) - out.append(None) - # Otherwise, we just run with it. - vals = _get_coord_float_values(coord, np.array(row), relative=relative) - out.append(vals) - # Now we can ascertain the line intended by None. - if fill_inds: - breakpoint() - # Then put the 4 points together. This gives us a len 4 array with rows - # as points from first line, then points from second. - points = np.stack(out, axis=-1).reshape(-1, 2) - origin_n_lines = _OriginNLines.from_point_list(points) - breakpoint() - - -def _get_mute_params(patch, kwargs, relative=True): + def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: + """Apply the mask to the output array.""" + coord_norms = self._get_normalized_array_coord(array, patch) + # Get padded arrays with 0 z values for cross product. + array_widths = ((0, 0), (0, 0), (0, 1)) + coord_z = np.pad( + coord_norms, pad_width=array_widths, mode="constant", constant_values=0 + ) + line_widths = (0, 1) + l1_z = np.pad(self.line1_norm, line_widths, mode="constant", constant_values=0) + l2_z = np.pad(self.line2_norm, line_widths, mode="constant", constant_values=0) + # The selected points will have different cross product sign and + # the same dot product sign. + cross1_z = np.cross(l1_z, coord_z)[:, :, 2] + cross2_z = np.cross(l2_z, coord_z)[:, :, 2] + ok_cross = cross1_z * cross2_z < 0 + # Get dot product with each line. + dot1 = np.sum(self.line1_norm[None, :] * coord_norms, axis=-1) + dot2 = np.sum(self.line2_norm[None, :] * coord_norms, axis=-1) + ok_dot = dot1 * dot2 > 0 + envelope = ok_cross & ok_dot + return envelope + + +def _get_mute_geometry(patch, kwargs, relative=True): """ Get (and validate) the muting parameters. @@ -112,68 +298,14 @@ def _get_mute_params(patch, kwargs, relative=True): axes = np.array([x[1] for x in dim_ax_vals]) val_list = [x[2] for x in dim_ax_vals] if len(dims) == 1: - vals = np.array(val_list) - if vals.size != 2: # Handle single dimension Mute. - msg = "Mute requires two boundaries when using a single dimension." - raise ParameterError(msg) + geometry = _MuteGeometry1D.from_params( + val_list, dims, axes, patch, relative=relative + ) elif len(dims) > 1: # Dealing with lines. . - vals = _get_2d_mute_lines(val_list, patch, dims, relative) - return dims, axes, vals - - -def _get_smooth_sigma(dims, smooth, patch): - """Get sigma values, in samples, for the gaussian kernel.""" - - def _broadcast_smooth_to_dims(dims, smooth): - """Broadcast smooth samples to dims length.""" - # First, we need to get the smooth parameters to line up with the dims. - # For a single value, just broadcast to dim length. - if not isinstance(smooth, Mapping): - vals = [smooth] * len(dims) - else: - # Otherwise each dimension's smooth must be specified. - if not set(dims) == set(smooth): - msg = ( - f"If a taper dictionary is used in Mute, it must have all " - f"the same keys as the dimensions. Kwarg dims are {dims} and" - f"taper keys are {list(smooth)}." - ) - raise ParameterError(msg) - vals = [smooth[dim] for dim in dims] - return vals - - def _convert_to_samples(smooth, dims, patch): - """Convert the smooth parameter to number of samples.""" - out = [] - for dim, val in zip(dims, smooth): - coord = patch.get_coord(dim) - if val is None: - out.append(0) - elif isinstance(val, int | np.integer): - out.append(val) - elif isinstance(val, float | np.floating): - if not 0 <= val <= 1: - msg = ( - f"Mute's smooth parameter for {dim} must be between 0 " - f"and 1 when using a floating point value." - ) - raise ParameterError(msg) - out.append(int(np.round(len(coord) * val))) - else: # should capture quantities - out.append(coord.get_sample_count(val)) - return out - - smooth_by_dims = _broadcast_smooth_to_dims(dims, smooth) - smooth_ints = _convert_to_samples(smooth_by_dims, dims, patch) - # Now finagle into input for scipy's gaussian smooth. - return smooth_ints - - -def _line_smooth(data, dims, axes, values, patch): - """ - Mute data between two lines. - """ - breakpoint() + geometry = _MuteGeometry2D.from_params( + val_list, dims, axes, patch, relative=relative + ) + return geometry @patch_function() @@ -292,20 +424,17 @@ def mute( - [`Patch.taper_range`](`dascore.Patch.taper_range`) - [`Patch.gaussian_filter`](`dacore.proc.filter.gaussian_filter`) """ - dims, axes, values = _get_mute_params(patch, kwargs, relative) - out = np.zeros_like(patch.data) if invert else np.ones_like(patch.data) + # Get geometry object to set up the problem. + geo = _get_mute_geometry(patch, kwargs, relative) + # Initialize the output array which shares the dimensionality of the + # patch (so it will broadcast) but is as flat as possible. + out_shape = [patch.shape[ax] if ax in geo.axes else 1 for ax in range(patch.ndim)] + out = np.zeros(out_shape) if invert else np.ones(out_shape) fill_val = 1 if invert else 0 # Easy path for 1D mute. - if len(dims) == 1: - coord = patch.get_coord(dims[0]) - args = (values[0][0], values[0][1]) - _, c_index = coord.select(args, relative=relative) - index = broadcast_for_index(out.ndim, axes[0], c_index, slice(None)) - out[index] = fill_val - else: - out = _line_smooth(out, dims, axes, values, patch) - # Apply smoothing. + out = geo._apply_mask(out, patch, fill_val) + # Apply smoothing if requested. if smooth is not None: - sigma = _get_smooth_sigma(dims, smooth, patch) - out = gaussian_filter(out, sigma=sigma, axes=tuple(axes)) + out = geo._apply_smoothing(out, smooth, patch) + return patch.update(data=patch.data * out) diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 4d529ef40..7b4ca25d2 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -895,3 +895,64 @@ def _robust_equality_check(obj1, obj2): return True finally: visited.remove(pair_id) + + +def get_2d_line_intersection(p1, p2, p3, p4): + """ + Return intersection point of two lines (p1,p2) and (p3,p4). + + Parameters + ---------- + p1, p2, p3, p4 + Each a pair of (x, y) coordinates. + + Returns + ------- + point + (x, y) intersection point, or None if lines are parallel. + """ + x1, y1 = p1 + x2, y2 = p2 + x3, y3 = p3 + x4, y4 = p4 + + denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4) + if denom == 0: + return None # Parallel or coincident + + px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denom + py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denom + return np.array([px, py]) + + +def vectors_same_direction(v1, v2, tol=1e-9): + """ + Determine if vectors (p2 - p1) and (p4 - p3) point in the same direction. + + Parameters + ---------- + v1 + First vector. + v2 + Second vector. + tol + Tolerance for considering vectors parallel. + + Returns + ------- + bool + True if vectors are parallel and oriented the same way. + """ + # Normalize to avoid magnitude issues + norm1 = np.linalg.norm(v1) + norm2 = np.linalg.norm(v2) + if norm1 < tol or norm2 < tol: + return False # one is degenerate + + v1 /= norm1 + v2 /= norm2 + + # Same direction: dot product positive + same_dir = np.dot(v1, v2) > 0 + + return same_dir diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 9203c7ffe..f3a8c982d 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -34,6 +34,31 @@ def _assert_coord_ranges( assert np.allclose(sub.data, 1) +def _assert_point_values( + patch, + dims=None, + points=(), + expected_values=(), + relative=True, +): + """Assert points in the patch are 0 or 1.""" + assert len(points) == len(expected_values) + dims = dims if dims is not None else patch.dims + axes = tuple(patch.get_axis(x) for x in dims) + + for vals, expected in zip(points, expected_values, strict=True): + # Get the index (in the data array) of the expected 0. + inds = ( + patch.get_coord(dims[num]).get_next_index( + vals[axes.index(num)], relative=relative + ) + if num in axes + else slice(None) + for num in range(patch.ndim) + ) + assert np.isclose(patch.data[tuple(inds)], expected) + + @pytest.fixture(scope="session") def patch_ones(random_patch): """Return a patch filled with ones.""" @@ -154,7 +179,7 @@ def test_single_quantity(self, patch_ones): class TestMuteLines: """Tests for muting between lines.""" - def test_generate_raise(self, patch_ones): + def test_point_raise(self, patch_ones): """A degenerate line (point) should raise.""" msg = "Line specified" with pytest.raises(ParameterError, match=msg): @@ -163,28 +188,26 @@ def test_generate_raise(self, patch_ones): distance=([0, 0], [0, 300]), ) - def test_no_common_point_raises(self, patch_ones): - """Two lines without a common point should raise.""" - msg = "No common point" - with pytest.raises(ParameterError, match=msg): - patch_ones.mute( - time=([10, 4], [0, 0.25]), - distance=([1, 6], [0, 300]), - ) + # def test_no_common_point_raises(self, patch_ones): + # """Two lines without a common point should raise.""" + # msg = "No common point" + # with pytest.raises(ParameterError, match=msg): + # patch_ones.mute( + # time=([10, 4], [0, 0.25]), + # distance=([1, 6], [0, 300]), + # ) def test_mute_lines(self, patch_ones): """Mute a rectangular region in 2D.""" - breakpoint() muted = patch_ones.mute( - time=([0, 0.375], [0, 0.25]), - distance=([0, 300], [0, 300]), + time=([0, 2], [0, 4]), + distance=([0, 100], [0, 100]), + ) + points = [(145, 4), (182, 6), (100, 3), (180, 3), (60, 6), (50, 1)] + expected = [1, 1, 1, 0, 0, 0] + _assert_point_values( + muted, ("time", "distance"), points=points, expected_values=expected ) - - # Check interior is muted - assert np.allclose(muted.data[10, 15], 0) - # Check corners are not muted - assert np.allclose(muted.data[0, 0], 1) - assert np.allclose(muted.data[-1, -1], 1) def test_mute_strips(self, patch_ones): """Mute strips along each dimension.""" From 617467e4aaebf8cc598c12d3498fb551fc805bb9 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Tue, 21 Oct 2025 11:48:41 +0100 Subject: [PATCH 05/15] work on muter --- dascore/proc/mute.py | 96 ++++++++++++++++++++++-------------- dascore/utils/misc.py | 39 ++------------- tests/test_proc/test_mute.py | 45 ++++++++--------- 3 files changed, 85 insertions(+), 95 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index acacb645c..f11622dd3 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -4,9 +4,11 @@ from collections import defaultdict from collections.abc import Mapping, Sized +from typing import ClassVar import numpy as np from numpy.typing import NDArray +from scipy.linalg import norm from scipy.ndimage import gaussian_filter import dascore as dc @@ -14,7 +16,6 @@ from dascore.exceptions import ParameterError from dascore.utils.misc import ( get_2d_line_intersection, - vectors_same_direction, ) from dascore.utils.models import DascoreBaseModel from dascore.utils.patch import get_dim_axis_value, patch_function @@ -138,6 +139,8 @@ class _MuteGeometry2D(_MuteGeometry): The first line divided by the norm. line2_norm The second line divided by the norm. + parallel + If True, the lines are parallel to each other. The mute region is selected by finding the cross product between each normalized (scaled) line @@ -147,37 +150,46 @@ class _MuteGeometry2D(_MuteGeometry): norm: NDArray[np.floating] line1_norm: NDArray[np.floating] line2_norm: NDArray[np.floating] + parallel: bool - @staticmethod - def _check_degenerate(line): - """Ensure a line isn't really a point else raise.""" - if len(np.unique(line, axis=0)) != 2: - msg = f"Line specified for mute {line} is degenerate!" - raise ParameterError(msg) + # For determining if any points are degenerate (norms close to 0) + # or parallel (dot products close to 1) + _tolerance: ClassVar[float] = 1e-9 @classmethod - def _params_from_point_list(cls, points, patch, dims): + def _get_line_params(cls, points, patch, dims): """Get the parameters from a list of points.""" - cls._check_degenerate(points[:2]) - cls._check_degenerate(points[2:]) - # We need to ensure one of the points is duplicated; this is the orign - origin = get_2d_line_intersection(*points) - norm = np.array([dc.to_float(patch.get_coord(x).coord_range()) for x in dims]) + tol = cls._tolerance + # Get the origin (where two lines intersect), (nan, nan) if parallel. + coord_norm = np.array( + [dc.to_float(patch.get_coord(x).coord_range()) for x in dims] + ) + origin = get_2d_line_intersection(*points) / coord_norm # Get vectors. Need to normalize in coord space and by l2. - v1 = ((points[1] - points[0]) - origin) / norm - v2 = ((points[3] - points[2]) - origin) / norm - v1_norm = v1 / np.linalg.norm(v1) - v2_norm = v2 / np.linalg.norm(v2) - # Ensure vectors are pointing in the same direction, otherwise the - # mute area is ambiguous. - if not vectors_same_direction(v1, v2): - msg = "Mute vectors must point in the same direction." + v1 = (points[1] - points[0]) / coord_norm + v2 = (points[3] - points[2]) / coord_norm + norm_v1, norm_v2 = norm(v1), norm(v2) + v1_norm, v2_norm = v1 / norm_v1, v2 / norm_v2 + # Check if any points are degenerate. + if (norm_v1 < tol) or (norm_v2 < tol): + msg = f"A line provided to mute ({v1} or {v2}) is degenerate!" raise ParameterError(msg) + # Determine if vectors are parallel and point in same direction + parallel = (1 - np.dot(v1, v2)) < tol + same_direction = np.dot(v1_norm, v2_norm) >= 0 + if not same_direction: + if not parallel: + msg = "Non-parallel mute vectors must point in the same direction." + raise ParameterError(msg) + else: + # If lines are parallel we can just reverse the direction of one. + v1 *= -1 out = dict( origin=origin, - norm=norm, - line1_norm=v1_norm, - line2_norm=v2_norm, + norm=coord_norm, + line1_norm=v1 / norm_v1, + line2_norm=v2 / norm_v2, + parallel=parallel, ) return out @@ -217,7 +229,7 @@ def _get_coord_float_values(coord, vals, relative): # Then put the 4 points together. This gives us a len 4 array with rows # as points from first line, then points from second. points = np.stack(out, axis=-1).reshape(-1, 2) - line_params = cls._params_from_point_list(points, patch, dims) + line_params = cls._get_line_params(points, patch, dims) kwargs = dict(dims=dims, axes=axes, relative=relative) | line_params return cls(**kwargs) @@ -230,13 +242,11 @@ def _get_normalized_array_coord(self, array, patch): for dim in self.dims: ax = patch.get_axis(dim) coord = patch.get_coord(dim) - # Get the coordinate values normalized to coord range. coord_vals = dc.to_float(coord.values) coord_range = dc.to_float(coord.coord_range()) if self.relative: coord_vals = coord_vals - dc.to_float(coord.min()) - # First get values in coord. norm_vals = (coord_vals - self.origin[ax]) / dc.to_float(coord_range) # We need to transform coordinate values to values between 0 and 1. @@ -244,7 +254,6 @@ def _get_normalized_array_coord(self, array, patch): coord_inds = [None] * array.ndim coord_inds[ax] = slice(None, len(coord)) norms = norm_vals[tuple(coord_inds)] - # Next, we set those values on an array with the same shape as array. inds = [slice(None)] * array.ndim inds[ax] = slice(None, len(coord)) @@ -255,13 +264,15 @@ def _get_normalized_array_coord(self, array, patch): out = np.stack(out, axis=-1) return out - def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: - """Apply the mask to the output array.""" - coord_norms = self._get_normalized_array_coord(array, patch) + def _cross_product_ok(self, coord_array): + """ + Determine if each point is in the region based on the cross product + requirement. + """ # Get padded arrays with 0 z values for cross product. array_widths = ((0, 0), (0, 0), (0, 1)) coord_z = np.pad( - coord_norms, pad_width=array_widths, mode="constant", constant_values=0 + coord_array, pad_width=array_widths, mode="constant", constant_values=0 ) line_widths = (0, 1) l1_z = np.pad(self.line1_norm, line_widths, mode="constant", constant_values=0) @@ -271,12 +282,24 @@ def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: cross1_z = np.cross(l1_z, coord_z)[:, :, 2] cross2_z = np.cross(l2_z, coord_z)[:, :, 2] ok_cross = cross1_z * cross2_z < 0 + return ok_cross + + def _dot_product_ok(self, coord_array): # Get dot product with each line. - dot1 = np.sum(self.line1_norm[None, :] * coord_norms, axis=-1) - dot2 = np.sum(self.line2_norm[None, :] * coord_norms, axis=-1) + dot1 = np.sum(self.line1_norm[None, :] * coord_array, axis=-1) + dot2 = np.sum(self.line2_norm[None, :] * coord_array, axis=-1) ok_dot = dot1 * dot2 > 0 - envelope = ok_cross & ok_dot - return envelope + return ok_dot + + def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: + """Apply the mask to the output array.""" + coord_norms = self._get_normalized_array_coord(array, patch) + ok_cross = self._cross_product_ok(coord_norms) + # For parallel lines only cross product is needed to define region. + if self.parallel: + return ok_cross + ok_dot = self._dot_product_ok(coord_norms) + return ok_cross & ok_dot def _get_mute_geometry(patch, kwargs, relative=True): @@ -359,7 +382,6 @@ def mute( >>> >>> patch = dc.get_example_patch().full(1) >>> - >>> >>> # Mute first 0.5s (relative to start by default) >>> muted = patch.mute(time=(0, 0.5)) >>> diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 7b4ca25d2..1a42b1096 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -909,7 +909,7 @@ def get_2d_line_intersection(p1, p2, p3, p4): Returns ------- point - (x, y) intersection point, or None if lines are parallel. + (x, y) intersection point. x and y are nan if lines are parallel. """ x1, y1 = p1 x2, y2 = p2 @@ -917,42 +917,9 @@ def get_2d_line_intersection(p1, p2, p3, p4): x4, y4 = p4 denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4) - if denom == 0: - return None # Parallel or coincident + if np.isclose(denom, 0): + np.array([np.nan, np.nan]) px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denom py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denom return np.array([px, py]) - - -def vectors_same_direction(v1, v2, tol=1e-9): - """ - Determine if vectors (p2 - p1) and (p4 - p3) point in the same direction. - - Parameters - ---------- - v1 - First vector. - v2 - Second vector. - tol - Tolerance for considering vectors parallel. - - Returns - ------- - bool - True if vectors are parallel and oriented the same way. - """ - # Normalize to avoid magnitude issues - norm1 = np.linalg.norm(v1) - norm2 = np.linalg.norm(v2) - if norm1 < tol or norm2 < tol: - return False # one is degenerate - - v1 /= norm1 - v2 /= norm2 - - # Same direction: dot product positive - same_dir = np.dot(v1, v2) > 0 - - return same_dir diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index f3a8c982d..937a7d5a2 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -175,30 +175,43 @@ def test_single_quantity(self, patch_ones): sub = muted1.select(time=(v1, v2), relative=True) assert not np.allclose(sub.data, 1) + def test_mute_strips(self, patch_ones): + """Mute strips along each dimension.""" + # Mute a time strip + muted_time = patch_ones.mute(time=(0.5, 1.0)) + # All distances should be affected equally + assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) + + # Mute a distance strip + muted_dist = patch_ones.mute(distance=(10, 20)) + # All times should be affected equally + assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) + class TestMuteLines: """Tests for muting between lines.""" def test_point_raise(self, patch_ones): """A degenerate line (point) should raise.""" - msg = "Line specified" + msg = "is degenerate" with pytest.raises(ParameterError, match=msg): patch_ones.mute( time=([0, 0], [0, 0.25]), distance=([0, 0], [0, 300]), ) - # def test_no_common_point_raises(self, patch_ones): - # """Two lines without a common point should raise.""" - # msg = "No common point" - # with pytest.raises(ParameterError, match=msg): - # patch_ones.mute( - # time=([10, 4], [0, 0.25]), - # distance=([1, 6], [0, 300]), - # ) + def test_not_same_direction(self, patch_ones): + """Create lines which do not point in the same directions.""" + match = "point in the same direction" + + with pytest.raises(ValueError, match=match): + patch_ones.mute( + time=[[0, 0], [1, -1]], + distance=[[1, -1], [-1, 1]], + ) def test_mute_lines(self, patch_ones): - """Mute a rectangular region in 2D.""" + """Mute non-parallel lines.""" muted = patch_ones.mute( time=([0, 2], [0, 4]), distance=([0, 100], [0, 100]), @@ -209,17 +222,5 @@ def test_mute_lines(self, patch_ones): muted, ("time", "distance"), points=points, expected_values=expected ) - def test_mute_strips(self, patch_ones): - """Mute strips along each dimension.""" - # Mute a time strip - muted_time = patch_ones.mute(time=(0.5, 1.0)) - # All distances should be affected equally - assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) - - # Mute a distance strip - muted_dist = patch_ones.mute(distance=(10, 20)) - # All times should be affected equally - assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) - # From 673e0ec50c65a16c3a3f4a41e56b555eaecb2154 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Tue, 21 Oct 2025 17:36:51 +0100 Subject: [PATCH 06/15] update on mute --- dascore/proc/mute.py | 42 ++++++++++++++++++++++-------------- tests/test_proc/test_mute.py | 29 +++++++++++++++---------- 2 files changed, 44 insertions(+), 27 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index f11622dd3..f142a9c91 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -128,7 +128,9 @@ class _MuteGeometry2D(_MuteGeometry): Parameters ---------- origin - The shared origin for the two points. + The origin for each point. If lines are not parallel, this is the + shared origin. If they are parallel the first value of each point + is used. line1 A numpy array of the un-normalized first line. line2 @@ -146,7 +148,7 @@ class _MuteGeometry2D(_MuteGeometry): normalized (scaled) line """ - origin: NDArray[np.floating] + origin: tuple(NDArray[np.floating], NDArray[np.floating]) norm: NDArray[np.floating] line1_norm: NDArray[np.floating] line2_norm: NDArray[np.floating] @@ -164,10 +166,12 @@ def _get_line_params(cls, points, patch, dims): coord_norm = np.array( [dc.to_float(patch.get_coord(x).coord_range()) for x in dims] ) - origin = get_2d_line_intersection(*points) / coord_norm - # Get vectors. Need to normalize in coord space and by l2. - v1 = (points[1] - points[0]) / coord_norm - v2 = (points[3] - points[2]) / coord_norm + origin = get_2d_line_intersection(*points) + # Get vectors with various stages of normalization. + # We need to first normalize in coord space and then by l2 norm. + l1 = points[1] - points[0] + l2 = points[3] - points[2] + v1, v2 = l1 / coord_norm, l2 / coord_norm norm_v1, norm_v2 = norm(v1), norm(v2) v1_norm, v2_norm = v1 / norm_v1, v2 / norm_v2 # Check if any points are degenerate. @@ -184,8 +188,13 @@ def _get_line_params(cls, points, patch, dims): else: # If lines are parallel we can just reverse the direction of one. v1 *= -1 + # Get the origin tuple (origin for l1, origin for l2) + if parallel: + origin_tuple = (points[0], points[2]) + else: + origin_tuple = (origin, origin) out = dict( - origin=origin, + origin=origin_tuple, norm=coord_norm, line1_norm=v1 / norm_v1, line2_norm=v2 / norm_v2, @@ -238,28 +247,29 @@ def _get_normalized_array_coord(self, array, patch): Get an array that matches the dimensionality of envelope but of its normalized, relative coordinates. """ - out = [] + out = [[], []] for dim in self.dims: ax = patch.get_axis(dim) coord = patch.get_coord(dim) # Get the coordinate values normalized to coord range. coord_vals = dc.to_float(coord.values) coord_range = dc.to_float(coord.coord_range()) - if self.relative: + if self.relative or self.parallel: coord_vals = coord_vals - dc.to_float(coord.min()) - # First get values in coord. - norm_vals = (coord_vals - self.origin[ax]) / dc.to_float(coord_range) + # First get values in coord. Only relative to origin if there is one. + origin = 0.0 if self.parallel else self.origin[ax] + norm_vals = (coord_vals / dc.to_float(coord_range)) - origin # We need to transform coordinate values to values between 0 and 1. # with the same dimensionality as array. coord_inds = [None] * array.ndim - coord_inds[ax] = slice(None, len(coord)) + coord_inds[ax] = slice(None) norms = norm_vals[tuple(coord_inds)] # Next, we set those values on an array with the same shape as array. inds = [slice(None)] * array.ndim inds[ax] = slice(None, len(coord)) carray = np.empty_like(array) carray[tuple(inds)] = norms - out.append(carray) + out[0].append(carray) # Return new axis as -1 so it will broadcast with lines. out = np.stack(out, axis=-1) return out @@ -277,8 +287,8 @@ def _cross_product_ok(self, coord_array): line_widths = (0, 1) l1_z = np.pad(self.line1_norm, line_widths, mode="constant", constant_values=0) l2_z = np.pad(self.line2_norm, line_widths, mode="constant", constant_values=0) - # The selected points will have different cross product sign and - # the same dot product sign. + # The selected points will have different cross product signs. Need to + # shift points relative to each line start if we have parallel lines. cross1_z = np.cross(l1_z, coord_z)[:, :, 2] cross2_z = np.cross(l2_z, coord_z)[:, :, 2] ok_cross = cross1_z * cross2_z < 0 @@ -293,7 +303,7 @@ def _dot_product_ok(self, coord_array): def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: """Apply the mask to the output array.""" - coord_norms = self._get_normalized_array_coord(array, patch) + cnorms_1, cnorms_2 = self._get_normalized_array_coord(array, patch) ok_cross = self._cross_product_ok(coord_norms) # For parallel lines only cross product is needed to define region. if self.parallel: diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 937a7d5a2..592cff6cb 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -48,15 +48,14 @@ def _assert_point_values( for vals, expected in zip(points, expected_values, strict=True): # Get the index (in the data array) of the expected 0. - inds = ( - patch.get_coord(dims[num]).get_next_index( - vals[axes.index(num)], relative=relative - ) - if num in axes - else slice(None) - for num in range(patch.ndim) - ) - assert np.isclose(patch.data[tuple(inds)], expected) + inds = [slice(None)] * patch.ndim + for val, dim, ax in zip(vals, dims, axes, strict=True): + ax = patch.get_axis(dim) + coord = patch.get_coord(dim) + inds[ax] = coord.get_next_index(val, relative=relative) + + patch_value = patch.data[tuple(inds)] + assert np.isclose(patch_value, expected) @pytest.fixture(scope="session") @@ -211,7 +210,7 @@ def test_not_same_direction(self, patch_ones): ) def test_mute_lines(self, patch_ones): - """Mute non-parallel lines.""" + """Test for muting non-parallel lines.""" muted = patch_ones.mute( time=([0, 2], [0, 4]), distance=([0, 100], [0, 100]), @@ -219,7 +218,15 @@ def test_mute_lines(self, patch_ones): points = [(145, 4), (182, 6), (100, 3), (180, 3), (60, 6), (50, 1)] expected = [1, 1, 1, 0, 0, 0] _assert_point_values( - muted, ("time", "distance"), points=points, expected_values=expected + muted, ("distance", "time"), points=points, expected_values=expected + ) + + def test_mute_lines_parallel(self, patch_ones): + """Test for muting parallel lines.""" + breakpoint() + muted = patch_ones.mute( + time=([0, 7.0], [1, 8.0]), + distance=([0, 300], [1, 301]), ) From f153a60eb8b33855fe6cceb27a2cecc2b4935819 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Tue, 21 Oct 2025 21:40:45 +0100 Subject: [PATCH 07/15] update --- dascore/proc/mute.py | 85 +++++++++++++++++++----------------- dascore/utils/misc.py | 7 ++- tests/test_proc/test_mute.py | 41 +++++++++++++---- 3 files changed, 83 insertions(+), 50 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index f142a9c91..cf505f4bd 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -38,7 +38,7 @@ def _mask_array( self, array: NDArray, ): - pass + """Apply the mask on the array for envelope calculation.""" def _apply_smoothing(self, array, smooth, patch): """Apply smoothing to the array.""" @@ -148,7 +148,7 @@ class _MuteGeometry2D(_MuteGeometry): normalized (scaled) line """ - origin: tuple(NDArray[np.floating], NDArray[np.floating]) + origin: tuple[NDArray[np.floating], NDArray[np.floating]] norm: NDArray[np.floating] line1_norm: NDArray[np.floating] line2_norm: NDArray[np.floating] @@ -173,7 +173,8 @@ def _get_line_params(cls, points, patch, dims): l2 = points[3] - points[2] v1, v2 = l1 / coord_norm, l2 / coord_norm norm_v1, norm_v2 = norm(v1), norm(v2) - v1_norm, v2_norm = v1 / norm_v1, v2 / norm_v2 + with np.errstate(divide="ignore", invalid="ignore"): + v1_norm, v2_norm = v1 / norm_v1, v2 / norm_v2 # Check if any points are degenerate. if (norm_v1 < tol) or (norm_v2 < tol): msg = f"A line provided to mute ({v1} or {v2}) is degenerate!" @@ -229,6 +230,7 @@ def _get_coord_float_values(coord, vals, relative): row = _get_coord_float_values(coord, vals, relative=relative) fill_inds[dim].append(row) out.append(None) + continue # Otherwise, we just run with it. vals = _get_coord_float_values(coord, np.array(row), relative=relative) out.append(vals) @@ -248,68 +250,74 @@ def _get_normalized_array_coord(self, array, patch): normalized, relative coordinates. """ out = [[], []] - for dim in self.dims: - ax = patch.get_axis(dim) - coord = patch.get_coord(dim) - # Get the coordinate values normalized to coord range. - coord_vals = dc.to_float(coord.values) - coord_range = dc.to_float(coord.coord_range()) - if self.relative or self.parallel: - coord_vals = coord_vals - dc.to_float(coord.min()) - # First get values in coord. Only relative to origin if there is one. - origin = 0.0 if self.parallel else self.origin[ax] - norm_vals = (coord_vals / dc.to_float(coord_range)) - origin - # We need to transform coordinate values to values between 0 and 1. - # with the same dimensionality as array. - coord_inds = [None] * array.ndim - coord_inds[ax] = slice(None) - norms = norm_vals[tuple(coord_inds)] - # Next, we set those values on an array with the same shape as array. - inds = [slice(None)] * array.ndim - inds[ax] = slice(None, len(coord)) - carray = np.empty_like(array) - carray[tuple(inds)] = norms - out[0].append(carray) + for onum, origin in enumerate(self.origin): + for dim in self.dims: + ax = patch.get_axis(dim) + coord = patch.get_coord(dim) + # Get the coordinate values normalized to coord range. + coord_vals = dc.to_float(coord.values) + coord_range = dc.to_float(coord.coord_range()) + if self.relative: + coord_vals -= dc.to_float(coord.min()) + # First get values in coord. Only relative to origin if there is one. + norm_vals = (coord_vals - origin[ax]) / dc.to_float(coord_range) + # We need to transform coordinate values to values between 0 and 1. + # with the same dimensionality as array. + coord_inds = [None] * array.ndim + coord_inds[ax] = slice(None) + norms = norm_vals[tuple(coord_inds)] + # Next, we set those values on an array with the same shape as array. + inds = [slice(None)] * array.ndim + inds[ax] = slice(None, len(coord)) + carray = np.empty_like(array) + carray[tuple(inds)] = norms + out[onum].append(carray) # Return new axis as -1 so it will broadcast with lines. - out = np.stack(out, axis=-1) + out = [np.stack(x, axis=-1) for x in out] return out - def _cross_product_ok(self, coord_array): + def _cross_product_ok(self, coord_array_1, coord_array_2): """ Determine if each point is in the region based on the cross product requirement. """ # Get padded arrays with 0 z values for cross product. array_widths = ((0, 0), (0, 0), (0, 1)) - coord_z = np.pad( - coord_array, pad_width=array_widths, mode="constant", constant_values=0 + coord_1_z = np.pad( + coord_array_1, pad_width=array_widths, mode="constant", constant_values=0 + ) + coord_2_z = np.pad( + coord_array_2, pad_width=array_widths, mode="constant", constant_values=0 ) line_widths = (0, 1) l1_z = np.pad(self.line1_norm, line_widths, mode="constant", constant_values=0) l2_z = np.pad(self.line2_norm, line_widths, mode="constant", constant_values=0) # The selected points will have different cross product signs. Need to # shift points relative to each line start if we have parallel lines. - cross1_z = np.cross(l1_z, coord_z)[:, :, 2] - cross2_z = np.cross(l2_z, coord_z)[:, :, 2] + cross1_z = np.cross(l1_z, coord_1_z)[:, :, 2] + cross2_z = np.cross(l2_z, coord_2_z)[:, :, 2] ok_cross = cross1_z * cross2_z < 0 return ok_cross - def _dot_product_ok(self, coord_array): + def _dot_product_ok(self, coord_array_1, coord_array_2): # Get dot product with each line. - dot1 = np.sum(self.line1_norm[None, :] * coord_array, axis=-1) - dot2 = np.sum(self.line2_norm[None, :] * coord_array, axis=-1) + dot1 = np.sum(self.line1_norm[None, :] * coord_array_1, axis=-1) + dot2 = np.sum(self.line2_norm[None, :] * coord_array_2, axis=-1) ok_dot = dot1 * dot2 > 0 return ok_dot def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: """Apply the mask to the output array.""" cnorms_1, cnorms_2 = self._get_normalized_array_coord(array, patch) - ok_cross = self._cross_product_ok(coord_norms) + ok_cross = self._cross_product_ok(cnorms_1, cnorms_2) # For parallel lines only cross product is needed to define region. if self.parallel: - return ok_cross - ok_dot = self._dot_product_ok(coord_norms) - return ok_cross & ok_dot + mask = ok_cross + else: + ok_dot = self._dot_product_ok(cnorms_1, cnorms_2) + mask = ok_cross & ok_dot + array[mask] = fill_value + return array def _get_mute_geometry(patch, kwargs, relative=True): @@ -468,5 +476,4 @@ def mute( # Apply smoothing if requested. if smooth is not None: out = geo._apply_smoothing(out, smooth, patch) - return patch.update(data=patch.data * out) diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 1a42b1096..525d1f7f3 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -920,6 +920,9 @@ def get_2d_line_intersection(p1, p2, p3, p4): if np.isclose(denom, 0): np.array([np.nan, np.nan]) - px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denom - py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denom + num_x = (x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4) + num_y = (x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4) + with np.errstate(divide="ignore", invalid="ignore"): + px = num_x / denom + py = num_y / denom return np.array([px, py]) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 592cff6cb..fa0d559cd 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -211,23 +211,46 @@ def test_not_same_direction(self, patch_ones): def test_mute_lines(self, patch_ones): """Test for muting non-parallel lines.""" - muted = patch_ones.mute( - time=([0, 2], [0, 4]), - distance=([0, 100], [0, 100]), - ) + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + dims = ("distance", "time") + muted = patch_ones.mute(time=time, distance=distance) + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + points = [(145, 4), (182, 6), (100, 3), (180, 3), (60, 6), (50, 1)] - expected = [1, 1, 1, 0, 0, 0] + expected = np.array([0, 0, 0, 1, 1, 1]) + _assert_point_values(muted, dims, points=points, expected_values=expected) _assert_point_values( - muted, ("distance", "time"), points=points, expected_values=expected + inverted, dims, points=points, expected_values=-expected + 1 ) def test_mute_lines_parallel(self, patch_ones): """Test for muting parallel lines.""" - breakpoint() + time = ([0, 7.0], [1, 8.0]) + distance = ([0, 300], [1, 301]) + muted = patch_ones.mute( + time=time, + distance=distance, + ) + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + + points = [(110, 3), (271, 7), (23, 1), (50, 6), (250, 1), (170, 3)] + expected = np.array([0, 0, 0, 1, 1, 1]) + dims = ("distance", "time") + _assert_point_values(muted, dims, points=points, expected_values=expected) + _assert_point_values( + inverted, dims, points=points, expected_values=-expected + 1 + ) + + def test_none_mutes(self, patch_ones): + """Ensure None can be used to specify vertical/horizontal lines.""" + time = (0, [1, 8.0]) + distance = (None, [1, 301]) muted = patch_ones.mute( - time=([0, 7.0], [1, 8.0]), - distance=([0, 300], [1, 301]), + time=time, + distance=distance, ) + # From a4e62d25217166609fba0bff2953c728a9e91e53 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Wed, 22 Oct 2025 17:11:18 +0100 Subject: [PATCH 08/15] implicit limits --- dascore/proc/mute.py | 148 +++++++++++++++++++++-------------- tests/test_proc/test_mute.py | 79 ++++++++++++++++--- 2 files changed, 158 insertions(+), 69 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index cf505f4bd..adf1c0d27 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -2,7 +2,6 @@ from __future__ import annotations -from collections import defaultdict from collections.abc import Mapping, Sized from typing import ClassVar @@ -127,7 +126,7 @@ class _MuteGeometry2D(_MuteGeometry): Parameters ---------- - origin + origins The origin for each point. If lines are not parallel, this is the shared origin. If they are parallel the first value of each point is used. @@ -148,7 +147,7 @@ class _MuteGeometry2D(_MuteGeometry): normalized (scaled) line """ - origin: tuple[NDArray[np.floating], NDArray[np.floating]] + origins: tuple[NDArray[np.floating], NDArray[np.floating]] norm: NDArray[np.floating] line1_norm: NDArray[np.floating] line2_norm: NDArray[np.floating] @@ -159,7 +158,7 @@ class _MuteGeometry2D(_MuteGeometry): _tolerance: ClassVar[float] = 1e-9 @classmethod - def _get_line_params(cls, points, patch, dims): + def _get_line_params(cls, points, patch, dims, fill_ind): """Get the parameters from a list of points.""" tol = cls._tolerance # Get the origin (where two lines intersect), (nan, nan) if parallel. @@ -175,40 +174,43 @@ def _get_line_params(cls, points, patch, dims): norm_v1, norm_v2 = norm(v1), norm(v2) with np.errstate(divide="ignore", invalid="ignore"): v1_norm, v2_norm = v1 / norm_v1, v2 / norm_v2 + # Then get array for easier manipulation + v_norms = np.stack([v1_norm, v2_norm], axis=0) + # Check if any points are degenerate. if (norm_v1 < tol) or (norm_v2 < tol): msg = f"A line provided to mute ({v1} or {v2}) is degenerate!" raise ParameterError(msg) # Determine if vectors are parallel and point in same direction - parallel = (1 - np.dot(v1, v2)) < tol - same_direction = np.dot(v1_norm, v2_norm) >= 0 + dot = np.dot(v_norms[0], v_norms[1]) + parallel = bool((1 - abs(dot)) < tol) + same_direction = bool(dot >= 0) if not same_direction: - if not parallel: + if parallel or fill_ind != -1: + preferred = fill_ind if fill_ind != -1 else 0 + # Just reverse order of line + v_norms[preferred] *= -1 + else: msg = "Non-parallel mute vectors must point in the same direction." raise ParameterError(msg) - else: - # If lines are parallel we can just reverse the direction of one. - v1 *= -1 + # Get the origin tuple (origin for l1, origin for l2) if parallel: origin_tuple = (points[0], points[2]) else: origin_tuple = (origin, origin) out = dict( - origin=origin_tuple, + origins=origin_tuple, norm=coord_norm, - line1_norm=v1 / norm_v1, - line2_norm=v2 / norm_v2, + line1_norm=v_norms[0], + line2_norm=v_norms[1], parallel=parallel, ) return out - @classmethod - def from_params(cls, vals, dims, axes, patch, relative=True): - """ - Return values which are two lines in absolute coordinate space, - (expressed as floats) - """ + @staticmethod + def _get_filled_value_list(patch, coords, value_list, relative): + """Create an array of points, swapping out implicit values.""" def _get_coord_float_values(coord, vals, relative): """Get the coordinate float values relative to start of coords.""" @@ -217,30 +219,52 @@ def _get_coord_float_values(coord, vals, relative): out = out - dc.to_float(coord.min()) return out - out = [] - # Keep track of the axis that need to be filled in (eg None, paired w/ float) - fill_inds = defaultdict(list) - for ind, (dim, row) in enumerate(zip(dims, vals)): - coord = patch.get_coord(dim) - # In the case of single values (eg None, ..., or some other) - # We need to just mark them and come back later. - if not isinstance(row, Sized): - # We need to get the actual value from the coord. - if row is not None and row is not Ellipsis: - row = _get_coord_float_values(coord, vals, relative=relative) - fill_inds[dim].append(row) - out.append(None) - continue - # Otherwise, we just run with it. - vals = _get_coord_float_values(coord, np.array(row), relative=relative) - out.append(vals) - # Now we can ascertain the line intended by None. - if fill_inds: - raise NotImplementedError("Working on it.") + # Output is dim by column and point by row. + out = np.full((4, 2), fill_value=np.nan, dtype=np.float64) + # Indicates an implicit value was used. + ifill_index = -1 + + for ind, (coord, row) in enumerate(zip(coords, value_list)): + coord = patch.get_coord(patch.dims[ind]) + # We iterate each pair because it might be an implicit value. + for pair_ind, pair in enumerate(row): + is_sized = isinstance(pair, Sized) + # The range is clearly defined. + if is_sized: + array_inds = (slice(pair_ind * 2, pair_ind * 2 + 2), ind) + out[array_inds] = _get_coord_float_values(coord, pair, relative) + continue + # This pair value is a place holder (None, ...) + elif not is_sized and (pair is None or pair is Ellipsis): + continue + # This is an implicit value; here is where things get crazy. + # We handle this by creating a new set of points to sub into + # the output array. This new set shares a first point with the + # other line, as well as the point on the implicit dimension. + # The value for the non-implicit dimension is held fixed. + else: + ifill_index = pair_ind + other_col_ind = 0 if ind == 1 else 1 + fixed_vals = _get_coord_float_values(coord, [pair], relative) + out[pair_ind * 2 : pair_ind * 2 + 2, ind] = fixed_vals + out[pair_ind * 2 : pair_ind * 2 + 2, other_col_ind] = np.array( + [0, 1] + ) + assert not np.any(np.isnan(out)) + return out, ifill_index + + @classmethod + def from_params(cls, vals, dims, axes, patch, relative=True): + """ + Return values which are two lines in absolute coordinate space, + (expressed as floats) + """ + coords = [patch.get_coord(x) for x in dims] + # Iterate over each row (consists of (start, stop)) + points, fill_ind = cls._get_filled_value_list(patch, coords, vals, relative) # Then put the 4 points together. This gives us a len 4 array with rows # as points from first line, then points from second. - points = np.stack(out, axis=-1).reshape(-1, 2) - line_params = cls._get_line_params(points, patch, dims) + line_params = cls._get_line_params(points, patch, dims, fill_ind) kwargs = dict(dims=dims, axes=axes, relative=relative) | line_params return cls(**kwargs) @@ -249,23 +273,30 @@ def _get_normalized_array_coord(self, array, patch): Get an array that matches the dimensionality of envelope but of its normalized, relative coordinates. """ + + def _nan_normalize(array): + """Normalize along last axis, handle norm close to 0.""" + array_norm = norm(array, axis=-1, keepdims=True) + array_norm[array_norm == 0] = -np.finfo(np.float64).min + return array / array_norm + out = [[], []] - for onum, origin in enumerate(self.origin): - for dim in self.dims: - ax = patch.get_axis(dim) - coord = patch.get_coord(dim) - # Get the coordinate values normalized to coord range. - coord_vals = dc.to_float(coord.values) - coord_range = dc.to_float(coord.coord_range()) - if self.relative: - coord_vals -= dc.to_float(coord.min()) - # First get values in coord. Only relative to origin if there is one. - norm_vals = (coord_vals - origin[ax]) / dc.to_float(coord_range) + for dim_num, dim in enumerate(self.dims): + ax = patch.get_axis(dim) + coord = patch.get_coord(dim) + # Get the coordinate values normalized to coord range. + coord_vals = dc.to_float(coord.values) + coord_range = dc.to_float(coord.coord_range()) + if self.relative: + coord_vals -= dc.to_float(coord.min()) + # Iterate over each origin. + for onum, origin in enumerate(self.origins): # We need to transform coordinate values to values between 0 and 1. # with the same dimensionality as array. - coord_inds = [None] * array.ndim - coord_inds[ax] = slice(None) - norms = norm_vals[tuple(coord_inds)] + norm_vals = (coord_vals - origin[dim_num]) / coord_range + bcast_inds = [None] * array.ndim + bcast_inds[ax] = slice(None) + norms = norm_vals[tuple(bcast_inds)] # Next, we set those values on an array with the same shape as array. inds = [slice(None)] * array.ndim inds[ax] = slice(None, len(coord)) @@ -273,8 +304,9 @@ def _get_normalized_array_coord(self, array, patch): carray[tuple(inds)] = norms out[onum].append(carray) # Return new axis as -1 so it will broadcast with lines. - out = [np.stack(x, axis=-1) for x in out] - return out + out_1 = [np.stack(x, axis=-1) for x in out] + out_norm = [_nan_normalize(x) for x in out_1] + return out_norm def _cross_product_ok(self, coord_array_1, coord_array_2): """ @@ -303,7 +335,7 @@ def _dot_product_ok(self, coord_array_1, coord_array_2): # Get dot product with each line. dot1 = np.sum(self.line1_norm[None, :] * coord_array_1, axis=-1) dot2 = np.sum(self.line2_norm[None, :] * coord_array_2, axis=-1) - ok_dot = dot1 * dot2 > 0 + ok_dot = (dot1 > 0) & (dot2 > 0) return ok_dot def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index fa0d559cd..96aa82bd8 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -190,6 +190,8 @@ def test_mute_strips(self, patch_ones): class TestMuteLines: """Tests for muting between lines.""" + _dims = ("distance", "time") + def test_point_raise(self, patch_ones): """A degenerate line (point) should raise.""" msg = "is degenerate" @@ -213,15 +215,14 @@ def test_mute_lines(self, patch_ones): """Test for muting non-parallel lines.""" time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) - dims = ("distance", "time") muted = patch_ones.mute(time=time, distance=distance) inverted = patch_ones.mute(time=time, distance=distance, invert=True) points = [(145, 4), (182, 6), (100, 3), (180, 3), (60, 6), (50, 1)] expected = np.array([0, 0, 0, 1, 1, 1]) - _assert_point_values(muted, dims, points=points, expected_values=expected) + _assert_point_values(muted, self._dims, points=points, expected_values=expected) _assert_point_values( - inverted, dims, points=points, expected_values=-expected + 1 + inverted, self._dims, points=points, expected_values=-expected + 1 ) def test_mute_lines_parallel(self, patch_ones): @@ -236,21 +237,77 @@ def test_mute_lines_parallel(self, patch_ones): points = [(110, 3), (271, 7), (23, 1), (50, 6), (250, 1), (170, 3)] expected = np.array([0, 0, 0, 1, 1, 1]) - dims = ("distance", "time") - _assert_point_values(muted, dims, points=points, expected_values=expected) + _assert_point_values(muted, self._dims, points=points, expected_values=expected) _assert_point_values( - inverted, dims, points=points, expected_values=-expected + 1 + inverted, self._dims, points=points, expected_values=-expected + 1 ) - def test_none_mutes(self, patch_ones): - """Ensure None can be used to specify vertical/horizontal lines.""" - time = (0, [1, 8.0]) - distance = (None, [1, 301]) + def test_implicit_with_positive_line(self, patch_ones): + """ + Ensure None can be used to specify vertical/horizontal lines with + another line with positive direction. + """ + time = (0, [0, 8.0]) + distance = (None, [0, 301]) muted = patch_ones.mute( time=time, distance=distance, ) + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + points = [(50, 6), (250, 2)] + expected = np.array([1, 0]) + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + _assert_point_values( + inverted, self._dims, points=points, expected_values=-expected + 1 + ) + def test_implicit_with_negative_line(self, patch_ones): + """ + Ensure None can be used to specify vertical/horizontal lines with + another line pointing in negative direction. + """ + time = (4, [4.0, 0]) + distance = (None, [150, 0]) + muted = patch_ones.mute( + time=time, + distance=distance, + ) + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + + points = [(50, 3), (10, 3), (50, 2), (100, 2)] + expected = np.array([0, 0, 0, 1]) + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + _assert_point_values( + inverted, self._dims, points=points, expected_values=-expected + 1 + ) -# + def test_two_implicit_parallel_lines(self, patch_ones): + """Ensure two implicit values can define parallel lines.""" + time = (0, 2) + distance = (None, None) + muted = patch_ones.mute( + time=time, + distance=distance, + ) + sub = muted.select(time=(2.1, ...), relative=True) + assert np.allclose(sub.data, 1) + + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + sub = inverted.select(time=(2.1, ...), relative=True) + assert np.allclose(sub.data, 0) + + def test_two_implicit_orthogonal_lines(self, patch_ones): + """Ensure two implicit values can define orthogonal lines.""" + time = (None, 2) + distance = (100, None) + muted = patch_ones.mute( + time=time, + distance=distance, + ) + sub = muted.select(time=(2.1, ...), distance=(101, ...), relative=True) + assert np.allclose(sub.data, 0) + + inverted = patch_ones.mute(time=time, distance=distance, invert=True) + sub = inverted.select(time=(2.1, ...), distance=(101, ...), relative=True) + assert np.allclose(sub.data, 1) From b3765539c0566e52c111d436972b42a7d4ed61a1 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Wed, 22 Oct 2025 18:22:52 +0100 Subject: [PATCH 09/15] progress on muter --- dascore/examples.py | 18 +++++++++--------- dascore/proc/mute.py | 12 ++++++------ dascore/proc/taper.py | 4 +--- tests/test_examples.py | 6 ------ tests/test_proc/test_mute.py | 9 +-------- 5 files changed, 17 insertions(+), 32 deletions(-) diff --git a/dascore/examples.py b/dascore/examples.py index 790231ae9..3fb24c7a6 100644 --- a/dascore/examples.py +++ b/dascore/examples.py @@ -441,7 +441,7 @@ def _ricker(time, delay): @register_func(EXAMPLE_PATCHES, key="delta_patch") def delta_patch( - dim: tuple[str, ...] | str = ("time", "distance"), + dim="time", shape=(10, 200), time_min="2020-01-01", time_step=1 / 250, @@ -457,21 +457,21 @@ def delta_patch( Parameters ---------- - dim + dim : str The dimension at the center of which to place the unit value. - Typically, ``"time"`` or ``"distance"``. - shape + Typically ``"time"`` or ``"distance"``. + shape : tuple of int The shape of the data as (distance, time). Defaults to (10, 200). This is used only if no existing ``patch`` is provided. - time_min + time_min : str or datetime64 The start time of the patch. - time_step + time_step : float The time step in seconds between samples. - distance_min + distance_min : float The minimum distance coordinate. - distance_step + distance_step : float The distance step in meters between samples. - patch + patch : dascore.Patch If provided, creates the delta patch based on this existing patch. Default is None. """ diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index adf1c0d27..f7c1b5d93 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -88,7 +88,7 @@ def _convert_to_samples(smooth, dims, patch): smooth_by_dims = _broadcast_smooth_to_dims(self.dims, smooth) smooth_ints = _convert_to_samples(smooth_by_dims, self.dims, patch) - # Now finagle into input for scipy's gaussian smooth. + # Now convert to input format for scipy's gaussian filter. return smooth_ints @@ -374,7 +374,7 @@ def _get_mute_geometry(patch, kwargs, relative=True): geometry = _MuteGeometry1D.from_params( val_list, dims, axes, patch, relative=relative ) - elif len(dims) > 1: # Dealing with lines. . + elif len(dims) > 1: # Dealing with lines. geometry = _MuteGeometry2D.from_params( val_list, dims, axes, patch, relative=relative ) @@ -401,7 +401,7 @@ def mute( patch The patch instance. smooth - Parameter controlling smoothing of the mute evenlope. Defines the sigma + Parameter controlling smoothing of the mute envelope. Defines the sigma Can be: - None: sharp mute - float (0.0-1.0): fraction of dimension range (e.g., 0.01 = 1%) @@ -446,7 +446,7 @@ def mute( >>> muted = patch.mute( ... time=(0, [0, 0.3]), ... distance=(None, [0, 300]), - ... taper=0.02, + ... smooth=0.02, ... ) >>> >>> # Mute late arrivals: from velocity line to end @@ -487,14 +487,14 @@ def mute( - Currently, mute doesn't support more than 2 dimensions. - For more control over boundary smoothing, use a patch with one values - then apply custom tapering/smooting before multiplying with the + then apply custom tapering/smoothing before multiplying with the original patch. See example section for more details. See Also -------- - [`Patch.select`](`dascore.Patch.select`) - [`Patch.taper_range`](`dascore.Patch.taper_range`) - - [`Patch.gaussian_filter`](`dacore.proc.filter.gaussian_filter`) + - [`Patch.gaussian_filter`](`dascore.proc.filter.gaussian_filter`) """ # Get geometry object to set up the problem. geo = _get_mute_geometry(patch, kwargs, relative) diff --git a/dascore/proc/taper.py b/dascore/proc/taper.py index 0ebae3d75..78933f4f7 100644 --- a/dascore/proc/taper.py +++ b/dascore/proc/taper.py @@ -153,9 +153,7 @@ def _get_taper_coord_inds(coord, values, relative, samples): # None or ... means min_val in first half of list else max_val out[num] = 0 if (num / len(out)) < 0.5 else len(coord) - 1 else: - out[num] = coord.get_next_index( - val, samples=samples, relative=relative, allow_out_of_bounds=True - ) + out[num] = coord.get_next_index(val, samples=samples, relative=relative) # Always need a len 4 sequence if len(out) == 2: out = [0, *out, len(coord)] diff --git a/tests/test_examples.py b/tests/test_examples.py index 19a5397b9..a964007af 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -73,12 +73,6 @@ def test_moveout(self): class TestDeltaPatch: """Tests for the delta_patch example.""" - @pytest.mark.parametrize("shape", ((10, 10), (100, 100), (1, 10))) - def test_shape(self, shape): - """Ensure the shape parameter controls the shape of the patch.""" - patch = dc.get_example_patch("delta_patch", shape=shape) - assert patch.shape == shape - @pytest.mark.parametrize("invalid_dim", ["inv_dim", "", None, 123, 1.1]) def test_delta_patch_invalid_dim(self, invalid_dim): """ diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 96aa82bd8..2708fc25f 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -5,7 +5,6 @@ import numpy as np import pytest -import dascore as dc from dascore.exceptions import ParameterError @@ -64,12 +63,6 @@ def patch_ones(random_patch): return random_patch.new(data=np.ones_like(random_patch.data)) -@pytest.fixture(scope="session") -def ricker_patch(): - """Return ricker moveout patch for velocity mute testing.""" - return dc.get_example_patch("ricker_moveout") - - class TestMuteBasics: """Basic tests for mute functionality.""" @@ -205,7 +198,7 @@ def test_not_same_direction(self, patch_ones): """Create lines which do not point in the same directions.""" match = "point in the same direction" - with pytest.raises(ValueError, match=match): + with pytest.raises(ParameterError, match=match): patch_ones.mute( time=[[0, 0], [1, -1]], distance=[[1, -1], [-1, 1]], From c6cbd52f48918ec985b98587e778c213b1eae1f6 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Wed, 22 Oct 2025 21:56:57 +0100 Subject: [PATCH 10/15] tests for absolute --- dascore/proc/mute.py | 35 ++-- tests/test_proc/test_basic.py | 10 ++ tests/test_proc/test_mute.py | 304 +++++++++++++++++++++++++++++++++- tests/test_proc/test_taper.py | 8 +- 4 files changed, 334 insertions(+), 23 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index f7c1b5d93..9de1e925e 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -41,8 +41,8 @@ def _mask_array( def _apply_smoothing(self, array, smooth, patch): """Apply smoothing to the array.""" - sigma = self._get_smooth_sigma(smooth, patch) - return gaussian_filter(array, sigma=sigma, axes=tuple(self.axes)) + sigma, axes = self._get_smooth_sigma(smooth, patch) + return gaussian_filter(array, sigma=sigma, axes=axes) def _get_smooth_sigma(self, smooth, patch): """Get sigma values, in samples, for the gaussian kernel.""" @@ -53,17 +53,20 @@ def _broadcast_smooth_to_dims(dims, smooth): # For a single value, just broadcast to dim length. if not isinstance(smooth, Mapping): vals = [smooth] * len(dims) + axes = self.axes else: - # Otherwise each dimension's smooth must be specified. - if not set(dims) == set(smooth): + # Otherwise, the smooth dict must be a subset of the dimensions. + if not set(smooth).issubset(set(dims)): msg = ( - f"If a taper dictionary is used in Mute, it must have all " - f"the same keys as the dimensions. Kwarg dims are {dims} and" - f"taper keys are {list(smooth)}." + f"If a smooth dictionary is used in Mute, it must be a " + f"subset of the dims in kwargs. Kwarg dims are {dims} and" + f"smooth keys are {list(smooth)}." ) raise ParameterError(msg) - vals = [smooth[dim] for dim in dims] - return vals + vals = [smooth[dim] for dim in dims if dim in smooth] + axes = [self.dims.index(x) for x in smooth] + + return vals, axes def _convert_to_samples(smooth, dims, patch): """Convert the smooth parameter to number of samples.""" @@ -86,10 +89,10 @@ def _convert_to_samples(smooth, dims, patch): out.append(coord.get_sample_count(val)) return out - smooth_by_dims = _broadcast_smooth_to_dims(self.dims, smooth) + smooth_by_dims, axes = _broadcast_smooth_to_dims(self.dims, smooth) smooth_ints = _convert_to_samples(smooth_by_dims, self.dims, patch) # Now convert to input format for scipy's gaussian filter. - return smooth_ints + return smooth_ints, axes class _MuteGeometry1D(_MuteGeometry): @@ -314,7 +317,8 @@ def _cross_product_ok(self, coord_array_1, coord_array_2): requirement. """ # Get padded arrays with 0 z values for cross product. - array_widths = ((0, 0), (0, 0), (0, 1)) + # Create padding for all dimensions: (0,0) for all but last, (0,1) for last + array_widths = [(0, 0)] * (coord_array_1.ndim - 1) + [(0, 1)] coord_1_z = np.pad( coord_array_1, pad_width=array_widths, mode="constant", constant_values=0 ) @@ -326,8 +330,9 @@ def _cross_product_ok(self, coord_array_1, coord_array_2): l2_z = np.pad(self.line2_norm, line_widths, mode="constant", constant_values=0) # The selected points will have different cross product signs. Need to # shift points relative to each line start if we have parallel lines. - cross1_z = np.cross(l1_z, coord_1_z)[:, :, 2] - cross2_z = np.cross(l2_z, coord_2_z)[:, :, 2] + # Use ellipsis to handle arbitrary dimensions, indexing last dim with -1 + cross1_z = np.cross(l1_z, coord_1_z)[..., 2] + cross2_z = np.cross(l2_z, coord_2_z)[..., 2] ok_cross = cross1_z * cross2_z < 0 return ok_cross @@ -436,7 +441,7 @@ def mute( >>> muted = patch.mute(time=(0, 0.5)) >>> >>> # Mute everything except middle section - >>> kept = patch.mute(time=(0.2, -0.2), mode="complement") + >>> kept = patch.mute(time=(0.2, -0.2), invert=True) >>> >>> # 1D Mute with smoothed absolute units for time. >>> muted = patch.mute(time=(0.2, 0.8), smooth=0.02 * dc.units.s) diff --git a/tests/test_proc/test_basic.py b/tests/test_proc/test_basic.py index e3d74dddc..af7731e29 100644 --- a/tests/test_proc/test_basic.py +++ b/tests/test_proc/test_basic.py @@ -824,3 +824,13 @@ def test_no_op(self, random_patch): """Ensure passing no dims does nothing.""" out = random_patch.flip() assert out is random_patch + + +class TestFull: + """Tests for creating patches filled with a single value.""" + + def test_full_1(self, random_patch): + """Ensure a patch can be created with 1s.""" + patch = random_patch.full(1.0) + patch.coords == random_patch.coords + assert np.allclose(patch.data, 1.0) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 2708fc25f..83643dea8 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -5,6 +5,7 @@ import numpy as np import pytest +import dascore as dc from dascore.exceptions import ParameterError @@ -14,7 +15,7 @@ def _get_testable_coord_values(coord, relative=False): start, stop = coord[start_ind], coord[stop_ind] if relative: start, stop = start - coord.min(), stop - coord.min() - return (start, stop) + return (dc.to_float(start), dc.to_float(stop)) def _assert_coord_ranges( @@ -63,6 +64,33 @@ def patch_ones(random_patch): return random_patch.new(data=np.ones_like(random_patch.data)) +@pytest.fixture(scope="module") +def patch_ones_3d(): + """Return a 3D patch filled with ones for testing.""" + data = np.ones((20, 30, 10)) + coords = { + "time": np.arange(20) * 0.1, + "distance": np.arange(30) * 10.0, + "depth": np.arange(10) * 5.0, + } + dims = ("time", "distance", "depth") + return dc.Patch(data=data, coords=coords, dims=dims) + + +@pytest.fixture(scope="module") +def patch_ones_4d(): + """Return a 4D patch filled with ones for testing.""" + data = np.ones((15, 20, 8, 6)) + coords = { + "time": np.arange(15) * 0.1, + "distance": np.arange(20) * 10.0, + "depth": np.arange(8) * 5.0, + "angle": np.arange(6) * 15.0, + } + dims = ("time", "distance", "depth", "angle") + return dc.Patch(data=data, coords=coords, dims=dims) + + class TestMuteBasics: """Basic tests for mute functionality.""" @@ -98,11 +126,12 @@ def test_1d_mute_no_taper(self, patch_ones): coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) muted = patch_ones.mute(time=(v1, v2), relative=True) + step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted, dim="time", zero_ranges=[(v1, v2)], - one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], + one_ranges=[(..., v1 - step), (v2 + step, ...)], relative=True, ) @@ -111,11 +140,12 @@ def test_mute_open_interval(self, patch_ones): coord = patch_ones.get_coord("distance") v1, v2 = _get_testable_coord_values(coord, relative=True) muted1 = patch_ones.mute(distance=(v2, ...), relative=True) + step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted1, dim="distance", zero_ranges=[(v2, ...)], - one_ranges=[(..., v2 - coord.step)], + one_ranges=[(..., v2 - step)], relative=True, ) @@ -124,11 +154,12 @@ def test_mute_absolute(self, patch_ones): coord = patch_ones.get_coord("distance") v1, v2 = _get_testable_coord_values(coord, relative=False) muted1 = patch_ones.mute(distance=(v1, v2), relative=False) + step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted1, dim="distance", zero_ranges=[(v1, v2)], - one_ranges=[(..., v1 - coord.step), (v2 + coord.step, ...)], + one_ranges=[(..., v1 - step), (v2 + step, ...)], relative=False, ) @@ -304,3 +335,268 @@ def test_two_implicit_orthogonal_lines(self, patch_ones): inverted = patch_ones.mute(time=time, distance=distance, invert=True) sub = inverted.select(time=(2.1, ...), distance=(101, ...), relative=True) assert np.allclose(sub.data, 1) + + def test_mute_lines_absolute(self, patch_ones_3d): + """Test 2D line mute with absolute coordinates (relative=False).""" + # Use 3D patch which has simple float coordinates (time: 0-1.9, distance: 0-290) + # Get the actual coordinate values from the patch + time_coord = patch_ones_3d.get_coord("time") + dist_coord = patch_ones_3d.get_coord("distance") + + time_min = dc.to_float(time_coord.min()) + dist_min = dc.to_float(dist_coord.min()) + + # Define two non-parallel lines using absolute coordinates + # Line 1: from (t=0, d=0) to (t=1.0, d=100) + # Line 2: from (t=0, d=0) to (t=1.5, d=100) + time = ([time_min, time_min + 1.0], [time_min, time_min + 1.5]) + distance = ([dist_min, dist_min + 100], [dist_min, dist_min + 100]) + + muted = patch_ones_3d.mute(time=time, distance=distance, relative=False) + inverted = patch_ones_3d.mute( + time=time, distance=distance, relative=False, invert=True + ) + + # Verify muting worked + assert np.any(muted.data == 0) + assert np.any(muted.data == 1) + assert np.any(inverted.data == 0) + assert np.any(inverted.data == 1) + + # Test some specific points using absolute coordinates + # Points that should be inside the muted region (between the two lines) + # Format: (distance, time) + # For 3D patch, selecting 2D point gives array along 3rd dimension + point_inside = (dist_min + 100, time_min + 1.2) + point_outside = (dist_min + 150, time_min + 0.5) + + # Get indices for inside point + dist_idx_in = dist_coord.get_next_index(point_inside[0], relative=False) + time_idx_in = time_coord.get_next_index(point_inside[1], relative=False) + # For 3D patch with dims (time, distance, depth), mute on time-distance + # affects all depths + assert np.allclose(muted.data[time_idx_in, dist_idx_in, :], 0) + assert np.allclose(inverted.data[time_idx_in, dist_idx_in, :], 1) + + # Get indices for outside point + dist_idx_out = dist_coord.get_next_index(point_outside[0], relative=False) + time_idx_out = time_coord.get_next_index(point_outside[1], relative=False) + assert np.allclose(muted.data[time_idx_out, dist_idx_out, :], 1) + assert np.allclose(inverted.data[time_idx_out, dist_idx_out, :], 0) + + +class TestMuteSmoothing: + """Tests for smoothing functionality in mute.""" + + def test_dict_smooth_2d_lines(self, patch_ones): + """Test smoothing with dict parameter for dimension-specific values.""" + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + smooth = {"time": 0.02, "distance": 5} + muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + + # With smoothing, muted region should have intermediate values + assert muted.shape == patch_ones.shape + # Check that we have values between 0 and 1 (not just sharp cutoff) + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_dict_smooth_with_none(self, patch_ones): + """Test dict with None value for one dimension.""" + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + smooth = {"time": 0.02, "distance": None} + muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + assert muted.shape == patch_ones.shape + + def test_dict_smooth_mismatched_keys_raises(self, patch_ones): + """Test that dict with mismatched keys raises appropriate error.""" + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + smooth = {"time": 0.02, "oriely": 0.2} # Missing distance key + msg = "a smooth dictionary" + with pytest.raises(ParameterError, match=msg): + patch_ones.mute(time=time, distance=distance, smooth=smooth) + + def test_smooth_creates_gradual_transition(self, patch_ones): + """Verify smoothing creates gradual transitions with intermediate values.""" + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted = patch_ones.mute(time=(v1, v2), relative=True, smooth=0.05) + + # Check that the transition region has values between 0 and 1 + transition_region = muted.select(time=(v1 - 0.2, v2 + 0.2), relative=True).data + # There should be gradual transition, not just 0s and 1s + unique_vals = np.unique(transition_region) + assert len(unique_vals) > 10 # Many intermediate values + assert np.any((transition_region > 0.01) & (transition_region < 0.99)) + + def test_smooth_with_invert(self, patch_ones): + """Test that smoothing works with invert=True.""" + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + smooth = 0.03 + muted = patch_ones.mute( + time=time, distance=distance, smooth=smooth, invert=True + ) + + # With smoothing and invert, should have intermediate values + assert muted.shape == patch_ones.shape + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_smooth_parallel_lines(self, patch_ones): + """Test smoothing with parallel lines.""" + time = ([0, 7.0], [1, 8.0]) + distance = ([0, 300], [1, 301]) + smooth = {"time": 5, "distance": 10} + muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + + assert muted.shape == patch_ones.shape + # Check for smooth transition + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_smooth_very_small_value(self, patch_ones): + """Test edge case with very small smooth value (1 sample).""" + coord = patch_ones.get_coord("time") + v1, v2 = _get_testable_coord_values(coord, relative=True) + muted = patch_ones.mute(time=(v1, v2), relative=True, smooth=1) + # Should still work, just with minimal smoothing + assert muted.shape == patch_ones.shape + + def test_smooth_preserves_shape(self, patch_ones): + """Test that smoothing preserves patch shape.""" + time = ([0, 2], [0, 4]) + distance = ([0, 100], [0, 100]) + smooth = 0.05 + muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + assert muted.shape == patch_ones.shape + + +class TestMute3D: + """Tests for muting on 3D patches.""" + + def test_1d_mute_time(self, patch_ones_3d): + """Test 1D mute along time dimension.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_3d.mute(time=(v1, v2), relative=True) + + # Check that the muted region is zeroed + sub_muted = muted.select(time=(v1, v2), relative=True) + assert np.allclose(sub_muted.data, 0) + + # Check that outside region is still ones + sub_kept = muted.select(time=(v2 + 0.1, ...), relative=True) + assert np.allclose(sub_kept.data, 1) + + def test_1d_mute_distance(self, patch_ones_3d): + """Test 1D mute along distance dimension.""" + muted = patch_ones_3d.mute(distance=(50, 150), relative=True) + sub_muted = muted.select(distance=(50, 150), relative=True) + assert np.allclose(sub_muted.data, 0) + + def test_2d_line_mute(self, patch_ones_3d): + """Test 2D line mute using time-distance plane.""" + time = ([0, 1.0], [0, 1.5]) + distance = ([0, 100], [0, 100]) + muted = patch_ones_3d.mute(time=time, distance=distance) + + # Should maintain 3D shape + assert muted.shape == patch_ones_3d.shape + assert muted.ndim == 3 + + # Check that some region is muted + assert np.any(muted.data == 0) + # Check that some region is not muted + assert np.any(muted.data == 1) + + def test_2d_line_mute_distance_depth(self, patch_ones_3d): + """Test 2D line mute using distance-depth plane.""" + distance = ([0, 100], [0, 200]) + depth = ([0, 20], [0, 30]) + muted = patch_ones_3d.mute(distance=distance, depth=depth) + + assert muted.shape == patch_ones_3d.shape + assert np.any(muted.data == 0) + assert np.any(muted.data == 1) + + def test_smoothing_3d(self, patch_ones_3d): + """Test that smooth parameter works correctly in 3D.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_3d.mute(time=(v1, v2), relative=True, smooth=0.05) + # With smoothing, should have intermediate values + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_invert_3d(self, patch_ones_3d): + """Test invert parameter in 3D.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_3d.mute(time=(v1, v2), relative=True, invert=True) + + # With invert, the selected region should be kept + sub_kept = muted.select(time=(v1, v2), relative=True) + assert np.allclose(sub_kept.data, 1) + + # Outside region should be zeroed + sub_zeroed = muted.select(time=(v2 + 0.1, ...), relative=True) + assert np.allclose(sub_zeroed.data, 0) + + +class TestMute4D: + """Tests for muting on 4D patches.""" + + def test_1d_mute_time(self, patch_ones_4d): + """Test 1D mute along time dimension in 4D patch.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_4d.mute(time=(v1, v2), relative=True) + + # Check that the muted region is zeroed + sub_muted = muted.select(time=(v1, v2), relative=True) + assert np.allclose(sub_muted.data, 0) + + # Check that outside region is still ones + sub_kept = muted.select(time=(v2 + 0.1, ...), relative=True) + assert np.allclose(sub_kept.data, 1) + + def test_1d_mute_angle(self, patch_ones_4d): + """Test 1D mute along angle dimension.""" + muted = patch_ones_4d.mute(angle=(20, 60), relative=True) + sub_muted = muted.select(angle=(20, 60), relative=True) + assert np.allclose(sub_muted.data, 0) + + def test_2d_line_mute_depth_angle(self, patch_ones_4d): + """Test 2D line mute using depth-angle plane.""" + depth = ([0, 15], [0, 25]) + angle = ([0, 30], [0, 60]) + muted = patch_ones_4d.mute(depth=depth, angle=angle) + + assert muted.shape == patch_ones_4d.shape + assert np.any(muted.data == 0) + assert np.any(muted.data == 1) + + def test_smoothing_4d_with_dict(self, patch_ones_4d): + """Test smooth parameter with dict in 4D.""" + time = ([0, 0.8], [0, 1.2]) + distance = ([0, 80], [0, 80]) + smooth = {"time": 0.02, "distance": 5} + muted = patch_ones_4d.mute(time=time, distance=distance, smooth=smooth) + # With smoothing, should have intermediate values + assert muted.shape == patch_ones_4d.shape + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_smoothing_1d_mute_in_4d(self, patch_ones_4d): + """Test smoothing with 1D mute in 4D patch.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_4d.mute(time=(v1, v2), relative=True, smooth=0.05) + # Check for gradual transition + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_invert_4d(self, patch_ones_4d): + """Test invert parameter in 4D.""" + v1, v2 = 0.3, 0.8 + muted = patch_ones_4d.mute(time=(v1, v2), relative=True, invert=True) + + # With invert, the selected region should be kept + sub_kept = muted.select(time=(v1, v2), relative=True) + assert np.allclose(sub_kept.data, 1) + + # Outside region should be zeroed + sub_zeroed = muted.select(time=(v2 + 0.1, ...), relative=True) + assert np.allclose(sub_zeroed.data, 0) diff --git a/tests/test_proc/test_taper.py b/tests/test_proc/test_taper.py index 600d28e0b..d56d5ddef 100644 --- a/tests/test_proc/test_taper.py +++ b/tests/test_proc/test_taper.py @@ -298,14 +298,14 @@ def test_two_value_invert_mutes_to_last_sample(self, patch_ones): # Use 2-value form with invert to mute from start_idx to end out = patch_ones.taper_range( - distance=(start_idx, end_idx), + distance=(start_idx - 1, start_idx, end_idx, end_idx), invert=False, samples=True, window_type="boxcar", ) # Values before start_idx should be 1 (unmuted) - assert np.allclose(out.data[start_idx - 1, :], 1) + assert np.allclose(out.data[start_idx - 10, :], 0) # Values from start_idx to end should be 0 (muted) - assert np.allclose(out.data[start_idx, :], 0) - assert np.allclose(out.data[end_idx, :], 0) + assert np.allclose(out.data[start_idx, :], 1) + assert np.allclose(out.data[end_idx - 1, :], 1) From 84362c270eb1ba624183469595c2eb6ad7454f05 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Thu, 23 Oct 2025 11:12:07 +0100 Subject: [PATCH 11/15] use abc --- dascore/proc/mute.py | 14 +++++++------- tests/test_proc/test_mute.py | 6 ------ 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 9de1e925e..0aba629b2 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -2,6 +2,7 @@ from __future__ import annotations +from abc import ABC, abstractmethod from collections.abc import Mapping, Sized from typing import ClassVar @@ -20,7 +21,7 @@ from dascore.utils.patch import get_dim_axis_value, patch_function -class _MuteGeometry(DascoreBaseModel): +class _MuteGeometry(ABC, DascoreBaseModel): """ Parent class for Mute Geometry. """ @@ -30,14 +31,13 @@ class _MuteGeometry(DascoreBaseModel): relative: bool = True @classmethod + @abstractmethod def from_params(cls, vals, dims, axes, patch, relative): """Initialize Mute Geometry from input parameters.""" - def _mask_array( - self, - array: NDArray, - ): - """Apply the mask on the array for envelope calculation.""" + @abstractmethod + def _apply_mask(self, array: NDArray, patch: dc.Patch, fill_value) -> NDArray: + """Apply the mask to the output array.""" def _apply_smoothing(self, array, smooth, patch): """Apply smoothing to the array.""" @@ -280,7 +280,7 @@ def _get_normalized_array_coord(self, array, patch): def _nan_normalize(array): """Normalize along last axis, handle norm close to 0.""" array_norm = norm(array, axis=-1, keepdims=True) - array_norm[array_norm == 0] = -np.finfo(np.float64).min + array_norm[array_norm == 0] = np.finfo(np.float64).eps return array / array_norm out = [[], []] diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 83643dea8..723fe4791 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -487,12 +487,6 @@ def test_1d_mute_time(self, patch_ones_3d): sub_kept = muted.select(time=(v2 + 0.1, ...), relative=True) assert np.allclose(sub_kept.data, 1) - def test_1d_mute_distance(self, patch_ones_3d): - """Test 1D mute along distance dimension.""" - muted = patch_ones_3d.mute(distance=(50, 150), relative=True) - sub_muted = muted.select(distance=(50, 150), relative=True) - assert np.allclose(sub_muted.data, 0) - def test_2d_line_mute(self, patch_ones_3d): """Test 2D line mute using time-distance plane.""" time = ([0, 1.0], [0, 1.5]) From 3bbffe0a53c4c9f04c06ab03ad725596f71f1c47 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Fri, 24 Oct 2025 11:28:46 +0100 Subject: [PATCH 12/15] deprecate notch in slope_filter, add slope_mute skeleton --- dascore/proc/filter.py | 19 ++++++--- dascore/proc/mute.py | 77 ++++++++++++++++++++++++++++------ tests/test_proc/test_filter.py | 11 +++-- 3 files changed, 86 insertions(+), 21 deletions(-) diff --git a/dascore/proc/filter.py b/dascore/proc/filter.py index b764e4dd7..ba0c306b6 100644 --- a/dascore/proc/filter.py +++ b/dascore/proc/filter.py @@ -8,6 +8,7 @@ from __future__ import annotations import sys +import warnings from collections.abc import Sequence import numpy as np @@ -438,13 +439,13 @@ def gaussian_filter( @patch_function() -@compose_docstring(sample_explanation=samples_arg_description) def slope_filter( patch: PatchType, filt: Sequence[float], dims: tuple[str, str] = ("distance", "time"), directional: bool = False, - notch: bool = False, + notch: bool | None = None, + invert: bool = False, ) -> PatchType: """ Filter the patch over certain slopes in the 2D Fourier domain. @@ -472,6 +473,8 @@ def slope_filter( This can be used for up/down or left/right separation, assuming a near-linear fiber layout. notch + Deprecated, use invert. + invert If True, the filter represents a notch, meaning the slopes specified by the inner `filt` parameters are attenuated rather than those outside of them. @@ -527,7 +530,7 @@ def _check_inputs(patch, filt, dims): msg = f"Cant apply slope filter. {missing} are missing from patch." raise ParameterError(msg) - def _get_taper_mask(filt, slope, notch): + def _get_taper_mask(filt, slope, invert): """Get a mask for applying taper and attenuation.""" fac = np.where( (slope >= filt[0]) & (slope <= filt[1]), @@ -540,7 +543,7 @@ def _get_taper_mask(filt, slope, notch): np.sin(0.5 * np.pi * (slope - filt[2]) / (filt[3] - filt[2])), fac, ) - fac = fac if notch else 1.0 - fac + fac = fac if invert else 1.0 - fac return fac def _get_slope_array(dft_patch, directional, freq_dims): @@ -596,7 +599,13 @@ def _maybe_transform_units(filt, dft_patch, freq_dims): slope = _get_slope_array(dft_patch, directional, freq_dims) filt = _maybe_transform_units(filt, dft_patch, freq_dims) - mask = _get_taper_mask(filt, slope, notch) + # TODO remove in dascore 0.2. + if notch is not None: + msg = "The `notch` parameter of slope filter is deprecated. Use invert." + warnings.warn(msg, DeprecationWarning) + invert = notch + + mask = _get_taper_mask(filt, slope, invert) new_data = dft_patch.data * mask out = dft_patch.update(data=new_data) if transformed: diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 0aba629b2..44d8a6a52 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -7,19 +7,38 @@ from typing import ClassVar import numpy as np +from numpy.linalg import norm from numpy.typing import NDArray -from scipy.linalg import norm from scipy.ndimage import gaussian_filter import dascore as dc from dascore.constants import PatchType from dascore.exceptions import ParameterError +from dascore.utils.docs import compose_docstring from dascore.utils.misc import ( get_2d_line_intersection, ) from dascore.utils.models import DascoreBaseModel from dascore.utils.patch import get_dim_axis_value, patch_function +_smooth_param = """ +smooth + Parameter controlling smoothing of the mute envelope. Defines the sigma + Can be: + - None: sharp mute + - float (0.0-1.0): fraction of dimension range (e.g., 0.01 = 1%) + which is applied independently to each dimension involved in the + mute. + - int: Indicates number of samples for each dimension. + - Quantity with units, indicates values along a single dimension. + Only applicable if a single dimension is specified. + - dict: {dim: taper_value} for dimension-specific smooth + values which can be any of the above. +""" + + +_smooth_type = None | float | int | tuple(float | int) | dict[str, float | int] + class _MuteGeometry(ABC, DascoreBaseModel): """ @@ -387,6 +406,7 @@ def _get_mute_geometry(patch, kwargs, relative=True): @patch_function() +@compose_docstring(smooth_param=_smooth_param) def mute( patch: PatchType, *, @@ -405,18 +425,7 @@ def mute( ---------- patch The patch instance. - smooth - Parameter controlling smoothing of the mute envelope. Defines the sigma - Can be: - - None: sharp mute - - float (0.0-1.0): fraction of dimension range (e.g., 0.01 = 1%) - which is applied independently to each dimension involved in the - mute. - - int: Indicates number of samples for each dimension. - - Quantity with units, indicates values along a single dimension. - Only applicable if a single dimension is specified. - - dict: {dim: taper_value} for dimension-specific smooth - values which can be any of the above. + {smooth_param} invert If True, invert the taper such that the values outside the defined region are set to 0. @@ -514,3 +523,45 @@ def mute( if smooth is not None: out = geo._apply_smoothing(out, smooth, patch) return patch.update(data=patch.data * out) + + +@patch_function() +@compose_docstring(_smooth_param=_smooth_param) +def slope_mute( + patch: PatchType, + slopes: tuple[float, float] | NDArray, + *, + dims: tuple[str, str] = ("distance", "time"), + smooth: float | None = None, + invert: bool = False, +) -> PatchType: + """ + Apply a mute between specified slopes (eg velocities). + + This facilitates common muting patterns for active source processing. + For more control, use [`Patch.mute`](`dascore.proc.mute.mute`). + + Parameters + ---------- + patch + The patch to filter. + slopes + A length 2 sequence which specifies the begining and ending slope + values. The dims parameter specifies how the slope is calculated. + dims + The dimensions used to determine slope. The first dim is in the + numerator and the second in the denominator. (eg distance, time) + represents a velocity since distance/time has units of |L|/|T| + (commonly m/s). + {_smooth_param} + invert + If True, invert the mute, meaning areas outside the given region + are muted + + See Also + -------- + - [`Patch.slope_filter`](`dascore.proc.filter.slope_filter`) + - [`Patch.mute`](`dascore.proc.mute.mute`) + + The [FK recipe](`docs/recipes/fk.qmd`) provides addtional examples. + """ diff --git a/tests/test_proc/test_filter.py b/tests/test_proc/test_filter.py index b9b264b48..e7e870ccb 100644 --- a/tests/test_proc/test_filter.py +++ b/tests/test_proc/test_filter.py @@ -403,13 +403,18 @@ def test_directional_filter(self, example_patch): assert isinstance(filtered_patch, dc.Patch) - def test_notch_filter(self, example_patch): - """Ensure notching can be performed with slope filter.""" + def test_invert_filter(self, example_patch): + """Ensure the slope filter can be inverted.""" filtered_patch = example_patch.slope_filter( - filt=[2e3, 2.2e3, 8e3, 2e4], notch=True + filt=[2e3, 2.2e3, 8e3, 2e4], invert=True ) assert isinstance(filtered_patch, dc.Patch) + def test_notch_deprecated(self, example_patch): + """Ensure using notch param issues deprecation warning.""" + with pytest.warns(DeprecationWarning): + example_patch.slope_filter(filt=[2e3, 2.2e3, 8e3, 2e4], notch=True) + def test_different_params_not_equal(self, example_patch): """Ensure filter is sensitive to different parameters.""" filtered_patch1 = example_patch.slope_filter(filt=[1e3, 1.5e3, 5e3, 1e4]) From e6664d5f5c4ec0c6491bf975615ac9678ad731d3 Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Sat, 25 Oct 2025 11:18:56 +0100 Subject: [PATCH 13/15] slope mute with tests --- dascore/core/patch.py | 1 + dascore/proc/__init__.py | 2 +- dascore/proc/mute.py | 41 ++++++- tests/test_proc/test_mute.py | 212 +++++++++++++++++++++++++++++++++++ 4 files changed, 252 insertions(+), 4 deletions(-) diff --git a/dascore/core/patch.py b/dascore/core/patch.py index 02bdf1518..f4135d295 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -440,6 +440,7 @@ def iresample(self, *args, **kwargs): taper = dascore.proc.taper taper_range = dascore.proc.taper_range mute = dascore.proc.mute + slope_mute = dascore.proc.slope_mute rolling = dascore.proc.rolling whiten = dascore.proc.whiten diff --git a/dascore/proc/__init__.py b/dascore/proc/__init__.py index 6172e8007..fffa72a96 100644 --- a/dascore/proc/__init__.py +++ b/dascore/proc/__init__.py @@ -12,7 +12,7 @@ from .resample import decimate, interpolate, resample from .rolling import rolling from .taper import taper, taper_range -from .mute import mute +from .mute import mute, slope_mute from .units import convert_units, set_units, simplify_units from .whiten import whiten from .hampel import hampel_filter diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 44d8a6a52..eec7d4c93 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -37,7 +37,7 @@ """ -_smooth_type = None | float | int | tuple(float | int) | dict[str, float | int] +_smooth_type = None | float | int | tuple[float | int, ...] | dict[str, float | int] class _MuteGeometry(ABC, DascoreBaseModel): @@ -558,10 +558,45 @@ def slope_mute( If True, invert the mute, meaning areas outside the given region are muted + Notes + ----- + - Assumes the mute origin is at the start of each of the selected dimensions. + You can use [`Patch.flip`](`dascore.proc.basic.flip`) to set the other + end of a dimension to be the origin. + See Also -------- - [`Patch.slope_filter`](`dascore.proc.filter.slope_filter`) - [`Patch.mute`](`dascore.proc.mute.mute`) - - The [FK recipe](`docs/recipes/fk.qmd`) provides addtional examples. """ + # Convert slopes to array and validate + slopes_array = np.asarray(slopes) + if slopes_array.shape != (2,): + msg = "slopes must be a tuple or array of length 2" + raise ParameterError(msg) + # Check for zero or negative slopes + if np.any(slopes_array <= 0): + msg = "slopes must be positive (greater than zero)" + raise ParameterError(msg) + # Get the coordinate ranges for the dimensions (in relative coordinates) + coord0 = patch.get_coord(dims[0]) + coord1 = patch.get_coord(dims[1]) + range0 = dc.to_float(coord0.coord_range()) + # The origin is always at the *start* of the coordinate. + origin = (dc.to_float(coord0[0]), dc.to_float(coord1[0])) + # Create endpoints for lines defined by each slope + endpoints = [] + for slope in slopes_array: + # Slope is rise over run; we just need to multiply by run + # Note: it is fine if points exceed patch dims; mute will handle it. + point = (origin[0] + range0, origin[1] + slope * range0) + endpoints.append(point) + # Build kwargs for mute + mute_kwargs = { + dims[0]: ([origin[0], endpoints[0][0]], [origin[0], endpoints[1][0]]), + dims[1]: ([origin[1], endpoints[0][1]], [origin[1], endpoints[1][1]]), + "smooth": smooth, + "invert": invert, + "relative": False, + } + return mute(patch, **mute_kwargs) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 723fe4791..3412d4e1b 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -594,3 +594,215 @@ def test_invert_4d(self, patch_ones_4d): # Outside region should be zeroed sub_zeroed = muted.select(time=(v2 + 0.1, ...), relative=True) assert np.allclose(sub_zeroed.data, 0) + + +class TestSlopeMute: + """Tests for slope_mute functionality.""" + + _dims = ("distance", "time") + + def test_slope_mute_basic(self, patch_ones): + """Test basic slope_mute with two velocities.""" + # Mute between two velocities (20 m/s and 30 m/s) + # These values are reasonable for the example patch which has + # distance range ~300m and time range ~8s + slopes = (20.0, 30.0) + muted = patch_ones.slope_mute(slopes=slopes) + + # Should return a patch with same shape + assert muted.shape == patch_ones.shape + + # Test specific points using _assert_point_values + # Format: (distance, time) in relative coordinates + # At time=4s: 20 m/s line is at 80m, 30 m/s line is at 120m + # Points between the two slope lines should be muted (0) + # Points outside should be unmuted (1) + points = [ + (50, 4.0), # Low velocity (12.5 m/s, below both lines) - unmuted + (100, 4.0), # Between the two velocities (25 m/s) - muted + (150, 4.0), # High velocity (37.5 m/s, above both lines) - unmuted + (80, 2.0), # Between slopes at earlier time (40 m/s) - unmuted + (50, 2.0), # Very low velocity (25 m/s) - muted + ] + expected = [1, 0, 1, 1, 0] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + + def test_slope_mute_with_array(self, patch_ones): + """Test slope_mute accepts numpy array.""" + slopes = np.array([20.0, 30.0]) + muted = patch_ones.slope_mute(slopes=slopes) + assert muted.shape == patch_ones.shape + + def test_slope_mute_wrong_length_raises(self, patch_ones): + """Test that wrong length slopes raises error.""" + with pytest.raises(ParameterError, match="length 2"): + patch_ones.slope_mute(slopes=(20.0,)) + with pytest.raises(ParameterError, match="length 2"): + patch_ones.slope_mute(slopes=(20.0, 30.0, 40.0)) + + def test_slope_mute_invert(self, patch_ones): + """Test slope_mute with invert parameter.""" + slopes = (20.0, 30.0) + muted = patch_ones.slope_mute(slopes=slopes, invert=False) + inverted = patch_ones.slope_mute(slopes=slopes, invert=True) + + # Test that regions are inverted + points = [ + (50, 4.0), # Outside mute region (low velocity) + (100, 4.0), # Inside mute region (between slopes) + (150, 4.0), # Outside mute region (high velocity) + ] + expected_muted = [1, 0, 1] + expected_inverted = [0, 1, 0] + _assert_point_values( + muted, self._dims, points=points, expected_values=expected_muted + ) + _assert_point_values( + inverted, self._dims, points=points, expected_values=expected_inverted + ) + + def test_slope_mute_with_smooth(self, patch_ones): + """Test slope_mute with smoothing parameter.""" + slopes = (20.0, 30.0) + muted = patch_ones.slope_mute(slopes=slopes, smooth=0.02) + + # With smoothing, should have intermediate values + assert muted.shape == patch_ones.shape + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_slope_mute_custom_dims(self, patch_ones_3d): + """Test slope_mute with custom dimension specification.""" + # Use time-distance instead of distance-time + # slope = time/distance (inverse velocity, or slowness) + slopes = (0.001, 0.002) # s/m + muted = patch_ones_3d.slope_mute(slopes=slopes, dims=("time", "distance")) + + assert muted.shape == patch_ones_3d.shape + # Verify muting occurred - should have both 0s and 1s + assert np.any(muted.data == 0) + assert np.any(muted.data == 1) + + # For slope = time/distance: at distance=200m (relative), mute region is + # time = 0.001*200 = 0.2s to 0.002*200 = 0.4s (relative) + # At time=0.3s relative (which is 15% of 2.0s range), dist=200m should be muted + time_idx = muted.get_coord("time").get_next_index(0.3, relative=True) + dist_idx = muted.get_coord("distance").get_next_index(200, relative=True) + # For 3D patch with dims (time, distance, depth), all depth values should + # be the same. + assert np.allclose(muted.data[time_idx, dist_idx, :], 0) + + def test_slope_mute_preserves_unmuted_dimensions(self, patch_ones_3d): + """Test that slope_mute only affects specified dimensions.""" + slopes = (20.0, 30.0) + muted = patch_ones_3d.slope_mute(slopes=slopes, dims=("distance", "time")) + + # All depth slices should be affected identically + assert muted.shape == patch_ones_3d.shape + assert np.array_equal(muted.data[:, :, 0], muted.data[:, :, 1]) + assert np.array_equal(muted.data[:, :, 0], muted.data[:, :, -1]) + + def test_slope_mute_steep_slopes(self, patch_ones): + """Test slope_mute with very steep slopes (high velocities).""" + # Use slopes near the max for the patch (37.5 m/s) + slopes = (35.0, 37.0) + muted = patch_ones.slope_mute(slopes=slopes) + + assert muted.shape == patch_ones.shape + # With very high velocities, most of patch should be unmuted + # Only a small wedge near upper right should be muted + points = [ + (250, 7.0), # High distance, late time, between slopes - muted + (100, 4.0), # Lower velocity - unmuted + (200, 2.0), # Far higher velocity - unmuted + ] + expected = [0, 1, 1] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + + def test_slope_mute_shallow_slopes(self, patch_ones): + """Test slope_mute with very shallow slopes (low velocities).""" + slopes = (5.0, 10.0) + muted = patch_ones.slope_mute(slopes=slopes) + + assert muted.shape == patch_ones.shape + # With very low velocities, mute region should be at low distances + points = [ + (20, 4.0), # Low distance, velocity=5 m/s - unmuted (below slopes) + (30, 4.0), # velocity=7.5 m/s, between slopes - muted + (50, 4.0), # velocity=12.5 m/s, above slopes - unmuted + (15, 2.0), # velocity=7.5 m/s, between slopes - muted + ] + expected = [1, 0, 1, 0] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + + def test_slope_mute_zero_slope_raises(self, patch_ones): + """Test that zero slope raises appropriate error.""" + slopes = (0.0, 30.0) + with pytest.raises(ParameterError, match="positive"): + patch_ones.slope_mute(slopes=slopes) + + def test_slope_mute_negative_slopes(self, patch_ones): + """Test slope_mute behavior with negative slopes.""" + slopes = (-20.0, -30.0) + # Negative slopes should cause error + with pytest.raises(ParameterError, match="positive"): + patch_ones.slope_mute(slopes=slopes) + + def test_slope_mute_matches_manual_mute(self, patch_ones): + """Test that slope_mute produces same result as manual mute call.""" + slopes = (20.0, 30.0) + + # Get coordinate ranges + dist_range = dc.to_float(patch_ones.get_coord("distance").coord_range()) + time_range = dc.to_float(patch_ones.get_coord("time").coord_range()) + + # Calculate endpoints manually + endpoints = [] + for slope in slopes: + dist_at_max_time = slope * time_range + if dist_at_max_time <= dist_range: + endpoint = (dist_at_max_time, time_range) + else: + endpoint = (dist_range, dist_range / slope) + endpoints.append(endpoint) + + # Manual mute call + manual_muted = patch_ones.mute( + distance=([0, endpoints[0][0]], [0, endpoints[1][0]]), + time=([0, endpoints[0][1]], [0, endpoints[1][1]]), + relative=True, + ) + + # slope_mute call + slope_muted = patch_ones.slope_mute(slopes=slopes) + + # Should produce identical results + assert np.allclose(manual_muted.data, slope_muted.data) + + def test_slope_mute_with_dict_smooth(self, patch_ones): + """Test slope_mute with dict smooth parameter.""" + slopes = (20.0, 30.0) + smooth = {"distance": 5, "time": 0.01} + muted = patch_ones.slope_mute(slopes=slopes, smooth=smooth) + + # Should have smooth transition + assert np.any((muted.data > 0.01) & (muted.data < 0.99)) + + def test_slope_mute_specific_velocity_wedge(self, patch_ones): + """Test muting specific velocity wedge with known points.""" + # Define precise velocity wedge + slopes = (15.0, 25.0) # m/s + muted = patch_ones.slope_mute(slopes=slopes) + + # Test points at specific locations + # At time=4.0s: + # - 15 m/s line is at distance=60m + # - 25 m/s line is at distance=100m + # So between 60-100m should be muted + points = [ + (50, 4.0), # Below 15 m/s line - unmuted + (80, 4.0), # Between 15-25 m/s - muted + (90, 4.0), # Between 15-25 m/s - muted + (110, 4.0), # Above 25 m/s line - unmuted + ] + expected = [1, 0, 0, 1] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) From 84a5ce62ff364b7427129f8a1fa8f4ea44e5575f Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Mon, 27 Oct 2025 17:32:10 +0000 Subject: [PATCH 14/15] slope_mute and benchmarks --- benchmarks/test_patch_benchmarks.py | 16 +++ dascore/core/patch.py | 2 +- dascore/proc/__init__.py | 2 +- dascore/proc/mute.py | 98 ++++++++++------ tests/test_proc/test_mute.py | 170 ++++++++++++++-------------- 5 files changed, 166 insertions(+), 122 deletions(-) diff --git a/benchmarks/test_patch_benchmarks.py b/benchmarks/test_patch_benchmarks.py index 4ce495da5..b01a67101 100644 --- a/benchmarks/test_patch_benchmarks.py +++ b/benchmarks/test_patch_benchmarks.py @@ -145,6 +145,22 @@ def test_wiener_filter(self, example_patch): patch = example_patch patch.wiener_filter(time=3, samples=True) + @pytest.mark.benchmark + def test_line_mute(self, example_patch): + """Time 2D line mute with velocity lines.""" + patch = example_patch + patch.line_mute( + time=(0, [0, 0.3]), + distance=(None, [0, 300]), + smooth=0.02, + ) + + @pytest.mark.benchmark + def test_slope_mute(self, example_patch): + """Time slope mute between velocities.""" + patch = example_patch + patch.slope_mute(slopes=(1000, 3000)) + class TestTransformBenchmarks: """Benchmarks for patch transform operations.""" diff --git a/dascore/core/patch.py b/dascore/core/patch.py index f4135d295..fbba00484 100644 --- a/dascore/core/patch.py +++ b/dascore/core/patch.py @@ -439,7 +439,7 @@ def iresample(self, *args, **kwargs): standardize = dascore.proc.standardize taper = dascore.proc.taper taper_range = dascore.proc.taper_range - mute = dascore.proc.mute + line_mute = dascore.proc.line_mute slope_mute = dascore.proc.slope_mute rolling = dascore.proc.rolling whiten = dascore.proc.whiten diff --git a/dascore/proc/__init__.py b/dascore/proc/__init__.py index fffa72a96..40fc327db 100644 --- a/dascore/proc/__init__.py +++ b/dascore/proc/__init__.py @@ -12,7 +12,7 @@ from .resample import decimate, interpolate, resample from .rolling import rolling from .taper import taper, taper_range -from .mute import mute, slope_mute +from .mute import line_mute, slope_mute from .units import convert_units, set_units, simplify_units from .whiten import whiten from .hampel import hampel_filter diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index eec7d4c93..9fbcc6157 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -309,8 +309,9 @@ def _nan_normalize(array): # Get the coordinate values normalized to coord range. coord_vals = dc.to_float(coord.values) coord_range = dc.to_float(coord.coord_range()) - if self.relative: - coord_vals -= dc.to_float(coord.min()) + # Since the points have already been converted to relative, + # even if they were absolute, we must do the same here. + coord_vals -= dc.to_float(coord.min()) # Iterate over each origin. for onum, origin in enumerate(self.origins): # We need to transform coordinate values to values between 0 and 1. @@ -407,7 +408,7 @@ def _get_mute_geometry(patch, kwargs, relative=True): @patch_function() @compose_docstring(smooth_param=_smooth_param) -def mute( +def line_mute( patch: PatchType, *, smooth=None, @@ -447,36 +448,36 @@ def mute( >>> patch = dc.get_example_patch().full(1) >>> >>> # Mute first 0.5s (relative to start by default) - >>> muted = patch.mute(time=(0, 0.5)) + >>> muted = patch.line_mute(time=(0, 0.5)) >>> >>> # Mute everything except middle section - >>> kept = patch.mute(time=(0.2, -0.2), invert=True) + >>> kept = patch.line_mute(time=(0.2, -0.2), invert=True) >>> >>> # 1D Mute with smoothed absolute units for time. - >>> muted = patch.mute(time=(0.2, 0.8), smooth=0.02 * dc.units.s) + >>> muted = patch.line_mute(time=(0.2, 0.8), smooth=0.02 * dc.units.s) >>> >>> # Classic first break mute: mute early arrivals >>> # Line from (t=0, d=0) to (t=0.3, d=300) defines velocity=1000 m/s - >>> muted = patch.mute( + >>> muted = patch.line_mute( ... time=(0, [0, 0.3]), ... distance=(None, [0, 300]), ... smooth=0.02, ... ) >>> >>> # Mute late arrivals: from velocity line to end - >>> muted = patch.mute( + >>> muted = patch.line_mute( ... time=([0, 0.3], None), ... distance=([0, 300], 0), ... ) >>> >>> # Mute wedge between two velocity lines - >>> muted = patch.mute( + >>> muted = patch.line_mute( ... time=([0, 0.375], [0, 0.25]), ... distance=([0, 300], [0, 300]), ... ) >>> >>> # Mute wedge outside two velocity lines - >>> muted = patch.mute( + >>> muted = patch.line_mute( ... time=([0, 0.375], [0, 0.25]), ... distance=([0, 300], [0, 300]), ... invert=True, @@ -484,7 +485,7 @@ def mute( >>> >>> # Apply custom tapering >>> ones = patch.full(1.0) - >>> envelope = ones.mute( + >>> envelope = ones.line_mute( ... time=([0, 0.375], [0, 0.25]), ... distance=([0, 300], [0, 300]), ... ) @@ -509,6 +510,7 @@ def mute( - [`Patch.select`](`dascore.Patch.select`) - [`Patch.taper_range`](`dascore.Patch.taper_range`) - [`Patch.gaussian_filter`](`dascore.proc.filter.gaussian_filter`) + - [`Patch.slope_mute`](`dascore.Patch.slope_mute`) """ # Get geometry object to set up the problem. geo = _get_mute_geometry(patch, kwargs, relative) @@ -546,7 +548,7 @@ def slope_mute( patch The patch to filter. slopes - A length 2 sequence which specifies the begining and ending slope + A length 2 sequence which specifies the beginning and ending slope values. The dims parameter specifies how the slope is calculated. dims The dimensions used to determine slope. The first dim is in the @@ -558,6 +560,30 @@ def slope_mute( If True, invert the mute, meaning areas outside the given region are muted + Examples + -------- + >>> import dascore as dc + >>> patch = dc.get_example_patch() + >>> + >>> # Mute data between velocities of 1000 m/s and 3000 m/s + >>> muted = patch.slope_mute(slopes=(1000, 3000)) + >>> + >>> # Keep only data between two velocities (invert the mute) + >>> kept = patch.slope_mute(slopes=(1500, 2500), invert=True) + >>> + >>> # Apply smoothing to mute boundaries (5% of dimension range) + >>> smooth_mute = patch.slope_mute(slopes=(1000, 3000), smooth=0.05) + >>> + >>> # Use different dimension order (e.g., time/distance for slowness) + >>> # This mutes between slownesses of 0.0003 s/m and 0.001 s/m + >>> slowness_mute = patch.slope_mute( + ... slopes=(0.0003, 0.001), dims=("time", "distance") + ... ) + >>> + >>> # Flip distance axis to set origin at far end before muting + >>> flipped = patch.flip("distance") + >>> muted_flipped = flipped.slope_mute(slopes=(1000, 3000)) + Notes ----- - Assumes the mute origin is at the start of each of the selected dimensions. @@ -567,36 +593,42 @@ def slope_mute( See Also -------- - [`Patch.slope_filter`](`dascore.proc.filter.slope_filter`) - - [`Patch.mute`](`dascore.proc.mute.mute`) + - [`Patch.line_mute`](`dascore.proc.mute.mute`) """ # Convert slopes to array and validate slopes_array = np.asarray(slopes) if slopes_array.shape != (2,): - msg = "slopes must be a tuple or array of length 2" + msg = "slopes must be a sequence of length 2" raise ParameterError(msg) # Check for zero or negative slopes - if np.any(slopes_array <= 0): - msg = "slopes must be positive (greater than zero)" + if np.any(slopes_array < 0): + msg = "slopes must be positive." raise ParameterError(msg) - # Get the coordinate ranges for the dimensions (in relative coordinates) - coord0 = patch.get_coord(dims[0]) - coord1 = patch.get_coord(dims[1]) - range0 = dc.to_float(coord0.coord_range()) - # The origin is always at the *start* of the coordinate. - origin = (dc.to_float(coord0[0]), dc.to_float(coord1[0])) - # Create endpoints for lines defined by each slope - endpoints = [] - for slope in slopes_array: - # Slope is rise over run; we just need to multiply by run - # Note: it is fine if points exceed patch dims; mute will handle it. - point = (origin[0] + range0, origin[1] + slope * range0) - endpoints.append(point) - # Build kwargs for mute + # Get the coordinate information + coord_x, coord_y = (patch.get_coord(x, require_sorted=True) for x in dims) + origin = (dc.to_float(coord_x[0]), dc.to_float(coord_y[0])) + # Here we take the full range as to allow it to be negative if the + # coord is reversed sorted. + range_x = dc.to_float(coord_x[-1] - coord_x[0]) + # Calculate endpoints for lines defined by each slope + dim0_vals = [[origin[0], origin[0] + range_x], [origin[0], origin[0] + range_x]] + dim1_vals = [[origin[1], None], [origin[1], None]] + for num, slope in enumerate(slopes_array): + # Special case for 0 and inf velocities + if np.isclose(slope, 0): + dim0_vals[num][1] = origin[0] + dim1_vals[num][1] = origin[1] + 1 + elif np.isinf(slope): + dim0_vals[num][1] = origin[0] + 1 + dim1_vals[num][1] = origin[1] + else: + # Slope is rise over run (e.g., distance/time for velocity) + dim1_vals[num][1] = origin[1] + range_x / slope mute_kwargs = { - dims[0]: ([origin[0], endpoints[0][0]], [origin[0], endpoints[1][0]]), - dims[1]: ([origin[1], endpoints[0][1]], [origin[1], endpoints[1][1]]), + dims[0]: dim0_vals, + dims[1]: dim1_vals, "smooth": smooth, "invert": invert, "relative": False, } - return mute(patch, **mute_kwargs) + return line_mute.func(patch, **mute_kwargs) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 3412d4e1b..11a18b8b0 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -91,33 +91,33 @@ def patch_ones_4d(): return dc.Patch(data=data, coords=coords, dims=dims) -class TestMuteBasics: +class TestLineMuteBasics: """Basic tests for mute functionality.""" def test_mute_no_kwargs_raises(self, random_patch): """Mute without dimension specifications should raise.""" with pytest.raises(ParameterError, match="one or two keyword"): - random_patch.mute() + random_patch.line_mute() def test_not_tuple_raises(self, random_patch): """Boundary must be tuple of length 2.""" with pytest.raises(ParameterError, match="two boundaries when using"): - random_patch.mute(time=5) + random_patch.line_mute(time=5) def test_tuple_wrong_length_raises(self, random_patch): """Tuple must have exactly 2 elements.""" with pytest.raises(ParameterError, match="two boundaries when using"): - random_patch.mute(time=(1, 2, 3)) + random_patch.line_mute(time=(1, 2, 3)) def test_smooth_bad_float(self, random_patch): """If a floating point value is used, it must be between 0 and 1.""" with pytest.raises(ParameterError, match="smooth parameter for"): - random_patch.mute(time=(1, 2), smooth=1.1) + random_patch.line_mute(time=(1, 2), smooth=1.1) with pytest.raises(ParameterError, match="smooth parameter for"): - random_patch.mute(time=(1, 2), smooth=-0.01) + random_patch.line_mute(time=(1, 2), smooth=-0.01) -class Test1DMute: +class Test1DLineMute: """Test 1D block mutes (single dimension).""" def test_1d_mute_no_taper(self, patch_ones): @@ -125,7 +125,7 @@ def test_1d_mute_no_taper(self, patch_ones): # Use relative=True (default) with numeric offset coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted = patch_ones.mute(time=(v1, v2), relative=True) + muted = patch_ones.line_mute(time=(v1, v2), relative=True) step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted, @@ -139,7 +139,7 @@ def test_mute_open_interval(self, patch_ones): """Mute using None for interval ends.""" coord = patch_ones.get_coord("distance") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted1 = patch_ones.mute(distance=(v2, ...), relative=True) + muted1 = patch_ones.line_mute(distance=(v2, ...), relative=True) step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted1, @@ -153,7 +153,7 @@ def test_mute_absolute(self, patch_ones): """Test absolute coordinates work.""" coord = patch_ones.get_coord("distance") v1, v2 = _get_testable_coord_values(coord, relative=False) - muted1 = patch_ones.mute(distance=(v1, v2), relative=False) + muted1 = patch_ones.line_mute(distance=(v1, v2), relative=False) step = dc.to_float(coord.step) _assert_coord_ranges( patch=muted1, @@ -168,7 +168,7 @@ def test_single_float_smooth(self, patch_ones): """Test that taper mute works with a single floating point value.""" coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=0.01) + muted1 = patch_ones.line_mute(time=(v1, v2), relative=True, smooth=0.01) # With taper, we can't assert exact zeros/ones, just check shape assert muted1.shape == patch_ones.shape # Check that middle region has been modified (not all ones) @@ -179,7 +179,7 @@ def test_single_int_smooth(self, patch_ones): """Test that smooth can be a single integer which means samples.""" coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=5) + muted1 = patch_ones.line_mute(time=(v1, v2), relative=True, smooth=5) # With taper, we can't assert exact zeros/ones, just check shape assert muted1.shape == patch_ones.shape # Check that middle region has been modified (not all ones) @@ -191,7 +191,7 @@ def test_single_quantity(self, patch_ones): coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) smooth_val = 0.01 * coord.units - muted1 = patch_ones.mute(time=(v1, v2), relative=True, smooth=smooth_val) + muted1 = patch_ones.line_mute(time=(v1, v2), relative=True, smooth=smooth_val) # With taper, we can't assert exact zeros/ones, just check shape assert muted1.shape == patch_ones.shape # Check that middle region has been modified (not all ones) @@ -201,12 +201,12 @@ def test_single_quantity(self, patch_ones): def test_mute_strips(self, patch_ones): """Mute strips along each dimension.""" # Mute a time strip - muted_time = patch_ones.mute(time=(0.5, 1.0)) + muted_time = patch_ones.line_mute(time=(0.5, 1.0)) # All distances should be affected equally assert np.allclose(muted_time.data[:, 10], muted_time.data[0, 10]) # Mute a distance strip - muted_dist = patch_ones.mute(distance=(10, 20)) + muted_dist = patch_ones.line_mute(distance=(10, 20)) # All times should be affected equally assert np.allclose(muted_dist.data[15, :], muted_dist.data[15, 0]) @@ -220,7 +220,7 @@ def test_point_raise(self, patch_ones): """A degenerate line (point) should raise.""" msg = "is degenerate" with pytest.raises(ParameterError, match=msg): - patch_ones.mute( + patch_ones.line_mute( time=([0, 0], [0, 0.25]), distance=([0, 0], [0, 300]), ) @@ -230,7 +230,7 @@ def test_not_same_direction(self, patch_ones): match = "point in the same direction" with pytest.raises(ParameterError, match=match): - patch_ones.mute( + patch_ones.line_mute( time=[[0, 0], [1, -1]], distance=[[1, -1], [-1, 1]], ) @@ -239,8 +239,8 @@ def test_mute_lines(self, patch_ones): """Test for muting non-parallel lines.""" time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) - muted = patch_ones.mute(time=time, distance=distance) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + muted = patch_ones.line_mute(time=time, distance=distance) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) points = [(145, 4), (182, 6), (100, 3), (180, 3), (60, 6), (50, 1)] expected = np.array([0, 0, 0, 1, 1, 1]) @@ -253,11 +253,11 @@ def test_mute_lines_parallel(self, patch_ones): """Test for muting parallel lines.""" time = ([0, 7.0], [1, 8.0]) distance = ([0, 300], [1, 301]) - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, ) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) points = [(110, 3), (271, 7), (23, 1), (50, 6), (250, 1), (170, 3)] expected = np.array([0, 0, 0, 1, 1, 1]) @@ -273,11 +273,11 @@ def test_implicit_with_positive_line(self, patch_ones): """ time = (0, [0, 8.0]) distance = (None, [0, 301]) - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, ) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) points = [(50, 6), (250, 2)] expected = np.array([1, 0]) @@ -293,11 +293,11 @@ def test_implicit_with_negative_line(self, patch_ones): """ time = (4, [4.0, 0]) distance = (None, [150, 0]) - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, ) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) points = [(50, 3), (10, 3), (50, 2), (100, 2)] expected = np.array([0, 0, 0, 1]) @@ -310,14 +310,14 @@ def test_two_implicit_parallel_lines(self, patch_ones): """Ensure two implicit values can define parallel lines.""" time = (0, 2) distance = (None, None) - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, ) sub = muted.select(time=(2.1, ...), relative=True) assert np.allclose(sub.data, 1) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) sub = inverted.select(time=(2.1, ...), relative=True) assert np.allclose(sub.data, 0) @@ -325,14 +325,14 @@ def test_two_implicit_orthogonal_lines(self, patch_ones): """Ensure two implicit values can define orthogonal lines.""" time = (None, 2) distance = (100, None) - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, ) sub = muted.select(time=(2.1, ...), distance=(101, ...), relative=True) assert np.allclose(sub.data, 0) - inverted = patch_ones.mute(time=time, distance=distance, invert=True) + inverted = patch_ones.line_mute(time=time, distance=distance, invert=True) sub = inverted.select(time=(2.1, ...), distance=(101, ...), relative=True) assert np.allclose(sub.data, 1) @@ -352,8 +352,8 @@ def test_mute_lines_absolute(self, patch_ones_3d): time = ([time_min, time_min + 1.0], [time_min, time_min + 1.5]) distance = ([dist_min, dist_min + 100], [dist_min, dist_min + 100]) - muted = patch_ones_3d.mute(time=time, distance=distance, relative=False) - inverted = patch_ones_3d.mute( + muted = patch_ones_3d.line_mute(time=time, distance=distance, relative=False) + inverted = patch_ones_3d.line_mute( time=time, distance=distance, relative=False, invert=True ) @@ -393,7 +393,7 @@ def test_dict_smooth_2d_lines(self, patch_ones): time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) smooth = {"time": 0.02, "distance": 5} - muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + muted = patch_ones.line_mute(time=time, distance=distance, smooth=smooth) # With smoothing, muted region should have intermediate values assert muted.shape == patch_ones.shape @@ -405,7 +405,7 @@ def test_dict_smooth_with_none(self, patch_ones): time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) smooth = {"time": 0.02, "distance": None} - muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + muted = patch_ones.line_mute(time=time, distance=distance, smooth=smooth) assert muted.shape == patch_ones.shape def test_dict_smooth_mismatched_keys_raises(self, patch_ones): @@ -415,13 +415,13 @@ def test_dict_smooth_mismatched_keys_raises(self, patch_ones): smooth = {"time": 0.02, "oriely": 0.2} # Missing distance key msg = "a smooth dictionary" with pytest.raises(ParameterError, match=msg): - patch_ones.mute(time=time, distance=distance, smooth=smooth) + patch_ones.line_mute(time=time, distance=distance, smooth=smooth) def test_smooth_creates_gradual_transition(self, patch_ones): """Verify smoothing creates gradual transitions with intermediate values.""" coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted = patch_ones.mute(time=(v1, v2), relative=True, smooth=0.05) + muted = patch_ones.line_mute(time=(v1, v2), relative=True, smooth=0.05) # Check that the transition region has values between 0 and 1 transition_region = muted.select(time=(v1 - 0.2, v2 + 0.2), relative=True).data @@ -435,7 +435,7 @@ def test_smooth_with_invert(self, patch_ones): time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) smooth = 0.03 - muted = patch_ones.mute( + muted = patch_ones.line_mute( time=time, distance=distance, smooth=smooth, invert=True ) @@ -448,7 +448,7 @@ def test_smooth_parallel_lines(self, patch_ones): time = ([0, 7.0], [1, 8.0]) distance = ([0, 300], [1, 301]) smooth = {"time": 5, "distance": 10} - muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + muted = patch_ones.line_mute(time=time, distance=distance, smooth=smooth) assert muted.shape == patch_ones.shape # Check for smooth transition @@ -458,7 +458,7 @@ def test_smooth_very_small_value(self, patch_ones): """Test edge case with very small smooth value (1 sample).""" coord = patch_ones.get_coord("time") v1, v2 = _get_testable_coord_values(coord, relative=True) - muted = patch_ones.mute(time=(v1, v2), relative=True, smooth=1) + muted = patch_ones.line_mute(time=(v1, v2), relative=True, smooth=1) # Should still work, just with minimal smoothing assert muted.shape == patch_ones.shape @@ -467,7 +467,7 @@ def test_smooth_preserves_shape(self, patch_ones): time = ([0, 2], [0, 4]) distance = ([0, 100], [0, 100]) smooth = 0.05 - muted = patch_ones.mute(time=time, distance=distance, smooth=smooth) + muted = patch_ones.line_mute(time=time, distance=distance, smooth=smooth) assert muted.shape == patch_ones.shape @@ -477,7 +477,7 @@ class TestMute3D: def test_1d_mute_time(self, patch_ones_3d): """Test 1D mute along time dimension.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_3d.mute(time=(v1, v2), relative=True) + muted = patch_ones_3d.line_mute(time=(v1, v2), relative=True) # Check that the muted region is zeroed sub_muted = muted.select(time=(v1, v2), relative=True) @@ -491,7 +491,7 @@ def test_2d_line_mute(self, patch_ones_3d): """Test 2D line mute using time-distance plane.""" time = ([0, 1.0], [0, 1.5]) distance = ([0, 100], [0, 100]) - muted = patch_ones_3d.mute(time=time, distance=distance) + muted = patch_ones_3d.line_mute(time=time, distance=distance) # Should maintain 3D shape assert muted.shape == patch_ones_3d.shape @@ -506,7 +506,7 @@ def test_2d_line_mute_distance_depth(self, patch_ones_3d): """Test 2D line mute using distance-depth plane.""" distance = ([0, 100], [0, 200]) depth = ([0, 20], [0, 30]) - muted = patch_ones_3d.mute(distance=distance, depth=depth) + muted = patch_ones_3d.line_mute(distance=distance, depth=depth) assert muted.shape == patch_ones_3d.shape assert np.any(muted.data == 0) @@ -515,14 +515,14 @@ def test_2d_line_mute_distance_depth(self, patch_ones_3d): def test_smoothing_3d(self, patch_ones_3d): """Test that smooth parameter works correctly in 3D.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_3d.mute(time=(v1, v2), relative=True, smooth=0.05) + muted = patch_ones_3d.line_mute(time=(v1, v2), relative=True, smooth=0.05) # With smoothing, should have intermediate values assert np.any((muted.data > 0.01) & (muted.data < 0.99)) def test_invert_3d(self, patch_ones_3d): """Test invert parameter in 3D.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_3d.mute(time=(v1, v2), relative=True, invert=True) + muted = patch_ones_3d.line_mute(time=(v1, v2), relative=True, invert=True) # With invert, the selected region should be kept sub_kept = muted.select(time=(v1, v2), relative=True) @@ -539,7 +539,7 @@ class TestMute4D: def test_1d_mute_time(self, patch_ones_4d): """Test 1D mute along time dimension in 4D patch.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_4d.mute(time=(v1, v2), relative=True) + muted = patch_ones_4d.line_mute(time=(v1, v2), relative=True) # Check that the muted region is zeroed sub_muted = muted.select(time=(v1, v2), relative=True) @@ -551,7 +551,7 @@ def test_1d_mute_time(self, patch_ones_4d): def test_1d_mute_angle(self, patch_ones_4d): """Test 1D mute along angle dimension.""" - muted = patch_ones_4d.mute(angle=(20, 60), relative=True) + muted = patch_ones_4d.line_mute(angle=(20, 60), relative=True) sub_muted = muted.select(angle=(20, 60), relative=True) assert np.allclose(sub_muted.data, 0) @@ -559,7 +559,7 @@ def test_2d_line_mute_depth_angle(self, patch_ones_4d): """Test 2D line mute using depth-angle plane.""" depth = ([0, 15], [0, 25]) angle = ([0, 30], [0, 60]) - muted = patch_ones_4d.mute(depth=depth, angle=angle) + muted = patch_ones_4d.line_mute(depth=depth, angle=angle) assert muted.shape == patch_ones_4d.shape assert np.any(muted.data == 0) @@ -570,7 +570,7 @@ def test_smoothing_4d_with_dict(self, patch_ones_4d): time = ([0, 0.8], [0, 1.2]) distance = ([0, 80], [0, 80]) smooth = {"time": 0.02, "distance": 5} - muted = patch_ones_4d.mute(time=time, distance=distance, smooth=smooth) + muted = patch_ones_4d.line_mute(time=time, distance=distance, smooth=smooth) # With smoothing, should have intermediate values assert muted.shape == patch_ones_4d.shape assert np.any((muted.data > 0.01) & (muted.data < 0.99)) @@ -578,14 +578,14 @@ def test_smoothing_4d_with_dict(self, patch_ones_4d): def test_smoothing_1d_mute_in_4d(self, patch_ones_4d): """Test smoothing with 1D mute in 4D patch.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_4d.mute(time=(v1, v2), relative=True, smooth=0.05) + muted = patch_ones_4d.line_mute(time=(v1, v2), relative=True, smooth=0.05) # Check for gradual transition assert np.any((muted.data > 0.01) & (muted.data < 0.99)) def test_invert_4d(self, patch_ones_4d): """Test invert parameter in 4D.""" v1, v2 = 0.3, 0.8 - muted = patch_ones_4d.mute(time=(v1, v2), relative=True, invert=True) + muted = patch_ones_4d.line_mute(time=(v1, v2), relative=True, invert=True) # With invert, the selected region should be kept sub_kept = muted.select(time=(v1, v2), relative=True) @@ -603,15 +603,10 @@ class TestSlopeMute: def test_slope_mute_basic(self, patch_ones): """Test basic slope_mute with two velocities.""" - # Mute between two velocities (20 m/s and 30 m/s) - # These values are reasonable for the example patch which has - # distance range ~300m and time range ~8s slopes = (20.0, 30.0) muted = patch_ones.slope_mute(slopes=slopes) - # Should return a patch with same shape assert muted.shape == patch_ones.shape - # Test specific points using _assert_point_values # Format: (distance, time) in relative coordinates # At time=4s: 20 m/s line is at 80m, 30 m/s line is at 120m @@ -645,7 +640,6 @@ def test_slope_mute_invert(self, patch_ones): slopes = (20.0, 30.0) muted = patch_ones.slope_mute(slopes=slopes, invert=False) inverted = patch_ones.slope_mute(slopes=slopes, invert=True) - # Test that regions are inverted points = [ (50, 4.0), # Outside mute region (low velocity) @@ -665,7 +659,6 @@ def test_slope_mute_with_smooth(self, patch_ones): """Test slope_mute with smoothing parameter.""" slopes = (20.0, 30.0) muted = patch_ones.slope_mute(slopes=slopes, smooth=0.02) - # With smoothing, should have intermediate values assert muted.shape == patch_ones.shape assert np.any((muted.data > 0.01) & (muted.data < 0.99)) @@ -676,15 +669,13 @@ def test_slope_mute_custom_dims(self, patch_ones_3d): # slope = time/distance (inverse velocity, or slowness) slopes = (0.001, 0.002) # s/m muted = patch_ones_3d.slope_mute(slopes=slopes, dims=("time", "distance")) - assert muted.shape == patch_ones_3d.shape # Verify muting occurred - should have both 0s and 1s assert np.any(muted.data == 0) assert np.any(muted.data == 1) - # For slope = time/distance: at distance=200m (relative), mute region is # time = 0.001*200 = 0.2s to 0.002*200 = 0.4s (relative) - # At time=0.3s relative (which is 15% of 2.0s range), dist=200m should be muted + # At time=0.3s relative (15% of 2.0s range), dist=200m should be muted time_idx = muted.get_coord("time").get_next_index(0.3, relative=True) dist_idx = muted.get_coord("distance").get_next_index(200, relative=True) # For 3D patch with dims (time, distance, depth), all depth values should @@ -706,7 +697,6 @@ def test_slope_mute_steep_slopes(self, patch_ones): # Use slopes near the max for the patch (37.5 m/s) slopes = (35.0, 37.0) muted = patch_ones.slope_mute(slopes=slopes) - assert muted.shape == patch_ones.shape # With very high velocities, most of patch should be unmuted # Only a small wedge near upper right should be muted @@ -722,7 +712,6 @@ def test_slope_mute_shallow_slopes(self, patch_ones): """Test slope_mute with very shallow slopes (low velocities).""" slopes = (5.0, 10.0) muted = patch_ones.slope_mute(slopes=slopes) - assert muted.shape == patch_ones.shape # With very low velocities, mute region should be at low distances points = [ @@ -734,12 +723,6 @@ def test_slope_mute_shallow_slopes(self, patch_ones): expected = [1, 0, 1, 0] _assert_point_values(muted, self._dims, points=points, expected_values=expected) - def test_slope_mute_zero_slope_raises(self, patch_ones): - """Test that zero slope raises appropriate error.""" - slopes = (0.0, 30.0) - with pytest.raises(ParameterError, match="positive"): - patch_ones.slope_mute(slopes=slopes) - def test_slope_mute_negative_slopes(self, patch_ones): """Test slope_mute behavior with negative slopes.""" slopes = (-20.0, -30.0) @@ -750,40 +733,30 @@ def test_slope_mute_negative_slopes(self, patch_ones): def test_slope_mute_matches_manual_mute(self, patch_ones): """Test that slope_mute produces same result as manual mute call.""" slopes = (20.0, 30.0) - # Get coordinate ranges dist_range = dc.to_float(patch_ones.get_coord("distance").coord_range()) - time_range = dc.to_float(patch_ones.get_coord("time").coord_range()) - # Calculate endpoints manually endpoints = [] for slope in slopes: - dist_at_max_time = slope * time_range - if dist_at_max_time <= dist_range: - endpoint = (dist_at_max_time, time_range) - else: - endpoint = (dist_range, dist_range / slope) - endpoints.append(endpoint) - + endpoints.append(dist_range / slope) # Manual mute call - manual_muted = patch_ones.mute( - distance=([0, endpoints[0][0]], [0, endpoints[1][0]]), - time=([0, endpoints[0][1]], [0, endpoints[1][1]]), + manual_muted = patch_ones.line_mute( + distance=([0, dist_range], [0, dist_range]), + time=([0, endpoints[0]], [0, endpoints[1]]), relative=True, ) - # slope_mute call slope_muted = patch_ones.slope_mute(slopes=slopes) - - # Should produce identical results - assert np.allclose(manual_muted.data, slope_muted.data) + # Should produce nearly identical results (rounding errors can produce + # minor differences) + equal = manual_muted.data == slope_muted.data + assert equal.sum() / equal.size > 0.9999 def test_slope_mute_with_dict_smooth(self, patch_ones): """Test slope_mute with dict smooth parameter.""" slopes = (20.0, 30.0) smooth = {"distance": 5, "time": 0.01} muted = patch_ones.slope_mute(slopes=slopes, smooth=smooth) - # Should have smooth transition assert np.any((muted.data > 0.01) & (muted.data < 0.99)) @@ -792,7 +765,6 @@ def test_slope_mute_specific_velocity_wedge(self, patch_ones): # Define precise velocity wedge slopes = (15.0, 25.0) # m/s muted = patch_ones.slope_mute(slopes=slopes) - # Test points at specific locations # At time=4.0s: # - 15 m/s line is at distance=60m @@ -806,3 +778,27 @@ def test_slope_mute_specific_velocity_wedge(self, patch_ones): ] expected = [1, 0, 0, 1] _assert_point_values(muted, self._dims, points=points, expected_values=expected) + + def test_0_velocity(self, patch_ones): + """Ensure 0 velocity is ok.""" + muted = patch_ones.slope_mute(slopes=(0, 100)) + points = [ + (250, 1.0), + (200, 1.0), + (250, 4.0), + (250, 5.0), + ] + expected = [1, 1, 0, 0] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) + + def test_inf_velocity(self, patch_ones): + """Ensure an infinite velocity is ok.""" + muted = patch_ones.slope_mute(slopes=(100, np.inf)) + points = [ + (250, 1.0), + (200, 1.0), + (250, 4.0), + (250, 5.0), + ] + expected = [0, 0, 1, 1] + _assert_point_values(muted, self._dims, points=points, expected_values=expected) From 37f386635185f65a4abf75055465e6a813b66a3c Mon Sep 17 00:00:00 2001 From: Derrick Chambers Date: Tue, 28 Oct 2025 10:01:37 +0000 Subject: [PATCH 15/15] address rabbit --- dascore/proc/filter.py | 8 ++++---- dascore/proc/mute.py | 4 ++-- dascore/utils/misc.py | 2 +- tests/test_proc/test_basic.py | 2 +- tests/test_proc/test_mute.py | 2 +- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/dascore/proc/filter.py b/dascore/proc/filter.py index ba0c306b6..60045ba03 100644 --- a/dascore/proc/filter.py +++ b/dascore/proc/filter.py @@ -501,7 +501,7 @@ def slope_filter( >>> patch_filtered = patch.slope_filter( ... filt=filt, ... directional=False, - ... notch=False + ... invert=False ... ) >>> # Plot results >>> fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8)) @@ -510,8 +510,8 @@ def slope_filter( >>> ax2 = patch_filtered.viz.waterfall(ax=ax2, scale=0.5) >>> _ = ax2.set_title('Filtered') >>> - >>> # Example 2: Notch filter - >>> patch_filtered = patch.slope_filter(filt=filt, notch=True) + >>> # Example 2: Inverted (notch) filter + >>> patch_filtered = patch.slope_filter(filt=filt, invert=True) >>> >>> # Example 3: specify units >>> filt = np.array([2e3,2.2e3,8e3,2e4]) * dc.get_unit("m/s") @@ -602,7 +602,7 @@ def _maybe_transform_units(filt, dft_patch, freq_dims): # TODO remove in dascore 0.2. if notch is not None: msg = "The `notch` parameter of slope filter is deprecated. Use invert." - warnings.warn(msg, DeprecationWarning) + warnings.warn(msg, DeprecationWarning, stacklevel=2) invert = notch mask = _get_taper_mask(filt, slope, invert) diff --git a/dascore/proc/mute.py b/dascore/proc/mute.py index 9fbcc6157..7893f75c8 100644 --- a/dascore/proc/mute.py +++ b/dascore/proc/mute.py @@ -90,7 +90,7 @@ def _broadcast_smooth_to_dims(dims, smooth): def _convert_to_samples(smooth, dims, patch): """Convert the smooth parameter to number of samples.""" out = [] - for dim, val in zip(dims, smooth): + for dim, val in zip(dims, smooth, strict=True): coord = patch.get_coord(dim) if val is None: out.append(0) @@ -246,7 +246,7 @@ def _get_coord_float_values(coord, vals, relative): # Indicates an implicit value was used. ifill_index = -1 - for ind, (coord, row) in enumerate(zip(coords, value_list)): + for ind, (coord, row) in enumerate(zip(coords, value_list, strict=True)): coord = patch.get_coord(patch.dims[ind]) # We iterate each pair because it might be an implicit value. for pair_ind, pair in enumerate(row): diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 525d1f7f3..d5f43599e 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -918,7 +918,7 @@ def get_2d_line_intersection(p1, p2, p3, p4): denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4) if np.isclose(denom, 0): - np.array([np.nan, np.nan]) + return np.array([np.nan, np.nan]) num_x = (x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4) num_y = (x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4) diff --git a/tests/test_proc/test_basic.py b/tests/test_proc/test_basic.py index af7731e29..ebbcdacf6 100644 --- a/tests/test_proc/test_basic.py +++ b/tests/test_proc/test_basic.py @@ -832,5 +832,5 @@ class TestFull: def test_full_1(self, random_patch): """Ensure a patch can be created with 1s.""" patch = random_patch.full(1.0) - patch.coords == random_patch.coords + assert patch.coords == random_patch.coords assert np.allclose(patch.data, 1.0) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py index 11a18b8b0..46ef09dbf 100644 --- a/tests/test_proc/test_mute.py +++ b/tests/test_proc/test_mute.py @@ -138,7 +138,7 @@ def test_1d_mute_no_taper(self, patch_ones): def test_mute_open_interval(self, patch_ones): """Mute using None for interval ends.""" coord = patch_ones.get_coord("distance") - v1, v2 = _get_testable_coord_values(coord, relative=True) + _v1, v2 = _get_testable_coord_values(coord, relative=True) muted1 = patch_ones.line_mute(distance=(v2, ...), relative=True) step = dc.to_float(coord.step) _assert_coord_ranges(