1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import os
7import tarfile
8import errno
9import numpy as num
10import time
13from pyrocko.guts import Object, Timestamp
14from pyrocko import gf, guts, util, pile
15from pyrocko.plot import gmtpy
17from .scenario import draw_scenario_gmt
18from .error import ScenarioError
20op = os.path
21guts_prefix = 'pf.scenario'
24def mtime(p):
25 return os.stat(p).st_mtime
28class ScenarioCollectionItem(Object):
29 scenario_id = gf.StringID.T()
30 time_created = Timestamp.T()
32 def __init__(self, **kwargs):
33 Object.__init__(self, **kwargs)
34 self._path = None
35 self._pile = None
36 self._engine = None
37 self._scenes = None
39 def set_base_path(self, path):
40 self._path = path
42 def get_base_path(self):
43 if self._path is None:
44 raise EnvironmentError('Base path not set!')
45 return self._path
47 def init_modelling(self, engine):
48 self._engine = engine
50 def get_path(self, *entry):
51 return op.join(*((self._path,) + entry))
53 def get_generator(self):
54 generator = guts.load(filename=self.get_path('generator.yaml'))
55 generator.init_modelling(self._engine)
56 return generator
58 def get_time_range(self):
59 return self.get_generator().get_time_range()
61 def have_waveforms(self, tmin, tmax):
62 p = self.get_waveform_pile()
63 trs_have = p.all(
64 tmin=tmin, tmax=tmax, load_data=False, degap=False)
66 return any(tr.data_len() > 0 for tr in trs_have)
68 def get_waveform_pile(self):
69 self.ensure_data()
71 if self._pile is None:
72 path_waveforms = self.get_path('waveforms')
73 util.ensuredir(path_waveforms)
74 fns = util.select_files(
75 [path_waveforms], show_progress=False)
77 self._pile = pile.Pile()
78 if fns:
79 self._pile.load_files(
80 fns, fileformat='mseed', show_progress=False)
82 return self._pile
84 def get_insar_scenes(self):
85 from kite import Scene
86 if self._scenes is None:
87 self._scenes = []
88 path_insar = self.get_path('insar')
89 util.ensuredir(path_insar)
91 fns = util.select_files([path_insar], include='\\.(npz)$',
92 show_progress=False)
93 for f in fns:
94 self._scenes.append(Scene.load(f))
96 return self._scenes
98 def get_gnss_campaigns(self):
99 return self.get_generator().get_gnss_campaigns()
101 def make_map(self, path_pdf):
102 draw_scenario_gmt(self.get_generator(), path_pdf)
104 def get_map(self, format='pdf'):
105 path_pdf = self.get_path('map.pdf')
107 if not op.exists(path_pdf):
108 self.make_map(path_pdf)
110 path = self.get_path('map.%s' % format)
112 outdated = op.exists(path) and mtime(path) < mtime(path_pdf)
113 if not op.exists(path) or outdated:
114 gmtpy.convert_graph(path_pdf, path)
116 return path
118 def ensure_data(self, tmin=None, tmax=None):
119 return self.get_generator().ensure_data(
120 self.get_path(), tmin, tmax)
122 def get_archive(self):
123 self.ensure_data()
125 path_tar = self.get_path('archive.tar')
126 if not op.exists(path_tar):
127 path_base = self.get_path()
128 path_waveforms = self.get_path('waveforms')
129 self.ensure_data()
131 fns = util.select_files(
132 [path_waveforms], show_progress=False)
134 f = tarfile.TarFile(path_tar, 'w')
135 for fn in fns:
136 fna = fn[len(path_base)+1:]
137 f.add(fn, fna)
139 f.close()
141 return path_tar
144class ScenarioCollection(object):
146 def __init__(self, path, engine):
147 self._scenario_suffix = 'scenario'
148 self._path = path
149 util.ensuredir(self._path)
150 self._engine = engine
151 self._load_scenarios()
153 def _load_scenarios(self):
154 scenarios = []
155 base_path = self.get_path()
156 for path_entry in os.listdir(base_path):
157 scenario_id, suffix = op.splitext(path_entry)
158 if suffix == '.' + self._scenario_suffix:
159 path = op.join(base_path, path_entry, 'scenario.yaml')
160 scenario = guts.load(filename=path)
161 assert scenario.scenario_id == scenario_id
162 scenario.set_base_path(op.join(base_path, path_entry))
163 scenario.init_modelling(self._engine)
164 scenarios.append(scenario)
166 self._scenarios = scenarios
167 self._scenarios.sort(key=lambda s: s.time_created)
169 def get_path(self, scenario_id=None, *entry):
170 if scenario_id is not None:
171 return op.join(self._path, '%s.%s' % (
172 scenario_id, self._scenario_suffix), *entry)
173 else:
174 return self._path
176 def add_scenario(self, scenario_id, scenario_generator):
178 if scenario_generator.seed is None:
179 scenario_generator = guts.clone(scenario_generator)
180 scenario_generator.seed = num.random.randint(1, 2**32-1)
182 path = self.get_path(scenario_id)
183 try:
184 os.mkdir(path)
185 except OSError as e:
186 if e.errno == errno.EEXIST:
187 raise ScenarioError(
188 'Scenario id is already in use: %s' % scenario_id)
189 else:
190 raise
192 scenario = ScenarioCollectionItem(
193 scenario_id=scenario_id,
194 time_created=util.to_time_float(time.time()))
196 scenario_path = self.get_path(scenario_id, 'scenario.yaml')
197 guts.dump(scenario, filename=scenario_path)
199 generator_path = self.get_path(scenario_id, 'generator.yaml')
200 guts.dump(scenario_generator, filename=generator_path)
202 scenario.set_base_path(self.get_path(scenario_id))
203 scenario.init_modelling(self._engine)
205 self._scenarios.append(scenario)
207 def list_scenarios(self, ilo=None, ihi=None):
208 return self._scenarios[ilo:ihi]
210 def get_scenario(self, scenario_id):
211 for scenario in self._scenarios:
212 if scenario.scenario_id == scenario_id:
213 return scenario
215 raise KeyError(scenario_id)