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

213 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2024-03-07 11:54 +0000

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Implementation of the scenario generator's main class. 

8''' 

9 

10import logging 

11import sys 

12import os 

13import os.path as op 

14import shutil 

15import random 

16from functools import wraps 

17 

18import numpy as num 

19 

20from pyrocko.guts import List 

21from pyrocko.plot import gmtpy 

22from pyrocko.gui.snuffler.marker import PhaseMarker 

23from pyrocko import pile, util, model 

24from pyrocko.dataset import topo 

25 

26from .error import ScenarioError, CannotCreatePath, LocationGenerationError 

27from .base import LocationGenerator 

28from .sources import SourceGenerator, DCSourceGenerator 

29from .targets import TargetGenerator, AVAILABLE_TARGETS 

30 

31logger = logging.getLogger('pyrocko.scenario') 

32guts_prefix = 'pf.scenario' 

33 

34km = 1000. 

35 

36 

37class ScenarioGenerator(LocationGenerator): 

38 ''' 

39 Generator for synthetic seismo-geodetic scenarios. 

40 ''' 

41 

42 target_generators = List.T( 

43 TargetGenerator.T(), 

44 default=[], 

45 help='Targets to spawn in the scenario.') 

46 

47 source_generator = SourceGenerator.T( 

48 default=DCSourceGenerator.D(), 

49 help='Sources to spawn in the scenario.') 

50 

51 def __init__(self, **kwargs): 

52 LocationGenerator.__init__(self, **kwargs) 

53 for itry in range(self.ntries): 

54 

55 try: 

56 self.get_stations() 

57 self.get_sources() 

58 return 

59 

60 except LocationGenerationError: 

61 self.retry() 

62 

63 raise ScenarioError( 

64 'Could not generate scenario in %i tries.' % self.ntries) 

65 

66 def init_modelling(self, engine): 

67 self._engine = engine 

68 

69 def get_engine(self): 

70 return self._engine 

71 

72 def get_sources(self): 

73 return self.source_generator.get_sources() 

74 

75 def get_events(self): 

76 return [s.pyrocko_event() for s in self.get_sources()] 

77 

78 def collect(collector): 

79 if not callable(collector): 

80 raise AttributeError('This method should not be called directly.') 

81 

82 @wraps(collector) 

83 def method(self, *args, **kwargs): 

84 result = [] 

85 func = collector(self, *args, **kwargs) 

86 for gen in self.target_generators: 

87 result.extend( 

88 func(gen) or []) 

89 return result 

90 

91 return method 

92 

93 @collect 

94 def get_stations(self): 

95 return lambda gen: gen.get_stations() 

96 

97 @collect 

98 def get_targets(self): 

99 return lambda gen: gen.get_targets() 

100 

101 @collect 

102 def get_waveforms(self, tmin=None, tmax=None): 

103 tmin, tmax = self._time_range_fill_defaults(tmin, tmax) 

104 return lambda gen: gen.get_waveforms( 

105 self._engine, self.get_sources(), 

106 tmin=tmin, tmax=tmax) 

107 

108 @collect 

109 def get_onsets(self, tmin=None, tmax=None): 

110 tmin, tmax = self._time_range_fill_defaults(tmin, tmax) 

111 return lambda gen: gen.get_onsets( 

112 self._engine, self.get_sources(), 

113 tmin=tmin, tmax=tmax) 

114 

115 @collect 

116 def get_insar_scenes(self, tmin=None, tmax=None): 

117 tmin, tmax = self._time_range_fill_defaults(tmin, tmax) 

118 return lambda gen: gen.get_insar_scenes( 

119 self._engine, self.get_sources(), 

120 tmin=tmin, tmax=tmax) 

121 

122 @collect 

123 def get_gnss_campaigns(self, tmin=None, tmax=None): 

124 tmin, tmax = self._time_range_fill_defaults(tmin, tmax) 

125 return lambda gen: gen.get_gnss_campaigns( 

126 self._engine, self.get_sources(), 

127 tmin=tmin, tmax=tmax) 

128 

129 def dump_data(self, path, overwrite=False): 

130 ''' 

