Coverage for kwave/kWaveSimulation_helper/scale_source_terms_func.py: 18%

130 statements  

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

1import math 

2import numpy as np 

3from warnings import warn 

4 

5from kwave import kWaveGrid 

6from kwave.ksource import kSource 

7from kwave.utils import matlab_mask 

8from kwave.utils import dotdict 

9 

10 

11def scale_source_terms_func( 

12 c0, dt, kgrid: kWaveGrid, source, 

13 p_source_pos_index, s_source_pos_index, u_source_pos_index, 

14 transducer_input_signal, 

15 flags: dotdict): 

16 """ 

17 Subscript for the first-order k-Wave simulation functions to scale source terms to the correct units. 

18 Args: 

19 Returns: 

20 

21 """ 

22 dx, dy, dz = kgrid.dx, kgrid.dy, kgrid.dz, 

23 

24 # get the dimension size 

25 N = kgrid.dim 

26 

27 if not check_conditions(flags.nonuniform_grid, flags.source_uy, flags.source_uz, flags.transducer_source): 

28 return 

29 

30 # ========================================================================= 

31 # PRESSURE SOURCES 

32 # ========================================================================= 

33 

34 apply_pressure_source_correction(flags.source_p, flags.use_w_source_correction_p, source, dt) 

35 

36 scale_pressure_source(flags.source_p, source, kgrid, N, c0, dx, dt, p_source_pos_index, flags.nonuniform_grid) 

37 

38 # ========================================================================= 

39 # STRESS SOURCES 

40 # ========================================================================= 

41 

42 scale_stress_sources(source, c0, flags, dt, dx, N, s_source_pos_index) 

43 

44 # ========================================================================= 

45 # VELOCITY SOURCES 

46 # ========================================================================= 

47 

48 apply_velocity_source_corrections(flags.use_w_source_correction_u, flags.source_ux, flags.source_uy, flags.source_uz, source, dt) 

49 

50 scale_velocity_sources(flags, source, kgrid, c0, dt, dx, dy, dz, u_source_pos_index) 

51 

52 # ========================================================================= 

53 # TRANSDUCER SOURCE 

54 # ========================================================================= 

55 transducer_input_signal = scale_transducer_source(flags.transducer_source, transducer_input_signal, 

56 c0, dt, dx, u_source_pos_index) 

57 return transducer_input_signal 

58 

59 

60def check_conditions(is_nonuniform_grid, is_source_uy, is_source_uz, is_transducer_source): 

61 """ 

62 check for non-uniform grid and give error for source terms that haven't yet been implemented 

63 Returns: 

64 

65 """ 

66 if is_nonuniform_grid and (is_source_uy or is_source_uz or is_transducer_source): 

67 warn('WARNING: source scaling not implemented for non-uniform grids with given source condition') 

68 return False 

69 return True 

70 

71 

72def apply_pressure_source_correction(is_source_p, use_w_source_correction_p, source, dt): 

73 """ 

74 apply k-space source correction expressed as a function of w 

75 Args: 

76 is_source_p: 

77 use_w_source_correction_p: 

78 source: 

79 dt: 

80 

81 Returns: 

82 

83 """ 

84 if is_source_p and use_w_source_correction_p: 

85 source.p = apply_source_correction(source.p, source.p_frequency_ref, dt) 

86 

87 

88def scale_pressure_source(is_source_p, source, kgrid, N, c0, dx, dt, p_source_pos_index, is_nonuniform_grid): 

89 """ 

90 scale the input pressure by 1/c0^2 (to convert to units of density), then 

91 by 1/N (to split the input across the split density field). If the 

92 pressure is injected as a mass source, also scale the pressure by 

93 2*dt*c0/dx to account for the time step and convert to units of [kg/(m^3 s)] 

94 Args: 

95 is_source_p: 

96 source: 

97 kgrid: 

98 N: 

99 c0: 

100 dx: 

101 dt: 

102 p_source_pos_index: 

103 is_nonuniform_grid: 

104 

105 Returns: 

106 

107 """ 

108 if not is_source_p: 

109 return 

110 

111 if source.p_mode == 'dirichlet': 

112 source.p = scale_pressure_source_dirichlet(source.p, c0, N, p_source_pos_index) 

