1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
7import sys
8import os
9import os.path as op
10import shutil
11import random
12from functools import wraps
14import numpy as num
16from pyrocko.guts import List
17from pyrocko.plot import gmtpy
18from pyrocko.gui.marker import PhaseMarker
19from pyrocko import pile, util, model
20from pyrocko.dataset import topo
22from .error import ScenarioError, CannotCreatePath, LocationGenerationError
23from .base import LocationGenerator
24from .sources import SourceGenerator, DCSourceGenerator
25from .targets import TargetGenerator, AVAILABLE_TARGETS
27logger = logging.getLogger('pyrocko.scenario')
28guts_prefix = 'pf.scenario'
30km = 1000.
33class ScenarioGenerator(LocationGenerator):
34 '''
35 Generator for synthetic seismo-geodetic scenarios.
36 '''
38 target_generators = List.T(
39 TargetGenerator.T(),
40 default=[],
41 help='Targets to spawn in the scenario.')
43 source_generator = SourceGenerator.T(
44 default=DCSourceGenerator.D(),
45 help='Sources to spawn in the scenario.')
47 def __init__(self, **kwargs):
48 LocationGenerator.__init__(self, **kwargs)
49 for itry in range(self.ntries):
51 try:
52 self.get_stations()
53 self.get_sources()
54 return
56 except LocationGenerationError:
57 self.retry()
59 raise ScenarioError(
60 'Could not generate scenario in %i tries.' % self.ntries)
62 def init_modelling(self, engine):
63 self._engine = engine
65 def get_engine(self):
66 return self._engine
68 def get_sources(self):
69 return self.source_generator.get_sources()
71 def get_events(self):
72 return [s.pyrocko_event() for s in self.get_sources()]
74 def collect(collector):
75 if not callable(collector):
76 raise AttributeError('This method should not be called directly.')
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
87 return method
89 @collect
90 def get_stations(self):
91 return lambda gen: gen.get_stations()
93 @collect
94 def get_targets(self):
95 return lambda gen: gen.get_targets()
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)
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)
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)
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)
125 def dump_data(self, path, overwrite=False):
126 '''
127 Invoke generators and dump the complete scenario.
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
136 Wrapper to call :py:meth:`prepare_data` followed by
137 :py:meth:`ensure_data` with default arguments.
138 '''
140 self.prepare_data(path, overwrite=False)
141 self.ensure_data(path)
143 def prepare_data(self, path, overwrite):
144 '''
145 Prepare directory for scenario content storage.
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 '''
155 for dentry in [
156 'meta',
157 'waveforms',
158 'insar',
159 'gnss']:
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)
168 for fentry in [
169 'events.txt',
170 'sources.yml',
171 'map.pdf',
172 'README.md']:
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)
181 @collect
182 def ensure_data(self, path, tmin=None, tmax=None):
184 '''
185 Generate and output scenario content to files, as needed.
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
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.
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 '''
204 tmin, tmax = self._time_range_fill_defaults(tmin, tmax)
205 self.source_generator.ensure_data(path)
207 meta_dir = op.join(path, 'meta')
208 util.ensuredir(meta_dir)
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)
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)
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)
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)
231 fn_readme = op.join(path, 'README.md')
232 if not op.exists(fn_readme):
233 dump_readme(fn_readme)
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)
240 return ensure_data
242 @collect
243 def _get_time_ranges(self):
244 return lambda gen: [gen.get_time_range(self.get_sources())]
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()
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
255 def get_pile(self, tmin=None, tmax=None):
256 p = pile.Pile()
258 trf = pile.MemTracesFile(None, self.get_waveforms(tmin, tmax))
259 p.add_file(trf)
260 return p
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
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.')
281 def draw_map(self, fns):
282 from pyrocko.plot import automap
284 if isinstance(fns, str):
285 fns = [fns]
287 lat, lon = self.get_center_latlon()
288 radius = self.get_radius()
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 )
303 self.source_generator.add_map_artists(m)
305 sources = self.get_sources()
306 for gen in self.target_generators:
307 gen.add_map_artists(self.get_engine(), sources, m)
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)
316 for fn in fns:
317 m.save(fn)
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')])
324 @property
325 def stores_missing(self):
326 return self.stores_wanted - set(self.get_engine().get_store_ids())
328 def ensure_gfstores(self, interactive=False):
329 if not self.stores_missing:
330 return
332 from pyrocko.gf import ws
334 engine = self.get_engine()
336 gf_store_superdirs = engine.store_superdirs
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
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)
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)
374 os.chdir(oldwd)
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``
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
406 if op.exists(path) and not force:
407 raise CannotCreatePath('Directory %s alread exists.' % path)
409 util.ensuredir(path)
410 fn = op.join(path, 'scenario.yml')
411 logger.debug('Writing new scenario to %s' % fn)
413 scenario = cls(
414 center_lat=center_lat,
415 center_lon=center_lon,
416 radius=radius)
418 scenario.seed = random.randint(1, 2**32-1)
420 scenario.target_generators.extend([t() for t in targets])
422 scenario.update_hierarchy()
424 scenario.dump(filename=fn)
425 scenario.prepare_data(path, overwrite=force)
427 return scenario
430def draw_scenario_gmt(generator, fn):
431 return generator.draw_map(fn)
434def dump_readme(path):
435 readme = '''# Pyrocko Earthquake Scenario
437The directory structure of a scenario is layed out as follows:
439## Map of the scenario
440A simple map is generated from `pyrocko.automap` in map.pdf
442## Earthquake Sources
444Can be found as events.txt and sources.yml hosts the pyrocko.gf sources.
446## Folder `meta`
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.
451## Folder `waveforms`
453Waveforms as mini-seed are stored here, segregated into days.
455## Folder `gnss`
457The GNSS campaign.yml is living here.
458Use `pyrocko.guts.load(filename='campaign.yml)` to load the campaign.
460## Folder `insar`
462Kite InSAR scenes for ascending and descending tracks are stored there.
463Use `kite.Scene.load(<filename>)` to inspect the scenes.
465'''
466 with open(path, 'w') as f:
467 f.write(readme)
469 return [path]