131 Invoke generators and dump the complete scenario. 

132 

133 :param path: Output directory 

134 :type path: str 

135 :param overwrite: If ``True`` remove all previously generated files 

136 and recreating the scenario content. If ``False`` stop if 

137 previously generated content is found. 

138 :type overwrite: bool 

139 

140 Wrapper to call :py:meth:`prepare_data` followed by 

141 :py:meth:`ensure_data` with default arguments. 

142 ''' 

143 

144 self.prepare_data(path, overwrite=False) 

145 self.ensure_data(path) 

146 

147 def prepare_data(self, path, overwrite): 

148 ''' 

149 Prepare directory for scenario content storage. 

150 

151 :param path: Output directory 

152 :type path: str 

153 :param overwrite: If ``True``, remove all previously generated files 

154 and recreate the scenario content. If ``False``, stop if 

155 previously generated content is found. 

156 :type overwrite: bool 

157 ''' 

158 

159 for dentry in [ 

160 'meta', 

161 'waveforms', 

162 'insar', 

163 'gnss']: 

164 

165 dpath = op.join(path, dentry) 

166 if op.exists(dpath): 

167 if overwrite: 

168 shutil.rmtree(dpath) 

169 else: 

170 raise CannotCreatePath('Path exists: %s' % dpath) 

171 

172 for fentry in [ 

173 'events.txt', 

174 'sources.yml', 

175 'map.pdf', 

176 'README.md']: 

177 

178 fpath = op.join(path, fentry) 

179 if op.exists(fpath): 

180 if overwrite: 

181 os.unlink(fpath) 

182 else: 

183 raise CannotCreatePath('Path exists: %s' % fpath) 

184 

185 @collect 

186 def ensure_data(self, path, tmin=None, tmax=None): 

187 

188 ''' 

189 Generate and output scenario content to files, as needed. 

190 

191 :param path: Output directory 

192 :type path: str 

193 :param tmin: Start of time interval to generate 

194 :type tmin: :py:func:`~pyrocko.util.get_time_float` 

195 :param tmax: End of time interval to generate 

196 :type tmax: :py:func:`~pyrocko.util.get_time_float` 

197 

198 This method only generates the files which are relevant for the 

199 given time interval, and which have not yet been generated. It is safe 

200 to call this method several times, for different time windows, as 

201 needed. 

202 

203 If no time interval is given, the origin times of the generated sources 

204 and signal propagation times are taken into account to estimate a 

205 reasonable default. 

