1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
5from __future__ import absolute_import, division, print_function
7import logging
8import sys
9import os
10import os.path as op
11import shutil
12import random
13from functools import wraps
15import numpy as num
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
23from .error import ScenarioError, CannotCreatePath, LocationGenerationError
24from .base import LocationGenerator
25from .sources import SourceGenerator, DCSourceGenerator
26from .targets import TargetGenerator, AVAILABLE_TARGETS
28logger = logging.getLogger('pyrocko.scenario')
29guts_prefix = 'pf.scenario'
31km = 1000.
34class ScenarioGenerator(LocationGenerator):
35 '''
36 Generator for synthetic seismo-geodetic scenarios.
37 '''
39 target_generators = List.T(
40 TargetGenerator.T(),
41 default=[],
42 help='Targets to spawn in the scenario.')
44 source_generator = SourceGenerator.T(
45 default=DCSourceGenerator.D(),
46 help='Sources to spawn in the scenario.')
48 def __init__(self, **kwargs):
49 LocationGenerator.__init__(self, **kwargs)
50 for itry in range(self.ntries):
52 try:
53 self.get_stations()
54 self.get_sources()
55 return
57 except LocationGenerationError:
58 self.retry()
60 raise ScenarioError(
61 'Could not generate scenario within %i tries.' % self.ntries)
63 def init_modelling(self, engine):
64 self._engine = engine
66 def get_engine(self):
67 return self._engine
69 def get_sources(self):
70 return self.source_generator.get_sources()
72 def get_events(self):
73 return [s.pyrocko_event() for s in self.get_sources()]
75 def collect(collector):
76 if not callable(collector):
77 raise AttributeError('This method should not be called directly.')
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
88 return method
90 @collect
91 def get_stations(self):
92 return lambda gen: gen.get_stations()
94 @collect
95 def get_targets(self):
96 return lambda gen: gen.get_targets()
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)
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)
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)
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)
126 def dump_data(self, path, overwrite=False):
127 '''
128 Invoke generators and dump the complete scenario.
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
137 Wrapper to call :py:meth:`prepare_data` followed by
138 :py:meth:`ensure_data` with default arguments.
139 '''
141 self.prepare_data(path, overwrite=False)
142 self.ensure_data(path)
144 def prepare_data(self, path, overwrite):
145 '''
146 Prepare directory for scenario content storage.
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 '''
156 for dentry in [
157 'meta',
158 'waveforms',
159 'insar',
160 'gnss']:
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)
169 for fentry in [
170 'events.txt',
171 'sources.yml',
172 'map.pdf',
173 'README.md']:
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)
182 @collect
183 def ensure_data(self, path, tmin=None, tmax=None):
185 '''
186 Generate and output scenario content to files, as needed.
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
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.
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 '''
205 tmin, tmax = self._time_range_fill_defaults(tmin, tmax)
206 self.source_generator.ensure_data(path)
208 meta_dir = op.join(path, 'meta')
209 util.ensuredir(meta_dir)
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)
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)
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)
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)
232 fn_readme = op.join(path, 'README.md')
233 if not op.exists(fn_readme):
234 dump_readme(fn_readme)
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)
241 return ensure_data
243 @collect
244 def _get_time_ranges(self):
245 return lambda gen: [gen.get_time_range(self.get_sources())]
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()
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
256 def get_pile(self, tmin=None, tmax=None):
257 p = pile.Pile()
259 trf = pile.MemTracesFile(None, self.get_waveforms(tmin, tmax))
260 p.add_file(trf)
261 return p
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
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.')
282 def draw_map(self, fns):
283 from pyrocko.plot import automap
285 if isinstance(fns, str):
286 fns = [fns]
288 lat, lon = self.get_center_latlon()
289 radius = self.get_radius()
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 )
304 self.source_generator.add_map_artists(m)
306 sources = self.get_sources()
307 for gen in self.target_generators:
308 gen.add_map_artists(self.get_engine(), sources, m)
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)
317 for fn in fns:
318 m.save(fn)
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')])
325 @property
326 def stores_missing(self):
327 return self.stores_wanted - set(self.get_engine().get_store_ids())
329 def ensure_gfstores(self, interactive=False):
330 if not self.stores_missing:
331 return
333 from pyrocko.gf import ws
335 engine = self.get_engine()
337 gf_store_superdirs = engine.store_superdirs
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
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)
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)
375 os.chdir(oldwd)
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``
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
407 if op.exists(path) and not force:
408 raise CannotCreatePath('Directory %s alread exists.' % path)
410 util.ensuredir(path)
411 fn = op.join(path, 'scenario.yml')
412 logger.debug('Writing new scenario to %s' % fn)
414 scenario = cls(
415 center_lat=center_lat,
416 center_lon=center_lon,
417 radius=radius)
419 scenario.seed = random.randint(1, 2**32-1)
421 scenario.target_generators.extend([t() for t in targets])
423 scenario.update_hierarchy()
425 scenario.dump(filename=fn)
426 scenario.prepare_data(path, overwrite=force)
428 return scenario
431def draw_scenario_gmt(generator, fn):
432 return generator.draw_map(fn)
435def dump_readme(path):
436 readme = '''# Pyrocko Earthquake Scenario
438The directory structure of a scenario is layed out as follows:
440## Map of the scenario
441A simple map is generated from `pyrocko.automap` in map.pdf
443## Earthquake Sources
445Can be found as events.txt and sources.yml hosts the pyrocko.gf sources.
447## Folder `meta`
449Contains stations.txt and StationXML files for waveforms as well as KML data.
450The responses are flat with gain of 1.0 at 1.0 Hz.
452## Folder `waveforms`
454Waveforms as mini-seed are stored here, segregated into days.
456## Folder `gnss`
458The GNSS campaign.yml is living here.
459Use `pyrocko.guts.load(filename='campaign.yml)` to load the campaign.
461## Folder `insar`
463Kite InSAR scenes for ascending and descending tracks are stored there.
464Use `kite.Scene.load(<filename>)` to inspect the scenes.
466'''
467 with open(path, 'w') as f:
468 f.write(readme)
470 return [path]