diff --git a/docs/source/references.bib b/docs/source/references.bib index 1ab5a820a..50410842b 100644 --- a/docs/source/references.bib +++ b/docs/source/references.bib @@ -55,6 +55,85 @@ @article{vignon2008capon year = {2008} } +@article{byram2015model, + author = {Byram, Brett and Dei, Kazuyuki and Tierney, Jaime and Dumont, Douglas}, + doi = {10.1109/TUFFC.2015.007004}, + journal = {{IEEE} Transactions on Ultrasonics, Ferroelectrics, and Frequency Control}, + number = {11}, + pages = {1913--1927}, + title = {{A} {M}odel and {R}egularization {S}cheme for {U}ltrasonic {B}eamforming + {C}lutter {R}eduction}, + volume = {62}, + year = {2015} +} + +@article{khan2021realtime, + author = {Khan, Christopher and Dei, Kazuyuki and Schlunk, Siegfried and Ozgun, Kathryn and + Byram, Brett}, + doi = {10.1109/TUFFC.2021.3056334}, + journal = {{IEEE} Transactions on Ultrasonics, Ferroelectrics, and Frequency Control}, + number = {6}, + pages = {2101--2116}, + title = {{A} {R}eal-{T}ime, {GPU}-{B}ased {I}mplementation of {A}perture {D}omain {M}odel + {I}mage {RE}construction}, + volume = {68}, + year = {2021} +} + +@article{byram2014ultrasonic, + author = {Byram, Brett and Jakovljevic, Marko}, + doi = {10.1109/TUFFC.2014.2928}, + journal = {{IEEE} Transactions on Ultrasonics, Ferroelectrics, and Frequency Control}, + number = {3}, + pages = {428--440}, + title = {{U}ltrasonic {M}ultipath and {B}eamforming {C}lutter {R}eduction: {A} {C}hirp + {M}odel {A}pproach}, + volume = {61}, + year = {2014} +} + +@article{dei2019computationally, + author = {Dei, Kazuyuki and Schlunk, Siegfried and Byram, Brett}, + doi = {10.1109/TUFFC.2019.2924824}, + journal = {{IEEE} Transactions on Ultrasonics, Ferroelectrics, and Frequency Control}, + number = {10}, + pages = {1546--1559}, + title = {{C}omputationally {E}fficient {I}mplementation of {A}perture {D}omain {M}odel + {I}mage {R}econstruction}, + volume = {66}, + year = {2019} +} + +@article{khan2020genre, + author = {Khan, Christopher and Byram, Brett}, + doi = {10.21105/joss.02644}, + journal = {Journal of Open Source Software}, + number = {54}, + pages = {2644}, + title = {{GENRE} ({GPU} {E}lastic-{N}et {RE}gression): {A} {CUDA}-{A}ccelerated {P}ackage + for {M}assively {P}arallel {L}inear {R}egression with {E}lastic-{N}et + {R}egularization}, + volume = {5}, + year = {2020} +} + +@article{shlens2014tutorial, + author = {Shlens, Jonathon}, + doi = {10.48550/arXiv.1404.2986}, + journal = {arXiv preprint}, + title = {{A} {T}utorial on {I}ndependent {C}omponent {A}nalysis}, + year = {2014} +} + +@inproceedings{cardoso1989source, + author = {Cardoso, Jean-Fran{\c{c}}ois}, + booktitle = {International Conference on Acoustics, Speech, and Signal Processing (ICASSP)}, + doi = {10.1109/ICASSP.1989.266878}, + pages = {2109--2112}, + title = {{S}ource {S}eparation {U}sing {H}igher {O}rder {M}oments}, + year = {1989} +} + @article{bottenus2018recovery, author = {Bottenus, Nick}, doi = {10.1109/TUFFC.2017.2773495}, diff --git a/tests/test_admire.py b/tests/test_admire.py new file mode 100644 index 000000000..ff1252b77 --- /dev/null +++ b/tests/test_admire.py @@ -0,0 +1,463 @@ +"""Tests for Aperture Domain Model Image REconstruction (ADMIRE).""" + +import dataclasses + +import numpy as np +import pytest + +from zea.beamform.admire import ( + ADMIREConfig, + _aperture_growth, + _fix_phase, + _resolve_calibration, + _ica_basis, + apply_admire, + elastic_net_ccd, + generate_admire_models, +) +from zea.internal.cache import cache_disabled + +from . import backend_equality_check + +SOUND_SPEED = 1540.0 +CENTER_FREQUENCY = 5e6 +WAVELENGTH = SOUND_SPEED / CENTER_FREQUENCY +PITCH = 0.3e-3 +BANDWIDTH = 0.6 + + +def _reference_ccd(X, y, alpha, lambda_scaling_factor, tolerance, max_iterations): + """Line-by-line port of the reference ``ccd_double_precision.c``.""" + n_obs, n_predictors = X.shape + lam = lambda_scaling_factor * np.sqrt(np.mean(y**2)) + std = np.sqrt(np.mean((y - y.mean()) ** 2)) + beta = np.zeros(n_predictors) + if std == 0: + return beta + lam = lam / std + residual = y / std + max_change, iteration = np.inf, 0 + while max_change >= tolerance and iteration < max_iterations: + max_change = 0.0 + for j in range(n_predictors): + previous = beta[j] + residual = residual + X[:, j] * previous + p_j = X[:, j] @ residual / n_obs + new = np.sign(p_j) * max(abs(p_j) - lam * alpha, 0.0) + new = new / (1 / n_obs + lam * (1 - alpha)) + beta[j] = new + residual = residual - X[:, j] * new + max_change = max(max_change, (previous - new) ** 2 / n_obs) + iteration += 1 + return beta * std + + +def _reference_fobi(X, fix_phase=lambda v: v): + """Port of the reference ``ica.m`` (FOBI) followed by ``pinv(W)``.""" + d, n = X.shape + X = X - X.mean(axis=1, keepdims=True) + covariance = X @ X.conj().T / (n - 1) + eigenvalues, E = np.linalg.eigh(covariance) + E = fix_phase(E) + whitening = np.diag(1 / np.sqrt(eigenvalues)) @ E.conj().T + X_w = whitening @ X + U, _, _ = np.linalg.svd((np.sum(X_w * X_w, axis=0) * X_w) @ X_w.conj().T) + return np.linalg.pinv(fix_phase(U) @ whitening) + + +@pytest.fixture(scope="module") +def iq_models(): + """Small ADMIRE models for analytic data around 12 mm depth.""" + depths = np.arange(11e-3, 13e-3, WAVELENGTH / 8) + with cache_disabled(): + return depths, generate_admire_models( + depths, + sound_speed=SOUND_SPEED, + center_frequency=CENTER_FREQUENCY, + pitch=PITCH, + n_elements=16, + f_number=2.0, + analytic=True, + ) + + +def _point_channel_data(depths, n_elements, x, z, analytic=True): + """Delayed channel data of a point scatterer under a 0 degree plane wave.""" + elements = (np.arange(n_elements) - (n_elements - 1) / 2) * PITCH + delays = (depths[:, None] + np.hypot(elements, depths[:, None])) / SOUND_SPEED + arrival = (z + np.hypot(elements - x, z)) / SOUND_SPEED + dt = delays - arrival + sigma_t = np.sqrt(8 * np.log(2)) / (2 * np.pi * BANDWIDTH * CENTER_FREQUENCY) + signal = np.exp(-(dt**2) / (2 * sigma_t**2)) * np.exp(2j * np.pi * CENTER_FREQUENCY * dt) + if analytic: + channels = np.stack([signal.real, signal.imag], axis=-1) + else: + channels = signal.real[..., None] + return channels[:, None].astype(np.float32) # (n_z, 1 line, n_el, n_ch) + + +@pytest.mark.parametrize("max_iterations", [2, 50]) +@backend_equality_check() +def test_elastic_net_ccd_matches_reference(max_iterations): + """The batched solver matches the reference coordinate descent, also with masked rows.""" + rng = np.random.default_rng(0) + n_obs, n_predictors, n_rhs, n_pad = 12, 20, 3, 4 + X = rng.standard_normal((n_obs, n_predictors)) + X /= np.linalg.norm(X, axis=0) + y = rng.standard_normal((n_rhs, n_obs)) + settings = dict(alpha=0.9, lambda_scaling_factor=0.0189, tolerance=1e-10) + expected = np.stack( + [_reference_ccd(X, y_i, max_iterations=max_iterations, **settings) for y_i in y] + ) + + # Zero-padded rows that are masked out must not change the fit + X_padded = np.concatenate([X, np.zeros((n_pad, n_predictors))]).astype(np.float32) + y_padded = np.concatenate([y, np.zeros((n_rhs, n_pad))], axis=1).astype(np.float32) + mask = np.r_[np.ones(n_obs), np.zeros(n_pad)].astype(np.float32) + beta = elastic_net_ccd(X_padded, y_padded, mask, max_iterations=max_iterations, **settings) + + np.testing.assert_allclose(np.asarray(beta), expected, atol=1e-4) + return beta + + +def test_elastic_net_ccd_stops_per_problem(): + """At a loose tolerance, each problem in a batch stops on its own, as in the reference.""" + rng = np.random.default_rng(0) + n_obs, n_predictors = 16, 32 + X = rng.standard_normal((n_obs, n_predictors)) + X /= np.linalg.norm(X, axis=0) + # Pure noise converges after one sweep; the structured right-hand sides need two + y = np.stack( + [ + X[:, :4] @ (scale * rng.standard_normal(4)) + rng.standard_normal(n_obs) + for scale in (0, 0.5, 2, 8, 20, 40) + ] + ) + settings = dict(alpha=0.9, lambda_scaling_factor=0.0189, tolerance=0.1, max_iterations=1000) + expected = np.stack([_reference_ccd(X, y_i, **settings) for y_i in y]) + + beta = elastic_net_ccd(X.astype(np.float32), y.astype(np.float32), np.ones(n_obs), **settings) + np.testing.assert_allclose(np.asarray(beta), expected, atol=1e-4) + + +def test_elastic_net_ccd_zero_data(): + """All-zero observations give all-zero coefficients instead of NaNs.""" + X = np.eye(4, dtype=np.float32) + beta = elastic_net_ccd( + X, np.zeros((2, 4), np.float32), np.ones(4, np.float32), 0.9, 0.0189, 0.1, 10 + ) + np.testing.assert_array_equal(np.asarray(beta), 0) + + +@pytest.mark.parametrize("keep_chunks", [True, False]) +def test_ica_basis_matches_reference(keep_chunks): + """The chunked FOBI basis matches the reference ICA with the same eigenvector phases.""" + rng = np.random.default_rng(1) + model = rng.standard_normal((6, 500)) + 1j * rng.standard_normal((6, 500)) + model = model * rng.uniform(0.5, 2, (6, 1)) + + basis = _ica_basis(lambda: iter(np.split(model, 5, axis=1)), 6, keep_chunks=keep_chunks) + np.testing.assert_allclose(basis, _reference_fobi(model, _fix_phase), atol=1e-8) + + +def test_ica_basis_independent_of_eigenvector_phase(monkeypatch): + """The basis does not depend on the phases that eigh and svd happen to return.""" + import zea.beamform.admire as admire + + rng = np.random.default_rng(2) + model = rng.standard_normal((6, 400)) + 1j * rng.standard_normal((6, 400)) + model = model * rng.uniform(0.5, 2, (6, 1)) + basis = _ica_basis(lambda: iter([model]), 6) + + eigh, svd = np.linalg.eigh, np.linalg.svd + phases = np.exp(1j * rng.uniform(0, 2 * np.pi, (2, 6))) + + def rotated_eigh(a): + values, vectors = eigh(a) + return values, vectors * phases[0] + + def rotated_svd(a): + u, s, vh = svd(a) + return u * phases[1], s, phases[1].conj()[:, None] * vh + + monkeypatch.setattr(admire.np.linalg, "eigh", rotated_eigh) + monkeypatch.setattr(admire.np.linalg, "svd", rotated_svd) + np.testing.assert_allclose(_ica_basis(lambda: iter([model]), 6), basis, atol=1e-8) + + +def test_generate_admire_models(iq_models): + """Generated models are consistent: band, unit-norm columns, masks and windows.""" + depths, models = iq_models + n_windows, n_freqs, n_el, n_predictors = models.models.shape + + assert n_el == 16 + assert n_windows == models.window_starts.size == models.aperture_mask.shape[0] + assert models.fft_length == 2 * models.window_length + np.testing.assert_array_equal(np.diff(models.window_starts), models.window_length) + assert models.window_starts[-1] + models.window_length <= depths.size + + half_band = 0.5 * 1.2 * BANDWIDTH * CENTER_FREQUENCY + assert np.all(np.abs(models.frequencies - CENTER_FREQUENCY) <= half_band) + assert n_freqs >= 1 + + # At f-number 2, 12 mm deep needs 20 elements: the full 16-element sub-aperture + assert np.all(models.aperture_mask) + + # Active columns have unit norm and live on the active elements only + norms = np.linalg.norm(models.models, axis=2) + assert np.all((np.abs(norms - 1) < 1e-5) | (norms == 0)) + inactive = ~models.aperture_mask[:, None, :, None] + assert np.all(np.where(inactive, models.models, 0) == 0) + + # With ICA, ROI and outer models each have at most one predictor per active element + n_active = models.aperture_mask.sum(axis=1)[:, None] + assert np.all(models.roi_mask.sum(axis=-1) <= n_active) + assert np.all(models.roi_mask.sum(axis=-1) > 0) + assert np.all((norms > 0).sum(axis=-1) <= 2 * n_active) + + +def test_generate_admire_models_parallel(iq_models): + """Generating windows in worker processes gives the same models.""" + depths, models = iq_models + with cache_disabled(): + parallel = generate_admire_models( + depths, + sound_speed=SOUND_SPEED, + center_frequency=CENTER_FREQUENCY, + pitch=PITCH, + n_elements=16, + f_number=2.0, + analytic=True, + n_workers=2, + ) + np.testing.assert_array_equal(parallel.models, models.models) + np.testing.assert_array_equal(parallel.roi_mask, models.roi_mask) + + +@pytest.mark.parametrize( + "center_frequency, expected", + [ + (7.8125e6, (1.295, 1.0, 0.855)), # table key 7813000 + (25e6 / 12, (1.32, 1.0, 0.845)), # table key 2083300 + (5e6, (1.29, 0.995, 0.85)), + (4.2e6, (1.135, 0.995, 0.925)), # not in the table: default values + ], +) +def test_reference_wavenumber_calibration(center_frequency, expected): + """The reference calibration is looked up by the nearest (rounded) table frequency.""" + np.testing.assert_allclose(_resolve_calibration("reference", center_frequency, 3), expected) + + +@pytest.mark.parametrize( + "z, n_el, expected", + [ + (5.7e-3, 32, np.arange(11, 21)), # ceil(5.7 mm / 0.3 mm / 2) = 10 elements + (1e-3, 32, np.arange(12, 20)), # at least min_num_elements + (30e-3, 32, np.arange(32)), # at most all elements + ], +) +def test_aperture_growth(z, n_el, expected): + """Aperture growth keeps the central depth / f-number elements, like the reference.""" + mask = _aperture_growth(z, n_el, PITCH, f_number=2.0, min_num_elements=8) + np.testing.assert_array_equal(np.flatnonzero(mask), expected) + assert np.all(_aperture_growth(z, n_el, PITCH, f_number=0, min_num_elements=8)) + + +def test_generate_admire_models_rejects_coarse_grid(): + """A grid too coarse for the signal band is rejected.""" + depths = np.arange(10e-3, 12e-3, WAVELENGTH) + with cache_disabled(), pytest.raises(ValueError, match="axial"): + generate_admire_models(depths, SOUND_SPEED, CENTER_FREQUENCY, PITCH, 8, analytic=True) + with cache_disabled(), pytest.raises(ValueError, match="Nyquist"): + generate_admire_models(depths, SOUND_SPEED, CENTER_FREQUENCY, PITCH, 8, analytic=False) + + +def test_apply_admire_suppresses_off_axis_clutter(iq_models): + """An on-axis scatterer is kept, an off-axis one is suppressed much more.""" + depths, models = iq_models + + # A converged fit; the reference tolerance stops after very few sweeps + models = dataclasses.replace(models, config=dataclasses.replace(models.config, tolerance=1e-3)) + + def retained_energy(x, z): + data = _point_channel_data(depths, models.n_elements, x, z) + out = np.asarray(apply_admire(data, models)) + das = data[..., 0].sum(-1) + 1j * data[..., 1].sum(-1) + admire = out[..., 0].sum(-1) + 1j * out[..., 1].sum(-1) + return np.sum(np.abs(admire) ** 2) / np.sum(np.abs(das) ** 2) + + on_axis = retained_energy(0.0, 12e-3) + off_axis = retained_energy(3e-3, 12e-3) + assert on_axis > 0.3 + assert off_axis < 0.5 * on_axis + + +def test_apply_admire_leaves_uncovered_samples(iq_models): + """Samples after the last STFT window are passed through unchanged.""" + depths, models = iq_models + rng = np.random.default_rng(2) + n_extra = 3 + data = rng.standard_normal((depths.size + n_extra, 2, models.n_elements, 2)).astype(np.float32) + out = np.asarray(apply_admire(data, models, window_batch_size=2)) + + assert out.shape == data.shape + end = models.window_starts[-1] + models.window_length + np.testing.assert_array_equal(out[end:], data[end:]) + assert not np.allclose(out[:end], data[:end]) + + +def test_apply_admire_window_batching(iq_models): + """Fitting the windows in batches gives the same result as all at once.""" + depths, models = iq_models + data = _point_channel_data(depths, models.n_elements, 1e-3, 12e-3) + np.testing.assert_allclose( + np.asarray(apply_admire(data, models, window_batch_size=3)), + np.asarray(apply_admire(data, models)), + atol=1e-5, + ) + + +def test_apply_admire_rf(): + """Real RF data is supported and stays real (one channel).""" + depths = np.arange(11.5e-3, 12.5e-3, WAVELENGTH / 8) + with cache_disabled(): + models = generate_admire_models( + depths, SOUND_SPEED, CENTER_FREQUENCY, PITCH, 12, f_number=2.0, analytic=False + ) + data = _point_channel_data(depths, 12, 0.0, 12e-3, analytic=False) + out = np.asarray(apply_admire(data, models)) + assert out.shape == data.shape + assert np.sum(out**2) > 0.3 * np.sum(data**2) + + +def test_apply_admire_checks_data(iq_models): + """Mismatched element count or data type raise a clear error.""" + depths, models = iq_models + with pytest.raises(ValueError, match="elements"): + apply_admire(np.zeros((depths.size, 1, 8, 2), np.float32), models) + with pytest.raises(ValueError, match="IQ"): + apply_admire(np.zeros((depths.size, 1, 16, 1), np.float32), models) + + +def _admire_inputs(n_el=16, n_x=5, n_tx=2): + """Grid, probe and parameters for an ADMIRE operation, with random aligned data.""" + probe_x = (np.arange(n_el) - (n_el - 1) / 2) * PITCH + probe_geometry = np.stack([probe_x, 0 * probe_x, 0 * probe_x], axis=1).astype(np.float32) + z = np.arange(11e-3, 12e-3, WAVELENGTH / 8) + x = probe_x[n_el // 2 - n_x // 2 : n_el // 2 - n_x // 2 + n_x] + X, Z = np.meshgrid(x, z) + grid = np.stack([X, 0 * X, Z], axis=-1).astype(np.float32) + rng = np.random.default_rng(3) + data = rng.standard_normal((1, n_tx, grid.shape[0] * n_x, n_el, 2)).astype(np.float32) + params = dict( + grid=grid, + probe_geometry=probe_geometry, + sound_speed=SOUND_SPEED, + center_frequency=0.0, # as left by zea.ops.Demodulate + demodulation_frequency=CENTER_FREQUENCY, + f_number=2.0, + ) + return data, params + + +def test_admire_operation(): + """The operation matches apply_admire on the compounded sub-aperture of each column.""" + from zea import ops + + data, params = _admire_inputs() + n_z, n_x = params["grid"].shape[:2] + with cache_disabled(): + operation = ops.ADMIRE(n_elements=9, with_batch_dim=True) + output = np.asarray(operation(data=data, **params)["data"]) + eager_operation = ops.ADMIRE(n_elements=9, with_batch_dim=True, jit_compile=False) + eager_output = np.asarray(eager_operation(data=data, **params)["data"]) + assert output.shape == (1, n_z * n_x, 2) + + # Columns sit on elements 6..10, so their 9-element sub-apertures are elements 2..14 + models = next(iter(operation._models_cache.values())) + compounded = data[0].sum(axis=0).reshape(n_z, n_x, -1, 2) + lines = np.stack([compounded[:, i, 2 + i : 11 + i] for i in range(n_x)], axis=1) + expected = np.asarray(apply_admire(lines, models)).sum(axis=2).reshape(n_z * n_x, 2) + np.testing.assert_allclose(eager_output[0], expected, rtol=1e-4, atol=1e-4) + # Compiled reductions may round differently (e.g. on GPU), which can shift the + # sweep at which the coarse CCD tolerance is met + np.testing.assert_allclose(output[0], expected, rtol=0, atol=1e-2) + + # A second call reuses the models and the compiled fit + operation(data=data, **params) + assert len(operation._models_cache) == 1 and len(operation._fit_cache) == 1 + + +def test_admire_in_beamform_requires_full_grid(): + """Patching splits the depth axis, which ADMIRE cannot handle.""" + from zea import ops + + data, params = _admire_inputs() + operation = ops.ADMIRE(n_elements=9, with_batch_dim=True) + with pytest.raises(ValueError, match="num_patches=1"): + operation(data=data[:, :, :-5], **params) + + +def test_beamform_admire_uses_full_grid(): + """Beamform with ADMIRE defaults to the full grid and rejects patching clearly.""" + from zea import ops + + beamform = ops.Beamform(beamformer="admire") + assert beamform.num_patches == 1 + assert not any(isinstance(op, ops.PatchedGrid) for op in beamform.operations) + assert beamform.get_dict()["params"]["num_patches"] == 1 + + for patching in (dict(num_patches=4), dict(patch_size=2048)): + with pytest.raises(ValueError, match="num_patches=1"): + ops.Beamform(beamformer="admire", **patching) + + +@pytest.mark.parametrize( + "geometry, message", + [ + ("cartesian", None), + ("tilted", "not cartesian"), # columns that are not straight down, as in a sector + ("off_array", "beyond the array"), + ], +) +def test_admire_warns_about_unsupported_geometry(geometry, message): + """Easy-to-spot geometries that ADMIRE does not model give a warning, once.""" + from unittest.mock import patch + + from zea import ops + + data, params = _admire_inputs() + grid = params["grid"].copy() + if geometry == "tilted": + grid[..., 0] += 0.2 * (grid[..., 2] - grid[0, 0, 2]) + elif geometry == "off_array": + grid[..., 0] += 16 * PITCH + params["grid"] = grid + + operation = ops.ADMIRE(n_elements=9, with_batch_dim=True) + with cache_disabled(), patch("zea.ops.pipeline.log.warning") as warning: + operation(data=data, **params) + operation(data=data, **params) + messages = [call.args[0] for call in warning.call_args_list] + if message is None: + assert not messages + else: + assert len(messages) == 1 and message in messages[0] + + +def test_admire_registered_as_beamformer(): + """ADMIRE is available to Beamform and keeps its settings in a config.""" + from zea import ops + + beamform = ops.Beamform(beamformer="admire", num_patches=1, n_elements=9) + admire = beamform.operations[1] + assert isinstance(admire, ops.ADMIRE) and admire.n_elements == 9 + + config = ADMIREConfig(alpha=0.5) + operation = ops.ADMIRE(n_elements=9, config=config) + params = operation.get_dict()["params"] + assert params["n_elements"] == 9 + assert params["config"] == dataclasses.asdict(config) + assert "models" not in params + restored = ops.get_ops("admire")(**params) + assert ADMIREConfig(**restored.config) == config diff --git a/tests/test_operations.py b/tests/test_operations.py index d411df6f1..6300abb40 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -982,10 +982,11 @@ def test_compute_time_to_peak(): # "able", registered by zea.models.able) is not comparable across backends: its # weights are randomly initialized per backend, and under torch its output requires # grad, which this harness cannot convert to numpy. Those live with their model. +# ADMIRE needs a real imaging geometry; it is tested in test_admire.py. _OPS_BEAMFORMERS = [ name for name in beamformer_registry.registered_names() - if beamformer_registry[name].__module__.startswith("zea.ops.") + if beamformer_registry[name].__module__.startswith("zea.ops.") and name != "admire" ] diff --git a/zea/beamform/__init__.py b/zea/beamform/__init__.py index a4df94475..70ab7614c 100644 --- a/zea/beamform/__init__.py +++ b/zea/beamform/__init__.py @@ -6,6 +6,7 @@ Modules ------- +- :mod:`zea.beamform.admire` -- Aperture Domain Model Image REconstruction (ADMIRE) models and fit. - :mod:`zea.beamform.beamformer` -- Main beamforming functions and time-of-flight correction. - :mod:`zea.beamform.delays` -- Delay calculation routines for plane wave and focused transmissions. - :mod:`zea.beamform.lens_correction` -- Lens-corrected delay computation. @@ -17,4 +18,4 @@ """ -from . import beamformer, delays, lens_correction, pfield, phantoms, pixelgrid +from . import admire, beamformer, delays, lens_correction, pfield, phantoms, pixelgrid diff --git a/zea/beamform/admire.py b/zea/beamform/admire.py new file mode 100644 index 000000000..8955ef0a1 --- /dev/null +++ b/zea/beamform/admire.py @@ -0,0 +1,966 @@ +"""Aperture Domain Model Image REconstruction (ADMIRE). + +ADMIRE is a model-based beamformer that suppresses off-axis and multipath +(reverberation) clutter. Time-delayed channel data is split into short axial +windows, each window is taken to the frequency domain, and at every frequency the +aperture-domain signal (one complex value per element) is fitted with a linear +model of the wavefronts that point scatterers would produce across the aperture. +Scatterers inside a small region of interest (ROI) around the window center are +kept; everything explained by scatterers *outside* the ROI is clutter and is +discarded. The model is fitted with elastic-net regularized least squares, solved +with cyclic coordinate descent. + +The implementation is split in two parts: + +- :func:`generate_admire_models` precomputes the models with NumPy. They depend + only on the imaging geometry (axial sampling, sound speed, center frequency, + element pitch and sub-aperture size), so they are computed once and reused. +- :func:`apply_admire` fits the models to channel data and reconstructs the + decluttered channel data with Keras ops, so it runs (and JIT compiles) on every + backend. + +The beamformer is usually used through :class:`zea.ops.ADMIRE`, which handles +the data layout of a zea pipeline. + +.. important:: + This is a ``zea`` implementation of ADMIRE. For the original MATLAB, C and + CUDA code of the Vanderbilt BEAM Lab, see + `VU-BEAM-Lab/ADMIRE `_. + +.. note:: + The model generation, its calibration constants and the model application + are ported from the reference code, Copyright 2020 Christopher Khan, + Kazuyuki Dei, Siegfried Schlunk and Brett Byram. The coordinate descent + solver is ported from ``ccd_double_precision.c``, Copyright 2020 Christopher + Khan. Both are licensed under the + `Apache License 2.0 `_; + see the reference repository's ``NOTICE`` file. The ICA is the FOBI code + (``ica.m``) from Shlens' tutorial, which the reference includes. + + Changes from the reference: the code is translated to NumPy (model + generation) and Keras ops (model fitting); all lines, windows and frequencies + are fitted in one batched solve; IQ data is supported in addition to RF data; + the phase of the ICA eigenvectors is fixed so the models are reproducible + across platforms; and the wavenumber calibration table is matched to the + nearest listed center frequency instead of exactly. + +.. citation:: byram2015model, dei2019computationally, khan2021realtime, khan2020genre, + byram2014ultrasonic, shlens2014tutorial, cardoso1989source + +""" + +import dataclasses +import math +from dataclasses import dataclass, field + +import numpy as np +from keras import ops +from scipy.special import erf + +from zea.func.tensor import vmap +from zea.internal.cache import cache_output + +#: Wavenumber calibration values of the reference implementation, by transducer +#: center frequency (Hz). Entry ``i`` scales the wavenumber of the ``i``-th selected +#: frequency of every STFT window. +REFERENCE_WAVENUMBER_CALIBRATION = { + 1818200: (1.32, 0.995, 0.835, 0.81), + 2083300: (1.32, 1.0, 0.845, 0.825), + 2500000: (1.315, 0.995, 0.835, 0.81), + 3000000: (1.31, 0.995, 0.835, 0.805), + 3125000: (1.305, 0.985, 0.825, 0.79), + 3500000: (1.29, 0.99, 0.835, 0.815), + 3600000: (1.315, 1.01, 0.855, 0.815), + 4000000: (1.29, 0.99, 0.845, 0.82), + 4500000: (1.325, 1.025, 0.87, 0.825), + 5000000: (1.29, 0.995, 0.85, 0.835), + 5208000: (1.145, 0.995, 0.895, 0.865), + 5500000: (1.28, 0.995, 0.855, 0.845), + 6500000: (1.31, 1.04, 0.89, 0.84), + 7500000: (1.275, 0.975, 0.835, 0.82), + 7800000: (1.295, 1.0, 0.855, 0.825), + 7813000: (1.295, 1.0, 0.855, 0.825), +} + +#: Calibration values the reference implementation uses for any other center frequency. +DEFAULT_WAVENUMBER_CALIBRATION = (1.135, 0.995, 0.925, 0.905) + + +@dataclass(frozen=True) +class ADMIREConfig: + """Settings of the ADMIRE model space and model fit. + + The defaults are those of the reference implementation. Most users only touch + the fit settings (``alpha``, ``lambda_scaling_factor``) and the aperture + settings; the model-space scaling factors and calibration offsets are + empirical and best left alone. + + Lateral and axial model-space positions are given as multiples of the lateral + resolution ``res_lat`` (estimated per window and frequency) and of the axial + resolution ``res_axl = 2 * res_lat``. Ranges are ``(min, spacing, max)``. + + Args: + bandwidth (float): Fractional bandwidth of the transmitted pulse. Sets the + STFT window length and the fitted frequency band. + alpha (float): Elastic-net mixing parameter in ``[0, 1]``: ``1`` is the + lasso, ``0`` is ridge regression. + lambda_scaling_factor (float): The regularization strength of each fit is + ``lambda_scaling_factor * rms(y)``. + max_iterations (int): Maximum number of coordinate descent sweeps. + tolerance (float): Coordinate descent stops once the largest change of a + (standardized) coefficient within a sweep, squared and divided by the + number of observations, falls below this value. + ica (bool): Replace each of the ROI and outer models by its independent + components (FOBI), which shrinks each to (at most) as many predictors + as there are elements. This is what makes ADMIRE tractable on a GPU; + ``False`` uses the full models, which can hold hundreds of thousands + of predictors. + min_num_elements (int): Smallest sub-aperture that aperture growth may + select. + pulse_multiplier (float): Scales the STFT window length (one pulse + full-width at half-maximum by default). + frequency_band_scaling (float): The fitted band is + ``f0 +- 0.5 * frequency_band_scaling * bandwidth * f0``. + wavenumber_calibration (str or tuple or None): Per-frequency wavenumber + scaling. ``"reference"`` looks the values up by center frequency like + the reference implementation; a tuple gives them explicitly; ``None`` + disables calibration. + cal_shift (float): Calibration depth offset in meters. + distance_offset_shift (float): Calibration distance offset in meters. + win_tune (float): Scales the half pulse length in the window amplitude model. + ellipsoid_constant_1 (float): Pads the ROI acceptance ellipse, in meters. + ellipsoid_constant_2 (float): Pads the ellipse that the outer model excludes, + in meters. + lateral_limit_offset (float): Extends the outer model laterally beyond the + sub-aperture, in meters. + roi_x (tuple): ROI lateral positions, in units of ``res_lat``. + roi_z (tuple): ROI depth positions around the window center, in units of + ``res_axl``. + roi_distance_offset_spacing (float): Spacing of the ROI distance offsets, in + wavelengths. + roi_distance_offset_limits (tuple): Min and max ROI distance offsets in meters. + outer_x (tuple): Outer-model lateral positions; min and max are multiples of + the lateral limit, the spacing is in units of ``res_lat``. + outer_z (tuple): Outer-model depth positions as + ``(min_z_center, min_res_axl, spacing_res_axl, max_z_center, max_res_axl, + max_z_center_extra)``; the minimum is + ``min_z_center * z_c + min_res_axl * res_axl + outer_z_offset`` and the + maximum ``(max_z_center + max_z_center_extra) * z_c + + max_res_axl * res_axl``. + outer_z_offset (float): Constant added to the minimum outer-model depth, in + meters. + outer_distance_offset_spacing (float): Spacing of the outer distance offsets, + in wavelengths. + outer_distance_offset_limits (tuple): Min and max outer distance offsets in + meters. + predictor_chunk_size (int): Number of predictors generated at once, which + bounds the memory use of model generation. + max_cache_bytes (int): Models up to this size are kept in memory between the + two passes of the ICA; larger ones are generated twice. + """ + + bandwidth: float = 0.6 + alpha: float = 0.9 + lambda_scaling_factor: float = 0.0189 + max_iterations: int = 100000 + tolerance: float = 0.1 + ica: bool = True + min_num_elements: int = 16 + pulse_multiplier: float = 1.0 + frequency_band_scaling: float = 1.2 + wavenumber_calibration: str | tuple | None = "reference" + cal_shift: float = 7.75e-6 + distance_offset_shift: float = 3.85e-5 + win_tune: float = 1.0 + ellipsoid_constant_1: float = 0.0 + ellipsoid_constant_2: float = 0.5e-3 + np.finfo(float).eps + lateral_limit_offset: float = 1e-3 + roi_x: tuple = (-0.5, 0.0179, 0.5) + roi_z: tuple = (-0.5, 0.1430, 0.5) + roi_distance_offset_spacing: float = 0.0485 + roi_distance_offset_limits: tuple = (-0.8e-3, 0.4e-3) + outer_x: tuple = (-1.0, 1.4228, 1.0) + outer_z: tuple = (0.0, 0.0, 0.7114, 1.0, 0.0, 0.05) + outer_z_offset: float = 0.0 + outer_distance_offset_spacing: float = 0.1211 + outer_distance_offset_limits: tuple = (-8e-3, 3.2e-3) + predictor_chunk_size: int = 32768 + max_cache_bytes: int = 2**30 + + def replace(self, **changes) -> "ADMIREConfig": + """Return a copy with some settings changed.""" + return dataclasses.replace(self, **changes) + + +@dataclass +class ADMIREModels: + """Precomputed ADMIRE models, as made by :func:`generate_admire_models`. + + There is one model per STFT window and selected frequency. Models are zero + padded to a common number of rows (elements) and columns (predictors), which + leaves the fit unchanged: padded columns get zero coefficients and padded rows + carry no data. + + Args: + models (np.ndarray): Complex models of shape + ``(n_windows, n_freqs, n_elements, n_predictors)``, with unit-norm columns. + roi_mask (np.ndarray): Boolean mask of shape + ``(n_windows, n_freqs, n_predictors)``, true for the ROI predictors used + to reconstruct the signal. + aperture_mask (np.ndarray): Boolean mask of shape ``(n_windows, n_elements)`` + with the elements that aperture growth keeps in each window. + window_starts (np.ndarray): Axial sample index where each window starts. + window_length (int): Number of axial samples per window. + fft_length (int): Zero-padded DFT length of a window. + frequency_bins (np.ndarray): DFT bin of each selected frequency. + frequencies (np.ndarray): Physical frequency of each selected bin in Hz. + element_positions (np.ndarray): Lateral element positions of the + sub-aperture, relative to its center, in meters. + analytic (bool): Whether the models are made for analytic (IQ) data, where + only the band around the center frequency is present, or for real RF + data, where negative frequencies mirror the positive ones. + config (ADMIREConfig): Settings used to generate the models. + """ + + models: np.ndarray + roi_mask: np.ndarray + aperture_mask: np.ndarray + window_starts: np.ndarray + window_length: int + fft_length: int + frequency_bins: np.ndarray + frequencies: np.ndarray + element_positions: np.ndarray + analytic: bool + config: ADMIREConfig = field(default_factory=ADMIREConfig) + + @property + def n_windows(self) -> int: + """Number of STFT windows.""" + return self.models.shape[0] + + @property + def n_elements(self) -> int: + """Number of elements in the sub-aperture.""" + return self.models.shape[2] + + +def _colon(start, step, stop): + """MATLAB's ``start:step:stop`` range.""" + if step == 0: + return np.array([start]) if start == stop else np.array([]) + n = math.floor((stop - start) / step + 1e-10) + 1 + return start + step * np.arange(max(n, 0)) + + +def _resolve_calibration(calibration, center_frequency, n_freqs): + """Return the wavenumber calibration value of each selected frequency.""" + if calibration is None: + return np.ones(n_freqs) + if isinstance(calibration, str): + if calibration != "reference": + raise ValueError( + f"Unknown wavenumber_calibration {calibration!r}, " + "expected 'reference', a tuple of values, or None." + ) + # The table keys are rounded (e.g. 7813000 for 7.8125 MHz), so match the nearest + nearest = min(REFERENCE_WAVENUMBER_CALIBRATION, key=lambda f: abs(f - center_frequency)) + if abs(nearest - center_frequency) <= 1e3: + calibration = REFERENCE_WAVENUMBER_CALIBRATION[nearest] + else: + calibration = DEFAULT_WAVENUMBER_CALIBRATION + calibration = np.asarray(calibration, dtype=float) + if n_freqs > calibration.size: + raise ValueError( + f"{n_freqs} frequencies are selected per STFT window, but only " + f"{calibration.size} wavenumber calibration values are available. Pass more " + "values with `wavenumber_calibration`, or set it to None." + ) + return calibration[:n_freqs] + + +def _select_frequencies( + fft_length, sampling_frequency, center_frequency, config: ADMIREConfig, analytic +): + """Select the DFT bins in the fitted band and their physical frequencies.""" + half_band = 0.5 * config.bandwidth * center_frequency * config.frequency_band_scaling + f_min, f_max = center_frequency - half_band, center_frequency + half_band + bins = np.arange(fft_length) + bin_frequencies = bins / fft_length * sampling_frequency + + if analytic: + if 2 * half_band >= sampling_frequency: + raise ValueError( + f"The fitted band ({2 * half_band / 1e6:.2f} MHz) does not fit in the axial " + f"sampling rate ({sampling_frequency / 1e6:.2f} MHz) of the grid. Use a " + "finer axial grid spacing." + ) + # IQ data is aliased: map every bin to its alias closest to the center frequency. + alias = np.round((center_frequency - bin_frequencies) / sampling_frequency) + frequencies = bin_frequencies + alias * sampling_frequency + else: + if f_max >= sampling_frequency / 2: + raise ValueError( + f"The fitted band reaches {f_max / 1e6:.2f} MHz, above the Nyquist frequency " + f"({sampling_frequency / 2e6:.2f} MHz) of the axial grid. Use a finer axial " + "grid spacing, or IQ data." + ) + # Only the positive half; the negative frequencies follow by conjugate symmetry. + frequencies = bin_frequencies[: fft_length // 2] + bins = bins[: fft_length // 2] + + selected = (frequencies >= f_min) & (frequencies <= f_max) + order = np.argsort(frequencies[selected]) + return bins[selected][order], frequencies[selected][order] + + +def _predictor_signals(x, z, offset, element_positions, k, z_center, window_depths, params): + """Aperture-domain signals of point-scatterer predictors. + + Port of ``generate_modeled_signal_for_predictor.m``. A predictor is a scatterer + at lateral position ``x`` and depth ``z`` whose transmit path is lengthened by + ``offset`` (which models multipath). Its signal across the aperture combines the + element directivity, the part of the pulse that falls within the STFT window, and + the phase of its perceived depth on each element after receive focusing. + + Args: + x, z, offset (np.ndarray): Predictor parameters of shape ``(P,)``. + element_positions (np.ndarray): Element positions of shape ``(n_el,)``. + k (float): Calibrated wavenumber. + z_center (float): Center depth of the STFT window. + window_depths (tuple): First and last depth of the STFT window. + params (dict): Pitch, sound speed, pulse width and calibration constants. + + Returns: + np.ndarray: Complex signals of shape ``(n_el, P)``. + """ + c = params["sound_speed"] + config: ADMIREConfig = params["config"] + e = element_positions[:, None] + offset = offset + config.distance_offset_shift + + # Element directivity (sinc of the element width times the obliquity factor) + theta = np.arctan2(e - x, z) + wavelength = 2 * np.pi / k + argument = np.pi * params["pitch"] * np.sin(theta) / wavelength + eps = np.finfo(float).eps + directivity = (np.sin(argument) + eps) / (argument + eps) * np.cos(theta) + + # Round-trip distance, and the depth at which the echo appears on each element + # after dynamic receive focusing (the distance offset sets the transmit path). + tau_n0 = (offset + 2 * z_center - (z - config.cal_shift)) / c + d0 = np.sqrt((e - x) ** 2 + (z - config.cal_shift) ** 2) + c * tau_n0 + perceived_depth = (d0**2 - e**2) / (2 * d0) + z_distance = d0 - (np.sqrt(e**2 + perceived_depth**2) - perceived_depth) + + # Fraction of a Gaussian pulse centered at z_distance that lies inside the window + min_travel = 2 * window_depths[0] + max_travel = 2 * window_depths[1] + half_pulse = params["half_pulse_width_samples"] / params["sampling_frequency"] * c / 2 + pulse_min = np.maximum(z_distance - half_pulse * config.win_tune, min_travel) + pulse_max = np.minimum(z_distance + half_pulse * config.win_tune, max_travel) + st = params["pulse_sigma_distance"] + gaussian = ( + 0.5 + * np.sqrt(np.pi) + * st + * (erf((z_distance - pulse_min) / st) - erf((z_distance - pulse_max) / st)) + ) + gaussian[(pulse_min > max_travel) | (pulse_max < min_travel)] = 0 + with np.errstate(divide="ignore", invalid="ignore"): + window_amplitude = np.sqrt(gaussian / gaussian.max(axis=0, keepdims=True)) + + # amplitude * exp(1j * phase), without the (slow) complex exponential + amplitude = window_amplitude * directivity + phase = z_distance * k + signals = np.empty(phase.shape, dtype=complex) + signals.real = amplitude * np.cos(phase) + signals.imag = amplitude * np.sin(phase) + return signals + + +def _predictor_grid(x_range, z_range, offset_range): + """All predictor parameters of a model space, in the reference ordering.""" + x = _colon(*x_range) + z = _colon(*z_range) + offsets = _colon(*offset_range) + # MATLAB's meshgrid + (:) flattens column-major + X, Z, offsets = np.meshgrid(x, z, offsets) + return X.ravel(order="F"), Z.ravel(order="F"), offsets.ravel(order="F") + + +def _model_chunks(x, z, offset, chunk_size, signal_fn): + """Yield the valid (nonzero, finite) predictor signals chunk by chunk.""" + for start in range(0, x.size, chunk_size): + sl = slice(start, start + chunk_size) + signals = signal_fn(x[sl], z[sl], offset[sl]) + valid = np.all(np.isfinite(signals), axis=0) & np.any(signals != 0, axis=0) + yield signals[:, valid] + + +def _ica_basis(chunks_fn, n_el, keep_chunks=True): + """Basis that replaces a model, from the reference's ICA (FOBI). + + Port of ``ica.m`` (Shlens, "A tutorial on independent component analysis", + 2014) followed by ``pinv(W)``, as in the reference implementation. The model + columns are the samples, so it is accumulated chunk by chunk without ever + holding the full model in memory. + + ``ica.m`` was written for real data. On complex models its result depends on the + phase of the eigenvectors that the linear algebra library returns, which is + arbitrary, so the reference gives different models on different platforms. Here + the phase of every eigenvector is fixed by a deterministic rule (see + :func:`_fix_phase`), so the models are reproducible; the reference would + produce the same models for that choice of phases. (Textbook complex FOBI + removes the ambiguity too, but its basis no longer separates on-axis signal + from clutter, so the reference formulation is kept.) + + Args: + chunks_fn (callable): Returns an iterator over chunks of model columns. + n_el (int): Number of rows (elements) of the model. + keep_chunks (bool): Keep the chunks in memory for the second pass instead + of generating them again. + + Returns: + np.ndarray: The ``(n_el, n_el)`` basis that replaces the model. + """ + # First pass: mean and covariance of the columns + n = 0 + total = np.zeros(n_el, dtype=complex) + outer = np.zeros((n_el, n_el), dtype=complex) + kept = [] + for chunk in chunks_fn(): + if keep_chunks: + kept.append(chunk) + n += chunk.shape[1] + total += chunk.sum(axis=1) + outer += chunk @ chunk.conj().T + if n < 2: + raise ValueError("An ADMIRE model has fewer than two valid predictors.") + mean = total / n + covariance = (outer - n * np.outer(mean, mean.conj())) / (n - 1) + + # Whitening, discarding numerically zero directions like MATLAB's pinv + eigenvalues, eigenvectors = np.linalg.eigh(covariance) + eigenvectors = _fix_phase(eigenvectors) + tol = n_el * np.finfo(float).eps * max(np.abs(eigenvalues).max(), np.finfo(float).tiny) + inv_sqrt = np.where(eigenvalues > tol, 1 / np.sqrt(np.maximum(eigenvalues, tol)), 0.0) + whitening = inv_sqrt[:, None] * eigenvectors.conj().T + + # Second pass: the fourth-order moment matrix + moment = np.zeros((n_el, n_el), dtype=complex) + for chunk in kept if keep_chunks else chunks_fn(): + whitened = whitening @ (chunk - mean[:, None]) + moment += (np.sum(whitened * whitened, axis=0) * whitened) @ whitened.conj().T + rotation, _, _ = np.linalg.svd(moment) + rotation = _fix_phase(rotation) + + unmixing = rotation @ whitening + return np.linalg.pinv(unmixing) + + +def _fix_phase(vectors): + """Rotate every column so that a fixed, asymmetric weighting of it is real and positive. + + Eigen- and singular vectors of complex matrices are only defined up to a + unit-modulus factor, which libraries choose differently. + """ + weights = np.linspace(1.0, 2.0, vectors.shape[0]) + return vectors * np.exp(-1j * np.angle(weights @ vectors)) + + +def _lateral_resolution(center_signal, n_el, pitch, wavelength, z_center): + """Lateral resolution estimate from the aperture signal of an on-axis scatterer.""" + n_fft = 2 ** (2 * math.floor(math.log2(n_el))) + spectrum = np.abs(np.fft.fft(np.abs(center_signal), n_fft)) + spectrum = spectrum / spectrum.max() + fwhm_width = np.sum(spectrum > 0.5) + return fwhm_width * (1 / pitch) / n_fft * wavelength * z_center + + +def _aperture_growth(z_center, n_el, pitch, f_number, min_num_elements): + """Elements kept by aperture growth (a fixed f-number) at a given depth.""" + if not f_number: + return np.ones(n_el, dtype=bool) + n_active = math.ceil(z_center / pitch / f_number) + n_active = 2 * math.ceil(n_active / 2) + n_active = min(max(n_active, min_num_elements), n_el) + offsets = np.arange(1 - n_active / 2, n_active / 2 + 1) + indices = np.ceil(n_el / 2 + offsets).astype(int) - 1 + mask = np.zeros(n_el, dtype=bool) + mask[indices] = True + return mask + + +def _models_for_window(z_center, window_depths, element_positions, ks, params): + """Generate the ROI and outer models of every selected frequency of a window.""" + config: ADMIREConfig = params["config"] + pitch = params["pitch"] + n_el = element_positions.size + results = [] + for k in ks: + wavelength = 2 * np.pi / k + + def signal_fn(x, z, offset, k=k): + return _predictor_signals( + x, z, offset, element_positions, k, z_center, window_depths, params + ) + + # Model-space sampling follows from the resolution at the window center + center = signal_fn(np.zeros(1), np.full(1, z_center), np.zeros(1))[:, 0] + res_lat = _lateral_resolution(center, n_el, pitch, wavelength, z_center) + res_axl = 2 * res_lat + + roi_x = tuple(s * res_lat for s in config.roi_x) + roi_z = ( + z_center + config.roi_z[0] * res_axl, + config.roi_z[1] * res_axl, + z_center + config.roi_z[2] * res_axl, + ) + roi_offsets = ( + config.roi_distance_offset_limits[0], + config.roi_distance_offset_spacing * wavelength, + config.roi_distance_offset_limits[1], + ) + x, z, offset = _predictor_grid(roi_x, roi_z, roi_offsets) + inside = (x / (res_lat + config.ellipsoid_constant_1)) ** 2 + ( + (z - z_center) / (res_axl + config.ellipsoid_constant_1) + ) ** 2 <= 1 + roi = (x[inside], z[inside], offset[inside]) + + lateral_limit = pitch * n_el / 2 + config.lateral_limit_offset + sz = config.outer_z + outer_x = ( + config.outer_x[0] * lateral_limit, + config.outer_x[1] * res_lat, + config.outer_x[2] * lateral_limit, + ) + outer_z = ( + sz[0] * z_center + sz[1] * res_axl + config.outer_z_offset, + sz[2] * res_axl, + sz[3] * z_center + sz[4] * res_axl + sz[5] * z_center, + ) + outer_offsets = ( + config.outer_distance_offset_limits[0], + config.outer_distance_offset_spacing * wavelength, + config.outer_distance_offset_limits[1], + ) + x, z, offset = _predictor_grid(outer_x, outer_z, outer_offsets) + outside = (x / (res_lat + config.ellipsoid_constant_2)) ** 2 + ( + (z - z_center) / (res_axl + config.ellipsoid_constant_2) + ) ** 2 > 1 + outer = (x[outside], z[outside], offset[outside]) + + parts = [] + for predictors in (roi, outer): + + def chunks_fn(predictors=predictors, signal_fn=signal_fn): + x, z, offset = predictors + return _model_chunks(x, z, offset, config.predictor_chunk_size, signal_fn) + + if config.ica: + n_bytes = predictors[0].size * n_el * np.dtype(complex).itemsize + keep = n_bytes <= config.max_cache_bytes + parts.append(_ica_basis(chunks_fn, n_el, keep_chunks=keep)) + else: + chunks = list(chunks_fn()) + parts.append(np.concatenate(chunks, axis=1)) + + model = np.concatenate(parts, axis=1) + norms = np.linalg.norm(model, axis=0) + model = np.where(norms > 0, model / np.where(norms > 0, norms, 1), 0) + is_roi = np.arange(model.shape[1]) < parts[0].shape[1] + results.append((model, is_roi)) + return results + + +def _single_threaded_blas(): + """Worker initializer: one BLAS thread per process, so workers don't oversubscribe.""" + try: + from threadpoolctl import threadpool_limits + except ImportError: + return + threadpool_limits(1) + + +@cache_output( + "depths", "sound_speed", "center_frequency", "pitch", "n_elements", "f_number", + "analytic", "start_depth", "end_depth", "config", verbose=True, +) # fmt: skip +def generate_admire_models( + depths, + sound_speed: float, + center_frequency: float, + pitch: float, + n_elements: int, + f_number: float = 2.0, + analytic: bool = True, + start_depth: float | None = None, + end_depth: float | None = None, + config: ADMIREConfig | None = None, + verbose: bool = False, + n_workers: int | None = None, +) -> ADMIREModels: + """Precompute the ADMIRE models for an imaging geometry. + + Port of ``ADMIRE_models_generation_main.m`` and the functions it calls. The + channel data that the models are applied to must be time-delayed (focused at + every depth) and sampled at the uniformly spaced ``depths``. + + Generation is done once per geometry, on the CPU, and takes from seconds to + minutes: deep windows have outer models with hundreds of thousands of + predictors. The result is cached on disk (see :mod:`zea.internal.cache`). + + Args: + depths (array-like): Uniformly spaced depths (m) of the axial samples of the + delayed channel data, i.e. the z-axis of the beamforming grid. + sound_speed (float): Speed of sound in m/s. + center_frequency (float): Transmit center frequency in Hz. + pitch (float): Element pitch in meters. + n_elements (int): Number of elements in the sub-aperture that is fitted for + each image line. + f_number (float): F-number of aperture growth; the sub-aperture is limited + to ``depth / f_number`` (but at least ``config.min_num_elements`` + elements). ``0`` disables aperture growth. + analytic (bool): ``True`` for analytic (IQ) data, ``False`` for real RF data. + start_depth (float, optional): Only windows starting at or below this depth + are processed. Defaults to the first depth. + end_depth (float, optional): Only windows ending at or above this depth are + processed. Defaults to the last depth. + config (ADMIREConfig, optional): Model and fit settings. Defaults to the + reference settings. + verbose (bool): Show a progress bar. + n_workers (int, optional): Number of processes that generate windows in + parallel. Defaults to one. Does not affect the result (or its cache key). + + Returns: + ADMIREModels: The models of every STFT window and selected frequency. + """ + config = ADMIREConfig() if config is None else config + depths = np.asarray(depths, dtype=float) + if depths.ndim != 1 or depths.size < 2: + raise ValueError("`depths` must be a 1D array with at least two depths.") + spacing = np.diff(depths) + dz = spacing.mean() + if dz <= 0 or not np.allclose(spacing, dz, rtol=1e-3, atol=0): + raise ValueError("ADMIRE requires uniformly spaced, increasing depths.") + sampling_frequency = sound_speed / (2 * dz) + + # STFT window: one pulse FWHM of a Gaussian pulse with the given bandwidth + sff = (config.bandwidth * center_frequency) ** 2 / (8 * np.log(2)) + pulse_fwhm = np.sqrt(8 * np.log(2)) * np.sqrt(1 / (4 * np.pi**2 * sff)) + window_length = math.ceil(config.pulse_multiplier * pulse_fwhm * sampling_frequency) + fft_length = 2 * window_length + + starts = np.arange(0, depths.size - window_length + 1, window_length) + start_depth = depths[0] if start_depth is None else start_depth + end_depth = depths[-1] if end_depth is None else end_depth + keep = (depths[starts] >= start_depth) & (depths[starts + window_length - 1] <= end_depth) + starts = starts[keep] + if starts.size == 0: + raise ValueError( + f"No STFT window of {window_length} samples fits in the depth range; the grid " + "is too short." + ) + + bins, frequencies = _select_frequencies( + fft_length, sampling_frequency, center_frequency, config, analytic + ) + if bins.size == 0: + raise ValueError("No DFT bin falls inside the fitted frequency band.") + ks = 2 * np.pi * frequencies / sound_speed + ks = ks * _resolve_calibration(config.wavenumber_calibration, center_frequency, ks.size) + + params = { + "sound_speed": sound_speed, + "pitch": pitch, + "sampling_frequency": sampling_frequency, + "half_pulse_width_samples": window_length, + "pulse_sigma_distance": np.sqrt(sound_speed**2 / (4 * np.pi**2 * sff)), + "config": config, + } + + element_positions = (np.arange(n_elements) - (n_elements - 1) / 2) * pitch + + windows, aperture_masks = [], [] + for start in starts: + window_depths = (depths[start], depths[start + window_length - 1]) + z_center = depths[start : start + window_length].mean() + aperture = _aperture_growth(z_center, n_elements, pitch, f_number, config.min_num_elements) + aperture_masks.append(aperture) + windows.append((z_center, window_depths, element_positions[aperture], ks, params)) + + if n_workers is not None and n_workers > 1: + import multiprocessing + from concurrent.futures import ProcessPoolExecutor + + # Spawn: the parent may hold a GPU context or threads that do not survive fork + context = multiprocessing.get_context("spawn") + with ProcessPoolExecutor( + min(n_workers, len(windows)), mp_context=context, initializer=_single_threaded_blas + ) as pool: + per_window = pool.map(_models_for_window, *zip(*windows)) + if verbose: + from tqdm import tqdm + + per_window = tqdm(per_window, total=len(windows), desc="Generating ADMIRE models") + per_window = list(per_window) + else: + iterator = windows + if verbose: + from tqdm import tqdm + + iterator = tqdm(windows, desc="Generating ADMIRE models") + per_window = [_models_for_window(*window) for window in iterator] + + n_predictors = max(model.shape[1] for window in per_window for model, _ in window) + models = np.zeros((starts.size, ks.size, n_elements, n_predictors), dtype=np.complex64) + roi_mask = np.zeros((starts.size, ks.size, n_predictors), dtype=bool) + for w, (window, aperture) in enumerate(zip(per_window, aperture_masks)): + for f, (model, is_roi) in enumerate(window): + models[w, f, aperture, : model.shape[1]] = model + roi_mask[w, f, : model.shape[1]] = is_roi + + return ADMIREModels( + models=models, + roi_mask=roi_mask, + aperture_mask=np.stack(aperture_masks), + window_starts=starts, + window_length=window_length, + fft_length=fft_length, + frequency_bins=bins, + frequencies=frequencies, + element_positions=element_positions, + analytic=analytic, + config=config, + ) + + +def elastic_net_ccd( + X, + y, + row_mask, + alpha: float, + lambda_scaling_factor: float, + tolerance: float, + max_iterations: int, +): + r"""Batched elastic-net regression with cyclic coordinate descent. + + Port of ``ccd_double_precision.c``. For every problem it minimizes + + .. math:: + + \frac{1}{2N} \lVert \tilde{y} - X \beta \rVert_2^2 + + \tilde{\lambda} \left( \alpha \lVert \beta \rVert_1 + + \frac{1 - \alpha}{2} \lVert \beta \rVert_2^2 \right), + + with :math:`\tilde{y} = y / \sigma_y` and :math:`\tilde{\lambda} = \lambda / + \sigma_y`, where :math:`\lambda` is ``lambda_scaling_factor`` times the RMS of + :math:`y`, and returns :math:`\sigma_y \beta`. The columns of ``X`` must have + unit norm (or be zero). Each problem stops after the first sweep in which the + largest change of its coefficients falls below ``tolerance``, as in the + reference; the loop runs until every problem has stopped. + + Args: + X (Tensor): Design matrices of shape ``(..., N, P)``. + y (Tensor): Observations of shape ``(..., B, N)``: ``B`` right-hand sides + share each design matrix. + row_mask (Tensor): Mask of shape ``(..., N)`` with the rows that hold + observations. Masked-out rows of ``X`` and ``y`` must be zero. + alpha (float): Elastic-net mixing parameter. + lambda_scaling_factor (float): Regularization strength relative to the RMS + of ``y``. + tolerance (float): Convergence tolerance. + max_iterations (int): Maximum number of sweeps over all coefficients. + + Returns: + Tensor: Coefficients of shape ``(..., B, P)``. + """ + X = ops.convert_to_tensor(X) + dtype = X.dtype + y = ops.convert_to_tensor(y, dtype=dtype) + row_mask = ops.cast(row_mask, dtype)[..., None, :] # (..., 1, N) + n_obs = ops.maximum(ops.sum(row_mask, axis=-1, keepdims=True), 1.0) # (..., 1, 1) + + # Standardize y by its standard deviation over the observed rows + mean = ops.sum(y, axis=-1, keepdims=True) / n_obs + std = ops.sqrt(ops.sum(row_mask * (y - mean) ** 2, axis=-1, keepdims=True) / n_obs) + lam = lambda_scaling_factor * ops.sqrt(ops.sum(y**2, axis=-1, keepdims=True) / n_obs) + valid = std > 0 + safe_std = ops.where(valid, std, 1.0) + residual = ops.where(valid, y / safe_std, 0.0) # (..., B, N) + lam = lam / safe_std # (..., B, 1) + + threshold = lam * alpha + denominator = 1.0 / n_obs + lam * (1.0 - alpha) + n_predictors = X.shape[-1] + X_columns = ops.moveaxis(X, -1, 0) # (P, ..., N) + beta = ops.zeros(y.shape[:-1] + (n_predictors,), dtype=dtype) + beta = ops.moveaxis(beta, -1, 0) # (P, ..., B) + + def coordinate_step(j, state): + beta, residual, max_change, active = state + x_j = ops.take(X_columns, j, axis=0)[..., None, :] # (..., 1, N) + beta_j = ops.take(beta, j, axis=0)[..., None] # (..., B, 1) + partial = residual + x_j * beta_j + p_j = ops.sum(x_j * partial, axis=-1, keepdims=True) / n_obs + new_beta_j = ops.sign(p_j) * ops.maximum(ops.abs(p_j) - threshold, 0.0) / denominator + # Problems that have converged keep their coefficients (and residual) unchanged + new_beta_j = ops.where(active, new_beta_j, beta_j) + residual = ops.where(active, partial - x_j * new_beta_j, residual) + change = (beta_j - new_beta_j) ** 2 / n_obs # (..., B, 1) + beta = ops.slice_update(beta, (j,) + (0,) * (len(beta.shape) - 1), new_beta_j[None, ..., 0]) + return beta, residual, ops.maximum(max_change, change), active + + def sweep(state): + beta, residual, active, iteration = state + beta, residual, max_change, _ = ops.fori_loop( + 0, n_predictors, coordinate_step, (beta, residual, ops.zeros_like(lam), active) + ) + # Like the reference, every problem stops on its own after its first sweep + # whose largest change falls below the tolerance + active = ops.logical_and(active, max_change >= tolerance) + return beta, residual, active, iteration + 1 + + def not_converged(beta, residual, active, iteration): + return ops.logical_and(ops.any(active), iteration < max_iterations) + + state = (beta, residual, valid, ops.zeros((), dtype="int32")) + state = ops.while_loop(not_converged, lambda *s: sweep(s), state) + beta = ops.moveaxis(state[0], 0, -1) + return beta * ops.where(valid, std, 0.0) + + +def _realify(models): + """``[[Re, -Im], [Im, Re]]`` real form of complex models ``(..., n_el, P)``.""" + re, im = np.real(models), np.imag(models) + top = np.concatenate([re, -im], axis=-1) + bottom = np.concatenate([im, re], axis=-1) + return np.concatenate([top, bottom], axis=-2).astype(np.float32) + + +def _dft_matrices(models: ADMIREModels): + """Forward and inverse DFT matrices restricted to the selected bins. + + Returns: + tuple: ``cos`` and ``sin`` of shape ``(n_freqs, window_length)`` and the + inverse weights of shape ``(n_freqs,)``. + """ + n = np.arange(models.window_length) + phase = 2 * np.pi * np.outer(models.frequency_bins, n) / models.fft_length + if models.analytic: + weights = np.ones(models.frequency_bins.size) + else: + # A positive bin also stands in for its conjugate negative twin + weights = np.where(models.frequency_bins == 0, 1.0, 2.0) + weights = weights / models.fft_length + return np.cos(phase).astype(np.float32), np.sin(phase).astype(np.float32), weights + + +def apply_admire(channel_data, models: ADMIREModels, window_batch_size: int | None = None): + """Remove clutter from time-delayed channel data with ADMIRE. + + Port of ``apply_ADMIRE_models_CPU.m``: each STFT window is taken to the + frequency domain, the models are fitted at every selected frequency, the signal + is rebuilt from the ROI predictors only and taken back to the time domain. + Frequencies outside the fitted band are removed. Samples that are not covered by + a window (at the end of the axial range) are passed through unchanged. + + Args: + channel_data (Tensor): Delayed channel data of shape + ``(n_z, n_lines, n_elements, n_ch)``: for every image line, the + sub-aperture centered on that line. ``n_ch`` is ``1`` for RF data and + ``2`` for IQ data, matching ``models.analytic``. + models (ADMIREModels): Models from :func:`generate_admire_models`. + window_batch_size (int, optional): Number of STFT windows fitted at once, + which bounds memory use. Defaults to all windows at once. + + Returns: + Tensor: Decluttered channel data with the shape of ``channel_data``. + """ + n_z, n_lines, n_el, n_ch = channel_data.shape + if n_el != models.n_elements: + raise ValueError( + f"The channel data has {n_el} elements per line, but the ADMIRE models were " + f"generated for {models.n_elements}." + ) + if models.analytic != (n_ch == 2): + kind = "IQ" if models.analytic else "RF" + raise ValueError( + f"The ADMIRE models were generated for {kind} data, but the channel data has " + f"n_ch={n_ch}." + ) + first = int(models.window_starts[0]) + length = models.window_length + covered = models.n_windows * length + if not np.array_equal(models.window_starts, first + length * np.arange(models.n_windows)): + raise ValueError("ADMIRE expects contiguous, non-overlapping STFT windows.") + if first + covered > n_z: + raise ValueError( + f"The ADMIRE models cover {first + covered} axial samples, but the channel " + f"data has only {n_z}." + ) + config = models.config + + data = ops.cast(channel_data, "float32") + real = data[..., 0] + imag = data[..., 1] if n_ch == 2 else ops.zeros_like(real) + + cos, sin, inverse_weights = _dft_matrices(models) + inverse_cos = (cos * inverse_weights[:, None]).astype(np.float32) + inverse_sin = (sin * inverse_weights[:, None]).astype(np.float32) + X = _realify(models.models) # (W, F, 2 n_el, 2 P) + roi = np.concatenate([models.roi_mask, models.roi_mask], axis=-1).astype(np.float32) + rows = np.concatenate([models.aperture_mask, models.aperture_mask], axis=-1) + rows = np.broadcast_to(rows[:, None], X.shape[:-1]).astype(np.float32) + + def to_windows(x): + x = x[first : first + covered] + return ops.reshape(x, (models.n_windows, length) + tuple(x.shape[1:])) + + def process(real_w, imag_w, X_w, roi_w, rows_w): + def forward(matrix, x): + return ops.einsum("fn,wnle->wfle", matrix, x) + + def inverse(matrix, x): + return ops.einsum("fn,wfle->wnle", matrix, x) + + # Forward DFT at the selected bins: (W, F, n_lines, n_el) + spec_re = forward(cos, real_w) + forward(sin, imag_w) + spec_im = forward(cos, imag_w) - forward(sin, real_w) + y = ops.concatenate([spec_re, spec_im], axis=-1) * rows_w[:, :, None, :] + + beta = elastic_net_ccd( + X_w, + y, + rows_w, + alpha=config.alpha, + lambda_scaling_factor=config.lambda_scaling_factor, + tolerance=config.tolerance, + max_iterations=config.max_iterations, + ) + fitted = ops.einsum("wfnp,wflp->wfln", X_w, beta * roi_w[:, :, None, :]) + fit_re, fit_im = fitted[..., :n_el], fitted[..., n_el:] + + # Inverse DFT from the selected bins back to the window samples + out_re = inverse(inverse_cos, fit_re) - inverse(inverse_sin, fit_im) + out_im = inverse(inverse_sin, fit_re) + inverse(inverse_cos, fit_im) + return out_re, out_im + + fn_args = (to_windows(real), to_windows(imag), X, roi, rows) + if window_batch_size is None or window_batch_size >= models.n_windows: + out_re, out_im = process(*fn_args) + else: + out_re, out_im = vmap(process, batch_size=window_batch_size, fn_supports_batch=True)( + *fn_args + ) + + def stitch(original, windows): + windows = ops.reshape(windows, (covered,) + tuple(original.shape[1:])) + return ops.concatenate([original[:first], windows, original[first + covered :]], axis=0) + + real = stitch(real, out_re) + if n_ch == 1: + return real[..., None] + return ops.stack([real, stitch(imag, out_im)], axis=-1) diff --git a/zea/ops/__init__.py b/zea/ops/__init__.py index 64449db85..49b41a4d8 100644 --- a/zea/ops/__init__.py +++ b/zea/ops/__init__.py @@ -231,6 +231,7 @@ from .base import Identity, Lambda, Mean, Operation, get_ops from .imports import import_ops_module, import_ops_modules from .pipeline import ( + ADMIRE, Beamform, CoherenceFactor, DelayAndSum, @@ -289,6 +290,7 @@ "CoherenceFactor", "GeneralizedCoherenceFactor", "MinimumVariance", + "ADMIRE", "Beamform", "Map", "PatchedGrid", diff --git a/zea/ops/pipeline.py b/zea/ops/pipeline.py index f45a321e6..bd3c77205 100644 --- a/zea/ops/pipeline.py +++ b/zea/ops/pipeline.py @@ -1,3 +1,4 @@ +import dataclasses import difflib import inspect import json @@ -314,6 +315,7 @@ def from_default( - "coherence_factor" - "generalized_coherence_factor" - "minimum_variance" + - "admire" (always beamforms the full grid at once: ``num_patches=1``) - "able" (available once :mod:`zea.models.able` is imported) Defaults to "delay_and_sum". num_patches (int, optional): Number of patches for the PatchedGrid operation. @@ -1408,6 +1410,7 @@ def __init__( - "coherence_factor" - "generalized_coherence_factor" - "minimum_variance" + - "admire" (always beamforms the full grid at once: ``num_patches=1``) Defaults to "delay_and_sum". num_patches (int, optional): Number of patches to split the grid into for patch-wise beamforming. If 1, no patching is performed. Prefer @@ -1442,6 +1445,21 @@ def __init__( "num_patches and patch_size are mutually exclusive. Please specify only one." ) + if beamformer == "admire": + # ADMIRE fits every image line over its full depth, so it cannot be patched + if patch_size is not None or num_patches not in (None, 1): + requested = ( + f"patch_size={patch_size}" + if patch_size is not None + else f"num_patches={num_patches}" + ) + raise ValueError( + "The ADMIRE beamformer fits every image line over its full depth, so it " + f"needs the whole grid at once and cannot be patched. Got {requested}; " + "use num_patches=1 (the default for ADMIRE) or leave it unset." + ) + num_patches = 1 + self.beamformer_type = beamformer self.num_patches = num_patches # Resolved here so a non-compact config records the size actually used. @@ -2114,6 +2132,293 @@ def call(self, grid=None, f_number=None, **kwargs): return {self.output_key: beamformed_data} +@beamformer_registry("admire") +@ops_registry("admire") +class ADMIRE(Operation): + """Aperture Domain Model Image REconstruction (ADMIRE) beamformer. + + ADMIRE suppresses off-axis and multipath (reverberation) clutter by fitting a + model of point-scatterer wavefronts to the aperture-domain signal of short axial + windows, frequency by frequency, and rebuilding the channel data from only the + scatterers in a small region around each window. The decluttered channels are + then summed as in :class:`DelayAndSum`. See :mod:`zea.beamform.admire` for the + method and :class:`~zea.beamform.admire.ADMIREConfig` for its settings. + + ADMIRE models image *lines*: for every column of the beamforming ``grid`` it + fits the ``n_elements`` neighbouring elements, centered as closely as the element + pitch allows on that column (elements beyond the array edge contribute zeros). It + therefore needs: + + - a linear array (``probe_geometry`` with uniformly spaced elements along x); + - a regular cartesian ``grid`` of shape ``(n_z, n_x, 3)`` with uniform axial + spacing, fine enough to sample the signal band (for IQ data the axial + sampling rate ``c / (2 dz)`` must exceed the fitted bandwidth; for RF data it + must exceed twice its highest frequency). For exact lateral registration put + the grid columns on the element positions (odd ``n_elements``) or halfway + between them (even ``n_elements``); + - the full grid at once: windows run along depth, so the pipeline must not be + patched. Use ``Beamform(beamformer="admire", num_patches=1)``. + + Transmits are compounded (summed) before ADMIRE is applied, so the method works + on the compounded channel data of plane-wave or diverging-wave sequences as well + as on scanline data. Aperture growth follows the pipeline ``f_number`` + parameter (``0`` disables it). + + The models depend on the geometry only; they are generated (on the CPU, which + takes seconds to minutes) the first time the operation sees a geometry, and + cached in memory and on disk. Pass precomputed ``models`` (see + :func:`~zea.beamform.admire.generate_admire_models`) to skip generation. Because + the models and the line layout are derived from concrete parameter values, the + operation itself runs eagerly and cannot be traced as part of a compiled + pipeline; the fit is JIT compiled on its own (when ``jit_compile`` is true, + the default) and reused for as long as the geometry stays the same. + + .. citation:: byram2015model, dei2019computationally, khan2021realtime + + Args: + n_elements (int, optional): Number of elements per image line. Defaults to + all elements of the probe. Model generation and fitting get slower with + more elements. + config (ADMIREConfig or dict, optional): Model and fit settings. Defaults to + the settings of the reference implementation. + models (ADMIREModels, optional): Precomputed models. When given, ``n_elements`` + and ``config`` are taken from them and no models are generated. + window_batch_size (int, optional): Number of STFT windows fitted at once, + which bounds memory use. Defaults to all windows at once. + **kwargs: Forwarded to :class:`~zea.ops.base.Operation`. + """ + + def __init__( + self, + n_elements=None, + config=None, + models=None, + window_batch_size=None, + **kwargs, + ): + from zea.beamform.admire import ADMIREConfig + + if n_elements is not None and (not isinstance(n_elements, int) or n_elements < 2): + raise ValueError(f"n_elements must be an integer of at least 2, got {n_elements!r}.") + if models is not None: + n_elements = models.n_elements + config = models.config + if isinstance(config, ADMIREConfig): + config = dataclasses.asdict(config) + if config is not None: + # Validate early; lists (e.g. from a JSON config) become tuples again + config = {k: tuple(v) if isinstance(v, list) else v for k, v in config.items()} + ADMIREConfig(**config) # ty: ignore[invalid-argument-type] + super().__init__( + input_data_type=DataTypes.ALIGNED_DATA, + output_data_type=DataTypes.BEAMFORMED_DATA, + jittable=False, + **kwargs, + ) + self.n_elements = n_elements + self.config = config + self.models = models + self.window_batch_size = window_batch_size + # The jit flag is consumed by the fit, since the operation itself runs eagerly + self._jit_fit = kwargs.get("jit_compile", True) + self._models_cache = {} + self._fit_cache = {} + self._checked_geometry = None + + def get_dict(self, compact=True): + """Serialize the operation, leaving out precomputed ``models``. + + Models are regenerated (or loaded from the disk cache) on first use. + """ + models = self.models + self.models = None + try: + return super().get_dict(compact=compact) + finally: + self.models = models + + @staticmethod + def _linear_array_pitch(probe_x): + """Element pitch of a uniform linear array.""" + spacing = np.diff(probe_x) + pitch = float(np.median(spacing)) + if pitch <= 0 or not np.allclose(spacing, pitch, rtol=1e-3, atol=0): + raise ValueError( + "ADMIRE requires a linear array with uniformly spaced elements along x." + ) + return pitch + + def _get_models( + self, depths, pitch, n_elements, sound_speed, center_frequency, f_number, analytic + ): + """Return the models for a geometry, generating them on first use.""" + from zea.beamform.admire import ADMIREConfig, generate_admire_models + + if self.models is not None: + return self.models + config = ADMIREConfig(**(self.config or {})) # ty: ignore[invalid-argument-type] + key = ( + depths.tobytes(), pitch, n_elements, sound_speed, center_frequency, f_number, + analytic, config, + ) # fmt: skip + if key not in self._models_cache: + log.info("Generating ADMIRE models; this can take a few minutes.") + self._models_cache = { + key: generate_admire_models( + depths, + sound_speed=sound_speed, + center_frequency=center_frequency, + pitch=pitch, + n_elements=n_elements, + f_number=f_number, + analytic=analytic, + config=config, + verbose=True, + ) + } + return self._models_cache[key] + + def call( + self, + grid=None, + probe_geometry=None, + sound_speed=None, + center_frequency=None, + demodulation_frequency=None, + f_number=None, + **kwargs, + ): + """Apply ADMIRE to TOF-corrected data and sum the channels. + + Args: + data (ops.Tensor): The TOF corrected input of shape + ``(n_tx, n_z * n_x, n_el, n_ch)`` with optional batch dimension. + grid (ops.Tensor): Beamforming grid of shape ``(n_z, n_x, 3)``. + probe_geometry (ops.Tensor): Element positions of shape ``(n_el, 3)``. + sound_speed (float): Speed of sound in m/s. + center_frequency (float): Transmit center frequency in Hz. + demodulation_frequency (float): Demodulation frequency in Hz. For IQ + data, it stands in for a zero ``center_frequency``, which is what + :class:`~zea.ops.Demodulate` leaves behind. + f_number (float): F-number of aperture growth. + + Returns: + dict: Dictionary containing the beamformed data of shape + ``(n_z * n_x, n_ch)`` with optional batch dimension. + """ + data = kwargs[self.key] + if grid is None or len(grid.shape) != 3: + raise ValueError("ADMIRE requires a 2D `grid` parameter of shape (n_z, n_x, 3).") + if probe_geometry is None: + raise ValueError("ADMIRE requires the `probe_geometry` parameter.") + n_z, n_x = grid.shape[0], grid.shape[1] + n_pix, n_el, n_ch = data.shape[-3:] + if n_pix != n_z * n_x: + raise ValueError( + f"ADMIRE needs the full grid of {n_z * n_x} pixels at once, but got {n_pix}. " + "Disable patching with `Beamform(beamformer='admire', num_patches=1)`." + ) + + try: + grid_np = ops.convert_to_numpy(grid).astype(float) + probe_x = ops.convert_to_numpy(probe_geometry)[:, 0].astype(float) + geometry = [ + float(ops.convert_to_numpy(v)) if v is not None else None + for v in (sound_speed, center_frequency, demodulation_frequency, f_number) + ] + except Exception as e: + raise RuntimeError( + "ADMIRE generates its models from concrete parameter values and cannot " + "run inside a traced function. Pass precomputed `models` (see " + "zea.beamform.admire.generate_admire_models) to compile it." + ) from e + sound_speed, center_frequency, demodulation_frequency, f_number = geometry + if n_ch == 2 and not center_frequency: + # Aligned IQ data is analytic and carries the carrier again (see tof_correction) + center_frequency = demodulation_frequency + if not sound_speed or not center_frequency: + raise ValueError( + "ADMIRE requires the `sound_speed` and `center_frequency` parameters (or, " + "for IQ data, `demodulation_frequency`)." + ) + + pitch = self._linear_array_pitch(probe_x) + n_sub = self.n_elements if self.n_elements is not None else n_el + + # For every grid column, the sub-aperture whose center lies closest to it + if self.models is not None: + relative = self.models.element_positions + else: # the positions generate_admire_models will use + relative = (np.arange(n_sub) - (n_sub - 1) / 2) * pitch + columns = grid_np[0, :, 0] + first = np.round((columns + relative[0] - probe_x[0]) / pitch).astype(int) + indices = first[:, None] + np.arange(n_sub) + indices = np.where((indices >= 0) & (indices < n_el), indices, n_el) # n_el: zeros + + # Check the geometry once, before the (slow) model generation + geometry_key = (grid_np.tobytes(), indices.tobytes()) + if geometry_key != self._checked_geometry: + self._warn_unsupported_geometry(grid_np, pitch, indices, n_el) + self._checked_geometry = geometry_key + + models = self._get_models( + grid_np[:, 0, 2], pitch, n_sub, sound_speed, center_frequency, f_number or 0.0, + analytic=n_ch == 2, + ) # fmt: skip + key = (id(models), indices.tobytes(), data.shape, self.with_batch_dim) + if key not in self._fit_cache: + fn = self._make_fit(models, indices, n_z, n_x, n_el, n_ch) + self._fit_cache = {key: jit(fn) if self._jit_fit else fn} + return {self.output_key: self._fit_cache[key](data)} + + @staticmethod + def _warn_unsupported_geometry(grid, pitch, indices, n_el): + """Warn about the easy-to-spot geometries that ADMIRE does not model. + + Not exhaustive: it catches non-cartesian grids (such as polar sector grids) + and grids that extend well beyond the array. + """ + x, z = grid[..., 0], grid[..., 2] + tol = 1e-3 * pitch + if not (np.allclose(x, x[:1], atol=tol) and np.allclose(z, z[:, :1], atol=tol)): + log.warning( + "ADMIRE models every grid column as an unsteered line straight below the " + "array, but this grid is not cartesian (for example a polar or sector grid). " + "The result will not be a valid ADMIRE image; use grid_type='cartesian'. " + "Steered lines, as in phased-array sector scans, are not supported." + ) + mostly_off = np.mean(indices == n_el, axis=1) > 0.5 + if np.mean(mostly_off) > 0.25: + log.warning( + f"For {np.mean(mostly_off):.0%} of the grid columns, more than half of the " + "ADMIRE sub-aperture lies beyond the array, where the channels are filled " + "with zeros. Limit the grid to the width of the array (xlims) or use a " + "smaller n_elements." + ) + + def _make_fit(self, models, indices, n_z, n_x, n_el, n_ch): + """Build the function that beamforms aligned data with ADMIRE.""" + from zea.beamform.admire import apply_admire + + n_lines, n_sub = indices.shape + # Flat gather indices into (n_x * (n_el + 1)): each column with a zero channel + flat_indices = (np.arange(n_x)[:, None] * (n_el + 1) + indices).ravel() + + def process_image(image): + image = ops.sum(ops.cast(image, "float32"), axis=0) # compound transmits + image = ops.reshape(image, (n_z, n_x, n_el, n_ch)) + image = ops.pad(image, [[0, 0], [0, 0], [0, 1], [0, 0]]) + image = ops.reshape(image, (n_z, n_x * (n_el + 1), n_ch)) + lines = ops.take(image, flat_indices, axis=1) + lines = ops.reshape(lines, (n_z, n_lines, n_sub, n_ch)) + lines = apply_admire(lines, models, window_batch_size=self.window_batch_size) + return ops.reshape(ops.sum(lines, axis=2), (n_z * n_x, n_ch)) + + if not self.with_batch_dim: + return process_image + return lambda data: ops.map(process_image, data) + + @ops_registry("refocus") class Refocus(Operation): r"""REFoCUS (Retrospective Encoding For Conventional Ultrasound Sequences).