1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5from __future__ import absolute_import, division, print_function 

6 

7import logging 

8import sys 

9import os 

10import os.path as op 

11import shutil 

12import random 

13from functools import wraps 

14 

15import numpy as num 

16 

17from pyrocko.guts import List 

18from pyrocko.plot import gmtpy 

19from pyrocko.gui.marker import PhaseMarker 

20from pyrocko import pile, util, model 

21from pyrocko.dataset import topo 

22 

23from .error import ScenarioError, CannotCreatePath, LocationGenerationError 

24from .base import LocationGenerator 

25from .sources import SourceGenerator, DCSourceGenerator 

26from .targets import TargetGenerator, AVAILABLE_TARGETS 

27 

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

29guts_prefix = 'pf.scenario' 

30 

31km = 1000. 

32 

33 

34class ScenarioGenerator(LocationGenerator): 

35 ''' 

36 Generator for synthetic seismo-geodetic scenarios. 

37 ''' 

38 

39 target_generators = List.T( 

40 TargetGenerator.T(), 

41 default=[], 

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

43 

44 source_generator = SourceGenerator.T( 

45 default=DCSourceGenerator.D(), 

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

47 

48 def __init__(self, **kwargs): 

49 LocationGenerator.__init__(self, **kwargs) 

50 for itry in range(self.ntries): 

51 

52 try: 

53 self.get_stations() 

54 self.get_sources() 

55 return 

56 

57 except LocationGenerationError: 

58 self.retry() 

59 

60 raise ScenarioError( 

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

62 

63 def init_modelling(self, engine): 

64 self._engine = engine 

65 

66 def get_engine(self): 

67 return self._engine 

68 

69 def get_sources(self): 

70 return self.source_generator.get_sources() 

71 

72 def get_events(self): 

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

74 

75 def collect(collector): 

76 if not callable(collector): 

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

78 

79 @wraps(collector) 

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

81 result = [] 

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

83 for gen in self.target_generators: 

84 result.extend( 

85 func(gen) or []) 

86 return result 

87 

88 return method 

89 

90 @collect 

91 def get_stations(self): 

92 return lambda gen: gen.get_stations() 

93 

94 @collect 

95 def get_targets(self): 

96 return lambda gen: gen.get_targets() 

97 

98 @collect 

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

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

101 return lambda gen: gen.get_waveforms( 

102 self._engine, self.get_sources(), 

103 tmin=tmin, tmax=tmax) 

104 

105 @collect 

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

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

108 return lambda gen: gen.get_onsets( 

109 self._engine, self.get_sources(), 

110 tmin=tmin, tmax=tmax) 

111 

112 @collect 

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

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

115 return lambda gen: gen.get_insar_scenes( 

116 self._engine, self.get_sources(), 

117 tmin=tmin, tmax=tmax) 

118 

119 @collect 

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

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

122 return lambda gen: gen.get_gnss_campaigns( 

123 self._engine, self.get_sources(), 

124 tmin=tmin, tmax=tmax) 

125 

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

127 ''' 

128 Invoke generators and dump the complete scenario. 

129 

130 :param path: Output directory 

131 :type path: str 

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

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

134 previously generated content is found. 

135 :type overwrite: bool 

136 

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

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

139 ''' 

140 

141 self.prepare_data(path, overwrite=False) 

142 self.ensure_data(path) 

143 

144 def prepare_data(self, path, overwrite): 

145 ''' 

146 Prepare directory for scenario content storage. 

147 

148 :param path: Output directory 

149 :type path: str 

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

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

152 previously generated content is found. 

153 :type overwrite: bool 

154 ''' 

155 

156 for dentry in [ 

157 'meta', 

158 'waveforms', 

159 'insar', 

160 'gnss']: 

161 

162 dpath = op.join(path, dentry) 

163 if op.exists(dpath): 

164 if overwrite: 

165 shutil.rmtree(dpath) 

166 else: 

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

168 

169 for fentry in [ 

170 'events.txt', 

171 'sources.yml', 

172 'map.pdf', 

173 'README.md']: 

174 

175 fpath = op.join(path, fentry) 

176 if op.exists(fpath): 

177 if overwrite: 

178 os.unlink(fpath) 

179 else: 

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

181 

182 @collect 

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

184 

185 ''' 

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

187 

188 :param path: Output directory 

189 :type path: str 

190 :param tmin: Start of time interval to generate 

191 :type tmin: timestamp, optional 

192 :param tmax: End of time interval to generate 

193 :type tmax: timestamp, optional 

194 

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

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

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

198 needed. 

199 

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

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

202 reasonable default. 

203 ''' 

204 

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

206 self.source_generator.ensure_data(path) 

207 

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

209 util.ensuredir(meta_dir) 

210 

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

212 if not op.exists(fn_stations): 

213 model.station.dump_stations( 

214 self.get_stations(), fn_stations) 

215 

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

217 if not op.exists(fn_stations_yaml): 

218 model.station.dump_stations_yaml( 

219 self.get_stations(), fn_stations_yaml) 

220 

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

222 if not op.exists(fn_stations_kml): 

223 model.station.dump_kml( 

224 self.get_stations(), fn_stations_kml) 

225 

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

227 if not op.exists(fn_markers): 

228 markers = self.get_onsets() 

229 if markers: 

230 PhaseMarker.save_markers(markers, fn_markers) 

231 

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

233 if not op.exists(fn_readme): 

234 dump_readme(fn_readme) 

235 

236 def ensure_data(gen): 

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

238 return gen.ensure_data( 

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

240 

241 return ensure_data 

242 

243 @collect 

244 def _get_time_ranges(self): 

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

246 

247 def get_time_range(self): 

248 ranges = num.array( 

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

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

251 

252 def _time_range_fill_defaults(self, tmin, tmax): 

253 stmin, stmax = self.get_time_range() 

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

255 

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

257 p = pile.Pile() 

258 

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

260 p.add_file(trf) 

261 return p 

262 

263 def make_map(self, filename): 

264 logger.info('Plotting scenario\'s map...') 

265 if not gmtpy.have_gmt(): 

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

267 return 

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

269 logger.info( 

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

271 'not implemented.') 

272 return 

273 

274 try: 

275 draw_scenario_gmt(self, filename) 

276 except gmtpy.GMTError: 

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

278 except topo.AuthenticationRequired: 

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

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

281 

282 def draw_map(self, fns): 

283 from pyrocko.plot import automap 

284 

285 if isinstance(fns, str): 

286 fns = [fns] 

287 

288 lat, lon = self.get_center_latlon() 

289 radius = self.get_radius() 

290 

291 m = automap.Map( 

292 width=30., 

293 height=30., 

294 lat=lat, 

295 lon=lon, 

296 radius=radius, 

297 show_topo=True, 

298 show_grid=True, 

299 show_rivers=True, 

300 color_wet=(216, 242, 254), 

301 color_dry=(238, 236, 230) 

302 ) 

303 

304 self.source_generator.add_map_artists(m) 

305 

306 sources = self.get_sources() 

307 for gen in self.target_generators: 

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

309 

310 # for patch in self.get_insar_patches(): 

311 # symbol_size = 50. 

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

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

314 # L=True, 

315 # *m.jxyr) 

316 

317 for fn in fns: 

318 m.save(fn) 

319 

320 @property 

321 def stores_wanted(self): 

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

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

324 

325 @property 

326 def stores_missing(self): 

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

328 

329 def ensure_gfstores(self, interactive=False): 

330 if not self.stores_missing: 

331 return 

332 

333 from pyrocko.gf import ws 

334 

335 engine = self.get_engine() 

336 

337 gf_store_superdirs = engine.store_superdirs 

338 

339 if interactive: 

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

341 '%s\n' 

342 'We can try to download the stores from ' 

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

344 'directories:' 

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

346 for idr, dr in enumerate(gf_store_superdirs): 

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

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

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

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

351 print('Canceled.') 

352 sys.exit(1) 

353 elif s == '': 

354 s = 0 

355 try: 

356 s = int(s) 

357 if s > len(gf_store_superdirs): 

358 raise ValueError 

359 except ValueError: 

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

361 sys.exit(1) 

362 else: 

363 s = 1 

364 

365 download_dir = gf_store_superdirs[s-1] 

366 util.ensuredir(download_dir) 

367 logger.info('Downloading Green\'s functions stores to %s' 

368 % download_dir) 

369 

370 oldwd = os.getcwd() 

371 for store in self.stores_missing: 

372 os.chdir(download_dir) 

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

374 

375 os.chdir(oldwd) 

376 

377 @classmethod 

378 def initialize( 

379 cls, path, 

380 center_lat=None, center_lon=None, radius=None, 

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

382 ''' 

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

384 

385 :param path: Path to create the scenerio in 

386 :type path: str 

387 :param center_lat: Center latitude [deg] 

388 :type center_lat: float, optional 

389 :param center_lon: Center longitude [deg] 

390 :type center_lon: float, optional 

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

392 :type radius: float, optional 

393 :param targets: Targets to throw into scenario, 

394 defaults to AVAILABLE_TARGETS 

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

396 objects, optional 

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

398 :type force: bool 

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

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

401 :type stationxml: str 

402 :returns: Scenario 

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

404 ''' 

405 import os.path as op 

406 

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

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

409 

410 util.ensuredir(path) 

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

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

413 

414 scenario = cls( 

415 center_lat=center_lat, 

416 center_lon=center_lon, 

417 radius=radius) 

418 

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

420 

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

422 

423 scenario.update_hierarchy() 

424 

425 scenario.dump(filename=fn) 

426 scenario.prepare_data(path, overwrite=force) 

427 

428 return scenario 

429 

430 

431def draw_scenario_gmt(generator, fn): 

432 return generator.draw_map(fn) 

433 

434 

435def dump_readme(path): 

436 readme = '''# Pyrocko Earthquake Scenario 

437 

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

439 

440## Map of the scenario 

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

442 

443## Earthquake Sources 

444 

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

446 

447## Folder `meta` 

448 

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

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

451 

452## Folder `waveforms` 

453 

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

455 

456## Folder `gnss` 

457 

458The GNSS campaign.yml is living here. 

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

460 

461## Folder `insar` 

462 

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

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

465 

466''' 

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

468 f.write(readme) 

469 

470 return [path]