113 else: 

114 if is_nonuniform_grid: 

115 source.p = scale_pressure_source_nonuniform_grid(source.p, kgrid, c0, N, dt, p_source_pos_index) 

116 

117 else: 

118 source.p = scale_pressure_source_uniform_grid(source.p, c0, N, dx, dt, p_source_pos_index) 

119 

120 

121def scale_pressure_source_dirichlet(source_p, c0, N, p_source_pos_index): 

122 if c0.size == 1: 

123 # compute the scale parameter based on the homogeneous 

124 # sound speed 

125 source_p = source_p / (N * (c0 ** 2)) 

126 

127 else: 

128 # compute the scale parameter seperately for each source 

129 # position based on the sound speed at that position 

130 for p_index in range(source_p.size[0]): 

131 source_p[p_index, :] = source_p[p_index, :] / (N * (c0[p_source_pos_index[p_index]] ** 2)) 

132 return source_p 

133 

134 

135def scale_pressure_source_nonuniform_grid(source_p, kgrid, c0, N, dt, p_source_pos_index): 

136 x = kgrid.x 

137 xn = kgrid.xn 

138 yn = kgrid.yn 

139 zn = kgrid.zn 

140 x_size, y_size, z_size = kgrid.size 

141 

142 # create empty matrix 

143 grid_point_sep = np.zeros(x.size) 

144 

145 # compute averaged grid point seperation map, the interior 

146 # points are calculated using the average distance to all 

147 # connected grid points (the edge values are not calculated 

148 # assuming there are no source points in the PML) 

149 if kgrid.dim == 1: 

150 grid_point_sep[1:-1] = x_size * (xn[2:, 0] - xn[0:-2, 0]) / 2 

151 elif kgrid.dim == 2: 

152 grid_point_sep[1:-1, 1:-1] = x_size * (xn[2:, 1:-1] - xn[0:-2, 1:- 1]) / 4 + \ 

153 y_size * (yn[1:-1, 2:] - yn[1:-1, 0:-2]) / 4 

154 elif kgrid.dim == 3: 

155 grid_point_sep[1:-1, 1:-1, 1:-1] = x_size * (xn[2:, 1:-1, 1:-1] - xn[0:-2, 1:-1, 1:-1]) / 6 + \ 

156 y_size * (yn[1:-1, 2:, 1:-1] - yn[1:-1, 0:-2, 1:- 1]) / 6 + \ 

157 z_size * (zn[1:-1, 1:-1, 2:] - zn[1:-1, 1:-1, 0:-2]) / 6 

158 

159 # compute and apply scale parameter 

160 for p_index in range(source_p.size[0]): 

161 if c0.size == 1: 

162 

163 # compute the scale parameter based on the homogeneous sound speed 

164 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * c0 * grid_point_sep[p_source_pos_index[p_index]])) 

165 

166 else: 

167 

168 # compute the scale parameter based on the sound speed at that position 

169 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * c0[p_source_pos_index[p_index]] * grid_point_sep[p_source_pos_index[p_index]])) 

170 return source_p 

171 

172 

173def scale_pressure_source_uniform_grid(source_p, c0, N, dx, dt, p_source_pos_index): 

174 if c0.size == 1: 

175 # compute the scale parameter based on the homogeneous 

176 # sound speed 

177 source_p = source_p * (2 * dt / (N * c0 * dx)) 

178 

179 else: 

180 # compute the scale parameter seperately for each source 

181 # position based on the sound speed at that position 

182 for p_index in range(source_p[:, 0].size): 

183 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * matlab_mask(c0, p_source_pos_index.flatten('F')[p_index]) * dx)) 

184 return source_p 

185 

186 

187def scale_stress_sources(source, c0, flags, dt, dx, N, s_source_pos_index): 

188 """ 

189 scale the stress source by 1/N to divide amoungst the split field 

190 components, and if source.s_mode is not set to 'dirichlet', also scale by 

191 2*dt*c0/dx to account for the time step and convert to units of 

192 [kg/(m^3 s)] (note dx is used in all dimensions) 

193 Args: 

194 source: 

195 c0: 

196 flags: 

197 dt: 

198 dx: 

199 N: 

200 s_source_pos_index: 

201 

202 Returns: 

203 

204 """ 

