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

1 import os

2

3 from scipy.io import savemat

4

5 from kwave import kWaveMedium, kWaveGrid, SimulationOptions

6 from kwave.utils import scale_time, TicToc, num_dim2, write_matrix, write_attributes

7 from kwave.utils import dotdict

8 import numpy as np

9

10

11 def save_to_disk_func(

12 kgrid: kWaveGrid, medium: kWaveMedium, source,

13 opt: SimulationOptions, values: dotdict, flags: dotdict):

14 # update command line status

15 print(' precomputation completed in ', scale_time(TicToc.toc()))

16 TicToc.tic()

17 print(' saving input files to disk...')

18

19 # check for a binary sensor mask or cuboid corners

20 # modified by Farid | disabled temporarily!

21 # assert self.binary_sensor_mask or self.cuboid_corners, \

22 # "Optional input ''SaveToDisk'' only supported for sensor masks defined as a binary matrix or the opposing corners of a rectangle (2D) or cuboid (3D)."

23

24 # =========================================================================

25 # VARIABLE LIST

26 # =========================================================================

27 integer_variables = dotdict()

28 float_variables = dotdict()

29

30 grab_integer_variables(integer_variables, kgrid, flags, medium)

31 grab_pml_size(integer_variables, opt)

32 grab_float_variables(float_variables, kgrid, opt, values, flags.elastic_code, flags.axisymmetric)

33

34 # overwrite z-values for 2D simulations

35 if kgrid.dim == 2:

36 integer_variables.Nz = 1

37 integer_variables.pml_z_size = 0

38

39 grab_medium_props(integer_variables, float_variables, medium, flags.elastic_code)

40 grab_source_props(integer_variables, float_variables, source,

41 values.u_source_pos_index, values.s_source_pos_index, values.p_source_pos_index,

42 values.transducer_input_signal, values.delay_mask)

43

44 grab_sensor_props(integer_variables, kgrid.dim, values.sensor_mask_index, values.record.cuboid_corners_list)

45 grab_nonuniform_grid_props(float_variables, kgrid, flags.nonuniform_grid)

46

47 # =========================================================================

48 # DATACAST AND SAVING

49 # =========================================================================

50

51 remove_z_dimension(float_variables, kgrid.dim)

52 save_file(opt.save_to_disk, integer_variables, float_variables, opt.hdf_compression_level)

53

54 # update command line status

55 print(' completed in ', scale_time(TicToc.toc()))

56

57

58 def grab_integer_variables(integer_variables, kgrid, flags, medium):

59 # integer variables used within the time loop for all codes

60

61 variables = dotdict({

62 'Nx': kgrid.Nx,

63 'Ny': kgrid.Ny,

64 'Nz': kgrid.Nz,

65 'Nt': kgrid.Nt,

66 'p_source_flag': flags.source_p,

67 'p0_source_flag': flags.source_p0,

68

69 'ux_source_flag': flags.source_ux,

70 'uy_source_flag': flags.source_uy,

71 'uz_source_flag': flags.source_uz,

72

73 'sxx_source_flag': flags.source_sxx,

74 'syy_source_flag': flags.source_syy,

75 'szz_source_flag': flags.source_szz,

76

77 'sxy_source_flag': flags.source_sxy,

78 'sxz_source_flag': flags.source_sxz,

79 'syz_source_flag': flags.source_syz,

80

81 'transducer_source_flag': flags.transducer_source,

82 'nonuniform_grid_flag': flags.nonuniform_grid,

83 'nonlinear_flag': medium.is_nonlinear(),

84 'absorbing_flag': None,

85 'elastic_flag': flags.elastic_code,

86 'axisymmetric_flag': flags.axisymmetric,

87

88 # create pseudonyms for the sensor flgs

89 # 0: binary mask indices

90 # 1: cuboid corners

91 'sensor_mask_type': flags.cuboid_corners

92 })

93 integer_variables.update(variables)

94

95

96 def grab_pml_size(integer_variables, opt):

97 # additional integer variables not used within time loop but stored directly to output file

98 integer_variables['pml_x_size'] = opt.pml_x_size

99 integer_variables['pml_y_size'] = opt.pml_y_size

