1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import logging 

7from pyrocko import squirrel as psq, trace, util 

8from pyrocko import pile as classic_pile 

9 

10logger = logging.getLogger('psq.pile') 

11 

12 

13def trace_callback_to_nut_callback(trace_callback): 

14 if trace_callback is None: 

15 return None 

16 

17 def nut_callback(nut): 

18 return trace_callback(nut.dummy_trace) 

19 

20 return nut_callback 

21 

22 

23class CodesDummyTrace(object): 

24 def __init__(self, codes): 

25 self.network, self.station, self.location, self.channel \ 

26 = self.nslc_id \ 

27 = codes[0:4] 

28 

29 

30def trace_callback_to_codes_callback(trace_callback): 

31 if trace_callback is None: 

32 return None 

33 

34 def codes_callback(codes): 

35 return trace_callback(CodesDummyTrace(codes)) 

36 

37 return codes_callback 

38 

39 

40class Pile(object): 

41 ''' 

42 :py:class:`pyrocko.pile.Pile` surrogate: waveform lookup, loading and 

43 caching. 

44 

45 This class emulates most of the older :py:class:`pyrocko.pile.Pile` methods 

46 by using calls to a :py:class:`pyrocko.squirrel.base.Squirrel` instance 

47 behind the scenes. 

48 

49 This interface can be used as a drop-in replacement for piles which are 

50 used in existing scripts and programs for efficient waveform data access. 

51 The Squirrel-based pile scales better for large datasets. Newer scripts 

52 should use Squirrel's native methods to avoid the emulation overhead. 

53 

54 .. note:: 

55 Many methods in the original pile implementation lack documentation, as 

56 do here. Read the source, Luke! 

57 ''' 

58 def __init__(self, squirrel=None): 

59 if squirrel is None: 

60 squirrel = psq.Squirrel() 

61 

62 self._squirrel = squirrel 

63 self._listeners = [] 

64 self._squirrel.get_database().add_listener( 

65 self._notify_squirrel_to_pile) 

66 

67 def _notify_squirrel_to_pile(self, event, *args): 

68 self.notify_listeners(event) 

69 

70 def add_listener(self, obj): 

71 self._listeners.append(util.smart_weakref(obj)) 

72 

73 def notify_listeners(self, what): 

74 for ref in self._listeners: 

75 obj = ref() 

76 if obj: 

77 obj.pile_changed(what, []) 

78 

79 def get_tmin(self): 

80 return self.tmin 

81 

82 def get_tmax(self): 

83 return self.tmax 

84 

85 def get_deltatmin(self): 

86 return self._squirrel.get_deltat_span('waveform')[0] 

87 

88 def get_deltatmax(self): 

89 return self._squirrel.get_deltat_span('waveform')[1] 

90 

91 @property 

92 def deltatmin(self): 

93 return self.get_deltatmin() 

94 

95 @property 

96 def deltatmax(self): 

97 return self.get_deltatmax() 

98 

99 @property 

100 def tmin(self): 

101 return self._squirrel.get_time_span('waveform')[0] 

102 

103 @property 

104 def tmax(self): 

105 return self._squirrel.get_time_span('waveform')[1] 

106 

107 @property 

108 def networks(self): 

109 return set( 

110 codes.network for codes in self._squirrel.get_codes('waveform')) 

111 

112 @property 

113 def stations(self): 

114 return set( 

115 codes.station for codes in self._squirrel.get_codes('waveform')) 

116 

117 @property 

118 def locations(self): 

119 return set( 

120 codes.location for codes in self._squirrel.get_codes('waveform')) 

121 

122 @property 

123 def channels(self): 

124 return set( 

125 codes.channel for codes in self._squirrel.get_codes('waveform')) 

126 

127 def is_relevant(self, tmin, tmax): 

128 ptmin, ptmax = self._squirrel.get_time_span( 

129 ['waveform', 'waveform_promise']) 

130 

131 if None in (ptmin, ptmax): 

132 return False 

133 

134 return tmax >= ptmin and ptmax >= tmin 

135 

136 def load_files( 

137 self, filenames, 

138 filename_attributes=None, 

139 fileformat='mseed', 

140 cache=None, 

141 show_progress=True, 

142 update_progress=None): 

143 

144 self._squirrel.add( 

145 filenames, kinds='waveform', format=fileformat) 

146 

147 def chop( 

148 self, tmin, tmax, 

149 nut_selector=None, 

150 snap=(round, round), 

151 include_last=False, 

152 load_data=True, 

153 accessor_id='default'): 

154 

155 nuts = self._squirrel.get_waveform_nuts(tmin=tmin, tmax=tmax) 

156 

157 if load_data: 

158 traces = [ 

159 self._squirrel.get_content(nut, 'waveform', accessor_id) 

160 

161 for nut in nuts if nut_selector is None or nut_selector(nut)] 

162 

