Use the depth axis to exclude the first 3D correlation peak

openpiv/pyprocess3D.py

@@ -159,7 +159,7 @@ def find_second_peak_3D(corr, i=None, j=None, z=None, width=2):
     jini = max(0, j - width)
     jfin = min(j + width + 1, corr.shape[1])
     zini = max(0, z - width)
-    zfin = min(z + width + 1, corr.shape[1])
+    zfin = min(z + width + 1, corr.shape[2])
 
     tmp[iini:ifin, jini:jfin, zini:zfin] = ma.masked
     (i, j, z), corr_max2 = find_first_peak(tmp)

openpiv/test/test_pyprocess3D.py

@@ -1,16 +1,51 @@
-from openpiv.pyprocess3D import extended_search_area_piv3D
 import numpy as np
-
-from skimage.util import random_noise
+import pytest
+from scipy.ndimage import shift as shift_img
 from skimage import img_as_ubyte
+from skimage.util import random_noise
 
-from scipy.ndimage import shift as shift_img
+from openpiv.pyprocess3D import (
+    extended_search_area_piv3D,
+    find_second_peak_3D,
+    sig2noise_ratio,
+)
 
 # import warnings
 
 threshold = 0.1
 
 
+@pytest.mark.parametrize("width", [0, 2])
+@pytest.mark.parametrize("explicit_peak", [False, True])
+@pytest.mark.parametrize("shape,first,second", [
+    ((7, 8, 20), (3, 4, 15), (3, 4, 5)),
+    ((7, 8, 20), (3, 4, 18), (3, 4, 5)),
+    ((7, 20, 8), (3, 10, 4), (3, 5, 4)),
+    ((9, 12, 12), (4, 6, 6), (4, 3, 6)),
+])
+def test_second_peak_uses_depth(shape, first, second, explicit_peak, width):
+    corr = np.zeros(shape)
+    corr[first] = 10.0
+    corr[second] = 4.0
+    original = corr.copy()
+    args = first if explicit_peak else ()
+
+    position, value = find_second_peak_3D(corr, *args, width=width)
+
+    assert position == second
+    assert value == 4.0
+    np.testing.assert_array_equal(corr, original)
+
+
+@pytest.mark.parametrize("depth", [16, 20])
+def test_signal_to_noise_excludes_first_peak_in_deep_volume(depth):
+    corr = np.zeros((7, 8, depth))
+    corr[3, 4, depth - 3] = 10.0
+    corr[3, 4, 5] = 4.0
+
+    assert sig2noise_ratio(corr, sig2noise_method="peak2peak") == 2.5
+
+
 def dist(u, shift):
     return np.mean(np.abs(u - shift))