add crosstalk simulation

Author: peterhollender

examples/tutorials/03_Solution_Generation_and_Analysis.py

@@ -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.
 
@@ -291,9 +290,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 +342,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,9 +139,21 @@ 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:

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/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."""