"""Lens corrected delay computation for ultrasound beamforming."""
from keras import ops
[docs]
def calculate_lens_corrected_delays(
grid,
t0_delays,
tx_apodizations,
probe_geometry,
initial_times,
sampling_frequency,
sound_speed,
n_tx,
n_el,
focus_distances,
polar_angles,
lens_sound_speed=1000,
lens_thickness=1e-3,
n_iter=2,
):
"""Compute the lens corrected delays IN SAMPLES for a given grid and transmit scheme.
Args:
grid (ndarray): The grid of shape (n_pixels, 3).
t0_delays (ndarray): The t0 delays of shape (n_tx, n_el).
tx_apodizations (ndarray): The transmit apodizations of shape (n_tx, n_el).
probe_geometry (ndarray): The probe geometry of shape (n_el, 3).
initial_times (ndarray): The initial times of shape (n_tx,).
sampling_frequency (float): The sampling frequency in Hz.
sound_speed (float): The speed of sound in the medium in m/s.
n_tx (int): The number of transmits.
n_el (int): The number of elements.
focus_distances (ndarray): The focus distances of shape (n_tx,).
polar_angles (ndarray): The polar angles of shape (n_tx,).
lens_sound_speed (float): The speed of sound in the lens in m/s.
lens_thickness (float): The thickness of the lens in meters.
n_iter (int): The number of iterations to run the Newton-Raphson method.
Returns:
ndarray: The transmit delays in samples of shape (n_pixels, n_tx).
ndarray: The receive delays in samples of shape (n_pixels, n_el).
"""
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
tx_delays = []
rx_delays = compute_lens_corrected_travel_times(
probe_geometry,
grid,
lens_thickness,
lens_sound_speed,
sound_speed,
n_iter=n_iter,
)
for tx in range(n_tx):
# Add a large offset to elements that are not used in the transmit to
# diqualify them from being the closest element
apod_offset = ops.where(tx_apodizations[tx] == 0, 10.0, 0)
tx_min = ops.min(rx_delays + t0_delays[tx] + apod_offset, axis=-1) + initial_times[tx]
tx_delays.append(tx_min)
tx_delays = ops.stack(tx_delays, axis=-1)
tx_delays *= sampling_frequency
rx_delays *= sampling_frequency
return tx_delays, rx_delays
[docs]
def compute_lens_corrected_travel_times(
element_pos, pixel_pos, lens_thickness, c_lens, c_medium, n_iter=1
):
"""Compute the travel time of the shortest path between the element and the pixel.
Args:
element_pos (ndarray): The position of the element of shape (n_elements, 3).
pixel_pos (ndarray): The position of the pixel of shape (n_pixels, 3).
lens_thickness (float): The thickness of the lens in meters.
c_lens (float): The speed of sound in the lens in m/s.
c_medium (float): The speed of sound in the medium in m/s.
n_iter (int): The number of iterations to run the Newton-Raphson method.
Returns:
ndarray: The travel times of shape (n_pixels, n_elements).
"""
pixel_pos = pixel_pos[:, None] - element_pos[None]
# Project the 3D problem to a 2D problem by shifting all pixels to have the element
# in the origin and then projecting the positions to the plane spanned by
# [pixel_x-element_x, pixel_z-element_z, 0], [0, 0, 1]
xs = ops.norm(pixel_pos[..., :2], axis=-1)
zs = pixel_pos[..., -1]
pixel_pos_2d = ops.stack([xs, zs], axis=-1)
element_pos_2d = ops.zeros((1, element_pos.shape[0], 2))
xl = compute_xl(
element_pos_2d,
pixel_pos_2d,
lens_thickness,
c_lens,
c_medium,
n_iter,
)
# Form the position of the lens crossing point by adding the z-coordinate of the
# lens to the point
pos_lenscrossing = ops.stack([xl, lens_thickness * ops.ones_like(xl)], axis=-1)
indices = ops.array([0, -1])
element_pos = ops.take(element_pos, indices, axis=-1)
# Compute the travel time of the shortest path
travel_time = compute_travel_time(
element_pos_2d, pos_lenscrossing, c_lens
) + compute_travel_time(pos_lenscrossing, pixel_pos_2d, c_medium)
return travel_time
[docs]
def compute_xl(element_pos_2d, pixel_pos_2d, lens_thickness, c_lens, c_medium, n_iter):
"""Computes the lateral point on the lens that the shortest path goes through based
on Fermat's principle.
Args:
element_pos_2d (float): The 2D position of the element.
pixel_pos_2d (float): The 2D position of the pixel.
lens_thickness (float): The thickness of the lens in meters.
c_lens (float): The speed of sound in the lens in m/s.
c_medium (float): The speed of sound in the medium in m/s.
n_iter (int): The number of iterations to run the Newton-Raphson method.
Returns:
float: The x-coordinate of the lateral point on the lens.
"""
xs = pixel_pos_2d[..., 0]
zs = pixel_pos_2d[..., 1]
xe = element_pos_2d[..., 0]
ze = element_pos_2d[..., 1]
# Apply Newton-Raphson method to find the lateral point on the lens that the shortest
# path goes through
xl_init = lens_thickness * (xs - xe) / (zs - ze) + xe
xl = xl_init
for _ in range(n_iter):
xl = xl + dxl(xe, ze, xl, xs, zs, lens_thickness, c_lens, c_medium)
# Clip the lateral point to be in between the element and the pixel
xl = ops.clip(xl, ops.minimum(xs, xe), ops.maximum(xs, xe))
return xl
[docs]
def dxl(xe, ze, xl, xs, zs, zl, c_lens, c_medium):
"""Computes the update step for the lateral point on the lens that the shortest path
using the Newton-Raphson method.
Notes
-----
This result was derived by defining the total travel time through the lens and the
medium as a function of the lateral point on the lens and then taking the
derivative. We then have a function whose root is the lateral point on the lens that
the shortest path goes through. We then compute the derivative and update the
lateral point on the lens using the Newton-Raphson method:
x_new = x - f(x) / f'(x).
"""
eps = 1e-6
numerator = -((xe - xl) / (c_lens * ops.sqrt((xe - xl) ** 2 + (ze - zl) ** 2))) + (
(xl - xs) / (c_medium * ops.sqrt((xl - xs) ** 2 + (zl - zs) ** 2)) + eps
)
denominator = (
-((xe - xl) ** 2 / (c_lens * ((xe - xl) ** 2 + (ze - zl) ** 2) ** (3 / 2) + eps))
+ (1 / (c_lens * ops.sqrt((xe - xl) ** 2 + (ze - zl) ** 2)))
- ((xl - xs) ** 2 / (c_medium * ((xl - xs) ** 2 + (zl - zs) ** 2) ** (3 / 2) + eps))
+ (1 / (c_medium * ops.sqrt((xl - xs) ** 2 + (zl - zs) ** 2) + eps))
)
result = -numerator / (denominator + eps)
# Handle NaNs
result = ops.nan_to_num(result)
# Clip the update step to prevent divergence
# This value is chosen to be small enough to prevent divergence but large enough to
# cover the distance accross a normal ultrasound aperture in a single step.
result = ops.clip(result, -10e-3, 10e-3)
return result
[docs]
def compute_travel_time(pos_a, pos_b, c):
"""Compute the travel time between two points."""
return ops.linalg.norm(pos_a - pos_b, axis=-1) / c