163 else: 

164 traces = [ 

165 trace.Trace(**nut.trace_kwargs) 

166 for nut in nuts if nut_selector is None or nut_selector(nut)] 

167 

168 self._squirrel.advance_accessor(accessor_id) 

169 

170 chopped = [] 

171 used_files = set() 

172 for tr in traces: 

173 if not load_data and tr.ydata is not None: 

174 tr = tr.copy(data=False) 

175 tr.ydata = None 

176 

177 try: 

178 chopped.append(tr.chop( 

179 tmin, tmax, 

180 inplace=False, 

181 snap=snap, 

182 include_last=include_last)) 

183 

184 except trace.NoData: 

185 pass 

186 

187 return chopped, used_files 

188 

189 def _process_chopped( 

190 self, chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin, 

191 tpad): 

192 

193 chopped.sort(key=lambda a: a.full_id) 

194 if degap: 

195 chopped = trace.degapper(chopped, maxgap=maxgap, maxlap=maxlap) 

196 

197 if not want_incomplete: 

198 chopped_weeded = [] 

199 for tr in chopped: 

200 emin = tr.tmin - (wmin-tpad) 

201 emax = tr.tmax + tr.deltat - (wmax+tpad) 

202 if (abs(emin) <= 0.5*tr.deltat and abs(emax) <= 0.5*tr.deltat): 

203 chopped_weeded.append(tr) 

204 

205 elif degap: 

206 if (0. < emin <= 5. * tr.deltat and 

207 -5. * tr.deltat <= emax < 0.): 

208 

209 tr.extend( 

210 wmin-tpad, 

211 wmax+tpad-tr.deltat, 

212 fillmethod='repeat') 

213 

214 chopped_weeded.append(tr) 

215 

216 chopped = chopped_weeded 

217 

218 for tr in chopped: 

219 tr.wmin = wmin 

220 tr.wmax = wmax 

221 

222 return chopped 

223 

224 def chopper( 

225 self, 

226 tmin=None, tmax=None, tinc=None, tpad=0., 

227 trace_selector=None, 

228 want_incomplete=True, degap=True, maxgap=5, maxlap=None, 

229 keep_current_files_open=False, accessor_id='default', 

230 snap=(round, round), include_last=False, load_data=True, 

231 style=None): 

232 

233 ''' 

234 Get iterator for shifting window wise data extraction from waveform 

235 archive. 

236 

237 :param tmin: start time (default uses start time of available data) 

238 :param tmax: end time (default uses end time of available data) 

239 :param tinc: time increment (window shift time) (default uses 

240 ``tmax-tmin``) 

241 :param tpad: padding time appended on either side of the data windows 

242 (window overlap is ``2*tpad``) 

243 :param trace_selector: filter callback taking 

244 :py:class:`pyrocko.trace.Trace` objects 

245 :param want_incomplete: if set to ``False``, gappy/incomplete traces 

246 are discarded from the results 

247 :param degap: whether to try to connect traces and to remove gaps and 

248 overlaps 

249 :param maxgap: maximum gap size in samples which is filled with 

250 interpolated samples when ``degap`` is ``True`` 

251 :param maxlap: maximum overlap size in samples which is removed when 

252 ``degap`` is ``True`` 

253 :param keep_current_files_open: whether to keep cached trace data in 

254 memory after the iterator has ended 

255 :param accessor_id: if given, used as a key to identify different 

256 points of extraction for the decision of when to release cached 

257 trace data (should be used when data is alternately extracted from 

258 more than one region / selection) 

259 :param snap: replaces Python's :py:func:`round` function which is used 

260 to determine indices where to start and end the trace data array 

261 :param include_last: whether to include last sample 

262 :param load_data: whether to load the waveform data. If set to 

263 ``False``, traces with no data samples, but with correct 

264 meta-information are returned 

265 :param style: set to ``'batch'`` to yield waveforms and information 

266 about the chopper state as :py:class:`pyrocko.pile.Batch` objects. 

267 By default lists of :py:class:`pyrocko.trace.Trace` objects are 

268 yielded. 

269 :returns: iterator providing extracted waveforms for each extracted 

270 window. See ``style`` argument for details. 

271 ''' 

272 

273 if tmin is None: 

274 if self.tmin is None: 

275 logger.warning("Pile's tmin is not set - pile may be empty.") 

276 return 

277 tmin = self.tmin + tpad 

278 

279 if tmax is None: 

280 if self.tmax is None: 

281 logger.warning("Pile's tmax is not set - pile may be empty.") 

282 return 

283 tmax = self.tmax - tpad 

284 

285 if tinc is None: 

286 tinc = tmax - tmin 

287 

288 if not self.is_relevant(tmin-tpad, tmax+tpad): 

289 return 

290 

291 nut_selector = trace_callback_to_nut_callback(trace_selector) 

292 

293 eps = tinc * 1e-6 

