Add zea.inverse: recover channel data from beamformed images - #511
Draft
sankethvedula wants to merge 20 commits into
Draft
Add zea.inverse: recover channel data from beamformed images#511sankethvedula wants to merge 20 commits into
sankethvedula wants to merge 20 commits into
Conversation
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR adds differentiable inverse beamforming with DAS and scatterer operators, matrix-free solvers, seeding, inversion, refinement, exports, documentation, notebook settings, and tests. It also improves legacy waveform-index parsing and updates a display docstring. ChangesInverse beamforming
Legacy file loading
Documentation cleanup
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
tristan-deep
added
enhancement
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
zea/data/legacy_file.py (1)
157-168: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHandle missing
tx_waveform_indicesto prevent iteration crashes.If a legacy file omits
tx_waveform_indices(e.g., becausewaveforms_one_wayis already correctly stored as an array of shape(n_tx, n_samples)),tx_waveform_indicesevaluates toNone. This will cause aTypeError: 'NoneType' object is not iterablewhen the list comprehension attempts to iterate over it.Guard the stack operations to ensure they only execute when
tx_waveform_indicesis present.🐛 Proposed fix
- if "waveforms_one_way" in scan_parameters: + if "waveforms_one_way" in scan_parameters and tx_waveform_indices is not None: waveforms_one_way_list = scan_parameters["waveforms_one_way"] scan_parameters["waveforms_one_way"] = np.stack( [waveforms_one_way_list[i] for i in tx_waveform_indices] ) - if "waveforms_two_way" in scan_parameters: + if "waveforms_two_way" in scan_parameters and tx_waveform_indices is not None: waveforms_two_way_list = scan_parameters["waveforms_two_way"] scan_parameters["waveforms_two_way"] = np.stack( [waveforms_two_way_list[i] for i in tx_waveform_indices] )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@zea/data/legacy_file.py` around lines 157 - 168, Guard the waveforms_one_way and waveforms_two_way stack operations so they run only when tx_waveform_indices is present. Preserve the existing behavior when indices are available, and leave already-correct waveform arrays unchanged when tx_waveform_indices is None.
🧹 Nitpick comments (1)
docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb (1)
509-520: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
FunctionTimerfor these benchmark timings.
time.time()can miss JAX async execution here, so the reported durations are lower than the actual runtime.zea.utils.FunctionTimeralready blocks on completion and usestime.perf_counter()internally.
docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb#L509-L520docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb#L678-L688🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb` around lines 509 - 520, The benchmark timing blocks at docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb lines 509-520 and 678-688 should use zea.utils.FunctionTimer instead of time.time(), wrapping the relevant inversion calls so JAX execution is synchronized and durations use perf_counter(). Update both sites consistently, preserving the existing result assignments, output labels, and reporting behavior.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@zea/inverse/inversion.py`:
- Around line 152-160: Update the inversion flow around DASOperator and
seed_scatterers so custom flatgrid configurations cannot seed against the
unrelated parameters.grid. Validate that the operator’s flatgrid and
parameters.grid have matching coordinates, or use an explicitly supplied
matching seeding grid/initial positions; reject mismatches before calling
seed_scatterers while preserving correct behavior for the default grid.
In `@zea/inverse/operators.py`:
- Line 288: Validate the normalized chunk size and magnitudes before the first
_pad_and_chunk call: reject chunk_size <= 0 and empty magnitudes with the
module’s standard validation error rather than allowing a ZeroDivisionError.
Update the initialization logic around self.chunk_size and the related handling
at the referenced later path, preserving valid positive chunk sizes and
non-empty magnitudes.
---
Outside diff comments:
In `@zea/data/legacy_file.py`:
- Around line 157-168: Guard the waveforms_one_way and waveforms_two_way stack
operations so they run only when tx_waveform_indices is present. Preserve the
existing behavior when indices are available, and leave already-correct waveform
arrays unchanged when tx_waveform_indices is None.
---
Nitpick comments:
In `@docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb`:
- Around line 509-520: The benchmark timing blocks at
docs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynb
lines 509-520 and 678-688 should use zea.utils.FunctionTimer instead of
time.time(), wrapping the relevant inversion calls so JAX execution is
synchronized and durations use perf_counter(). Update both sites consistently,
preserving the existing result assignments, output labels, and reporting
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: dac933a4-fc41-4f8c-8f01-ebaf7d8d75ee
⛔ Files ignored due to path filters (2)
docs/source/notebooks/pipeline/inverse_beamforming_carotid_prebf.gifis excluded by!**/*.gifdocs/source/notebooks/pipeline/inverse_beamforming_cirs_prebf.gifis excluded by!**/*.gif
📒 Files selected for processing (13)
docs/source/_autosummary/zea.rstdocs/source/notebooks/pipeline/inverse_beamforming_example.ipynbdocs/source/notebooks/pipeline/inverse_beamforming_real_scans_example.ipynbtests/data/test_file.pytests/test_inverse.pytests/test_notebooks.pyzea/__init__.pyzea/data/legacy_file.pyzea/inverse/__init__.pyzea/inverse/inversion.pyzea/inverse/operators.pyzea/inverse/seeding.pyzea/inverse/solvers.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
zea/inverse/__init__.py (1)
56-60: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDefine an explicit
__all__for the package API.Without it, wildcard imports expose the implementation submodules as public names and make the export surface depend on internal imports. Restore an explicit list of supported exports, or confirm that this expanded wildcard API is intentional.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@zea/inverse/__init__.py` around lines 56 - 60, Define an explicit __all__ in the zea.inverse package initializer listing only the supported public symbols, including the intended inversion functions and classes, operators, seeding utilities, solvers, and any intentionally public submodules. Exclude implementation-only imports so wildcard imports remain stable and independent of internal module imports.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@zea/display.py`:
- Line 24: Update the return-value documentation for the affected display
function to state that it returns either an ndarray or a PIL.Image.Image,
depending on the pillow parameter; preserve the existing output range
description.
---
Nitpick comments:
In `@zea/inverse/__init__.py`:
- Around line 56-60: Define an explicit __all__ in the zea.inverse package
initializer listing only the supported public symbols, including the intended
inversion functions and classes, operators, seeding utilities, solvers, and any
intentionally public submodules. Exclude implementation-only imports so wildcard
imports remain stable and independent of internal module imports.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: ec640d25-f901-4541-ba3c-b7d642d7492f
📒 Files selected for processing (7)
tests/data/test_file.pytests/test_inverse.pyzea/data/legacy_file.pyzea/display.pyzea/inverse/__init__.pyzea/inverse/inversion.pyzea/inverse/operators.py
🚧 Files skipped from review as they are similar to previous changes (4)
- zea/inverse/inversion.py
- zea/data/legacy_file.py
- tests/test_inverse.py
- zea/inverse/operators.py
|
|
||
| Returns: | ||
| image (ndarray): Output 8 bit image(s) [0, 255]. | ||
| ndarray: Output 8 bit image(s) [0, 255]. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document both possible return types.
Line [24] says the function always returns an ndarray, but the default pillow=True path returns Image.fromarray(...). Document ndarray or PIL.Image.Image, depending on pillow.
Proposed fix
- ndarray: Output 8 bit image(s) [0, 255].
+ ndarray or PIL.Image.Image: Output 8 bit image(s) [0, 255],
+ depending on the value of ``pillow``.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| ndarray: Output 8 bit image(s) [0, 255]. | |
| ndarray or PIL.Image.Image: Output 8 bit image(s) [0, 255], | |
| depending on the value of ``pillow``. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@zea/display.py` at line 24, Update the return-value documentation for the
affected display function to state that it returns either an ndarray or a
PIL.Image.Image, depending on the pillow parameter; preserve the existing output
range description.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/data/test_file.py (1)
2460-2479: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConstruct the compatibility fixture without the current writer.
This test claims to verify reading an existing non-snake-case file, but
FileSpec.save()creates that malformed file. This couples reader compatibility to—and implicitly permits—the writer bypassing the naming restriction. Save a valid fixture, then rename the HDF5 key directly before loading it, as done intest_custom_element_non_snake_case_still_readable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/data/test_file.py` around lines 2460 - 2479, The test_custom_data_metadata_non_snake_case_still_readable fixture currently uses FileSpec.save(), so it does not represent a legacy malformed file. Replace the writer-generated “My Overlay” setup with a valid fixture, then rename the relevant HDF5 key directly before opening it with File, following the approach in test_custom_element_non_snake_case_still_readable. Keep the existing read assertions unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/data/test_file.py`:
- Around line 2460-2479: The
test_custom_data_metadata_non_snake_case_still_readable fixture currently uses
FileSpec.save(), so it does not represent a legacy malformed file. Replace the
writer-generated “My Overlay” setup with a valid fixture, then rename the
relevant HDF5 key directly before opening it with File, following the approach
in test_custom_element_non_snake_case_still_readable. Keep the existing read
assertions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 0d52a803-e630-49f3-9346-12c755dae993
📒 Files selected for processing (1)
tests/data/test_file.py
Collaborator
|
Hey @sankethvedula nice work! I had a read through the code and had some initial questions / comments, curious what your thoughts are!
I'm thinking maybe the simplest approach would just be to enable the frequency-domain simulator to use a waveform stored in a scan either? Maybe we'd just need to extent
|
Some legacy (pre-v0.1.0) jaxus-style files store scan/tx_waveform_indices as float64; legacy_scan used them directly to index the per-waveform dataset lists, raising an IndexError. Cast them to int before indexing, and remove a stray unassigned np.stack(...) line that would raise a NameError for files that only contain two-way waveforms. The same files may also carry lens parameters (lens_correction, lens_thickness, lens_sound_speed) in the scan group; the current format keeps lens parameters with the probe, so drop them instead of failing ScanSpec validation with an unexpected-keyword error. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
New subpackage that inverts the DAS beamformer: given a post-beamformed image, recover the pre-beamformed channel data by expressing beamforming as a differentiable linear operator and solving a least-squares problem. - operators.py: DASOperator wraps the zea.beamform primitives (calculate_delays, apply_delays, fnumber_mask) into a linear channel-data -> image map with an autodiff adjoint; ScattererSimulator is a time-domain point-scatterer forward model using the scan's stored two-way waveforms and the same delay model as the beamformer. Both iterate over transmits and scatterer chunks with ops.scan and keras.remat, so memory under automatic differentiation stays bounded for scans with hundreds of transmits. - solvers.py: matrix-free CGLS and linear_adjoint (adjoint of a linear operator as the gradient of <A(x), y> via zea.backend.AutoGrad). - seeding.py: seed_scatterers samples scatterer positions from the image envelope. - inversion.py: invert_direct (minimum-norm pseudo-inverse; fits the image but not the physical channel data) and invert_scatterers (scatterer prior regularizing the DAS nullspace, with optional joint Adam refinement via zea.backend.optimizer.adam). Everything is written with keras.ops and runs on all Keras backends. The scatterer-prior inversion follows the off-grid scatterer model of van de Schaft et al. (arXiv:2407.02285). The beamforming operator was verified against an independent jaxus-based DAS implementation on a simulated CIRS phantom (relative error 4e-4, correlation 1.000000). Includes tests (adjoint identity, CGLS vs lstsq/pinv, echo arrival times, seeding, end-to-end inversions) and an introductory example notebook under docs/source/notebooks/pipeline/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Applies zea.inverse to two real acquisitions and reproduces the results of the standalone das-inverse study it was ported from: a simulated CIRS point-scatterer phantom (direct pseudo-inverse vs scatterer prior) and an in-vivo carotid scan (149 single-element transmits, 30k scatterers). Executed at full scale on an NVIDIA L40S; recovered/reference pre-BF correlations: CIRS direct 0.560/0.557, CIRS scatterers 0.855/0.838, carotid 0.059/0.057. Registered in tests/test_notebooks.py with reduced iteration counts and the carotid section disabled for CI (411 MB download). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Tests (the CI tests job runs with all backends installed): - Cast operator inputs (grid, positions, waveforms, magnitudes, images) to float32: numpy-built parameters arrive as float64, which jax silently demotes but tensorflow/torch propagate into dtype errors in the delay computation. - Return the per-transmit channel as both carry and stacked output of the simulator's transmit scan: keras's tensorflow scan requires the stacked outputs to match the carry's shape and dtype. - Pass a static length to every ops.scan and compile AD-derived adjoint/ gradient functions without XLA on tensorflow (via a small helper): XLA cannot size the gradient accumulators of the scan's while loop. None of this changes jax numerics, so the executed example notebooks remain valid. Docs (sphinx -W treats warnings as errors): - Document InversionResult fields with an Args section like other zea dataclasses; the Attributes section duplicated autodoc's field entries. - Drop __all__ from zea/inverse/__init__.py and import the submodules instead, matching zea.beamform/zea.data: with autosummary_ignore_module_all = False, __all__ suppressed the submodule toctree (orphan pages) and documented every re-export twice (ambiguous cross-references). - Fix the Returns section of zea.display.to_8bit to napoleon's 'type: description' form; the named form made 'image' a type cross-reference that InversionResult.image turned ambiguous. Verified: 19/19 tests green under jax, tensorflow and torch main-process backends with the cross-backend workers active; docs build green with -W. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- legacy_scan: keep waveforms as-is when tx_waveform_indices is absent (files that already store one waveform per transmit) instead of failing on iterating None; regression test with array-stored waveforms. - invert_scatterers: reject operators built on a custom flatgrid with a clear ValueError — seeding samples positions from parameters.grid, so a mismatched grid would silently seed at unrelated coordinates (or fail on reshape). Custom grids should seed manually and use ScattererSimulator with cgls directly. - ScattererSimulator: validate chunk_size > 0 at construction and non-empty magnitudes in __call__ instead of surfacing a ZeroDivisionError from the chunking arithmetic. 19/19 tests green under jax, tensorflow and torch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The inverse beamforming notebooks link the original study's repository with a repo-root URL, which is not covered by the existing permalink ignores. GitHub rate-limits anonymous requests from CI runners (the link checks fine from other hosts), so ignore it like the other non-permalink GitHub links. This was the remaining test-docs-build failure: docs-build itself passes since the -W fixes, and a clean docs-build plus full linkcheck both pass locally at this head. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sankethvedula
force-pushed
the
feature/inverse-beamforming
branch
from
July 21, 2026 07:50
dbc01e6 to
f918b97
Compare
tof_correction accepts fnum_window_fn, but the operation always used the rect default. Expose it as a constructor argument (default unchanged) so pipelines can select e.g. a tukey window.
Replace the hand-rolled delays/gather/mask implementation and its rematerialized transmit scan with the regular zea pipeline: Pipeline([PatchedGrid([TOFCorrection(fnum_window_fn), DelayAndSum()])]). The delay model now runs through the exact code path the rest of zea uses, and memory is bounded by processing the imaging grid in pixel patches (a num_patches knob on the operator) rather than by checkpointing transmits. jit_options=None keeps the operator a pure function, so the inversion drivers still jit and differentiate matvec/rmatvec as a whole. Full-scale smoke on a 149-transmit in-vivo scan matches the previous implementation exactly: measured-image correlation 1.000000 against the das-inverse reference and identical CGLS correlations (0.568 CIRS, 0.038 carotid), within the same memory budget on a 48 GB L40S.
tristan-deep
marked this pull request as draft
July 21, 2026 08:34
Add get_measured_pulse_spectrum_fn, which evaluates the DTFT of a sampled waveform at arbitrary frequencies, and let simulate_rf (and the Simulate op) accept a pulse_spectrum_fn override. This drives the frequency-domain simulator with the waveforms stored in a scan (parameters.waveforms_two_way) instead of the parametric four-period pulse, following the same time-zero convention as zea.inverse.ScattererSimulator (the beamformer adds t_peak). Cross-validating the two simulators on the same 104-scatterer phantom with the same two-cycle pulse yields matching B-mode artifact levels (background within a few dB), confirming they implement consistent physics once waveform-matched.
A controlled comparison (time- vs frequency-domain simulator, rect vs Hanning transmit apodization, +-8.6 vs +-5 degree steering) shows the granular background and lateral streaks in the fish B-mode are side lobes of compounding only three plane waves with a short two-cycle pulse: with a four-period pulse the background drops ~5 dB, while apodization and steering change it by <1 dB. Note this in the notebook next to the figure; the short pulse is kept because recovering realistic broadband channel data is the point of the example.
Replace the idealized 64-element/3.125 MHz toy setup with a realistic plane-wave acquisition: the Verasonics L11-5v probe preset (128 elements, 0.3 mm pitch, 6.25 MHz) at 4x sampling, five steered plane waves, a two-way waveform shaped by convolving a 1.5-cycle excitation twice with a Gaussian impulse response matching the probe bandwidth, diffuse tissue speckle underneath the fish phantom, and receiver noise at ~30 dB channel SNR. The measured image now reads like a routine clinical B-mode. Recovering channel data behind speckle and noise is honestly harder than in a sterile point-target scene (pseudo-inverse 0.56 vs scatterer prior 0.66 pre-beamforming correlation); the takeaways relate this to the point-phantom vs in-vivo gap in the real-scans notebook. The notebook gains a parameters cell, registered with reduced values in tests/test_notebooks.py for CI (verified via pytest: 2m45s on CPU).
The adjoint was constructed lazily on first use and cached. When the first call happened inside a jitted function (invert_direct compiles the adjoint as rmatvec), the cached closure captured traced values, and any later jitted inversion on the same operator raised jax's UnexpectedTracerError. Build it once in __init__ instead — the closure then always holds concrete arrays — and add a regression test that runs invert_direct twice on a fresh operator.
The docs environment pins KERAS_BACKEND=numpy (docs/Makefile), and 'make docs-test' in CI executes the docstring examples there. The linear_adjoint example computes a real gradient through AutoGrad, which supports only jax/tensorflow/torch, so the final line can never pass in the docs environment — mark it '# doctest: +SKIP' like the repo's other backend-dependent examples. The behavior itself stays covered by tests/test_inverse.py::test_linear_adjoint_matches_matrix_transpose. This was the actual cause of the test-docs-build CI failures (the doctest stage runs between docs-build and linkcheck).
Three ordered tutorials under docs/source/notebooks/inversion/, as suggested in review: the synthetic fish walkthrough, the recorded CIRS phantom scan, and the in-vivo carotid scan — the last two split out of the former real-scans notebook, each self-contained and executed at full scale (CIRS: direct 0.560/1.000 and scatterer prior 0.857/0.996 vs references 0.557/1.000 and 0.838/0.996; carotid: 0.059/0.942 vs 0.057/0.946). The fish tutorial gains a 'breaking the inverse crime' control that regenerates the data with the frequency-domain simulator driven by the measured waveform (scatterer prior 0.515 vs 0.659 on matched-model data); a refresh of its executed outputs with the added pseudo-inverse baseline follows. Wires the section into examples.rst and the notebook redirects, adds an 'inversion' group to the notebook-tests matrix, updates the zea.inverse docstring links, and registers CI parameters (CIRS reduced, carotid and the frequency-domain control skipped on CI). Note: the previous push unintentionally carried the half-staged notebook move (stale index), which orphaned the moved notebook and broke test collection against a registered-but-deleted notebook — this commit restores a consistent tree.
Rename the tutorials folder to notebooks/inverse (matching the zea.inverse subpackage) with short filenames — fish_example, cirs_example, carotid_example, and matching *_prebf.gif names — and update badges, cross-links, docs toctrees, redirects, CI matrix group and notebook registrations accordingly. Drop the frequency-domain control section from the fish tutorial: the capability stays in zea.simulator (unit-tested on all backends) and the cross-simulator evidence is recorded in the PR discussion, but the chunked per-element simulation ran for hours — too heavy for a tutorial. Run the recorded-scan inversions with 40 CGLS iterations instead of the original study's 70: on these ill-posed problems the iteration count acts as the regularization parameter, and 40 captures nearly all of the attainable correlation at a fraction of the runtime (CIRS scatterer prior 0.842 vs reference 0.838 in 17 s instead of ~8 min; carotid 0.060 vs reference 0.057 in 22 min instead of 43). All three notebooks re-executed at full scale; the carotid tutorial now also points to the zea-native zeahub/zea-carotid-2023 dataset (the Zenodo scan is kept for direct comparability with the das-inverse study).
The per-chunk echoes tensor was laid out (chunk, n_ax, n_el) and contracted over its leading (strided) axis, which XLA lowers to a transposing reduction — on a memory-bandwidth-bound kernel that roughly doubles the traffic, and its VJP is worse still. Switch to the element-major layout (n_el, n_ax, chunk) and contract the contiguous scatterer axis instead. Benchmark (H100, carotid scan, 149 transmits, 30k scatterers, chunk 4096, jitted, steady-state): forward 11.76 s -> 7.45 s, AD adjoint 19.47 s -> 7.67 s; a CGLS iteration drops from 31.2 s to 15.1 s. Same numerics up to float reassociation; all tests pass.
Collaborator
|
Hi @sankethvedula, thanks for this contribution, looks really good! Could you try to merge the interesting parts into a single example notebook? The multiple notebooks are appreciated as reference material (maybe hosted in a different repository, with zea version pinned at some point), but we try to keep the notebook examples in the zea docs curated. When looking at the difference in your notebooks, it is mostly various datasets, which you can easily show in one notebook I think. I see some broken link formatting in some of the notebooks as well. Lastly, try to follow the format and notebook layout / structure of some of the other notebooks. Besides showing this new inverse feature, the notebooks also serve the role of good zea practices. We'd like to reuse as many tools within the library to keep the notebooks minimal and highlighting the specific thing we want to show (in this case the inverse modeling!) Once you have a single notebook we can help with that as well to clean that up! :) Good technical comments by @OisinNolan as well. We can host a sample on zeahub (see other notebook examples). Maybe there is already an existing dataset on zeahub you'd like and otherwise we can upload one. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were made?
This PR adds
zea.inverse, a new subpackage that inverts the DAS beamformer: given a post-beamformed image, it recovers the pre-beamformed channel data by expressing beamforming as a differentiable linear operator and solving a least-squares problem, optionally regularized with a point-scatterer prior.zea/inverse/operators.py—DASOperator: the DAS beamformer as a linear channel-data → image map, built as a regularzea.Pipeline(PatchedGrid([TOFCorrection(fnum_window_fn), DelayAndSum()])), so it shares zea's delay model (incl. lens correction) with everything else; grid-pixel patching (anum_patchesknob) bounds memory, and the adjoint comes for free via autodiff.ScattererSimulator: a time-domain point-scatterer forward model driven by the scan's stored two-way waveforms, consistent with the beamformer's delay model.zea/inverse/solvers.py— matrix-freecglsandlinear_adjoint(adjoint of a linear operator as the gradient of ⟨A(x), y⟩ viazea.backend.autograd.AutoGrad).zea/inverse/seeding.py—seed_scatterers: samples scatterer positions from the image envelope.zea/inverse/inversion.py— high-level drivers:invert_direct(minimum-norm pseudo-inverse) andinvert_scatterers(scatterer prior with optional joint Adam refinement).TOFCorrectionaccepts anfnum_window_fn(default unchanged); the frequency-domain simulator can run with measured waveforms viaget_measured_pulse_spectrum_fn+ apulse_spectrum_fnoverride onsimulate_rf/Simulate; a legacy-loader fix (floattx_waveform_indices, jaxus-style lens fields, files with only two-way waveforms);zea.inverseregistered as a lazy submodule and in the docs.Three tutorials in a new Inversion docs section (
docs/source/notebooks/inverse/), ordered from controlled to clinical:fish_example.ipynb— a realistic synthetic acquisition (Verasonics L11-5v preset, 4× sampling, five plane waves, a two-way pulse shaped by the probe's bandwidth, diffuse tissue speckle, 30 dB receiver noise). Ground truth is known exactly: the pseudo-inverse fits the image at 0.999 while recovering only 0.56 of the channel data; the scatterer prior recovers 0.66.cirs_example.ipynb— a recorded scan of a simulated CIRS phantom, reproducing the das-inverse study: pseudo-inverse 0.556/1.000 (reference 0.557/1.000), scatterer prior 0.842/0.993 (reference 0.838/0.996).carotid_example.ipynb— in-vivo carotid (149 single-element transmits): scatterer prior 0.060/0.931 (reference 0.057/0.946) — the honest ceiling: a compounded image constrains speckle statistics, not the channel micro-realization.The recorded-scan tutorials run 40 CGLS iterations instead of the study's 70: on these ill-posed problems the iteration count acts as the regularization parameter, and 40 captures nearly all of the attainable correlation at a fraction of the runtime (CIRS: 17 s instead of ~8 min). Each tutorial renders a recorded-vs-recovered channel-data GIF, and the carotid one points to
zeahub/zea-carotid-2023for zea-native carotid data (the Zenodo scan is kept for direct comparability with the study).Why were these changes made?
Inverse beamforming enables working backwards from images to channel data — useful for studying what information beamforming discards, for generating physically consistent channel data from image-domain priors, and as a building block for optimization-based reconstruction. The scatterer-prior inversion follows the off-grid scatterer model of van de Schaft et al., arXiv:2407.02285. Everything is written with
keras.opsand runs on all Keras backends.How were the changes validated?
tests/test_inverse.py,tests/test_simulator.py, legacy-loader regression tests): adjoint identity, operator linearity, CGLS vslstsq/pinv, simulator echo arrival times, simulator–beamformer consistency, measured-waveform spectra vs FFT, seeding properties, repeated-jit regression, and end-to-end inversions — green on jax, tensorflow and torch.DASOperatorverified against the independent jaxus-based DAS implementation of the original study: measured-image correlation 1.000000.inversegroup, docs build with-W, doctests, and linkcheck.How to test the changes
or start with the introductory tutorial:
docs/source/notebooks/inverse/fish_example.ipynb.Summary by CodeRabbit
zea.inverse, with API autosummaries updated accordingly.