1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5from __future__ import absolute_import, division, print_function 

6 

7import os 

8import tarfile 

9import errno 

10import numpy as num 

11import time 

12 

13 

14from pyrocko.guts import Object, Timestamp 

15from pyrocko import gf, guts, util, pile 

16from pyrocko.plot import gmtpy 

17 

18from .scenario import draw_scenario_gmt 

19from .error import ScenarioError 

20 

21op = os.path 

22guts_prefix = 'pf.scenario' 

23 

24 

25def mtime(p): 

26 return os.stat(p).st_mtime 

27 

28 

29class ScenarioCollectionItem(Object): 

30 scenario_id = gf.StringID.T() 

31 time_created = Timestamp.T() 

32 

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 

39 

40 def set_base_path(self, path): 

41 self._path = path 

42 

43 def get_base_path(self): 

44 if self._path is None: 

45 raise EnvironmentError('Base path not set!') 

46 return self._path 

47 

48 def init_modelling(self, engine): 

49 self._engine = engine 

50 

51 def get_path(self, *entry): 

52 return op.join(*((self._path,) + entry)) 

53 

54 def get_generator(self): 

55 generator = guts.load(filename=self.get_path('generator.yaml')) 

56 generator.init_modelling(self._engine) 

57 return generator 

58 

59 def get_time_range(self): 

60 return self.get_generator().get_time_range() 

61 

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) 

66 

67 return any(tr.data_len() > 0 for tr in trs_have) 

68 

69 def get_waveform_pile(self): 

70 self.ensure_data() 

71 

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) 

77 

78 self._pile = pile.Pile() 

79 if fns: 

80 self._pile.load_files( 

81 fns, fileformat='mseed', show_progress=False) 

82 

83 return self._pile 

84 

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) 

91 

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)) 

96 

97 return self._scenes 

98 

99 def get_gnss_campaigns(self): 

100 return self.get_generator().get_gnss_campaigns() 

101 

102 def make_map(self, path_pdf): 

103 draw_scenario_gmt(self.get_generator(), path_pdf) 

104 

105 def get_map(self, format='pdf'): 

106 path_pdf = self.get_path('map.pdf') 

107 

108 if not op.exists(path_pdf): 

109 self.make_map(path_pdf) 

110 

111 path = self.get_path('map.%s' % format) 

112 

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) 

116 

117 return path 

118 

119 def ensure_data(self, tmin=None, tmax=None): 

120 return self.get_generator().ensure_data( 

121 self.get_path(), tmin, tmax) 

122 

123 def get_archive(self): 

124 self.ensure_data() 

125 

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() 

131 

132 fns = util.select_files( 

133 [path_waveforms], show_progress=False) 

134 

135 f = tarfile.TarFile(path_tar, 'w') 

136 for fn in fns: 

137 fna = fn[len(path_base)+1:] 

138 f.add(fn, fna) 

139 

140 f.close() 

141 

142 return path_tar 

143 

144 

145class ScenarioCollection(object): 

146 

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() 

153 

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) 

166 

167 self._scenarios = scenarios 

168 self._scenarios.sort(key=lambda s: s.time_created) 

169 

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 

176 

177 def add_scenario(self, scenario_id, scenario_generator): 

178 

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) 

182 

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 

192 

193 scenario = ScenarioCollectionItem( 

194 scenario_id=scenario_id, 

195 time_created=util.to_time_float(time.time())) 

196 

197 scenario_path = self.get_path(scenario_id, 'scenario.yaml') 

198 guts.dump(scenario, filename=scenario_path) 

199 

200 generator_path = self.get_path(scenario_id, 'generator.yaml') 

201 guts.dump(scenario_generator, filename=generator_path) 

202 

203 scenario.set_base_path(self.get_path(scenario_id)) 

204 scenario.init_modelling(self._engine) 

205 

206 self._scenarios.append(scenario) 

207 

208 def list_scenarios(self, ilo=None, ihi=None): 

209 return self._scenarios[ilo:ihi] 

210 

211 def get_scenario(self, scenario_id): 

212 for scenario in self._scenarios: 

213 if scenario.scenario_id == scenario_id: 

214 return scenario 

215 

216 raise KeyError(scenario_id)