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
10import fnmatch
12from ..model import QuantityType, separator
13from .. import error
15from pyrocko.guts import Object, String, Duration, Float, clone
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'[^%s]*' % separator
51 elif c == '?':
52 d = r'[^%s]' % separator
53 elif c == '.':
54 d = separator
55 else:
56 d = c
58 dd.append(d)
59 reg = ''.join(dd)
60 return reg
63_compiled_patterns = {}
66def compiled(pattern):
67 if pattern not in _compiled_patterns:
68 rpattern = re.compile(fnmatch.translate(pattern), re.I)
69 _compiled_patterns[pattern] = rpattern
70 else:
71 rpattern = _compiled_patterns[pattern]
73 return rpattern
76class Filtering(Object):
78 def filter(self, it):
79 return list(it)
82class RegexFiltering(Object):
83 pattern = String.T(default=r'(.*)')
85 def __init__(self, **kwargs):
86 Filtering.__init__(self, **kwargs)
87 self._compiled_pattern = re.compile(self.pattern)
89 def filter(self, it):
90 return [
91 x for x in it if self._compiled_pattern.fullmatch(x)]
94class Grouping(Object):
96 def key(self, codes):
97 return codes
100class RegexGrouping(Grouping):
101 pattern = String.T(default=r'(.*)')
103 def __init__(self, **kwargs):
104 Grouping.__init__(self, **kwargs)
105 self._compiled_pattern = re.compile(self.pattern)
107 def key(self, codes):
108 return self._compiled_pattern.fullmatch(codes).groups()
111class ComponentGrouping(RegexGrouping):
112 pattern = String.T(default=_cglob_translate('(*.*.*.*.*)?(.*)'))
115class Naming(Object):
116 suffix = String.T(default='')
118 def get_name(self, base):
119 return base + self.suffix
122class RegexNaming(Naming):
123 pattern = String.T(default=r'(.*)')
124 replacement = String.T(default=r'\1')
126 def __init__(self, **kwargs):
127 Naming.__init__(self, **kwargs)
128 self._compiled_pattern = re.compile(self.pattern)
130 def get_name(self, base):
131 return self._compiled_pattern.sub(
132 self.replacement, base) + self.suffix
135class ReplaceComponentNaming(RegexNaming):
136 pattern = String.T(default=_cglob_translate('(*.*.*.*.*)?(.*)'))
137 replacement = String.T(default=r'\1{component}\2')
140def lsplit_codes(lcodes):
141 return [tuple(codes.split(separator)) for codes in lcodes]
144class Operator(Object):
146 filtering = Filtering.T(default=Filtering.D())
147 grouping = Grouping.T(default=Grouping.D())
148 naming = Naming.T(default=Naming.D())
150 def __init__(self, **kwargs):
151 Object.__init__(self, **kwargs)
152 self._output_names_cache = {}
153 self._groups = {}
154 self._available = []
156 @property
157 def name(self):
158 return self.__class__.__name__
160 def describe(self):
161 return self.name
163 def iter_mappings(self):
164 for k, group in self._groups.items():
165 if group[1] is None:
166 group[1] = sorted(group[0])
168 yield (
169 lsplit_codes(group[1]),
170 lsplit_codes(group[2]))
172 def update_mappings(self, available, registry):
173 available = list(available)
174 removed, added = odiff(self._available, available)
176 filt = self.filtering.filter
177 gkey = self.grouping.key
178 groups = self._groups
180 need_update = set()
182 def deregister(group):
183 for codes in group[2]:
184 del registry[codes]
186 def register(group):
187 for codes in group[2]:
188 if codes in registry:
189 logger.warning(
190 'duplicate operator output codes: %s' % codes)
191 registry[codes] = (self, group)
193 for codes in filt(removed):
194 k = gkey(codes)
195 groups[k][0].remove(codes)
196 need_update.add(k)
198 for codes in filt(added):
199 k = gkey(codes)
200 if k not in groups:
201 groups[k] = [set(), None, ()]
203 groups[k][0].add(codes)
204 need_update.add(k)
206 for k in need_update:
207 group = groups[k]
208 deregister(group)
209 group[1] = tuple(sorted(group[0]))
210 if not group[1]:
211 del groups[k]
212 else:
213 group[2] = self._out_codes(group[1])
214 register(group)
216 self._available = available
218 def _out_codes(self, in_codes):
219 return [self.naming.get_name(codes) for codes in in_codes]
221 def get_channels(self, squirrel, group, *args, **kwargs):
222 _, in_codes, out_codes = group
223 assert len(in_codes) == 1 and len(out_codes) == 1
224 in_codes_tup = in_codes[0].split(separator)
225 channels = squirrel.get_channels(codes=in_codes_tup, **kwargs)
226 agn, net, sta, loc, cha, ext = out_codes[0].split(separator)
227 channels_out = []
228 for channel in channels:
229 channel_out = clone(channel)
230 channel_out.set_codes(
231 agency=agn,
232 network=net,
233 station=sta,
234 location=loc,
235 channel=cha,
236 extra=ext)
237 channels_out.append(channel_out)
239 return channels_out
241 def get_waveforms(self, squirrel, group, **kwargs):
242 _, in_codes, out_codes = group
243 assert len(in_codes) == 1 and len(out_codes) == 1
245 in_codes_tup = in_codes[0].split(separator)
246 trs = squirrel.get_waveforms(codes=in_codes_tup, **kwargs)
247 agn, net, sta, loc, cha, ext = out_codes[0].split(separator)
248 for tr in trs:
249 tr.set_codes(
250 agency=agn,
251 network=net,
252 station=sta,
253 location=loc,
254 channel=cha,
255 extra=ext)
257 return trs
259 # def update_waveforms(self, squirrel, tmin, tmax, codes):
260 # if codes is None:
261 # for _, in_codes, out_codes in self._groups.values():
262 # for codes in
265class Parameters(Object):
266 pass
269class RestitutionParameters(Parameters):
270 frequency_min = Float.T()
271 frequency_max = Float.T()
272 frequency_taper_factor = Float.T(default=1.5)
273 time_taper_factor = Float.T(default=2.0)
276class Restitution(Operator):
277 naming = Naming(suffix='R{quantity}')
278 quantity = QuantityType.T(default='displacement')
280 @property
281 def name(self):
282 return 'Restitution(%s)' % self.quantity[0]
284 def _out_codes(self, group):
285 return [
286 self.naming.get_name(codes).format(
287 quantity=self.quantity[0])
288 for codes in group]
290 def get_tpad(self, params):
291 return params.time_taper_factor / params.frequency_min
293 def get_waveforms(
294 self, squirrel, codes, params, tmin, tmax, **kwargs):
296 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
297 assert self is self_
298 assert len(in_codes) == 1 and len(out_codes) == 1
299 in_codes_tup = tuple(in_codes[0].split(separator))
301 tpad = self.get_tpad(params)
303 tmin_raw = tmin - tpad
304 tmax_raw = tmax + tpad
306 trs = squirrel.get_waveforms(
307 codes=in_codes_tup, tmin=tmin_raw, tmax=tmax_raw, **kwargs)
309 try:
310 resp = squirrel.get_response(
311 codes=in_codes_tup,
312 tmin=tmin_raw,
313 tmax=tmax_raw).get_effective(self.quantity)
315 except error.NotAvailable:
316 return []
318 freqlimits = (
319 params.frequency_min / params.frequency_taper_factor,
320 params.frequency_min,
321 params.frequency_max,
322 params.frequency_max * params.frequency_taper_factor)
324 agn, net, sta, loc, cha, ext = out_codes[0].split(separator)
325 trs_rest = []
326 for tr in trs:
327 tr_rest = tr.transfer(
328 tfade=tpad,
329 freqlimits=freqlimits,
330 transfer_function=resp,
331 invert=True)
333 tr_rest.set_codes(
334 agency=agn,
335 network=net,
336 station=sta,
337 location=loc,
338 channel=cha,
339 extra=ext)
341 trs_rest.append(tr_rest)
343 return trs_rest
346class Shift(Operator):
347 naming = Naming(suffix='S')
348 delay = Duration.T()
351class Transform(Operator):
352 grouping = Grouping.T(default=ComponentGrouping.D())
353 naming = ReplaceComponentNaming(suffix='T{system}')
355 def _out_codes(self, group):
356 return [
357 self.naming.get_name(group[0]).format(
358 component=c, system=self.components.lower())
359 for c in self.components]
362class ToENZ(Transform):
363 components = 'ENZ'
365 def get_waveforms(
366 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
368 trs = squirrel.get_waveforms(
369 codes=lsplit_codes(in_codes), tmin=tmin, tmax=tmax, **kwargs)
371 for tr in trs:
372 print(tr)
375class ToRTZ(Transform):
376 components = 'RTZ'
377 backazimuth = Float.T()
380class ToLTQ(Transform):
381 components = 'LTQ'
384class Composition(Operator):
385 g = Operator.T()
386 f = Operator.T()
388 def __init__(self, g, f, **kwargs):
389 Operator.__init__(self, g=g, f=f, **kwargs)
391 @property
392 def name(self):
393 return '(%s ○ %s)' % (self.g.name, self.f.name)
396__all__ = [
397 'Operator',
398 'RestitutionParameters',
399 'Restitution',
400 'Shift',
401 'ToENZ',
402 'ToRTZ',
403 'ToLTQ',
404 'Composition']