zea 0.0.9__tar.gz → 0.0.10__tar.gz

{zea-0.0.9 → zea-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zea
-Version: 0.0.9
+Version: 0.0.10
 Summary: A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework.
 License-File: LICENSE
 Keywords: ultrasound,machine learning,beamforming
@@ -44,8 +44,6 @@ Requires-Dist: jax ; extra == "backends"
 Requires-Dist: jax[cuda12-pip] (>=0.4.26) ; extra == "jax"
 Requires-Dist: keras (>=3.12)
 Requires-Dist: matplotlib (>=3.8)
-Requires-Dist: mock ; extra == "dev"
-Requires-Dist: mock ; extra == "docs"
 Requires-Dist: myst-parser ; extra == "dev"
 Requires-Dist: myst-parser ; extra == "docs"
 Requires-Dist: nbsphinx ; extra == "dev"
@@ -76,8 +74,6 @@ Requires-Dist: simpleitk (>=2.2.1) ; extra == "dev"
 Requires-Dist: simpleitk (>=2.2.1) ; extra == "tests"
 Requires-Dist: sphinx ; extra == "dev"
 Requires-Dist: sphinx ; extra == "docs"
-Requires-Dist: sphinx-argparse ; extra == "dev"
-Requires-Dist: sphinx-argparse ; extra == "docs"
 Requires-Dist: sphinx-autobuild ; extra == "dev"
 Requires-Dist: sphinx-autobuild ; extra == "docs"
 Requires-Dist: sphinx-autodoc-typehints ; extra == "dev"
@@ -88,6 +84,8 @@ Requires-Dist: sphinx-reredirects ; extra == "dev"
 Requires-Dist: sphinx-reredirects ; extra == "docs"
 Requires-Dist: sphinx_design ; extra == "dev"
 Requires-Dist: sphinx_design ; extra == "docs"
+Requires-Dist: sphinxcontrib-autoprogram ; extra == "dev"
+Requires-Dist: sphinxcontrib-autoprogram ; extra == "docs"
 Requires-Dist: sphinxcontrib-bibtex ; extra == "dev"
 Requires-Dist: sphinxcontrib-bibtex ; extra == "docs"
 Requires-Dist: tensorflow ; extra == "backends"

{zea-0.0.9 → zea-0.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "zea"
-version = "0.0.9"
+version = "0.0.10"
 description = "A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework."
 authors = [
     { name = "Tristan Stevens", email = "t.s.w.stevens@tue.nl" },
@@ -65,10 +65,9 @@ dev = [
     "sphinx-autodoc-typehints",
     "sphinx-copybutton",
     "sphinx_design",
-    "sphinx-argparse",
+    "sphinxcontrib-autoprogram",
     "sphinx-reredirects",
     "sphinxcontrib-bibtex",
-    "mock",
     "myst-parser",
     "nbsphinx",
     "furo",
@@ -97,10 +96,9 @@ docs = [
     "sphinx-autodoc-typehints",
     "sphinx-copybutton",
     "sphinx_design",
-    "sphinx-argparse",
+    "sphinxcontrib-autoprogram",
     "sphinx-reredirects",
     "sphinxcontrib-bibtex",
-    "mock",
     "myst-parser",
     "nbsphinx",
     "furo",

{zea-0.0.9 → zea-0.0.10}/zea/__init__.py

@@ -2,12 +2,16 @@
 
 import importlib.util
 import os
+from importlib.metadata import PackageNotFoundError, version
 
 from . import log
 
-# dynamically add __version__ attribute (see pyproject.toml)
-# __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.9"
+try:
+    # dynamically add __version__ attribute (see pyproject.toml)
+    __version__ = version("zea")
+except PackageNotFoundError:
+    # Package is not installed (e.g., running from source)
+    __version__ = "dev"
 
 
 def _bootstrap_backend():
@@ -80,8 +84,10 @@ def _bootstrap_backend():
     log.info(f"Using backend {keras_backend()!r}")
 
 
-# call and clean up namespace
-_bootstrap_backend()
+# Skip backend bootstrap when building on ReadTheDocs
+if os.environ.get("READTHEDOCS") != "True":
+    _bootstrap_backend()
+
 del _bootstrap_backend
 
 from . import (

{zea-0.0.9 → zea-0.0.10}/zea/agent/masks.py

@@ -4,6 +4,8 @@ Mask generation utilities.
 These masks are used as a measurement operator for focused scan-line subsampling.
 """
 
+from __future__ import annotations
+
 from typing import List
 
 import keras
@@ -117,11 +119,21 @@ def initial_equispaced_lines(
         Tensor: k-hot-encoded line vector of shape (n_possible_actions).
             Needs to be converted to image size.
     """
+    assert n_actions > 0, "Number of actions must be > 0."
+    assert n_possible_actions > 0, "Number of possible actions must be > 0."
+    assert n_actions <= n_possible_actions, (
+        "Number of actions must be less than or equal to number of possible actions."
+    )
+
     if assert_equal_spacing:
         _assert_equal_spacing(n_actions, n_possible_actions)
-        selected_indices = ops.arange(0, n_possible_actions, n_possible_actions // n_actions)
-    else:
-        selected_indices = ops.linspace(0, n_possible_actions - 1, n_actions, dtype="int32")
+
+    # Distribute indices as evenly as possible
+    # This approach ensures spacing differs by at most 1 when not divisible
+    step = n_possible_actions / n_actions
+    selected_indices = ops.cast(
+        ops.round(ops.arange(0, n_actions, dtype="float32") * step), "int32"
+    )
 
     return indices_to_k_hot(selected_indices, n_possible_actions, dtype=dtype)
 

{zea-0.0.9 → zea-0.0.10}/zea/agent/selection.py

@@ -96,6 +96,7 @@ class GreedyEntropy(LinesActionModel):
         std_dev: float = 1,
         num_lines_to_update: int = 5,
         entropy_sigma: float = 1.0,
+        average_entropy_across_batch: bool = False,
     ):
         """Initialize the GreedyEntropy action selection model.
 
@@ -110,6 +111,10 @@ class GreedyEntropy(LinesActionModel):
                 to update. Must be odd.
             entropy_sigma (float, optional): The standard deviation of the Gaussian
                 Mixture components used to approximate the posterior.
+            average_entropy_across_batch (bool, optional): Whether to average entropy
+                across the batch when selecting lines. This can be useful when
+                selecting planes in 3D imaging, where the batch dimension represents
+                a third spatial dimension. Defaults to False.
         """
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 
@@ -117,6 +122,7 @@ class GreedyEntropy(LinesActionModel):
         # of the selected line is set to 0 once it's been selected.
         assert num_lines_to_update % 2 == 1, "num_samples must be odd."
         self.num_lines_to_update = num_lines_to_update
+        self.average_entropy_across_batch = average_entropy_across_batch
 
         # see here what I mean by upside_down_gaussian:
         # https://colab.research.google.com/drive/1CQp_Z6nADzOFsybdiH5Cag0vtVZjjioU?usp=sharing
@@ -153,7 +159,7 @@ class GreedyEntropy(LinesActionModel):
         assert particles.shape[1] > 1, "The entropy cannot be approximated using a single particle."
 
         if n_possible_actions is None:
-            n_possible_actions = particles.shape[-1]
+            n_possible_actions = ops.shape(particles)[-1]
 
         # TODO: I think we only need to compute the lower triangular
         # of this matrix, since it's symmetric
@@ -164,7 +170,8 @@ class GreedyEntropy(LinesActionModel):
         # Vertically stack all columns corresponding with the same line
         # This way we can just sum across the height axis and get the entropy
         # for each pixel in a given line
-        batch_size, n_particles, _, height, _ = gaussian_error_per_pixel_i_j.shape
+        batch_size, n_particles, _, height, _ = ops.shape(gaussian_error_per_pixel_i_j)
+
         gaussian_error_per_pixel_stacked = ops.transpose(
             ops.reshape(
                 ops.transpose(gaussian_error_per_pixel_i_j, (0, 1, 2, 4, 3)),
@@ -274,6 +281,8 @@ class GreedyEntropy(LinesActionModel):
 
         pixelwise_entropy = self.compute_pixelwise_entropy(particles)
         linewise_entropy = ops.sum(pixelwise_entropy, axis=1)
+        if self.average_entropy_across_batch:
+            linewise_entropy = ops.expand_dims(ops.mean(linewise_entropy, axis=0), axis=0)
 
         # Greedily select best line, reweight entropies, and repeat
         all_selected_lines = []
@@ -334,7 +343,7 @@ class EquispacedLines(LinesActionModel):
         n_possible_actions: int,
         img_width: int,
         img_height: int,
-        assert_equal_spacing=True,
+        assert_equal_spacing: bool = True,
     ):
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 

{zea-0.0.9 → zea-0.0.10}/zea/backend/__init__.py

@@ -59,8 +59,21 @@ def _import_torch():
         return None
 
 
+def _get_backend():
+    try:
+        backend_result = keras.backend.backend()
+        if isinstance(backend_result, str):
+            return backend_result
+        else:
+            # to handle mocked backends during testing
+            return None
+    except Exception:
+        return None
+
+
 tf_mod = _import_tf()
 jax_mod = _import_jax()
+backend = _get_backend()
 
 
 def tf_function(func=None, jit_compile=False, **kwargs):
@@ -184,7 +197,7 @@ class on_device:
         self._context.__exit__(exc_type, exc_val, exc_tb)
 
 
-if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
+if backend in [None, "tensorflow", "jax", "numpy"]:
 
     def func_on_device(func, device, *args, **kwargs):
         """Moves all tensor arguments of a function to a specified device before calling it.
@@ -199,7 +212,8 @@ if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
         """
         with on_device(device):
             return func(*args, **kwargs)
-elif keras.backend.backend() == "torch":
+
+elif backend == "torch":
     from zea.backend.torch import func_on_device
 else:
-    raise ValueError(f"Unsupported backend: {keras.backend.backend()}")
+    raise ValueError(f"Unsupported backend: {backend}")

{zea-0.0.9 → zea-0.0.10}/zea/beamform/beamformer.py

@@ -4,7 +4,7 @@ import keras
 import numpy as np
 from keras import ops
 
-from zea.beamform.lens_correction import calculate_lens_corrected_delays
+from zea.beamform.lens_correction import compute_lens_corrected_travel_times
 from zea.func.tensor import vmap
 
 
@@ -62,6 +62,7 @@ def tof_correction(
     focus_distances,
     t_peak,
     tx_waveform_indices,
+    transmit_origins,
     apply_lens_correction=False,
     lens_thickness=1e-3,
     lens_sound_speed=1000,
@@ -87,6 +88,7 @@ def tof_correction(
             Shape `(n_waveforms,)`.
         tx_waveform_indices (ops.Tensor): The indices of the waveform used for each
             transmit of shape `(n_tx,)`.
+        transmit_origins (ops.Tensor): Transmit origins of shape (n_tx, 3).
         apply_lens_correction (bool, optional): Whether to apply lens correction to
             time-of-flights. This makes it slower, but more accurate in the near-field.
             Defaults to False.
@@ -120,8 +122,7 @@ def tof_correction(
     # rxdel has shape (n_el, n_pix)
     # --------------------------------------------------------------------
 
-    delay_fn = calculate_lens_corrected_delays if apply_lens_correction else calculate_delays
-    txdel, rxdel = delay_fn(
+    txdel, rxdel = calculate_delays(
         flatgrid,
         t0_delays,
         tx_apodizations,
@@ -133,10 +134,12 @@ def tof_correction(
         n_el,
         focus_distances,
         polar_angles,
-        t_peak=t_peak,
-        tx_waveform_indices=tx_waveform_indices,
-        lens_thickness=lens_thickness,
-        lens_sound_speed=lens_sound_speed,
+        t_peak,
+        tx_waveform_indices,
+        transmit_origins,
+        apply_lens_correction,
+        lens_thickness,
+        lens_sound_speed,
     )
 
     n_pix = ops.shape(flatgrid)[0]
@@ -207,7 +210,11 @@ def calculate_delays(
     polar_angles,
     t_peak,
     tx_waveform_indices,
-    **kwargs,
+    transmit_origins,
+    apply_lens_correction=False,
+    lens_thickness=None,
+    lens_sound_speed=None,
+    n_iter=2,
 ):
     """Calculates the delays in samples to every pixel in the grid.
 
@@ -242,6 +249,16 @@ def calculate_delays(
             `(n_waveforms,)`.
         tx_waveform_indices (Tensor): The indices of the waveform used for each
             transmit of shape `(n_tx,)`.
+        transmit_origins (Tensor): Transmit origins of shape (n_tx, 3).
+        apply_lens_correction (bool, optional): Whether to apply lens correction to
+            time-of-flights. This makes it slower, but more accurate in the near-field.
+            Defaults to False.
+        lens_thickness (float, optional): Thickness of the lens in meters. Used for
+            lens correction.
+        lens_sound_speed (float, optional): Speed of sound in the lens in m/s. Used
+            for lens correction.
+        n_iter (int, optional): Number of iterations for the Newton-Raphson method
+            used in lens correction. Defaults to 2.
 
 
     Returns:
@@ -252,38 +269,56 @@ def calculate_delays(
             `(n_pix, n_el)`.
     """
 
-    def _tx_distances(polar_angles, t0_delays, tx_apodizations, focus_distances):
-        return distance_Tx_generic(
-            grid,
-            t0_delays,
-            tx_apodizations,
+    # Validate input shapes
+    for arr in [t0_delays, grid, tx_apodizations, probe_geometry]:
+        assert arr.ndim == 2
+    assert probe_geometry.shape[0] == n_el
+    assert t0_delays.shape[0] == n_tx
+
+    if not apply_lens_correction:
+        # Compute receive distances in meters of shape (n_pix, n_el)
+        rx_distances = distance_Rx(grid, probe_geometry)
+
+        # Convert distances to delays in seconds
+        rx_delays = rx_distances / sound_speed
+    else:
+        # Compute lens-corrected travel times from each element to each pixel
+        assert lens_thickness is not None, "lens_thickness must be provided for lens correction."
+        assert lens_sound_speed is not None, (
+            "lens_sound_speed must be provided for lens correction."
+        )
+        rx_delays = compute_lens_corrected_travel_times(
             probe_geometry,
-            focus_distances,
-            polar_angles,
+            grid,
+            lens_thickness,
+            lens_sound_speed,
             sound_speed,
+            n_iter=n_iter,
         )
 
-    tx_distances = vmap(_tx_distances)(polar_angles, t0_delays, tx_apodizations, focus_distances)
-    tx_distances = ops.transpose(tx_distances, (1, 0))
-    # tx_distances shape is now (n_pix, n_tx)
+    # Compute transmit delays
+    tx_delays = vmap(transmit_delays, in_axes=(None, 0, 0, None, 0, 0, 0, None, 0), out_axes=1)(
+        grid,
+        t0_delays,
+        tx_apodizations,
+        rx_delays,
+        focus_distances,
+        polar_angles,
+        initial_times,
+        None,
+        transmit_origins,
+    )
 
-    # Compute receive distances
-    def _rx_distances(probe_geometry):
-        return distance_Rx(grid, probe_geometry)
+    # Add the offset to the transmit peak time
+    tx_delays += ops.take(t_peak, tx_waveform_indices)[None]
 
-    rx_distances = vmap(_rx_distances)(probe_geometry)
-    rx_distances = ops.transpose(rx_distances, (1, 0))
-    # rx_distances shape is now (n_pix, n_el)
+    # TODO: nan to num needed?
+    # tx_delays = ops.nan_to_num(tx_delays, nan=0.0, posinf=0.0, neginf=0.0)
+    # rx_delays = ops.nan_to_num(rx_delays, nan=0.0, posinf=0.0, neginf=0.0)
 
-    # Compute the delays [in samples] from the distances
-    # The units here are ([m]/[m/s]-[s])*[1/s] resulting in a unitless quantity
-    # TODO: Add pulse width to transmit delays
-    tx_delays = (
-        tx_distances / sound_speed
-        - initial_times[None]
-        + ops.take(t_peak, tx_waveform_indices)[None]
-    ) * sampling_frequency
-    rx_delays = (rx_distances / sound_speed) * sampling_frequency
+    # Convert from seconds to samples
+    tx_delays *= sampling_frequency
+    rx_delays *= sampling_frequency
 
     return tx_delays, rx_delays
 
@@ -414,7 +449,7 @@ def complex_rotate(iq, theta):
 def distance_Rx(grid, probe_geometry):
     """Computes distance to user-defined pixels from elements.
 
-    Expects all inputs to be numpy arrays specified in SI units.
+    Expects all inputs to be arrays specified in SI units.
 
     Args:
         grid (ops.Tensor): Pixel positions in x,y,z of shape `(n_pix, 3)`.
@@ -425,83 +460,117 @@ def distance_Rx(grid, probe_geometry):
             `(n_pix, n_el)`.
     """
     # Get norm of distance vector between elements and pixels via broadcasting
-    dist = ops.linalg.norm(grid - probe_geometry[None, ...], axis=-1)
+    dist = ops.linalg.norm(grid[:, None, :] - probe_geometry[None, :, :], axis=-1)
     return dist
 
 
-def distance_Tx_generic(
+def transmit_delays(
     grid,
     t0_delays,
     tx_apodization,
-    probe_geometry,
+    rx_delays,
     focus_distance,
     polar_angle,
-    sound_speed=1540,
+    initial_time,
+    azimuth_angle=None,
+    transmit_origin=None,
 ):
-    """Generic transmit distance calculation.
+    """
+    Computes the transmit delay from transmission to each pixel in the grid.
+
+    Uses the first-arrival time for pixels before the focus (or virtual source)
+    and the last-arrival time for pixels beyond the focus.
 
-    Computes distance to user-defined pixels for generic transmits based on
-    the t0_delays.
+    The receive delays can be precomputed since they do not depend on the
+    transmit parameters.
 
     Args:
-        grid (ops.Tensor): Flattened tensor of pixel positions in x,y,z of shape
-            `(n_pix, 3)`
-        t0_delays (ops.Tensor): The transmit delays in seconds of shape `(n_el,)`,
-            shifted such that the smallest delay is 0. Defaults to None.
-        tx_apodization (ops.Tensor): The transmit apodizations of shape
-            `(n_el,)`.
-        probe_geometry (ops.Tensor): The positions of the transducer elements of shape
-            `(n_el, 3)`.
+        grid (ops.Tensor): Flattened tensor of pixel positions in x,y,z of shape `(n_pix, 3)`
+        t0_delays (Tensor): The transmit delays in seconds of shape (n_el,).
+        tx_apodization (Tensor): The transmit apodization of shape (n_el,).
+        rx_delays (Tensor): The travel times in seconds from elements to pixels
+            of shape (n_pix, n_el).
         focus_distance (float): The focus distance in meters.
         polar_angle (float): The polar angle in radians.
-        sound_speed (float): The speed of sound in m/s. Defaults to 1540.
+        initial_time (float): The initial time for this transmit in seconds.
+        azimuth_angle (float, optional): The azimuth angle in radians. Defaults to 0.0.
+        transmit_origin (ops.Tensor, optional): The origin of the transmit beam of shape (3,).
+            If None, defaults to (0, 0, 0). Defaults to None.
 
     Returns:
-        Tensor: Distance from each pixel to each element in meters
-        of shape `(n_pix,)`
+        Tensor: The transmit delays of shape `(n_pix,)`.
     """
-    # Get the individual x, y, and z components of the pixel coordinates
-    x = grid[:, 0]
-    y = grid[:, 1]
-    z = grid[:, 2]
-
-    # Reshape x, y, and z to shape (n_pix, 1)
-    x = x[..., None]
-    y = y[..., None]
-    z = z[..., None]
-
-    # Get the individual x, y, and z coordinates of the elements and add a
-    # dummy dimension at the beginning to shape (1, n_el).
-    ele_x = probe_geometry[None, :, 0]
-    ele_y = probe_geometry[None, :, 1]
-    ele_z = probe_geometry[None, :, 2]
-
-    # Compute the differences dx, dy, and dz of shape (n_pix, n_el)
-    dx = x - ele_x
-    dy = y - ele_y
-    dz = z - ele_z
-
-    # Define an infinite offset for elements that do not fire to not consider them in
-    # the transmit distance calculation.
+    # Add a large offset for elements that are not used in the transmit to
+    # disqualify them from being the closest element
     offset = ops.where(tx_apodization == 0, np.inf, 0.0)
 
-    # Compute the distance between the elements and the pixels of shape
-    # (n_pix, n_el)
-    dist = t0_delays[None] * sound_speed + ops.sqrt(dx**2 + dy**2 + dz**2)
-
-    # Compute the z-coordinate of the focal point
-    focal_z = ops.cos(polar_angle) * focus_distance
+    # Compute total travel time from t=0 to each pixel via each element
+    # rx_delays has shape (n_pix, n_el)
+    # t0_delays has shape (n_el,)
+    total_times = rx_delays + t0_delays[None, :]
+
+    if azimuth_angle is None:
+        azimuth_angle = ops.zeros_like(polar_angle)
+
+    # Set origin to (0, 0, 0) if not provided
+    if transmit_origin is None:
+        transmit_origin = ops.zeros(3, dtype=grid.dtype)
+
+    # Compute the 3D position of the focal point
+    # The beam direction vector
+    beam_direction = ops.stack(
+        [
+            ops.sin(polar_angle) * ops.cos(azimuth_angle),
+            ops.sin(polar_angle) * ops.sin(azimuth_angle),
+            ops.cos(polar_angle),
+        ]
+    )
 
-    # Compute the effective distance of the pixels to the wavefront by computing the
-    # largest distance over all the elements when the pixel is behind the virtual
-    # source and the smallest distance otherwise.
-    dist = ops.where(
-        ops.cast(ops.sign(focus_distance), "float32") * (grid[:, 2] - focal_z) <= 0.0,
-        ops.min(dist + offset[None], 1),
-        ops.max(dist - offset[None], 1),
+    # Handle plane wave case where focus_distance is set to zero
+    # We use np.inf to consider the first wavefront arrival for all pixels
+    focus_distance = ops.where(focus_distance == 0.0, np.inf, focus_distance)
+
+    # Compute focal point position: origin + focus_distance * beam_direction
+    # For negative focus_distance (diverging/virtual source), this is behind the origin
+    focal_point = transmit_origin + focus_distance * beam_direction  # shape (3,)
+
+    # Deal with plane wave case where focus_distance is infinite and beam_direction is zero
+    # (np.inf * 0.0 -> nan) so we convert nan to zero
+    focal_point = ops.where(ops.isnan(focal_point), 0.0, focal_point)
+
+    # Compute the position of each pixel relative to the focal point
+    pixel_relative_to_focus = grid - focal_point[None, :]  # shape (n_pix, 3)
+
+    # Project onto the beam direction to determine if pixel is before or after focus
+    # Positive projection means pixel is in the direction of beam propagation (beyond focus)
+    # Negative projection means pixel is behind the focus (before focus)
+    projection_along_beam = ops.sum(
+        pixel_relative_to_focus * beam_direction[None, :], axis=-1
+    )  # shape (n_pix,)
+
+    # For focused waves (positive focus_distance):
+    #   - Use min time for pixels before focus (projection < 0)
+    #   - Use max time for pixels beyond focus (projection > 0)
+    # For diverging waves (negative focus_distance, virtual source):
+    #   - The sign of focus_distance flips the logic
+    #   - Use min time for pixels between transducer and virtual source
+    #   - Use max time for pixels beyond transducer
+    is_before_focus = ops.cast(ops.sign(focus_distance), "float32") * projection_along_beam < 0.0
+
+    # Compute the effective time of the pixels to the wavefront by computing the
+    # smallest time over all elements (first wavefront arrival) for pixels before
+    # the focus, and the largest time (last wavefront contribution) for pixels
+    # beyond the focus.
+    tx_delay = ops.where(
+        is_before_focus,
+        ops.min(total_times + offset[None, :], axis=-1),
+        ops.max(total_times - offset[None, :], axis=-1),
     )
 
-    return dist
+    # Subtract the initial time offset for this transmit
+    tx_delay = tx_delay - initial_time
+
+    return tx_delay
 
 
 def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):

{zea-0.0.9 → zea-0.0.10}/zea/beamform/delays.py

@@ -52,7 +52,7 @@ def compute_t0_delays_planewave(probe_geometry, polar_angles, azimuth_angles=0,
 
 
 def compute_t0_delays_focused(
-    origins,
+    transmit_origins,
     focus_distances,
     probe_geometry,
     polar_angles,
@@ -63,12 +63,12 @@ def compute_t0_delays_focused(
     the first element fires at t=0.
 
     Args:
-        origins (np.ndarray): The origin of the focused transmit of shape (n_tx, 3,).
-        focus_distance (float): The distance to the focus.
+        transmit_origins (np.ndarray): The origin of the focused transmit of shape (n_tx, 3,).
+        focus_distances (np.ndarray): The distance to the focus for each transmit of shape (n_tx,).
         probe_geometry (np.ndarray): The positions of the elements in the array of
             shape (element, 3).
-        polar_angles (np.ndarray): The polar angles of the planewave in radians of shape (n_tx,).
-        azimuth_angles (np.ndarray, optional): The azimuth angles of the planewave in
+        polar_angles (np.ndarray): The polar angles in radians of shape (n_tx,).
+        azimuth_angles (np.ndarray, optional): The azimuth angles in
             radians of shape (n_tx,).
         sound_speed (float, optional): The speed of sound. Defaults to 1540.
 
@@ -79,12 +79,15 @@ def compute_t0_delays_focused(
     assert polar_angles.shape == (n_tx,), (
         f"polar_angles must have length n_tx = {n_tx}. Got length {len(polar_angles)}."
     )
-    assert origins.shape == (n_tx, 3), (
-        f"origins must have shape (n_tx, 3). Got shape {origins.shape}."
+    assert transmit_origins.shape == (n_tx, 3), (
+        f"transmit_origins must have shape (n_tx, 3). Got shape {transmit_origins.shape}."
     )
     assert probe_geometry.shape[1] == 3 and probe_geometry.ndim == 2, (
         f"probe_geometry must have shape (element, 3). Got shape {probe_geometry.shape}."
     )
+    assert focus_distances.shape == (n_tx,), (
+        f"focus_distances must have length n_tx = {n_tx}. Got length {len(focus_distances)}."
+    )
 
     # Convert single angles to arrays for broadcasting
     polar_angles = np.atleast_1d(polar_angles)
@@ -107,12 +110,12 @@ def compute_t0_delays_focused(
     )
 
     # Add a new dimension for broadcasting
-    # The shape is now (n_tx, n_el, 3)
+    # The shape is now (n_tx, 1, 3)
     v = np.expand_dims(v, axis=1)
 
     # Compute the location of the virtual source by adding the focus distance
     # to the origin along the wave vectors.
-    virtual_sources = origins[:, None] + focus_distances[:, None, None] * v
+    virtual_sources = transmit_origins[:, None] + focus_distances[:, None, None] * v
 
     # Compute the distances between the virtual sources and each element
     dist = np.linalg.norm(virtual_sources - probe_geometry, axis=-1)