100 integer_variables['pml_z_size'] = opt.pml_z_size

101

102

103 def grab_float_variables(float_variables: dotdict, kgrid, opt, values, is_elastic_code, is_axisymmetric):

104 # single precision variables not used within time loop but stored directly

105 # to the output file for all files

106 variables = dotdict({

107 'dx': kgrid.dx,

108 'dy': kgrid.dy,

109 'dz': kgrid.dz,

110

111 'pml_x_alpha': opt.pml_x_alpha,

112 'pml_y_alpha': opt.pml_y_alpha,

113 'pml_z_alpha': opt.pml_z_alpha

114 })

115 float_variables.update(variables)

116

117 if is_elastic_code: # pragma: no cover

118 grab_elastic_code_variables(float_variables, kgrid, values)

119 elif is_axisymmetric:

120 grab_axisymmetric_variables(float_variables, values)

121 else:

122 # single precision variables used within the time loop

123 float_variables['dt'] = values.dt

124 float_variables['c0'] = values.c0

125 float_variables['c_ref'] = values.c_ref

126 float_variables['rho0'] = values.rho0

127 float_variables['rho0_sgx'] = values.rho0_sgx

128 float_variables['rho0_sgy'] = values.rho0_sgy

129 float_variables['rho0_sgz'] = values.rho0_sgz

130

131

132 def grab_elastic_code_variables(float_variables, kgrid, values): # pragma: no cover

133 # single precision variables used within the time loop

134 float_variables['dt'] = None

135 float_variables['c_ref'] = None

136 float_variables['lambda'] = None

137 float_variables['mu'] = None

138

139 float_variables['rho0_sgx'] = None

140 float_variables['rho0_sgy'] = None

141 float_variables['rho0_sgz'] = None

142

143 float_variables['mu_sgxy'] = None

144 float_variables['mu_sgxz'] = None

145 float_variables['mu_sgyz'] = None

146

147 # create shift variables used for calculating u_non_staggered and I outputs

148 x_shift_neg = np.fft.ifftshift(np.exp(-1j * kgrid.k_vec.x * kgrid.dx / 2))

149 y_shift_neg = np.fft.ifftshift(np.exp(-1j * kgrid.k_vec.y * kgrid.dy / 2)).T

150 z_shift_neg = np.transpose(np.fft.ifftshift(np.exp(-1j * kgrid.k_vec.z * kgrid.dz / 2)), (1, 2, 0))

151

152 # create reduced variables for use with real-to-complex FFT

153 Nz = kgrid.Nz if kgrid.dim != 2 else 1

154 Nx_r = kgrid.Nx // 2 + 1

155 Ny_r = kgrid.Ny // 2 + 1

156 Nz_r = Nz // 2 + 1

157

158 ddx_k_shift_pos = values.ddx_k_shift_pos

159 ddx_k_shift_neg = values.ddx_k_shift_neg

160

161 float_variables['ddx_k_shift_pos_r'] = ddx_k_shift_pos[:Nx_r]

162 float_variables['ddy_k_shift_pos'] = None

163 float_variables['ddz_k_shift_pos'] = None

164

165 float_variables['ddx_k_shift_neg_r'] = ddx_k_shift_neg[:Nx_r]

166 float_variables['ddy_k_shift_neg'] = None

167 float_variables['ddz_k_shift_neg'] = None

168

169 float_variables['x_shift_neg_r'] = x_shift_neg[:Nx_r]

170 float_variables['y_shift_neg_r'] = y_shift_neg[:Ny_r]

171 float_variables['z_shift_neg_r'] = z_shift_neg[:Nz_r]

172

173 del x_shift_neg

174

175 float_variables['pml_x'] = None

176 float_variables['pml_y'] = None

177 float_variables['pml_z'] = None

178

179 float_variables['pml_x_sgx'] = None

180 float_variables['pml_y_sgy'] = None

181 float_variables['pml_z_sgz'] = None

182

183 float_variables['mpml_x_sgx'] = None

184 float_variables['mpml_y_sgy'] = None

185 float_variables['mpml_z_sgz'] = None

186

187 float_variables['mpml_x'] = None

188 float_variables['mpml_y'] = None

189 float_variables['mpml_z'] = None

190

