Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/scenario/targets/waveform.py: 91%

249 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-10-06 06:59 +0000

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Synthetic seismic waveform generator. 

8''' 

9 

10import hashlib 

11import math 

12import logging 

13import numpy as num 

14 

15from os import path as op 

16from functools import reduce 

17 

18from pyrocko.guts import StringChoice, Float, List, Bool 

19from pyrocko.gui.snuffler.marker import PhaseMarker, EventMarker 

20from pyrocko import gf, model, util, trace, io 

21from pyrocko.response import DifferentiationResponse 

22from pyrocko.io.io_common import FileSaveError 

23from pyrocko import pile 

24 

25from ..station import StationGenerator, RandomStationGenerator 

26from .base import TargetGenerator, NoiseGenerator 

27from ..error import ScenarioError 

28 

29 

30DEFAULT_STORE_ID = 'global_2s' 

31 

32logger = logging.getLogger('pyrocko.scenario.targets.waveform') 

33guts_prefix = 'pf.scenario' 

34 

35 

36class WaveformNoiseGenerator(NoiseGenerator): 

37 

38 def get_time_increment(self, deltat): 

39 return deltat * 1024 

40 

41 def get_intersecting_snippets(self, deltat, codes, tmin, tmax): 

42 raise NotImplementedError() 

43 

44 def add_noise(self, tr): 

45 for ntr in self.get_intersecting_snippets( 

46 tr.deltat, tr.nslc_id, tr.tmin, tr.tmax): 

47 tr.add(ntr) 

48 

49 

50class WhiteNoiseGenerator(WaveformNoiseGenerator): 

51 

52 scale = Float.T(default=1e-6) 

53 

54 def get_seed_offset2(self, deltat, iw, codes): 

55 m = hashlib.sha1(('%e %i %s.%s.%s.%s' % ((deltat, iw) + codes)) 

56 .encode('utf8')) 

57 return int(m.hexdigest(), base=16) % 10000000 

58 

59 def get_intersecting_snippets(self, deltat, codes, tmin, tmax): 

60 tinc = self.get_time_increment(deltat) 

61 iwmin = int(math.floor(tmin / tinc)) 

62 iwmax = int(math.floor(tmax / tinc)) 

63 

64 trs = [] 

65 for iw in range(iwmin, iwmax+1): 

66 seed_offset = self.get_seed_offset2(deltat, iw, codes) 

67 rstate = self.get_rstate(seed_offset) 

68 

69 n = int(round(tinc // deltat)) 

70 

71 trs.append(trace.Trace( 

72 codes[0], codes[1], codes[2], codes[3], 

73 deltat=deltat, 

74 tmin=iw*tinc, 

75 ydata=rstate.normal(loc=0, scale=self.scale, size=n))) 

76 

77 return trs 

78 

79 

80class WaveformGenerator(TargetGenerator): 

81 

82 station_generators = List.T( 

83 StationGenerator.T(), 

84 default=[RandomStationGenerator.D()], 

85 help='List of StationGenerators.') 

86 

87 noise_generator = WaveformNoiseGenerator.T( 

88 default=WhiteNoiseGenerator.D(), 

89 help='Add Synthetic noise on the waveforms.') 

90 

91 store_id = gf.StringID.T( 

92 default=DEFAULT_STORE_ID, 

93 help='The GF store to use for forward-calculations.') 

94 

95 seismogram_quantity = StringChoice.T( 

96 choices=['displacement', 'velocity', 'acceleration', 'counts'], 

97 default='displacement') 

98 

99 vmin_cut = Float.T( 

100 default=2000., 

101 help='Minimum velocity to seismic velocity to consider in the model.') 

102 vmax_cut = Float.T( 

103 default=8000., 

104 help='Maximum velocity to seismic velocity to consider in the model.') 

105 

106 fmin = Float.T( 

107 default=0.01, 

108 help='Minimum frequency/wavelength to resolve in the' 

109 ' synthetic waveforms.') 

110 

111 tabulated_phases = List.T( 

112 gf.meta.TPDef.T(), optional=True, 

113 help='Define seismic phases to be calculated.') 

114 

115 tabulated_phases_from_store = Bool.T( 

116 default=False, 

117 help='Calculate seismic phase arrivals for all travel-time tables ' 

118 'defined in GF store.') 

119 

120 tabulated_phases_noise_scale = Float.T( 

121 default=0.0, 

122 help='Standard deviation of normally distributed noise added to ' 

123 'calculated phase arrivals.') 

124 

125 taper = trace.Taper.T( 

126 optional=True, 

127 help='Time domain taper applied to synthetic waveforms.') 

128 

129 compensate_synthetic_offsets = Bool.T( 

130 default=False, 

131 help='Center synthetic trace amplitudes using mean of waveform tips.') 

132 

133 tinc = Float.T( 

134 optional=True, 

135 help='Time increment of waveforms.') 

136 

137 continuous = Bool.T( 

138 default=True, 

139 help='Only produce traces that intersect with events.') 

140 

141 def __init__(self, *args, **kwargs): 

142 super(WaveformGenerator, self).__init__(*args, **kwargs) 

143 self._targets = [] 

144 self._piles = {} 

145 

146 def _get_pile(self, path): 

147 apath = op.abspath(path) 

148 assert op.isdir(apath) 

149 

150 if apath not in self._piles: 

151 fns = util.select_files( 

152 [apath], show_progress=False) 

153 

154 p = pile.Pile() 

155 if fns: 

156 p.load_files(fns, fileformat='mseed', show_progress=False) 

157 

158 self._piles[apath] = p 

159 

160 return self._piles[apath] 

161 

162 def get_stations(self): 

163 stations = [] 

164 for station_generator in self.station_generators: 

165 stations.extend(station_generator.get_stations()) 

166 return stations 

167 

168 def get_distance_range(self, sources): 

169 distances = num.array( 

170 [sg.get_distance_range(sources) 

171 for sg in self.station_generators]) 

172 return (distances[:, 0].min(), distances[:, 1].max()) 

173 

174 def get_targets(self): 

175 if self._targets: 

176 return self._targets 

177 

178 for station in self.get_stations(): 

179 channel_data = [] 

180 channels = station.get_channels() 

181 if channels: 

182 for channel in channels: 

183 channel_data.append([ 

184 channel.name, 

185 channel.azimuth, 

186 channel.dip]) 

187 

188 else: 

189 for c_name in ['BHZ', 'BHE', 'BHN']: 

190 channel_data.append([ 

191 c_name, 

192 model.guess_azimuth_from_name(c_name), 

193 model.guess_dip_from_name(c_name)]) 

194 

195 for c_name, c_azi, c_dip in channel_data: 

196 

197 target = gf.Target( 

198 codes=( 

199 station.network, 

200 station.station, 

201 station.location, 

202 c_name), 

203 quantity='displacement', 

204 lat=station.lat, 

205 lon=station.lon, 

206 north_shift=station.north_shift, 

207 east_shift=station.east_shift, 

208 depth=station.depth, 

209 store_id=self.store_id, 

210 optimization='enable', 

211 interpolation='nearest_neighbor', 

212 azimuth=c_azi, 

213 dip=c_dip) 

214 

215 self._targets.append(target) 

216 

217 return self._targets 

218 

219 def get_time_range(self, sources): 

220 dmin, dmax = self.get_distance_range(sources) 

221 

222 times = num.array([source.time for source in sources], 

223 dtype=util.get_time_dtype()) 

224 

225 tmin_events = num.min(times) 

226 tmax_events = num.max(times) 

227 

228 tmin = tmin_events + dmin / self.vmax_cut - 10.0 / self.fmin 

229 tmax = tmax_events + dmax / self.vmin_cut + 10.0 / self.fmin 

230 

231 return tmin, tmax 

232 

233 def get_codes_to_deltat(self, engine, sources): 

234 deltats = {} 

235 

236 targets = self.get_targets() 

237 for source in sources: 

238 for target in targets: 

239 deltats[target.codes] = engine.get_store( 

240 target.store_id).config.deltat 

241 

242 return deltats 

243 

244 def get_useful_time_increment(self, engine, sources): 

245 _, dmax = self.get_distance_range(sources) 

246 tinc = dmax / self.vmin_cut + 2.0 / self.fmin 

247 

248 deltats = set(self.get_codes_to_deltat(engine, sources).values()) 

249 deltat = reduce(util.lcm, deltats) 

250 tinc = int(round(tinc / deltat)) * deltat 

251 return tinc 

252 

253 def get_relevant_sources(self, sources, tmin, tmax): 

254 dmin, dmax = self.get_distance_range(sources) 

255 trange = tmax - tmin 

256 tmax_pad = trange + tmax + dmin / self.vmax_cut 

257 tmin_pad = tmin - (dmax / self.vmin_cut + trange) 

258 

259 return [s for s in sources if s.time < tmax_pad and s.time > tmin_pad] 

260 

261 def get_waveforms(self, engine, sources, tmin, tmax): 

262 

263 sources_relevant = self.get_relevant_sources(sources, tmin, tmax) 

264 if not (self.continuous or sources_relevant): 

265 return [] 

266 

267 trs = {} 

268 tts = util.time_to_str 

269 

270 for nslc, deltat in self.get_codes_to_deltat(engine, sources).items(): 

271 tr_tmin = round(tmin / deltat) * deltat 

272 tr_tmax = (round(tmax / deltat)-1) * deltat 

273 nsamples = int(round((tr_tmax - tr_tmin) / deltat)) + 1 

274 

275 tr = trace.Trace( 

276 *nslc, 

277 tmin=tr_tmin, 

278 ydata=num.zeros(nsamples), 

279 deltat=deltat) 

280 

281 self.noise_generator.add_noise(tr) 

282 

283 trs[nslc] = tr 

284 

285 logger.debug('Forward modelling waveforms between %s - %s...' 

286 % (tts(tmin, format='%Y-%m-%d_%H-%M-%S'), 

287 tts(tmax, format='%Y-%m-%d_%H-%M-%S'))) 

288 

289 if not sources_relevant: 

290 return list(trs.values()) 

291 

292 targets = self.get_targets() 

293 response = engine.process(sources_relevant, targets) 

294 for source, target, res in response.iter_results( 

295 get='results'): 

296 

297 if isinstance(res, gf.SeismosizerError): 

298 logger.warning( 

299 'Out of bounds! \nTarget: %s\nSource: %s\n' % ( 

300 '.'.join(target.codes)), source) 

301 continue 

302 

303 tr = res.trace.pyrocko_trace() 

304 

305 candidate = trs[target.codes] 

306 if not candidate.overlaps(tr.tmin, tr.tmax): 

307 continue 

308 

309 if self.compensate_synthetic_offsets: 

310 tr.ydata -= (num.mean(tr.ydata[-3:-1]) + 

311 num.mean(tr.ydata[1:3])) / 2. 

312 

313 if self.taper: 

314 tr.taper(self.taper) 

315 

316 resp = self.get_transfer_function(target.codes) 

317 if resp: 

318 tr = tr.transfer(transfer_function=resp) 

319 

320 candidate.add(tr) 

321 trs[target.codes] = candidate 

322 

323 return list(trs.values()) 

324 

325 def get_onsets(self, engine, sources, *args, **kwargs): 

326 

327 targets = {t.codes[:3]: t for t in self.get_targets()} 

328 

329 markers = [] 

330 for source in sources: 

331 ev = source.pyrocko_event() 

332 markers.append(EventMarker(ev)) 

333 for nsl, target in targets.items(): 

334 store = engine.get_store(target.store_id) 

335 if self.tabulated_phases: 

336 tabulated_phases = self.tabulated_phases 

337 

338 elif self.tabulated_phases_from_store: 

339 tabulated_phases = store.config.tabulated_phases 

340 else: 

341 tabulated_phases = [] 

342 

343 for phase in tabulated_phases: 

344 t = store.t(phase.id, source, target) 

345 if not t: 

346 continue 

347 

348 noise_scale = self.tabulated_phases_noise_scale 

349 if noise_scale != 0.0: 

350 t += num.random.normal(scale=noise_scale) 

351 

352 t += source.time 

353 markers.append( 

354 PhaseMarker( 

355 phasename=phase.id, 

356 tmin=t, 

357 tmax=t, 

358 event=source.pyrocko_event(), 

359 nslc_ids=(nsl+('*',),) 

360 ) 

361 ) 

362 return markers 

363 

364 def get_transfer_function(self, codes): 

365 if self.seismogram_quantity == 'displacement': 

366 return None 

367 elif self.seismogram_quantity == 'velocity': 

368 return DifferentiationResponse(1) 

369 elif self.seismogram_quantity == 'acceleration': 

370 return DifferentiationResponse(2) 

371 elif self.seismogram_quantity == 'counts': 

372 raise NotImplementedError() 

373 

374 def ensure_data(self, engine, sources, path, tmin=None, tmax=None): 

375 self.ensure_waveforms(engine, sources, path, tmin, tmax) 

376 self.ensure_responses(path) 

377 

378 def ensure_waveforms(self, engine, sources, path, tmin=None, tmax=None): 

379 

380 path_waveforms = op.join(path, 'waveforms') 

381 util.ensuredir(path_waveforms) 

382 

383 p = self._get_pile(path_waveforms) 

384 

385 nslc_ids = set(target.codes for target in self.get_targets()) 

386 

387 def have_waveforms(tmin, tmax): 

388 trs_have = p.all( 

389 tmin=tmin, tmax=tmax, 

390 load_data=False, degap=False, 

391 trace_selector=lambda tr: tr.nslc_id in nslc_ids) 

392 

393 return any(tr.data_len() > 0 for tr in trs_have) 

394 

395 def add_files(paths): 

396 p.load_files(paths, fileformat='mseed', show_progress=False) 

397 

398 path_traces = op.join( 

399 path_waveforms, 

400 '%(wmin_year)s', 

401 '%(wmin_month)s', 

402 '%(wmin_day)s', 

403 'waveform_%(network)s_%(station)s_' + 

404 '%(location)s_%(channel)s_%(tmin)s_%(tmax)s.mseed') 

405 

406 tmin_all, tmax_all = self.get_time_range(sources) 

407 tmin = tmin if tmin is not None else tmin_all 

408 tmax = tmax if tmax is not None else tmax_all 

409 tts = util.time_to_str 

410 

411 tinc = self.tinc or self.get_useful_time_increment(engine, sources) 

412 tmin = num.floor(tmin / tinc) * tinc 

413 tmax = num.ceil(tmax / tinc) * tinc 

414 

415 nwin = int(round((tmax - tmin) / tinc)) 

416 

417 pbar = None 

418 try: 

419 for iwin in range(nwin): 

420 tmin_win = tmin + iwin*tinc 

421 tmax_win = tmin + (iwin+1)*tinc 

422 

423 if have_waveforms(tmin_win, tmax_win): 

424 continue 

425 

426 if pbar is None: 

427 pbar = util.progressbar( 

428 'Generating waveforms', (nwin-iwin)) 

429 

430 pbar.update(iwin) 

431 

432 trs = self.get_waveforms(engine, sources, tmin_win, tmax_win) 

433 

434 try: 

435 wpaths = io.save( 

436 trs, path_traces, 

437 additional=dict( 

438 wmin_year=tts(tmin_win, format='%Y'), 

439 wmin_month=tts(tmin_win, format='%m'), 

440 wmin_day=tts(tmin_win, format='%d'), 

441 wmin=tts(tmin_win, format='%Y-%m-%d_%H-%M-%S'), 

442 wmax_year=tts(tmax_win, format='%Y'), 

443 wmax_month=tts(tmax_win, format='%m'), 

444 wmax_day=tts(tmax_win, format='%d'), 

445 wmax=tts(tmax_win, format='%Y-%m-%d_%H-%M-%S'))) 

446 

447 for wpath in wpaths: 

448 logger.debug('Generated file: %s' % wpath) 

449 

450 add_files(wpaths) 

451 

452 except FileSaveError as e: 

453 raise ScenarioError(str(e)) 

454 

455 finally: 

456 if pbar is not None: 

457 pbar.finish() 

458 

459 def ensure_responses(self, path): 

460 from pyrocko.io import stationxml 

461 

462 path_responses = op.join(path, 'meta') 

463 util.ensuredir(path_responses) 

464 

465 fn_stationxml = op.join(path_responses, 'waveform_response.xml') 

466 i = 0 

467 while op.exists(fn_stationxml): 

468 fn_stationxml = op.join( 

469 path_responses, 'waveform_response-%s.xml' % i) 

470 i += 1 

471 

472 logger.debug('Writing waveform meta information to StationXML...') 

473 

474 stations = self.get_stations() 

475 sxml = stationxml.FDSNStationXML.from_pyrocko_stations(stations) 

476 

477 sunit = { 

478 'displacement': 'M', 

479 'velocity': 'M/S', 

480 'acceleration': 'M/S**2', 

481 'counts': 'COUNTS'}[self.seismogram_quantity] 

482 

483 response = stationxml.Response( 

484 instrument_sensitivity=stationxml.Sensitivity( 

485 value=1., 

486 frequency=1., 

487 input_units=stationxml.Units(sunit), 

488 output_units=stationxml.Units('COUNTS')), 

489 stage_list=[]) 

490 

491 for net, station, channel in sxml.iter_network_station_channels(): 

492 channel.response = response 

493 

494 sxml.dump_xml(filename=fn_stationxml) 

495 

496 def add_map_artists(self, engine, sources, automap): 

497 automap.add_stations(self.get_stations())