206 ''' 

207 

208 tmin, tmax = self._time_range_fill_defaults(tmin, tmax) 

209 self.source_generator.ensure_data(path) 

210 

211 meta_dir = op.join(path, 'meta') 

212 util.ensuredir(meta_dir) 

213 

214 fn_stations = op.join(meta_dir, 'stations.txt') 

215 if not op.exists(fn_stations): 

216 model.station.dump_stations( 

217 self.get_stations(), fn_stations) 

218 

219 fn_stations_yaml = op.join(meta_dir, 'stations.yml') 

220 if not op.exists(fn_stations_yaml): 

221 model.station.dump_stations_yaml( 

222 self.get_stations(), fn_stations_yaml) 

223 

224 fn_stations_kml = op.join(meta_dir, 'stations.kml') 

225 if not op.exists(fn_stations_kml): 

226 model.station.dump_kml( 

227 self.get_stations(), fn_stations_kml) 

228 

229 fn_markers = op.join(meta_dir, 'markers.txt') 

230 if not op.exists(fn_markers): 

231 markers = self.get_onsets() 

232 if markers: 

233 PhaseMarker.save_markers(markers, fn_markers) 

234 

235 fn_readme = op.join(path, 'README.md') 

236 if not op.exists(fn_readme): 

237 dump_readme(fn_readme) 

238 

239 def ensure_data(gen): 

240 logger.info('Creating files from %s...' % gen.__class__.__name__) 

241 return gen.ensure_data( 

242 self._engine, self.get_sources(), path, tmin=tmin, tmax=tmax) 

243 

244 return ensure_data 

245 

246 @collect 

247 def _get_time_ranges(self): 

248 return lambda gen: [gen.get_time_range(self.get_sources())] 

249 

250 def get_time_range(self): 

251 ranges = num.array( 

252 self._get_time_ranges(), dtype=util.get_time_dtype()) 

253 return ranges.min(), ranges.max() 

254 

255 def _time_range_fill_defaults(self, tmin, tmax): 

256 stmin, stmax = self.get_time_range() 

257 return stmin if tmin is None else tmin, stmax if tmax is None else tmax 

258 

259 def get_pile(self, tmin=None, tmax=None): 

260 p = pile.Pile() 

261 

262 trf = pile.MemTracesFile(None, self.get_waveforms(tmin, tmax)) 

263 p.add_file(trf) 

264 return p 

265 

266 def make_map(self, filename): 

267 logger.info("Plotting scenario's map...") 

268 if not gmtpy.have_gmt(): 

269 logger.warning('Cannot plot map, GMT is not installed.') 

270 return 

271 if self.radius is None or self.radius > 5000*km: 

272 logger.info( 

273 'Drawing map for scenarios with a radius > 5000 km is ' 

274 'not implemented.') 

275 return 

276 

277 try: 

278 draw_scenario_gmt(self, filename) 

279 except gmtpy.GMTError: 

280 logger.warning('GMT threw an error, could not plot map.') 

281 except topo.AuthenticationRequired: 

282 logger.warning('Cannot download topography data (authentication ' 

283 'required). Could not plot map.') 

284 

285 def draw_map(self, fns): 

286 from pyrocko.plot import automap 

287 

288 if isinstance(fns, str): 

289 fns = [fns] 

290 

291 lat, lon = self.get_center_latlon() 

292 radius = self.get_radius() 

293 

294 m = automap.Map( 

295 width=30., 

296 height=30., 

297 lat=lat, 

298 lon=lon, 

299 radius=radius, 

300 show_topo=True, 

301 show_grid=True, 

302 show_rivers=True, 

303 color_wet=(216, 242, 254), 

304 color_dry=(238, 236, 230) 

305 ) 

306 

307 self.source_generator.add_map_artists(m) 

308 

309 sources = self.get_sources() 

310 for gen in self.target_generators: 

311 gen.add_map_artists(self.get_engine(), sources, m) 

312 

313 # for patch in self.get_insar_patches(): 

314 # symbol_size = 50. 

315 # coords = num.array(patch.get_corner_coordinates()) 

316 # m.gmt.psxy(in_rows=num.fliplr(coords), 

317 # L=True, 

318 # *m.jxyr) 

319 

320 for fn in fns: 

321 m.save(fn) 

322 

323 @property 

324 def stores_wanted(self): 

325 return set([gen.store_id for gen in self.target_generators 

326 if hasattr(gen, 'store_id')]) 

327 

328 @property 

329 def stores_missing(self): 

330 return self.stores_wanted - set(self.get_engine().get_store_ids()) 

331 

332 def ensure_gfstores(self, interactive=False): 

333 if not self.stores_missing: 

334 return 

335 

336 from pyrocko.gf import ws 

337 

338 engine = self.get_engine() 

339 

340 gf_store_superdirs = engine.store_superdirs 

341 

342 if interactive: 

343 print("We could not find the following Green's function stores:\n" 

344 '%s\n' 

345 'We can try to download the stores from ' 

346 'http://kinherd.org into one of the following ' 

347 'directories:' 

348 % '\n'.join(' ' + s for s in self.stores_missing)) 

349 for idr, dr in enumerate(gf_store_superdirs): 

350 print(' %d. %s' % ((idr+1), dr)) 

351 s = input('\nInto which directory should we download the GF ' 

352 'store(s)?\nDefault 1, (C)ancel: ') 

353 if s in ['c', 'C']: 

354 print('Canceled.') 

355 sys.exit(1) 

356 elif s == '': 

357 s = 0 

358 try: 

359 s = int(s) 

360 if s > len(gf_store_superdirs): 

361 raise ValueError 

362 except ValueError: 

363 print('Invalid selection: %s' % s) 

364 sys.exit(1) 

365 else: 

366 s = 1 

367 

368 download_dir = gf_store_superdirs[s-1] 

369 util.ensuredir(download_dir) 

370 logger.info("Downloading Green's functions stores to %s" 

371 % download_dir) 

372 

373 oldwd = os.getcwd() 

374 for store in self.stores_missing: 

375 os.chdir(download_dir) 

376 ws.download_gf_store(site='kinherd', store_id=store) 

377 

378 os.chdir(oldwd) 

379 

380 @classmethod 

381 def initialize( 

382 cls, path, 

383 center_lat=None, center_lon=None, radius=None, 

384 targets=AVAILABLE_TARGETS, stationxml=None, force=False): 

385 ''' 