191

192 def grab_axisymmetric_variables(float_variables, values):

193 # single precision variables used within the time loop

194 float_variables['dt'] = values.dt

195 float_variables['c0'] = values.c0

196 float_variables['c_ref'] = values.c_ref

197 float_variables['rho0'] = values.rho0

198 float_variables['rho0_sgx'] = values.rho0_sgx

199 float_variables['rho0_sgy'] = values.rho0_sgy

200

201

202 def grab_medium_props(integer_variables, float_variables, medium, is_elastic_code):

203 # =========================================================================

204 # VARIABLES USED IN NONLINEAR SIMULATIONS

205 # =========================================================================

206 if medium.is_nonlinear():

207 float_variables['BonA'] = medium.BonA

208

209 # =========================================================================

210 # VARIABLES USED IN ABSORBING SIMULATIONS

211 # =========================================================================

212

213 # set absorbing flag

214 if medium.absorbing:

215 integer_variables.absorbing_flag = 2 if medium.stokes else 1

216 else:

217 integer_variables.absorbing_flag = 0

218

219 if medium.absorbing:

220 if is_elastic_code: # pragma: no cover

221 # add to the variable list

222 float_variables['chi'] = None

223 float_variables['eta'] = None

224 float_variables['eta_sgxy'] = None

225 float_variables['eta_sgxz'] = None

226 float_variables['eta_sgyz'] = None

227 else:

228 float_variables['alpha_coeff'] = medium.alpha_coeff

229 float_variables['alpha_power'] = medium.alpha_power

230

231

232 def grab_source_props(integer_variables, float_variables, source,

233 u_source_pos_index, s_source_pos_index, p_source_pos_index,

234 transducer_input_signal, delay_mask):

235 # =========================================================================

236 # SOURCE VARIABLES

237 # =========================================================================

238 # source modes and indicies

239 # - these are only defined if the source flgs are > 0

240 # - the source mode describes whether the source will be added or replaced

241 # - the source indicies describe which grid points act as the source

242 # - the u_source_index is reused for any of the u sources and the transducer source

243

244 grab_velocity_source_props(integer_variables, source, u_source_pos_index)

245 grab_stress_source_props(integer_variables, source, s_source_pos_index)

246 grab_pressure_source_props(integer_variables, source, p_source_pos_index, u_source_pos_index)

247 grab_time_varying_source_props(integer_variables, float_variables, source, transducer_input_signal, delay_mask)

248

249

250 def grab_velocity_source_props(integer_variables, source, u_source_pos_index):

251 # velocity source

252 if any(integer_variables.get(k) for k in ['ux_source_flag', 'uy_source_flag', 'uz_source_flag']):

253 integer_variables['u_source_mode'] = {

254 'dirichlet': 0,

255 'additive-no-correction': 1,

256 'additive': 2,

257 }[source.u_mode]

258

259 if integer_variables.ux_source_flag:

260 u_source_many = num_dim2(source.ux) > 1

261 elif integer_variables.uy_source_flag:

262 u_source_many = num_dim2(source.uy) > 1

263 elif integer_variables.uz_source_flag:

264 u_source_many = num_dim2(source.uz) > 1

265 integer_variables['u_source_many'] = u_source_many

266

267 integer_variables.u_source_index = u_source_pos_index

268

269

270 def grab_stress_source_props(integer_variables, source, s_source_pos_index):

271 # stress source

272 if integer_variables.sxx_source_flag or integer_variables.syy_source_flag or integer_variables.szz_source_flag \

273 or integer_variables.sxy_source_flag or integer_variables.sxz_source_flag or integer_variables.syz_source_flag:

274 integer_variables.s_source_mode = source.s_mode != 'dirichlet'

275 if integer_variables.sxx_source_flag:

276 s_source_many = num_dim2(source.sxx) > 1

277 elif integer_variables.syy_source_flag:

278 s_source_many = num_dim2(source.syy) > 1

279 elif integer_variables.szz_source_flag:

280 s_source_many = num_dim2(source.szz) > 1

281 elif integer_variables.sxy_source_flag:

282 s_source_many = num_dim2(source.sxy) > 1

283 elif integer_variables.sxz_source_flag:

