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
12from .. import error
14from pyrocko.guts import Object, String, Duration, Float, clone
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 Grouping(Object):
82 def key(self, codes):
83 return codes
86class RegexGrouping(Grouping):
87 pattern = String.T(default=r'(.*)')
89 def __init__(self, **kwargs):
90 Grouping.__init__(self, **kwargs)
91 self._compiled_pattern = re.compile(self.pattern)
93 def key(self, codes):
94 return self._compiled_pattern.fullmatch(str(codes)).groups()
97class ComponentGrouping(RegexGrouping):
98 pattern = String.T(default=_cglob_translate('(*.*.*.*.*)?(.*)'))
101class Translation(Object):
103 def translate(self, codes):
104 return codes
107class AddSuffixTranslation(Translation):
108 suffix = String.T(default='')
110 def translate(self, codes):
111 return codes.replace(extra=codes.extra + self.suffix)
114class RegexTranslation(AddSuffixTranslation):
115 pattern = String.T(default=r'(.*)')
116 replacement = String.T(default=r'\1')
118 def __init__(self, **kwargs):
119 AddSuffixTranslation.__init__(self, **kwargs)
120 self._compiled_pattern = re.compile(self.pattern)
122 def translate(self, codes):
123 return codes.__class__(
124 self._compiled_pattern.sub(self.replacement, str(codes)))
127class ReplaceComponentTranslation(RegexTranslation):
128 pattern = String.T(default=_cglob_translate('(*.*.*.*.*)?(.*)'))
129 replacement = String.T(default=r'\1{component}\2')
132class Operator(Object):
134 filtering = Filtering.T(default=Filtering.D())
135 grouping = Grouping.T(default=Grouping.D())
136 translation = Translation.T(default=Translation.D())
138 def __init__(self, **kwargs):
139 Object.__init__(self, **kwargs)
140 self._output_names_cache = {}
141 self._groups = {}
142 self._available = []
144 @property
145 def name(self):
146 return self.__class__.__name__
148 def describe(self):
149 return self.name
151 def iter_mappings(self):
152 for k, group in self._groups.items():
153 if group[1] is None:
154 group[1] = sorted(group[0])
156 yield (group[1], group[2])
158 def update_mappings(self, available, registry):
159 available = list(available)
160 removed, added = odiff(self._available, available)
162 filt = self.filtering.filter
163 gkey = self.grouping.key
164 groups = self._groups
166 need_update = set()
168 def deregister(group):
169 for codes in group[2]:
170 del registry[codes]
172 def register(group):
173 for codes in group[2]:
174 if codes in registry:
175 logger.warning(
176 'duplicate operator output codes: %s' % codes)
177 registry[codes] = (self, group)
179 for codes in filt(removed):
180 k = gkey(codes)
181 groups[k][0].remove(codes)
182 need_update.add(k)
184 for codes in filt(added):
185 k = gkey(codes)
186 if k not in groups:
187 groups[k] = [set(), None, ()]
189 groups[k][0].add(codes)
190 need_update.add(k)
192 for k in need_update:
193 group = groups[k]
194 deregister(group)
195 group[1] = tuple(sorted(group[0]))
196 if not group[1]:
197 del groups[k]
198 else:
199 group[2] = self._out_codes(group[1])
200 register(group)
202 self._available = available
204 def _out_codes(self, in_codes):
205 return [self.translation.translate(codes) for codes in in_codes]
207 def get_channels(self, squirrel, group, *args, **kwargs):
208 _, in_codes, out_codes = group
209 assert len(in_codes) == 1 and len(out_codes) == 1
210 channels = squirrel.get_channels(codes=in_codes[0], **kwargs)
211 channels_out = []
212 for channel in channels:
213 channel_out = clone(channel)
214 channel_out.codes = out_codes[0]
215 channels_out.append(channel_out)
217 return channels_out
219 def get_waveforms(self, squirrel, group, **kwargs):
220 _, in_codes, out_codes = group
221 assert len(in_codes) == 1 and len(out_codes) == 1
223 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs)
224 for tr in trs:
225 tr.set_codes(*out_codes[0])
227 return trs
229 # def update_waveforms(self, squirrel, tmin, tmax, codes):
230 # if codes is None:
231 # for _, in_codes, out_codes in self._groups.values():
232 # for codes in
235class Parameters(Object):
236 pass
239class RestitutionParameters(Parameters):
240 frequency_min = Float.T()
241 frequency_max = Float.T()
242 frequency_taper_factor = Float.T(default=1.5)
243 time_taper_factor = Float.T(default=2.0)
246class Restitution(Operator):
247 translation = AddSuffixTranslation(suffix='R{quantity}')
248 quantity = QuantityType.T(default='displacement')
250 @property
251 def name(self):
252 return 'Restitution(%s)' % self.quantity[0]
254 def _out_codes(self, group):
255 return [
256 codes.__class__(str(self.translation.translate(codes)).format(
257 quantity=self.quantity[0]))
258 for codes in group]
260 def get_tpad(self, params):
261 return params.time_taper_factor / params.frequency_min
263 def get_waveforms(
264 self, squirrel, codes, params, tmin, tmax, **kwargs):
266 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
267 assert self is self_
268 assert len(in_codes) == 1 and len(out_codes) == 1
270 tpad = self.get_tpad(params)
272 tmin_raw = tmin - tpad
273 tmax_raw = tmax + tpad
275 trs = squirrel.get_waveforms(
276 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs)
278 try:
279 resp = squirrel.get_response(
280 codes=in_codes[0],
281 tmin=tmin_raw,
282 tmax=tmax_raw).get_effective(self.quantity)
284 except error.NotAvailable:
285 return []
287 freqlimits = (
288 params.frequency_min / params.frequency_taper_factor,
289 params.frequency_min,
290 params.frequency_max,
291 params.frequency_max * params.frequency_taper_factor)
293 trs_rest = []
294 for tr in trs:
295 tr_rest = tr.transfer(
296 tfade=tpad,
297 freqlimits=freqlimits,
298 transfer_function=resp,
299 invert=True)
301 tr_rest.set_codes(*out_codes[0])
303 trs_rest.append(tr_rest)
305 return trs_rest
308class Shift(Operator):
309 translation = AddSuffixTranslation(suffix='S')
310 delay = Duration.T()
313class Transform(Operator):
314 grouping = Grouping.T(default=ComponentGrouping.D())
315 translation = ReplaceComponentTranslation(suffix='T{system}')
317 def _out_codes(self, group):
318 return [
319 self.translation.translate(group[0]).format(
320 component=c, system=self.components.lower())
321 for c in self.components]
324class ToENZ(Transform):
325 components = 'ENZ'
327 def get_waveforms(
328 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
330 trs = squirrel.get_waveforms(
331 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs)
333 for tr in trs:
334 print(tr)
337class ToRTZ(Transform):
338 components = 'RTZ'
339 backazimuth = Float.T()
342class ToLTQ(Transform):
343 components = 'LTQ'
346class Composition(Operator):
347 g = Operator.T()
348 f = Operator.T()
350 def __init__(self, g, f, **kwargs):
351 Operator.__init__(self, g=g, f=f, **kwargs)
353 @property
354 def name(self):
355 return '(%s ○ %s)' % (self.g.name, self.f.name)
358__all__ = [
359 'Operator',
360 'RestitutionParameters',
361 'Restitution',
362 'Shift',
363 'ToENZ',
364 'ToRTZ',
365 'ToLTQ',
366 'Composition']