add crosstalk simulation

#430

Author: peterhollender

examples/tutorials/03_Solution_Generation_and_Analysis.py

@@ -113,8 +113,8 @@
 # With the `target`, `transducer`, and `protocol` defined, we can now calculate the `Solution`.
 # The `calc_solution` method returns:
 # *   `solution`: The `Solution` object containing delays, apodizations, voltage, etc.
-# *   `sim_res`: A `SimResult` object if `simulate=True`.
-# *   `analysis`: A `SolutionAnalysis` object if `simulate=True` and `scale=False` (or `analysis` if `scale=True`).
+# *   `sim_res`: An `xa.Dataset` simulation result if `simulate=True`, aggregated over the focal points.
+# *   `analysis`: A `SolutionAnalysis` object if `simulate=True`.
 #
 # Setting `simulate=True` will run an acoustic simulation.
 # Setting `scale=True` will attempt to scale the output pressure to match `target_pressure` defined in the focal pattern or protocol, and returns a `scaled_analysis`.
@@ -162,7 +162,7 @@
 )
 solution2.apodizations[0,:128] = 0
 params = protocol1.sim_setup.setup_sim_scene(protocol1.seg_method)
-simulation_result = solution2.simulate(params=params, use_gpu=True)
+simulation_result = solution2.simulate(params=params)
 
 # %%
 simulation_result['p_min'].sel(focal_point_index=0).sel(y=0).plot.imshow()

src/openlifu/plan/protocol.py

@@ -22,7 +22,6 @@
 from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from openlifu.plan.target_constraints import TargetConstraints
 from openlifu.util.annotations import OpenLIFUFieldData
-from openlifu.util.checkgpu import gpu_available
 from openlifu.util.json import PYFUSEncoder
 from openlifu.virtual_fit import VirtualFitOptions
 from openlifu.xdc import Transducer
@@ -250,8 +249,8 @@ def calc_solution(
         sim_options: sim.SimSetup | None = None,
         analysis_options: SolutionAnalysisOptions | None = None,
         on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
-        use_gpu: bool | None = None,
-        voltage: float = 1.0
+        voltage: float = 1.0,
+        _force_cpu: bool = False
     ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:
         """Calculate the solution and aggregated k-wave simulation outputs.
 
@@ -267,8 +266,18 @@ def calc_solution(
                 The subject scan (Default: None).
                 It is expected to be in the simulation grid coordinates (lat, ele, ax).
                 If None, a default simulation grid will be used.
-            session: db.Session
-                A session used to define solution_id (Default: None).
+            params: xa.Dataset | None = None
+                The simulation parameters to use for the k-wave simulation (Default: None, which uses sim_options.setup_sim_scene on the volume).
+                If the params are provided, the volume is not required.
+                Params should include xa.DataArrays for the following:
+                    'sound_speed': speed of sound in the medium (m/s)
+                    'density': density of the medium (kg/m^3)
+                    'attenuation': attenuation coefficient of the medium (dB/cm/MHz)
+                    and can also include (though they aren't used in the simulation)
+                    'specific_heat': specific heat capacity of the medium (J/kg/K)
+                    'thermal_conductivity': thermal conductivity of the medium (W/m/K)
+            session: Session | None = None
+                the current session object (used for assigning a matching solution ID only)
             simulate: bool
                 Enable solution simulation (Default: true).
             scale: bool
@@ -291,9 +300,6 @@ def calc_solution(
             scaled_solution_analysis: SolutionAnalysis
                 This is the resulting rescaled analysis, if scale is enabled.
         """
-        if use_gpu is None:
-            use_gpu = gpu_available()
-
         if sim_options is None:
             sim_options = self.sim_setup
         if analysis_options is None:
@@ -346,7 +352,7 @@ def calc_solution(
             simulation_result = solution.simulate(
                 params=params,
                 sim_options=sim_options,
-                use_gpu=use_gpu)
+                _force_cpu=_force_cpu)
             solution.simulation_result = simulation_result
 
         # optionally scale the solution with simulation result

src/openlifu/plan/solution.py

@@ -139,22 +139,32 @@ def num_foci(self) -> int:
     def simulate(self,
         params: xa.Dataset,
         sim_options: SimSetup | None = None,
-        use_gpu: bool | None = None) -> xa.Dataset:
+        _force_cpu: bool = False) -> xa.Dataset:
+        """Run a simulation for this solution and store the result in the `simulation_result` attribute.
 
-        if use_gpu is None:
+        Args:
+            params: The simulation parameters as an xarray Dataset.
+            sim_options: Optional simulation setup options.
+            _force_cpu: Whether to force the simulation to run on the CPU. If True, the GPU will not be used even if available.
+
+        Returns:
+            The simulation result as an xarray Dataset.
+        """
+
+        if _force_cpu:
+            use_gpu = False
+        else:
             use_gpu = gpu_available()
 
         if sim_options is None:
             sim_options = SimSetup()
