Fix batching (#192)

agencyenterprise · Jan 2, 2025 · c5e4616 · c5e4616
1 parent dfb468f
commit c5e4616
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/src/neurotechdevkit/imaging/beamform.py b/src/neurotechdevkit/imaging/beamform.py
@@ -90,9 +90,15 @@ def beamform_delay_and_sum(
     )
 
     # Beamform
-    iq_signals_column_vec = iq_signals.reshape(-1, 1)
+    if iq_signals.ndim == 2:
+        iq_signals_column_vec = iq_signals.reshape(-1, 1)
+    else:
+        iq_signals_column_vec = iq_signals.reshape(-1, num_echoes)
     beamformed_iq_signals = das_matrix @ iq_signals_column_vec
-    beamformed_iq_signals = beamformed_iq_signals.reshape(x.shape)
+    if iq_signals.ndim == 2:
+        beamformed_iq_signals = beamformed_iq_signals.reshape(x.shape)
+    else:
+        beamformed_iq_signals = beamformed_iq_signals.reshape(x.shape + (-1,))
 
     return beamformed_iq_signals