205 source.sxx = scale_stress_source(source, c0, flags.source_sxx, flags.source_p0, source.sxx, dt, N, dx, s_source_pos_index) 

206 source.syy = scale_stress_source(source, c0, flags.source_syy, flags.source_p0, source.syy, dt, N, dx, s_source_pos_index) 

207 source.szz = scale_stress_source(source, c0, flags.source_szz, flags.source_p0, source.szz, dt, N, dx, s_source_pos_index) 

208 source.sxy = scale_stress_source(source, c0, flags.source_sxy, True, source.sxy, dt, N, dx, s_source_pos_index) 

209 source.sxz = scale_stress_source(source, c0, flags.source_sxz, True, source.sxz, dt, N, dx, s_source_pos_index) 

210 source.syz = scale_stress_source(source, c0, flags.source_syz, True, source.syz, dt, N, dx, s_source_pos_index) 

211 

212 

213def scale_stress_source(source, c0, is_source_exists, is_p0_exists, source_val, dt, N, dx, s_source_pos_index): 

214 if is_source_exists: 

215 if source.s_mode == 'dirichlet' or is_p0_exists: 

216 source_val = source_val / N 

217 else: 

218 if c0.size == 1: 

219 

220 # compute the scale parameter based on the homogeneous sound 

221 # speed 

222 source_val = source_val * (2 * dt * c0 / (N * dx)) 

223 

224 else: 

225 

226 # compute the scale parameter seperately for each source 

227 # position based on the sound speed at that position 

228 for s_index in range(source_val.size[0]): 

229 source_val[s_index, :] = source_val[s_index, :] * (2 * dt * c0[s_source_pos_index[s_index]] / (N * dx)) 

230 return source_val 

231 

232 

233def apply_velocity_source_corrections( 

234 use_w_source_correction_u: bool, is_ux_exists: bool, is_uy_exists: bool, is_uz_exists: bool, 

235 source: kSource, dt: float 

236): 

237 """ 

238 apply k-space source correction expressed as a function of w 

239 Args: 

240 use_w_source_correction_u: 

241 is_ux_exists: 

242 is_uy_exists: 

243 is_uz_exists: 

244 source: 

245 dt: 

246 

247 Returns: 

248 

249 """ 

250 if not use_w_source_correction_u: 

251 return 

252 

253 if is_ux_exists: 

254 source.ux = apply_source_correction(source.ux, source.u_frequency_ref, dt) 

255 

256 if is_uy_exists: 

257 source.uy = apply_source_correction(source.uy, source.u_frequency_ref, dt) 

258 

259 if is_uz_exists: 

260 source.uz = apply_source_correction(source.uz, source.u_frequency_ref, dt) 

261 

262 

263def apply_source_correction(source_val, frequency_ref, dt): 

264 return source_val * math.cos(2 * math.pi * frequency_ref * dt/2) 

265 

266 

267def scale_velocity_sources(flags, source, kgrid, c0, dt, dx, dy, dz, u_source_pos_index): 

268 source.ux = scale_velocity_source_x(flags.source_ux, source.u_mode, source.ux, kgrid, c0, dt, dx, u_source_pos_index, flags.nonuniform_grid) 

269 source.uy = scale_velocity_source(flags.source_uy, source.u_mode, source.uy, c0, dt, u_source_pos_index, dy) 

270 source.uz = scale_velocity_source(flags.source_uz, source.u_mode, source.uz, c0, dt, u_source_pos_index, dz) 

271 

272 

273def scale_velocity_source_x(is_source_ux, source_u_mode, source_val, kgrid, c0, dt, dx, u_source_pos_index, is_nonuniform_grid): 

274 """ 

275 if source.u_mode is not set to 'dirichlet', scale the x-direction 

276 velocity source terms by 2*dt*c0/dx to account for the time step and 

277 convert to units of [m/s^2] 

278 Returns: 

279 

280 """ 

281 if not is_source_ux or source_u_mode == 'dirichlet': 

282 return 

283 

284 if is_nonuniform_grid: 

