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

1import numpy as np 

2 

3 

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: 

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 

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. 

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