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
16guts_prefix = 'squirrel.ops'
18logger = logging.getLogger('psq.ops')
21def odiff(a, b):
22 ia = ib = 0
23 only_a = []
24 only_b = []
25 while ia < len(a) or ib < len(b):
26 # TODO remove when finished with implementation
27 if ia > 0:
28 assert a[ia] > a[ia-1]
29 if ib > 0:
30 assert b[ib] > b[ib-1]
32 if ib == len(b) or (ia < len(a) and a[ia] < b[ib]):
33 only_a.append(a[ia])
34 ia += 1
35 elif ia == len(a) or (ib < len(b) and a[ia] > b[ib]):
36 only_b.append(b[ib])
37 ib += 1
38 elif a[ia] == b[ib]:
39 ia += 1
40 ib += 1
42 return only_a, only_b
45def _cglob_translate(creg):
46 dd = []
47 for c in creg:
48 if c == '*':
49 d = r'[^.]*'
50 elif c == '?':
51 d = r'[^.]'
52 elif c == '.':
53 d = r'\.'
54 else:
55 d = c
57 dd.append(d)
58 reg = ''.join(dd)
59 return reg
62class Filtering(Object):
64 def filter(self, it):
65 return list(it)
68class RegexFiltering(Object):
69 pattern = String.T(default=r'(.*)')
71 def __init__(self, **kwargs):
72 Filtering.__init__(self, **kwargs)
73 self._compiled_pattern = re.compile(self.pattern)
75 def filter(self, it):
76 return [
77 x for x in it if self._compiled_pattern.fullmatch(x)]
80class CodesPatternFiltering(Object):
81 codes = List.T(CodesNSLCE.T(), optional=True)
83 def filter(self, it):
84 if self.codes is None:
85 return list(it)
86 else:
87 return [
88 x for x in it
89 if any(match_codes(sc, x) for sc in self.codes)]
92class Grouping(Object):
94 def key(self, codes):
95 return codes
98class RegexGrouping(Grouping):
99 pattern = String.T(default=r'(.*)')
101 def __init__(self, **kwargs):
102 Grouping.__init__(self, **kwargs)
103 self._compiled_pattern = re.compile(self.pattern)
105 def key(self, codes):
106 return self._compiled_pattern.fullmatch(codes.safe_str).groups()
109class NetworkGrouping(RegexGrouping):
110 '''
111 Group by *network* code.
112 '''
113 pattern = String.T(default=_cglob_translate('(*).*.*.*.*'))
116class StationGrouping(RegexGrouping):
117 '''
118 Group by *network.station* codes.
119 '''
120 pattern = String.T(default=_cglob_translate('(*.*).*.*.*'))
123class LocationGrouping(RegexGrouping):
124 '''
125 Group by *network.station.location* codes.
126 '''
127 pattern = String.T(default=_cglob_translate('(*.*.*).*.*'))
130class ChannelGrouping(RegexGrouping):
131 '''
132 Group by *network.station.location.channel* codes.
134 This effectively groups all processings of a channel, which may differ in
135 the *extra* codes attribute.
136 '''
137 pattern = String.T(default=_cglob_translate('(*.*.*.*).*'))
140class SensorGrouping(RegexGrouping):
141 '''
142 Group by *network.station.location.sensor* and *extra* codes.
144 For *sensor* all but the last character of the channel code (indicating the
145 component) are used. This effectively groups all components of a sensor,
146 or processings of a sensor.
147 '''
148 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
151class Translation(Object):
153 def translate(self, codes):
154 return codes
157class AddSuffixTranslation(Translation):
158 suffix = String.T(default='')
160 def translate(self, codes):
161 return codes.replace(extra=codes.extra + self.suffix)
164class RegexTranslation(AddSuffixTranslation):
165 pattern = String.T(default=r'(.*)')
166 replacement = String.T(default=r'\1')
168 def __init__(self, **kwargs):
169 AddSuffixTranslation.__init__(self, **kwargs)
170 self._compiled_pattern = re.compile(self.pattern)
172 def translate(self, codes):
173 return codes.__class__(
174 self._compiled_pattern.sub(self.replacement, codes.safe_str))
177class ReplaceComponentTranslation(RegexTranslation):
178 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
179 replacement = String.T(default=r'\1{component}\2')
182def deregister(registry, group):
183 for codes in group[2]:
184 del registry[codes]
187def register(registry, operator, group):
188 for codes in group[2]:
189 if codes in registry:
190 logger.warning(
191 'duplicate operator output codes: %s' % codes)
192 registry[codes] = (operator, group)
195class Operator(Object):
197 filtering = Filtering.T(default=Filtering.D())
198 grouping = Grouping.T(default=Grouping.D())
199 translation = Translation.T(default=Translation.D())
201 def __init__(self, **kwargs):
202 Object.__init__(self, **kwargs)
203 self._output_names_cache = {}
204 self._groups = {}
205 self._available = []
207 @property
208 def name(self):
209 return self.__class__.__name__
211 def describe(self):
212 return self.name
214 def iter_mappings(self):
215 for k, group in self._groups.items():
216 yield (group[1], group[2])
218 def iter_in_codes(self):
219 for k, group in self._groups.items():
220 yield group[1]
222 def iter_out_codes(self):
223 for k, group in self._groups.items():
224 yield group[2]
226 def update_mappings(self, available, registry=None):
227 available = list(available)
228 removed, added = odiff(self._available, available)
230 filt = self.filtering.filter
231 gkey = self.grouping.key
232 groups = self._groups
234 need_update = set()
236 for codes in filt(removed):
237 k = gkey(codes)
238 groups[k][0].remove(codes)
239 need_update.add(k)
241 for codes in filt(added):
242 k = gkey(codes)
243 if k not in groups:
244 groups[k] = [set(), None, ()]
246 groups[k][0].add(codes)
247 need_update.add(k)
249 for k in need_update:
250 group = groups[k]
251 if registry is not None:
252 deregister(registry, group)
254 group[1] = tuple(sorted(group[0]))
255 if not group[1]:
256 del groups[k]
257 else:
258 group[2] = self._out_codes(group[1])
259 if registry is not None:
260 register(registry, self, group)
262 self._available = available
264 def _out_codes(self, in_codes):
265 return [self.translation.translate(codes) for codes in in_codes]
267 def get_channels(self, squirrel, group, *args, **kwargs):
268 _, in_codes, out_codes = group
269 assert len(in_codes) == 1 and len(out_codes) == 1
270 channels = squirrel.get_channels(codes=in_codes[0], **kwargs)
271 channels_out = []
272 for channel in channels:
273 channel_out = clone(channel)
274 channel_out.codes = out_codes[0]
275 channels_out.append(channel_out)
277 return channels_out
279 def get_waveforms(self, squirrel, group, **kwargs):
280 _, in_codes, out_codes = group
281 assert len(in_codes) == 1 and len(out_codes) == 1
283 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs)
284 for tr in trs:
285 tr.set_codes(*out_codes[0])
287 return trs
289 # def update_waveforms(self, squirrel, tmin, tmax, codes):
290 # if codes is None:
291 # for _, in_codes, out_codes in self._groups.values():
292 # for codes in
295class Parameters(Object):
296 pass
299class RestitutionParameters(Parameters):
300 frequency_min = Float.T()
301 frequency_max = Float.T()
302 frequency_taper_factor = Float.T(default=1.5)
303 time_taper_factor = Float.T(default=2.0)
306class Restitution(Operator):
307 translation = AddSuffixTranslation(suffix='R{quantity}')
308 quantity = QuantityType.T(default='displacement')
310 @property
311 def name(self):
312 return 'Restitution(%s)' % self.quantity[0]
314 def _out_codes(self, group):
315 return [
316 codes.__class__(self.translation.translate(codes).safe_str.format(
317 quantity=self.quantity[0]))
318 for codes in group]
320 def get_tpad(self, params):
321 return params.time_taper_factor / params.frequency_min
323 def get_waveforms(
324 self, squirrel, codes, params, tmin, tmax, **kwargs):
326 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
327 assert self is self_
328 assert len(in_codes) == 1 and len(out_codes) == 1
330 tpad = self.get_tpad(params)
332 tmin_raw = tmin - tpad
333 tmax_raw = tmax + tpad
335 trs = squirrel.get_waveforms(
336 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs)
338 try:
339 resp = squirrel.get_response(
340 codes=in_codes[0],
341 tmin=tmin_raw,
342 tmax=tmax_raw).get_effective(self.quantity)
344 except error.NotAvailable:
345 return []
347 freqlimits = (
348 params.frequency_min / params.frequency_taper_factor,
349 params.frequency_min,
350 params.frequency_max,
351 params.frequency_max * params.frequency_taper_factor)
353 trs_rest = []
354 for tr in trs:
355 tr_rest = tr.transfer(
356 tfade=tpad,
357 freqlimits=freqlimits,
358 transfer_function=resp,
359 invert=True)
361 tr_rest.set_codes(*out_codes[0])
363 trs_rest.append(tr_rest)
365 return trs_rest
368class Shift(Operator):
369 translation = AddSuffixTranslation(suffix='S')
370 delay = Duration.T()
373class Transform(Operator):
374 grouping = Grouping.T(default=SensorGrouping.D())
375 translation = ReplaceComponentTranslation(suffix='T{system}')
377 def _out_codes(self, group):
378 return [
379 self.translation.translate(group[0]).format(
380 component=c, system=self.components.lower())
381 for c in self.components]
384class ToENZ(Transform):
385 components = 'ENZ'
387 def get_waveforms(
388 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
390 trs = squirrel.get_waveforms(
391 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs)
393 for tr in trs:
394 print(tr)
397class ToRTZ(Transform):
398 components = 'RTZ'
399 backazimuth = Float.T()
402class ToLTQ(Transform):
403 components = 'LTQ'
406class Composition(Operator):
407 g = Operator.T()
408 f = Operator.T()
410 def __init__(self, g, f, **kwargs):
411 Operator.__init__(self, g=g, f=f, **kwargs)
413 @property
414 def name(self):
415 return '(%s ○ %s)' % (self.g.name, self.f.name)
418__all__ = [
419 'Grouping',
420 'RegexGrouping',
421 'NetworkGrouping',
422 'StationGrouping',
423 'LocationGrouping',
424 'SensorGrouping',
425 'ChannelGrouping',
426 'Operator',
427 'RestitutionParameters',
428 'Restitution',
429 'Shift',
430 'ToENZ',
431 'ToRTZ',
432 'ToLTQ',
433 'Composition']