1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function
8import logging
9import re
11from ..model import QuantityType, match_codes, CodesNSLCE
12from .. import error
14from pyrocko.guts import Object, String, Duration, Float, clone, List
15from pyrocko import model
17guts_prefix = 'squirrel.ops'
19logger = logging.getLogger('psq.ops')
22def odiff(a, b):
23 ia = ib = 0
24 only_a = []
25 only_b = []
26 while ia < len(a) or ib < len(b):
27 # TODO remove when finished with implementation
28 if ia > 0:
29 assert a[ia] > a[ia-1]
30 if ib > 0:
31 assert b[ib] > b[ib-1]
33 if ib == len(b) or (ia < len(a) and a[ia] < b[ib]):
34 only_a.append(a[ia])
35 ia += 1
36 elif ia == len(a) or (ib < len(b) and a[ia] > b[ib]):
37 only_b.append(b[ib])
38 ib += 1
39 elif a[ia] == b[ib]:
40 ia += 1
41 ib += 1
43 return only_a, only_b
46def _cglob_translate(creg):
47 dd = []
48 for c in creg:
49 if c == '*':
50 d = r'[^.]*'
51 elif c == '?':
52 d = r'[^.]'
53 elif c == '.':
54 d = r'\.'
55 else:
56 d = c
58 dd.append(d)
59 reg = ''.join(dd)
60 return reg
63class Filter(Object):
65 def filter(self, it, squirrel=None, tmin=None, tmax=None):
66 return list(it)
69class RegexFilter(Filter):
70 pattern = String.T(default=r'(.*)')
72 def __init__(self, **kwargs):
73 Filter.__init__(self, **kwargs)
74 self._compiled_pattern = re.compile(self.pattern)
76 def filter(self, it, squirrel=None, tmin=None, tmax=None):
77 return [
78 x for x in it if self._compiled_pattern.fullmatch(x)]
81class CodesPatternFilter(Filter):
82 codes = List.T(CodesNSLCE.T(), optional=True)
84 def filter(self, it, squirrel=None, tmin=None, tmax=None):
85 if self.codes is None:
86 return list(it)
87 else:
88 return [
89 x for x in it
90 if any(match_codes(sc, x) for sc in self.codes)]
93class DistanceFilter(Filter):
94 location = model.Location.T()
95 distance_min = Float.T(optional=True)
96 distance_max = Float.T(optional=True)
98 def filter(self, it, squirrel=None, tmin=None, tmax=None):
99 locations = squirrel.get_location_pool(tmin=tmin, tmax=tmax)
101 dist_min = self.distance_min
102 dist_max = self.distance_max
103 filtered = []
104 for codes in it:
105 loc = locations.get(codes)
106 if loc is not None:
107 dist = self.location.distance_to(loc)
108 if (dist_min is None or dist_min <= dist) \
109 and (dist_max is None or dist <= dist_max):
111 filtered.append(codes)
113 return filtered
116class Grouping(Object):
118 def key(self, codes):
119 return codes
122class RegexGrouping(Grouping):
123 pattern = String.T(default=r'(.*)')
125 def __init__(self, **kwargs):
126 Grouping.__init__(self, **kwargs)
127 self._compiled_pattern = re.compile(self.pattern)
129 def key(self, codes):
130 return self._compiled_pattern.fullmatch(codes.safe_str).groups()
133class NetworkGrouping(RegexGrouping):
134 '''
135 Group by *network* code.
136 '''
137 pattern = String.T(default=_cglob_translate('(*).*.*.*.*'))
140class StationGrouping(RegexGrouping):
141 '''
142 Group by *network.station* codes.
143 '''
144 pattern = String.T(default=_cglob_translate('(*.*).*.*.*'))
147class LocationGrouping(RegexGrouping):
148 '''
149 Group by *network.station.location* codes.
150 '''
151 pattern = String.T(default=_cglob_translate('(*.*.*).*.*'))
154class ChannelGrouping(RegexGrouping):
155 '''
156 Group by *network.station.location.channel* codes.
158 This effectively groups all processings of a channel, which may differ in
159 the *extra* codes attribute.
160 '''
161 pattern = String.T(default=_cglob_translate('(*.*.*.*).*'))
164class SensorGrouping(RegexGrouping):
165 '''
166 Group by *network.station.location.sensor* and *extra* codes.
168 For *sensor* all but the last character of the channel code (indicating the
169 component) are used. This effectively groups all components of a sensor,
170 or processings of a sensor.
171 '''
172 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
175class Translation(Object):
177 def translate(self, codes):
178 return codes
181class AddSuffixTranslation(Translation):
182 suffix = String.T(default='')
184 def translate(self, codes):
185 return codes.replace(extra=codes.extra + self.suffix)
188class RegexTranslation(AddSuffixTranslation):
189 pattern = String.T(default=r'(.*)')
190 replacement = String.T(default=r'\1')
192 def __init__(self, **kwargs):
193 AddSuffixTranslation.__init__(self, **kwargs)
194 self._compiled_pattern = re.compile(self.pattern)
196 def translate(self, codes):
197 return codes.__class__(
198 self._compiled_pattern.sub(self.replacement, codes.safe_str))
201class ReplaceComponentTranslation(RegexTranslation):
202 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
203 replacement = String.T(default=r'\1{component}\2')
206def deregister(registry, group):
207 for codes in group[2]:
208 del registry[codes]
211def register(registry, operator, group):
212 for codes in group[2]:
213 if codes in registry:
214 logger.warning(
215 'duplicate operator output codes: %s' % codes)
216 registry[codes] = (operator, group)
219class Operator(Object):
221 filter = Filter.T(default=Filter.D())
222 grouping = Grouping.T(default=Grouping.D())
223 translation = Translation.T(default=Translation.D())
225 def __init__(self, **kwargs):
226 Object.__init__(self, **kwargs)
227 self._output_names_cache = {}
228 self._groups = {}
229 self._available = []
231 @property
232 def name(self):
233 return self.__class__.__name__
235 def describe(self):
236 return self.name
238 def iter_mappings(self):
239 for k, group in self._groups.items():
240 yield (group[1], group[2])
242 def iter_in_codes(self):
243 for k, group in self._groups.items():
244 yield group[1]
246 def iter_out_codes(self):
247 for k, group in self._groups.items():
248 yield group[2]
250 def update_mappings(self, available, squirrel, tmin, tmax, registry=None):
251 available = list(available)
252 removed, added = odiff(self._available, available)
254 filt = self.filter.filter
255 gkey = self.grouping.key
256 groups = self._groups
258 need_update = set()
260 for codes in filt(removed, squirrel, tmin, tmax):
261 k = gkey(codes)
262 groups[k][0].remove(codes)
263 need_update.add(k)
265 for codes in filt(added, squirrel, tmin, tmax):
266 k = gkey(codes)
267 if k not in groups:
268 groups[k] = [set(), None, ()]
270 groups[k][0].add(codes)
271 need_update.add(k)
273 for k in need_update:
274 group = groups[k]
275 if registry is not None:
276 deregister(registry, group)
278 group[1] = tuple(sorted(group[0]))
279 if not group[1]:
280 del groups[k]
281 else:
282 group[2] = self._out_codes(group[1])
283 if registry is not None:
284 register(registry, self, group)
286 self._available = available
288 def _out_codes(self, in_codes):
289 return [self.translation.translate(codes) for codes in in_codes]
291 def get_channels(self, squirrel, group, *args, **kwargs):
292 _, in_codes, out_codes = group
293 assert len(in_codes) == 1 and len(out_codes) == 1
294 channels = squirrel.get_channels(codes=in_codes[0], **kwargs)
295 channels_out = []
296 for channel in channels:
297 channel_out = clone(channel)
298 channel_out.codes = out_codes[0]
299 channels_out.append(channel_out)
301 return channels_out
303 def get_waveforms(self, squirrel, group, **kwargs):
304 _, in_codes, out_codes = group
305 assert len(in_codes) == 1 and len(out_codes) == 1
307 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs)
308 for tr in trs:
309 tr.set_codes(*out_codes[0])
311 return trs
313 # def update_waveforms(self, squirrel, tmin, tmax, codes):
314 # if codes is None:
315 # for _, in_codes, out_codes in self._groups.values():
316 # for codes in
319class Parameters(Object):
320 pass
323class RestitutionParameters(Parameters):
324 frequency_min = Float.T()
325 frequency_max = Float.T()
326 frequency_taper_factor = Float.T(default=1.5)
327 time_taper_factor = Float.T(default=2.0)
330class Restitution(Operator):
331 translation = AddSuffixTranslation(suffix='R{quantity}')
332 quantity = QuantityType.T(default='displacement')
334 @property
335 def name(self):
336 return 'Restitution(%s)' % self.quantity[0]
338 def _out_codes(self, group):
339 return [
340 codes.__class__(self.translation.translate(codes).safe_str.format(
341 quantity=self.quantity[0]))
342 for codes in group]
344 def get_tpad(self, params):
345 return params.time_taper_factor / params.frequency_min
347 def get_waveforms(
348 self, squirrel, codes, params, tmin, tmax, **kwargs):
350 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
351 assert self is self_
352 assert len(in_codes) == 1 and len(out_codes) == 1
354 tpad = self.get_tpad(params)
356 tmin_raw = tmin - tpad
357 tmax_raw = tmax + tpad
359 trs = squirrel.get_waveforms(
360 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs)
362 try:
363 resp = squirrel.get_response(
364 codes=in_codes[0],
365 tmin=tmin_raw,
366 tmax=tmax_raw).get_effective(self.quantity)
368 except error.NotAvailable:
369 return []
371 freqlimits = (
372 params.frequency_min / params.frequency_taper_factor,
373 params.frequency_min,
374 params.frequency_max,
375 params.frequency_max * params.frequency_taper_factor)
377 trs_rest = []
378 for tr in trs:
379 tr_rest = tr.transfer(
380 tfade=tpad,
381 freqlimits=freqlimits,
382 transfer_function=resp,
383 invert=True)
385 tr_rest.set_codes(*out_codes[0])
387 trs_rest.append(tr_rest)
389 return trs_rest
392class Shift(Operator):
393 translation = AddSuffixTranslation(suffix='S')
394 delay = Duration.T()
397class Transform(Operator):
398 grouping = Grouping.T(default=SensorGrouping.D())
399 translation = ReplaceComponentTranslation(suffix='T{system}')
401 def _out_codes(self, group):
402 return [
403 self.translation.translate(group[0]).format(
404 component=c, system=self.components.lower())
405 for c in self.components]
408class ToENZ(Transform):
409 components = 'ENZ'
411 def get_waveforms(
412 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
414 trs = squirrel.get_waveforms(
415 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs)
417 for tr in trs:
418 print(tr)
421class ToRTZ(Transform):
422 components = 'RTZ'
423 backazimuth = Float.T()
426class ToLTQ(Transform):
427 components = 'LTQ'
430class Composition(Operator):
431 g = Operator.T()
432 f = Operator.T()
434 def __init__(self, g, f, **kwargs):
435 Operator.__init__(self, g=g, f=f, **kwargs)
437 @property
438 def name(self):
439 return '(%s ○ %s)' % (self.g.name, self.f.name)
442__all__ = [
443 'Filter',
444 'RegexFilter',
445 'CodesPatternFilter',
446 'DistanceFilter',
447 'Grouping',
448 'RegexGrouping',
449 'NetworkGrouping',
450 'StationGrouping',
451 'LocationGrouping',
452 'SensorGrouping',
453 'ChannelGrouping',
454 'Operator',
455 'RestitutionParameters',
456 'Restitution',
457 'Shift',
458 'ToENZ',
459 'ToRTZ',
460 'ToLTQ',
461 'Composition']