1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6from __future__ import absolute_import, print_function 

7 

8import os.path as op 

9import logging 

10import time 

11try: 

12 import cPickle as pickle 

13except ImportError: 

14 import pickle 

15 

16from pyrocko import util 

17from pyrocko.guts import String, Dict, Duration, dump_all 

18 

19from .base import Source 

20from ..model import ehash 

21from ..lock import LockDir 

22 

23guts_prefix = 'squirrel' 

24 

25logger = logging.getLogger('psq.client.catalog') 

26 

27 

28class Link(object): 

29 def __init__(self, tmin, tmax, tmodified, nevents=-1, content_id=None): 

30 self.tmin = tmin 

31 self.tmax = tmax 

32 self.tmodified = tmodified 

33 self.nevents = nevents 

34 self.content_id = content_id 

35 

36 def __str__(self): 

37 return 'span %s - %s, access %s, nevents %i' % ( 

38 util.tts(self.tmin), 

39 util.tts(self.tmax), 

40 util.tts(self.tmodified), 

41 self.nevents) 

42 

43 

44class NoSuchCatalog(Exception): 

45 pass 

46 

47 

48def get_catalog(name): 

49 if name == 'geofon': 

50 from pyrocko.client.geofon import Geofon 

51 return Geofon() 

52 elif name == 'gcmt': 

53 from pyrocko.client.globalcmt import GlobalCMT 

54 return GlobalCMT() 

55 else: 

56 raise NoSuchCatalog(name) 

57 

58 

59class CatalogSource(Source): 

60 ''' 

61 Squirrel data-source to transparently access online earthquake catalogs. 

62 

63 The catalog source maintains and synchronizes a partial copy of the online 

64 catalog, e.g. of all events above a certain magnitude. The time span for 

65 which the local copy of the catalog should be up-to date is maintained 

66 automatically be Squirrel. Data is loaded and updated in chunks as 

67 needed in a just-in-time fashion. Data validity can optionally expire after 

68 a given period of time and new data can be treated to be preliminary. 

69 In both cases information will be refreshed as needed. 

70 ''' 

71 

72 catalog = String.T( 

73 help='Catalog name.') 

74 

75 query_args = Dict.T( 

76 String.T(), String.T(), 

77 optional=True, 

78 help='Common arguments, which are appended to all queries, e.g. to ' 

79 'constrain location, depth or magnitude ranges.') 

80 

81 expires = Duration.T( 

82 optional=True, 

83 help='Expiration time [s]. Information older than this will be ' 

84 'refreshed, i.e. queried again.') 

85 

86 anxious = Duration.T( 

87 optional=True, 

88 help='Anxiety period [s]. Information will be treated as preliminary ' 

89 'if it was younger than this at the time of its retrieval. ' 

90 'Preliminary information is refreshed on each query relevant ' 

91 'to it.') 

92 

93 cache_path = String.T( 

94 optional=True, 

95 help='Directory path where the partial local copy of the catalog is ' 

96 'kept. By default the Squirrel environment\'s cache directory is ' 

97 'used.') 

98 

99 def __init__(self, catalog, query_args=None, **kwargs): 

100 Source.__init__(self, catalog=catalog, query_args=query_args, **kwargs) 

101 

102 self._hash = self.make_hash() 

103 self._nevents_query_hint = 1000 

104 self._nevents_chunk_hint = 5000 

105 self._tquery = 3600.*24. 

106 self._tquery_limits = (3600., 3600.*24.*365.) 

107 

108 def describe(self): 

109 return 'catalog:%s:%s' % (self.catalog, self.get_hash()) 

110 

111 def setup(self, squirrel, check=True): 

112 self._force_query_age_max = self.anxious 

113 self._catalog = get_catalog(self.catalog) 

114 

115 self._cache_path = op.join( 

116 self.cache_path or squirrel._cache_path, 

117 'catalog', 

118 self.get_hash()) 

119 

120 util.ensuredir(self._cache_path) 

121 

122 def make_hash(self): 

123 s = self.catalog 

124 if self.query_args is not None: 

125 s += ','.join( 

126 '%s:%s' % (k, self.query_args[k]) 

127 for k in sorted(self.query_args.keys())) 

128 else: 

129 s += 'noqueryargs' 

130 

131 return ehash(s) 

132 

133 def get_hash(self): 

134 return self._hash 

135 

136 def update_event_inventory(self, squirrel, constraint=None): 

137 

138 with LockDir(self._cache_path): 

139 self._load_chain() 

140 

141 assert constraint is not None 

142 if constraint is not None: 

143 tmin, tmax = constraint.tmin, constraint.tmax 

144 

145 tmin_sq, tmax_sq = squirrel.get_time_span() 

146 

147 if tmin is None: 

148 tmin = tmin_sq 

149 

150 if tmax is None: 

151 tmax = tmax_sq 

152 

153 if tmin is None or tmax is None: 

154 logger.warning( 

155 'Cannot query catalog source "%s" without time ' 

156 'constraint. Could not determine appropriate time ' 

157 'constraint from current data holdings (no data?).' 

158 % self.catalog) 

159 

160 return 

161 

162 if tmin >= tmax: 

163 return 

164 

165 tnow = time.time() 

166 modified = False 

167 

168 if not self._chain: 

169 self._chain = [Link(tmin, tmax, tnow)] 

