Coverage for kwave/kWaveSimulation_helper/expand_grid_matrices.py: 14%

121 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-24 11:55 -0700

1from kwave import kWaveGrid, kWaveMedium, SimulationOptions, NotATransducer 

2from kwave.data import Array 

3from kwave.utils import matlab_find, expand_matrix, get_smallest_possible_type 

4from kwave.utils import dotdict 

5import numpy as np 

6 

7 

8def expand_grid_matrices( 

9 kgrid: kWaveGrid, medium: kWaveMedium, source, sensor, opt: SimulationOptions, 

10 values: dotdict, flags: dotdict): 

11 # update command line status 

12 print(' expanding computational grid...') 

13 

14 ##################### 

15 # Grab values 

16 ##################### 

17 

18 # retaining the values for kgrid time array 

19 pml_size = [opt.pml_x_size, opt.pml_y_size, opt.pml_z_size] 

20 pml_size = Array(pml_size[:kgrid.dim]) 

21 

22 p_source_pos_index = values.p_source_pos_index 

23 u_source_pos_index = values.u_source_pos_index 

24 s_source_pos_index = values.s_source_pos_index 

25 

26 is_source_sensor_same = (isinstance(sensor, NotATransducer) and sensor == source) 

27 

28 ##################### 

29 # Expand Structures 

30 ##################### 

31 

32 # expand the computational grid, replacing the original grid 

33 kgrid = expand_kgrid(kgrid, flags.axisymmetric, pml_size.numpy()) 

34 

35 expand_size = calculate_expand_size(kgrid, flags.axisymmetric, pml_size) 

36 

37 # update the data type in case adding the PML requires additional index precision 

38 total_grid_points = kgrid.total_grid_points 

39 index_data_type = get_smallest_possible_type(total_grid_points, 'uint', default='double') 

40 

41 expand_sensor(sensor, expand_size, flags.use_sensor, flags.blank_sensor) 

42 

43 # TODO why it is not self.record ? "self" 

44 record = expand_cuboid_corner_list(flags.cuboid_corners, kgrid, pml_size) 

45 

46 expand_medium(medium, expand_size) 

47 

48 p_source_pos_index, u_source_pos_index, s_source_pos_index = expand_source( 

49 source, is_source_sensor_same, flags, expand_size, index_data_type, 

50 p_source_pos_index, u_source_pos_index, s_source_pos_index 

51 ) 

52 

53 expand_directivity_angle(kgrid, sensor, expand_size, flags.use_sensor, flags.compute_directivity) 

54 

55 print_grid_size(kgrid) 

56 

57 return kgrid, index_data_type, p_source_pos_index, u_source_pos_index, s_source_pos_index 

58 

59 

60def expand_kgrid(kgrid, is_axisymmetric, pml_size): 

61 Nt_temp, dt_temp = kgrid.Nt, kgrid.dt 

62 

63 if kgrid.dim == 1: 

64 new_size = kgrid.N + 2 * pml_size 

65 elif kgrid.dim == 2: 

66 if is_axisymmetric: 

67 new_size = [kgrid.Nx + 2 * pml_size[0], kgrid.Ny + pml_size[1]] 

68 else: 

69 new_size = kgrid.N + 2 * pml_size 

70 elif kgrid.dim == 3: 

71 new_size = kgrid.N + 2 * pml_size 

72 else: 

73 raise NotImplementedError 

74 

75 kgrid = kWaveGrid(new_size, kgrid.spacing) 

76 # re-assign original time array 

77 kgrid.setTime(Nt_temp, dt_temp) 

78 

79 return kgrid 

80 

81 

82def calculate_expand_size(kgrid, is_axisymmetric, pml_size): 

83 # set the PML size for use with expandMatrix, don't expand the inner radial 

84 # dimension if using the axisymmetric code 

85 if kgrid.dim == 1: 

86 expand_size = pml_size[0] 

87 elif kgrid.dim == 2: 

88 if is_axisymmetric: 

89 expand_size = [pml_size[0], pml_size[0], 0, pml_size[1]] 

90 else: 

91 expand_size = pml_size 

92 elif kgrid.dim == 3: 

93 expand_size = pml_size 

94 else: 

95 raise NotImplementedError 

96 return np.array(expand_size) 

97 

98 

