1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import logging 

7import sys 

8import os 

9import os.path as op 

10import shutil 

11import random 

12from functools import wraps 

13 

14import numpy as num 

15 

16from pyrocko.guts import List 

17from pyrocko.plot import gmtpy 

18from pyrocko.gui.snuffler.marker import PhaseMarker 

19from pyrocko import pile, util, model 

20from pyrocko.dataset import topo 

21 

22from .error import ScenarioError, CannotCreatePath, LocationGenerationError 

23from .base import LocationGenerator 

24from .sources import SourceGenerator, DCSourceGenerator 

25from .targets import TargetGenerator, AVAILABLE_TARGETS 

26 

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

28guts_prefix = 'pf.scenario' 

29 

30km = 1000. 

31 

32 

33class ScenarioGenerator(LocationGenerator): 

34 ''' 

35 Generator for synthetic seismo-geodetic scenarios. 

36 ''' 

37 

38 target_generators = List.T( 

39 TargetGenerator.T(), 

40 default=[], 

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

42 

43 source_generator = SourceGenerator.T( 

44 default=DCSourceGenerator.D(), 

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

46 

47 def __init__(self, **kwargs): 

48 LocationGenerator.__init__(self, **kwargs) 

49 for itry in range(self.ntries): 

50 

51 try: 

52 self.get_stations() 

53 self.get_sources() 

54 return 

55 

56 except LocationGenerationError: 

57 self.retry() 

58 

59 raise ScenarioError( 

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

61 

62 def init_modelling(self, engine): 

63 self._engine = engine 

64 

65 def get_engine(self): 

66 return self._engine 

67 

68 def get_sources(self): 

69 return self.source_generator.get_sources() 

70 

71 def get_events(self): 

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

73 

74 def collect(collector): 

75 if not callable(collector): 

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

77 

78 @wraps(collector) 

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

80 result = [] 

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

82 for gen in self.target_generators: 

83 result.extend( 

84 func(gen) or []) 

85 return result 

86 

87 return method 

88 

89 @collect 

90 def get_stations(self): 

91 return lambda gen: gen.get_stations() 

92 

93 @collect 

94 def get_targets(self): 

95 return lambda gen: gen.get_targets() 

96 

97 @collect 

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

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

100 return lambda gen: gen.get_waveforms( 

101 self._engine, self.get_sources(), 

102 tmin=tmin, tmax=tmax) 

103 

104 @collect 

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

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

107 return lambda gen: gen.get_onsets( 

108 self._engine, self.get_sources(), 

109 tmin=tmin, tmax=tmax) 

110 

111 @collect 

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

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

114 return lambda gen: gen.get_insar_scenes( 

115 self._engine, self.get_sources(), 

116 tmin=tmin, tmax=tmax) 

117 

118 @collect 

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

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

121 return lambda gen: gen.get_gnss_campaigns( 

122 self._engine, self.get_sources(), 

123 tmin=tmin, tmax=tmax) 

124 

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

126 ''' 

127 Invoke generators and dump the complete scenario. 

128 

129 :param path: Output directory 

130 :type path: str 

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

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

133 previously generated content is found. 

134 :type overwrite: bool 

135 

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

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

138 ''' 

139 

140 self.prepare_data(path, overwrite=False) 

141 self.ensure_data(path) 

142 

143 def prepare_data(self, path, overwrite): 

144 ''' 

145 Prepare directory for scenario content storage. 

146 

147 :param path: Output directory 

148 :type path: str 

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

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

151 previously generated content is found. 

152 :type overwrite: bool 

153 ''' 

154 

155 for dentry in [ 

156 'meta', 

157 'waveforms', 

158 'insar', 

159 'gnss']: 

160 

161 dpath = op.join(path, dentry) 

162 if op.exists(dpath): 

163 if overwrite: 

164 shutil.rmtree(dpath) 

165 else: 

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

167 

168 for fentry in [ 

169 'events.txt', 

170 'sources.yml', 

171 'map.pdf', 

172 'README.md']: 

173 

174 fpath = op.join(path, fentry) 

175 if op.exists(fpath): 

176 if overwrite: 

177 os.unlink(fpath) 

178 else: 

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

180 

181 @collect 

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

183 

184 ''' 

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

186 

187 :param path: Output directory 

188 :type path: str 

189 :param tmin: Start of time interval to generate 

190 :type tmin: timestamp, optional 

191 :param tmax: End of time interval to generate 

192 :type tmax: timestamp, optional 

193 

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

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

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

197 needed. 

198 

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

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

201 reasonable default. 

202 ''' 

203 

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

205 self.source_generator.ensure_data(path) 

206 

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

208 util.ensuredir(meta_dir) 

209 

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

211 if not op.exists(fn_stations): 

212 model.station.dump_stations( 

213 self.get_stations(), fn_stations) 

214 

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

216 if not op.exists(fn_stations_yaml): 

217 model.station.dump_stations_yaml( 

218 self.get_stations(), fn_stations_yaml) 

219 

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

221 if not op.exists(fn_stations_kml): 

222 model.station.dump_kml( 

223 self.get_stations(), fn_stations_kml) 

224 

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

226 if not op.exists(fn_markers): 

227 markers = self.get_onsets() 

228 if markers: 

229 PhaseMarker.save_markers(markers, fn_markers) 

230 

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

232 if not op.exists(fn_readme): 

233 dump_readme(fn_readme) 

234 

235 def ensure_data(gen): 

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

237 return gen.ensure_data( 

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

239 

240 return ensure_data 

241 

242 @collect 

243 def _get_time_ranges(self): 

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

245 

246 def get_time_range(self): 

247 ranges = num.array( 

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

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

250 

251 def _time_range_fill_defaults(self, tmin, tmax): 

252 stmin, stmax = self.get_time_range() 

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

254 

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

256 p = pile.Pile() 

257 

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

259 p.add_file(trf) 

260 return p 

261 

262 def make_map(self, filename): 

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

264 if not gmtpy.have_gmt(): 

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

266 return 

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

268 logger.info( 

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

270 'not implemented.') 

271 return 

272 

273 try: 

274 draw_scenario_gmt(self, filename) 

275 except gmtpy.GMTError: 

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

277 except topo.AuthenticationRequired: 

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

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

280 

281 def draw_map(self, fns): 

282 from pyrocko.plot import automap 

283 

284 if isinstance(fns, str): 

285 fns = [fns] 

286 

287 lat, lon = self.get_center_latlon() 

288 radius = self.get_radius() 

289 

290 m = automap.Map( 

291 width=30., 

292 height=30., 

293 lat=lat, 

294 lon=lon, 

295 radius=radius, 

296 show_topo=True, 

297 show_grid=True, 

298 show_rivers=True, 

299 color_wet=(216, 242, 254), 

300 color_dry=(238, 236, 230) 

301 ) 

302 

303 self.source_generator.add_map_artists(m) 

304 

305 sources = self.get_sources() 

306 for gen in self.target_generators: 

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

308 

309 # for patch in self.get_insar_patches(): 

310 # symbol_size = 50. 

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

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

313 # L=True, 

314 # *m.jxyr) 

315 

316 for fn in fns: 

317 m.save(fn) 

318 

319 @property 

320 def stores_wanted(self): 

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

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

323 

324 @property 

325 def stores_missing(self): 

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

327 

328 def ensure_gfstores(self, interactive=False): 

329 if not self.stores_missing: 

330 return 

331 

332 from pyrocko.gf import ws 

333 

334 engine = self.get_engine() 

335 

336 gf_store_superdirs = engine.store_superdirs 

337 

338 if interactive: 

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

340 '%s\n' 

341 'We can try to download the stores from ' 

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

343 'directories:' 

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

345 for idr, dr in enumerate(gf_store_superdirs): 

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

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

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

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

350 print('Canceled.') 

351 sys.exit(1) 

352 elif s == '': 

353 s = 0 

354 try: 

355 s = int(s) 

356 if s > len(gf_store_superdirs): 

357 raise ValueError 

358 except ValueError: 

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

360 sys.exit(1) 

361 else: 

362 s = 1 

363 

364 download_dir = gf_store_superdirs[s-1] 

365 util.ensuredir(download_dir) 

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

367 % download_dir) 

368 

369 oldwd = os.getcwd() 

370 for store in self.stores_missing: 

371 os.chdir(download_dir) 

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

373 

374 os.chdir(oldwd) 

375 

376 @classmethod 

377 def initialize( 

378 cls, path, 

379 center_lat=None, center_lon=None, radius=None, 

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

381 ''' 

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

383 

384 :param path: Path to create the scenerio in 

385 :type path: str 

386 :param center_lat: Center latitude [deg] 

387 :type center_lat: float, optional 

388 :param center_lon: Center longitude [deg] 

389 :type center_lon: float, optional 

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

391 :type radius: float, optional 

392 :param targets: Targets to throw into scenario, 

393 defaults to AVAILABLE_TARGETS 

394 :type targets: list of :class:`pyrocko.scenario.TargetGenerator` 

395 objects, optional 

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

397 :type force: bool 

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

399 :class:`pyrocko.scenario.targets.WaveformGenerator`. 

400 :type stationxml: str 

401 :returns: Scenario 

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

403 ''' 

404 import os.path as op 

405 

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

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

408 

409 util.ensuredir(path) 

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

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

412 

413 scenario = cls( 

414 center_lat=center_lat, 

415 center_lon=center_lon, 

416 radius=radius) 

417 

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

419 

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

421 

422 scenario.update_hierarchy() 

423 

424 scenario.dump(filename=fn) 

425 scenario.prepare_data(path, overwrite=force) 

426 

427 return scenario 

428 

429 

430def draw_scenario_gmt(generator, fn): 

431 return generator.draw_map(fn) 

432 

433 

434def dump_readme(path): 

435 readme = '''# Pyrocko Earthquake Scenario 

436 

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

438 

439## Map of the scenario 

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

441 

442## Earthquake Sources 

443 

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

445 

446## Folder `meta` 

447 

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

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

450 

451## Folder `waveforms` 

452 

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

454 

455## Folder `gnss` 

456 

457The GNSS campaign.yml is living here. 

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

459 

460## Folder `insar` 

461 

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

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

464 

465''' 

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

467 f.write(readme) 

468 

469 return [path]