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[1:5]
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(codes[1] for codes in self._squirrel.get_codes('waveform'))
112 @property
113 def stations(self):
114 return set(codes[2] for codes in self._squirrel.get_codes('waveform'))
116 @property
117 def locations(self):
118 return set(codes[3] for codes in self._squirrel.get_codes('waveform'))
120 @property
121 def channels(self):
122 return set(codes[4] for codes in self._squirrel.get_codes('waveform'))
124 def is_relevant(self, tmin, tmax):
125 ptmin, ptmax = self._squirrel.get_time_span(
126 ['waveform', 'waveform_promise'])
128 if None in (ptmin, ptmax):
129 return False
131 return tmax >= ptmin and ptmax >= tmin
133 def load_files(
134 self, filenames,
135 filename_attributes=None,
136 fileformat='mseed',
137 cache=None,
138 show_progress=True,
139 update_progress=None):
141 self._squirrel.add(
142 filenames, kinds='waveform', format=fileformat)
144 def chop(
145 self, tmin, tmax,
146 nut_selector=None,
147 snap=(round, round),
148 include_last=False,
149 load_data=True,
150 accessor_id='default'):
152 nuts = self._squirrel.get_waveform_nuts(tmin=tmin, tmax=tmax)
154 if load_data:
155 traces = [
156 self._squirrel.get_content(nut, 'waveform', accessor_id)
158 for nut in nuts if nut_selector is None or nut_selector(nut)]
160 else:
161 traces = [
162 trace.Trace(**nut.trace_kwargs)
163 for nut in nuts if nut_selector is None or nut_selector(nut)]
165 self._squirrel.advance_accessor(accessor_id)
167 chopped = []
168 used_files = set()
169 for tr in traces:
170 if not load_data and tr.ydata is not None:
171 tr = tr.copy(data=False)
172 tr.ydata = None
174 try:
175 chopped.append(tr.chop(
176 tmin, tmax,
177 inplace=False,
178 snap=snap,
179 include_last=include_last))
181 except trace.NoData:
182 pass
184 return chopped, used_files
186 def _process_chopped(
187 self, chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
188 tpad):
190 chopped.sort(key=lambda a: a.full_id)
191 if degap:
192 chopped = trace.degapper(chopped, maxgap=maxgap, maxlap=maxlap)
194 if not want_incomplete:
195 chopped_weeded = []
196 for tr in chopped:
197 emin = tr.tmin - (wmin-tpad)
198 emax = tr.tmax + tr.deltat - (wmax+tpad)
199 if (abs(emin) <= 0.5*tr.deltat and abs(emax) <= 0.5*tr.deltat):
200 chopped_weeded.append(tr)
202 elif degap:
203 if (0. < emin <= 5. * tr.deltat and
204 -5. * tr.deltat <= emax < 0.):
206 tr.extend(
207 wmin-tpad,
208 wmax+tpad-tr.deltat,
209 fillmethod='repeat')
211 chopped_weeded.append(tr)
213 chopped = chopped_weeded
215 for tr in chopped:
216 tr.wmin = wmin
217 tr.wmax = wmax
219 return chopped
221 def chopper(
222 self,
223 tmin=None, tmax=None, tinc=None, tpad=0.,
224 group_selector=None, trace_selector=None,
225 want_incomplete=True, degap=True, maxgap=5, maxlap=None,
226 keep_current_files_open=False, accessor_id='default',
227 snap=(round, round), include_last=False, load_data=True,
228 style=None):
230 '''
231 Get iterator for shifting window wise data extraction from waveform
232 archive.
234 :param tmin: start time (default uses start time of available data)
235 :param tmax: end time (default uses end time of available data)
236 :param tinc: time increment (window shift time) (default uses
237 ``tmax-tmin``)
238 :param tpad: padding time appended on either side of the data windows
239 (window overlap is ``2*tpad``)
240 :param group_selector: *ignored in squirrel-based pile*
241 :param trace_selector: filter callback taking
242 :py:class:`pyrocko.trace.Trace` objects
243 :param want_incomplete: if set to ``False``, gappy/incomplete traces
244 are discarded from the results
245 :param degap: whether to try to connect traces and to remove gaps and
246 overlaps
247 :param maxgap: maximum gap size in samples which is filled with
248 interpolated samples when ``degap`` is ``True``
249 :param maxlap: maximum overlap size in samples which is removed when
250 ``degap`` is ``True``
251 :param keep_current_files_open: whether to keep cached trace data in
252 memory after the iterator has ended
253 :param accessor_id: if given, used as a key to identify different
254 points of extraction for the decision of when to release cached
255 trace data (should be used when data is alternately extracted from
256 more than one region / selection)
257 :param snap: replaces Python's :py:func:`round` function which is used
258 to determine indices where to start and end the trace data array
259 :param include_last: whether to include last sample
260 :param load_data: whether to load the waveform data. If set to
261 ``False``, traces with no data samples, but with correct
262 meta-information are returned
263 :param style: set to ``'batch'`` to yield waveforms and information
264 about the chopper state as :py:class:`pyrocko.pile.Batch` objects. By
265 default lists of :py:class:`pyrocko.trace.Trace` objects are yielded.
266 :returns: iterator providing extracted waveforms for each extracted
267 window. See ``style`` argument for details.
268 '''
270 if tmin is None:
271 if self.tmin is None:
272 logger.warning('Pile\'s tmin is not set - pile may be empty.')
273 return
274 tmin = self.tmin + tpad
276 if tmax is None:
277 if self.tmax is None:
278 logger.warning('Pile\'s tmax is not set - pile may be empty.')
279 return
280 tmax = self.tmax - tpad
282 if tinc is None:
283 tinc = tmax - tmin
285 if not self.is_relevant(tmin-tpad, tmax+tpad):
286 return
288 nut_selector = trace_callback_to_nut_callback(trace_selector)
290 eps = tinc * 1e-6
291 if tinc != 0.0:
292 nwin = int(((tmax - eps) - tmin) / tinc) + 1
293 else:
294 nwin = 1
296 for iwin in range(nwin):
297 wmin, wmax = tmin+iwin*tinc, min(tmin+(iwin+1)*tinc, tmax)
299 chopped, used_files = self.chop(
300 wmin-tpad, wmax+tpad, nut_selector, snap,
301 include_last, load_data, accessor_id)
303 processed = self._process_chopped(
304 chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
305 tpad)
307 if style == 'batch':
308 yield classic_pile.Batch(
309 tmin=wmin,
310 tmax=wmax,
311 i=iwin,
312 n=nwin,
313 traces=processed)
315 else:
316 yield processed
318 if not keep_current_files_open:
319 self._squirrel.clear_accessor(accessor_id, 'waveform')
321 def chopper_grouped(self, gather, progress=None, *args, **kwargs):
322 raise NotImplementedError
324 def reload_modified(self):
325 self._squirrel.reload()
327 def iter_traces(
328 self,
329 load_data=False,
330 return_abspath=False,
331 group_selector=None,
332 trace_selector=None):
334 '''
335 Iterate over all traces in pile.
337 :param load_data: whether to load the waveform data, by default empty
338 traces are yielded
339 :param return_abspath: if ``True`` yield tuples containing absolute
340 file path and :py:class:`pyrocko.trace.Trace` objects
341 :param group_selector: *ignored in squirre-based pile*
342 :param trace_selector: filter callback taking
343 :py:class:`pyrocko.trace.Trace` objects
345 '''
346 assert not load_data
347 assert not return_abspath
349 nut_selector = trace_callback_to_nut_callback(trace_selector)
351 for nut in self._squirrel.get_waveform_nuts():
352 if nut_selector is None or nut_selector(nut):
353 yield trace.Trace(**nut.trace_kwargs)
355 def gather_keys(self, gather, selector=None):
356 codes_gather = trace_callback_to_codes_callback(gather)
357 codes_selector = trace_callback_to_codes_callback(selector)
358 return self._squirrel._gather_codes_keys(
359 'waveform', codes_gather, codes_selector)
361 def snuffle(self, **kwargs):
362 '''Visualize it.
364 :param stations: list of `pyrocko.model.Station` objects or ``None``
365 :param events: list of `pyrocko.model.Event` objects or ``None``
366 :param markers: list of `pyrocko.gui_util.Marker` objects or ``None``
367 :param ntracks: float, number of tracks to be shown initially
368 (default: 12)
369 :param follow: time interval (in seconds) for real time follow mode or
370 ``None``
371 :param controls: bool, whether to show the main controls (default:
372 ``True``)
373 :param opengl: bool, whether to use opengl (default: ``False``)
374 '''
376 from pyrocko.gui.snuffler import snuffle
377 snuffle(self, **kwargs)
379 def add_file(self, mtf):
380 if isinstance(mtf, classic_pile.MemTracesFile):
381 name = self._squirrel.add_volatile_waveforms(mtf.get_traces())
382 mtf._squirrel_name = name
383 else:
384 assert False
386 def remove_file(self, mtf):
387 if isinstance(mtf, classic_pile.MemTracesFile) \
388 and getattr(mtf, '_squirrel_name', False):
390 self._squirrel.remove(mtf._squirrel_name)
391 mtf._squirrel_name = None
393 def is_empty(self):
394 return 'waveform' not in self._squirrel.get_kinds()
396 def get_update_count(self):
397 return 0
400def get_cache(_):
401 return None