99def expand_medium(medium: kWaveMedium, expand_size): 

100 # enlarge the sound speed grids by exting the edge values into the expanded grid 

101 medium.sound_speed = np.atleast_1d(medium.sound_speed) 

102 if medium.sound_speed.size > 1: 

103 medium.sound_speed = expand_matrix(medium.sound_speed, expand_size) 

104 

105 # enlarge the grid of density by exting the edge values into the expanded grid 

106 medium.density = np.atleast_1d(medium.density) 

107 if medium.density.size > 1: 

108 medium.density = expand_matrix(medium.density, expand_size) 

109 

110 # for key in ['alpha_coeff', 'alpha_coeff_compression', 'alpha_coeff_shear', 'BonA']: 

111 for key in ['alpha_coeff', 'BonA']: 

112 # enlarge the grid of medium[key] if given 

113 attr = getattr(medium, key) 

114 if attr is not None and np.atleast_1d(attr).size > 1: 

115 attr = expand_matrix(np.atleast_1d(attr), expand_size) 

116 setattr(medium, key, attr) 

117 

118 # enlarge the absorption filter mask if given 

119 if medium.alpha_filter is not None: 

120 medium.alpha_filter = expand_matrix(medium.alpha_filter, expand_size, 0) 

121 

122 

123def expand_source( 

124 source, is_source_sensor_same, flags, expand_size, index_data_type, 

125 p_source_pos_index, u_source_pos_index, s_source_pos_index): 

126 

127 p_source_pos_index = expand_pressure_sources(source, expand_size, flags.source_p0, flags.source_p, 

128 index_data_type, p_source_pos_index) 

129 

130 u_source_pos_index = expand_velocity_sources( 

131 source, expand_size, is_source_sensor_same, index_data_type, u_source_pos_index, 

132 flags.source_ux, flags.source_uy, flags.source_uz, flags.transducer_source 

133 ) 

134 

135 s_source_pos_index = expand_stress_sources(source, expand_size, flags, index_data_type, s_source_pos_index) 

136 

137 return p_source_pos_index, u_source_pos_index, s_source_pos_index 

138 

139 

140def expand_pressure_sources(source, expand_size, is_source_p0, is_source_p, index_data_type, p_source_pos_index): 

141 # enlarge the initial pressure if given 

142 if is_source_p0: 

143 source.p0 = expand_matrix(source.p0, expand_size, 0) 

144 

145 # enlarge the pressure source mask if given 

146 if is_source_p: 

147 

148 # enlarge the pressure source mask 

149 source.p_mask = expand_matrix(source.p_mask, expand_size, 0) 

150 

151 # create an indexing variable corresponding to the source elements 

152 # and convert the data type deping on the number of indices 

153 p_source_pos_index = matlab_find(source.p_mask).astype(index_data_type) 

154 return p_source_pos_index 

155 

156 

157def expand_velocity_sources( 

158 source, expand_size, is_source_sensor_same, index_data_type, u_source_pos_index, 

159 is_source_ux, is_source_uy, is_source_uz, is_transducer_source 

160): 

161 """ 

162 enlarge the velocity source mask if given 

163 Args: 

164 source: 

165 expand_size: 

166 is_source_sensor_same: 

167 index_data_type: 

168 u_source_pos_index: 

169 is_source_ux: 

170 is_source_uy: 

171 is_source_uz: 

172 is_transducer_source: 

173 

174 Returns: 

175 

176 """ 

177 if is_source_ux or is_source_uy or is_source_uz or is_transducer_source: 

178 

179 # update the source indexing variable 

180 if isinstance(source, NotATransducer): 

181 

182 # check if the sensor is also the same transducer, if so, don't expand the grid again 

183 if not is_source_sensor_same: 

184 

185 # expand the transducer mask 

186 source.expand_grid(expand_size) 

187 

188 # get the new active elements mask 

189 active_elements_mask = source.active_elements_mask 

190 

191 # update the indexing variable corresponding to the active elements 

192 u_source_pos_index = matlab_find(active_elements_mask) 

193 else: 

194 

195 # enlarge the velocity source mask 

196 source.u_mask = expand_matrix(source.u_mask, expand_size, 0) 

197 

198 # create an indexing variable corresponding to the source elements 

