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 NetworkStationGrouping(RegexGrouping):
102 pattern = String.T(default=_cglob_translate('(*.*).*.*.*'))
105class NetworkStationLocationGrouping(RegexGrouping):
106 pattern = String.T(default=_cglob_translate('(*.*.*).*.*'))
109class Translation(Object):
111 def translate(self, codes):
112 return codes
115class AddSuffixTranslation(Translation):
116 suffix = String.T(default='')
118 def translate(self, codes):
119 return codes.replace(extra=codes.extra + self.suffix)
122class RegexTranslation(AddSuffixTranslation):
123 pattern = String.T(default=r'(.*)')
124 replacement = String.T(default=r'\1')
126 def __init__(self, **kwargs):
127 AddSuffixTranslation.__init__(self, **kwargs)
128 self._compiled_pattern = re.compile(self.pattern)
130 def translate(self, codes):
131 return codes.__class__(
132 self._compiled_pattern.sub(self.replacement, str(codes)))
135class ReplaceComponentTranslation(RegexTranslation):
136 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
137 replacement = String.T(default=r'\1{component}\2')
140class Operator(Object):
142 filtering = Filtering.T(default=Filtering.D())
143 grouping = Grouping.T(default=Grouping.D())
144 translation = Translation.T(default=Translation.D())
146 def __init__(self, **kwargs):
147 Object.__init__(self, **kwargs)
148 self._output_names_cache = {}
149 self._groups = {}
150 self._available = []
152 @property
153 def name(self):
154 return self.__class__.__name__
156 def describe(self):
157 return self.name
159 def iter_mappings(self):
160 for k, group in self._groups.items():
161 if group[1] is None:
162 group[1] = sorted(group[0])
164 yield (group[1], group[2])
166 def update_mappings(self, available, registry):
167 available = list(available)
168 removed, added = odiff(self._available, available)
170 filt = self.filtering.filter
171 gkey = self.grouping.key
172 groups = self._groups
174 need_update = set()
176 def deregister(group):
177 for codes in group[2]:
178 del registry[codes]
180 def register(group):
181 for codes in group[2]:
182 if codes in registry:
183 logger.warning(
184 'duplicate operator output codes: %s' % codes)
185 registry[codes] = (self, group)
187 for codes in filt(removed):
188 k = gkey(codes)
189 groups[k][0].remove(codes)
190 need_update.add(k)
192 for codes in filt(added):
193 k = gkey(codes)
194 if k not in groups:
195 groups[k] = [set(), None, ()]
197 groups[k][0].add(codes)
198 need_update.add(k)
200 for k in need_update:
201 group = groups[k]
202 deregister(group)
203 group[1] = tuple(sorted(group[0]))
204 if not group[1]:
205 del groups[k]
206 else:
207 group[2] = self._out_codes(group[1])
208 register(group)
210 self._available = available
212 def _out_codes(self, in_codes):
213 return [self.translation.translate(codes) for codes in in_codes]
215 def get_channels(self, squirrel, group, *args, **kwargs):
216 _, in_codes, out_codes = group
217 assert len(in_codes) == 1 and len(out_codes) == 1
218 channels = squirrel.get_channels(codes=in_codes[0], **kwargs)
219 channels_out = []
220 for channel in channels:
221 channel_out = clone(channel)
222 channel_out.codes = out_codes[0]
223 channels_out.append(channel_out)
225 return channels_out
227 def get_waveforms(self, squirrel, group, **kwargs):
228 _, in_codes, out_codes = group
229 assert len(in_codes) == 1 and len(out_codes) == 1
231 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs)
232 for tr in trs:
233 tr.set_codes(*out_codes[0])
235 return trs
237 # def update_waveforms(self, squirrel, tmin, tmax, codes):
238 # if codes is None:
239 # for _, in_codes, out_codes in self._groups.values():
240 # for codes in
243class Parameters(Object):
244 pass
247class RestitutionParameters(Parameters):
248 frequency_min = Float.T()
249 frequency_max = Float.T()
250 frequency_taper_factor = Float.T(default=1.5)
251 time_taper_factor = Float.T(default=2.0)
254class Restitution(Operator):
255 translation = AddSuffixTranslation(suffix='R{quantity}')
256 quantity = QuantityType.T(default='displacement')
258 @property
259 def name(self):
260 return 'Restitution(%s)' % self.quantity[0]
262 def _out_codes(self, group):
263 return [
264 codes.__class__(str(self.translation.translate(codes)).format(
265 quantity=self.quantity[0]))
266 for codes in group]
268 def get_tpad(self, params):
269 return params.time_taper_factor / params.frequency_min
271 def get_waveforms(
272 self, squirrel, codes, params, tmin, tmax, **kwargs):
274 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
275 assert self is self_
276 assert len(in_codes) == 1 and len(out_codes) == 1
278 tpad = self.get_tpad(params)
280 tmin_raw = tmin - tpad
281 tmax_raw = tmax + tpad
283 trs = squirrel.get_waveforms(
284 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs)
286 try:
287 resp = squirrel.get_response(
288 codes=in_codes[0],
289 tmin=tmin_raw,
290 tmax=tmax_raw).get_effective(self.quantity)
292 except error.NotAvailable:
293 return []
295 freqlimits = (
296 params.frequency_min / params.frequency_taper_factor,
297 params.frequency_min,
298 params.frequency_max,
299 params.frequency_max * params.frequency_taper_factor)
301 trs_rest = []
302 for tr in trs:
303 tr_rest = tr.transfer(
304 tfade=tpad,
305 freqlimits=freqlimits,
306 transfer_function=resp,
307 invert=True)
309 tr_rest.set_codes(*out_codes[0])
311 trs_rest.append(tr_rest)
313 return trs_rest
316class Shift(Operator):
317 translation = AddSuffixTranslation(suffix='S')
318 delay = Duration.T()
321class Transform(Operator):
322 grouping = Grouping.T(default=ComponentGrouping.D())
323 translation = ReplaceComponentTranslation(suffix='T{system}')
325 def _out_codes(self, group):
326 return [
327 self.translation.translate(group[0]).format(
328 component=c, system=self.components.lower())
329 for c in self.components]
332class ToENZ(Transform):
333 components = 'ENZ'
335 def get_waveforms(
336 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
338 trs = squirrel.get_waveforms(
339 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs)
341 for tr in trs:
342 print(tr)
345class ToRTZ(Transform):
346 components = 'RTZ'
347 backazimuth = Float.T()
350class ToLTQ(Transform):
351 components = 'LTQ'
354class Composition(Operator):
355 g = Operator.T()
356 f = Operator.T()
358 def __init__(self, g, f, **kwargs):
359 Operator.__init__(self, g=g, f=f, **kwargs)
361 @property
362 def name(self):
363 return '(%s ○ %s)' % (self.g.name, self.f.name)
366__all__ = [
367 'Grouping',
368 'RegexGrouping',
369 'ComponentGrouping',
370 'NetworkStationGrouping',
371 'NetworkStationLocationGrouping',
372 'Operator',
373 'RestitutionParameters',
374 'Restitution',
375 'Shift',
376 'ToENZ',
377 'ToRTZ',
378 'ToLTQ',
379 'Composition']