1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
7import weakref
8from pyrocko import squirrel as psq, trace
9from pyrocko import pile as classic_pile
11logger = logging.getLogger('psq.pile')
14def trace_callback_to_nut_callback(trace_callback):
15 if trace_callback is None:
16 return None
18 def nut_callback(nut):
19 return trace_callback(nut.dummy_trace)
21 return nut_callback
24class CodesDummyTrace(object):
25 def __init__(self, codes):
26 self.network, self.station, self.location, self.channel \
27 = self.nslc_id \
28 = codes[0:4]
31def trace_callback_to_codes_callback(trace_callback):
32 if trace_callback is None:
33 return None
35 def codes_callback(codes):
36 return trace_callback(CodesDummyTrace(codes))
38 return codes_callback
41class Pile(object):
42 '''
43 :py:class:`pyrocko.pile.Pile` surrogate: waveform lookup, loading and
44 caching.
46 This class emulates most of the older :py:class:`pyrocko.pile.Pile` methods
47 by using calls to a :py:class:`pyrocko.squirrel.base.Squirrel` instance
48 behind the scenes.
50 This interface can be used as a drop-in replacement for piles which are
51 used in existing scripts and programs for efficient waveform data access.
52 The Squirrel-based pile scales better for large datasets. Newer scripts
53 should use Squirrel's native methods to avoid the emulation overhead.
55 .. note::
56 Many methods in the original pile implementation lack documentation, as
57 do here. Read the source, Luke!
58 '''
59 def __init__(self, squirrel=None):
60 if squirrel is None:
61 squirrel = psq.Squirrel()
63 self._squirrel = squirrel
64 self._listeners = []
65 self._squirrel.get_database().add_listener(
66 self._notify_squirrel_to_pile)
68 def _notify_squirrel_to_pile(self, event, *args):
69 self.notify_listeners(event)
71 def add_listener(self, obj):
72 self._listeners.append(weakref.ref(obj))
74 def notify_listeners(self, what):
75 for ref in self._listeners:
76 obj = ref()
77 if obj:
78 obj.pile_changed(what)
80 def get_tmin(self):
81 return self.tmin
83 def get_tmax(self):
84 return self.tmax
86 def get_deltatmin(self):
87 return self._squirrel.get_deltat_span('waveform')[0]
89 def get_deltatmax(self):
90 return self._squirrel.get_deltat_span('waveform')[1]
92 @property
93 def deltatmin(self):
94 return self.get_deltatmin()
96 @property
97 def deltatmax(self):
98 return self.get_deltatmax()
100 @property
101 def tmin(self):
102 return self._squirrel.get_time_span('waveform')[0]
104 @property
105 def tmax(self):
106 return self._squirrel.get_time_span('waveform')[1]
108 @property
109 def networks(self):
110 return set(
111 codes.network for codes in self._squirrel.get_codes('waveform'))
113 @property
114 def stations(self):
115 return set(
116 codes.station for codes in self._squirrel.get_codes('waveform'))
118 @property
119 def locations(self):
120 return set(
121 codes.location for codes in self._squirrel.get_codes('waveform'))
123 @property
124 def channels(self):
125 return set(
126 codes.channel for codes in self._squirrel.get_codes('waveform'))
128 def is_relevant(self, tmin, tmax):
129 ptmin, ptmax = self._squirrel.get_time_span(
130 ['waveform', 'waveform_promise'])
132 if None in (ptmin, ptmax):
133 return False
135 return tmax >= ptmin and ptmax >= tmin
137 def load_files(
138 self, filenames,
139 filename_attributes=None,
140 fileformat='mseed',
141 cache=None,
142 show_progress=True,
143 update_progress=None):
145 self._squirrel.add(
146 filenames, kinds='waveform', format=fileformat)
148 def chop(
149 self, tmin, tmax,
150 nut_selector=None,
151 snap=(round, round),
152 include_last=False,
153 load_data=True,
154 accessor_id='default'):
156 nuts = self._squirrel.get_waveform_nuts(tmin=tmin, tmax=tmax)
158 if load_data:
159 traces = [
160 self._squirrel.get_content(nut, 'waveform', accessor_id)
162 for nut in nuts if nut_selector is None or nut_selector(nut)]
164 else:
165 traces = [
166 trace.Trace(**nut.trace_kwargs)
167 for nut in nuts if nut_selector is None or nut_selector(nut)]
169 self._squirrel.advance_accessor(accessor_id)
171 chopped = []
172 used_files = set()
173 for tr in traces:
174 if not load_data and tr.ydata is not None:
175 tr = tr.copy(data=False)
176 tr.ydata = None
178 try:
179 chopped.append(tr.chop(
180 tmin, tmax,
181 inplace=False,
182 snap=snap,
183 include_last=include_last))
185 except trace.NoData:
186 pass
188 return chopped, used_files
190 def _process_chopped(
191 self, chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
192 tpad):
194 chopped.sort(key=lambda a: a.full_id)
195 if degap:
196 chopped = trace.degapper(chopped, maxgap=maxgap, maxlap=maxlap)
198 if not want_incomplete:
199 chopped_weeded = []
200 for tr in chopped:
201 emin = tr.tmin - (wmin-tpad)
202 emax = tr.tmax + tr.deltat - (wmax+tpad)
203 if (abs(emin) <= 0.5*tr.deltat and abs(emax) <= 0.5*tr.deltat):
204 chopped_weeded.append(tr)
206 elif degap:
207 if (0. < emin <= 5. * tr.deltat and
208 -5. * tr.deltat <= emax < 0.):
210 tr.extend(
211 wmin-tpad,
212 wmax+tpad-tr.deltat,
213 fillmethod='repeat')
215 chopped_weeded.append(tr)
217 chopped = chopped_weeded
219 for tr in chopped:
220 tr.wmin = wmin
221 tr.wmax = wmax
223 return chopped
225 def chopper(
226 self,
227 tmin=None, tmax=None, tinc=None, tpad=0.,
228 group_selector=None, trace_selector=None,
229 want_incomplete=True, degap=True, maxgap=5, maxlap=None,
230 keep_current_files_open=False, accessor_id='default',
231 snap=(round, round), include_last=False, load_data=True,
232 style=None):
234 '''
235 Get iterator for shifting window wise data extraction from waveform
236 archive.
238 :param tmin: start time (default uses start time of available data)
239 :param tmax: end time (default uses end time of available data)
240 :param tinc: time increment (window shift time) (default uses
241 ``tmax-tmin``)
242 :param tpad: padding time appended on either side of the data windows
243 (window overlap is ``2*tpad``)
244 :param group_selector: *ignored in squirrel-based pile*
245 :param trace_selector: filter callback taking
246 :py:class:`pyrocko.trace.Trace` objects
247 :param want_incomplete: if set to ``False``, gappy/incomplete traces
248 are discarded from the results
249 :param degap: whether to try to connect traces and to remove gaps and
250 overlaps
251 :param maxgap: maximum gap size in samples which is filled with
252 interpolated samples when ``degap`` is ``True``
253 :param maxlap: maximum overlap size in samples which is removed when
254 ``degap`` is ``True``
255 :param keep_current_files_open: whether to keep cached trace data in
256 memory after the iterator has ended
257 :param accessor_id: if given, used as a key to identify different
258 points of extraction for the decision of when to release cached
259 trace data (should be used when data is alternately extracted from
260 more than one region / selection)
261 :param snap: replaces Python's :py:func:`round` function which is used
262 to determine indices where to start and end the trace data array
263 :param include_last: whether to include last sample
264 :param load_data: whether to load the waveform data. If set to
265 ``False``, traces with no data samples, but with correct
266 meta-information are returned
267 :param style: set to ``'batch'`` to yield waveforms and information
268 about the chopper state as :py:class:`pyrocko.pile.Batch` objects.
269 By default lists of :py:class:`pyrocko.trace.Trace` objects are
270 yielded.
271 :returns: iterator providing extracted waveforms for each extracted
272 window. See ``style`` argument for details.
273 '''
275 if tmin is None:
276 if self.tmin is None:
277 logger.warning('Pile\'s tmin is not set - pile may be empty.')
278 return
279 tmin = self.tmin + tpad
281 if tmax is None:
282 if self.tmax is None:
283 logger.warning('Pile\'s tmax is not set - pile may be empty.')
284 return
285 tmax = self.tmax - tpad
287 if tinc is None:
288 tinc = tmax - tmin
290 if not self.is_relevant(tmin-tpad, tmax+tpad):
291 return
293 nut_selector = trace_callback_to_nut_callback(trace_selector)
295 eps = tinc * 1e-6
296 if tinc != 0.0:
297 nwin = int(((tmax - eps) - tmin) / tinc) + 1
298 else:
299 nwin = 1
301 for iwin in range(nwin):
302 wmin, wmax = tmin+iwin*tinc, min(tmin+(iwin+1)*tinc, tmax)
304 chopped, used_files = self.chop(
305 wmin-tpad, wmax+tpad, nut_selector, snap,
306 include_last, load_data, accessor_id)
308 processed = self._process_chopped(
309 chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
310 tpad)
312 if style == 'batch':
313 yield classic_pile.Batch(
314 tmin=wmin,
315 tmax=wmax,
316 i=iwin,
317 n=nwin,
318 traces=processed)
320 else:
321 yield processed
323 if not keep_current_files_open:
324 self._squirrel.clear_accessor(accessor_id, 'waveform')
326 def chopper_grouped(self, gather, progress=None, *args, **kwargs):
327 raise NotImplementedError
329 def reload_modified(self):
330 self._squirrel.reload()
332 def iter_traces(
333 self,
334 load_data=False,
335 return_abspath=False,
336 group_selector=None,
337 trace_selector=None):
339 '''
340 Iterate over all traces in pile.
342 :param load_data: whether to load the waveform data, by default empty
343 traces are yielded
344 :param return_abspath: if ``True`` yield tuples containing absolute
345 file path and :py:class:`pyrocko.trace.Trace` objects
346 :param group_selector: *ignored in squirre-based pile*
347 :param trace_selector: filter callback taking
348 :py:class:`pyrocko.trace.Trace` objects
350 '''
351 assert not load_data
352 assert not return_abspath
354 nut_selector = trace_callback_to_nut_callback(trace_selector)
356 for nut in self._squirrel.get_waveform_nuts():
357 if nut_selector is None or nut_selector(nut):
358 yield trace.Trace(**nut.trace_kwargs)
360 def gather_keys(self, gather, selector=None):
361 codes_gather = trace_callback_to_codes_callback(gather)
362 codes_selector = trace_callback_to_codes_callback(selector)
363 return self._squirrel._gather_codes_keys(
364 'waveform', codes_gather, codes_selector)
366 def snuffle(self, **kwargs):
367 '''Visualize it.
369 :param stations: list of `pyrocko.model.Station` objects or ``None``
370 :param events: list of `pyrocko.model.Event` objects or ``None``
371 :param markers: list of `pyrocko.gui_util.Marker` objects or ``None``
372 :param ntracks: float, number of tracks to be shown initially
373 (default: 12)
374 :param follow: time interval (in seconds) for real time follow mode or
375 ``None``
376 :param controls: bool, whether to show the main controls (default:
377 ``True``)
378 :param opengl: bool, whether to use opengl (default: ``False``)
379 '''
381 from pyrocko.gui.snuffler import snuffle
382 snuffle(self, **kwargs)
384 def add_file(self, mtf):
385 if isinstance(mtf, classic_pile.MemTracesFile):
386 name = self._squirrel.add_volatile_waveforms(mtf.get_traces())
387 mtf._squirrel_name = name
388 else:
389 assert False
391 def remove_file(self, mtf):
392 if isinstance(mtf, classic_pile.MemTracesFile) \
393 and getattr(mtf, '_squirrel_name', False):
395 self._squirrel.remove(mtf._squirrel_name)
396 mtf._squirrel_name = None
398 def is_empty(self):
399 return 'waveform' not in self._squirrel.get_kinds()
401 def get_update_count(self):
402 return 0
405def get_cache(_):
406 return None