199 u_source_pos_index = matlab_find(source.u_mask) 

200 

201 # convert the data type deping on the number of indices 

202 u_source_pos_index = u_source_pos_index.astype(index_data_type) 

203 return u_source_pos_index 

204 

205 

206def expand_stress_sources(source, expand_size, flags, index_data_type, s_source_pos_index): 

207 # enlarge the stress source mask if given 

208 if flags.source_sxx or flags.source_syy or flags.source_szz or flags.source_sxy or flags.source_sxz or flags.source_syz: 

209 # enlarge the velocity source mask 

210 source.s_mask = expand_matrix(source.s_mask, expand_size, 0) 

211 

212 # create an indexing variable corresponding to the source elements 

213 s_source_pos_index = matlab_find(source.s_mask != 0) 

214 

215 # convert the data type deping on the number of indices 

216 s_source_pos_index = s_source_pos_index.astype(index_data_type) 

217 return s_source_pos_index 

218 

219 

220def expand_directivity_angle(kgrid, sensor, expand_size, is_use_sensor, is_compute_directivity): 

221 """ 

222 enlarge the directivity angle if given (2D only) 

223 Args: 

224 kgrid: 

225 sensor: 

226 expand_size: 

227 is_use_sensor: 

228 is_compute_directivity: 

229 

230 Returns: 

231 

232 """ 

233 if is_use_sensor and kgrid.dim == 2 and is_compute_directivity: 

234 # enlarge the directivity angle 

235 sensor.directivity.angle = expand_matrix(sensor.directivity.angle, expand_size, 0) 

236 # re-assign the wavenumber vectors 

237 sensor.directivity.wavenumbers = np.hstack((kgrid.ky.T, kgrid.kx.T)) 

238 

239 

240def print_grid_size(kgrid): 

241 """ 

242 update command line status 

243 Args: 

244 kgrid: 

245 

246 Returns: 

247 

248 """ 

249 k_Nx, k_Ny, k_Nz = kgrid.Nx, kgrid.Ny, kgrid.Nz 

250 if kgrid.dim == 1: 

251 print(' computational grid size:', int(k_Nx), 'grid points') 

252 elif kgrid.dim == 2: 

253 print(' computational grid size:', int(k_Nx), 'by', int(k_Ny), 'grid points') 

254 elif kgrid.dim == 3: 

255 print(' computational grid size:', int(k_Nx), 'by', int(k_Ny), 'by', int(k_Nz), 'grid points') 

256 

257 

258def expand_cuboid_corner_list(is_cuboid_list, kgrid, pml_size: Array): 

259 """ 

260 add the PML size to cuboid corner indices if using a cuboid sensor mask 

261 Args: 

262 is_cuboid_list: 

263 kgrid: 

264 

265 Returns: 

266 

267 """ 

268 if not is_cuboid_list: 

269 return 

270 

271 record = dotdict() 

272 if kgrid.dim == 1: 

273 record.cuboid_corners_list = record.cuboid_corners_list + pml_size.x 

274 elif kgrid.dim == 2: 

275 record.cuboid_corners_list[[0, 2], :] = record.cuboid_corners_list[[0, 2], :] + pml_size.x 

276 record.cuboid_corners_list[[1, 3], :] = record.cuboid_corners_list[[1, 3], :] + pml_size.y 

277 elif kgrid.dim == 3: 

278 record.cuboid_corners_list[[0, 3], :] = record.cuboid_corners_list[[0, 3], :] + pml_size.x 

279 record.cuboid_corners_list[[1, 4], :] = record.cuboid_corners_list[[1, 4], :] + pml_size.y 

280 record.cuboid_corners_list[[2, 5], :] = record.cuboid_corners_list[[2, 5], :] + pml_size.z 

281 return record 

282 

283 

284def expand_sensor(sensor, expand_size, is_use_sensor, is_blank_sensor): 

285 """ 

286 enlarge the sensor mask (for Cartesian sensor masks and cuboid corners, 

287 this has already been converted to a binary mask for display in inputChecking) 

288 Args: 

289 sensor: 

290 expand_size: 

291 is_use_sensor: 

292 is_blank_sensor: 

293 

294 Returns: 

295 

296 """ 

297 if is_use_sensor and not is_blank_sensor: 

298 sensor.expand_grid(expand_size)