-        # check before if target is within bounds
         simulation_outputs_to_stack: List[xa.Dataset] = []
         simulation_cycles = np.round(self.pulse.duration * self.pulse.frequency)
         for focidx, focus in enumerate(self.foci):
-            self.logger.info(f"Beamform for focus {focus}...")
             delays = self.delays[focidx, :]
             apodization = self.apodizations[focidx, :]
             simulation_output_xarray = None
-            self.logger.info(f"Simulate for focus {focus}...")
+            self.logger.info(f"Running k-Wave simulation for focus {focus}...")
             run_simulation_kwargs = {
                     "arr": self.transducer,
                     "params": params,
@@ -422,15 +432,14 @@ def get_ita(self, intensity: xa.DataArray | None = None, units: str = "mW/cm^2")
         Calculate the intensity-time-area product for a treatment solution.
 
         Args:
-            output: A struct for simulation results from the treatment.
             intensity: xa.DataArray | None
                 If provided, use this intensity data array instead of the one from the simulation result.
             units: str
                 Target units. Default "mW/cm^2".
 
-
         Returns:
-            A Solution instance with the calculated intensity value.
+            xa.DataArray
+                The intensity-time-area product as an xarray DataArray.
         """
         if intensity is not None:
             intensity_scaled = rescale_data_arr(intensity, units)

src/openlifu/sim/kwave_if.py

@@ -95,9 +95,43 @@ def run_simulation(arr: xdc.Transducer,
                    return_kwave_outputs: bool = False,
                    return_kwave_inputs: bool = False,
                    sensor_record: List[str] = ['p_max', 'p_min'],
-                   _source: kSource|None = None,
-                   _sensor: kSensor|None = None
+                   _source = None,
+                   _sensor = None
 ):
+    """ Run a k-wave simulation for the given transducer array and parameters.
+    Args:
+        arr: The transducer array to simulate.
+        params: The simulation parameters as an xarray Dataset. Must include 'sound_speed', '
+        density', and 'attenuation' variables with appropriate units.
+        delays: Optional array of time delays for each element in the transducer array, in seconds. If None, no delays will be applied.
+        apod: Optional array of apodization values for each
+            element in the transducer array. If None, no apodization will be applied.
+        freq: The frequency of the input signal in Hz. Default is 1 MHz.
+        cycles: The number of cycles in the input signal. Default is 20.
+        amplitude: The amplitude of the input signal. Default is 1.
+        dt: The time step for the simulation in seconds. If 0, it will be automatically calculated based on the CFL condition.
+        t_end: The total time for the simulation in seconds. If 0, it will be automatically calculated based on the input signal duration and the CFL condition.
+        cfl: The Courant-Friedrichs-Lewy (CFL) number for the simulation. Default is 0.5.
+        bli_tolerance: The tolerance for the boundary layer integral method used in k-wave. Default
+            is 0.05.
+        upsampling_rate: The upsampling rate for the boundary layer integral method used in k-wave
+            Default is 5.
+        gpu: Whether to use GPU for the simulation. Default is True. If False, the simulation will run on the CPU. Note that running on CPU may be very slow for large simulations.
+        ref_values_only: Whether to use only the reference values for the medium properties (sound speed, density, attenuation) instead of the full spatial maps. Default is False. Setting this to True can significantly speed up the simulation, but will not capture any spatial variations in the medium properties.
+        return_kwave_outputs: Whether to return the raw outputs from k-wave in addition to the processed xarray Dataset. Default is False.
+        return_kwave_inputs: Whether to return the inputs to k-wave (kgrid, source, sensor, medium) in addition to the processed xarray Dataset. Default is False.
+        sensor_record: List of strings specifying which k-wave sensor outputs to record. Can include '
+        p_max', 'p_min', and 'p'. Default is ['p_max', 'p_min'].
+
+    Additional Args:
+        _source: Optional kSource object to use for the simulation. If None, a source will be created based on the transducer array and input signal.
+        _sensor: Optional kSensor object to use for the simulation. If None, a sensor
+
+    Returns:
+        An xarray Dataset containing the simulation results, with variables corresponding to the requested sensor outputs and
+        coordinates corresponding to the spatial dimensions of the simulation. If return_kwave_outputs is True, also returns a dictionary containing the raw outputs from k-wave. If return_kwave_inputs is True, also returns a dictionary containing the inputs to k-wave.
+
+            """
     from kwave.kspaceFirstOrder3D import kspaceFirstOrder3D
     from kwave.options.simulation_execution_options import SimulationExecutionOptions
     from kwave.options.simulation_options import SimulationOptions
@@ -111,10 +145,7 @@ def run_simulation(arr: xdc.Transducer,
         raise ValueError("All dimensions must have the same units")
     scl = getunitconversion(units[0], 'm')
     array_offset =[-float(coord.mean())*scl for coord in params.coords.values()]
-    karray = get_karray(arr,
-                        translation=array_offset,
-                        bli_tolerance=bli_tolerance,
-                        upsampling_rate=upsampling_rate)
+
     medium = get_medium(params, ref_values_only=ref_values_only)
     if _sensor is not None:
         sensor = _sensor
@@ -126,7 +157,31 @@ def run_simulation(arr: xdc.Transducer,
         source = _source
     else:
         source_mat = arr.calc_output(input_signal, kgrid.dt, delays, apod)
-        source = get_source(kgrid, karray, source_mat)
+    if arr.crosstalk_frac != 0:
+        # Simulate crosstalk by adding additional elements to the array for each element that
+        #   is within the crosstalk distance, with the signal scaled by the crosstalk fraction.
+        #   This is a simple model of crosstalk and may not capture all the complexities of real
+        #   crosstalk, but it can be useful for testing and simulation purposes.
+        crosstalk_arr = arr.copy()
+        crosstalk_mat = source_mat
+        positions = arr.get_positions(units="m")
+        for src_idx in range(arr.numelements()):
+            for dst_idx in range(arr.numelements()):
+                if src_idx == dst_idx:
+                    continue
+                src_pos = np.array(positions[src_idx])
+                dst_pos = np.array(positions[dst_idx])
+                dist = np.linalg.norm(src_pos - dst_pos)
+                if dist <= arr.crosstalk_dist:
+                    crosstalk_arr.elements += [arr.elements[dst_idx].copy()]
+                    crosstalk_mat = np.vstack((crosstalk_mat, arr.crosstalk_frac*source_mat[src_idx,:]))
+        arr = crosstalk_arr
+        source_mat = crosstalk_mat
+    karray = get_karray(arr,
+                        translation=array_offset,
+                        bli_tolerance=bli_tolerance,
+                        upsampling_rate=upsampling_rate)
+    source = get_source(kgrid, karray, source_mat)
     logging.info("Running simulation")
     simulation_options = SimulationOptions(
                             pml_auto=True,

src/openlifu/sim/sim_setup.py

@@ -45,7 +45,7 @@ class SimSetup(DictMixin):
     c0: Annotated[float, OpenLIFUFieldData("Speed of Sound (m/s)", "Reference speed of sound for converting distance to time")] = 1500.0
     """Reference speed of sound for converting distance to time"""
 
-    cfl: Annotated[float, OpenLIFUFieldData("CFL number", "Courant-Friedrichs-Lewy number")] = 0.5
+    cfl: Annotated[float, OpenLIFUFieldData("CFL number", "Courant-Friedrichs-Lewy number")] = 0.3
     """Courant-Friedrichs-Lewy number"""
 
     options: Annotated[dict[str, str], OpenLIFUFieldData("Simulation options", "Additional simulation options")] = field(default_factory=dict)

src/openlifu/xdc/transducer.py

@@ -58,6 +58,12 @@ class Transducer:
     sensitivity: Annotated[float | None, OpenLIFUFieldData("Sensitivity", "Sensitivity of the element (Pa/V)")] = None
     """Sensitivity of the element (Pa/V)"""
 
+    crosstalk_frac: Annotated[float, OpenLIFUFieldData("Crosstalk fraction", "Fraction of the signal that leaks into other elements due to crosstalk")] = 0.0
+    """Fraction of the signal that leaks into other elements due to crosstalk"""
+
+    crosstalk_dist: Annotated[float, OpenLIFUFieldData("Crosstalk distance", "Distance within which elements experience crosstalk")] = 0.0
+    """Distance within which elements experience crosstalk"""
+
     impulse_response: Annotated[np.ndarray | None, OpenLIFUFieldData("Impulse response", "Impulse response of the element")] = None
     """Impulse response of the element, can be a single value or an array of values. If an array, `impulse_dt` must be set to the time step of the impulse response. Is convolved with the input signal."""
 

tests/test_protocol.py

@@ -3,12 +3,9 @@
 import logging
 from pathlib import Path
 
-import numpy as np
 import pytest
-import xarray as xa
-from pytest_mock import MockerFixture
 
-from openlifu import Point, Protocol, Transducer
+from openlifu import Protocol, Transducer
 from openlifu.bf.focal_patterns import Wheel
 from openlifu.db import Session
 from openlifu.plan.protocol import OnPulseMismatchAction
@@ -109,46 +106,3 @@ def test_fix_pulse_mismatch(
             assert example_protocol.sequence.pulse_count == 2*num_foci
         elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
             assert example_protocol.sequence.pulse_count == num_foci
-
-@pytest.mark.parametrize("use_gpu", [True, False, None])
-@pytest.mark.parametrize("gpu_is_available", [True, False])
-def test_calc_solution_use_gpu(
-    mocker:MockerFixture,
-    example_protocol:Protocol,
-    example_transducer:Transducer,
-    use_gpu:bool | None,
-    gpu_is_available:bool,
-):
-    """Test that the correct value of use_gpu is passed to the simulation runner"""
-    example_simulation_output = xa.Dataset(
-        {
-            'p_min': xa.DataArray(data=np.empty((3, 2, 3)), dims=["x", "y", "z"], attrs={'units': "Pa"}),
-            'p_max': xa.DataArray(data=np.empty((3, 2, 3)),dims=["x", "y", "z"],attrs={'units': "Pa"}),
-            'intensity': xa.DataArray(data=np.empty((3, 2, 3)),dims=["x", "y", "z"],attrs={'units': "W/cm^2"}),
-        },
-        coords={
-            'x': xa.DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
-            'y': xa.DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "m"}),
-            'z': xa.DataArray(dims=["z"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
-        },
-    )
-    mocker.patch(
-        "openlifu.plan.protocol.gpu_available",
-        return_value = gpu_is_available,
-    )
-    run_simulation_mock = mocker.patch(
-        "openlifu.plan.solution.run_simulation",
-        return_value = example_simulation_output
-    )
-    example_protocol.calc_solution(
-        target = Point(),
-        transducer = example_transducer,
-        simulate = True,
-        scale = False,
-        use_gpu=use_gpu,
-    )
-    args, kwargs = run_simulation_mock.call_args
-    if use_gpu is None:
-        assert kwargs['gpu'] == gpu_is_available
-    else:
-        assert kwargs['gpu'] == use_gpu

tests/test_solution.py

@@ -8,6 +8,7 @@
 import pytest
 import xarray as xa
 from helpers import dataclasses_are_equal
+from pytest_mock import MockerFixture
 
 from openlifu import Point, Pulse, Sequence, Solution, Transducer
 from openlifu.bf.focal_patterns import SinglePoint
@@ -318,3 +319,60 @@ def test_solution_analyze_ratios(example_solution: Solution):
     assert analysis_case7.mainlobe_isppa_Wcm2[0] == 0.0
     assert analysis_case7.sidelobe_isppa_Wcm2[0] == 0.0
     assert np.isnan(analysis_case7.sidelobe_to_mainlobe_intensity_ratio[0]) # 0 / 0 == nan
+
+
+@pytest.mark.parametrize("force_cpu", [True, False])
+@pytest.mark.parametrize("gpu_is_available", [True, False])
+def test_solution_simulate_gpu_handling(
+    mocker: MockerFixture,
+    example_solution: Solution,
+    force_cpu: bool,
+    gpu_is_available: bool,
+):
+    """Test that the correct GPU flag is passed to run_simulation in solution.simulate"""
+    example_simulation_output = xa.Dataset(
+        {
+            'p_min': xa.DataArray(data=np.empty((3, 2, 3)), dims=["x", "y", "z"], attrs={'units': "Pa"}),
+            'p_max': xa.DataArray(data=np.empty((3, 2, 3)), dims=["x", "y", "z"], attrs={'units': "Pa"}),
+            'intensity': xa.DataArray(data=np.empty((3, 2, 3)), dims=["x", "y", "z"], attrs={'units': "W/cm^2"}),
+        },
+        coords={
+            'x': xa.DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
+            'y': xa.DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "m"}),
+            'z': xa.DataArray(dims=["z"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
+        },
+    )
+    mocker.patch(
+        "openlifu.plan.solution.gpu_available",
+        return_value=gpu_is_available,
+    )
+    run_simulation_mock = mocker.patch(
+        "openlifu.plan.solution.run_simulation",
+        return_value=example_simulation_output
+    )
+
+    # Create params dataset for simulation
+    params = xa.Dataset(
+        {
+            'density': xa.DataArray(data=np.ones((3, 2, 3)), dims=["x", "y", "z"], attrs={'units': "kg/m^3"}),
+            'sound_speed': xa.DataArray(data=np.ones((3, 2, 3)) * 1500, dims=["x", "y", "z"], attrs={'units': "m/s"}),
+        },
+        coords={
+            'x': xa.DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
+            'y': xa.DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "m"}),
+            'z': xa.DataArray(dims=["z"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
+        },
+    )
+
+    example_solution.simulate(
+        params=params,
+        _force_cpu=force_cpu,
+    )
+
+    args, kwargs = run_simulation_mock.call_args
+    # If force_cpu is True, GPU should always be False
+    # Otherwise, GPU should match gpu_available()
+    if force_cpu:
+        assert not(kwargs['gpu'])
+    else:
+        assert kwargs['gpu'] == gpu_is_available