1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
5'''
6Utility functions for Pyrocko.
8.. _time-handling-mode:
10High precision time handling mode
11.................................
13Pyrocko can treat timestamps either as standard double precision (64 bit)
14floating point values, or as high precision float (:py:class:`numpy.float128`
15or :py:class:`numpy.float96`, whichever is available), aliased here as
16:py:class:`pyrocko.util.hpfloat`. High precision time stamps are required when
17handling data with sub-millisecond precision, i.e. kHz/MHz data streams and
18event catalogs derived from such data.
20Not all functions in Pyrocko and in programs depending on Pyrocko may work
21correctly with high precision times. Therefore, Pyrocko's high precision time
22handling mode has to be actively activated by user config, command line option
23or enforced within a certain script/program.
25The default high precision time handling mode can be configured globally with
26the user config variable
27:py:gattr:`~pyrocko.config.Config.use_high_precision_time`. Calling the
28function :py:func:`use_high_precision_time` overrides the default from the
29config file. This function may be called at startup of a program/script
30requiring a specific time handling mode.
32To create a valid time stamp for use in Pyrocko (e.g. in
33:py:class:`~pyrocko.model.Event` or :py:class:`~pyrocko.trace.Trace` objects),
34use:
36.. code-block :: python
38 import time
39 from pyrocko import util
41 # By default using mode selected in user config, override with:
42 # util.use_high_precision_time(True) # force high precision mode
43 # util.use_high_precision_time(False) # force low precision mode
45 t1 = util.str_to_time('2020-08-27 10:22:00')
46 t2 = util.str_to_time('2020-08-27 10:22:00.111222')
47 t3 = util.to_time_float(time.time())
49 # To get the appropriate float class, use:
51 time_float = util.get_time_float()
52 # -> float, numpy.float128 or numpy.float96
53 [isinstance(t, time_float) for t in [t1, t2, t3]]
54 # -> [True, True, True]
56 # Shortcut:
57 util.check_time_class(t1)
59Module content
60..............
61'''
63import time
64import logging
65import os
66import sys
67import re
68import calendar
69import math
70import fnmatch
71import inspect
72import weakref
73try:
74 import fcntl
75except ImportError:
76 fcntl = None
77import optparse
78import os.path as op
79import errno
81import numpy as num
82from scipy import signal
83import pyrocko
84from pyrocko import dummy_progressbar
87from urllib.parse import urlencode, quote, unquote # noqa
88from urllib.request import Request, build_opener, HTTPDigestAuthHandler # noqa
89from urllib.request import urlopen as _urlopen # noqa
90from urllib.error import HTTPError, URLError # noqa
93try:
94 import certifi
95 import ssl
96 g_ssl_context = ssl.create_default_context(cafile=certifi.where())
97except ImportError:
98 g_ssl_context = None
101class URLErrorSSL(URLError):
103 def __init__(self, urlerror):
104 self.__dict__ = urlerror.__dict__.copy()
106 def __str__(self):
107 return (
108 'Requesting web resource failed and the problem could be '
109 'related to SSL. Python standard libraries on some older '
110 'systems (like Ubuntu 14.04) are known to have trouble '
111 'with some SSL setups of today\'s servers: %s'
112 % URLError.__str__(self))
115def urlopen_ssl_check(*args, **kwargs):
116 try:
117 return urlopen(*args, **kwargs)
118 except URLError as e:
119 if str(e).find('SSL') != -1:
120 raise URLErrorSSL(e)
121 else:
122 raise
125def urlopen(*args, **kwargs):
127 if 'context' not in kwargs and g_ssl_context is not None:
128 kwargs['context'] = g_ssl_context
130 return _urlopen(*args, **kwargs)
133try:
134 long
135except NameError:
136 long = int
139force_dummy_progressbar = False
142try:
143 from pyrocko import util_ext
144except ImportError:
145 util_ext = None
148logger = logging.getLogger('pyrocko.util')
151try:
152 num_full = num.full
153except AttributeError:
154 def num_full(shape, fill_value, dtype=None, order='C'):
155 a = num.empty(shape, dtype=dtype, order=order)
156 a.fill(fill_value)
157 return a
159try:
160 num_full_like = num.full_like
161except AttributeError:
162 def num_full_like(arr, fill_value, dtype=None, order='K', subok=True):
163 a = num.empty_like(arr, dtype=dtype, order=order, subok=subok)
164 a.fill(fill_value)
165 return a
168g_setup_logging_args = 'pyrocko', 'warning'
171def setup_logging(programname='pyrocko', levelname='warning'):
172 '''
173 Initialize logging.
175 :param programname: program name to be written in log
176 :param levelname: string indicating the logging level ('debug', 'info',
177 'warning', 'error', 'critical')
179 This function is called at startup by most pyrocko programs to set up a
180 consistent logging format. This is simply a shortcut to a call to
181 :py:func:`logging.basicConfig()`.
182 '''
184 global g_setup_logging_args
185 g_setup_logging_args = (programname, levelname)
187 levels = {'debug': logging.DEBUG,
188 'info': logging.INFO,
189 'warning': logging.WARNING,
190 'error': logging.ERROR,
191 'critical': logging.CRITICAL}
193 logging.basicConfig(
194 level=levels[levelname],
195 format=programname+':%(name)-25s - %(levelname)-8s - %(message)s')
198def subprocess_setup_logging_args():
199 '''
200 Get arguments from previous call to setup_logging.
202 These can be sent down to a worker process so it can setup its logging
203 in the same way as the main process.
204 '''
205 return g_setup_logging_args
208def data_file(fn):
209 return os.path.join(os.path.split(__file__)[0], 'data', fn)
212class DownloadError(Exception):
213 pass
216class PathExists(DownloadError):
217 pass
220def _download(url, fpath, username=None, password=None,
221 force=False, method='download', stats=None,
222 status_callback=None, entries_wanted=None,
223 recursive=False, header=None):
225 import requests
226 from requests.auth import HTTPBasicAuth
227 from requests.exceptions import HTTPError as req_HTTPError
229 requests.adapters.DEFAULT_RETRIES = 5
230 urljoin = requests.compat.urljoin
232 session = requests.Session()
233 session.header = header
234 session.auth = None if username is None\
235 else HTTPBasicAuth(username, password)
237 status = {
238 'ntotal_files': 0,
239 'nread_files': 0,
240 'ntotal_bytes_all_files': 0,
241 'nread_bytes_all_files': 0,
242 'ntotal_bytes_current_file': 0,
243 'nread_bytes_current_file': 0,
244 'finished': False
245 }
247 try:
248 url_to_size = {}
250 if callable(status_callback):
251 status_callback(status)
253 if not recursive and url.endswith('/'):
254 raise DownloadError(
255 'URL: %s appears to be a directory'
256 ' but recurvise download is False' % url)
258 if recursive and not url.endswith('/'):
259 url += '/'
261 def parse_directory_tree(url, subdir=''):
262 r = session.get(urljoin(url, subdir))
263 r.raise_for_status()
265 entries = re.findall(r'href="([a-zA-Z0-9_.-]+/?)"', r.text)
267 files = sorted(set(subdir + fn for fn in entries
268 if not fn.endswith('/')))
270 if entries_wanted is not None:
271 files = [fn for fn in files
272 if (fn in wanted for wanted in entries_wanted)]
274 status['ntotal_files'] += len(files)
276 dirs = sorted(set(subdir + dn for dn in entries
277 if dn.endswith('/')
278 and dn not in ('./', '../')))
280 for dn in dirs:
281 files.extend(parse_directory_tree(
282 url, subdir=dn))
284 return files
286 def get_content_length(url):
287 if url not in url_to_size:
288 r = session.head(url, headers={'Accept-Encoding': ''})
290 content_length = r.headers.get('content-length', None)
291 if content_length is None:
292 logger.debug('Could not get HTTP header '
293 'Content-Length for %s' % url)
295 content_length = None
297 else:
298 content_length = int(content_length)
299 status['ntotal_bytes_all_files'] += content_length
300 if callable(status_callback):
301 status_callback(status)
303 url_to_size[url] = content_length
305 return url_to_size[url]
307 def download_file(url, fn):
308 logger.info('starting download of %s...' % url)
309 ensuredirs(fn)
311 fsize = get_content_length(url)
312 status['ntotal_bytes_current_file'] = fsize
313 status['nread_bytes_current_file'] = 0
314 if callable(status_callback):
315 status_callback(status)
317 r = session.get(url, stream=True, timeout=5)
318 r.raise_for_status()
320 frx = 0
321 fn_tmp = fn + '.%i.temp' % os.getpid()
322 with open(fn_tmp, 'wb') as f:
323 for d in r.iter_content(chunk_size=1024):
324 f.write(d)
325 frx += len(d)
327 status['nread_bytes_all_files'] += len(d)
328 status['nread_bytes_current_file'] += len(d)
329 if callable(status_callback):
330 status_callback(status)
332 os.rename(fn_tmp, fn)
334 if fsize is not None and frx != fsize:
335 logger.warning(
336 'HTTP header Content-Length: %i bytes does not match '
337 'download size %i bytes' % (fsize, frx))
339 logger.info('finished download of %s' % url)
341 status['nread_files'] += 1
342 if callable(status_callback):
343 status_callback(status)
345 if recursive:
346 if op.exists(fpath) and not force:
347 raise PathExists('path %s already exists' % fpath)
349 files = parse_directory_tree(url)
351 dsize = 0
352 for fn in files:
353 file_url = urljoin(url, fn)
354 dsize += get_content_length(file_url) or 0
356 if method == 'calcsize':
357 return dsize
359 else:
360 for fn in files:
361 file_url = urljoin(url, fn)
362 download_file(file_url, op.join(fpath, fn))
364 else:
365 status['ntotal_files'] += 1
366 if callable(status_callback):
367 status_callback(status)
369 fsize = get_content_length(url)
370 if method == 'calcsize':
371 return fsize
372 else:
373 download_file(url, fpath)
375 except req_HTTPError as e:
376 logging.warning("http error: %s" % e)
377 raise DownloadError('could not download file(s) from: %s' % url)
379 finally:
380 status['finished'] = True
381 if callable(status_callback):
382 status_callback(status)
383 session.close()
386def download_file(
387 url, fpath, username=None, password=None, status_callback=None,
388 **kwargs):
389 return _download(
390 url, fpath, username, password,
391 recursive=False,
392 status_callback=status_callback,
393 **kwargs)
396def download_dir(
397 url, fpath, username=None, password=None, status_callback=None,
398 **kwargs):
400 return _download(
401 url, fpath, username, password,
402 recursive=True,
403 status_callback=status_callback,
404 **kwargs)
407class HPFloatUnavailable(Exception):
408 pass
411class dummy_hpfloat(object):
412 def __init__(self, *args, **kwargs):
413 raise HPFloatUnavailable(
414 'NumPy lacks support for float128 or float96 data type on this '
415 'platform.')
418if hasattr(num, 'float128'):
419 hpfloat = num.float128
421elif hasattr(num, 'float96'):
422 hpfloat = num.float96
424else:
425 hpfloat = dummy_hpfloat
428g_time_float = None
429g_time_dtype = None
432class TimeFloatSettingError(Exception):
433 pass
436def use_high_precision_time(enabled):
437 '''
438 Globally force a specific time handling mode.
440 See :ref:`High precision time handling mode <time-handling-mode>`.
442 :param enabled: enable/disable use of high precision time type
443 :type enabled: bool
445 This function should be called before handling/reading any time data.
446 It can only be called once.
448 Special attention is required when using multiprocessing on a platform
449 which does not use fork under the hood. In such cases, the desired setting
450 must be set also in the subprocess.
451 '''
452 _setup_high_precision_time_mode(enabled_app=enabled)
455def _setup_high_precision_time_mode(enabled_app=False):
456 global g_time_float
457 global g_time_dtype
459 if not (g_time_float is None and g_time_dtype is None):
460 raise TimeFloatSettingError(
461 'Cannot set time handling mode: too late, it has already been '
462 'fixed by an earlier call.')
464 from pyrocko import config
466 conf = config.config()
467 enabled_config = conf.use_high_precision_time
469 enabled_env = os.environ.get('PYROCKO_USE_HIGH_PRECISION_TIME', None)
470 if enabled_env is not None:
471 try:
472 enabled_env = int(enabled_env) == 1
473 except ValueError:
474 raise TimeFloatSettingError(
475 'Environment variable PYROCKO_USE_HIGH_PRECISION_TIME '
476 'should be set to 0 or 1.')
478 enabled = enabled_config
479 mode_from = 'config variable `use_high_precision_time`'
480 notify = enabled
482 if enabled_env is not None:
483 if enabled_env != enabled:
484 notify = True
485 enabled = enabled_env
486 mode_from = 'environment variable `PYROCKO_USE_HIGH_PRECISION_TIME`'
488 if enabled_app is not None:
489 if enabled_app != enabled:
490 notify = True
491 enabled = enabled_app
492 mode_from = 'application override'
494 logger.debug('''
495Pyrocko high precision time mode selection (latter override earlier):
496 config: %s
497 env: %s
498 app: %s
499 -> enabled: %s'''.lstrip() % (
500 enabled_config, enabled_env, enabled_app, enabled))
502 if notify:
503 logger.info('Pyrocko high precision time mode %s by %s.' % (
504 'activated' if enabled else 'deactivated',
505 mode_from))
507 if enabled:
508 g_time_float = hpfloat
509 g_time_dtype = hpfloat
510 else:
511 g_time_float = float
512 g_time_dtype = num.float64
515def get_time_float():
516 '''
517 Get the effective float class for timestamps.
519 See :ref:`High precision time handling mode <time-handling-mode>`.
521 :returns: :py:class:`float` or :py:class:`hpfloat`, depending on the
522 current time handling mode
523 '''
524 global g_time_float
526 if g_time_float is None:
527 _setup_high_precision_time_mode()
529 return g_time_float
532def get_time_dtype():
533 '''
534 Get effective NumPy float class to handle timestamps.
536 See :ref:`High precision time handling mode <time-handling-mode>`.
537 '''
539 global g_time_dtype
541 if g_time_dtype is None:
542 _setup_high_precision_time_mode()
544 return g_time_dtype
547def to_time_float(t):
548 '''
549 Convert float to valid time stamp in the current time handling mode.
551 See :ref:`High precision time handling mode <time-handling-mode>`.
552 '''
553 return get_time_float()(t)
556class TimestampTypeError(ValueError):
557 pass
560def check_time_class(t, error='raise'):
561 '''
562 Type-check variable against current time handling mode.
564 See :ref:`High precision time handling mode <time-handling-mode>`.
565 '''
567 if t == 0.0:
568 return
570 if not isinstance(t, get_time_float()):
571 message = (
572 'Timestamp %g is of type %s but should be of type %s with '
573 'Pyrocko\'s currently selected time handling mode.\n\n'
574 'See https://pyrocko.org/docs/current/library/reference/util.html'
575 '#high-precision-time-handling-mode' % (
576 t, type(t), get_time_float()))
578 if error == 'raise':
579 raise TimestampTypeError(message)
580 elif error == 'warn':
581 logger.warning(message)
582 else:
583 assert False
586class Stopwatch(object):
587 '''
588 Simple stopwatch to measure elapsed wall clock time.
590 Usage::
592 s = Stopwatch()
593 time.sleep(1)
594 print s()
595 time.sleep(1)
596 print s()
597 '''
599 def __init__(self):
600 self.start = time.time()
602 def __call__(self):
603 return time.time() - self.start
606def wrap(text, line_length=80):
607 '''
608 Paragraph and list-aware wrapping of text.
609 '''
611 text = text.strip('\n')
612 at_lineend = re.compile(r' *\n')
613 at_para = re.compile(r'((^|(\n\s*)?\n)(\s+[*] )|\n\s*\n)')
615 paragraphs = at_para.split(text)[::5]
616 listindents = at_para.split(text)[4::5]
617 newlist = at_para.split(text)[3::5]
619 listindents[0:0] = [None]
620 listindents.append(True)
621 newlist.append(None)
623 det_indent = re.compile(r'^ *')
625 outlines = []
626 for ip, p in enumerate(paragraphs):
627 if not p:
628 continue
630 if listindents[ip] is None:
631 _indent = det_indent.findall(p)[0]
632 findent = _indent
633 else:
634 findent = listindents[ip]
635 _indent = ' ' * len(findent)
637 ll = line_length - len(_indent)
638 llf = ll
640 oldlines = [s.strip() for s in at_lineend.split(p.rstrip())]
641 p1 = ' '.join(oldlines)
642 possible = re.compile(r'(^.{1,%i}|.{1,%i})( |$)' % (llf, ll))
643 for imatch, match in enumerate(possible.finditer(p1)):
644 parout = match.group(1)
645 if imatch == 0:
646 outlines.append(findent + parout)
647 else:
648 outlines.append(_indent + parout)
650 if ip != len(paragraphs)-1 and (
651 listindents[ip] is None
652 or newlist[ip] is not None
653 or listindents[ip+1] is None):
655 outlines.append('')
657 return outlines
660def ewrap(lines, width=80, indent=''):
661 lines = list(lines)
662 if not lines:
663 return ''
664 fwidth = max(len(s) for s in lines)
665 nx = max(1, (80-len(indent)) // (fwidth+1))
666 i = 0
667 rows = []
668 while i < len(lines):
669 rows.append(indent + ' '.join(x.ljust(fwidth) for x in lines[i:i+nx]))
670 i += nx
672 return '\n'.join(rows)
675class BetterHelpFormatter(optparse.IndentedHelpFormatter):
677 def __init__(self, *args, **kwargs):
678 optparse.IndentedHelpFormatter.__init__(self, *args, **kwargs)
680 def format_option(self, option):
681 '''
682 From IndentedHelpFormatter but using a different wrap method.
683 '''
685 help_text_position = 4 + self.current_indent
686 help_text_width = self.width - help_text_position
688 opts = self.option_strings[option]
689 opts = "%*s%s" % (self.current_indent, "", opts)
690 if option.help:
691 help_text = self.expand_default(option)
693 if self.help_position + len(help_text) + 1 <= self.width:
694 lines = [
695 '',
696 '%-*s %s' % (self.help_position, opts, help_text),
697 '']
698 else:
699 lines = ['']
700 lines.append(opts)
701 lines.append('')
702 if option.help:
703 help_lines = wrap(help_text, help_text_width)
704 lines.extend(["%*s%s" % (help_text_position, "", line)
705 for line in help_lines])
706 lines.append('')
708 return "\n".join(lines)
710 def format_description(self, description):
711 if not description:
712 return ''
714 if self.current_indent == 0:
715 lines = []
716 else:
717 lines = ['']
719 lines.extend(wrap(description, self.width - self.current_indent))
720 if self.current_indent == 0:
721 lines.append('\n')
723 return '\n'.join(
724 ['%*s%s' % (self.current_indent, '', line) for line in lines])
727class ProgressBar:
728 def __init__(self, label, n):
729 from pyrocko.progress import progress
730 self._context = progress.view()
731 self._context.__enter__()
732 self._task = progress.task(label, n)
734 def update(self, i):
735 self._task.update(i)
737 def finish(self):
738 self._task.done()
739 if self._context:
740 self._context.__exit__()
741 self._context = None
744def progressbar(label, maxval):
745 if force_dummy_progressbar:
746 return dummy_progressbar.ProgressBar(maxval=maxval).start()
748 return ProgressBar(label, maxval)
751def progress_beg(label):
752 '''
753 Notify user that an operation has started.
755 :param label: name of the operation
757 To be used in conjuction with :py:func:`progress_end`.
758 '''
760 sys.stderr.write(label)
761 sys.stderr.flush()
764def progress_end(label=''):
765 '''
766 Notify user that an operation has ended.
768 :param label: name of the operation
770 To be used in conjuction with :py:func:`progress_beg`.
771 '''
773 sys.stderr.write(' done. %s\n' % label)
774 sys.stderr.flush()
777class ArangeError(Exception):
778 pass
781def arange2(start, stop, step, dtype=float, epsilon=1e-6, error='raise'):
782 '''
783 Return evenly spaced numbers over a specified interval.
785 Like :py:func:`numpy.arange` but returning floating point numbers by
786 default and with defined behaviour when stepsize is inconsistent with
787 interval bounds. It is considered inconsistent if the difference between
788 the closest multiple of ``step`` and ``stop`` is larger than ``epsilon *
789 step``. Inconsistencies are handled according to the ``error`` parameter.
790 If it is set to ``'raise'`` an exception of type :py:exc:`ArangeError` is
791 raised. If it is set to ``'round'``, ``'floor'``, or ``'ceil'``, ``stop``
792 is silently changed to the closest, the next smaller, or next larger
793 multiple of ``step``, respectively.
794 '''
796 assert error in ('raise', 'round', 'floor', 'ceil')
798 start = dtype(start)
799 stop = dtype(stop)
800 step = dtype(step)
802 rnd = {'floor': math.floor, 'ceil': math.ceil}.get(error, round)
804 n = int(rnd((stop - start) / step)) + 1
805 stop_check = start + (n-1) * step
807 if error == 'raise' and abs(stop_check - stop) > step * epsilon:
808 raise ArangeError(
809 'inconsistent range specification: start=%g, stop=%g, step=%g'
810 % (start, stop, step))
812 x = num.arange(n, dtype=dtype)
813 x *= step
814 x += start
815 return x
818def polylinefit(x, y, n_or_xnodes):
819 '''
820 Fit piece-wise linear function to data.
822 :param x,y: arrays with coordinates of data
823 :param n_or_xnodes: int, number of segments or x coordinates of polyline
825 :returns: `(xnodes, ynodes, rms_error)` arrays with coordinates of
826 polyline, root-mean-square error
827 '''
829 x = num.asarray(x)
830 y = num.asarray(y)
832 if isinstance(n_or_xnodes, int):
833 n = n_or_xnodes
834 xmin = x.min()
835 xmax = x.max()
836 xnodes = num.linspace(xmin, xmax, n+1)
837 else:
838 xnodes = num.asarray(n_or_xnodes)
839 n = xnodes.size - 1
841 assert len(x) == len(y)
842 assert n > 0
844 ndata = len(x)
845 a = num.zeros((ndata+(n-1), n*2))
846 for i in range(n):
847 xmin_block = xnodes[i]
848 xmax_block = xnodes[i+1]
849 if i == n-1: # don't loose last point
850 indices = num.where(
851 num.logical_and(xmin_block <= x, x <= xmax_block))[0]
852 else:
853 indices = num.where(
854 num.logical_and(xmin_block <= x, x < xmax_block))[0]
856 a[indices, i*2] = x[indices]
857 a[indices, i*2+1] = 1.0
859 w = float(ndata)*100.
860 if i < n-1:
861 a[ndata+i, i*2] = xmax_block*w
862 a[ndata+i, i*2+1] = 1.0*w
863 a[ndata+i, i*2+2] = -xmax_block*w
864 a[ndata+i, i*2+3] = -1.0*w
866 d = num.concatenate((y, num.zeros(n-1)))
867 model = num.linalg.lstsq(a, d, rcond=-1)[0].reshape((n, 2))
869 ynodes = num.zeros(n+1)
870 ynodes[:n] = model[:, 0]*xnodes[:n] + model[:, 1]
871 ynodes[1:] += model[:, 0]*xnodes[1:] + model[:, 1]
872 ynodes[1:n] *= 0.5
874 rms_error = num.sqrt(num.mean((num.interp(x, xnodes, ynodes) - y)**2))
876 return xnodes, ynodes, rms_error
879def plf_integrate_piecewise(x_edges, x, y):
880 '''
881 Calculate definite integral of piece-wise linear function on intervals.
883 Use trapezoidal rule to calculate definite integral of a piece-wise linear
884 function for a series of consecutive intervals. ``x_edges`` and ``x`` must
885 be sorted.
887 :param x_edges: array with edges of the intervals
888 :param x,y: arrays with coordinates of piece-wise linear function's
889 control points
890 '''
892 x_all = num.concatenate((x, x_edges))
893 ii = num.argsort(x_all)
894 y_edges = num.interp(x_edges, x, y)
895 y_all = num.concatenate((y, y_edges))
896 xs = x_all[ii]
897 ys = y_all[ii]
898 y_all[ii[1:]] = num.cumsum((xs[1:] - xs[:-1]) * 0.5 * (ys[1:] + ys[:-1]))
899 return num.diff(y_all[-len(y_edges):])
902def diff_fd_1d_4o(dt, data):
903 '''
904 Approximate first derivative of an array (forth order, central FD).
906 :param dt: sampling interval
907 :param data: NumPy array with data samples
909 :returns: NumPy array with same shape as input
911 Interior points are approximated to fourth order, edge points to first
912 order right- or left-sided respectively, points next to edge to second
913 order central.
914 '''
915 import scipy.signal
917 ddata = num.empty_like(data, dtype=float)
919 if data.size >= 5:
920 ddata[2:-2] = scipy.signal.lfilter(
921 [-1., +8., 0., -8., 1.], [1.], data)[4:] / (12.*dt)
923 if data.size >= 3:
924 ddata[1] = (data[2] - data[0]) / (2. * dt)
925 ddata[-2] = (data[-1] - data[-3]) / (2. * dt)
927 if data.size >= 2:
928 ddata[0] = (data[1] - data[0]) / dt
929 ddata[-1] = (data[-1] - data[-2]) / dt
931 if data.size == 1:
932 ddata[0] = 0.0
934 return ddata
937def diff_fd_1d_2o(dt, data):
938 '''
939 Approximate first derivative of an array (second order, central FD).
941 :param dt: sampling interval
942 :param data: NumPy array with data samples
944 :returns: NumPy array with same shape as input
946 Interior points are approximated to second order, edge points to first
947 order right- or left-sided respectively.
949 Uses :py:func:`numpy.gradient`.
950 '''
952 return num.gradient(data, dt)
955def diff_fd_2d_4o(dt, data):
956 '''
957 Approximate second derivative of an array (forth order, central FD).
959 :param dt: sampling interval
960 :param data: NumPy array with data samples
962 :returns: NumPy array with same shape as input
964 Interior points are approximated to fourth order, next-to-edge points to
965 second order, edge points repeated.
966 '''
967 import scipy.signal
969 ddata = num.empty_like(data, dtype=float)
971 if data.size >= 5:
972 ddata[2:-2] = scipy.signal.lfilter(
973 [-1., +16., -30., +16., -1.], [1.], data)[4:] / (12.*dt**2)
975 if data.size >= 3:
976 ddata[:2] = (data[2] - 2.0 * data[1] + data[0]) / dt**2
977 ddata[-2:] = (data[-1] - 2.0 * data[-2] + data[-3]) / dt**2
979 if data.size < 3:
980 ddata[:] = 0.0
982 return ddata
985def diff_fd_2d_2o(dt, data):
986 '''
987 Approximate second derivative of an array (second order, central FD).
989 :param dt: sampling interval
990 :param data: NumPy array with data samples
992 :returns: NumPy array with same shape as input
994 Interior points are approximated to second order, edge points repeated.
995 '''
996 import scipy.signal
998 ddata = num.empty_like(data, dtype=float)
1000 if data.size >= 3:
1001 ddata[1:-1] = scipy.signal.lfilter(
1002 [1., -2., 1.], [1.], data)[2:] / (dt**2)
1004 ddata[0] = ddata[1]
1005 ddata[-1] = ddata[-2]
1007 if data.size < 3:
1008 ddata[:] = 0.0
1010 return ddata
1013def diff_fd(n, order, dt, data):
1014 '''
1015 Approximate 1st or 2nd derivative of an array.
1017 :param n: 1 for first derivative, 2 for second
1018 :param order: order of the approximation 2 and 4 are supported
1019 :param dt: sampling interval
1020 :param data: NumPy array with data samples
1022 :returns: NumPy array with same shape as input
1024 This is a frontend to the functions :py:func:`diff_fd_1d_2o`,
1025 :py:func:`diff_fd_1d_4o`, :py:func:`diff_fd_2d_2o`, and
1026 :py:func:`diff_fd_2d_4o`.
1028 Raises :py:exc:`ValueError` for unsupported `n` or `order`.
1029 '''
1031 funcs = {
1032 1: {2: diff_fd_1d_2o, 4: diff_fd_1d_4o},
1033 2: {2: diff_fd_2d_2o, 4: diff_fd_2d_4o}}
1035 try:
1036 funcs_n = funcs[n]
1037 except KeyError:
1038 raise ValueError(
1039 'pyrocko.util.diff_fd: '
1040 'Only 1st and 2sd derivatives are supported.')
1042 try:
1043 func = funcs_n[order]
1044 except KeyError:
1045 raise ValueError(
1046 'pyrocko.util.diff_fd: '
1047 'Order %i is not supported for %s derivative. Supported: %s' % (
1048 order, ['', '1st', '2nd'][n],
1049 ', '.join('%i' % order for order in sorted(funcs_n.keys()))))
1051 return func(dt, data)
1054class GlobalVars(object):
1055 reuse_store = dict()
1056 decitab_nmax = 0
1057 decitab = {}
1058 decimate_fir_coeffs = {}
1059 decimate_fir_remez_coeffs = {}
1060 decimate_iir_coeffs = {}
1061 re_frac = None
1064def decimate_coeffs(q, n=None, ftype='iir'):
1066 q = int(q)
1068 if n is None:
1069 if ftype == 'fir':
1070 n = 30
1071 elif ftype == 'fir-remez':
1072 n = 45*q
1073 else:
1074 n = 8
1076 if ftype == 'fir':
1077 coeffs = GlobalVars.decimate_fir_coeffs
1078 if (n, 1./q) not in coeffs:
1079 coeffs[n, 1./q] = signal.firwin(n+1, .75/q, window='hamming')
1081 b = coeffs[n, 1./q]
1082 return b, [1.], n
1084 elif ftype == 'fir-remez':
1085 coeffs = GlobalVars.decimate_fir_remez_coeffs
1086 if (n, 1./q) not in coeffs:
1087 coeffs[n, 1./q] = signal.remez(
1088 n+1, (0., .75/q, 1./q, 1.),
1089 (1., 0.), fs=2, weight=(1, 50))
1090 b = coeffs[n, 1./q]
1091 return b, [1.], n
1093 else:
1094 coeffs = GlobalVars.decimate_iir_coeffs
1095 if (n, 0.05, 0.8/q) not in coeffs:
1096 coeffs[n, 0.05, 0.8/q] = signal.cheby1(n, 0.05, 0.8/q)
1098 b, a = coeffs[n, 0.05, 0.8/q]
1099 return b, a, n
1102def decimate(x, q, n=None, ftype='iir', zi=None, ioff=0):
1103 '''
1104 Downsample the signal x by an integer factor q, using an order n filter
1106 By default, an order 8 Chebyshev type I filter is used or a 30 point FIR
1107 filter with hamming window if ftype is 'fir'.
1109 :param x: the signal to be downsampled (1D :class:`numpy.ndarray`)
1110 :param q: the downsampling factor
1111 :param n: order of the filter (1 less than the length of the filter for a
1112 `fir` filter)
1113 :param ftype: type of the filter; can be `iir`, `fir` or `fir-remez`
1115 :returns: the downsampled signal (1D :class:`numpy.ndarray`)
1117 '''
1119 b, a, n = decimate_coeffs(q, n, ftype)
1121 if zi is None or zi is True:
1122 zi_ = num.zeros(max(len(a), len(b))-1, dtype=float)
1123 else:
1124 zi_ = zi
1126 y, zf = signal.lfilter(b, a, x, zi=zi_)
1128 if zi is not None:
1129 return y[n//2+ioff::q].copy(), zf
1130 else:
1131 return y[n//2+ioff::q].copy()
1134class UnavailableDecimation(Exception):
1135 '''
1136 Exception raised by :py:func:`decitab` for unavailable decimation factors.
1137 '''
1139 pass
1142def gcd(a, b, epsilon=1e-7):
1143 '''
1144 Greatest common divisor.
1145 '''
1147 while b > epsilon*a:
1148 a, b = b, a % b
1150 return a
1153def lcm(a, b):
1154 '''
1155 Least common multiple.
1156 '''
1158 return a*b // gcd(a, b)
1161def mk_decitab(nmax=100):
1162 '''
1163 Make table with decimation sequences.
1165 Decimation from one sampling rate to a lower one is achieved by a
1166 successive application of :py:func:`decimation` with small integer
1167 downsampling factors (because using large downampling factors can make the
1168 decimation unstable or slow). This function sets up a table with downsample
1169 sequences for factors up to ``nmax``.
1170 '''
1172 tab = GlobalVars.decitab
1173 for i in range(1, 10):
1174 for j in range(1, i+1):
1175 for k in range(1, j+1):
1176 for l_ in range(1, k+1):
1177 for m in range(1, l_+1):
1178 p = i*j*k*l_*m
1179 if p > nmax:
1180 break
1181 if p not in tab:
1182 tab[p] = (i, j, k, l_, m)
1183 if i*j*k*l_ > nmax:
1184 break
1185 if i*j*k > nmax:
1186 break
1187 if i*j > nmax:
1188 break
1189 if i > nmax:
1190 break
1192 GlobalVars.decitab_nmax = nmax
1195def zfmt(n):
1196 return '%%0%ii' % (int(math.log10(n - 1)) + 1)
1199def _year_to_time(year):
1200 tt = (year, 1, 1, 0, 0, 0)
1201 return to_time_float(calendar.timegm(tt))
1204def _working_year(year):
1205 try:
1206 tt = (year, 1, 1, 0, 0, 0)
1207 t = calendar.timegm(tt)
1208 tt2_ = time.gmtime(t)
1209 tt2 = tuple(tt2_)[:6]
1210 if tt != tt2:
1211 return False
1213 s = '%i-01-01 00:00:00' % year
1214 s2 = time.strftime('%Y-%m-%d %H:%M:%S', tt2_)
1215 if s != s2:
1216 return False
1218 t3 = str_to_time(s2, format='%Y-%m-%d %H:%M:%S')
1219 s3 = time_to_str(t3, format='%Y-%m-%d %H:%M:%S')
1220 if s3 != s2:
1221 return False
1223 except Exception:
1224 return False
1226 return True
1229def working_system_time_range(year_min_lim=None, year_max_lim=None):
1230 '''
1231 Check time range supported by the systems's time conversion functions.
1233 Returns system time stamps of start of year of first/last fully supported
1234 year span. If this is before 1900 or after 2100, return first/last century
1235 which is fully supported.
1237 :returns: ``(tmin, tmax, year_min, year_max)``
1238 '''
1240 year0 = 2000
1241 year_min = year0
1242 year_max = year0
1244 itests = list(range(101))
1245 for i in range(19):
1246 itests.append(200 + i*100)
1248 for i in itests:
1249 year = year0 - i
1250 if year_min_lim is not None and year < year_min_lim:
1251 year_min = year_min_lim
1252 break
1253 elif not _working_year(year):
1254 break
1255 else:
1256 year_min = year
1258 for i in itests:
1259 year = year0 + i
1260 if year_max_lim is not None and year > year_max_lim:
1261 year_max = year_max_lim
1262 break
1263 elif not _working_year(year + 1):
1264 break
1265 else:
1266 year_max = year
1268 return (
1269 _year_to_time(year_min),
1270 _year_to_time(year_max),
1271 year_min, year_max)
1274g_working_system_time_range = None
1277def get_working_system_time_range():
1278 '''
1279 Caching variant of :py:func:`working_system_time_range`.
1280 '''
1282 global g_working_system_time_range
1283 if g_working_system_time_range is None:
1284 g_working_system_time_range = working_system_time_range()
1286 return g_working_system_time_range
1289def is_working_time(t):
1290 tmin, tmax, _, _ = get_working_system_time_range()
1291 return tmin <= t <= tmax
1294def julian_day_of_year(timestamp):
1295 '''
1296 Get the day number after the 1st of January of year in ``timestamp``.
1298 :returns: day number as int
1299 '''
1301 return time.gmtime(int(timestamp)).tm_yday
1304def hour_start(timestamp):
1305 '''
1306 Get beginning of hour for any point in time.
1308 :param timestamp: time instant as system timestamp (in seconds)
1310 :returns: instant of hour start as system timestamp
1311 '''
1313 tt = time.gmtime(int(timestamp))
1314 tts = tt[0:4] + (0, 0)
1315 return to_time_float(calendar.timegm(tts))
1318def day_start(timestamp):
1319 '''
1320 Get beginning of day for any point in time.
1322 :param timestamp: time instant as system timestamp (in seconds)
1324 :returns: instant of day start as system timestamp
1325 '''
1327 tt = time.gmtime(int(timestamp))
1328 tts = tt[0:3] + (0, 0, 0)
1329 return to_time_float(calendar.timegm(tts))
1332def month_start(timestamp):
1333 '''
1334 Get beginning of month for any point in time.
1336 :param timestamp: time instant as system timestamp (in seconds)
1338 :returns: instant of month start as system timestamp
1339 '''
1341 tt = time.gmtime(int(timestamp))
1342 tts = tt[0:2] + (1, 0, 0, 0)
1343 return to_time_float(calendar.timegm(tts))
1346def year_start(timestamp):
1347 '''
1348 Get beginning of year for any point in time.
1350 :param timestamp: time instant as system timestamp (in seconds)
1352 :returns: instant of year start as system timestamp
1353 '''
1355 tt = time.gmtime(int(timestamp))
1356 tts = tt[0:1] + (1, 1, 0, 0, 0)
1357 return to_time_float(calendar.timegm(tts))
1360def iter_days(tmin, tmax):
1361 '''
1362 Yields begin and end of days until given time span is covered.
1364 :param tmin,tmax: input time span
1366 :yields: tuples with (begin, end) of days as system timestamps
1367 '''
1369 t = day_start(tmin)
1370 while t < tmax:
1371 tend = day_start(t + 26*60*60)
1372 yield t, tend
1373 t = tend
1376def iter_months(tmin, tmax):
1377 '''
1378 Yields begin and end of months until given time span is covered.
1380 :param tmin,tmax: input time span
1382 :yields: tuples with (begin, end) of months as system timestamps
1383 '''
1385 t = month_start(tmin)
1386 while t < tmax:
1387 tend = month_start(t + 24*60*60*33)
1388 yield t, tend
1389 t = tend
1392def iter_years(tmin, tmax):
1393 '''
1394 Yields begin and end of years until given time span is covered.
1396 :param tmin,tmax: input time span
1398 :yields: tuples with (begin, end) of years as system timestamps
1399 '''
1401 t = year_start(tmin)
1402 while t < tmax:
1403 tend = year_start(t + 24*60*60*369)
1404 yield t, tend
1405 t = tend
1408def today():
1409 return day_start(time.time())
1412def tomorrow():
1413 return day_start(time.time() + 24*60*60)
1416def decitab(n):
1417 '''
1418 Get integer decimation sequence for given downampling factor.
1420 :param n: target decimation factor
1422 :returns: tuple with downsampling sequence
1423 '''
1425 if n > GlobalVars.decitab_nmax:
1426 mk_decitab(n*2)
1427 if n not in GlobalVars.decitab:
1428 raise UnavailableDecimation('ratio = %g' % n)
1429 return GlobalVars.decitab[n]
1432def ctimegm(s, format="%Y-%m-%d %H:%M:%S"):
1433 '''
1434 Convert string representing UTC time to system time.
1436 :param s: string to be interpreted
1437 :param format: format string passed to :py:func:`strptime`
1439 :returns: system time stamp
1441 Interpretes string with format ``'%Y-%m-%d %H:%M:%S'``, using strptime.
1443 .. note::
1444 This function is to be replaced by :py:func:`str_to_time`.
1445 '''
1447 return calendar.timegm(time.strptime(s, format))
1450def gmctime(t, format="%Y-%m-%d %H:%M:%S"):
1451 '''
1452 Get string representation from system time, UTC.
1454 Produces string with format ``'%Y-%m-%d %H:%M:%S'``, using strftime.
1456 .. note::
1457 This function is to be repaced by :py:func:`time_to_str`.
1458 '''
1460 return time.strftime(format, time.gmtime(t))
1463def gmctime_v(t, format="%a, %d %b %Y %H:%M:%S"):
1464 '''
1465 Get string representation from system time, UTC. Same as
1466 :py:func:`gmctime` but with a more verbose default format.
1468 .. note::
1469 This function is to be replaced by :py:func:`time_to_str`.
1470 '''
1472 return time.strftime(format, time.gmtime(t))
1475def gmctime_fn(t, format="%Y-%m-%d_%H-%M-%S"):
1476 '''
1477 Get string representation from system time, UTC. Same as
1478 :py:func:`gmctime` but with a default usable in filenames.
1480 .. note::
1481 This function is to be replaced by :py:func:`time_to_str`.
1482 '''
1484 return time.strftime(format, time.gmtime(t))
1487class TimeStrError(Exception):
1488 pass
1491class FractionalSecondsMissing(TimeStrError):
1492 '''
1493 Exception raised by :py:func:`str_to_time` when the given string lacks
1494 fractional seconds.
1495 '''
1497 pass
1500class FractionalSecondsWrongNumberOfDigits(TimeStrError):
1501 '''
1502 Exception raised by :py:func:`str_to_time` when the given string has an
1503 incorrect number of digits in the fractional seconds part.
1504 '''
1506 pass
1509def _endswith_n(s, endings):
1510 for ix, x in enumerate(endings):
1511 if s.endswith(x):
1512 return ix
1513 return -1
1516def str_to_time(s, format='%Y-%m-%d %H:%M:%S.OPTFRAC'):
1517 '''
1518 Convert string representing UTC time to floating point system time.
1520 :param s: string representing UTC time
1521 :param format: time string format
1522 :returns: system time stamp as floating point value
1524 Uses the semantics of :py:func:`time.strptime` but allows for fractional
1525 seconds. If the format ends with ``'.FRAC'``, anything after a dot is
1526 interpreted as fractional seconds. If the format ends with ``'.OPTFRAC'``,
1527 the fractional part, including the dot is made optional. The latter has the
1528 consequence, that the time strings and the format may not contain any other
1529 dots. If the format ends with ``'.xFRAC'`` where x is 1, 2, or 3, it is
1530 ensured, that exactly that number of digits are present in the fractional
1531 seconds.
1532 '''
1534 time_float = get_time_float()
1536 if util_ext is not None:
1537 try:
1538 t, tfrac = util_ext.stt(s, format)
1539 except util_ext.UtilExtError as e:
1540 raise TimeStrError(
1541 '%s, string=%s, format=%s' % (str(e), s, format))
1543 return time_float(t)+tfrac
1545 fracsec = 0.
1546 fixed_endings = '.FRAC', '.1FRAC', '.2FRAC', '.3FRAC'
1548 iend = _endswith_n(format, fixed_endings)
1549 if iend != -1:
1550 dotpos = s.rfind('.')
1551 if dotpos == -1:
1552 raise FractionalSecondsMissing(
1553 'string=%s, format=%s' % (s, format))
1555 if iend > 0 and iend != (len(s)-dotpos-1):
1556 raise FractionalSecondsWrongNumberOfDigits(
1557 'string=%s, format=%s' % (s, format))
1559 format = format[:-len(fixed_endings[iend])]
1560 fracsec = float(s[dotpos:])
1561 s = s[:dotpos]
1563 elif format.endswith('.OPTFRAC'):
1564 dotpos = s.rfind('.')
1565 format = format[:-8]
1566 if dotpos != -1 and len(s[dotpos:]) > 1:
1567 fracsec = float(s[dotpos:])
1569 if dotpos != -1:
1570 s = s[:dotpos]
1572 try:
1573 return time_float(calendar.timegm(time.strptime(s, format))) \
1574 + fracsec
1575 except ValueError as e:
1576 raise TimeStrError('%s, string=%s, format=%s' % (str(e), s, format))
1579stt = str_to_time
1582def str_to_time_fillup(s):
1583 '''
1584 Default :py:func:`str_to_time` with filling in of missing values.
1586 Allows e.g. `'2010-01-01 00:00:00'` as `'2010-01-01 00:00'`,
1587 `'2010-01-01 00'`, ..., or `'2010'`.
1588 '''
1590 if len(s) in (4, 7, 10, 13, 16):
1591 s += '0000-01-01 00:00:00'[len(s):]
1593 return str_to_time(s)
1596def time_to_str(t, format='%Y-%m-%d %H:%M:%S.3FRAC'):
1597 '''
1598 Get string representation for floating point system time.
1600 :param t: floating point system time
1601 :param format: time string format
1602 :returns: string representing UTC time
1604 Uses the semantics of :py:func:`time.strftime` but additionally allows for
1605 fractional seconds. If ``format`` contains ``'.xFRAC'``, where ``x`` is a
1606 digit between 1 and 9, this is replaced with the fractional part of ``t``
1607 with ``x`` digits precision.
1608 '''
1610 if pyrocko.grumpy > 0:
1611 check_time_class(t, 'warn' if pyrocko.grumpy == 1 else 'raise')
1613 if isinstance(format, int):
1614 if format > 0:
1615 format = '%Y-%m-%d %H:%M:%S.' + '%iFRAC' % format
1616 else:
1617 format = '%Y-%m-%d %H:%M:%S'
1619 if util_ext is not None:
1620 t0 = num.floor(t)
1621 try:
1622 return util_ext.tts(int(t0), float(t - t0), format)
1623 except util_ext.UtilExtError as e:
1624 raise TimeStrError(
1625 '%s, timestamp=%f, format=%s' % (str(e), t, format))
1627 if not GlobalVars.re_frac:
1628 GlobalVars.re_frac = re.compile(r'\.[1-9]FRAC')
1629 GlobalVars.frac_formats = dict(
1630 [('.%sFRAC' % x, '%.'+x+'f') for x in '123456789'])
1632 ts = float(num.floor(t))
1633 tfrac = t-ts
1635 m = GlobalVars.re_frac.search(format)
1636 if m:
1637 sfrac = (GlobalVars.frac_formats[m.group(0)] % tfrac)
1638 if sfrac[0] == '1':
1639 ts += 1.
1641 format, nsub = GlobalVars.re_frac.subn(sfrac[1:], format, 1)
1643 return time.strftime(format, time.gmtime(ts))
1646tts = time_to_str
1647_abbr_weekday = 'Mon Tue Wed Thu Fri Sat Sun'.split()
1648_abbr_month = 'Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec'.split()
1651def mystrftime(fmt=None, tt=None, milliseconds=0):
1652 # Needed by Snuffler for the time axis. In other cases `time_to_str`
1653 # should be used.
1655 if fmt is None:
1656 fmt = '%Y-%m-%d %H:%M:%S .%r'
1658 # Get these two locale independent, needed by Snuffler.
1659 # Setting the locale seems to have no effect.
1660 fmt = fmt.replace('%a', _abbr_weekday[tt.tm_wday])
1661 fmt = fmt.replace('%b', _abbr_month[tt.tm_mon-1])
1663 fmt = fmt.replace('%r', '%03i' % int(round(milliseconds)))
1664 fmt = fmt.replace('%u', '%06i' % int(round(milliseconds*1000)))
1665 fmt = fmt.replace('%n', '%09i' % int(round(milliseconds*1000000)))
1666 return time.strftime(fmt, tt)
1669def gmtime_x(timestamp):
1670 etimestamp = float(num.floor(timestamp))
1671 tt = time.gmtime(etimestamp)
1672 ms = (timestamp-etimestamp)*1000
1673 return tt, ms
1676def plural_s(n):
1677 if not isinstance(n, int):
1678 n = len(n)
1680 return 's' if n != 1 else ''
1683def ensuredirs(dst):
1684 '''
1685 Create all intermediate path components for a target path.
1687 :param dst: target path
1689 The leaf part of the target path is not created (use :py:func:`ensuredir`
1690 if a the target path is a directory to be created).
1691 '''
1693 d, x = os.path.split(dst.rstrip(os.sep))
1694 dirs = []
1695 while d and not os.path.exists(d):
1696 dirs.append(d)
1697 d, x = os.path.split(d)
1699 dirs.reverse()
1701 for d in dirs:
1702 try:
1703 os.mkdir(d)
1704 except OSError as e:
1705 if not e.errno == errno.EEXIST:
1706 raise
1709def ensuredir(dst):
1710 '''
1711 Create directory and all intermediate path components to it as needed.
1713 :param dst: directory name
1715 Nothing is done if the given target already exists.
1716 '''
1718 if os.path.exists(dst):
1719 return
1721 dst.rstrip(os.sep)
1723 ensuredirs(dst)
1724 try:
1725 os.mkdir(dst)
1726 except OSError as e:
1727 if not e.errno == errno.EEXIST:
1728 raise
1731def reuse(x):
1732 '''
1733 Get unique instance of an object.
1735 :param x: hashable object
1736 :returns: reference to x or an equivalent object
1738 Cache object ``x`` in a global dict for reuse, or if x already
1739 is in that dict, return a reference to it.
1740 '''
1742 grs = GlobalVars.reuse_store
1743 if x not in grs:
1744 grs[x] = x
1745 return grs[x]
1748def deuse(x):
1749 grs = GlobalVars.reuse_store
1750 if x in grs:
1751 del grs[x]
1754class Anon(object):
1755 '''
1756 Dict-to-object utility.
1758 Any given arguments are stored as attributes.
1760 Example::
1762 a = Anon(x=1, y=2)
1763 print a.x, a.y
1764 '''
1766 def __init__(self, **dict):
1767 for k in dict:
1768 self.__dict__[k] = dict[k]
1771def iter_select_files(
1772 paths,
1773 include=None,
1774 exclude=None,
1775 selector=None,
1776 show_progress=True,
1777 pass_through=None):
1779 '''
1780 Recursively select files (generator variant).
1782 See :py:func:`select_files`.
1783 '''
1785 if show_progress:
1786 progress_beg('selecting files...')
1788 ngood = 0
1789 check_include = None
1790 if include is not None:
1791 rinclude = re.compile(include)
1793 def check_include(path):
1794 m = rinclude.search(path)
1795 if not m:
1796 return False
1798 if selector is None:
1799 return True
1801 infos = Anon(**m.groupdict())
1802 return selector(infos)
1804 check_exclude = None
1805 if exclude is not None:
1806 rexclude = re.compile(exclude)
1808 def check_exclude(path):
1809 return not bool(rexclude.search(path))
1811 if check_include and check_exclude:
1813 def check(path):
1814 return check_include(path) and check_exclude(path)
1816 elif check_include:
1817 check = check_include
1819 elif check_exclude:
1820 check = check_exclude
1822 else:
1823 check = None
1825 if isinstance(paths, str):
1826 paths = [paths]
1828 for path in paths:
1829 if pass_through and pass_through(path):
1830 if check is None or check(path):
1831 yield path
1833 elif os.path.isdir(path):
1834 for (dirpath, dirnames, filenames) in os.walk(path):
1835 dirnames.sort()
1836 filenames.sort()
1837 for filename in filenames:
1838 path = op.join(dirpath, filename)
1839 if check is None or check(path):
1840 yield os.path.abspath(path)
1841 ngood += 1
1842 else:
1843 if check is None or check(path):
1844 yield os.path.abspath(path)
1845 ngood += 1
1847 if show_progress:
1848 progress_end('%i file%s selected.' % (ngood, plural_s(ngood)))
1851def select_files(
1852 paths, include=None, exclude=None, selector=None, show_progress=True,
1853 regex=None):
1855 '''
1856 Recursively select files.
1858 :param paths: entry path names
1859 :param include: pattern for conditional inclusion
1860 :param exclude: pattern for conditional exclusion
1861 :param selector: callback for conditional inclusion
1862 :param show_progress: if True, indicate start and stop of processing
1863 :param regex: alias for ``include`` (backwards compatibility)
1864 :returns: list of path names
1866 Recursively finds all files under given entry points ``paths``. If
1867 parameter ``include`` is a regular expression, only files with matching
1868 path names are included. If additionally parameter ``selector`` is given a
1869 callback function, only files for which the callback returns ``True`` are
1870 included. The callback should take a single argument. The callback is
1871 called with a single argument, an object, having as attributes, any named
1872 groups given in ``include``.
1874 Examples
1876 To find all files ending in ``'.mseed'`` or ``'.msd'``::
1878 select_files(paths,
1879 include=r'\\.(mseed|msd)$')
1881 To find all files ending with ``'$Year.$DayOfYear'``, having set 2009 for
1882 the year::
1884 select_files(paths,
1885 include=r'(?P<year>\\d\\d\\d\\d)\\.(?P<doy>\\d\\d\\d)$',
1886 selector=(lambda x: int(x.year) == 2009))
1887 '''
1888 if regex is not None:
1889 assert include is None
1890 include = regex
1892 return list(iter_select_files(
1893 paths, include, exclude, selector, show_progress))
1896def base36encode(number, alphabet='0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'):
1897 '''
1898 Convert positive integer to a base36 string.
1899 '''
1901 if not isinstance(number, (int, long)):
1902 raise TypeError('number must be an integer')
1903 if number < 0:
1904 raise ValueError('number must be positive')
1906 # Special case for small numbers
1907 if number < 36:
1908 return alphabet[number]
1910 base36 = ''
1911 while number != 0:
1912 number, i = divmod(number, 36)
1913 base36 = alphabet[i] + base36
1915 return base36
1918def base36decode(number):
1919 '''
1920 Decode base36 endcoded positive integer.
1921 '''
1923 return int(number, 36)
1926class UnpackError(Exception):
1927 '''
1928 Exception raised when :py:func:`unpack_fixed` encounters an error.
1929 '''
1931 pass
1934ruler = ''.join(['%-10i' % i for i in range(8)]) \
1935 + '\n' + '0123456789' * 8 + '\n'
1938def unpack_fixed(format, line, *callargs):
1939 '''
1940 Unpack fixed format string, as produced by many fortran codes.
1942 :param format: format specification
1943 :param line: string to be processed
1944 :param callargs: callbacks for callback fields in the format
1946 The format is described by a string of comma-separated fields. Each field
1947 is defined by a character for the field type followed by the field width. A
1948 questionmark may be appended to the field description to allow the argument
1949 to be optional (The data string is then allowed to be filled with blanks
1950 and ``None`` is returned in this case).
1952 The following field types are available:
1954 ==== ================================================================
1955 Type Description
1956 ==== ================================================================
1957 A string (full field width is extracted)
1958 a string (whitespace at the beginning and the end is removed)
1959 i integer value
1960 f floating point value
1961 @ special type, a callback must be given for the conversion
1962 x special field type to skip parts of the string
1963 ==== ================================================================
1964 '''
1966 ipos = 0
1967 values = []
1968 icall = 0
1969 for form in format.split(','):
1970 form = form.strip()
1971 optional = form[-1] == '?'
1972 form = form.rstrip('?')
1973 typ = form[0]
1974 ln = int(form[1:])
1975 s = line[ipos:ipos+ln]
1976 cast = {
1977 'x': None,
1978 'A': str,
1979 'a': lambda x: x.strip(),
1980 'i': int,
1981 'f': float,
1982 '@': 'extra'}[typ]
1984 if cast == 'extra':
1985 cast = callargs[icall]
1986 icall += 1
1988 if cast is not None:
1989 if optional and s.strip() == '':
1990 values.append(None)
1991 else:
1992 try:
1993 values.append(cast(s))
1994 except Exception:
1995 mark = [' '] * 80
1996 mark[ipos:ipos+ln] = ['^'] * ln
1997 mark = ''.join(mark)
1998 raise UnpackError(
1999 'Invalid cast to type "%s" at position [%i:%i] of '
2000 'line: \n%s%s\n%s' % (
2001 typ, ipos, ipos+ln, ruler, line.rstrip(), mark))
2003 ipos += ln
2005 return values
2008_pattern_cache = {}
2011def _nslc_pattern(pattern):
2012 if pattern not in _pattern_cache:
2013 rpattern = re.compile(fnmatch.translate(pattern), re.I)
2014 _pattern_cache[pattern] = rpattern
2015 else:
2016 rpattern = _pattern_cache[pattern]
2018 return rpattern
2021def match_nslc(patterns, nslc):
2022 '''
2023 Match network-station-location-channel code against pattern or list of
2024 patterns.
2026 :param patterns: pattern or list of patterns
2027 :param nslc: tuple with (network, station, location, channel) as strings
2029 :returns: ``True`` if the pattern matches or if any of the given patterns
2030 match; or ``False``.
2032 The patterns may contain shell-style wildcards: \\*, ?, [seq], [!seq].
2034 Example::
2036 match_nslc('*.HAM3.*.BH?', ('GR', 'HAM3', '', 'BHZ')) # -> True
2037 '''
2039 if isinstance(patterns, str):
2040 patterns = [patterns]
2042 if not isinstance(nslc, str):
2043 s = '.'.join(nslc)
2044 else:
2045 s = nslc
2047 for pattern in patterns:
2048 if _nslc_pattern(pattern).match(s):
2049 return True
2051 return False
2054def match_nslcs(patterns, nslcs):
2055 '''
2056 Get network-station-location-channel codes that match given pattern or any
2057 of several given patterns.
2059 :param patterns: pattern or list of patterns
2060 :param nslcs: list of (network, station, location, channel) tuples
2062 See also :py:func:`match_nslc`
2063 '''
2065 matching = []
2066 for nslc in nslcs:
2067 if match_nslc(patterns, nslc):
2068 matching.append(nslc)
2070 return matching
2073class Timeout(Exception):
2074 pass
2077def create_lockfile(fn, timeout=None, timewarn=10.):
2078 t0 = time.time()
2080 while True:
2081 try:
2082 f = os.open(fn, os.O_CREAT | os.O_WRONLY | os.O_EXCL)
2083 os.close(f)
2084 return
2086 except OSError as e:
2087 if e.errno in (errno.EEXIST, 13): # 13 occurs on windows
2088 pass # retry
2089 else:
2090 raise
2092 tnow = time.time()
2094 if timeout is not None and tnow - t0 > timeout:
2095 raise Timeout(
2096 'Timeout (%gs) occured while waiting to get exclusive '
2097 'access to: %s' % (timeout, fn))
2099 if timewarn is not None and tnow - t0 > timewarn:
2100 logger.warning(
2101 'Waiting since %gs to get exclusive access to: %s' % (
2102 timewarn, fn))
2104 timewarn *= 2
2106 time.sleep(0.01)
2109def delete_lockfile(fn):
2110 os.unlink(fn)
2113class Lockfile(Exception):
2115 def __init__(self, path, timeout=5, timewarn=10.):
2116 self._path = path
2117 self._timeout = timeout
2118 self._timewarn = timewarn
2120 def __enter__(self):
2121 create_lockfile(
2122 self._path, timeout=self._timeout, timewarn=self._timewarn)
2123 return None
2125 def __exit__(self, type, value, traceback):
2126 delete_lockfile(self._path)
2129class SoleError(Exception):
2130 '''
2131 Exception raised by objects of type :py:class:`Sole`, when an concurrent
2132 instance is running.
2133 '''
2135 pass
2138class Sole(object):
2139 '''
2140 Use POSIX advisory file locking to ensure that only a single instance of a
2141 program is running.
2143 :param pid_path: path to lockfile to be used
2145 Usage::
2147 from pyrocko.util import Sole, SoleError, setup_logging
2148 import os
2150 setup_logging('my_program')
2152 pid_path = os.path.join(os.environ['HOME'], '.my_program_lock')
2153 try:
2154 sole = Sole(pid_path)
2156 except SoleError, e:
2157 logger.fatal( str(e) )
2158 sys.exit(1)
2159 '''
2161 def __init__(self, pid_path):
2162 self._pid_path = pid_path
2163 self._other_running = False
2164 ensuredirs(self._pid_path)
2165 self._lockfile = None
2166 self._os = os
2168 if not fcntl:
2169 raise SoleError(
2170 'Python standard library module "fcntl" not available on '
2171 'this platform.')
2173 self._fcntl = fcntl
2175 try:
2176 self._lockfile = os.open(self._pid_path, os.O_CREAT | os.O_WRONLY)
2177 except Exception:
2178 raise SoleError(
2179 'Cannot open lockfile (path = %s)' % self._pid_path)
2181 try:
2182 fcntl.lockf(self._lockfile, fcntl.LOCK_EX | fcntl.LOCK_NB)
2184 except IOError:
2185 self._other_running = True
2186 try:
2187 f = open(self._pid_path, 'r')
2188 pid = f.read().strip()
2189 f.close()
2190 except Exception:
2191 pid = '?'
2193 raise SoleError('Other instance is running (pid = %s)' % pid)
2195 try:
2196 os.ftruncate(self._lockfile, 0)
2197 os.write(self._lockfile, '%i\n' % os.getpid())
2198 os.fsync(self._lockfile)
2200 except Exception:
2201 # the pid is only stored for user information, so this is allowed
2202 # to fail
2203 pass
2205 def __del__(self):
2206 if not self._other_running:
2207 if self._lockfile is not None:
2208 self._fcntl.lockf(self._lockfile, self._fcntl.LOCK_UN)
2209 self._os.close(self._lockfile)
2210 try:
2211 self._os.unlink(self._pid_path)
2212 except Exception:
2213 pass
2216re_escapequotes = re.compile(r"(['\\])")
2219def escapequotes(s):
2220 return re_escapequotes.sub(r"\\\1", s)
2223class TableWriter(object):
2224 '''
2225 Write table of space separated values to a file.
2227 :param f: file like object
2229 Strings containing spaces are quoted on output.
2230 '''
2232 def __init__(self, f):
2233 self._f = f
2235 def writerow(self, row, minfieldwidths=None):
2237 '''
2238 Write one row of values to underlying file.
2240 :param row: iterable of values
2241 :param minfieldwidths: minimum field widths for the values
2243 Each value in in ``row`` is converted to a string and optionally padded
2244 with blanks. The resulting strings are output separated with blanks. If
2245 any values given are strings and if they contain whitespace, they are
2246 quoted with single quotes, and any internal single quotes are
2247 backslash-escaped.
2248 '''
2250 out = []
2252 for i, x in enumerate(row):
2253 w = 0
2254 if minfieldwidths and i < len(minfieldwidths):
2255 w = minfieldwidths[i]
2257 if isinstance(x, str):
2258 if re.search(r"\s|'", x):
2259 x = "'%s'" % escapequotes(x)
2261 x = x.ljust(w)
2262 else:
2263 x = str(x).rjust(w)
2265 out.append(x)
2267 self._f.write(' '.join(out).rstrip() + '\n')
2270class TableReader(object):
2272 '''
2273 Read table of space separated values from a file.
2275 :param f: file-like object
2277 This uses Pythons shlex module to tokenize lines. Should deal correctly
2278 with quoted strings.
2279 '''
2281 def __init__(self, f):
2282 self._f = f
2283 self.eof = False
2285 def readrow(self):
2286 '''
2287 Read one row from the underlying file, tokenize it with shlex.
2289 :returns: tokenized line as a list of strings.
2290 '''
2292 line = self._f.readline()
2293 if not line:
2294 self.eof = True
2295 return []
2296 line.strip()
2297 if line.startswith('#'):
2298 return []
2300 return qsplit(line)
2303def gform(number, significant_digits=3):
2304 '''
2305 Pretty print floating point numbers.
2307 Align floating point numbers at the decimal dot.
2309 ::
2311 | -d.dde+xxx|
2312 | -d.dde+xx |
2313 |-ddd. |
2314 | -dd.d |
2315 | -d.dd |
2316 | -0.ddd |
2317 | -0.0ddd |
2318 | -0.00ddd |
2319 | -d.dde-xx |
2320 | -d.dde-xxx|
2321 | nan|
2324 The formatted string has length ``significant_digits * 2 + 6``.
2325 '''
2327 no_exp_range = (pow(10., -1),
2328 pow(10., significant_digits))
2329 width = significant_digits+significant_digits-1+1+1+5
2331 if (no_exp_range[0] <= abs(number) < no_exp_range[1]) or number == 0.:
2332 s = ('%#.*g' % (significant_digits, number)).rstrip('0')
2333 else:
2334 s = '%.*E' % (significant_digits-1, number)
2335 s = (' '*(-s.find('.')+(significant_digits+1))+s).ljust(width)
2336 if s.strip().lower() == 'nan':
2337 s = 'nan'.rjust(width)
2338 return s
2341def human_bytesize(value):
2343 exts = 'Bytes kB MB GB TB PB EB ZB YB'.split()
2345 if value == 1:
2346 return '1 Byte'
2348 for i, ext in enumerate(exts):
2349 x = float(value) / 1000**i
2350 if round(x) < 10. and not value < 1000:
2351 return '%.1f %s' % (x, ext)
2352 if round(x) < 1000.:
2353 return '%.0f %s' % (x, ext)
2355 return '%i Bytes' % value
2358re_compatibility = re.compile(
2359 r'!pyrocko\.(trace|gf\.(meta|seismosizer)|fomosto\.' +
2360 r'(dummy|poel|qseis|qssp))\.'
2361)
2364def pf_is_old(fn):
2365 oldstyle = False
2366 with open(fn, 'r') as f:
2367 for line in f:
2368 if re_compatibility.search(line):
2369 oldstyle = True
2371 return oldstyle
2374def pf_upgrade(fn):
2375 need = pf_is_old(fn)
2376 if need:
2377 fn_temp = fn + '.temp'
2379 with open(fn, 'r') as fin:
2380 with open(fn_temp, 'w') as fout:
2381 for line in fin:
2382 line = re_compatibility.sub('!pf.', line)
2383 fout.write(line)
2385 os.rename(fn_temp, fn)
2387 return need
2390def read_leap_seconds(tzfile='/usr/share/zoneinfo/right/UTC'):
2391 '''
2392 Extract leap second information from tzdata.
2394 Based on example at http://stackoverflow.com/questions/19332902/\
2395 extract-historic-leap-seconds-from-tzdata
2397 See also 'man 5 tzfile'.
2398 '''
2400 from struct import unpack, calcsize
2401 out = []
2402 with open(tzfile, 'rb') as f:
2403 # read header
2404 fmt = '>4s c 15x 6l'
2405 (magic, format, ttisgmtcnt, ttisstdcnt, leapcnt, timecnt,
2406 typecnt, charcnt) = unpack(fmt, f.read(calcsize(fmt)))
2407 assert magic == 'TZif'.encode('US-ASCII'), 'Not a timezone file'
2409 # skip over some uninteresting data
2410 fmt = '>%(timecnt)dl %(timecnt)dB %(ttinfo)s %(charcnt)ds' % dict(
2411 timecnt=timecnt, ttinfo='lBB'*typecnt, charcnt=charcnt)
2412 f.read(calcsize(fmt))
2414 # read leap-seconds
2415 fmt = '>2l'
2416 for i in range(leapcnt):
2417 tleap, nleap = unpack(fmt, f.read(calcsize(fmt)))
2418 out.append((tleap-nleap+1, nleap))
2420 return out
2423class LeapSecondsError(Exception):
2424 pass
2427class LeapSecondsOutdated(LeapSecondsError):
2428 pass
2431class InvalidLeapSecondsFile(LeapSecondsOutdated):
2432 pass
2435def parse_leap_seconds_list(fn):
2436 data = []
2437 texpires = None
2438 try:
2439 t0 = int(round(str_to_time('1900-01-01 00:00:00')))
2440 except TimeStrError:
2441 t0 = int(round(str_to_time('1970-01-01 00:00:00'))) - 2208988800
2443 tnow = int(round(time.time()))
2445 if not op.exists(fn):
2446 raise LeapSecondsOutdated('no leap seconds file found')
2448 try:
2449 with open(fn, 'rb') as f:
2450 for line in f:
2451 if line.strip().startswith(b'<!DOCTYPE'):
2452 raise InvalidLeapSecondsFile('invalid leap seconds file')
2454 if line.startswith(b'#@'):
2455 texpires = int(line.split()[1]) + t0
2456 elif line.startswith(b'#') or len(line) < 5:
2457 pass
2458 else:
2459 toks = line.split()
2460 t = int(toks[0]) + t0
2461 nleap = int(toks[1]) - 10
2462 data.append((t, nleap))
2464 except IOError:
2465 raise LeapSecondsError('cannot read leap seconds file %s' % fn)
2467 if texpires is None or tnow > texpires:
2468 raise LeapSecondsOutdated('leap seconds list is outdated')
2470 return data
2473def read_leap_seconds2():
2474 from pyrocko import config
2475 conf = config.config()
2476 fn = conf.leapseconds_path
2477 url = conf.leapseconds_url
2478 # check for outdated default URL
2479 if url == 'http://www.ietf.org/timezones/data/leap-seconds.list':
2480 url = 'https://www.ietf.org/timezones/data/leap-seconds.list'
2481 logger.info(
2482 'Leap seconds default URL is now: %s\nUsing new default.' % url)
2484 for i in range(3):
2485 try:
2486 return parse_leap_seconds_list(fn)
2488 except LeapSecondsOutdated:
2489 try:
2490 logger.info('updating leap seconds list...')
2491 download_file(url, fn)
2493 except Exception as e:
2494 raise LeapSecondsError(
2495 'cannot download leap seconds list from %s to %s (%s)'
2496 % (url, fn, e))
2498 raise LeapSecondsError('Could not retrieve/read leap seconds file.')
2501def gps_utc_offset(t_utc):
2502 '''
2503 Time offset t_gps - t_utc for a given t_utc.
2504 '''
2505 ls = read_leap_seconds2()
2506 i = 0
2507 if t_utc < ls[0][0]:
2508 return ls[0][1] - 1 - 9
2510 while i < len(ls) - 1:
2511 if ls[i][0] <= t_utc and t_utc < ls[i+1][0]:
2512 return ls[i][1] - 9
2513 i += 1
2515 return ls[-1][1] - 9
2518def utc_gps_offset(t_gps):
2519 '''
2520 Time offset t_utc - t_gps for a given t_gps.
2521 '''
2522 ls = read_leap_seconds2()
2524 if t_gps < ls[0][0] + ls[0][1] - 9:
2525 return - (ls[0][1] - 1 - 9)
2527 i = 0
2528 while i < len(ls) - 1:
2529 if ls[i][0] + ls[i][1] - 9 <= t_gps \
2530 and t_gps < ls[i+1][0] + ls[i+1][1] - 9:
2531 return - (ls[i][1] - 9)
2532 i += 1
2534 return - (ls[-1][1] - 9)
2537def make_iload_family(iload_fh, doc_fmt='FMT', doc_yielded_objects='FMT'):
2538 import itertools
2539 import glob
2540 from pyrocko.io.io_common import FileLoadError
2542 def iload_filename(filename, **kwargs):
2543 try:
2544 with open(filename, 'rb') as f:
2545 for cr in iload_fh(f, **kwargs):
2546 yield cr
2548 except FileLoadError as e:
2549 e.set_context('filename', filename)
2550 raise
2552 def iload_dirname(dirname, **kwargs):
2553 for entry in os.listdir(dirname):
2554 fpath = op.join(dirname, entry)
2555 if op.isfile(fpath):
2556 for cr in iload_filename(fpath, **kwargs):
2557 yield cr
2559 def iload_glob(pattern, **kwargs):
2561 for fn in glob.iglob(pattern):
2562 for cr in iload_filename(fn, **kwargs):
2563 yield cr
2565 def iload(source, **kwargs):
2566 if isinstance(source, str):
2567 if op.isdir(source):
2568 return iload_dirname(source, **kwargs)
2569 elif op.isfile(source):
2570 return iload_filename(source, **kwargs)
2571 else:
2572 return iload_glob(source, **kwargs)
2574 elif hasattr(source, 'read'):
2575 return iload_fh(source, **kwargs)
2576 else:
2577 return itertools.chain.from_iterable(
2578 iload(subsource, **kwargs) for subsource in source)
2580 iload_filename.__doc__ = '''
2581 Read %s information from named file.
2582 ''' % doc_fmt
2584 iload_dirname.__doc__ = '''
2585 Read %s information from directory of %s files.
2586 ''' % (doc_fmt, doc_fmt)
2588 iload_glob.__doc__ = '''
2589 Read %s information from files matching a glob pattern.
2590 ''' % doc_fmt
2592 iload.__doc__ = '''
2593 Load %s information from given source(s)
2595 The ``source`` can be specified as the name of a %s file, the name of a
2596 directory containing %s files, a glob pattern of %s files, an open
2597 filehandle or an iterator yielding any of the forementioned sources.
2599 This function behaves as a generator yielding %s objects.
2600 ''' % (doc_fmt, doc_fmt, doc_fmt, doc_fmt, doc_yielded_objects)
2602 for f in iload_filename, iload_dirname, iload_glob, iload:
2603 f.__module__ = iload_fh.__module__
2605 return iload_filename, iload_dirname, iload_glob, iload
2608class Inconsistency(Exception):
2609 pass
2612def consistency_check(list_of_tuples, message='values differ:'):
2613 '''
2614 Check for inconsistencies.
2616 Given a list of tuples, check that all tuple elements except for first one
2617 match. E.g. ``[('STA.N', 55.3, 103.2), ('STA.E', 55.3, 103.2)]`` would be
2618 valid because the coordinates at the two channels are the same.
2619 '''
2621 if len(list_of_tuples) >= 2:
2622 if any(t[1:] != list_of_tuples[0][1:] for t in list_of_tuples[1:]):
2623 raise Inconsistency('%s\n' % message + '\n'.join(
2624 ' %s: %s' % (t[0], ', '.join(str(x) for x in t[1:]))
2625 for t in list_of_tuples))
2628class defaultzerodict(dict):
2629 def __missing__(self, k):
2630 return 0
2633def mostfrequent(x):
2634 c = defaultzerodict()
2635 for e in x:
2636 c[e] += 1
2638 return sorted(list(c.keys()), key=lambda k: c[k])[-1]
2641def consistency_merge(list_of_tuples,
2642 message='values differ:',
2643 error='raise',
2644 merge=mostfrequent):
2646 assert error in ('raise', 'warn', 'ignore')
2648 if len(list_of_tuples) == 0:
2649 raise Exception('cannot merge empty sequence')
2651 try:
2652 consistency_check(list_of_tuples, message)
2653 return list_of_tuples[0][1:]
2654 except Inconsistency as e:
2655 if error == 'raise':
2656 raise
2658 elif error == 'warn':
2659 logger.warning(str(e))
2661 return tuple([merge(x) for x in list(zip(*list_of_tuples))[1:]])
2664def short_to_list(nmax, it):
2665 import itertools
2667 if isinstance(it, list):
2668 return it
2670 li = []
2671 for i in range(nmax+1):
2672 try:
2673 li.append(next(it))
2674 except StopIteration:
2675 return li
2677 return itertools.chain(li, it)
2680def parse_md(f):
2681 try:
2682 with open(op.join(
2683 op.dirname(op.abspath(f)),
2684 'README.md'), 'r') as readme:
2685 mdstr = readme.read()
2686 except IOError as e:
2687 return 'Failed to get README.md: %s' % e
2689 # Remve the title
2690 mdstr = re.sub(r'^# .*\n?', '', mdstr)
2691 # Append sphinx reference to `pyrocko.` modules
2692 mdstr = re.sub(r'`pyrocko\.(.*)`', r':py:mod:`pyrocko.\1`', mdstr)
2693 # Convert Subsections to toc-less rubrics
2694 mdstr = re.sub(r'## (.*)\n', r'.. rubric:: \1\n', mdstr)
2695 return mdstr
2698def mpl_show(plt):
2699 import matplotlib
2700 if matplotlib.get_backend().lower() == 'agg':
2701 logger.warning('Cannot show() when using matplotlib "agg" backend')
2702 else:
2703 plt.show()
2706g_re_qsplit = re.compile(
2707 r'"([^"\\]*(?:\\.[^"\\]*)*)"|\'([^\'\\]*(?:\\.[^\'\\]*)*)\'|(\S+)')
2708g_re_qsplit_sep = {}
2711def get_re_qsplit(sep):
2712 if sep is None:
2713 return g_re_qsplit
2714 else:
2715 if sep not in g_re_qsplit_sep:
2716 assert len(sep) == 1
2717 assert sep not in '\'"'
2718 esep = re.escape(sep)
2719 g_re_qsplit_sep[sep] = re.compile(
2720 r'"([^"\\]*(?:\\.[^"\\]*)*)"|\'([^\'\\]*(?:\\.[^\'\\]*)*)\'|'
2721 + r'([^' + esep + r']+|(?<=' + esep + r')(?=' + esep + r')|^(?=' + esep + r')|(?<=' + esep + r')$)') # noqa
2722 return g_re_qsplit_sep[sep]
2725g_re_trivial = re.compile(r'\A[^\'"\s]+\Z')
2726g_re_trivial_sep = {}
2729def get_re_trivial(sep):
2730 if sep is None:
2731 return g_re_trivial
2732 else:
2733 if sep not in g_re_qsplit_sep:
2734 assert len(sep) == 1
2735 assert sep not in '\'"'
2736 esep = re.escape(sep)
2737 g_re_trivial_sep[sep] = re.compile(r'\A[^\'"' + esep + r']+\Z')
2739 return g_re_trivial_sep[sep]
2742g_re_escape_s = re.compile(r'([\\\'])')
2743g_re_unescape_s = re.compile(r'\\([\\\'])')
2744g_re_escape_d = re.compile(r'([\\"])')
2745g_re_unescape_d = re.compile(r'\\([\\"])')
2748def escape_s(s):
2749 '''
2750 Backslash-escape single-quotes and backslashes.
2752 Example: ``Jack's`` => ``Jack\\'s``
2754 '''
2755 return g_re_escape_s.sub(r'\\\1', s)
2758def unescape_s(s):
2759 '''
2760 Unescape backslash-escaped single-quotes and backslashes.
2762 Example: ``Jack\\'s`` => ``Jack's``
2763 '''
2764 return g_re_unescape_s.sub(r'\1', s)
2767def escape_d(s):
2768 '''
2769 Backslash-escape double-quotes and backslashes.
2771 Example: ``"Hello \\O/"`` => ``\\"Hello \\\\O/\\"``
2772 '''
2773 return g_re_escape_d.sub(r'\\\1', s)
2776def unescape_d(s):
2777 '''
2778 Unescape backslash-escaped double-quotes and backslashes.
2780 Example: ``\\"Hello \\\\O/\\"`` => ``"Hello \\O/"``
2781 '''
2782 return g_re_unescape_d.sub(r'\1', s)
2785def qjoin_s(it, sep=None):
2786 '''
2787 Join sequence of strings into a line, single-quoting non-trivial strings.
2789 Example: ``["55", "Sparrow's Island"]`` => ``"55 'Sparrow\\\\'s Island'"``
2790 '''
2791 re_trivial = get_re_trivial(sep)
2793 if sep is None:
2794 sep = ' '
2796 return sep.join(
2797 w if re_trivial.search(w) else "'%s'" % escape_s(w) for w in it)
2800def qjoin_d(it, sep=None):
2801 '''
2802 Join sequence of strings into a line, double-quoting non-trivial strings.
2804 Example: ``['55', 'Pete "The Robot" Smith']`` =>
2805 ``'55' "Pete \\\\"The Robot\\\\" Smith"'``
2806 '''
2807 re_trivial = get_re_trivial(sep)
2808 if sep is None:
2809 sep = ' '
2811 return sep.join(
2812 w if re_trivial.search(w) else '"%s"' % escape_d(w) for w in it)
2815def qsplit(s, sep=None):
2816 '''
2817 Split line into list of strings, allowing for quoted strings.
2819 Example: ``"55 'Sparrow\\\\'s Island'"`` =>
2820 ``["55", "Sparrow's Island"]``,
2821 ``'55' "Pete \\\\"The Robot\\\\" Smith"'`` =>
2822 ``['55', 'Pete "The Robot" Smith']``
2823 '''
2824 re_qsplit = get_re_qsplit(sep)
2825 return [
2826 (unescape_d(x[0]) or unescape_s(x[1]) or x[2])
2827 for x in re_qsplit.findall(s)]
2830g_have_warned_threadpoolctl = False
2833class threadpool_limits_dummy(object):
2835 def __init__(self, *args, **kwargs):
2836 pass
2838 def __enter__(self):
2839 global g_have_warned_threadpoolctl
2841 if not g_have_warned_threadpoolctl:
2842 logger.warning(
2843 'Cannot control number of BLAS threads because '
2844 '`threadpoolctl` module is not available. You may want to '
2845 'install `threadpoolctl`.')
2847 g_have_warned_threadpoolctl = True
2849 return self
2851 def __exit__(self, type, value, traceback):
2852 pass
2855def get_threadpool_limits():
2856 '''
2857 Try to import threadpoolctl.threadpool_limits, provide dummy if not avail.
2858 '''
2860 try:
2861 from threadpoolctl import threadpool_limits
2862 return threadpool_limits
2864 except ImportError:
2865 return threadpool_limits_dummy
2868def smart_weakref(obj, callback=None):
2869 if inspect.ismethod(obj):
2870 return weakref.WeakMethod(obj, callback)
2871 else:
2872 return weakref.ref(obj, callback)