Coverage for kwave/utils/pml.py: 5%
52 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-03-24 12:06 -0700
1import numpy as np
3from kwave.utils.math import largest_prime_factor
6def get_pml(
7 Nx: int, dx: float, dt: float, c: float, pml_size: int, pml_alpha: float, staggered: bool, dimension: int, axisymmetric: bool = False
8) -> np.ndarray:
9 """
10 Returns a 1D perfectly matched layer variable based on the given size and absorption coefficient.
12 This function calculates a 1D perfectly matched layer (PML) variable based on the specified size and absorption coefficient.
13 It uses the given parameters to create an absorption profile, which is then exponentiated and reshaped in the desired direction.
14 If the axisymmetric argument is set to True, the axial side of the radial PML will not be added.
16 Args:
17 Nx: The number of grid points in the x direction.
18 dx: The spacing between grid points in the x direction.
19 dt: The time step size.
20 c: The wave speed in the medium.
21 pml_size: The size of the PML layer in grid points.
22 pml_alpha: The absorption coefficient of the PML layer.
23 staggered: Whether to use a staggered grid for calculating the varying components of the PML.
24 dimension: The dimension of the PML (1, 2, or 3).
25 axisymmetric: Whether to use axisymmetry when calculating the PML. Defaults to False.
27 Returns:
28 A 1D numpy array representing the PML variable.
29 """
30 # define x-axis
31 Nx = int(Nx)
32 pml_size = int(pml_size)
33 x = np.arange(1, pml_size + 1)
35 # create absorption profile
36 if staggered:
37 pml_left = pml_alpha * (c / dx) * ((((x + 0.5) - pml_size - 1) / (0 - pml_size)) ** 4)
38 pml_right = pml_alpha * (c / dx) * (((x + 0.5) / pml_size) ** 4)
39 else:
40 pml_left = pml_alpha * (c / dx) * (((x - pml_size - 1) / (0 - pml_size)) ** 4)
41 pml_right = pml_alpha * (c / dx) * ((x / pml_size) ** 4)
43 # exponentiate and add the components of the pml to the total function
44 pml_left = np.exp(-pml_left * dt / 2)
45 pml_right = np.exp(-pml_right * dt / 2)
46 pml = np.ones((1, Nx))
47 if not axisymmetric:
48 pml[:, :pml_size] = pml_left
49 pml[:, Nx - pml_size :] = pml_right
51 # reshape the pml vector to be in the desired direction
52 if dimension == 1:
53 pml = pml.T
54 elif dimension == 3:
55 pml = np.reshape(pml, (1, 1, Nx))
56 return pml
57 # ------------
58 # Other forms:
59 # ------------
60 # Use this to include an extra unity point:
61 # pml_left = pml_alpha*(c/dx)* ( (x - pml_size) ./ (1 - pml_size) ).^2;
62 # pml_right = pml_alpha*(c/dx)* ( (x - 1) ./ (pml_size - 1) ).^2;
63 # Staggered grid equivalents:
64 # pml_left = pml_alpha*(c/dx)* ( ((x + 0.5) - pml_size) ./ (1 - pml_size) ).^2;
65 # pml_right = pml_alpha*(c/dx)* ( ((x + 0.5) - 1) ./ (pml_size - 1) ).^2;
68def get_optimal_pml_size(grid_size, pml_range=None, axisymmetric=None):
69 """
70 get_optimal_pml_size finds the size of the perfectly matched layer (PML)
71 that gives an overall grid size with the smallest prime factors when
72 using the first-order simulation functions in k-Wave with the
73 optional input 'PMLInside', false. Choosing grid sizes with small
74 prime factors can have a significant impact on the computational
75 speed, as the code computes spatial gradients using the fast Fourier
76 transform (FFT).
77 Args:
78 grid_size: Grid size defined as a one (1D), two (2D), or three (3D) element vector. Alternatively, can be an
79 object of the kWaveGrid class defining the Cartesian and k-space grid fields.
80 pml_range: Two element vector specifying the minimum and maximum PML size (default = [10, 40]).
81 axisymmetric: If using the axisymmetric code, string specifying the radial symmetry. Allowable inputs are 'WSWA'
82 and 'WSWS' (default = ''). This is important as the axisymmetric code only applies to the
83 PML to the outside edge in the radial dimension.
85 Returns:
86 PML size that gives the overall grid with the smallest prime factors.
88 """
89 # check if grid size is given as kgrid, and extract grid size
90 from kwave.kgrid import kWaveGrid
92 if isinstance(grid_size, kWaveGrid):
93 grid_size = grid_size.N
95 # assign grid size
96 grid_dim = len(grid_size)
98 # check grid size is 1, 2, or 3
99 assert 1 <= grid_dim <= 3, "Grid dimensions must be given as a 1, 2, or 3 element vector."
101 # check for pml_range input
102 if pml_range is None:
103 pml_range = [10, 40]
105 # force integer
106 pml_range = np.round(pml_range).astype(int)
108 # check for positive values
109 assert np.all(pml_range >= 0), "Optional input pml_range must be positive."
111 # check for correct length
112 assert len(pml_range) == 2, "Optional input pml_range must be a two element vector."
114 # check for monotonic
115 assert pml_range[1] > pml_range[0], "The second value for pml_range must be greater than the first."
117 # check for axisymmetric input
118 if axisymmetric is None:
119 axisymmetric = False
121 # check for correct string
122 assert not isinstance(axisymmetric, str) or axisymmetric.startswith(
123 ("WSWA", "WSWS")
124 ), "Optional input axisymmetric must be set to ''WSWA'' or ''WSWS''."
126 # check for correct dimensions
127 if isinstance(axisymmetric, str) and grid_dim != 2:
128 raise ValueError("Optional input axisymmetric is only valid for 2D grid sizes.")
130 # create array of PML values to search
131 pml_size = np.arange(pml_range[0], pml_range[1] + 1)
133 # extract the largest prime factor for each dimension for each pml size
134 facs = np.zeros((grid_dim, len(pml_size)))
135 for dim in range(0, grid_dim):
136 for index in range(0, len(pml_size)):
137 if isinstance(axisymmetric, str) and dim == 2:
138 if axisymmetric == "WSWA":
139 facs[dim, index] = largest_prime_factor((grid_size[dim] + pml_size[index]) * 4)
140 if axisymmetric == "WSWS":
141 facs[dim, index] = largest_prime_factor((grid_size[dim] + pml_size[index]) * 2 - 2)
142 else:
143 facs[dim, index] = largest_prime_factor(grid_size[dim] + 2 * pml_size[index])
145 # get best dimension size
146 ind_opt = np.argmin(facs, 1)
148 # assign output
149 pml_sz = pml_size[ind_opt]
151 return pml_sz