1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
7import re
9from ..model import QuantityType, match_codes, CodesNSLCE
10from .. import error
12from pyrocko.guts import Object, String, Duration, Float, clone, List
14guts_prefix = 'squirrel.ops'
16logger = logging.getLogger('psq.ops')
19def odiff(a, b):
20 ia = ib = 0
21 only_a = []
22 only_b = []
23 while ia < len(a) or ib < len(b):
24 # TODO remove when finished with implementation
25 if ia > 0:
26 assert a[ia] > a[ia-1]
27 if ib > 0:
28 assert b[ib] > b[ib-1]
30 if ib == len(b) or (ia < len(a) and a[ia] < b[ib]):
31 only_a.append(a[ia])
32 ia += 1
33 elif ia == len(a) or (ib < len(b) and a[ia] > b[ib]):
34 only_b.append(b[ib])
35 ib += 1
36 elif a[ia] == b[ib]:
37 ia += 1
38 ib += 1
40 return only_a, only_b
43def _cglob_translate(creg):
44 dd = []
45 for c in creg:
46 if c == '*':
47 d = r'[^.]*'
48 elif c == '?':
49 d = r'[^.]'
50 elif c == '.':
51 d = r'\.'
52 else:
53 d = c
55 dd.append(d)
56 reg = ''.join(dd)
57 return reg
60class Filtering(Object):
62 def filter(self, it):
63 return list(it)
66class RegexFiltering(Object):
67 pattern = String.T(default=r'(.*)')
69 def __init__(self, **kwargs):
70 Filtering.__init__(self, **kwargs)
71 self._compiled_pattern = re.compile(self.pattern)
73 def filter(self, it):
74 return [
75 x for x in it if self._compiled_pattern.fullmatch(x)]
78class CodesPatternFiltering(Object):
79 codes = List.T(CodesNSLCE.T(), optional=True)
81 def filter(self, it):
82 if self.codes is None:
83 return list(it)
84 else:
85 return [
86 x for x in it
87 if any(match_codes(sc, x) for sc in self.codes)]
90class Grouping(Object):
92 def key(self, codes):
93 return codes
96class RegexGrouping(Grouping):
97 pattern = String.T(default=r'(.*)')
99 def __init__(self, **kwargs):
100 Grouping.__init__(self, **kwargs)
101 self._compiled_pattern = re.compile(self.pattern)
103 def key(self, codes):
104 return self._compiled_pattern.fullmatch(codes.safe_str).groups()
107class NetworkGrouping(RegexGrouping):
108 '''
109 Group by *network* code.
110 '''
111 pattern = String.T(default=_cglob_translate('(*).*.*.*.*'))
114class StationGrouping(RegexGrouping):
115 '''
116 Group by *network.station* codes.
117 '''
118 pattern = String.T(default=_cglob_translate('(*.*).*.*.*'))
121class LocationGrouping(RegexGrouping):
122 '''
123 Group by *network.station.location* codes.
124 '''
125 pattern = String.T(default=_cglob_translate('(*.*.*).*.*'))
128class ChannelGrouping(RegexGrouping):
129 '''
130 Group by *network.station.location.channel* codes.
132 This effectively groups all processings of a channel, which may differ in
133 the *extra* codes attribute.
134 '''
135 pattern = String.T(default=_cglob_translate('(*.*.*.*).*'))
138class SensorGrouping(RegexGrouping):
139 '''
140 Group by *network.station.location.sensor* and *extra* codes.
142 For *sensor* all but the last character of the channel code (indicating the
143 component) are used. This effectively groups all components of a sensor,
144 or processings of a sensor.
145 '''
146 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
149class Translation(Object):
151 def translate(self, codes):
152 return codes
155class AddSuffixTranslation(Translation):
156 suffix = String.T(default='')
158 def translate(self, codes):
159 return codes.replace(extra=codes.extra + self.suffix)
162class RegexTranslation(AddSuffixTranslation):
163 pattern = String.T(default=r'(.*)')
164 replacement = String.T(default=r'\1')
166 def __init__(self, **kwargs):
167 AddSuffixTranslation.__init__(self, **kwargs)
168 self._compiled_pattern = re.compile(self.pattern)
170 def translate(self, codes):
171 return codes.__class__(
172 self._compiled_pattern.sub(self.replacement, codes.safe_str))
175class ReplaceComponentTranslation(RegexTranslation):
176 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)'))
177 replacement = String.T(default=r'\1{component}\2')
180def deregister(registry, group):
181 for codes in group[2]:
182 del registry[codes]
185def register(registry, operator, group):
186 for codes in group[2]:
187 if codes in registry:
188 logger.warning(
189 'duplicate operator output codes: %s' % codes)
190 registry[codes] = (operator, group)
193class Operator(Object):
195 filtering = Filtering.T(default=Filtering.D())
196 grouping = Grouping.T(default=Grouping.D())
197 translation = Translation.T(default=Translation.D())
199 def __init__(self, **kwargs):
200 Object.__init__(self, **kwargs)
201 self._output_names_cache = {}
202 self._groups = {}
203 self._available = []
205 @property
206 def name(self):
207 return self.__class__.__name__
209 def describe(self):
210 return self.name
212 def iter_mappings(self):
213 for k, group in self._groups.items():
214 yield (group[1], group[2])
216 def iter_in_codes(self):
217 for k, group in self._groups.items():
218 yield group[1]
220 def iter_out_codes(self):
221 for k, group in self._groups.items():
222 yield group[2]
224 def update_mappings(self, available, registry=None):
225 available = list(available)
226 removed, added = odiff(self._available, available)
228 filt = self.filtering.filter
229 gkey = self.grouping.key
230 groups = self._groups
232 need_update = set()
234 for codes in filt(removed):
235 k = gkey(codes)
236 groups[k][0].remove(codes)
237 need_update.add(k)
239 for codes in filt(added):
240 k = gkey(codes)
241 if k not in groups:
242 groups[k] = [set(), None, ()]
244 groups[k][0].add(codes)
245 need_update.add(k)
247 for k in need_update:
248 group = groups[k]
249 if registry is not None:
250 deregister(registry, group)
252 group[1] = tuple(sorted(group[0]))
253 if not group[1]:
254 del groups[k]
255 else:
256 group[2] = self._out_codes(group[1])
257 if registry is not None:
258 register(registry, self, group)
260 self._available = available
262 def _out_codes(self, in_codes):
263 return [self.translation.translate(codes) for codes in in_codes]
265 def get_channels(self, squirrel, group, *args, **kwargs):
266 _, in_codes, out_codes = group
267 assert len(in_codes) == 1 and len(out_codes) == 1
268 channels = squirrel.get_channels(codes=in_codes[0], **kwargs)
269 channels_out = []
270 for channel in channels:
271 channel_out = clone(channel)
272 channel_out.codes = out_codes[0]
273 channels_out.append(channel_out)
275 return channels_out
277 def get_waveforms(self, squirrel, group, **kwargs):
278 _, in_codes, out_codes = group
279 assert len(in_codes) == 1 and len(out_codes) == 1
281 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs)
282 for tr in trs:
283 tr.set_codes(*out_codes[0])
285 return trs
287 # def update_waveforms(self, squirrel, tmin, tmax, codes):
288 # if codes is None:
289 # for _, in_codes, out_codes in self._groups.values():
290 # for codes in
293class Parameters(Object):
294 pass
297class RestitutionParameters(Parameters):
298 frequency_min = Float.T()
299 frequency_max = Float.T()
300 frequency_taper_factor = Float.T(default=1.5)
301 time_taper_factor = Float.T(default=2.0)
304class Restitution(Operator):
305 translation = AddSuffixTranslation(suffix='R{quantity}')
306 quantity = QuantityType.T(default='displacement')
308 @property
309 def name(self):
310 return 'Restitution(%s)' % self.quantity[0]
312 def _out_codes(self, group):
313 return [
314 codes.__class__(self.translation.translate(codes).safe_str.format(
315 quantity=self.quantity[0]))
316 for codes in group]
318 def get_tpad(self, params):
319 return params.time_taper_factor / params.frequency_min
321 def get_waveforms(
322 self, squirrel, codes, params, tmin, tmax, **kwargs):
324 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes)
325 assert self is self_
326 assert len(in_codes) == 1 and len(out_codes) == 1
328 tpad = self.get_tpad(params)
330 tmin_raw = tmin - tpad
331 tmax_raw = tmax + tpad
333 trs = squirrel.get_waveforms(
334 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs)
336 try:
337 resp = squirrel.get_response(
338 codes=in_codes[0],
339 tmin=tmin_raw,
340 tmax=tmax_raw).get_effective(self.quantity)
342 except error.NotAvailable:
343 return []
345 freqlimits = (
346 params.frequency_min / params.frequency_taper_factor,
347 params.frequency_min,
348 params.frequency_max,
349 params.frequency_max * params.frequency_taper_factor)
351 trs_rest = []
352 for tr in trs:
353 tr_rest = tr.transfer(
354 tfade=tpad,
355 freqlimits=freqlimits,
356 transfer_function=resp,
357 invert=True)
359 tr_rest.set_codes(*out_codes[0])
361 trs_rest.append(tr_rest)
363 return trs_rest
366class Shift(Operator):
367 translation = AddSuffixTranslation(suffix='S')
368 delay = Duration.T()
371class Transform(Operator):
372 grouping = Grouping.T(default=SensorGrouping.D())
373 translation = ReplaceComponentTranslation(suffix='T{system}')
375 def _out_codes(self, group):
376 return [
377 self.translation.translate(group[0]).format(
378 component=c, system=self.components.lower())
379 for c in self.components]
382class ToENZ(Transform):
383 components = 'ENZ'
385 def get_waveforms(
386 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs):
388 trs = squirrel.get_waveforms(
389 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs)
391 for tr in trs:
392 print(tr)
395class ToRTZ(Transform):
396 components = 'RTZ'
397 backazimuth = Float.T()
400class ToLTQ(Transform):
401 components = 'LTQ'
404class Composition(Operator):
405 g = Operator.T()
406 f = Operator.T()
408 def __init__(self, g, f, **kwargs):
409 Operator.__init__(self, g=g, f=f, **kwargs)
411 @property
412 def name(self):
413 return '(%s ○ %s)' % (self.g.name, self.f.name)
416__all__ = [
417 'Grouping',
418 'RegexGrouping',
419 'NetworkGrouping',
420 'StationGrouping',
421 'LocationGrouping',
422 'SensorGrouping',
423 'ChannelGrouping',
424 'Operator',
425 'RestitutionParameters',
426 'Restitution',
427 'Shift',
428 'ToENZ',
429 'ToRTZ',
430 'ToLTQ',
431 'Composition']