285 source_val = scale_velocity_source_nonuniform(is_source_ux, source_u_mode, kgrid, source_val, 

286 c0, dt, u_source_pos_index) 

287 else: 

288 source_val = scale_velocity_source(is_source_ux, source_u_mode, source_val, c0, dt, u_source_pos_index, dx) 

289 return source_val 

290 

291 

292def scale_velocity_source(is_source, source_u_mode, source_val, c0, dt, u_source_pos_index, d_direction): 

293 """ 

294 if source.u_mode is not set to 'dirichlet', scale the d_direction 

295 velocity source terms by 2*dt*c0/dz to account for the time step and 

296 convert to units of [m/s^2] 

297 Args: 

298 is_source: 

299 source_u_mode: 

300 source_val: 

301 c0: 

302 dt: 

303 u_source_pos_index: 

304 d_direction: 

305 

306 Returns: 

307 

308 """ 

309 if not is_source or source_u_mode == 'dirichlet': 

310 return source_val 

311 

312 if c0.size == 1: 

313 # compute the scale parameter based on the homogeneous sound speed 

314 source_val = source_val * (2 * c0 * dt / d_direction) 

315 else: 

316 # compute the scale parameter seperately for each source position 

317 # based on the sound speed at that position 

318 for u_index in range(source_val.size[0]): 

319 source_val[u_index, :] = source_val[u_index, :] * (2 * c0(u_source_pos_index[u_index]) * dt / d_direction) 

320 return source_val 

321 

322 

323def scale_velocity_source_nonuniform(is_source, source_u_mode, kgrid, source_val, c0, dt, u_source_pos_index): 

324 """ 

325 if source.u_mode is not set to 'dirichlet', scale the d_direction 

326 velocity source terms by 2*dt*c0/dz to account for the time step and 

327 convert to units of [m/s^2] 

328 Args: 

329 is_source: 

330 source_u_mode: 

331 kgrid: 

332 source_val: 

333 c0: 

334 dt: 

335 u_source_pos_index: 

336 

337 Returns: 

338 

339 """ 

340 if not is_source or source_u_mode == 'dirichlet': 

341 return source_val 

342 

343 # create empty matrix 

344 x = kgrid.x 

345 xn = kgrid.xn 

346 x_size = kgrid.size[0] 

347 grid_point_sep = np.zeros_like(x) 

348 

349 # compute averaged grid point seperation map, the interior 

350 # points are calculated using the average distance to all 

351 # connected grid points (the edge values are not calculated 

352 # assuming there are no source points in the PML) 

353 grid_point_sep[1:- 1, :, :] = x_size * (xn[2:, :, :] - xn[1:- 2, :, :]) / 2 

354 

355 # compute and apply scale parameter 

356 for u_index in range(source_val.size[0]): 

357 if c0.size == 1: 

358 # compute the scale parameter based on the homogeneous sound speed 

359 source_val[u_index, :] = source_val[u_index, :] * (2 * c0 * dt / (grid_point_sep[u_source_pos_index[u_index]])) 

360 else: 

361 # compute the scale parameter based on the sound speed at that position 

362 source_val[u_index, :] = source_val[u_index, :] * (2 * c0[u_source_pos_index[u_index]] * dt / (grid_point_sep[u_source_pos_index[u_index]])) 

363 return source_val 

364 

365 

366def scale_transducer_source(is_transducer_source, transducer_input_signal, c0, dt, dx, u_source_pos_index): 

367 """ 

368 scale the transducer source term by 2*dt*c0/dx to account for the time 

369 step and convert to units of [m/s^2] 

370 Args: 

371 is_transducer_source: 

372 transducer_input_signal: 

373 c0: 

374 dt: 

375 dx: 

376 u_source_pos_index: 

377 

378 Returns: 

379 

380 """ 

381 if is_transducer_source: 

382 if c0.size == 1: 

383 transducer_input_signal = transducer_input_signal * (2 * c0 * dt / dx) 

384 else: 

385 # compute the scale parameter based on the average sound speed at the 

386 # transducer positions (only one input signal is used to drive the transducer) 

387 transducer_input_signal = transducer_input_signal * (2 * (np.mean(c0.flatten(order='F')[u_source_pos_index - 1])) * dt / dx) 

388 return transducer_input_signal