294 if tinc != 0.0: 

295 nwin = int(((tmax - eps) - tmin) / tinc) + 1 

296 else: 

297 nwin = 1 

298 

299 for iwin in range(nwin): 

300 wmin, wmax = tmin+iwin*tinc, min(tmin+(iwin+1)*tinc, tmax) 

301 

302 chopped, used_files = self.chop( 

303 wmin-tpad, wmax+tpad, nut_selector, snap, 

304 include_last, load_data, accessor_id) 

305 

306 processed = self._process_chopped( 

307 chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin, 

308 tpad) 

309 

310 if style == 'batch': 

311 yield classic_pile.Batch( 

312 tmin=wmin, 

313 tmax=wmax, 

314 i=iwin, 

315 n=nwin, 

316 traces=processed) 

317 

318 else: 

319 yield processed 

320 

321 if not keep_current_files_open: 

322 self._squirrel.clear_accessor(accessor_id, 'waveform') 

323 

324 def chopper_grouped(self, gather, progress=None, *args, **kwargs): 

325 keys = self.gather_keys(gather) 

326 if len(keys) == 0: 

327 return 

328 

329 outer_trace_selector = None 

330 if 'trace_selector' in kwargs: 

331 outer_trace_selector = kwargs['trace_selector'] 

332 

333 # the use of this gather-cache makes it impossible to modify the pile 

334 # during chopping 

335 pbar = None 

336 try: 

337 if progress is not None: 

338 pbar = util.progressbar(progress, len(keys)) 

339 

340 for ikey, key in enumerate(keys): 

341 def tsel(tr): 

342 return gather(tr) == key and ( 

343 outer_trace_selector is None or 

344 outer_trace_selector(tr)) 

345 

346 kwargs['trace_selector'] = tsel 

347 

348 for traces in self.chopper(*args, **kwargs): 

349 yield traces 

350 

351 if pbar: 

352 pbar.update(ikey+1) 

353 

354 finally: 

355 if pbar: 

356 pbar.finish() 

357 

358 def reload_modified(self): 

359 self._squirrel.reload() 

360 

361 def iter_traces( 

362 self, 

363 load_data=False, 

364 return_abspath=False, 

365 trace_selector=None): 

366 

367 ''' 

368 Iterate over all traces in pile. 

369 

370 :param load_data: whether to load the waveform data, by default empty 

371 traces are yielded 

372 :param return_abspath: if ``True`` yield tuples containing absolute 

373 file path and :py:class:`pyrocko.trace.Trace` objects 

374 :param trace_selector: filter callback taking 

375 :py:class:`pyrocko.trace.Trace` objects 

376 

377 ''' 

378 assert not load_data 

379 assert not return_abspath 

380 

381 nut_selector = trace_callback_to_nut_callback(trace_selector) 

382 

383 for nut in self._squirrel.get_waveform_nuts(): 

384 if nut_selector is None or nut_selector(nut): 

385 yield trace.Trace(**nut.trace_kwargs) 

386 

387 def gather_keys(self, gather, selector=None): 

388 codes_gather = trace_callback_to_codes_callback(gather) 

389 codes_selector = trace_callback_to_codes_callback(selector) 

390 return self._squirrel._gather_codes_keys( 

391 'waveform', codes_gather, codes_selector) 

392 

393 def snuffle(self, **kwargs): 

394 '''Visualize it. 

395 

396 :param stations: list of `pyrocko.model.Station` objects or ``None`` 

397 :param events: list of `pyrocko.model.Event` objects or ``None`` 

398 :param markers: list of `pyrocko.gui_util.Marker` objects or ``None`` 

399 :param ntracks: float, number of tracks to be shown initially 

400 (default: 12) 

401 :param follow: time interval (in seconds) for real time follow mode or 

402 ``None`` 

403 :param controls: bool, whether to show the main controls (default: 

404 ``True``) 

405 :param opengl: bool, whether to use opengl (default: ``False``) 

406 ''' 

407 

408 from pyrocko.gui.snuffler import snuffle 

409 snuffle(self, **kwargs) 

410 

411 def add_file(self, mtf): 

412 if isinstance(mtf, classic_pile.MemTracesFile): 

413 name = self._squirrel.add_volatile_waveforms(mtf.get_traces()) 

414 mtf._squirrel_name = name 

415 else: 

416 assert False 

417 

418 def remove_file(self, mtf): 

419 if isinstance(mtf, classic_pile.MemTracesFile) \ 

420 and getattr(mtf, '_squirrel_name', False): 

421 

422 self._squirrel.remove(mtf._squirrel_name) 

423 mtf._squirrel_name = None 

424 

425 def is_empty(self): 

426 return 'waveform' not in self._squirrel.get_kinds() 

427 

428 def get_update_count(self): 

429 return 0 

430 

431 

432def get_cache(_): 

433 return None