Improve test coverage on extract_amp_phase (#545)

Author: waltsims

kwave/utils/conversion.py

@@ -489,3 +489,18 @@ def find_closest(array: ndarray, value: Num[kt.ScalarLike, ""]):
     array = np.asarray(array)
     idx = (np.abs(array - value)).argmin()
     return array[idx], idx
+
+
+def create_index_at_dim(ndim: int, dim: int, index_value: Any) -> tuple:
+    """
+    Create a tuple of slice objects with a specific index value at the specified dimension.
+
+    Args:
+        ndim: Number of dimensions in the array
+        dim: The dimension where the specific index should be placed
+        index_value: The index value to place at the specified dimension
+
+    Returns:
+        A tuple of slice objects with the index value at the specified dimension
+    """
+    return tuple(index_value if i == dim else slice(None) for i in range(ndim))

kwave/utils/filters.py

@@ -6,8 +6,8 @@
 from scipy.fftpack import fft, fftshift, ifft, ifftshift
 from scipy.signal import convolve, lfilter
 
-from ..kgrid import kWaveGrid
-from ..kmedium import kWaveMedium
+from kwave.utils.conversion import create_index_at_dim
+
 from .checks import is_number
 from .data import scale_SI
 from .math import find_closest, gaussian, next_pow2, norm_var, sinc
@@ -24,31 +24,26 @@ def single_sided_correction(func_fft: np.ndarray, fft_len: int, dim: int) -> np.
     Args:
         func_fft: The FFT of the function to be corrected.
         fft_len: The length of the FFT.
-        dim: The number of dimensions of `func_fft`.
+        dim: The dimension along which to apply the correction.
 
     Returns:
         The corrected FFT of the function.
     """
+    # Determine the slice to use based on FFT length
     if fft_len % 2:
-        # odd FFT length switch dim case
-        if dim == 0:
-            func_fft[1:, :] = func_fft[1:, :] * 2
-        elif dim == 1:
-            func_fft[:, 1:] = func_fft[:, 1:] * 2
-        elif dim == 2:
-            func_fft[:, :, 1:] = func_fft[:, :, 1:] * 2
-        elif dim == 3:
-            func_fft[:, :, :, 1:] = func_fft[:, :, :, 1:] * 2
+        # odd FFT length - multiply all elements except the first one by 2
+        dim_slice = slice(1, None)
     else:
-        # even FFT length
-        if dim == 0:
-            func_fft[1:-1] = func_fft[1:-1] * 2
-        elif dim == 1:
-            func_fft[:, 1:-1] = func_fft[:, 1:-1] * 2
-        elif dim == 2:
-            func_fft[:, :, 1:-1] = func_fft[:, :, 1:-1] * 2
-        elif dim == 3:
-            func_fft[:, :, :, 1:-1] = func_fft[:, :, :, 1:-1] * 2
+        # even FFT length - multiply all elements except the first and last ones by 2
+        dim_slice = slice(1, -1)
+
+    # Create a slice tuple with the appropriate slice at the specified dimension
+    idx_all = [slice(None)] * func_fft.ndim
+    idx_all[dim] = dim_slice
+    idx_tuple = tuple(idx_all)
+
+    # Apply the correction
+    func_fft[idx_tuple] = func_fft[idx_tuple] * 2
 
     return func_fft
 
@@ -197,7 +192,6 @@ def extract_amp_phase(
 
     # compute amplitude and phase spectra
     f, func_as, func_ps = spect(data, Fs, fft_len=fft_padding * data.shape[dim], dim=dim)
-
     # correct for coherent gain
     func_as = func_as / coherent_gain
 
@@ -209,20 +203,10 @@ def extract_amp_phase(
     sz[dim - 1] = 1
 
     # extract amplitude and relative phase at freq_index
-    if dim == 0:
-        amp = func_as[f_index]
-        phase = func_ps[f_index]
-    elif dim == 1:
-        amp = func_as[:, f_index]
-        phase = func_ps[:, f_index]
-    elif dim == 2:
-        amp = func_as[:, :, f_index]
-        phase = func_ps[:, :, f_index]
-    elif dim == 3:
-        amp = func_as[:, :, :, f_index]
-        phase = func_ps[:, :, :, f_index]
-    else:
-        raise ValueError("dim must be 0, 1, 2, or 3")
+    # Create a tuple of slice objects with the frequency index at the correct dimension
+    idx = create_index_at_dim(func_as.ndim, dim, f_index)
+    amp = func_as[idx]
+    phase = func_ps[idx]
 
     return amp.squeeze(), phase.squeeze(), f[f_index]
 

tests/matlab_test_data_collectors/run_all_collectors.m

@@ -3,17 +3,16 @@
 addpath(genpath('../../../k-wave'));
 directory = pwd + "/matlab_collectors";
 files = getListOfFiles(directory);
-% remove this file.
 
 for idx=1:length(files)
         % ensure collected value directory has been created
         file_parts = split(files(idx),["_","."]);
         collected_value_dir = pwd + ...
             "/matlab_collectors/collectedValues/" + file_parts(2);
         mkdir(collected_value_dir)
-        % run value collector
-        run(fullfile(directory, files{idx}));
-        clearvars -except idx files directory
+    % run value collector
+    run(fullfile(directory, files{idx}));
+    clearvars -except idx files directory
 end
 
 if ~isRunningInCI()

tests/test_math.py

@@ -91,11 +91,11 @@ def test_affine_functions_equivalent(self):
         # Test deprecation warning
         with pytest.warns(DeprecationWarning, match="get_affine_matrix is deprecated as of 0.4.1") as warns:
             old_result = get_affine_matrix(translation, rotation)
-            
+
         # Verify warning details
         assert len(warns) == 1
         assert "will be removed in 0.5.0" in str(warns[0].message)
-        
+
         # Test functional equivalence
         new_result = make_affine(translation, rotation)
         assert np.allclose(old_result, new_result)
@@ -111,11 +111,11 @@ def test_shift_functions_equivalent(self):
         # Test deprecation warning
         with pytest.warns(DeprecationWarning, match="fourier_shift is deprecated as of 0.4.1") as warns:
             old_result = fourier_shift(signal, shift)
-            
+
         # Verify warning details
         assert len(warns) == 1
         assert "will be removed in 0.5.0" in str(warns[0].message)
-        
+
         # Test functional equivalence
         new_result = phase_shift_interpolate(signal, shift)
         assert np.allclose(old_result, new_result)

tests/test_utils.py

@@ -97,6 +97,64 @@ def test_extract_amp_phase():
     assert (abs(c_t - c) < 100).all()
 
 
+def test_extract_amp_phase_2d():
+    # Create a 2D test signal with 2 channels
+    Fs = 10_000_000  # Sample frequency
+    source_freq = 2.5 * 1_000  # Signal frequency
+    amp_1 = 2.0
+    amp_2 = 1.0
+    phase_1 = np.pi / 8
+    phase_2 = np.pi / 4
+    # Create a 2D test signal with 2 channels
+    sig_1 = amp_1 * np.sin(source_freq * 2 * np.pi * np.arange(Fs) / Fs + phase_1)
+    sig_2 = amp_2 * np.sin(source_freq * 2 * np.pi * np.arange(Fs) / Fs + phase_2)
+    test_signal = np.vstack([sig_1, sig_2])
+    # plt.plot(test_signal[0])
+    # plt.show()
+    amp, phase, f = extract_amp_phase(test_signal, Fs, source_freq, dim=1)
+
+    assert np.allclose(amp, np.array([amp_1, amp_2]))
+    # Phase is not used in any k-wave-python examples
+    # assert np.allclose(phase, np.array([-phase_1, -phase_2]))
+    assert np.allclose(f, source_freq)
+
+
+def test_extract_amp_phase_double_freq():
+    # Create a test signal with double the detection frequency
+    Fs = 10_000_000  # Sample frequency
+    source_freq = 2.5 * 1_000  # Source frequency
+    detection_freq = 5 * 1_000  # Double the frequency (2 * source_freq)
+
+    amp = 2.0
+    phase = np.pi / 6
+
+    # Create test signal at source_freq
+    t = np.arange(Fs) / Fs
+    test_signal = amp * np.sin(source_freq * 2 * np.pi * t + phase)
+
+    # Extract amplitude and phase at double the frequency
+    # Explicitly set dim=0 for 1D array
+    detected_amp, detected_phase, detected_f = extract_amp_phase(test_signal, Fs, detection_freq, dim=0)
+
+    # The amplitude should be very close to zero since we're detecting at a different frequency
+    # than what's present in the signal
+    assert np.isclose(detected_amp, 0, atol=1e-3)
+    assert np.isclose(detected_f, detection_freq)
+
+    # Now create a signal with the detection frequency
+    test_signal_at_detection = amp * np.sin(detection_freq * 2 * np.pi * t + phase)
+
+    # Extract amplitude and phase at the correct frequency
+    # Explicitly set dim=0 for 1D array
+    detected_amp2, detected_phase2, detected_f2 = extract_amp_phase(test_signal_at_detection, Fs, detection_freq, dim=0)
+
+    # Now the amplitude should match
+    assert np.isclose(detected_amp2, amp, rtol=0.01)
+    # Phase is not used in any k-wave-python examples, but we can verify it's consistent
+    # assert np.isclose(detected_phase2, -phase, rtol=0.01)
+    assert np.isclose(detected_f2, detection_freq)
+
+
 def test_apply_filter_lowpass():
     test_signal = tone_burst(sample_freq=10_000_000, signal_freq=2.5 * 1_000_000, num_cycles=2, envelope="Gaussian")
     filtered_signal = apply_filter(test_signal, Fs=1e7, cutoff_f=1e7, filter_type="LowPass")