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

1import logging 

2 

3import numpy as np 

4 

5from kwave.data import Vector 

6from kwave.kgrid import kWaveGrid 

7from kwave.kmedium import kWaveMedium 

8from kwave.ktransducer import NotATransducer 

9from kwave.options.simulation_options import SimulationOptions 

10from kwave.utils.data import get_smallest_possible_type 

11from kwave.utils.dotdictionary import dotdict 

12from kwave.utils.matlab import matlab_find 

13from kwave.utils.matrix import expand_matrix 

14 

15 

16def expand_grid_matrices(kgrid: kWaveGrid, medium: kWaveMedium, source, sensor, opt: SimulationOptions, values: dotdict, flags: dotdict): 

17 # update command line status 

18 logging.log(logging.INFO, " expanding computational grid...") 

19 

20 ##################### 

21 # Grab values 

22 ##################### 

23 

24 # retaining the values for kgrid time array 

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

26 pml_size = Vector(pml_size[: kgrid.dim]) 

27 

28 p_source_pos_index = values.p_source_pos_index 

29 u_source_pos_index = values.u_source_pos_index 

30 s_source_pos_index = values.s_source_pos_index 

31 

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

33 

34 ##################### 

35 # Expand Structures 

36 ##################### 

37 

38 # expand the computational grid, replacing the original grid 

39 kgrid = expand_kgrid(kgrid, flags.axisymmetric, pml_size) 

40 

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

42 

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

44 total_grid_points = kgrid.total_grid_points 

45 index_data_type = get_smallest_possible_type(total_grid_points, "uint", default="double") 

46 

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

48 

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

50 record = expand_cuboid_corner_list(flags.cuboid_corners, kgrid, pml_size) # noqa: F841 

51 

52 expand_medium(medium, expand_size) 

53 

54 p_source_pos_index, u_source_pos_index, s_source_pos_index = expand_source( 

55 source, is_source_sensor_same, flags, expand_size, index_data_type, p_source_pos_index, u_source_pos_index, s_source_pos_index 

56 ) 

57 

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

59 

60 print_grid_size(kgrid) 

61 

62 return kgrid, index_data_type, p_source_pos_index, u_source_pos_index, s_source_pos_index 

63 

64 

65def expand_kgrid(kgrid, is_axisymmetric, pml_size): 

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

67 

68 pml_size = pml_size.squeeze() 

69 

70 if kgrid.dim == 1: 

71 new_size = kgrid.N + 2 * pml_size 

72 elif kgrid.dim == 2: 

73 if is_axisymmetric: 

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

75 else: 

76 new_size = kgrid.N + 2 * pml_size 

77 elif kgrid.dim == 3: 

78 new_size = kgrid.N + 2 * pml_size 

79 else: 

80 raise NotImplementedError 

81 

82 kgrid = kWaveGrid(new_size, kgrid.spacing) 

83 # re-assign original time array 

84 kgrid.setTime(Nt_temp, dt_temp) 

85 

86 return kgrid 

87 

88 

89def calculate_expand_size(kgrid, is_axisymmetric, pml_size): 

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

91 # dimension if using the axisymmetric code 

92 if kgrid.dim == 1: 

93 expand_size = pml_size[0] 

94 elif kgrid.dim == 2: 

95 if is_axisymmetric: 

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

97 else: 

98 expand_size = pml_size 

99 elif kgrid.dim == 3: 

100 expand_size = pml_size 

101 else: 

102 raise NotImplementedError 

103 return np.array(expand_size) 

104 

105 

106def expand_medium(medium: kWaveMedium, expand_size): 

107 # enlarge the sound speed grids by extending the edge values into the expanded grid 

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

109 if medium.sound_speed.size > 1: 

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

111 

112 # enlarge the grid of density by extending the edge values into the expanded grid 

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

114 if medium.density.size > 1: 

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

116 

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

118 for key in ["alpha_coeff", "BonA"]: 

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

120 attr = getattr(medium, key) 

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

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

123 setattr(medium, key, attr) 

124 

125 # enlarge the absorption filter mask if given 

126 if medium.alpha_filter is not None: 

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

128 

129 

130def expand_source( 

131 source, is_source_sensor_same, flags, expand_size, index_data_type, p_source_pos_index, u_source_pos_index, s_source_pos_index 

132): 

133 p_source_pos_index = expand_pressure_sources(source, expand_size, flags.source_p0, flags.source_p, index_data_type, p_source_pos_index) 

134 

