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

1 from kwave import kWaveGrid, kWaveMedium, SimulationOptions, NotATransducer

2 from kwave.data import Array

3 from kwave.utils import matlab_find, expand_matrix, get_smallest_possible_type

4 from kwave.utils import dotdict

5 import numpy as np

6

7

8 def 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

60 def 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

82 def 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

99 def 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

123 def 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

140 def 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

157 def 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

206 def 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

220 def 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

240 def 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

258 def 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

284 def 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)