Example of karray source (#193)

* WIP commit with two examples

* Working source example.

 - Pressure field is transposed. (might be related to open issue #173)
 - Needs to be cleaned up.

* sensor example running

* Change from open PR for 2D sim io

* Update branch to pass CI

* Pass Ruff CI

* Minor reformatting and adding future todos

* Add testing for linear_array_transducer setup

* Fix array as sensor example

* turn off simulation in tests

* how did that get in there?

* reorder sensor data to complete example

* pass ruff

* fix signatures and docs

* passes CI?

* make parse_executable_output static method

* clean up test_linear_array.py

* remove license header

* Use vectors

* Update checks to use Iterable type

* fix errors

* update variable names to clarify code

* clean up imports and magic numbers

* clean up commented code

* delete unused code

* clean up example file after review

* update linear array example indexing

* Clean up make cart circle test

* clean up ruff error

* clean up unused import for ruff

examples/bmode_reconstruction_example.py

@@ -18,7 +18,8 @@
 from kwave.utils.dotdictionary import dotdict
 from kwave.utils.signals import tone_burst
 
-if __name__ == '__main__':
+
+def main():
     # pathname for the input and output files
     pathname = gettempdir()
     phantom_data_path = 'phantom_data.mat'
@@ -110,7 +111,7 @@
 
         # set the input settings
         input_filename = f'example_input_{scan_line_index}.h5'
-        input_file_full_path = os.path.join(pathname, input_filename)
+        input_file_full_path = os.path.join(pathname, input_filename) # noqa: F841
         # set the input settings
         simulation_options = SimulationOptions(
             pml_inside=False,
@@ -160,3 +161,7 @@
 
     logging.log(logging.INFO, "Beamforming channel data and reconstructing the image...")
     beamform(channel_data)
+
+
+if __name__ == '__main__':
+    main()

examples/example_array_as_a_sensor.py

@@ -0,0 +1,149 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from kwave.data import Vector
+from kwave.kgrid import kWaveGrid
+from kwave.kmedium import kWaveMedium
+from kwave.ksensor import kSensor
+from kwave.ksource import kSource
+from kwave.kspaceFirstOrder2D import kspace_first_order_2d_gpu
+from kwave.options.simulation_execution_options import SimulationExecutionOptions
+from kwave.options.simulation_options import SimulationOptions
+from kwave.utils.conversion import cart2grid
+from kwave.utils.kwave_array import kWaveArray
+from kwave.utils.mapgen import make_cart_circle, make_disc
+from kwave.utils.signals import reorder_binary_sensor_data
+
+
+def main():
+    # create empty array
+    karray = kWaveArray()
+
+    # define arc properties
+    radius = 100e-3  # [m]
+    diameter = 8e-3  # [m]
+    ring_radius = 50e-3  # [m]
+    num_elements = 20
+
+    # orient all elements towards the center of the grid
+    focus_pos = Vector([0, 0])  # [m]
+
+    element_pos = make_cart_circle(ring_radius, num_elements, focus_pos)
+
+    for idx in range(num_elements):
+        karray.add_arc_element(element_pos[:, idx], radius, diameter, focus_pos)
+
+    # grid properties
+    N = Vector([256, 256])
+    d = Vector([0.5e-3, 0.5e-3])
+    kgrid = kWaveGrid(N, d)
+
+    # medium properties
+    medium = kWaveMedium(sound_speed=1500)
+
+    # time array
+    kgrid.makeTime(medium.sound_speed)
+
+    source = kSource()
+    x_offset = 20  # [pixels]
+    # make a small disc in the top left of the domain
+    source.p0 = make_disc(N, Vector([N.x / 4 + x_offset, N.y / 4]), 4)
+    source.p0[99:119, 59:199] = 1
+    logical_p0 = source.p0.astype(bool)
+    sensor = kSensor()
+    sensor.mask = element_pos
+    simulation_options = SimulationOptions(
+        save_to_disk=True,
+        data_cast='single',
+    )
+
+    execution_options = SimulationExecutionOptions(is_gpu_simulation=True)
+    output = kspace_first_order_2d_gpu(kgrid, source, sensor, medium, simulation_options, execution_options)
+        # TODO (walter): This should be done by kspaceFirstOrder
+    _, _, reorder_index = cart2grid(kgrid, element_pos)
+    sensor_data_point = reorder_binary_sensor_data(output['p'].T, reorder_index=reorder_index)
+
+    # assign binary mask from karray to the source mask
+    sensor.mask = karray.get_array_binary_mask(kgrid)
+
+    output = kspace_first_order_2d_gpu(kgrid, source, sensor, medium, simulation_options, execution_options)
+    sensor_data = output['p'].T
+    combined_sensor_data = karray.combine_sensor_data(kgrid, sensor_data)
+
+    # =========================================================================
+    # VISUALIZATION
+    # =========================================================================
+
+    # create pml mask (re-use default size of 20 grid points from simulation_options)
+    pml_size = simulation_options.pml_x_size  # 20 [grid_points] 
+    pml_mask = np.zeros((N.x, N.y), dtype=bool)
+    pml_mask[:pml_size, :] = 1
+    pml_mask[:, :pml_size] = 1
+    pml_mask[-pml_size:, :] = 1
+    pml_mask[:, -pml_size:] = 1
+
+    # Plot source, sensor, and pml masks
+
+    # Assign unique values to each mask
+    sensor_val = sensor.mask * 1
+    logical_p0_val = logical_p0 * 2
+    pml_mask_val = pml_mask * 3
+
+    # Combine masks
+    combined_mask = sensor_val + logical_p0_val + pml_mask_val
+    combined_mask = np.flipud(combined_mask) 
+
+    # Define custom colormap
+    colors = [
+        (1, 1, 1),  # White (Background)
+        (233/255, 131/255, 0/255),  # Orange (Sensor) 
+        (254/255, 221/255, 92/255),  # Yellow (Sources) 
+        (0.8, 0.8, 0.8),  # Light Grey (PML Mask)
+    ]
+    cmap = plt.matplotlib.colors.ListedColormap(colors)
+
+    fig, ax = plt.subplots()
+    c = ax.pcolormesh(combined_mask, cmap=cmap, shading='auto')
+    plt.axis('image')
+
+    # Define labels for the colorbar
+    labels = {
+        0: 'None',
+        1: 'Sensor',
+        2: 'Initial pressure p0',
+        3: 'PML Mask',
+    }
+
+    colorbar = plt.colorbar(c, ticks=list(labels.keys()), ax=ax)
+    colorbar.ax.set_yticklabels(list(labels.values()))
+
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    # Plot recorded sensor data
+    fig, (ax1, ax2)= plt.subplots(ncols=1, nrows=2)
+    ax1.imshow(sensor_data_point, aspect='auto')
+    ax1.set_xlabel(r'Time [$\mu$s]')
+    ax1.set_ylabel('Detector Number')
+    ax1.set_title('Cartesian point detectors')
+
+    ax2.imshow(combined_sensor_data, aspect='auto')
+    ax2.set_xlabel(r'Time [$\mu$s]')
+    ax2.set_ylabel('Detector Number')
+    ax2.set_title('Arc detectors')
+
+    fig.subplots_adjust(hspace=0.5)
+
+    # Plot a trace from the recorded sensor data
+    fig = plt.figure()
+    plt.plot(kgrid.t_array.squeeze() * 1e6, sensor_data_point[0, :], label='Cartesian point detectors')
+    plt.plot(kgrid.t_array.squeeze() * 1e6, combined_sensor_data[0, :], label='Arc detectors')
+    plt.xlabel(r'Time [$\mu$s]')
+    plt.ylabel('Pressure [pa]')
+    plt.legend()
+
+    plt.show()
+
+
+if __name__ == '__main__':
+    main()

examples/example_array_as_a_source.py

@@ -0,0 +1,166 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import colors
+from matplotlib.animation import FuncAnimation
+
+from kwave.data import Vector
+from kwave.kgrid import kWaveGrid
+from kwave.kmedium import kWaveMedium
+from kwave.ksensor import kSensor
+from kwave.ksource import kSource
+from kwave.kspaceFirstOrder2D import kspace_first_order_2d_gpu
+from kwave.options.simulation_execution_options import SimulationExecutionOptions
+from kwave.options.simulation_options import SimulationOptions
+from kwave.utils.kwave_array import kWaveArray
+from kwave.utils.signals import tone_burst
+
+
+def main():
+    # =========================================================================
+    # DEFINE KWAVEARRAY
+    # =========================================================================
+
+    # create empty array
+    karray = kWaveArray()
+
+    # define arc properties
+    radius = 50e-3  # [m]
+    diameter = 30e-3  # [m]
+    focus_pos = [-20e-3, 0]  # [m]
+
+    # add arc-shaped element
+    elem_pos = [10e-3, -40e-3]  # [m]
+    karray.add_arc_element(elem_pos, radius, diameter, focus_pos)
+
+    # add arc-shaped element
+    elem_pos = [20e-3, 0]  # [m]
+    karray.add_arc_element(elem_pos, radius, diameter, focus_pos)
+
+    # add arc-shaped element
+    elem_pos = [10e-3, 40e-3]  # [m]
+    karray.add_arc_element(elem_pos, radius, diameter, focus_pos)
+
+    # move the array down 10 mm, and rotate by 10 degrees (this moves all the
+    # elements together)
+    karray.set_array_position([10e-3, 0], 10)
+
+    # =========================================================================
+    # DEFINE GRID PROPERTIES
+    # =========================================================================
+
+    # grid properties
+    Nx = 256
+    dx = 0.5e-3
+    Ny = 256
+    dy = 0.5e-3
+    kgrid = kWaveGrid(Vector([Nx, Ny]), Vector([dx, dy]))
+
+    # medium properties
+    medium = kWaveMedium(sound_speed=1500)
+
+    # time array
+    kgrid.makeTime(medium.sound_speed)
+
+    # =========================================================================
+    # SIMULATION
+    # =========================================================================
+
+    # assign binary mask from karray to the source mask
+    source_p_mask = karray.get_array_binary_mask(kgrid)
+
+    # set source signals, one for each physical array element
+    f1 = 100e3
+    f2 = 200e3
+    f3 = 500e3
+    sig1 = tone_burst(1 / kgrid.dt, f1, 3).squeeze()
+    sig2 = tone_burst(1 / kgrid.dt, f2, 5).squeeze()
+    sig3 = tone_burst(1 / kgrid.dt, f3, 5).squeeze()
+
+    # combine source signals into one array
+    source_signal = np.zeros((3, max(len(sig1), len(sig2))))
+    source_signal[0, :len(sig1)] = sig1
+    source_signal[1, :len(sig2)] = sig2
+    source_signal[2, :len(sig3)] = sig3
+
+    # get distributed source signals (this automatically returns a weighted
+    # source signal for each grid point that forms part of the source)
+    source_p = karray.get_distributed_source_signal(kgrid, source_signal)
+
+    simulation_options = SimulationOptions(
+        save_to_disk=True,
+        data_cast='single',
+    )
+    execution_options = SimulationExecutionOptions(is_gpu_simulation=True)
+    # run k-Wave simulation (no sensor is used for this example)
+    # TODO: I would say proper behaviour would be to return the entire pressure field if sensor is None
+    sensor = kSensor()
+    sensor.mask = np.ones((Nx, Ny), dtype=bool)
+
+    source = kSource()
+    source.p_mask = source_p_mask
+    source.p = source_p
+
+    p = kspace_first_order_2d_gpu(kgrid, source, sensor, medium, simulation_options, execution_options)
+
+    p_field = np.reshape(p['p'], (kgrid.Nt, Nx, Ny))
+    p_field = np.transpose(p_field, (0, 2, 1))
+    # =========================================================================
+    # VISUALIZATION
+    # =========================================================================
+
+    # Normalize frames based on the maximum value over all frames
+    max_value = np.max(p_field)
+    normalized_frames = p_field / max_value
+
+    # Create a custom colormap (replace 'viridis' with your preferred colormap)
+    cmap = plt.get_cmap('viridis')
+
+    # Create a figure and axis
+    fig, ax = plt.subplots()
+
+    # Create an empty image with the first normalized frame
+    image = ax.imshow(normalized_frames[0], cmap=cmap, norm=colors.Normalize(vmin=0, vmax=1))
+
+    # Function to update the image for each frame
+    def update(frame):
+        image.set_data(normalized_frames[frame])
+        ax.set_title(f'Frame {frame + 1}/{kgrid.Nt}')
+        return [image]
+
+    # Create the animation
+    ani = FuncAnimation(fig, update, frames=kgrid.Nt, interval=100)  # Adjust interval as needed (in milliseconds)
+
+    # Save the animation as a video file (e.g., MP4)
+    video_filename = 'output_video1.mp4'
+    ani.save('/tmp/' + video_filename, writer='ffmpeg', fps=30)  # Adjust FPS as needed
+
+    # Show the animation (optional)
+    plt.show()
+
+    # Show the animation (optional)
+    plt.show()
+
+    # create pml mask (default size in 2D is 20 grid points)
+    pml_size = 20
+    pml_mask = np.zeros((Nx, Ny), dtype=bool)
+    pml_mask[:pml_size, :] = 1
+    pml_mask[:, :pml_size] = 1
+    pml_mask[-pml_size:, :] = 1
+    pml_mask[:, -pml_size:] = 1
+
+    # plot source and pml masks
+    plt.figure()
+    plt.imshow(np.logical_not(np.squeeze(source.p_mask | pml_mask)), aspect='auto', cmap='gray')
+    plt.xlabel('x-position [m]')
+    plt.ylabel('y-position [m]')
+    plt.title('Source and PML Masks')
+    plt.show()
+
+    # overlay the physical source positions
+    plt.figure()
+    # TODO: missing karray.plot_array(show=True)
+    # karray.plot_array(show=True)
+
+
+if __name__ == "__main__":
+    main()