135 u_source_pos_index = expand_velocity_sources( 

136 source, 

137 expand_size, 

138 is_source_sensor_same, 

139 index_data_type, 

140 u_source_pos_index, 

141 flags.source_ux, 

142 flags.source_uy, 

143 flags.source_uz, 

144 flags.transducer_source, 

145 ) 

146 

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

148 

149 return p_source_pos_index, u_source_pos_index, s_source_pos_index 

150 

151 

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

153 # enlarge the initial pressure if given 

154 if is_source_p0: 

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

156 

157 # enlarge the pressure source mask if given 

158 if is_source_p: 

159 # enlarge the pressure source mask 

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

161 

162 # create an indexing variable corresponding to the source elements 

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

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

165 return p_source_pos_index 

166 

167 

168def expand_velocity_sources( 

169 source, 

170 expand_size, 

171 is_source_sensor_same, 

172 index_data_type, 

173 u_source_pos_index, 

174 is_source_ux, 

175 is_source_uy, 

176 is_source_uz, 

177 is_transducer_source, 

178): 

179 """ 

180 enlarge the velocity source mask if given 

181 Args: 

182 source: 

183 expand_size: 

184 is_source_sensor_same: 

185 index_data_type: 

186 u_source_pos_index: 

187 is_source_ux: 

188 is_source_uy: 

189 is_source_uz: 

190 is_transducer_source: 

191 

192 Returns: 

193 

194 """ 

195 if is_source_ux or is_source_uy or is_source_uz or is_transducer_source: 

196 # update the source indexing variable 

197 if isinstance(source, NotATransducer): 

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

199 if not is_source_sensor_same: 

200 # expand the transducer mask 

201 source.expand_grid(expand_size) 

202 

203 # get the new active elements mask 

204 active_elements_mask = source.active_elements_mask 

205 

206 # update the indexing variable corresponding to the active elements 

207 u_source_pos_index = matlab_find(active_elements_mask) 

208 else: 

209 # enlarge the velocity source mask 

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

211 

212 # create an indexing variable corresponding to the source elements 

213 u_source_pos_index = matlab_find(source.u_mask) 

214 

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

216 u_source_pos_index = u_source_pos_index.astype(index_data_type) 

217 return u_source_pos_index 

218 

219 

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

221 # enlarge the stress source mask if given 

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

223 # enlarge the velocity source mask 

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

225 

226 # create an indexing variable corresponding to the source elements 

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

228 

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

230 s_source_pos_index = s_source_pos_index.astype(index_data_type) 

231 return s_source_pos_index 

232 

233 

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

235 """ 

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

237 Args: 

238 kgrid: 

239 sensor: 

240 expand_size: 

241 is_use_sensor: 

242 is_compute_directivity: 

243 

244 Returns: 

245 

246 """ 

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

248 # enlarge the directivity angle 

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

250 # re-assign the wavenumber vectors 

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

252 

253 

254def print_grid_size(kgrid): 

255 """ 

256 update command line status 

257 Args: 

258 kgrid: 

259 

260 Returns: 

261 

262 """ 

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

264 if kgrid.dim == 1: 

265 logging.log(logging.INFO, " computational grid size:", int(k_Nx), "grid points") 

266 elif kgrid.dim == 2: 

267 logging.log(logging.INFO, " computational grid size:", int(k_Nx), "by", int(k_Ny), "grid points") 

268 elif kgrid.dim == 3: 

269 logging.log(logging.INFO, " computational grid size:", int(k_Nx), "by", int(k_Ny), "by", int(k_Nz), "grid points") 

270 

271 

272def expand_cuboid_corner_list(is_cuboid_list, kgrid, pml_size: Vector): 

273 """ 

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

275 Args: 

276 is_cuboid_list: 

277 kgrid: 

278 

279 Returns: 

280 

281 """ 

282 if not is_cuboid_list: 

283 return 

284 

285 record = dotdict() 

286 if kgrid.dim == 1: 

287 record.cuboid_corners_list = record.cuboid_corners_list + pml_size.x 

288 elif kgrid.dim == 2: 

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

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

291 elif kgrid.dim == 3: 

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

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

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

295 return record 

296 

297 

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

299 """ 

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

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

302 Args: 

303 sensor: 

304 expand_size: 

305 is_use_sensor: 

306 is_blank_sensor: 

307 

308 Returns: 

309 

310 """ 

311 if is_use_sensor and not is_blank_sensor: 

312 sensor.expand_grid(expand_size)