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

1 import math

2 import numpy as np

3 from warnings import warn

4

5 from kwave import kWaveGrid

6 from kwave.ksource import kSource

7 from kwave.utils import matlab_mask

8 from kwave.utils import dotdict

9

10

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

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

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

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

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

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

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

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

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

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

263 def apply_source_correction(source_val, frequency_ref, dt):

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

265

266

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

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

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

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

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