284 s_source_many = num_dim2(source.sxz) > 1

285 elif integer_variables.syz_source_flag:

286 s_source_many = num_dim2(source.syz) > 1

287 integer_variables.s_source_many = s_source_many

288 integer_variables.s_source_index = s_source_pos_index

289

290

291 def grab_pressure_source_props(integer_variables, source, p_source_pos_index, u_source_pos_index):

292 # pressure source

293 if integer_variables.p_source_flag:

294 integer_variables.p_source_mode = {

295 'dirichlet': 0,

296 'additive-no-correction': 1,

297 'additive': 2,

298 }[source.p_mode]

299 integer_variables.p_source_many = num_dim2(source.p) > 1

300 integer_variables.p_source_index = p_source_pos_index

301

302 # transducer source

303 if integer_variables.transducer_source_flag:

304 integer_variables.u_source_index = u_source_pos_index

305

306

307 def grab_time_varying_source_props(integer_variables, float_variables, source, transducer_input_signal, delay_mask):

308 # time varying source variables

309 # - these are only defined if the source flgs are > 0

310 # - these are the actual source values

311 # - these are indexed as (position_index, time_index)

312 if integer_variables.ux_source_flag:

313 float_variables.ux_source_input = source.ux

314

315 if integer_variables.uy_source_flag:

316 float_variables.uy_source_input = source.uy

317

318 if integer_variables.uz_source_flag:

319 float_variables.uz_source_input = source.uz

320

321 if integer_variables.sxx_source_flag:

322 float_variables.sxx_source_input = source.sxx

323

324 if integer_variables.syy_source_flag:

325 float_variables.syy_source_input = source.syy

326

327 if integer_variables.szz_source_flag:

328 float_variables.szz_source_input = source.szz

329

330 if integer_variables.sxy_source_flag:

331 float_variables.sxy_source_input = source.sxy

332

333 if integer_variables.sxz_source_flag:

334 float_variables.sxz_source_input = source.sxz

335

336 if integer_variables.syz_source_flag:

337 float_variables.syz_source_input = source.syz

338

339 if integer_variables.p_source_flag:

340 float_variables.p_source_input = source.p

341

342 if integer_variables.transducer_source_flag:

343 float_variables.transducer_source_input = transducer_input_signal

344 integer_variables.delay_mask = delay_mask

345

346 # initial pressure source variable

347 # - this is only defined if the p0 source flag is 1

348 # - this defines the initial pressure everywhere (there is no indicies)

349 if integer_variables.p0_source_flag:

350 float_variables.p0_source_input = source.p0

351

352

353 def grab_sensor_props(integer_variables, kgrid_dim, sensor_mask_index, cuboid_corners_list):

354 # =========================================================================

355 # SENSOR VARIABLES

356 # =========================================================================

357

358 if integer_variables.sensor_mask_type == 0:

359 # mask is defined as a list of grid indices

360 integer_variables.sensor_mask_index = sensor_mask_index

361

362 elif integer_variables.sensor_mask_type == 1:

363

364 cuboid_corners_list = cuboid_corners_list

365 # mask is defined as a list of cuboid corners

366 if kgrid_dim == 2:

367 sensor_mask_corners = np.ones((6, cuboid_corners_list.shape[1]))

368 sensor_mask_corners[0, :] = cuboid_corners_list[0, :]

369 sensor_mask_corners[1, :] = cuboid_corners_list[1, :]

370 sensor_mask_corners[3, :] = cuboid_corners_list[2, :]

371 sensor_mask_corners[4, :] = cuboid_corners_list[3, :]

372 else:

373 sensor_mask_corners = cuboid_corners_list

374 integer_variables.sensor_mask_corners = sensor_mask_corners

375

376 else:

377 raise NotImplementedError('unknown option for sensor_mask_type')

378

379

380 def grab_nonuniform_grid_props(float_variables, kgrid, is_nonuniform_grid):

381 # =========================================================================

382 # VARIABLES USED FOR NONUNIFORM GRIDS

383 # =========================================================================

384

385 # set nonuniform flag and variables

386 # - these are only defined if nonuniform_grid_flag is 1

387 # - these are applied using the bsxfun formulation

388 if not is_nonuniform_grid:

389 return

390

391 dxudxn = kgrid.dudn.x

392 if np.array(dxudxn).size == 1:

393 dxudxn = np.ones((kgrid.Nx, 1))

394 float_variables['dxudxn'] = dxudxn

395

396 dyudyn = kgrid.dudn.y

397 if np.array(dyudyn).size == 1:

398 dyudyn = np.ones((1, kgrid.Ny))

399 float_variables['dyudyn'] = dyudyn

400

401 dzudzn = kgrid.dudn.z

402 if np.array(dzudzn).size == 1:

403 dzudzn = np.ones((1, 1, kgrid.Nz))

404 float_variables['dzudzn'] = dzudzn

405

406 dxudxn_sgx = kgrid.dudn_sg.x

407 if np.array(dxudxn).size == 1:

408 dxudxn_sgx = np.ones((kgrid.Nx, 1))

409 float_variables['dxudxn_sgx'] = dxudxn_sgx

410

411 dyudyn_sgy = kgrid.dudn_sg.y

412 if np.array(dyudyn).size == 1:

413 dyudyn_sgy = np.ones((1, kgrid.Ny))

414 float_variables['dyudyn_sgy'] = dyudyn_sgy

415

416 dzudzn_sgz = kgrid.dudn_sg.z

417 if np.array(dzudzn).size == 1:

418 dzudzn_sgz = np.ones((1, 1, kgrid.Nz))

419 float_variables['dzudzn_sgz'] = dzudzn_sgz

420

421

422 def remove_z_dimension(float_variables, kgrid_dim):

423 # remove z-dimension variables for saving 2D files

424 if kgrid_dim == 2:

425 for k in list(float_variables.keys()):

426 if 'z' in k:

427 del float_variables[k]

428

429

430 def enforce_filename_standards(filepath):

431 # check for HDF5 filename extension

432 filename_ext = os.path.splitext(filepath)[1]

433

434 # use .h5 as default if no extension is given

435 if len(filename_ext) == 0:

436 filename_ext = '.h5'

437 filepath = filepath + '.h5'

438 return filepath, filename_ext

439

440

441 def save_file(filepath, integer_variables, float_variables, hdf_compression_level):

442 filepath, filename_ext = enforce_filename_standards(filepath)

443

444 # save file

445 if filename_ext == '.h5':

446 save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level)

447

448 elif filename_ext == '.mat':

449 save_mat_file(filepath, integer_variables, float_variables)

450 else:

451 # throw error for unknown filetype

452 raise NotImplementedError('unknown file extension for ''SaveToDisk'' filename')

453

454

455 def save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level):

456 # ----------------

457 # SAVE HDF5 FILE

458 # ----------------

459

460 # check if file exists, and delete if it does (the hdf5 library will

461 # give an error if the file already exists)

462 if os.path.exists(filepath):

463 os.remove(filepath)

464

465 # change all the variables to be in single precision (float in C++),

466 # then add to HDF5 File

467 for key, value in float_variables.items():

468 # cast matrix to single precision

469 value = np.array(value, dtype=np.float32)

470 write_matrix(filepath, value, key, hdf_compression_level)

471 del value

472

473 # change all the index variables to be in 64-bit unsigned integers

474 # (long in C++), then add to HDF5 file

475 for key, value in integer_variables.items():

476 # cast matrix to 64-bit unsigned integer

477 value = np.array(value, dtype=np.uint64)

478 write_matrix(filepath, value, key, hdf_compression_level)

479 del value

480

481 # set additional file attributes

482 write_attributes(filepath, legacy=True) # TODO: update to currently breaking code after references are updated

483

484

485 def save_mat_file(filepath, integer_variables, float_variables):

486 # ----------------

487 # SAVE .MAT FILE

488 # ----------------

489

490 # change all the variables to be in single precision (float in C++)

491 for key, value in float_variables.items():

492 float_variables[key] = np.array(value, dtype=np.float32)

493

494 for key, value in integer_variables.items():

495 integer_variables[key] = np.array(value, dtype=np.uin64)

496

497 # save the input variables to disk as a MATLAB binary file

498 float_variables = dict(**float_variables, **integer_variables)

499 savemat(filepath, float_variables)