Update simulation_options.py - ensure that a 2d pml can be set (#328)

kwave/options/simulation_options.py

@@ -254,26 +254,21 @@ def option_factory(kgrid: "kWaveGrid", options: SimulationOptions):
                     raise ValueError("Optional input ''pml_size'' must be a single numerical value.")
 
         if kgrid.dim == 1:
-            options.pml_x_alpha = 2
             options.pml_x_size = options.pml_size if options.pml_size else 20
             options.plot_scale = [-1.1, 1.1]
         elif kgrid.dim == 2:
-            options.pml_x_alpha = 2
-            options.pml_y_alpha = options.pml_x_alpha
-            if options.pml_size is None:
-                options.pml_x_size = 20
-                options.pml_y_size = 20
-            else:
-                options.pml_x_size = options.pml_size[0]
-                options.pml_y_size = options.pml_x_size
+            if (options.pml_size is not None): 
+                if (len(options.pml_size) == kgrid.dim):
+                    options.pml_x_size, options.pml_y_size = np.asarray(options.pml_size, dtype=int).ravel()
+                else:
+                    options.pml_x_size, options.pml_y_size = (options.pml_size[0], options.pml_size[0])
+            else: 
+                options.pml_x_size, options.pml_y_size = (20, 20)
             options.plot_scale = [-1, 1]
         elif kgrid.dim == 3:
-            if options.pml_size is not None and len(options.pml_size) == kgrid.dim:
-                options.pml_x_size, options.pml_y_size, options.pml_z_size = options.pml_size.ravel()
+            if ((options.pml_size is not None) and (len(options.pml_size) == kgrid.dim)):
+                options.pml_x_size, options.pml_y_size, options.pml_z_size = np.asarray(options.pml_size).ravel()
             else:
-                options.pml_x_alpha = 2
-                options.pml_y_alpha = options.pml_x_alpha
-                options.pml_z_alpha = options.pml_x_alpha
                 if options.pml_size is None:
                     options.pml_x_size = 10
                     options.pml_y_size = 10

tests/test_pmlutils.py

@@ -1,4 +1,18 @@
+import pytest
+
+
+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 kspaceFirstOrder2DC, kspace_first_order_2d_gpu
+from kwave.options.simulation_execution_options import SimulationExecutionOptions
+from kwave.options.simulation_options import SimulationOptions
+from kwave.utils.mapgen import make_disc, make_cart_circle
+
 from kwave.utils.pml import get_optimal_pml_size
 
 
@@ -56,3 +70,86 @@ def test_get_optimal_pml_size_3D():
     kgrid.setTime(round(t_end / dt) + 1, dt)
 
     assert (32 == get_optimal_pml_size(kgrid)).all()
+
+def test_pml_sizes_2d():
+
+    nx = 128  # number of grid points in the x (row) direction
+    x = 128e-3  # size of the domain in the x direction [m]
+    dx = x / nx  # grid point spacing in the x direction [m]
+    ny = 132  # number of grid points in the y (column) direction
+    y = 132e-3  # size of the domain in the y direction [m]
+    dy = y / ny  # grid point spacing in the y direction [m]
+
+    grid_size = Vector([nx, ny]) # [grid points]
+    grid_spacing = Vector([dx, dy])  # [m]
+
+    # time array
+    dt = 2e-9  # [s]
+    t_end = 300e-9  # [s]
+    # create the computational grid
+    kgrid = kWaveGrid(grid_size, grid_spacing)
+    # create the time array
+    kgrid.setTime(round(t_end / dt) + 1, dt)
+
+    medium = kWaveMedium(sound_speed=1500, alpha_coeff=0.75, alpha_power=1.5)
+
+    # create initial pressure distribution using make_disc
+    disc_magnitude = 5 # [Pa]
+    disc_pos = Vector([50, 50])  # [grid points]
+    disc_radius = 8    # [grid points]
+    disc_1 = disc_magnitude * make_disc(grid_size, disc_pos, disc_radius)
+
+    disc_magnitude = 3 # [Pa]
+    disc_pos = Vector([80, 60])  # [grid points]
+    disc_radius = 5    # [grid points]
+    disc_2 = disc_magnitude * make_disc(grid_size, disc_pos, disc_radius)
+
+    source = kSource()
+    source.p0 = disc_1 + disc_2
+
+    # define a centered circular sensor
+    sensor_radius = 4e-3   # [m]
+    num_sensor_points = 50
+    sensor_mask = make_cart_circle(sensor_radius, num_sensor_points)
+    sensor = kSensor(sensor_mask)
+
+    # these tests passes a list with three entries for the sizes of the pml 
+    # but to a 2d simulation, and so raises an error.
+    with pytest.raises(ValueError):
+        simulation_options = SimulationOptions(pml_size=[18,18,18])
+        _ = kspace_first_order_2d_gpu(medium=medium,
+                                      kgrid=kgrid,
+                                      source=source,
+                                      sensor=sensor,
+                                      simulation_options=simulation_options,
+                                      execution_options=SimulationExecutionOptions())
+
+    with pytest.raises(ValueError):
+        simulation_options = SimulationOptions(pml_size=[18,18,18]) 
+        _ = kspaceFirstOrder2DC(medium=medium,
+                                kgrid=kgrid,
+                                source=source,
+                                sensor=sensor,
+                                simulation_options=simulation_options,
+                                execution_options=SimulationExecutionOptions())
+
+    pml_size: int = 19
+    simulation_options = SimulationOptions(pml_size=pml_size)
+    options = SimulationOptions.option_factory(kgrid, simulation_options)
+    assert ((options.pml_x_size == pml_size) and (options.pml_y_size == pml_size)), \
+        "pml sizes incorrect when passing int"
+
+    pml_sizes = [21, 22]
+    simulation_options = SimulationOptions(pml_size=pml_sizes)
+    options = SimulationOptions.option_factory(kgrid, simulation_options)
+    assert ((options.pml_x_size == pml_sizes[0]) and (options.pml_y_size == pml_sizes[1])), \
+        "pml sizes incorrect when passing list"
+
+    pml_sizes = None
+    pml_default: int = 20
+    simulation_options = SimulationOptions()
+    options = SimulationOptions.option_factory(kgrid, simulation_options)
+    assert ((options.pml_x_size == pml_default) and (options.pml_y_size == pml_default)), \
+        "pml sizes incorrect when not defining sizes"
+
+