allow align_to_coord with non-overlapping traces

d-chambers · d-chambers · commit 9a01cf190e8d · 2025-11-22T20:45:15.000-08:00
diff --git a/dascore/proc/align.py b/dascore/proc/align.py
@@ -189,7 +189,7 @@ def _get_aligned_coords(patch, dim_name, meta):
     return patch.coords.update(**{dim_name: new_coord})
 
 
-def _get_shift_indices(coord_vals, dim, reverse, samples):
+def _get_shift_indices(coord_vals, dim, reverse, samples, mode):
     """
     Get the indices  for shifting.
 
@@ -200,11 +200,19 @@ def _get_shift_indices(coord_vals, dim, reverse, samples):
     # Therefore, we need to insure positive values go into get_next_index.
     abs_vals = np.abs(coord_vals)
     sign_vals = np.sign(coord_vals)
-    inds_abs = dim.get_next_index(
-        abs_vals,
-        samples=samples,
-        relative=False if samples else True,
-    )
+    try:
+        inds_abs = dim.get_next_index(
+            abs_vals,
+            samples=samples,
+            relative=False if samples else True,
+            allow_out_of_bounds=mode == "full",
+        )
+    except ValueError:
+        msg = (
+            f"Trace shift with align_to_coord results in some traces with no "
+            f"overlaps. This is only possible with mode = 'full' not {mode}."
+        )
+        raise ParameterError(msg)
     # Reverse index if needed. This way the reference stays the same.
     inds = inds_abs * (sign_vals.astype(np.int64) * (-1 if reverse else 1))
     return inds
@@ -263,7 +271,8 @@ def align_to_coord(
         Determines the output shape of the patch. Options are:
         "full" - Regardless of shift, all original data are preserved.
             This can result in patches with many fill values along the
-            aligned dimension.
+            aligned dimension. It also allows cases where some traces are
+            do not overlap at all.
         "same" - The patch will retain its shape, however, only one trace
             (and traces that weren't shifted) will remain complete. Parts
             of shifted traces will be discarded.
@@ -370,7 +379,7 @@ def align_to_coord(
     coord_dims = patch.coords.dim_map[coord_name]
     coord_axes = tuple(patch.dims.index(x) for x in coord_dims)
     # Get the metadata about shift and the indices for shifting.
-    inds = _get_shift_indices(coord_vals, dim, reverse, samples)
+    inds = _get_shift_indices(coord_vals, dim, reverse, samples, mode)
     meta = _calculate_shift_info(mode, inds, dim, dim_axis, patch.shape)
     # Apply shifts to data
     shifted_data = _apply_shifts_to_data(
diff --git a/tests/test_proc/test_align.py b/tests/test_proc/test_align.py
@@ -113,6 +113,17 @@ def patch_with_mixed_shifts(simple_patch_for_alignment):
     return patch.update_coords(shift_time_mixed=("distance", mixed_shifts))
 
 
+@pytest.fixture(scope="class")
+def patch_non_overlapping_shifts(random_patch):
+    """Patch with shifts that cause some traces to not overlap."""
+    sub = (
+        random_patch.select(distance=slice(0, 3), samples=True)
+        .select(time=slice(0, 10), samples=True)
+        .update_coords(shift=("distance", np.array([0, 8, 16])))
+    )
+    return sub
+
+
 class TestAlignToCoordValidation:
     """Tests for align_to_coord validation logic."""
 
@@ -375,3 +386,16 @@ def test_reverse_coordinate_accuracy(self, patch_with_known_shifts):
         assert np.isclose(
             reversed_patch.get_coord("time").step, original_time_coord.step
         )
+
+    def test_align_no_overlap(self, patch_non_overlapping_shifts):
+        """Test that a patch with shifts that cause traces to not overlap."""
+        patch = patch_non_overlapping_shifts
+        msg = "some traces with no overlaps"
+        # Mode = valid and same should fail.
+        with pytest.raises(ParameterError, match=msg):
+            patch.align_to_coord(time="shift", mode="valid")
+        with pytest.raises(ParameterError, match=msg):
+            patch.align_to_coord(time="shift", mode="same")
+        # But full should work.
+        new = patch.align_to_coord(time="shift", mode="full")
+        assert new.dropna("time").data.size == 0
