Coverage for kwave/utils/pmlutils.py: 6%
52 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-24 11:55 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-24 11:55 -0700
1import numpy as np
4def get_pml(Nx, dx, dt, c, pml_size, pml_alpha, staggered, dimension, axisymmetric=False):
5 """
6 getPML returns a 1D perfectly matched layer variable based on the given size and absorption coefficient.
7 Args:
8 Nx:
9 dx:
10 dt:
11 c:
12 pml_size:
13 pml_alpha:
14 staggered:
15 dimension:
16 axisymmetric:
18 Returns:
20 """
21 # define x-axis
22 Nx = int(Nx)
23 pml_size = int(pml_size)
24 x = np.arange(1, pml_size + 1)
26 # create absorption profile
27 if staggered:
29 # calculate the varying components of the pml using a staggered grid
30 pml_left = pml_alpha * (c / dx) * (( ((x + 0.5) - pml_size - 1) / (0 - pml_size) ) ** 4)
31 pml_right = pml_alpha * (c / dx) * (( (x + 0.5) / pml_size ) ** 4)
33 else:
35 # calculate the varying components of the pml using a regular grid
36 pml_left = pml_alpha * (c / dx) * (( (x - pml_size - 1) / (0 - pml_size) ) ** 4)
37 pml_right = pml_alpha * (c / dx) * (( x / pml_size ) ** 4)
39 # exponentiation
40 pml_left = np.exp(-pml_left * dt / 2)
41 pml_right = np.exp(-pml_right * dt / 2)
43 # add the components of the pml to the total function, not adding the axial
44 # side of the radial PML if axisymmetric
45 pml = np.ones((1, Nx))
46 if not axisymmetric:
47 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
58 # ------------
59 # Other forms:
60 # ------------
61 # Use this to include an extra unity point:
62 # pml_left = pml_alpha*(c/dx)* ( (x - pml_size) ./ (1 - pml_size) ).^2;
63 # pml_right = pml_alpha*(c/dx)* ( (x - 1) ./ (pml_size - 1) ).^2;
64 # Staggered grid equivalents:
65 # pml_left = pml_alpha*(c/dx)* ( ((x + 0.5) - pml_size) ./ (1 - pml_size) ).^2;
66 # pml_right = pml_alpha*(c/dx)* ( ((x + 0.5) - 1) ./ (pml_size - 1) ).^2;
69def getOptimalPMLSize(grid_size, pml_range=None, axisymmetric=None):
70 """
71 getOptimalPMLSize finds the size of the perfectly matched layer (PML)
72 that gives an overall grid size with the smallest prime factors when
73 using the first-order simulation functions in k-Wave with the
74 optional input 'PMLInside', false. Choosing grid sizes with small
75 prime factors can have a significant impact on the computational
76 speed, as the code computes spatial gradients using the fast Fourier
77 transform (FFT).
78 Args:
79 grid_size: Grid size defined as a one (1D), two (2D), or three (3D) element vector. Alternatively, can be an
80 object of the kWaveGrid class defining the Cartesian and k-space grid fields.
81 pml_range: Two element vector specifying the minimum and maximum PML size (default = [10, 40]).
82 axisymmetric: If using the axisymmetric code, string specifying the radial symmetry. Allowable inputs are 'WSWA'
83 and 'WSWS' (default = ''). This is important as the axisymmetric code only applies to the
84 PML to the outside edge in the radial dimension.
86 Returns:
87 pml_sz: PML size that gives the overall grid with the smallest prime factors.
89 """
90 # check if grid size is given as kgrid, and extract grid size
91 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(('WSWA', 'WSWS')), \
123 "Optional input axisymmetric must be set to ''WSWA'' or ''WSWS''."
125 # check for correct dimensions
126 if isinstance(axisymmetric, str) and grid_dim != 2:
127 raise ValueError('Optional input axisymmetric is only valid for 2D grid sizes.')
129 # create array of PML values to search
130 pml_size = np.arange(pml_range[0], pml_range[1] + 1)
132 # extract the largest prime factor for each dimension for each pml size
133 facs = np.zeros((grid_dim, len(pml_size)))
134 from kwave.utils import largest_prime_factor
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