170 modified = True 

171 else: 

172 if tmin < self._chain[0].tmin: 

173 self._chain[0:0] = [Link(tmin, self._chain[0].tmin, tnow)] 

174 modified = True 

175 if self._chain[-1].tmax < tmax: 

176 self._chain.append(Link(self._chain[-1].tmax, tmax, tnow)) 

177 modified = True 

178 

179 chain = [] 

180 remove = [] 

181 for link in self._chain: 

182 if tmin < link.tmax and link.tmin < tmax \ 

183 and self._outdated(link, tnow): 

184 

185 if link.content_id: 

186 remove.append( 

187 self._get_events_file_path(link.content_id)) 

188 

189 tmin_query = max(link.tmin, tmin) 

190 tmax_query = min(link.tmax, tmax) 

191 

192 if link.tmin < tmin_query: 

193 chain.append(Link(link.tmin, tmin_query, tnow)) 

194 

195 if tmin_query < tmax_query: 

196 for link in self._iquery(tmin_query, tmax_query, tnow): 

197 chain.append(link) 

198 

199 if tmax_query < link.tmax: 

200 chain.append(Link(tmax_query, link.tmax, tnow)) 

201 

202 modified = True 

203 

204 else: 

205 chain.append(link) 

206 

207 if modified: 

208 self._chain = chain 

209 self._dump_chain() 

210 squirrel.remove(remove) 

211 

212 add = [] 

213 for link in self._chain: 

214 if link.content_id: 

215 add.append(self._get_events_file_path(link.content_id)) 

216 

217 squirrel.add(add, kinds=['event'], format='yaml') 

218 

219 def _iquery(self, tmin, tmax, tmodified): 

220 

221 nwant = self._nevents_query_hint 

222 tlim = self._tquery_limits 

223 

224 t = tmin 

225 tpack_min = tmin 

226 

227 events = [] 

228 while t < tmax: 

229 tmin_query = t 

230 tmax_query = min(t + self._tquery, tmax) 

231 

232 events_new = self._query(tmin_query, tmax_query) 

233 nevents_new = len(events_new) 

234 events.extend(events_new) 

235 while len(events) > int(self._nevents_chunk_hint * 1.5): 

236 tpack_max = events[self._nevents_chunk_hint].time 

237 yield self._pack( 

238 events[:self._nevents_chunk_hint], 

239 tpack_min, tpack_max, tmodified) 

240 

241 tpack_min = tpack_max 

242 events[:self._nevents_query_hint] = [] 

243 

244 t += self._tquery 

245 

246 if tmax_query != tmax: 

247 if nevents_new < 5: 

248 self._tquery *= 10.0 

249 

250 elif not (nwant // 2 < nevents_new < nwant * 2): 

251 self._tquery /= float(nevents_new) / float(nwant) 

252 

253 self._tquery = max(tlim[0], min(self._tquery, tlim[1])) 

254 

255 if self._force_query_age_max is not None: 

256 tsplit = tmodified - self._force_query_age_max 

257 if tpack_min < tsplit < tmax: 

258 events_older = [] 

259 events_newer = [] 

260 for ev in events: 

261 if ev.time < tsplit: 

262 events_older.append(ev) 

263 else: 

264 events_newer.append(ev) 

265 

266 yield self._pack(events_older, tpack_min, tsplit, tmodified) 

267 yield self._pack(events_newer, tsplit, tmax, tmodified) 

268 return 

269 

270 yield self._pack(events, tpack_min, tmax, tmodified) 

271 

272 def _pack(self, events, tmin, tmax, tmodified): 

273 if events: 

274 content_id = ehash( 

275 self.get_hash() + ' %r %r %r' % (tmin, tmax, tmodified)) 

276 path = self._get_events_file_path(content_id) 

277 dump_all(events, filename=path) 

278 else: 

279 content_id = None 

280 

281 return Link(tmin, tmax, tmodified, len(events), content_id) 

282 

283 def _query(self, tmin, tmax): 

284 logger.info('Querying catalog "%s" for time span %s - %s.' % ( 

285 self.catalog, util.tts(tmin), util.tts(tmax))) 

286 

287 return self._catalog.get_events( 

288 (tmin, tmax), 

289 **(self.query_args or {})) 

290 

291 def _outdated(self, link, tnow): 

292 if link.nevents == -1: 

293 return True 

294 

295 if self._force_query_age_max \ 

296 and link.tmax + self._force_query_age_max > link.tmodified: 

297 

298 return True 

299 

300 if self.expires is not None \ 

301 and link.tmodified < tnow - self.expires: 

302 

303 return True 

304 

305 return False 

306 

307 def _get_events_file_path(self, fhash): 

308 return op.join(self._cache_path, fhash + '.pf') 

309 

310 def _get_chain_file_path(self): 

311 return op.join(self._cache_path, 'chain.pickle') 

312 

313 def _load_chain(self): 

314 path = self._get_chain_file_path() 

315 if op.exists(path): 

316 with open(path, 'rb') as f: 

317 self._chain = pickle.load(f) 

318 else: 

319 self._chain = [] 

320 

321 def _dump_chain(self): 

322 with open(self._get_chain_file_path(), 'wb') as f: 

323 pickle.dump(self._chain, f, protocol=2) 

324 

325 

326__all__ = [ 

327 'CatalogSource' 

328]