Patch.where (#550)

d-chambers · web-flow · commit b115a7672876 · 2025-09-23T21:21:23.000+01:00
* add patch.where

* fix doctest

* mind the rabbit
diff --git a/dascore/core/patch.py b/dascore/core/patch.py
@@ -421,6 +421,7 @@ def iselect(self, *args, **kwargs):
     resample = dascore.proc.resample
     pad = dascore.proc.pad
     roll = dascore.proc.roll
+    where = dascore.proc.where
 
     @deprecate(
         "patch.iresample is deprecated. Please use patch.resample " "with samples=True",
diff --git a/dascore/proc/basic.py b/dascore/proc/basic.py
@@ -10,6 +10,7 @@
 from scipy.fft import next_fast_len
 
 import dascore as dc
+from dascore.compat import array
 from dascore.constants import PatchType
 from dascore.core.attrs import PatchAttrs
 from dascore.core.coordmanager import CoordManager, get_coord_manager
@@ -19,6 +20,7 @@
 from dascore.utils.misc import _get_nullish
 from dascore.utils.models import ArrayLike
 from dascore.utils.patch import (
+    align_patch_coords,
     get_dim_axis_value,
     patch_function,
 )
@@ -688,3 +690,66 @@ def roll(patch, samples=False, update_coord=False, **kwargs):
         patch = patch.update_coords(**{dim: new_coord})
 
     return patch.new(data=roll_arr)
+
+
+@patch_function()
+def where(
+    patch: PatchType, cond: ArrayLike | PatchType, other: Any | PatchType = np.nan
+) -> PatchType:
+    """
+    Return elements from patch where condition is True, else fill with other.
+
+    Parameters
+    ----------
+    patch
+        The input patch
+    cond
+        Condition array. Should be a boolean array with the same shape as patch data,
+        or a patch with boolean data that is broadcastable to the patch's shape.
+    other
+        Value to use for locations where cond is False. Can be a scalar value,
+        array, or patch that is broadcastable to the patch's shape. Default is NaN.
+
+    Returns
+    -------
+    PatchType
+        A new patch with values from patch where cond is True, and other elsewhere.
+
+    Examples
+    --------
+    >>> import dascore as dc
+    >>> import numpy as np
+    >>> patch = dc.get_example_patch()
+    >>>
+    >>> # Where data > 0 fill with original patch values else nan.
+    >>> condition = patch.data > 0
+    >>> out = patch.where(condition)
+    >>>
+    >>> # Use another patch as condition
+    >>> threshold = patch.data.mean()
+    >>> boolean_patch = patch.new(data=(patch.data > threshold))
+    >>> out = patch.where(boolean_patch, other=0)
+    >>>
+    >>> # Replace values below threshold with 0
+    >>> out = patch.where(patch.data > patch.data.mean(), other=0)
+    """
+    cls = patch.__class__  # Use this so it works with subclasses
+    # Align patch and cond
+    if isinstance(cond, cls):
+        patch, cond = align_patch_coords(patch, cond)
+    # Align patch and other, may need to re-align cond
+    if isinstance(other, cls):
+        patch, other = align_patch_coords(patch, other)
+        if isinstance(cond, cls):
+            patch, cond = align_patch_coords(patch, cond)
+
+    cond_array, other_array = array(cond), array(other)
+
+    # Ensure condition is boolean
+    if not np.issubdtype(cond_array.dtype, np.bool_):
+        msg = "Condition must be a boolean array or patch with boolean data"
+        raise ValueError(msg)
+
+    # Use numpy.where to apply condition
+    new_data = np.where(cond_array, patch.data, other_array)
+    return patch.new(data=new_data)
diff --git a/tests/test_proc/test_basic.py b/tests/test_proc/test_basic.py
@@ -598,3 +598,177 @@ def test_dist_coord_roll(self, random_patch):
             random_patch.coords.get_array("distance")[0]
             == rolled_patch.coords.get_array("distance")[value]
         )
+
+
+class TestWhere:
+    """Tests for the where method of Patch."""
+
+    def test_where_with_boolean_array(self, random_patch):
+        """Test where with a boolean array condition."""
+        condition = random_patch.data > random_patch.data.mean()
+        result = random_patch.where(condition)
+
+        # Check that the result has the same shape
+        assert result.shape == random_patch.shape
+
+        # Check that values where condition is True are preserved
+        assert np.allclose(result.data[condition], random_patch.data[condition])
+
+        # Check that values where condition is False are NaN
+        assert np.all(np.isnan(result.data[~condition]))
+
+    def test_where_with_other_value(self, random_patch):
+        """Test where with a replacement value."""
+        condition = random_patch.data > 0
+        other_value = -999
+        result = random_patch.where(condition, other=other_value)
+
+        # Check that values where condition is True are preserved
+        assert np.allclose(result.data[condition], random_patch.data[condition])
+
+        # Check that values where condition is False are replaced
+        assert np.all(result.data[~condition] == other_value)
+
+    def test_where_with_patch_condition(self, random_patch):
+        """Test where with another patch as condition."""
+        boolean_data = (random_patch.data > random_patch.data.mean()).astype(bool)
+        boolean_patch = random_patch.new(data=boolean_data)
+        result = random_patch.where(boolean_patch, other=0)
+
+        # Check that the result has the same shape
+        assert result.shape == random_patch.shape
+
+        # Check that values where condition is True are preserved
+        true_mask = boolean_data
+        assert np.allclose(result.data[true_mask], random_patch.data[true_mask])
+
+        # Check that values where condition is False are 0
+        false_mask = ~boolean_data
+        assert np.all(result.data[false_mask] == 0)
+
+    def test_where_preserves_metadata(self, random_patch):
+        """Test that where preserves patch metadata."""
+        condition = random_patch.data > 0
+        result = random_patch.where(condition)
+
+        # Check that coordinates are preserved
+        assert result.coords == random_patch.coords
+
+        # Check that dimensions are preserved
+        assert result.dims == random_patch.dims
+
+        # Check that attributes are preserved (except history)
+        assert result.attrs.model_dump(
+            exclude={"history"}
+        ) == random_patch.attrs.model_dump(exclude={"history"})
+
+    def test_where_non_boolean_condition_raises(self, random_patch):
+        """Test that non-boolean condition raises ValueError."""
+        non_boolean_condition = random_patch.data  # Not boolean
+
+        with pytest.raises(ValueError, match="Condition must be a boolean array"):
+            random_patch.where(non_boolean_condition)
+
+    def test_where_broadcasts_condition(self, random_patch):
+        """Test that condition can be broadcast to patch shape."""
+        # Create a condition that can be broadcast to the full shape
+        # Create a boolean array that matches the first dimension
+        condition = np.ones(random_patch.shape[0], dtype=bool)
+        condition[0] = False  # Make first element False
+
+        # This should broadcast across the second dimension
+        result = random_patch.where(condition[:, np.newaxis], other=-1)
+        assert result.shape == random_patch.shape
+
+        # Check that first row is all -1 and others are preserved
+        assert np.all(result.data[0, :] == -1)
+        assert np.allclose(result.data[1:, :], random_patch.data[1:, :])
+
+    def test_where_with_broadcastable_patch_other(self, random_patch):
+        """Test where with a broadcastable patch as other parameter."""
+        # Get the actual dimensions of the patch to create the right broadcasting
+        broadcastable_patch1 = random_patch.mean("distance").squeeze()
+        broadcastable_patch2 = random_patch.mean("time").squeeze()
+
+        # Create condition
+        condition = random_patch.data > random_patch.data.mean()
+
+        for castable in [broadcastable_patch1, broadcastable_patch2]:
+            result = random_patch.where(condition, other=castable)
+            assert result.shape == random_patch.shape
+            # Check that values where condition is True are preserved
+            assert np.allclose(result.data[condition], random_patch.data[condition])
+            # Check that values where condition is False come from the broadcasted other
+            false_mask = ~condition
+            # The exact values depend on how the broadcasting worked
+            assert np.all(
+                ~np.isnan(result.data[false_mask])
+            )  # Should have valid values
+
+    def test_where_with_misaligned_coords(self, random_patch):
+        """Test where with condition patch having misaligned coordinates."""
+        # Create a subset of the original patch with partial overlap
+        time_coord = random_patch.coords.get_array("time")
+        # Take only part of the time coordinates to create a partial overlap
+        partial_time = time_coord[10:20]  # Use a subset
+
+        # Create a boolean condition patch with partial time coordinates
+        shifted_patch = random_patch.new(
+            coords={
+                "time": partial_time,
+                "distance": random_patch.coords.get_array("distance"),
+            },
+            data=(random_patch.data[:, 10:20] > random_patch.data[:, 10:20].mean()),
+        )
+
+        # This should work with coordinate alignment (union)
+        result = random_patch.where(shifted_patch, other=0)
+
+        # The result should have the union of coordinates and correct shape
+        assert result is not None
+        assert isinstance(result.data, np.ndarray)
+        # After alignment, coords should have the overlapping range
+        result_time = result.coords.get_array("time")
+        partial_time_len = len(partial_time)
+        assert len(result_time) == partial_time_len
+
+    def test_where_both_cond_and_other_misaligned(self, random_patch):
+        """Test where with both condition and other patches having misaligned coords."""
+        # Create condition patch with partial time overlap (first part)
+        time_coord = random_patch.coords.get_array("time")
+        cond_time = time_coord[5:15]  # indices 5-14
+
+        condition_patch = random_patch.new(
+            coords={
+                "time": cond_time,
+                "distance": random_patch.coords.get_array("distance"),
+            },
+            data=(random_patch.data[:, 5:15] > random_patch.data[:, 5:15].mean()),
+        )
+
+        # Create other patch with different partial time overlap (shifted range)
+        other_time = time_coord[8:18]  # indices 8-17, overlaps with condition
+        other_patch = random_patch.new(
+            coords={
+                "time": other_time,
+                "distance": random_patch.coords.get_array("distance"),
+            },
+            data=random_patch.data[:, 8:18] * 0.5,  # Use different values
+        )
+
+        # This should work with coordinate alignment handling both patches
+        result = random_patch.where(condition_patch, other=other_patch)
+
+        # The result should have coordinates that are the intersection of all three
+        assert result is not None
+        assert isinstance(result.data, np.ndarray)
+
+        # After alignment, the time coordinate should be the intersection
+        result_time = result.coords.get_array("time")
+        # The intersection of [5:15], [8:18], and full range should be [8:15]
+        expected_overlap_len = 7  # indices 8, 9, 10, 11, 12, 13, 14
+        assert len(result_time) == expected_overlap_len
+
+        # Verify the actual time values match the expected overlap
+        expected_time_values = time_coord[8:15]
+        assert np.array_equal(result_time, expected_time_values)
