[fix] match vbeam parameters (#22)

* Changes that seem to help with tolerance? double-check

* Add single-transmit test?

* Was missing a tukey-alpha=0 that was causing an issue

* FIx grid shape typo

* Fix the atol rtol issues

* Plot the difference map in db instead of fraction

* Fix lint

* Some cleanup

* Some cleanup

* Clean up some prints

* Minor fix

* More cleanup

---------

Co-authored-by: Charles Guan <3221512+charlesincharge@users.noreply.github.com>

Author: charlesbmi

pyproject.toml

@@ -186,6 +186,9 @@ preview = true
 filterwarnings = [
     "error",  # make warnings errors
     "ignore:.*This will add latency due to CPU<->GPU memory transfers.*:UserWarning",
+    # vbeam warnings
+    "ignore:point_position will be overwritten by the scan.:UserWarning",
+    "ignore:Both point_position and scan are set. Scan will be used.:UserWarning",
 ]
 markers = [
     # see conftest.py for default-addition of cuda marker

src/mach/_vis.py

@@ -1,6 +1,5 @@
 """Visualization utilities for test-diagnostics."""
 
-from copy import copy
 from pathlib import Path
 from typing import Optional
 
@@ -95,7 +94,7 @@ def save_debug_figures(
     if reference_result is not None:
         fig, axes = plt.subplots(2, 2, figsize=(12, 10))
 
-        max_value = max(np.max(np.abs(our_img)), np.max(np.abs(ref_img)))
+        max_value = max(np.max(np.abs(our_img)), np.max(np.abs(ref_img)), 1e-12)
 
         # Our result - convert to dB
         our_img_db = db(our_img / max_value)
@@ -131,15 +130,14 @@ def save_debug_figures(
         cbar3 = plt.colorbar(im3, ax=axes[1, 0])
         cbar3.set_label("Linear")
 
-        # Relative difference
-        rel_diff = np.abs(diff_img) / (np.abs(ref_img) + 1e-10)
-        cmap = copy(plt.cm.hot)
-        cmap.set_over("blue", 1.0)
-        im4 = axes[1, 1].imshow(rel_diff.T, aspect="auto", origin="upper", cmap=cmap, extent=extent, vmin=0, vmax=1)
-        axes[1, 1].set_title("Relative Difference")
+        # Relative difference in dB
+        diff_db = db(diff_img / max_value)
+        im4 = axes[1, 1].imshow(diff_db.T, aspect="auto", origin="upper", cmap="hot", extent=extent, vmin=-140, vmax=0)
+        axes[1, 1].set_title("Difference (dB, 0dB = max(ref, our))")
         axes[1, 1].set_xlabel("Lateral [cm]")
         axes[1, 1].set_ylabel("Depth [cm]")
-        plt.colorbar(im4, ax=axes[1, 1], extend="max")
+        cbar4 = plt.colorbar(im4, ax=axes[1, 1])
+        cbar4.set_label("dB")
 
     else:
         fig, ax = plt.subplots(1, 1, figsize=(8, 6))

src/mach/experimental.py

@@ -79,7 +79,6 @@ def beamform(
     for transmit_idx in range(n_transmits):
         # Extract single-transmit data
         single_channel_data = channel_data[transmit_idx]
-        single_rx_start_s = rx_start_s[transmit_idx]
 
         # Call single-transmit beamform
         _ = kernel.beamform(
@@ -88,7 +87,7 @@ def beamform(
             scan_coords_m,
             tx_wave_arrivals_s[transmit_idx],
             out=out,
-            rx_start_s=single_rx_start_s,
+            rx_start_s=rx_start_s,
             sampling_freq_hz=sampling_freq_hz,
             f_number=f_number,
             sound_speed_m_s=sound_speed_m_s,

src/mach/io/uff.py

@@ -46,6 +46,7 @@ def compute_tx_wave_arrivals_s(
         scan_coords_m: Output positions for beamforming (N, 3) in meters
         speed_of_sound: Speed of sound in the medium
         origin: Wave origin point, defaults to [0, 0, 0]
+            Can also be parsed from ultrasound_angles_to_cartesian(channel_data.sequence[idx].origin)
         xp: Array namespace (optional, will be inferred from scan_coords_m)
 
     Returns:
@@ -159,9 +160,12 @@ def create_beamforming_setup(channel_data, scan, f_number: float = 1.7, xp=None)
     # Compute transmit arrivals for all transmits
     directions = extract_wave_directions(channel_data.sequence, xp or np)
     tx_wave_arrivals_s = compute_tx_wave_arrivals_s(directions, scan_coords_m, speed_of_sound, xp=xp)
+    # further delay each transmit by the delay of the wave
+    tx_wave_arrivals_s = tx_wave_arrivals_s + extract_sequence_delays(channel_data.sequence, xp)[:, None]
 
-    # Extract timing information for all transmits
-    rx_delays = extract_sequence_delays(channel_data.sequence, xp)
+    # Account for initial_time offset (this is how vbeam handles it)
+    # The initial_time represents when the first sample was acquired relative to t=0
+    rx_start_s = float(channel_data.initial_time)
 
     return {
         "channel_data": signal,
@@ -173,7 +177,7 @@ def create_beamforming_setup(channel_data, scan, f_number: float = 1.7, xp=None)
         "sampling_freq_hz": sampling_freq_hz,
         "sound_speed_m_s": speed_of_sound,
         "modulation_freq_hz": modulation_frequency,
-        "rx_start_s": rx_delays,
+        "rx_start_s": rx_start_s,
     }
 
 
@@ -198,8 +202,7 @@ def create_single_transmit_beamforming_setup(
 
     # Extract single transmit data
     single_setup = multi_setup.copy()
-    single_setup["channel_data"] = multi_setup["channel_data"][wave_index]  # Remove transmit dimension
-    single_setup["tx_wave_arrivals_s"] = multi_setup["tx_wave_arrivals_s"][wave_index]  # Remove transmit dimension
-    single_setup["rx_start_s"] = multi_setup["rx_start_s"][wave_index]
+    single_setup["channel_data"] = multi_setup["channel_data"][wave_index]
+    single_setup["tx_wave_arrivals_s"] = multi_setup["tx_wave_arrivals_s"][wave_index]
 
     return single_setup

tests/compare/test_vbeam.py

@@ -37,9 +37,9 @@
 from vbeam.scan import LinearScan
 from vbeam.wavefront import PlaneWavefront, ReflectedWavefront
 
-from mach import experimental
+from mach import experimental, kernel
 from mach._vis import save_debug_figures
-from mach.io.uff import create_beamforming_setup
+from mach.io.uff import create_beamforming_setup, create_single_transmit_beamforming_setup
 
 # ============================================================================
 # Fixtures
@@ -178,14 +178,44 @@ def vbeam_setup_uff(
 
 
 @pytest.fixture(scope="module")
-def picmus_phantom_resolution_beamform_kwargs(
+def mach_beamform_kwargs(
     picmus_phantom_resolution_channel_data: ChannelData, picmus_phantom_resolution_scan: Scan
 ) -> dict:
     """mach kwargs for UFF data."""
     return create_beamforming_setup(
         channel_data=picmus_phantom_resolution_channel_data,
         scan=picmus_phantom_resolution_scan,
-        xp=cp if HAS_CUPY else None,
+        xp=cp if HAS_CUPY else np,
+    )
+
+
+# ============================================================================
+# Single Transmit Test Fixtures
+# ============================================================================
+
+
+@pytest.fixture(scope="module", params=[0, 1, 10, 37, 74])
+def transmit_idx(request):
+    """Parametrized transmit index for single-transmit testing."""
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def vbeam_setup_uff_single_transmit(vbeam_setup_uff: SignalForPointSetup, transmit_idx: int) -> SignalForPointSetup:
+    """Create a single-transmit vbeam setup from the full UFF setup."""
+    return vbeam_setup_uff.slice["transmits", transmit_idx]
+
+
+@pytest.fixture(scope="module")
+def mach_single_transmit_kwargs(
+    picmus_phantom_resolution_channel_data: ChannelData, picmus_phantom_resolution_scan: Scan, transmit_idx: int
+) -> dict:
+    """mach kwargs for UFF data."""
+    return create_single_transmit_beamforming_setup(
+        channel_data=picmus_phantom_resolution_channel_data,
+        scan=picmus_phantom_resolution_scan,
+        wave_index=transmit_idx,
+        xp=cp if HAS_CUPY else np,
     )
 
 
@@ -261,13 +291,66 @@ def vbeam_beamform():
 
 @pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available")
 @pytest.mark.filterwarnings("ignore:array is not contiguous, rearranging will add latency:UserWarning")
-def test_mach_matches_vbeam(
-    picmus_phantom_resolution_beamform_kwargs, vbeam_setup_uff: SignalForPointSetup, output_dir
+def test_mach_matches_vbeam_single_transmit(
+    mach_single_transmit_kwargs: dict,
+    vbeam_setup_uff_single_transmit: SignalForPointSetup,
+    transmit_idx: int,
+    output_dir,
 ):
+    """Test mach vs vbeam on a single plane wave transmit to isolate core beamforming differences."""
+    grid_shape = vbeam_setup_uff_single_transmit.scan.shape
+
+    # Run mach single-transmit beamforming using kernel.beamform directly
+    gpu_result = kernel.beamform(**mach_single_transmit_kwargs, tukey_alpha=0.0)
+    result = cp.asnumpy(gpu_result)
+    # Reshape to (x, z)
+    result = result.reshape(grid_shape)
+
+    # Verify basic properties
+    assert np.isfinite(result).all()
+
+    # Run vbeam single-transmit beamforming
+    beamformer = get_das_beamformer(
+        vbeam_setup_uff_single_transmit,
+        compensate_for_apodization_overlap=False,
+        log_compress=False,
+        scan_convert=False,
+    )
+    vbeam_result_jax = beamformer(**vbeam_setup_uff_single_transmit.data).block_until_ready()
+    vbeam_result = np.asarray(vbeam_result_jax)
+
+    # Save debug output if requested
+    if output_dir is not None:
+        output_dir = output_dir / "single_transmit_comparison" / f"transmit_{transmit_idx}"
+        save_debug_figures(
+            our_result=np.abs(result),
+            reference_result=np.abs(vbeam_result),
+            grid_shape=grid_shape,
+            x_axis=vbeam_setup_uff_single_transmit.scan.x,
+            z_axis=vbeam_setup_uff_single_transmit.scan.z,
+            output_dir=output_dir,
+            test_name=f"single_transmit_{transmit_idx}",
+            our_label="mach",
+            reference_label="vbeam",
+        )
+
+    np.testing.assert_allclose(
+        actual=result,
+        desired=vbeam_result,
+        atol=0.01,
+        rtol=1 / 100,
+        err_msg=f"mach single transmit {transmit_idx} results do not match vbeam within expected tolerances",
+    )
+
+
+@pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available")
+@pytest.mark.filterwarnings("ignore:array is not contiguous, rearranging will add latency:UserWarning")
+def test_mach_matches_vbeam(mach_beamform_kwargs, vbeam_setup_uff: SignalForPointSetup, output_dir):
     """Validate mach against vbeam output on a PICMUS UFF data file."""
     grid_shape = vbeam_setup_uff.scan.shape
 
-    gpu_result = experimental.beamform(**picmus_phantom_resolution_beamform_kwargs)
+    # Match our custom vbeam apodization settings
+    gpu_result = experimental.beamform(**mach_beamform_kwargs, tukey_alpha=0.0)
     result = cp.asnumpy(gpu_result)
     # Reshape to (x, z)
     result = result.reshape(grid_shape)
@@ -286,7 +369,7 @@ def test_mach_matches_vbeam(
     vbeam_result_jax = beamformer(**vbeam_setup_uff.data).block_until_ready()
     vbeam_result = np.asarray(vbeam_result_jax)
 
-    # Compare magnitudes because we handle the phase slightly differently from vbeam
+    # Also show magnitude comparison for reference
     vbeam_magnitude = np.abs(vbeam_result)
     cuda_magnitude = np.abs(result)
 
@@ -306,14 +389,12 @@ def test_mach_matches_vbeam(
         )
         print("Saved debug figures to", output_dir)
 
-    # Validate mach against vbeam
-    # TODO: may want to further fine-tune these tolerances
     np.testing.assert_allclose(
-        actual=cuda_magnitude,
-        desired=vbeam_magnitude,
-        atol=10,
-        rtol=0.3,
-        err_msg="mach results do not match vbeam within expected tolerances",
+        actual=result,
+        desired=vbeam_result,
+        atol=0.01,
+        rtol=1 / 100,
+        err_msg="mach complex results do not match vbeam within expected tolerances (with scaling correction)",
     )