1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function
8import re
9import logging
10import os.path as op
12import numpy as num
14from pyrocko import io, trace, util
15from pyrocko.progress import progress
16from pyrocko.has_paths import Path, HasPaths
17from pyrocko.guts import Dict, String, Choice, Float, List, Timestamp, \
18 StringChoice, IntChoice, Defer, load_all, clone
19from pyrocko.squirrel.dataset import Dataset
20from pyrocko.squirrel.client.local import LocalData
21from pyrocko.squirrel.tool import common
22from pyrocko.squirrel.error import ToolError
24tts = util.time_to_str
26guts_prefix = 'jackseis'
27logger = logging.getLogger('psq.cli.jackseis')
30def make_task(*args):
31 return progress.task(*args, logger=logger)
34class JackseisError(ToolError):
35 pass
38class Chain(object):
39 def __init__(self, node, parent=None):
40 self.node = node
41 self.parent = parent
43 def mcall(self, name, *args, **kwargs):
44 ret = []
45 if self.parent is not None:
46 ret.append(self.parent.mcall(name, *args, **kwargs))
48 ret.append(getattr(self.node, name)(*args, **kwargs))
49 return ret
51 def fcall(self, name, *args, **kwargs):
52 v = getattr(self.node, name)(*args, **kwargs)
53 if v is None and self.parent is not None:
54 return self.parent.fcall(name, *args, **kwargs)
55 else:
56 return v
58 def get(self, name):
59 v = getattr(self.node, name)
60 if v is None and self.parent is not None:
61 return self.parent.get(name)
62 else:
63 return v
65 def dget(self, name, k):
66 v = getattr(self.node, name).get(k, None)
67 if v is None and self.parent is not None:
68 return self.parent.dget(name, k)
69 else:
70 return v
73class OutputFormatChoice(StringChoice):
74 choices = io.allowed_formats('save')
77class OutputDataTypeChoice(StringChoice):
78 choices = ['int32', 'int64', 'float32', 'float64']
79 name_to_dtype = {
80 'int32': num.int32,
81 'int64': num.int64,
82 'float32': num.float32,
83 'float64': num.float64}
86class Converter(HasPaths):
88 in_dataset = Dataset.T(optional=True)
89 in_path = String.T(optional=True)
90 in_paths = List.T(String.T(optional=True))
92 codes = String.T(optional=True)
94 rename = Dict.T(
95 String.T(),
96 Choice.T([
97 String.T(),
98 Dict.T(String.T(), String.T())]))
99 tmin = Timestamp.T(optional=True)
100 tmax = Timestamp.T(optional=True)
101 tinc = Float.T(optional=True)
103 downsample = Float.T(optional=True)
105 out_path = Path.T(optional=True)
106 out_sds_path = Path.T(optional=True)
107 out_format = OutputFormatChoice.T(optional=True)
108 out_data_type = OutputDataTypeChoice.T(optional=True)
109 out_mseed_record_length = IntChoice.T(
110 optional=True,
111 choices=list(io.mseed.VALID_RECORD_LENGTHS))
112 out_mseed_steim = IntChoice.T(
113 optional=True,
114 choices=[1, 2])
116 parts = List.T(Defer('Converter.T'))
118 @classmethod
119 def add_arguments(cls, p):
120 common.add_query_arguments(p, without=['time', 'codes'])
122 p.add_argument(
123 '--downsample',
124 dest='downsample',
125 type=float,
126 metavar='RATE',
127 help='Downsample to RATE [Hz].')
129 p.add_argument(
130 '--out-path',
131 dest='out_path',
132 metavar='TEMPLATE',
133 help='Set output path to TEMPLATE. Available placeholders '
134 'are %%n: network, %%s: station, %%l: location, %%c: '
135 'channel, %%b: begin time, %%e: end time, %%j: julian day of '
136 'year. The following additional placeholders use the window '
137 'begin and end times rather than trace begin and end times '
138 '(to suppress producing many small files for gappy traces), '
139 '%%(wmin_year)s, %%(wmin_month)s, %%(wmin_day)s, %%(wmin)s, '
140 '%%(wmin_jday)s, %%(wmax_year)s, %%(wmax_month)s, '
141 '%%(wmax_day)s, %%(wmax)s, %%(wmax_jday)s. '
142 'Example: --output=\'data/%%s/trace-%%s-%%c.mseed\'')
144 p.add_argument(
145 '--out-sds-path',
146 dest='out_sds_path',
147 metavar='PATH',
148 help='Set output path to create SDS (https://www.seiscomp.de'
149 '/seiscomp3/doc/applications/slarchive/SDS.html), rooted at '
150 'the given path. Implies --tinc=86400. '
151 'Example: --output-sds-path=data/sds')
153 p.add_argument(
154 '--out-format',
155 dest='out_format',
156 choices=io.allowed_formats('save'),
157 metavar='FORMAT',
158 help='Set output file format. Choices: %s' % io.allowed_formats(
159 'save', 'cli_help', 'mseed'))
161 p.add_argument(
162 '--out-data-type',
163 dest='out_data_type',
164 choices=OutputDataTypeChoice.choices,
165 metavar='DTYPE',
166 help='Set numerical data type. Choices: %s. The output file '
167 'format must support the given type. By default, the data '
168 'type is kept unchanged.' % ', '.join(
169 OutputDataTypeChoice.choices))
171 p.add_argument(
172 '--out-mseed-record-length',
173 dest='out_mseed_record_length',
174 type=int,
175 choices=io.mseed.VALID_RECORD_LENGTHS,
176 metavar='INT',
177 help='Set the mseed record length in bytes. Choices: %s. '
178 'Default is 4096 bytes, which is commonly used for archiving.'
179 % ', '.join(str(b) for b in io.mseed.VALID_RECORD_LENGTHS))
181 p.add_argument(
182 '--out-mseed-steim',
183 dest='out_mseed_steim',
184 type=int,
185 choices=(1, 2),
186 metavar='INT',
187 help='Set the mseed STEIM compression. Choices: 1 or 2. '
188 'Default is STEIM-2, which can compress full range int32. '
189 'Note: STEIM-2 is limited to 30 bit dynamic range.')
191 @classmethod
192 def from_arguments(cls, args):
193 kwargs = common.squirrel_query_from_arguments(args)
195 obj = cls(
196 downsample=args.downsample,
197 out_format=args.out_format,
198 out_path=args.out_path,
199 out_sds_path=args.out_sds_path,
200 out_data_type=args.out_data_type,
201 out_mseed_record_length=args.out_mseed_record_length,
202 out_mseed_steim=args.out_mseed_steim,
203 **kwargs)
205 obj.validate()
206 return obj
208 def add_dataset(self, sq):
209 if self.in_dataset is not None:
210 sq.add_dataset(self.in_dataset)
212 if self.in_path is not None:
213 ds = Dataset(sources=[LocalData(paths=[self.in_path])])
214 ds.set_basepath_from(self)
215 sq.add_dataset(ds)
217 if self.in_paths:
218 ds = Dataset(sources=[LocalData(paths=self.in_paths)])
219 ds.set_basepath_from(self)
220 sq.add_dataset(ds)
222 def get_effective_rename_rules(self, chain):
223 d = {}
224 for k in ['network', 'station', 'location', 'channel']:
225 v = chain.dget('rename', k)
226 if isinstance(v, str):
227 m = re.match(r'/([^/]+)/([^/]*)/', v)
228 if m:
229 try:
230 v = (re.compile(m.group(1)), m.group(2))
231 except Exception:
232 raise JackseisError(
233 'Invalid replacement pattern: /%s/' % m.group(1))
235 d[k] = v
237 return d
239 def get_effective_out_path(self):
240 nset = sum(x is not None for x in (
241 self.out_path,
242 self.out_sds_path))
244 if nset > 1:
245 raise JackseisError(
246 'More than one out of [out_path, out_sds_path] set.')
248 is_sds = False
249 if self.out_path:
250 out_path = self.out_path
252 elif self.out_sds_path:
253 out_path = op.join(
254 self.out_sds_path,
255 '%(wmin_year)s/%(network)s/%(station)s/%(channel)s.D'
256 '/%(network)s.%(station)s.%(location)s.%(channel)s.D'
257 '.%(wmin_year)s.%(wmin_jday)s')
258 is_sds = True
259 else:
260 out_path = None
262 if out_path is not None:
263 return self.expand_path(out_path), is_sds
264 else:
265 return None
267 def do_rename(self, rules, tr):
268 rename = {}
269 for k in ['network', 'station', 'location', 'channel']:
270 v = rules.get(k, None)
271 if isinstance(v, str):
272 rename[k] = v
273 elif isinstance(v, dict):
274 try:
275 oldval = getattr(tr, k)
276 rename[k] = v[oldval]
277 except KeyError:
278 raise ToolError(
279 'No mapping defined for %s code "%s".' % (k, oldval))
281 elif isinstance(v, tuple):
282 pat, repl = v
283 oldval = getattr(tr, k)
284 newval, n = pat.subn(repl, oldval)
285 if n:
286 rename[k] = newval
288 tr.set_codes(**rename)
290 def convert(self, args, chain=None):
291 if chain is None:
292 defaults = clone(g_defaults)
293 defaults.set_basepath_from(self)
294 chain = Chain(defaults)
296 chain = Chain(self, chain)
298 if self.parts:
299 task = make_task('Jackseis parts')
300 for part in task(self.parts):
301 part.convert(args, chain)
303 del task
305 else:
306 sq = common.squirrel_from_selection_arguments(args)
308 cli_overrides = Converter.from_arguments(args)
309 cli_overrides.set_basepath('.')
311 chain = Chain(cli_overrides, chain)
313 chain.mcall('add_dataset', sq)
315 tmin = chain.get('tmin')
316 tmax = chain.get('tmax')
317 tinc = chain.get('tinc')
318 codes = chain.get('codes')
319 downsample = chain.get('downsample')
320 out_path, is_sds = chain.fcall('get_effective_out_path') \
321 or (None, False)
323 if is_sds and tinc != 86400.:
324 logger.warning('Setting time window to 1 day to generate SDS.')
325 tinc = 86400.0
327 out_format = chain.get('out_format')
328 out_data_type = chain.get('out_data_type')
330 target_deltat = None
331 if downsample is not None:
332 target_deltat = 1.0 / float(downsample)
334 save_kwargs = {}
335 if out_format == 'mseed':
336 save_kwargs['record_length'] = chain.get(
337 'out_mseed_record_length')
338 save_kwargs['steim'] = chain.get(
339 'out_mseed_steim')
341 tpad = 0.0
342 if target_deltat is not None:
343 tpad += target_deltat * 50.
345 task = None
346 rename_rules = self.get_effective_rename_rules(chain)
347 for batch in sq.chopper_waveforms(
348 tmin=tmin, tmax=tmax, tpad=tpad, tinc=tinc,
349 codes=codes,
350 snap_window=True):
352 if task is None:
353 task = make_task('Jackseis blocks', batch.n)
355 task.update(
356 batch.i,
357 '%s - %s' % (
358 util.time_to_str(batch.tmin),
359 util.time_to_str(batch.tmax)))
361 twmin = batch.tmin
362 twmax = batch.tmax
364 traces = batch.traces
366 if target_deltat is not None:
367 out_traces = []
368 for tr in traces:
369 try:
370 tr.downsample_to(
371 target_deltat, snap=True, demean=False)
373 out_traces.append(tr)
375 except (trace.TraceTooShort, trace.NoData):
376 pass
378 traces = out_traces
380 for tr in traces:
381 self.do_rename(rename_rules, tr)
383 if out_data_type:
384 for tr in traces:
385 tr.ydata = tr.ydata.astype(
386 OutputDataTypeChoice.name_to_dtype[out_data_type])
388 chopped_traces = []
389 for tr in traces:
390 try:
391 otr = tr.chop(twmin, twmax, inplace=False)
392 chopped_traces.append(otr)
393 except trace.NoData:
394 pass
396 traces = chopped_traces
398 if out_path is not None:
399 try:
400 io.save(
401 traces, out_path,
402 format=out_format,
403 overwrite=args.force,
404 additional=dict(
405 wmin_year=tts(twmin, format='%Y'),
406 wmin_month=tts(twmin, format='%m'),
407 wmin_day=tts(twmin, format='%d'),
408 wmin_jday=tts(twmin, format='%j'),
409 wmin=tts(twmin, format='%Y-%m-%d_%H-%M-%S'),
410 wmax_year=tts(twmax, format='%Y'),
411 wmax_month=tts(twmax, format='%m'),
412 wmax_day=tts(twmax, format='%d'),
413 wmax_jday=tts(twmax, format='%j'),
414 wmax=tts(twmax, format='%Y-%m-%d_%H-%M-%S')),
415 **save_kwargs)
417 except io.FileSaveError as e:
418 raise JackseisError(str(e))
420 else:
421 for tr in traces:
422 print(tr.summary)
424 if task:
425 task.done()
428g_defaults = Converter(
429 out_mseed_record_length=4096,
430 out_format='mseed',
431 out_mseed_steim=2)
434def setup(subparsers):
435 p = common.add_parser(
436 subparsers, 'jackseis',
437 help='Convert waveform archive data.')
439 common.add_selection_arguments(p)
441 p.add_argument(
442 '--config',
443 dest='config_path',
444 metavar='NAME',
445 help='File containing `jackseis2.Converter` settings.')
447 p.add_argument(
448 '--force',
449 dest='force',
450 action='store_true',
451 default=False,
452 help='Force overwriting of existing files.')
454 Converter.add_arguments(p)
456 return p
459def call(parser, args):
460 if args.config_path:
461 try:
462 converters = load_all(filename=args.config_path)
463 except Exception as e:
464 raise ToolError(str(e))
466 for converter in converters:
467 if not isinstance(converter, Converter):
468 raise ToolError(
469 'Config file should only contain '
470 '`jackseis.Converter` objects.')
472 converter.set_basepath(op.dirname(args.config_path))
474 else:
475 converter = Converter()
476 converter.set_basepath('.')
477 converters = [converter]
479 with progress.view():
480 task = make_task('Jackseis jobs')
481 for converter in task(converters):
482 converter.convert(args)