Coverage for kwave/utils/pmlutils.py: 6%

1 import numpy as np

2

3

4 def 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:

17

18 Returns:

19

20 """

21 # define x-axis

22 Nx = int(Nx)

23 pml_size = int(pml_size)

24 x = np.arange(1, pml_size + 1)

25

26 # create absorption profile

27 if staggered:

28

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)

32

33 else:

34

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)

38

39 # exponentiation

40 pml_left = np.exp(-pml_left * dt / 2)

41 pml_right = np.exp(-pml_right * dt / 2)

42

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

48

49 pml[:, Nx - pml_size:] = pml_right

50

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 # ------------

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;

67

68

69 def 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.

85

86 Returns:

87 pml_sz: PML size that gives the overall grid with the smallest prime factors.

88

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

94

95 # assign grid size

96 grid_dim = len(grid_size)

97

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.'

100

101 # check for pml_range input

102 if pml_range is None:

103 pml_range = [10, 40]

104

105 # force integer

106 pml_range = np.round(pml_range).astype(int)

107

108 # check for positive values

109 assert np.all(pml_range >= 0), 'Optional input pml_range must be positive.'

110

111 # check for correct length

112 assert len(pml_range) == 2, 'Optional input pml_range must be a two element vector.'

113

114 # check for monotonic

115 assert pml_range[1] > pml_range[0], 'The second value for pml_range must be greater than the first.'

116

117 # check for axisymmetric input

118 if axisymmetric is None:

119 axisymmetric = False

120

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''."

124

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.')

128

129 # create array of PML values to search

130 pml_size = np.arange(pml_range[0], pml_range[1] + 1)

131

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

144

145 # get best dimension size

146 ind_opt = np.argmin(facs, 1)

147

148 # assign output

149 pml_sz = pml_size[ind_opt]

150

151 return pml_sz