Time-reverse simulation for phased-array delay design (#105)

* Add PointSources that have no direction or aperture

* Draft time-reverse simulation to calculate delays

* Update comment on point source

* Enable math latex rendering with Arithmatex

* Update docs/examples/plot_time_reverse.py

Co-authored-by: Diogo de Lucena <[email protected]>

* Fix up documentation - address PR comments

* Include target error calculation in plot_time_reverse.py

* lint/spellcheck

* Consolidate flags for source rendering type into a single Enum

* Add Enum to dictionary

* Fix spellcheck error / typo

* Add thought exercise for time reverse simulation

* Rename SourceType to SourceRenderType to be more specific

---------

Co-authored-by: Diogo de Lucena <[email protected]>

docs/examples/plot_time_reverse.py

@@ -0,0 +1,209 @@
+# -*- coding: utf-8 -*-
+"""
+Time-reverse simulation for phased array
+====================================================================
+
+The skull adds aberrations to the beam propagation; phased arrays can compensate
+for those by having different delays for each element, but estimating these
+delays can be challenging.
+One method to estimate the delays is a "time reverse" simulation:
+https://koreascience.kr/article/JAKO200612242715181.pdf
+This notebook demonstrates the "time reverse" method to estimate the delays. The
+notebook sets up a scenario with a phased array source and a target and then
+runs a simulation with the source and target reversed to calculate the delays.
+Finally, it uses the calculated delays to perform a forward-time simulation.
+
+Note: In this notebook, we refer to the "true" target as the eventual brain
+region we would like to stimulate, and the "true" source as the placement of
+the ultrasound probes. We refer to the "reversed" or "simulated" target and
+point-source as the values defined in our simulation, which are reversed from
+the physical setup to help calculate values.
+"""
+
+# %%
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import neurotechdevkit as ndk
+
+# Parameters
+SCENARIO_NAME = "scenario-2-2d-v0"
+NUM_ELEMENTS = 20
+ELEMENT_WIDTH = 1.2e-3
+
+
+# %%
+# Helper function to make the scenario with a PhasedArraySource
+def make_scenario(element_delays=None) -> ndk.scenarios.Scenario:
+    true_scenario = ndk.make(SCENARIO_NAME)
+
+    # define a phased-array source
+    default_source = true_scenario.get_default_source()
+    true_source = ndk.sources.PhasedArraySource2D(
+        element_delays=element_delays,
+        position=default_source.position,
+        direction=default_source.unit_direction,
+        num_elements=NUM_ELEMENTS,
+        pitch=default_source.aperture / NUM_ELEMENTS,
+        element_width=ELEMENT_WIDTH,
+        num_points=1000,
+    )
+
+    true_scenario.add_source(true_source)
+    return true_scenario
+
+
+# %%
+# ## Set up and visualize the forward scenario
+true_scenario = make_scenario()
+assert isinstance(true_scenario, ndk.scenarios.Scenario2D)
+true_scenario.render_layout()
+
+
+# %%
+# ## Simulate the time-reverse scenario
+# Place a point source at the true target, and simulate a pulse.
+# The point source is visualized as a gray dot.
+
+# Reinitialize the scenario
+reversed_scenario = ndk.make(SCENARIO_NAME)
+# and reverse the source
+point_source = ndk.sources.PointSource2D(
+    position=true_scenario.target.center,
+)
+reversed_scenario.add_source(point_source)
+
+assert isinstance(reversed_scenario, ndk.scenarios.Scenario2D)
+reversed_scenario.render_layout()
+
+
+# %%
+result = reversed_scenario.simulate_pulse()
+assert isinstance(result, ndk.results.PulsedResult2D)
+result.render_pulsed_simulation_animation()
+
+
+# %% Calculate the time-reverse delays
+# We calculate how long it took for the point-source pulse to reach each of
+# the true array elements. Here, we coarsely approximate these delays by
+# finding the pressure argmax at each element's nearest-neighbor coordinates.
+
+# Map array elements onto the nearest pixels in our simulation
+def map_coordinates_to_indices(coordinates, origin, dx):
+    indices = np.round((coordinates - origin) / dx).astype(int)
+    return indices
+
+
+# Get the pressure time-series of these elements
+[true_source] = true_scenario.sources
+assert isinstance(true_source, ndk.sources.PhasedArraySource2D)
+element_indices = map_coordinates_to_indices(
+    true_source.element_positions,
+    reversed_scenario.origin,
+    reversed_scenario.dx,
+)
+pressure_at_elements = result.wavefield[element_indices[:, 0], element_indices[:, 1]]
+
+# Calculate the time of arrival for each element
+element_reverse_delays = np.argmax(pressure_at_elements, axis=1) * result.effective_dt
+plt.plot(element_reverse_delays, marker="o")
+plt.xlabel("element index")
+plt.ylabel("delay [s]")
+
+
+# %%
+# Visually inspecting the earlier scenario layout, these results seem reasonable.
+# The expected delay \(t_d\) is approximately:
+#
+# $$
+# t_d \approx \frac{||x_{source} - x_{target}||_2}{c_{water}} \approx
+# \frac{0.07 \text{ m}}{1500 \text{ m/s}} \approx 47 \mu s
+# $$
+#
+
+
+# %%
+# ## Use delays in forward-time simulation
+# Next, let's validate these delays by using them in a normal forward-time
+# simulation.
+# We simulate the original scenario, setting the pulse delays as calculated.
+
+# Elements that took longer to reach should now be pulsed first,
+# so we invert the values
+element_delays = element_reverse_delays.max() - element_reverse_delays
+
+true_scenario = make_scenario(element_delays=element_delays)
+result = true_scenario.simulate_pulse()
+assert isinstance(result, ndk.results.PulsedResult2D)
+result.render_pulsed_simulation_animation()
+
+
+# %%
+# The pulse should focus on the true target.
+
+
+# %%
+# ### Simulate steady-state
+# Another way to visualize the simulation is to check that the steady-state
+# pressure (within the skull) peaks near the target.
+
+# Re-initialize scenario to clear previous simulation
+true_scenario = make_scenario(element_delays=element_delays)
+steady_state_result = true_scenario.simulate_steady_state()
+assert isinstance(steady_state_result, ndk.results.SteadyStateResult2D)
+steady_state_result.render_steady_state_amplitudes()
+
+
+# %%
+# We want to visualize and find the maximum pressure within the brain, so let's
+# mask out everything else.
+steady_state_pressure = steady_state_result.get_steady_state()
+# Only consider the brain region
+steady_state_pressure[~true_scenario.get_layer_mask("brain")] = np.nan
+steady_state_result.steady_state = steady_state_pressure
+
+steady_state_result.render_steady_state_amplitudes()
+
+
+# %%
+# We can also calculate how far the "time reverse" estimate is from the true
+# target.
+max_pressure_flat_idx = np.nanargmax(steady_state_pressure)
+max_pressure_idx = np.unravel_index(max_pressure_flat_idx, steady_state_pressure.shape)
+max_pressure_idx
+
+grid = steady_state_result.traces.grid.space.grid
+focal_point = np.array(
+    [
+        grid[0][max_pressure_idx[0]],
+        grid[1][max_pressure_idx[1]],
+    ]
+)
+# The backend grid is in different coordinates from the scenario grid, so we
+# need to shift it.
+focal_point += true_scenario.origin
+
+print("target center:", true_scenario.target.center)
+print("beam focal point:", focal_point)
+error_distance = np.linalg.norm(true_scenario.target.center - focal_point)
+print("error [m]:", error_distance)
+print("error [mm]:", error_distance * 1000)
+
+
+# %%
+# ## Reasons for target mismatch
+# The time-reverse simulation is not an exact solution for the forward-time
+# design. Other factors, like the angle of incidence at the boundary of two
+# materials, will be different in the time reverse vs forward-time.
+#
+# ### Exercise
+# Do you think the time-reverse simulation will work better or worse for deeper
+# targets? How about if the transducer was positioned next to a different part
+# of the skull that is flatter?
+
+
+# %%
+# ### Acknowledgments
+# Thanks to Sergio Jiménez-Gambín and Samuel Blackman for pointing us to the
+# "time reverse" simulation method.

docs/javascript/mathjax.js

@@ -0,0 +1,17 @@
+window.MathJax = {
+  tex: {
+    inlineMath: [["\\(", "\\)"]],
+    displayMath: [["\\[", "\\]"]],
+    processEscapes: true,
+    processEnvironments: true
+  },
+  options: {
+    ignoreHtmlClass: ".*|",
+    processHtmlClass: "arithmatex"
+  }
+};
+
+document$.subscribe(() => {
+  MathJax.typesetPromise()
+})
+

mkdocs.yml

@@ -47,6 +47,11 @@ nav:
 extra_css:
 - css/mkdocstrings.css
 
+extra_javascript:
+  - javascripts/mathjax.js
+  - https://polyfill.io/v3/polyfill.min.js?features=es6
+  - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
+
 markdown_extensions:
   - pymdownx.highlight:
       anchor_linenums: true

src/neurotechdevkit/rendering/__init__.py

@@ -6,7 +6,7 @@
     save_animation,
     video_only_output,
 )
-from ._source import source_should_be_flat
+from ._source import SourceRenderType
 from .layers import (
     SourceDrawingParams,
     draw_material_outlines,
@@ -41,7 +41,7 @@
     "make_animation",
     "video_only_output",
     "display_video_file",
-    "source_should_be_flat",
+    "SourceRenderType",
     "configure_matplotlib_for_embedded_animation",
     "save_animation",
 ]