diff --git a/examples/at_circular_piston_3D/at_circular_piston_3D.py b/examples/at_circular_piston_3D/at_circular_piston_3D.py index 4d7bc139..807507c6 100644 --- a/examples/at_circular_piston_3D/at_circular_piston_3D.py +++ b/examples/at_circular_piston_3D/at_circular_piston_3D.py @@ -48,20 +48,9 @@ # GRID # -------------------- -# calculate the grid spacing based on the PPW and F0 -dx: float = c0 / (ppw * source_f0) # [m] - -# compute the size of the grid -# is round_even needed? -Nx: int = round_even(axial_size / dx) -Ny: int = round_even(lateral_size / dx) -Nz: int = Ny - -grid_size_points = Vector([Nx, Ny, Nz]) -grid_spacing_meters = Vector([dx, dx, dx]) - -# create the k-space grid -kgrid = kWaveGrid(grid_size_points, grid_spacing_meters) +kgrid = kWaveGrid.from_domain( + dimensions=np.array([axial_size, lateral_size, lateral_size]), frequency=source_f0, sound_speed_min=c0, points_per_wavelength=ppw +) # compute points per temporal period ppp: int = round(ppw / cfl) @@ -113,8 +102,8 @@ sensor = kSensor() # set sensor mask to record central plane, not including the source point -sensor.mask = np.zeros((Nx, Ny, Nz), dtype=bool) -sensor.mask[1:, :, Nz // 2] = True +sensor.mask = np.zeros(kgrid.N, dtype=bool) +sensor.mask[1:, :, kgrid.Nz // 2] = True # record the pressure sensor.record = ["p"] @@ -143,10 +132,10 @@ amp, _, _ = extract_amp_phase(sensor_data["p"].T, 1.0 / kgrid.dt, source_f0, dim=1, fft_padding=1, window="Rectangular") # reshape data -amp = np.reshape(amp, (Nx - 1, Ny), order="F") +amp = np.reshape(amp, (kgrid.Nx - 1, kgrid.Ny), order="F") # extract pressure on axis -amp_on_axis = amp[:, Ny // 2] +amp_on_axis = amp[:, kgrid.Ny // 2] # define axis vectors for plotting x_vec = kgrid.x_vec[1:, :] - kgrid.x_vec[0] @@ -161,7 +150,7 @@ # define radius and axis a: float = source_diam / 2.0 -x_max: float = (Nx - 1) * dx +x_max: float = (kgrid.Nx - 1) * kgrid.dx delta_x: float = x_max / 10000.0 x_ref: float = np.arange(0.0, x_max + delta_x, delta_x, dtype=float) @@ -194,7 +183,7 @@ # plot the source mask (pml is outside the grid in this example) fig2, ax2 = plt.subplots(1, 1) ax2.pcolormesh( - 1e3 * np.squeeze(kgrid.y_vec), 1e3 * np.squeeze(kgrid.x_vec), np.flip(source.p_mask[:, :, Nz // 2], axis=0), shading="nearest" + 1e3 * np.squeeze(kgrid.y_vec), 1e3 * np.squeeze(kgrid.x_vec), np.flip(source.p_mask[:, :, kgrid.Nz // 2], axis=0), shading="nearest" ) ax2.set(xlabel="y [mm]", ylabel="x [mm]", title="Source Mask") diff --git a/examples/at_focused_bowl_3D/at_focused_bowl_3D.py b/examples/at_focused_bowl_3D/at_focused_bowl_3D.py index 83909d7e..2b273ff5 100644 --- a/examples/at_focused_bowl_3D/at_focused_bowl_3D.py +++ b/examples/at_focused_bowl_3D/at_focused_bowl_3D.py @@ -57,18 +57,9 @@ # -------------------- # calculate the grid spacing based on the PPW and F0 -dx: float = c0 / (ppw * source_f0) # [m] - -# compute the size of the grid -Nx: int = round_even(axial_size / dx) + source_x_offset -Ny: int = round_even(lateral_size / dx) -Nz: int = Ny - -grid_size_points = Vector([Nx, Ny, Nz]) -grid_spacing_meters = Vector([dx, dx, dx]) - -# create the k-space grid -kgrid = kWaveGrid(grid_size_points, grid_spacing_meters) +kgrid = kWaveGrid.from_domain( + dimensions=[axial_size, lateral_size, lateral_size], frequency=source_f0, sound_speed_min=c0, points_per_wavelength=ppw +) # compute points per temporal period ppp: int = round(ppw / cfl) @@ -124,8 +115,8 @@ sensor = kSensor() # set sensor mask to record central plane, not including the source point -sensor.mask = np.zeros((Nx, Ny, Nz), dtype=bool) -sensor.mask[(source_x_offset + 1) : -1, :, Nz // 2] = True +sensor.mask = np.zeros(kgrid.N, dtype=bool) +sensor.mask[(source_x_offset + 1) : -1, :, kgrid.Nz // 2] = True # record the pressure sensor.record = ["p"] @@ -155,10 +146,10 @@ amp, _, _ = extract_amp_phase(sensor_data["p"].T, 1.0 / kgrid.dt, source_f0, dim=1, fft_padding=1, window="Rectangular") # reshape data -amp = np.reshape(amp, (Nx - (source_x_offset + 2), Ny), order="F") +amp = np.reshape(amp, (kgrid.Nx - (source_x_offset + 2), kgrid.Ny), order="F") # extract pressure on axis -amp_on_axis = amp[:, Ny // 2] +amp_on_axis = amp[:, kgrid.Ny // 2] # define axis vectors for plotting x_vec = kgrid.x_vec[(source_x_offset + 1) : -1, :] - kgrid.x_vec[source_x_offset] @@ -171,8 +162,9 @@ # calculate the wavenumber knumber = 2.0 * np.pi * source_f0 / c0 -# define axis -x_max = Nx * dx +# define the position vectors +x_ref = np.arange(0.0, kgrid.x_size, kgrid.dx) +x_max = (kgrid.Nx - 1) * kgrid.dx delta_x = x_max / 10000.0 x_ref = np.arange(0.0, x_max + delta_x, delta_x) @@ -210,14 +202,14 @@ ax2a.pcolormesh( 1e3 * np.squeeze(kgrid.y_vec), 1e3 * np.squeeze(kgrid.x_vec), - np.flip(source.p_mask[:, :, Nz // 2], axis=0), + np.flip(source.p_mask[:, :, kgrid.Nz // 2], axis=0), shading="nearest", ) ax2a.set(xlabel="y [mm]", ylabel="x [mm]", title="Source Mask") ax2b.pcolormesh( 1e3 * np.squeeze(kgrid.y_vec), 1e3 * np.squeeze(kgrid.x_vec), - np.flip(grid_weights[:, :, Nz // 2], axis=0), + np.flip(grid_weights[:, :, kgrid.Nz // 2], axis=0), shading="nearest", ) ax2b.set(xlabel="y [mm]", ylabel="x [mm]", title="Off-Grid Source Weights") diff --git a/kwave/kgrid.py b/kwave/kgrid.py index 3476c7ab..df895b2d 100644 --- a/kwave/kgrid.py +++ b/kwave/kgrid.py @@ -37,6 +37,8 @@ def __init__(self, N, spacing): N, spacing = np.atleast_1d(N), np.atleast_1d(spacing) # if inputs are lists assert N.ndim == 1 and spacing.ndim == 1 # ensure no multidimensional lists assert (1 <= N.size <= 3) and (1 <= spacing.size <= 3) # ensure valid dimensionality + if spacing.size == 1: + spacing = spacing * np.ones(N.size) assert N.size == spacing.size, "Size list N and spacing list do not have the same size." self.N = N.astype(int) #: grid size in each dimension [grid points] @@ -699,3 +701,86 @@ def k_dtt(self, dtt_type): # Not tested for correctness! # define product of implied period M = Mx * My * Mz return k, M + + @classmethod + def from_geometry(cls, dimensions, min_element_width, points_per_wavelength=10): + """ + Create a kWaveGrid based on domain dimensions and the smallest resolvable geometry element. + + Args: + dimensions: List or array of physical domain sizes [m] + min_element_width: Width of the smallest resolvable geometry element [m] + points_per_wavelength: Number of points per wavelength (default=10) + + Returns: + kWaveGrid instance with appropriate grid size and spacing + """ + # Validate input parameters + dimensions = np.atleast_1d(dimensions) + if not np.all(dimensions > 0): + raise ValueError("Domain dimensions must be positive") + if not min_element_width > 0: + raise ValueError("Minimum element width must be positive") + + # Calculate grid spacing based on minimum element width + # Ensure at least points_per_wavelength points across the smallest element + grid_spacing = min_element_width / points_per_wavelength + + # Create a list of grid spacings with the same length as domain_size + grid_spacing_list = [grid_spacing] * dimensions.size + + # Calculate grid size + N = np.ceil(dimensions / grid_spacing).astype(int) + + # Create grid instance + grid = cls(N=N, spacing=grid_spacing_list) + + # Note: Time parameters are left as "auto" + # The user can set them later using makeTime method + + return grid + + @classmethod + def from_domain(cls, dimensions, frequency, sound_speed_min, sound_speed_max=None, points_per_wavelength=10, cfl=None): + """ + Create a kWaveGrid based on physical dimensions and acoustic properties. + + Args: + dimensions: List or array of physical domain sizes [m] + frequency: Transmit frequency [Hz] + sound_speed_min: Minimum sound speed in the medium [m/s] + sound_speed_max: Maximum sound speed in the medium [m/s] (default=sound_speed_min) + points_per_wavelength: Number of points per wavelength (default=10) + cfl: CFL number (default=cls.CFL_DEFAULT) + + Returns: + kWaveGrid instance with appropriate grid size, spacing, and time parameters + """ + # Validate input parameters + dimensions = np.atleast_1d(dimensions) + if not np.all(dimensions > 0): + raise ValueError("Dimensions must be positive") + if not frequency > 0: + raise ValueError("Frequency must be positive") + if not sound_speed_min > 0: + raise ValueError("Sound speed must be positive") + + # Use sound_speed_min for sound_speed_max if not provided + if sound_speed_max is None: + sound_speed_max = sound_speed_min + if not sound_speed_max > 0: + raise ValueError("Sound speed must be positive") + + # Calculate wavelength + wavelength = sound_speed_min / frequency + + # Calculate grid spacing + grid_spacing = wavelength / points_per_wavelength + + # Calculate grid size + N = np.ceil(dimensions / grid_spacing).astype(int) + + # Create grid instance + grid = cls(N=N, spacing=grid_spacing) + + return grid diff --git a/tests/test_kgrid.py b/tests/test_kgrid.py new file mode 100644 index 00000000..44480b51 --- /dev/null +++ b/tests/test_kgrid.py @@ -0,0 +1,152 @@ +import numpy as np +import pytest + +from kwave.kgrid import kWaveGrid + + +def test_from_geometry(): + # Test 1D grid + dimensions = [0.1] # 10cm domain + min_element_width = 0.001 # 1mm minimum element + grid = kWaveGrid.from_geometry(dimensions, min_element_width) + assert grid.dim == 1 + assert grid.dx == 0.0001 # 0.1mm spacing (min_element_width/10) + assert grid.Nx == 1000 # 10cm/0.1mm + + # Test 2D grid + dimensions = [0.1, 0.2] # 10cm x 20cm domain + min_element_width = 0.001 # 1mm minimum element + grid = kWaveGrid.from_geometry(dimensions, min_element_width) + assert grid.dim == 2 + assert grid.dx == 0.0001 # 0.1mm spacing + assert grid.dy == 0.0001 # 0.1mm spacing + assert grid.Nx == 1000 # 10cm/0.1mm + assert grid.Ny == 2000 # 20cm/0.1mm + + # Test 3D grid + dimensions = [0.1, 0.2, 0.3] # 10cm x 20cm x 30cm domain + min_element_width = 0.001 # 1mm minimum element + grid = kWaveGrid.from_geometry(dimensions, min_element_width) + assert grid.dim == 3 + assert grid.dx == 0.0001 # 0.1mm spacing + assert grid.dy == 0.0001 # 0.1mm spacing + assert grid.dz == 0.0001 # 0.1mm spacing + assert grid.Nx == 1000 # 10cm/0.1mm + assert grid.Ny == 2000 # 20cm/0.1mm + assert grid.Nz == 3000 # 30cm/0.1mm + + # Test custom points_per_wavelength + dimensions = [0.1] # 10cm domain + min_element_width = 0.001 # 1mm minimum element + points_per_wavelength = 20 # Double the default + grid = kWaveGrid.from_geometry(dimensions, min_element_width, points_per_wavelength=points_per_wavelength) + assert grid.dx == 0.00005 # 0.05mm spacing + assert grid.Nx == 2000 # 10cm/0.05mm + + # Test error cases + with pytest.raises(ValueError): + kWaveGrid.from_geometry([-0.1], 0.01) # Negative dimension + with pytest.raises(ValueError): + kWaveGrid.from_geometry([0.1], -0.01) # Negative element width + + +def test_from_domain(): + # Test 1D grid based on at_circular_piston_3D example + dimensions = [0.032] # 32mm domain + frequency = 1e6 # 1MHz + sound_speed = 1500 # 1500 m/s + ppw = 3 # points per wavelength + grid = kWaveGrid.from_domain(dimensions, frequency, sound_speed, points_per_wavelength=ppw) + + wavelength = sound_speed / frequency # 1.5mm + expected_spacing = wavelength / ppw # 0.5mm + expected_points = int(np.ceil(dimensions[0] / expected_spacing)) + + assert grid.dim == 1 + assert np.isclose(grid.dx, expected_spacing) + assert grid.Nx == expected_points + + # Test 2D grid with different sound speeds + dimensions = [0.032, 0.023] # 32mm x 23mm domain (from example) + frequency = 1e6 # 1MHz + sound_speed_min = 1500 # 1500 m/s + sound_speed_max = 2000 # 2000 m/s + grid = kWaveGrid.from_domain(dimensions, frequency, sound_speed_min, sound_speed_max, points_per_wavelength=ppw) + + wavelength = sound_speed_min / frequency # 1.5mm + expected_spacing = wavelength / ppw # 0.5mm + expected_points_x = int(np.ceil(dimensions[0] / expected_spacing)) + expected_points_y = int(np.ceil(dimensions[1] / expected_spacing)) + + assert grid.dim == 2 + assert np.isclose(grid.dx, expected_spacing) + assert np.isclose(grid.dy, expected_spacing) + assert grid.Nx == expected_points_x + assert grid.Ny == expected_points_y + + # Test error cases + with pytest.raises(ValueError): + kWaveGrid.from_domain([-0.1], 1e6, 1500) # Negative dimension + with pytest.raises(ValueError): + kWaveGrid.from_domain([0.1], -1e6, 1500) # Negative frequency + with pytest.raises(ValueError): + kWaveGrid.from_domain([0.1], 1e6, -1500) # Negative sound speed + + +def test_total_grid_points(): + # Test 1D grid + grid = kWaveGrid([10], [0.1]) + assert grid.total_grid_points == 10 + + # Test 2D grid + grid = kWaveGrid([10, 20], [0.1, 0.1]) + assert grid.total_grid_points == 200 + + # Test 3D grid + grid = kWaveGrid([10, 20, 30], [0.1, 0.1, 0.1]) + assert grid.total_grid_points == 6000 + + +def test_kx_ky_kz_properties(): + # Test 1D grid + grid = kWaveGrid([10], [0.1]) + assert np.array_equal(grid.kx, grid.k_vec.x) + assert np.isnan(grid.ky) + assert np.isnan(grid.kz) + + # Test 2D grid + grid = kWaveGrid([10, 20], [0.1, 0.1]) + expected_kx = np.tile(grid.k_vec.x, (1, 20)) + expected_ky = np.tile(grid.k_vec.y.T, (10, 1)) + assert np.array_equal(grid.kx, expected_kx) + assert np.array_equal(grid.ky, expected_ky) + assert np.isnan(grid.kz) + + # Test 3D grid + grid = kWaveGrid([10, 20, 30], [0.1, 0.1, 0.1]) + expected_kx = np.tile(grid.k_vec.x[:, :, None], (1, 20, 30)) + expected_ky = np.tile(grid.k_vec.y[None, :, :], (10, 1, 30)) + expected_kz = np.tile(grid.k_vec.z.T[None, :, :], (10, 20, 1)) + assert np.array_equal(grid.kx, expected_kx) + assert np.array_equal(grid.ky, expected_ky) + assert np.array_equal(grid.kz, expected_kz) + + +def test_size_properties(): + # Test 1D grid + grid = kWaveGrid([10], [0.1]) + assert grid.x_size == 1.0 # 10 * 0.1 + assert grid.y_size == 0.0 # Not applicable in 1D + assert grid.z_size == 0.0 # Not applicable in 1D + + # Test 2D grid + grid = kWaveGrid([10, 20], [0.1, 0.1]) + assert grid.x_size == 1.0 # 10 * 0.1 + assert grid.y_size == 2.0 # 20 * 0.1 + assert grid.z_size == 0.0 # Not applicable in 2D + + # Test 3D grid + grid = kWaveGrid([10, 20, 30], [0.1, 0.1, 0.1]) + assert grid.x_size == 1.0 # 10 * 0.1 + assert grid.y_size == 2.0 # 20 * 0.1 + assert grid.z_size == 3.0 # 30 * 0.1