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----------
6'''
7Implementation of the scenario generator's main class.
8'''
10import logging
11import sys
12import os
13import os.path as op
14import shutil
15import random
16from functools import wraps
18import numpy as num
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
26from .error import ScenarioError, CannotCreatePath, LocationGenerationError
27from .base import LocationGenerator
28from .sources import SourceGenerator, DCSourceGenerator
29from .targets import TargetGenerator, AVAILABLE_TARGETS
31logger = logging.getLogger('pyrocko.scenario')
32guts_prefix = 'pf.scenario'
34km = 1000.
37class ScenarioGenerator(LocationGenerator):
38 '''
39 Generator for synthetic seismo-geodetic scenarios.
40 '''
42 target_generators = List.T(
43 TargetGenerator.T(),
44 default=[],
45 help='Targets to spawn in the scenario.')
47 source_generator = SourceGenerator.T(
48 default=DCSourceGenerator.D(),
49 help='Sources to spawn in the scenario.')
51 def __init__(self, **kwargs):
52 LocationGenerator.__init__(self, **kwargs)
53 for itry in range(self.ntries):
55 try:
56 self.get_stations()
57 self.get_sources()
58 return
60 except LocationGenerationError:
61 self.retry()
63 raise ScenarioError(
64 'Could not generate scenario in %i tries.' % self.ntries)
66 def init_modelling(self, engine):
67 self._engine = engine
69 def get_engine(self):
70 return self._engine
72 def get_sources(self):
73 return self.source_generator.get_sources()
75 def get_events(self):
76 return [s.pyrocko_event() for s in self.get_sources()]
78 def collect(collector):
79 if not callable(collector):
80 raise AttributeError('This method should not be called directly.')
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
91 return method
93 @collect
94 def get_stations(self):
95 return lambda gen: gen.get_stations()
97 @collect
98 def get_targets(self):
99 return lambda gen: gen.get_targets()
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)
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)
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)
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)
129 def dump_data(self, path, overwrite=False):
130 '''
131 Invoke generators and dump the complete scenario.
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
140 Wrapper to call :py:meth:`prepare_data` followed by
141 :py:meth:`ensure_data` with default arguments.
142 '''
144 self.prepare_data(path, overwrite=False)
145 self.ensure_data(path)
147 def prepare_data(self, path, overwrite):
148 '''
149 Prepare directory for scenario content storage.
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 '''
159 for dentry in [
160 'meta',
161 'waveforms',
162 'insar',
163 'gnss']:
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)
172 for fentry in [
173 'events.txt',
174 'sources.yml',
175 'map.pdf',
176 'README.md']:
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)
185 @collect
186 def ensure_data(self, path, tmin=None, tmax=None):
188 '''
189 Generate and output scenario content to files, as needed.
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`
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.
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 '''
208 tmin, tmax = self._time_range_fill_defaults(tmin, tmax)
209 self.source_generator.ensure_data(path)
211 meta_dir = op.join(path, 'meta')
212 util.ensuredir(meta_dir)
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)
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)
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)
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)
235 fn_readme = op.join(path, 'README.md')
236 if not op.exists(fn_readme):
237 dump_readme(fn_readme)
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)
244 return ensure_data
246 @collect
247 def _get_time_ranges(self):
248 return lambda gen: [gen.get_time_range(self.get_sources())]
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()
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
259 def get_pile(self, tmin=None, tmax=None):
260 p = pile.Pile()
262 trf = pile.MemTracesFile(None, self.get_waveforms(tmin, tmax))
263 p.add_file(trf)
264 return p
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
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.')
285 def draw_map(self, fns):
286 from pyrocko.plot import automap
288 if isinstance(fns, str):
289 fns = [fns]
291 lat, lon = self.get_center_latlon()
292 radius = self.get_radius()
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 )
307 self.source_generator.add_map_artists(m)
309 sources = self.get_sources()
310 for gen in self.target_generators:
311 gen.add_map_artists(self.get_engine(), sources, m)
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)
320 for fn in fns:
321 m.save(fn)
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')])
328 @property
329 def stores_missing(self):
330 return self.stores_wanted - set(self.get_engine().get_store_ids())
332 def ensure_gfstores(self, interactive=False):
333 if not self.stores_missing:
334 return
336 from pyrocko.gf import ws
338 engine = self.get_engine()
340 gf_store_superdirs = engine.store_superdirs
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
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)
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)
378 os.chdir(oldwd)
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``
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
410 if op.exists(path) and not force:
411 raise CannotCreatePath('Directory %s alread exists.' % path)
413 util.ensuredir(path)
414 fn = op.join(path, 'scenario.yml')
415 logger.debug('Writing new scenario to %s' % fn)
417 scenario = cls(
418 center_lat=center_lat,
419 center_lon=center_lon,
420 radius=radius)
422 scenario.seed = random.randint(1, 2**32-1)
424 scenario.target_generators.extend([t() for t in targets])
426 scenario.update_hierarchy()
428 scenario.dump(filename=fn)
429 scenario.prepare_data(path, overwrite=force)
431 return scenario
434def draw_scenario_gmt(generator, fn):
435 return generator.draw_map(fn)
438def dump_readme(path):
439 readme = '''# Pyrocko Earthquake Scenario
441The directory structure of a scenario is layed out as follows:
443## Map of the scenario
444A simple map is generated from `pyrocko.automap` in map.pdf
446## Earthquake Sources
448Can be found as events.txt and sources.yml hosts the pyrocko.gf sources.
450## Folder `meta`
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.
455## Folder `waveforms`
457Waveforms as mini-seed are stored here, segregated into days.
459## Folder `gnss`
461The GNSS campaign.yml is living here.
462Use `pyrocko.guts.load(filename='campaign.yml)` to load the campaign.
464## Folder `insar`
466Kite InSAR scenes for ascending and descending tracks are stored there.
467Use `kite.Scene.load(<filename>)` to inspect the scenes.
469'''
470 with open(path, 'w') as f:
471 f.write(readme)
473 return [path]