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 483b1c8e5..fbba00484 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: """ @@ -438,6 +439,8 @@ def iresample(self, *args, **kwargs): standardize = dascore.proc.standardize taper = dascore.proc.taper taper_range = dascore.proc.taper_range + 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 944d5ab04..40fc327db 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 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/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/filter.py b/dascore/proc/filter.py index b764e4dd7..60045ba03 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. @@ -498,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)) @@ -507,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") @@ -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, stacklevel=2) + 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 new file mode 100644 index 000000000..7893f75c8 --- /dev/null +++ b/dascore/proc/mute.py @@ -0,0 +1,634 @@ +"""Processing for muting (zeroing) patch data in specified regions.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sized +from typing import ClassVar + +import numpy as np +from numpy.linalg import norm +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.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): + """ + Parent class for Mute Geometry. + """ + + dims: tuple[str, ...] + axes: tuple[int, ...] + relative: bool = True + + @classmethod + @abstractmethod + def from_params(cls, vals, dims, axes, patch, relative): + """Initialize Mute Geometry from input parameters.""" + + @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.""" + 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.""" + + 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) + axes = self.axes + else: + # Otherwise, the smooth dict must be a subset of the dimensions. + if not set(smooth).issubset(set(dims)): + msg = ( + 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 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.""" + out = [] + for dim, val in zip(dims, smooth, strict=True): + 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, 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, axes + + +class _MuteGeometry1D(_MuteGeometry): + """ + Private container to manage 1D Mute Geometry. + """ + + lims: tuple + + @classmethod + 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 + ---------- + 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. + 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. + 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 + """ + + origins: tuple[NDArray[np.floating], NDArray[np.floating]] + norm: NDArray[np.floating] + line1_norm: NDArray[np.floating] + line2_norm: NDArray[np.floating] + parallel: bool + + # 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 _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. + coord_norm = np.array( + [dc.to_float(patch.get_coord(x).coord_range()) for x in dims] + ) + 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) + 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 + 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 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) + + # 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( + origins=origin_tuple, + norm=coord_norm, + line1_norm=v_norms[0], + line2_norm=v_norms[1], + parallel=parallel, + ) + return out + + @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.""" + out = dc.to_float(np.array(vals)) + if not relative: + out = out - dc.to_float(coord.min()) + return out + + # 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, 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): + 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. + line_params = cls._get_line_params(points, patch, dims, fill_ind) + 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. + """ + + 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).eps + return array / array_norm + + out = [[], []] + 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()) + # 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. + # with the same dimensionality as array. + 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)) + 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_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): + """ + Determine if each point is in the region based on the cross product + requirement. + """ + # Get padded arrays with 0 z values for cross product. + # 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 + ) + 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. + # 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 + + 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 > 0) & (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(cnorms_1, cnorms_2) + # For parallel lines only cross product is needed to define region. + if self.parallel: + 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): + """ + Get (and validate) the muting parameters. + + 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: + 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]) + val_list = [x[2] for x in dim_ax_vals] + if len(dims) == 1: + geometry = _MuteGeometry1D.from_params( + val_list, dims, axes, patch, relative=relative + ) + elif len(dims) > 1: # Dealing with lines. + geometry = _MuteGeometry2D.from_params( + val_list, dims, axes, patch, relative=relative + ) + return geometry + + +@patch_function() +@compose_docstring(smooth_param=_smooth_param) +def line_mute( + patch: PatchType, + *, + smooth=None, + invert: bool = False, + relative: bool = True, + **kwargs, +) -> PatchType: + """ + 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 or arrays of values. + + Parameters + ---------- + patch + The patch instance. + {smooth_param} + 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. + **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 + -------- + >>> import dascore as dc + >>> from scipy.ndimage import gaussian_filter + >>> + >>> patch = dc.get_example_patch().full(1) + >>> + >>> # Mute first 0.5s (relative to start by default) + >>> muted = patch.line_mute(time=(0, 0.5)) + >>> + >>> # Mute everything except middle section + >>> kept = patch.line_mute(time=(0.2, -0.2), invert=True) + >>> + >>> # 1D Mute with smoothed absolute units for time. + >>> 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.line_mute( + ... time=(0, [0, 0.3]), + ... distance=(None, [0, 300]), + ... smooth=0.02, + ... ) + >>> + >>> # Mute late arrivals: from velocity line to end + >>> muted = patch.line_mute( + ... time=([0, 0.3], None), + ... distance=([0, 300], 0), + ... ) + >>> + >>> # Mute wedge between two velocity lines + >>> muted = patch.line_mute( + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), + ... ) + >>> + >>> # Mute wedge outside two velocity lines + >>> muted = patch.line_mute( + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), + ... invert=True, + ... ) + >>> + >>> # Apply custom tapering + >>> ones = patch.full(1.0) + >>> envelope = ones.line_mute( + ... time=([0, 0.375], [0, 0.25]), + ... distance=([0, 300], [0, 300]), + ... ) + >>> # 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 + + Notes + ----- + - 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 boundary smoothing, use a patch with one values + 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`](`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) + # 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. + out = geo._apply_mask(out, patch, fill_val) + # Apply smoothing if requested. + 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 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 + 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 + + 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. + 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.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 sequence of length 2" + raise ParameterError(msg) + # Check for zero or negative slopes + if np.any(slopes_array < 0): + msg = "slopes must be positive." + raise ParameterError(msg) + # 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]: dim0_vals, + dims[1]: dim1_vals, + "smooth": smooth, + "invert": invert, + "relative": False, + } + return line_mute.func(patch, **mute_kwargs) diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py index 4d529ef40..d5f43599e 100644 --- a/dascore/utils/misc.py +++ b/dascore/utils/misc.py @@ -895,3 +895,34 @@ 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. x and y are nan 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 np.isclose(denom, 0): + 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) + with np.errstate(divide="ignore", invalid="ignore"): + px = num_x / denom + py = num_y / denom + return np.array([px, py]) 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_basic.py b/tests/test_proc/test_basic.py index e3d74dddc..ebbcdacf6 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) + assert patch.coords == random_patch.coords + assert np.allclose(patch.data, 1.0) 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]) diff --git a/tests/test_proc/test_mute.py b/tests/test_proc/test_mute.py new file mode 100644 index 000000000..46ef09dbf --- /dev/null +++ b/tests/test_proc/test_mute.py @@ -0,0 +1,804 @@ +"""Tests for mute processing function.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dascore as dc +from dascore.exceptions import ParameterError + + +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 + start, stop = coord[start_ind], coord[stop_ind] + if relative: + start, stop = start - coord.min(), stop - coord.min() + return (dc.to_float(start), dc.to_float(stop)) + + +def _assert_coord_ranges( + patch, + dim, + zero_ranges, + one_ranges, + relative=True, +): + """Assert that the expected values occur in the patch.""" + for zrange in zero_ranges: + 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) + 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 = [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") +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="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 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.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.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.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.line_mute(time=(1, 2), smooth=1.1) + with pytest.raises(ParameterError, match="smooth parameter for"): + random_patch.line_mute(time=(1, 2), smooth=-0.01) + + +class Test1DLineMute: + """Test 1D block mutes (single dimension).""" + + def test_1d_mute_no_taper(self, patch_ones): + """Mute first portion of time dimension.""" + # 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.line_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 - step), (v2 + step, ...)], + relative=True, + ) + + 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.line_mute(distance=(v2, ...), relative=True) + step = dc.to_float(coord.step) + _assert_coord_ranges( + patch=muted1, + dim="distance", + zero_ranges=[(v2, ...)], + one_ranges=[(..., v2 - step)], + relative=True, + ) + + 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.line_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 - step), (v2 + step, ...)], + relative=False, + ) + + # 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.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) + sub = muted1.select(time=(v1, v2), relative=True) + assert not np.allclose(sub.data, 1) + + 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.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) + 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.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) + 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.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.line_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.""" + + _dims = ("distance", "time") + + def test_point_raise(self, patch_ones): + """A degenerate line (point) should raise.""" + msg = "is degenerate" + with pytest.raises(ParameterError, match=msg): + patch_ones.line_mute( + time=([0, 0], [0, 0.25]), + distance=([0, 0], [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(ParameterError, match=match): + patch_ones.line_mute( + time=[[0, 0], [1, -1]], + distance=[[1, -1], [-1, 1]], + ) + + 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.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]) + _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_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.line_mute( + time=time, + distance=distance, + ) + 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]) + _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_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.line_mute( + time=time, + distance=distance, + ) + inverted = patch_ones.line_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.line_mute( + time=time, + distance=distance, + ) + 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]) + _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.line_mute( + time=time, + distance=distance, + ) + sub = muted.select(time=(2.1, ...), relative=True) + assert np.allclose(sub.data, 1) + + 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) + + 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.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.line_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.line_mute(time=time, distance=distance, relative=False) + inverted = patch_ones_3d.line_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.line_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.line_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.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.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 + # 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.line_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.line_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.line_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.line_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.line_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_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.line_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.line_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.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.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) + 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.line_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.line_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.line_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.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)) + + 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.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.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) + 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 TestSlopeMute: + """Tests for slope_mute functionality.""" + + _dims = ("distance", "time") + + def test_slope_mute_basic(self, patch_ones): + """Test basic slope_mute with two velocities.""" + 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 (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_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()) + # Calculate endpoints manually + endpoints = [] + for slope in slopes: + endpoints.append(dist_range / slope) + # Manual mute call + 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 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)) + + 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) + + 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) diff --git a/tests/test_proc/test_taper.py b/tests/test_proc/test_taper.py index 80b887634..d56d5ddef 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 - 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 - 10, :], 0) + + # Values from start_idx to end should be 0 (muted) + assert np.allclose(out.data[start_idx, :], 1) + assert np.allclose(out.data[end_idx - 1, :], 1)