Add class functions to simplify kGrid creation

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")
 

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")

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

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