From b6edf30b6f4954e6676991ad07e6652422162447 Mon Sep 17 00:00:00 2001
From: Derrick Chambers <d-chambers@users.noreply.github.com>
Date: Thu, 13 Jun 2024 15:19:41 -0600
Subject: [PATCH] fix correlate doc build (#402)

---
 dascore/proc/correlate.py         | 9 ++++-----
 docs/recipes/correlate.qmd        | 4 +++-
 tests/test_proc/test_correlate.py | 2 +-
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/dascore/proc/correlate.py b/dascore/proc/correlate.py
index 9af48ac8..db0817e2 100644
--- a/dascore/proc/correlate.py
+++ b/dascore/proc/correlate.py
@@ -70,9 +70,7 @@ def correlate_shift(patch, dim, undo_weighting=True):
     data = np.fft.fftshift(patch.data, axes=axis)
     if undo_weighting:
         data = data / to_float(coord.step)
-    # so it appears (from testing) there is one lest sample on the positive
-    # side.
-    step = to_float(coord.step)
+    step = coord.step
     new_start = -np.ceil((len(coord) - 1) / 2) * step
     new_end = np.ceil((len(coord) - 1) / 2) * step
     new_coord = dc.get_coord(start=new_start, stop=new_end, step=step)
@@ -182,8 +180,9 @@ def correlate(
     -----
     1 - The cross-correlation is performed in the frequency domain.
 
-    2 - The output dimension is opposite of the one specified in kwargs, has
-    the units of float, and the string "lag_" prepended. For example, "lag_time".
+    2 - The output dimension is opposite of the one specified in kwargs and
+      shares a name with the original coord except the string "lag_" is
+      prepended. For example, "lag_time".
     """
     if lag is not None:
         msg = "Correlate lag is deprecated. Simply use on the output patch."
diff --git a/docs/recipes/correlate.qmd b/docs/recipes/correlate.qmd
index 06b1fb91..3ac316f8 100644
--- a/docs/recipes/correlate.qmd
+++ b/docs/recipes/correlate.qmd
@@ -38,5 +38,7 @@ patch = dc.get_example_patch(
 
 corr = patch.correlate(distance=3, samples=True)
 
-corr.viz.waterfall();
+# Note we squeeze the last dimension to get 2D patch
+
+corr.squeeze().viz.waterfall();
 ```
diff --git a/tests/test_proc/test_correlate.py b/tests/test_proc/test_correlate.py
index 41ba5067..8b559ecc 100644
--- a/tests/test_proc/test_correlate.py
+++ b/tests/test_proc/test_correlate.py
@@ -24,7 +24,7 @@ def test_auto_correlation(self, random_dft_patch):
         # ensure the max value happens at zero lag time.
         time_ax = auto_patch.dims.index("lag_time")
         argmax = np.argmax(random_dft_patch.data, axis=time_ax)
-        assert np.allclose(coord_array[argmax], 0)
+        assert np.all(coord_array[argmax] == dc.to_timedelta64(0))
 
 
 class TestCorrelateInternal: