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 os
8import tarfile
9import errno
10import numpy as num
11import time
14from pyrocko.guts import Object, Timestamp
15from pyrocko import gf, guts, util, pile
16from pyrocko.plot import gmtpy
18from .scenario import draw_scenario_gmt
19from .error import ScenarioError
21op = os.path
22guts_prefix = 'pf.scenario'
25def mtime(p):
26 return os.stat(p).st_mtime
29class ScenarioCollectionItem(Object):
30 scenario_id = gf.StringID.T()
31 time_created = Timestamp.T()
33 def __init__(self, **kwargs):
34 Object.__init__(self, **kwargs)
35 self._path = None
36 self._pile = None
37 self._engine = None
38 self._scenes = None
40 def set_base_path(self, path):
41 self._path = path
43 def get_base_path(self):
44 if self._path is None:
45 raise EnvironmentError('Base path not set!')
46 return self._path
48 def init_modelling(self, engine):
49 self._engine = engine
51 def get_path(self, *entry):
52 return op.join(*((self._path,) + entry))
54 def get_generator(self):
55 generator = guts.load(filename=self.get_path('generator.yaml'))
56 generator.init_modelling(self._engine)
57 return generator
59 def get_time_range(self):
60 return self.get_generator().get_time_range()
62 def have_waveforms(self, tmin, tmax):
63 p = self.get_waveform_pile()
64 trs_have = p.all(
65 tmin=tmin, tmax=tmax, load_data=False, degap=False)
67 return any(tr.data_len() > 0 for tr in trs_have)
69 def get_waveform_pile(self):
70 self.ensure_data()
72 if self._pile is None:
73 path_waveforms = self.get_path('waveforms')
74 util.ensuredir(path_waveforms)
75 fns = util.select_files(
76 [path_waveforms], show_progress=False)
78 self._pile = pile.Pile()
79 if fns:
80 self._pile.load_files(
81 fns, fileformat='mseed', show_progress=False)
83 return self._pile
85 def get_insar_scenes(self):
86 from kite import Scene
87 if self._scenes is None:
88 self._scenes = []
89 path_insar = self.get_path('insar')
90 util.ensuredir(path_insar)
92 fns = util.select_files([path_insar], regex='\\.(npz)$',
93 show_progress=False)
94 for f in fns:
95 self._scenes.append(Scene.load(f))
97 return self._scenes
99 def get_gnss_campaigns(self):
100 return self.get_generator().get_gnss_campaigns()
102 def make_map(self, path_pdf):
103 draw_scenario_gmt(self.get_generator(), path_pdf)
105 def get_map(self, format='pdf'):
106 path_pdf = self.get_path('map.pdf')
108 if not op.exists(path_pdf):
109 self.make_map(path_pdf)
111 path = self.get_path('map.%s' % format)
113 outdated = op.exists(path) and mtime(path) < mtime(path_pdf)
114 if not op.exists(path) or outdated:
115 gmtpy.convert_graph(path_pdf, path)
117 return path
119 def ensure_data(self, tmin=None, tmax=None):
120 return self.get_generator().ensure_data(
121 self.get_path(), tmin, tmax)
123 def get_archive(self):
124 self.ensure_data()
126 path_tar = self.get_path('archive.tar')
127 if not op.exists(path_tar):
128 path_base = self.get_path()
129 path_waveforms = self.get_path('waveforms')
130 self.ensure_data()
132 fns = util.select_files(
133 [path_waveforms], show_progress=False)
135 f = tarfile.TarFile(path_tar, 'w')
136 for fn in fns:
137 fna = fn[len(path_base)+1:]
138 f.add(fn, fna)
140 f.close()
142 return path_tar
145class ScenarioCollection(object):
147 def __init__(self, path, engine):
148 self._scenario_suffix = 'scenario'
149 self._path = path
150 util.ensuredir(self._path)
151 self._engine = engine
152 self._load_scenarios()
154 def _load_scenarios(self):
155 scenarios = []
156 base_path = self.get_path()
157 for path_entry in os.listdir(base_path):
158 scenario_id, suffix = op.splitext(path_entry)
159 if suffix == '.' + self._scenario_suffix:
160 path = op.join(base_path, path_entry, 'scenario.yaml')
161 scenario = guts.load(filename=path)
162 assert scenario.scenario_id == scenario_id
163 scenario.set_base_path(op.join(base_path, path_entry))
164 scenario.init_modelling(self._engine)
165 scenarios.append(scenario)
167 self._scenarios = scenarios
168 self._scenarios.sort(key=lambda s: s.time_created)
170 def get_path(self, scenario_id=None, *entry):
171 if scenario_id is not None:
172 return op.join(self._path, '%s.%s' % (
173 scenario_id, self._scenario_suffix), *entry)
174 else:
175 return self._path
177 def add_scenario(self, scenario_id, scenario_generator):
179 if scenario_generator.seed is None:
180 scenario_generator = guts.clone(scenario_generator)
181 scenario_generator.seed = num.random.randint(1, 2**32-1)
183 path = self.get_path(scenario_id)
184 try:
185 os.mkdir(path)
186 except OSError as e:
187 if e.errno == errno.EEXIST:
188 raise ScenarioError(
189 'Scenario id is already in use: %s' % scenario_id)
190 else:
191 raise
193 scenario = ScenarioCollectionItem(
194 scenario_id=scenario_id,
195 time_created=util.to_time_float(time.time()))
197 scenario_path = self.get_path(scenario_id, 'scenario.yaml')
198 guts.dump(scenario, filename=scenario_path)
200 generator_path = self.get_path(scenario_id, 'generator.yaml')
201 guts.dump(scenario_generator, filename=generator_path)
203 scenario.set_base_path(self.get_path(scenario_id))
204 scenario.init_modelling(self._engine)
206 self._scenarios.append(scenario)
208 def list_scenarios(self, ilo=None, ihi=None):
209 return self._scenarios[ilo:ihi]
211 def get_scenario(self, scenario_id):
212 for scenario in self._scenarios:
213 if scenario.scenario_id == scenario_id:
214 return scenario
216 raise KeyError(scenario_id)