386 Initialize a Scenario and create a ``scenario.yml`` 

387 

388 :param path: Path to create the scenerio in 

389 :type path: str 

390 :param center_lat: Center latitude [deg] 

391 :type center_lat: float 

392 :param center_lon: Center longitude [deg] 

393 :type center_lon: float 

394 :param radius: Scenario's radius in [m] 

395 :type radius: float 

396 :param targets: Targets to throw into scenario, 

397 defaults to AVAILABLE_TARGETS 

398 :type targets: list of 

399 :class:`pyrocko.scenario.targets.base.TargetGenerator` objects 

400 :param force: If set to ``True``, overwrite directory 

401 :type force: bool 

402 :param stationxml: path to a StationXML to be used by the 

403 :class:`pyrocko.scenario.targets.waveform.WaveformGenerator`. 

404 :type stationxml: str 

405 :returns: Scenario 

406 :rtype: :class:`pyrocko.scenario.scenario.ScenarioGenerator` 

407 ''' 

408 import os.path as op 

409 

410 if op.exists(path) and not force: 

411 raise CannotCreatePath('Directory %s alread exists.' % path) 

412 

413 util.ensuredir(path) 

414 fn = op.join(path, 'scenario.yml') 

415 logger.debug('Writing new scenario to %s' % fn) 

416 

417 scenario = cls( 

418 center_lat=center_lat, 

419 center_lon=center_lon, 

420 radius=radius) 

421 

422 scenario.seed = random.randint(1, 2**32-1) 

423 

424 scenario.target_generators.extend([t() for t in targets]) 

425 

426 scenario.update_hierarchy() 

427 

428 scenario.dump(filename=fn) 

429 scenario.prepare_data(path, overwrite=force) 

430 

431 return scenario 

432 

433 

434def draw_scenario_gmt(generator, fn): 

435 return generator.draw_map(fn) 

436 

437 

438def dump_readme(path): 

439 readme = '''# Pyrocko Earthquake Scenario 

440 

441The directory structure of a scenario is layed out as follows: 

442 

443## Map of the scenario 

444A simple map is generated from `pyrocko.automap` in map.pdf 

445 

446## Earthquake Sources 

447 

448Can be found as events.txt and sources.yml hosts the pyrocko.gf sources. 

449 

450## Folder `meta` 

451 

452Contains stations.txt and StationXML files for waveforms as well as KML data. 

453The responses are flat with gain 1.0 at 1.0 Hz. 

454 

455## Folder `waveforms` 

456 

457Waveforms as mini-seed are stored here, segregated into days. 

458 

459## Folder `gnss` 

460 

461The GNSS campaign.yml is living here. 

462Use `pyrocko.guts.load(filename='campaign.yml)` to load the campaign. 

463 

464## Folder `insar` 

465 

466Kite InSAR scenes for ascending and descending tracks are stored there. 

467Use `kite.Scene.load(<filename>)` to inspect the scenes. 

468 

469''' 

470 with open(path, 'w') as f: 

471 f.write(readme) 

472 

473 return [path]