|
# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------
from __future__ import absolute_import, division, print_function
import os
import logging
import time
import weakref
import copy
import re
import sys
import operator
import math
import hashlib
try:
import cPickle as pickle
except ImportError:
import pickle
from . import avl
from . import trace, io, util
from . import config
from .trace import degapper
show_progress_force_off = False
version_salt = 'v1-'
def ehash(s):
return hashlib.sha1((version_salt + s).encode('utf8')).hexdigest()
def cmp(a, b):
return int(a > b) - int(a < b)
def sl(s):
return [str(x) for x in sorted(s)]
class Counter(dict):
def __missing__(self, k):
return 0
def update(self, other):
for k, v in other.items():
self[k] += v
def subtract(self, other):
for k, v in other.items():
self[k] -= v
if self[k] <= 0:
del self[k]
def subtract1(self, k):
self[k] -= 1
if self[k] <= 0:
del self[k]
def fix_unicode_copy(counter, func):
counter_new = Counter()
for k in counter:
counter_new[func(k)] = counter[k]
return counter_new
pjoin = os.path.join
logger = logging.getLogger('pyrocko.pile')
def avl_remove_exact(avltree, element):
ilo, ihi = avltree.span(element)
for i in range(ilo, ihi):
if avltree[i] is element:
avltree.remove_at(i)
return
raise ValueError(
'avl_remove_exact(avltree, element): element not in avltree')
def cmpfunc(key):
if isinstance(key, str):
# special cases; these run about 50% faster than the generic one on
# Python 2.5
if key == 'tmin':
return lambda a, b: cmp(a.tmin, b.tmin)
if key == 'tmax':
return lambda a, b: cmp(a.tmax, b.tmax)
key = operator.attrgetter(key)
return lambda a, b: cmp(key(a), key(b))
g_dummys = {}
def get_dummy(key):
if key not in g_dummys:
class Dummy(object):
def __init__(self, k):
setattr(self, key, k)
g_dummys[key] = Dummy
return g_dummys[key]
class Sorted(object):
def __init__(self, values=[], key=None):
self._set_key(key)
self._avl = avl.new(values, self._cmp)
def _set_key(self, key):
self._key = key
self._cmp = cmpfunc(key)
if isinstance(key, str):
self._dummy = get_dummy(key)
def __getstate__(self):
state = list(self._avl.iter()), self._key
return state
def __setstate__(self, state):
l, key = state
self._set_key(key)
self._avl = avl.from_iter(iter(l), len(l))
def insert(self, value):
self._avl.insert(value)
def remove(self, value):
avl_remove_exact(self._avl, value)
def insert_many(self, values):
for value in values:
self._avl.insert(value)
def remove_many(self, values):
for value in values:
avl_remove_exact(self._avl, value)
def __iter__(self):
return iter(self._avl)
def with_key_in(self, kmin, kmax):
omin, omax = self._dummy(kmin), self._dummy(kmax)
ilo, ihi = self._avl.span(omin, omax)
return self._avl[ilo:ihi]
def min(self):
return self._avl.min()
def max(self):
return self._avl.max()
def __len__(self):
return len(self._avl)
class TracesFileCache(object):
'''Manages trace metainformation cache.
For each directory with files containing traces, one cache file is
maintained to hold the trace metainformation of all files which are
contained in the directory.
'''
caches = {}
def __init__(self, cachedir):
'''Create new cache.
:param cachedir: directory to hold the cache files.
'''
self.cachedir = cachedir
self.dircaches = {}
self.modified = set()
util.ensuredir(self.cachedir)
def get(self, abspath):
'''Try to get an item from the cache.
:param abspath: absolute path of the object to retrieve
:returns: a stored object is returned or None if nothing could be
found.
'''
dircache = self._get_dircache_for(abspath)
if abspath in dircache:
return dircache[abspath]
return None
def put(self, abspath, tfile):
'''Put an item into the cache.
:param abspath: absolute path of the object to be stored
:param tfile: object to be stored
'''
cachepath = self._dircachepath(abspath)
# get lock on cachepath here
dircache = self._get_dircache(cachepath)
dircache[abspath] = tfile
self.modified.add(cachepath)
def dump_modified(self):
'''Save any modifications to disk.'''
for cachepath in self.modified:
self._dump_dircache(self.dircaches[cachepath], cachepath)
# unlock
self.modified = set()
def clean(self):
'''Weed out missing files from the disk caches.'''
self.dump_modified()
for fn in os.listdir(self.cachedir):
if len(fn) == 40:
cache = self._load_dircache(pjoin(self.cachedir, fn))
self._dump_dircache(cache, pjoin(self.cachedir, fn))
def _get_dircache_for(self, abspath):
return self._get_dircache(self._dircachepath(abspath))
def _get_dircache(self, cachepath):
if cachepath not in self.dircaches:
if os.path.isfile(cachepath):
self.dircaches[cachepath] = self._load_dircache(cachepath)
else:
self.dircaches[cachepath] = {}
return self.dircaches[cachepath]
def _dircachepath(self, abspath):
cachefn = ehash(os.path.dirname(abspath))
return pjoin(self.cachedir, cachefn)
def _load_dircache(self, cachefilename):
with open(cachefilename, 'rb') as f:
cache = pickle.load(f)
# weed out files which no longer exist
for fn in list(cache.keys()):
if not os.path.isfile(fn):
del cache[fn]
for v in cache.values():
v.trees_from_content(v.traces)
for tr in v.traces:
tr.file = v
# fix Py2 codes to not include unicode when the cache file
# was created with Py3
if not isinstance(tr.station, str):
tr.prune_from_reuse_cache()
tr.set_codes(
str(tr.network),
str(tr.station),
str(tr.location),
str(tr.channel))
v.data_use_count = 0
v.data_loaded = False
v.fix_unicode_codes()
return cache
def _dump_dircache(self, cache, cachefilename):
if not cache:
if os.path.exists(cachefilename):
os.remove(cachefilename)
return
# make a copy without the parents and the binsearch trees
cache_copy = {}
for fn in cache.keys():
trf = copy.copy(cache[fn])
trf.parent = None
trf.by_tmin = None
trf.by_tmax = None
trf.by_tlen = None
trf.by_mtime = None
trf.data_use_count = 0
trf.data_loaded = False
traces = []
for tr in trf.traces:
tr = tr.copy(data=False)
tr.ydata = None
tr.meta = None
tr.file = trf
traces.append(tr)
trf.traces = traces
cache_copy[fn] = trf
tmpfn = cachefilename+'.%i.tmp' % os.getpid()
with open(tmpfn, 'wb') as f:
pickle.dump(cache_copy, f, protocol=2)
os.rename(tmpfn, cachefilename)
def get_cache(cachedir):
'''Get global TracesFileCache object for given directory.'''
if cachedir not in TracesFileCache.caches:
TracesFileCache.caches[cachedir] = TracesFileCache(cachedir)
return TracesFileCache.caches[cachedir]
def loader(
filenames, fileformat, cache, filename_attributes,
show_progress=True, update_progress=None):
if show_progress_force_off:
show_progress = False
class Progress(object):
def __init__(self, label, n):
self._label = label
self._n = n
self._bar = None
if show_progress:
self._bar = util.progressbar(label, self._n)
if update_progress:
update_progress(label, 0, self._n)
def update(self, i):
if self._bar:
if i < self._n-1:
self._bar.update(i)
else:
self._bar.finish()
self._bar = None
abort = False
if update_progress:
abort = update_progress(self._label, i, self._n)
return abort
if not filenames:
logger.warning('No files to load from')
return
regex = None
if filename_attributes:
regex = re.compile(filename_attributes)
progress = Progress('Looking at files', len(filenames))
failures = []
to_load = []
for i, filename in enumerate(filenames):
try:
abspath = os.path.abspath(filename)
substitutions = None
if regex:
m = regex.search(filename)
if not m:
raise FilenameAttributeError(
"Cannot get attributes with pattern '%s' "
"from path '%s'" % (filename_attributes, filename))
substitutions = {}
for k in m.groupdict():
if k in ('network', 'station', 'location', 'channel'):
substitutions[k] = m.groupdict()[k]
mtime = os.stat(filename)[8]
tfile = None
if cache:
tfile = cache.get(abspath)
mustload = (
not tfile or
(tfile.format != fileformat and fileformat != 'detect') or
tfile.mtime != mtime or
substitutions is not None)
to_load.append((mustload, mtime, abspath, substitutions, tfile))
except (OSError, FilenameAttributeError) as xerror:
failures.append(abspath)
logger.warning(xerror)
abort = progress.update(i+1)
if abort:
progress.update(len(filenames))
return
progress.update(len(filenames))
to_load.sort(key=lambda x: x[2])
nload = len([1 for x in to_load if x[0]])
iload = 0
count_all = False
if nload < 0.01*len(to_load):
nload = len(to_load)
count_all = True
if to_load:
progress = Progress('Scanning files', nload)
for (mustload, mtime, abspath, substitutions, tfile) in to_load:
try:
if mustload:
tfile = TracesFile(
None, abspath, fileformat,
substitutions=substitutions, mtime=mtime)
if cache and not substitutions:
cache.put(abspath, tfile)
if not count_all:
iload += 1
if count_all:
iload += 1
except (io.FileLoadError, OSError) as xerror:
failures.append(abspath)
logger.warning(xerror)
else:
yield tfile
abort = progress.update(iload+1)
if abort:
break
progress.update(nload)
if failures:
logger.warning(
'The following file%s caused problems and will be ignored:\n' %
util.plural_s(len(failures)) + '\n'.join(failures))
if cache:
cache.dump_modified()
def tlen(x):
return x.tmax-x.tmin
class TracesGroup(object):
'''Trace container base class.
Base class for Pile, SubPile, and TracesFile, i.e. anything containing
a collection of several traces. A TracesGroup object maintains lookup sets
of some of the traces meta-information, as well as a combined time-range
of its contents.
'''
def __init__(self, parent):
self.parent = parent
self.empty()
self.nupdates = 0
self.abspath = None
def set_parent(self, parent):
self.parent = parent
def get_parent(self):
return self.parent
def empty(self):
self.networks, self.stations, self.locations, self.channels, \
self.nslc_ids, self.deltats = [Counter() for x in range(6)]
self.by_tmin = Sorted([], 'tmin')
self.by_tmax = Sorted([], 'tmax')
self.by_tlen = Sorted([], tlen)
self.by_mtime = Sorted([], 'mtime')
self.tmin, self.tmax = None, None
self.deltatmin, self.deltatmax = None, None
def trees_from_content(self, content):
self.by_tmin = Sorted(content, 'tmin')
self.by_tmax = Sorted(content, 'tmax')
self.by_tlen = Sorted(content, tlen)
self.by_mtime = Sorted(content, 'mtime')
self.adjust_minmax()
def fix_unicode_codes(self):
for net in self.networks:
if isinstance(net, str):
return
self.networks = fix_unicode_copy(self.networks, str)
self.stations = fix_unicode_copy(self.stations, str)
self.locations = fix_unicode_copy(self.locations, str)
self.channels = fix_unicode_copy(self.channels, str)
self.nslc_ids = fix_unicode_copy(
self.nslc_ids, lambda k: tuple(str(x) for x in k))
def add(self, content):
'''
Add content to traces group and update indices.
Accepts :py:class:`pyrocko.trace.Trace` objects and
:py:class:`pyrocko.pile.TracesGroup` objects.
'''
if isinstance(content, (trace.Trace, TracesGroup)):
content = [content]
for c in content:
if isinstance(c, TracesGroup):
self.networks.update(c.networks)
self.stations.update(c.stations)
self.locations.update(c.locations)
self.channels.update(c.channels)
self.nslc_ids.update(c.nslc_ids)
self.deltats.update(c.deltats)
self.by_tmin.insert_many(c.by_tmin)
self.by_tmax.insert_many(c.by_tmax)
self.by_tlen.insert_many(c.by_tlen)
self.by_mtime.insert_many(c.by_mtime)
elif isinstance(c, trace.Trace):
self.networks[c.network] += 1
self.stations[c.station] += 1
self.locations[c.location] += 1
self.channels[c.channel] += 1
self.nslc_ids[c.nslc_id] += 1
self.deltats[c.deltat] += 1
self.by_tmin.insert(c)
self.by_tmax.insert(c)
self.by_tlen.insert(c)
self.by_mtime.insert(c)
self.adjust_minmax()
self.nupdates += 1
self.notify_listeners('add')
if self.parent is not None:
self.parent.add(content)
def remove(self, content):
'''
Remove content to traces group and update indices.
'''
if isinstance(content, (trace.Trace, TracesGroup)):
content = [content]
for c in content:
if isinstance(c, TracesGroup):
self.networks.subtract(c.networks)
self.stations.subtract(c.stations)
self.locations.subtract(c.locations)
self.channels.subtract(c.channels)
self.nslc_ids.subtract(c.nslc_ids)
self.deltats.subtract(c.deltats)
self.by_tmin.remove_many(c.by_tmin)
self.by_tmax.remove_many(c.by_tmax)
self.by_tlen.remove_many(c.by_tlen)
self.by_mtime.remove_many(c.by_mtime)
elif isinstance(c, trace.Trace):
self.networks.subtract1(c.network)
self.stations.subtract1(c.station)
self.locations.subtract1(c.location)
self.channels.subtract1(c.channel)
self.nslc_ids.subtract1(c.nslc_id)
self.deltats.subtract1(c.deltat)
self.by_tmin.remove(c)
self.by_tmax.remove(c)
self.by_tlen.remove(c)
self.by_mtime.remove(c)
self.adjust_minmax()
self.nupdates += 1
self.notify_listeners('remove')
if self.parent is not None:
self.parent.remove(content)
def relevant(self, tmin, tmax, group_selector=None, trace_selector=None):
'''Return list of :py:class:`pyrocko.trace.Trace` objects where given
arguments ``tmin`` and ``tmax`` match.
:param tmin: start time
:param tmax: end time
:param group_selector: lambda expression taking group dict of regex
match object as a single argument and which returns true or false
to keep or reject a file (default: ``None``)
:param trace_selector: lambda expression taking group dict of regex
match object as a single argument and which returns true or false
to keep or reject a file (default: ``None``)
'''
if not self.by_tmin or not self.is_relevant(
tmin, tmax, group_selector):
return []
return [tr for tr in self.by_tmin.with_key_in(tmin-self.tlenmax, tmax)
if tr.is_relevant(tmin, tmax, trace_selector)]
def adjust_minmax(self):
if self.by_tmin:
self.tmin = self.by_tmin.min().tmin
self.tmax = self.by_tmax.max().tmax
t = self.by_tlen.max()
self.tlenmax = t.tmax - t.tmin
self.mtime = self.by_mtime.max().mtime
deltats = list(self.deltats.keys())
self.deltatmin = min(deltats)
self.deltatmax = max(deltats)
else:
self.tmin = None
self.tmax = None
self.tlenmax = None
self.mtime = None
self.deltatmin = None
self.deltatmax = None
def notify_listeners(self, what):
pass
def get_update_count(self):
return self.nupdates
def overlaps(self, tmin, tmax):
return self.tmin is not None \
and tmax >= self.tmin and self.tmax >= tmin
def is_relevant(self, tmin, tmax, group_selector=None):
if self.tmin is None or self.tmax is None:
return False
return tmax >= self.tmin and self.tmax >= tmin and (
group_selector is None or group_selector(self))
class MemTracesFile(TracesGroup):
'''This is needed to make traces without an actual disc file to be inserted
into a Pile.'''
def __init__(self, parent, traces):
TracesGroup.__init__(self, parent)
self.add(traces)
self.mtime = time.time()
def add(self, traces):
if isinstance(traces, trace.Trace):
traces = [traces]
for tr in traces:
tr.file = self
TracesGroup.add(self, traces)
def load_headers(self, mtime=None):
pass
def load_data(self):
pass
def use_data(self):
pass
def drop_data(self):
pass
def reload_if_modified(self):
return False
def iter_traces(self):
for tr in self.by_tmin:
yield tr
def get_traces(self):
return list(self.by_tmin)
def gather_keys(self, gather, selector=None):
keys = set()
for tr in self.by_tmin:
if selector is None or selector(tr):
keys.add(gather(tr))
return keys
def __str__(self):
s = 'MemTracesFile\n'
s += 'file mtime: %s\n' % util.time_to_str(self.mtime)
s += 'number of traces: %i\n' % len(self.by_tmin)
s += 'timerange: %s - %s\n' % (
util.time_to_str(self.tmin), util.time_to_str(self.tmax))
s += 'networks: %s\n' % ', '.join(sl(self.networks.keys()))
s += 'stations: %s\n' % ', '.join(sl(self.stations.keys()))
s += 'locations: %s\n' % ', '.join(sl(self.locations.keys()))
s += 'channels: %s\n' % ', '.join(sl(self.channels.keys()))
s += 'deltats: %s\n' % ', '.join(sl(self.deltats.keys()))
return s
class TracesFile(TracesGroup):
def __init__(
self, parent, abspath, format,
substitutions=None, mtime=None):
TracesGroup.__init__(self, parent)
self.abspath = abspath
self.format = format
self.traces = []
self.data_loaded = False
self.data_use_count = 0
self.substitutions = substitutions
self.load_headers(mtime=mtime)
self.mtime = mtime
def load_headers(self, mtime=None):
logger.debug('loading headers from file: %s' % self.abspath)
if mtime is None:
self.mtime = os.stat(self.abspath)[8]
def kgen(tr):
return (tr.mtime, tr.tmin, tr.tmax) + tr.nslc_id
self.remove(self.traces)
ks = set()
for tr in io.load(self.abspath,
format=self.format,
getdata=False,
substitutions=self.substitutions):
k = kgen(tr)
if k not in ks:
ks.add(k)
self.traces.append(tr)
tr.file = self
self.add(self.traces)
self.data_loaded = False
self.data_use_count = 0
def load_data(self, force=False):
file_changed = False
if not self.data_loaded or force:
logger.debug('loading data from file: %s' % self.abspath)
def kgen(tr):
return (tr.mtime, tr.tmin, tr.tmax) + tr.nslc_id
traces_ = io.load(self.abspath, format=self.format, getdata=True,
substitutions=self.substitutions)
# prevent adding duplicate snippets from corrupt mseed files
k_loaded = set()
traces = []
for tr in traces_:
k = kgen(tr)
if k not in k_loaded:
k_loaded.add(k)
traces.append(tr)
k_current_d = dict((kgen(tr), tr) for tr in self.traces)
k_current = set(k_current_d)
k_new = k_loaded - k_current
k_delete = k_current - k_loaded
k_unchanged = k_current & k_loaded
for tr in self.traces[:]:
if kgen(tr) in k_delete:
self.remove(tr)
self.traces.remove(tr)
tr.file = None
file_changed = True
for tr in traces:
if kgen(tr) in k_new:
tr.file = self
self.traces.append(tr)
self.add(tr)
file_changed = True
for tr in traces:
if kgen(tr) in k_unchanged:
ctr = k_current_d[kgen(tr)]
ctr.ydata = tr.ydata
self.data_loaded = True
if file_changed:
logger.debug('reloaded (file may have changed): %s' % self.abspath)
return file_changed
def use_data(self):
if not self.data_loaded:
raise Exception('Data not loaded')
self.data_use_count += 1
def drop_data(self):
if self.data_loaded:
if self.data_use_count == 1:
logger.debug('forgetting data of file: %s' % self.abspath)
for tr in self.traces:
tr.drop_data()
self.data_loaded = False
self.data_use_count -= 1
else:
self.data_use_count = 0
def reload_if_modified(self):
mtime = os.stat(self.abspath)[8]
if mtime != self.mtime:
logger.debug(
'mtime=%i, reloading file: %s' % (mtime, self.abspath))
self.mtime = mtime
if self.data_loaded:
self.load_data(force=True)
else:
self.load_headers()
return True
return False
def iter_traces(self):
for tr in self.traces:
yield tr
def gather_keys(self, gather, selector=None):
keys = set()
for tr in self.by_tmin:
if selector is None or selector(tr):
keys.add(gather(tr))
return keys
def __str__(self):
s = 'TracesFile\n'
s += 'abspath: %s\n' % self.abspath
s += 'file mtime: %s\n' % util.time_to_str(self.mtime)
s += 'number of traces: %i\n' % len(self.traces)
s += 'timerange: %s - %s\n' % (
util.time_to_str(self.tmin), util.time_to_str(self.tmax))
s += 'networks: %s\n' % ', '.join(sl(self.networks.keys()))
s += 'stations: %s\n' % ', '.join(sl(self.stations.keys()))
s += 'locations: %s\n' % ', '.join(sl(self.locations.keys()))
s += 'channels: %s\n' % ', '.join(sl(self.channels.keys()))
s += 'deltats: %s\n' % ', '.join(sl(self.deltats.keys()))
return s
class FilenameAttributeError(Exception):
pass
class SubPile(TracesGroup):
def __init__(self, parent):
TracesGroup.__init__(self, parent)
self.files = []
self.empty()
def add_file(self, file):
self.files.append(file)
file.set_parent(self)
self.add(file)
def remove_file(self, file):
self.files.remove(file)
file.set_parent(None)
self.remove(file)
def remove_files(self, files):
for file in files:
self.files.remove(file)
file.set_parent(None)
self.remove(files)
def gather_keys(self, gather, selector=None):
keys = set()
for file in self.files:
keys |= file.gather_keys(gather, selector)
return keys
def iter_traces(
self,
load_data=False,
return_abspath=False,
group_selector=None,
trace_selector=None):
for file in self.files:
if group_selector and not group_selector(file):
continue
must_drop = False
if load_data:
file.load_data()
file.use_data()
must_drop = True
for tr in file.iter_traces():
if trace_selector and not trace_selector(tr):
continue
if return_abspath:
yield file.abspath, tr
else:
yield tr
if must_drop:
file.drop_data()
def iter_files(self):
for file in self.files:
yield file
def reload_modified(self):
modified = False
for file in self.files:
modified |= file.reload_if_modified()
return modified
def __str__(self):
s = 'SubPile\n'
s += 'number of files: %i\n' % len(self.files)
s += 'timerange: %s - %s\n' % (
util.time_to_str(self.tmin), util.time_to_str(self.tmax))
s += 'networks: %s\n' % ', '.join(sl(self.networks.keys()))
s += 'stations: %s\n' % ', '.join(sl(self.stations.keys()))
s += 'locations: %s\n' % ', '.join(sl(self.locations.keys()))
s += 'channels: %s\n' % ', '.join(sl(self.channels.keys()))
s += 'deltats: %s\n' % ', '.join(sl(self.deltats.keys()))
return s
class Pile(TracesGroup):
'''Waveform archive lookup, data loading and caching infrastructure.'''
def __init__(self):
TracesGroup.__init__(self, None)
self.subpiles = {}
self.open_files = {}
self.listeners = []
self.abspaths = set()
def add_listener(self, obj):
self.listeners.append(weakref.ref(obj))
def notify_listeners(self, what):
for ref in self.listeners:
obj = ref()
if obj:
obj.pile_changed(what)
def load_files(
self, filenames,
filename_attributes=None,
fileformat='mseed',
cache=None,
show_progress=True,
update_progress=None):
load = loader(
filenames, fileformat, cache, filename_attributes,
show_progress=show_progress,
update_progress=update_progress)
self.add_files(load)
def add_files(self, files):
for file in files:
self.add_file(file)
def add_file(self, file):
if file.abspath is not None and file.abspath in self.abspaths:
logger.warning('File already in pile: %s' % file.abspath)
return
if file.deltatmin is None:
logger.warning('Sampling rate of all traces are zero in file: %s' %
file.abspath)
return
subpile = self.dispatch(file)
subpile.add_file(file)
if file.abspath is not None:
self.abspaths.add(file.abspath)
def remove_file(self, file):
subpile = file.get_parent()
if subpile is not None:
subpile.remove_file(file)
if file.abspath is not None:
self.abspaths.remove(file.abspath)
def remove_files(self, files):
subpile_files = {}
for file in files:
subpile = file.get_parent()
if subpile not in subpile_files:
subpile_files[subpile] = []
subpile_files[subpile].append(file)
for subpile, files in subpile_files.items():
subpile.remove_files(files)
for file in files:
if file.abspath is not None:
self.abspaths.remove(file.abspath)
def dispatch_key(self, file):
dt = int(math.floor(math.log(file.deltatmin)))
return dt
def dispatch(self, file):
k = self.dispatch_key(file)
if k not in self.subpiles:
self.subpiles[k] = SubPile(self)
return self.subpiles[k]
def get_deltats(self):
return list(self.deltats.keys())
def chop(
self, tmin, tmax,
group_selector=None,
trace_selector=None,
snap=(round, round),
include_last=False,
load_data=True):
chopped = []
used_files = set()
traces = self.relevant(tmin, tmax, group_selector, trace_selector)
if load_data:
files_changed = False
for tr in traces:
if tr.file and tr.file not in used_files:
if tr.file.load_data():
files_changed = True
if tr.file is not None:
used_files.add(tr.file)
if files_changed:
traces = self.relevant(
tmin, tmax, group_selector, trace_selector)
for tr in traces:
if not load_data and tr.ydata is not None:
tr = tr.copy(data=False)
tr.ydata = None
try:
chopped.append(tr.chop(
tmin, tmax,
inplace=False,
snap=snap,
include_last=include_last))
except trace.NoData:
pass
return chopped, used_files
def _process_chopped(
self, chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
tpad):
chopped.sort(key=lambda a: a.full_id)
if degap:
chopped = degapper(chopped, maxgap=maxgap, maxlap=maxlap)
if not want_incomplete:
chopped_weeded = []
for tr in chopped:
emin = tr.tmin - (wmin-tpad)
emax = tr.tmax + tr.deltat - (wmax+tpad)
if (abs(emin) <= 0.5*tr.deltat and abs(emax) <= 0.5*tr.deltat):
chopped_weeded.append(tr)
elif degap:
if (0. < emin <= 5. * tr.deltat and
-5. * tr.deltat <= emax < 0.):
tr.extend(
wmin-tpad,
wmax+tpad-tr.deltat,
fillmethod='repeat')
chopped_weeded.append(tr)
chopped = chopped_weeded
for tr in chopped:
tr.wmin = wmin
tr.wmax = wmax
return chopped
def chopper(
self,
tmin=None, tmax=None, tinc=None, tpad=0.,
group_selector=None, trace_selector=None,
want_incomplete=True, degap=True, maxgap=5, maxlap=None,
keep_current_files_open=False, accessor_id=None,
snap=(round, round), include_last=False, load_data=True):
'''
Get iterator for shifting window wise data extraction from waveform
archive.
:param tmin: start time (default uses start time of available data)
:param tmax: end time (default uses end time of available data)
:param tinc: time increment (window shift time) (default uses
``tmax-tmin``)
:param tpad: padding time appended on either side of the data windows
(window overlap is ``2*tpad``)
:param group_selector: filter callback taking :py:class:`TracesGroup`
objects
:param trace_selector: filter callback taking
:py:class:`pyrocko.trace.Trace` objects
:param want_incomplete: if set to ``False``, gappy/incomplete traces
are discarded from the results
:param degap: whether to try to connect traces and to remove gaps and
overlaps
:param maxgap: maximum gap size in samples which is filled with
interpolated samples when ``degap`` is ``True``
:param maxlap: maximum overlap size in samples which is removed when
``degap`` is ``True``
:param keep_current_files_open: whether to keep cached trace data in
memory after the iterator has ended
:param accessor_id: if given, used as a key to identify different
points of extraction for the decision of when to release cached
trace data (should be used when data is alternately extracted from
more than one region / selection)
:param snap: replaces Python's :py:func:`round` function which is used
to determine indices where to start and end the trace data array
:param include_last: whether to include last sample
:param load_data: whether to load the waveform data. If set to
``False``, traces with no data samples, but with correct
meta-information are returned
:returns: itererator yielding a list of :py:class:`pyrocko.trace.Trace`
objects for every extracted time window
'''
if tmin is None:
if self.tmin is None:
logger.warning('Pile\'s tmin is not set - pile may be empty.')
return
tmin = self.tmin + tpad
if tmax is None:
if self.tmax is None:
logger.warning('Pile\'s tmax is not set - pile may be empty.')
return
tmax = self.tmax - tpad
if tinc is None:
tinc = tmax - tmin
if not self.is_relevant(tmin-tpad, tmax+tpad, group_selector):
return
if accessor_id not in self.open_files:
self.open_files[accessor_id] = set()
open_files = self.open_files[accessor_id]
iwin = 0
while True:
chopped = []
wmin, wmax = tmin+iwin*tinc, min(tmin+(iwin+1)*tinc, tmax)
eps = tinc*1e-6
if wmin >= tmax-eps:
break
chopped, used_files = self.chop(
wmin-tpad, wmax+tpad, group_selector, trace_selector, snap,
include_last, load_data)
for file in used_files - open_files:
# increment datause counter on newly opened files
file.use_data()
open_files.update(used_files)
processed = self._process_chopped(
chopped, degap, maxgap, maxlap, want_incomplete, wmax, wmin,
tpad)
yield processed
unused_files = open_files - used_files
while unused_files:
file = unused_files.pop()
file.drop_data()
open_files.remove(file)
iwin += 1
if not keep_current_files_open:
while open_files:
file = open_files.pop()
file.drop_data()
def all(self, *args, **kwargs):
'''
Shortcut to aggregate :py:meth:`chopper` output into a single list.
'''
alltraces = []
for traces in self.chopper(*args, **kwargs):
alltraces.extend(traces)
return alltraces
def iter_all(self, *args, **kwargs):
for traces in self.chopper(*args, **kwargs):
for tr in traces:
yield tr
def chopper_grouped(self, gather, progress=None, *args, **kwargs):
keys = self.gather_keys(gather)
if len(keys) == 0:
return
outer_group_selector = None
if 'group_selector' in kwargs:
outer_group_selector = kwargs['group_selector']
outer_trace_selector = None
if 'trace_selector' in kwargs:
outer_trace_selector = kwargs['trace_selector']
# the use of this gather-cache makes it impossible to modify the pile
# during chopping
gather_cache = {}
pbar = None
if progress is not None:
pbar = util.progressbar(progress, len(keys))
for ikey, key in enumerate(keys):
def tsel(tr):
return gather(tr) == key and (outer_trace_selector is None or
outer_trace_selector(tr))
def gsel(gr):
if gr not in gather_cache:
gather_cache[gr] = gr.gather_keys(gather)
return key in gather_cache[gr] and (
outer_group_selector is None or outer_group_selector(gr))
kwargs['trace_selector'] = tsel
kwargs['group_selector'] = gsel
for traces in self.chopper(*args, **kwargs):
yield traces
if pbar:
pbar.update(ikey+1)
if pbar:
pbar.finish()
def gather_keys(self, gather, selector=None):
keys = set()
for subpile in self.subpiles.values():
keys |= subpile.gather_keys(gather, selector)
return sorted(keys)
def iter_traces(
self,
load_data=False,
return_abspath=False,
group_selector=None,
trace_selector=None):
'''Iterate over all traces in pile.
:param load_data: whether to load the waveform data, by default empty
traces are yielded
:param return_abspath: if ``True`` yield tuples containing absolute
file path and :py:class:`pyrocko.trace.Trace` objects
:param group_selector: filter callback taking :py:class:`TracesGroup`
objects
:param trace_selector: filter callback taking
:py:class:`pyrocko.trace.Trace` objects
Example; yields only traces, where the station code is 'HH1'::
test_pile = pile.make_pile('/local/test_trace_directory')
for t in test_pile.iter_traces(
trace_selector=lambda tr: tr.station=='HH1'):
print t
'''
for subpile in self.subpiles.values():
if not group_selector or group_selector(subpile):
for tr in subpile.iter_traces(load_data, return_abspath,
group_selector, trace_selector):
yield tr
def iter_files(self):
for subpile in self.subpiles.values():
for file in subpile.iter_files():
yield file
def reload_modified(self):
modified = False
for subpile in self.subpiles.values():
modified |= subpile.reload_modified()
return modified
def get_tmin(self):
return self.tmin
def get_tmax(self):
return self.tmax
def get_deltatmin(self):
return self.deltatmin
def get_deltatmax(self):
return self.deltatmax
def is_empty(self):
return self.tmin is None and self.tmax is None
def __str__(self):
if self.tmin is not None and self.tmax is not None:
tmin = util.time_to_str(self.tmin)
tmax = util.time_to_str(self.tmax)
s = 'Pile\n'
s += 'number of subpiles: %i\n' % len(self.subpiles)
s += 'timerange: %s - %s\n' % (tmin, tmax)
s += 'networks: %s\n' % ', '.join(sl(self.networks.keys()))
s += 'stations: %s\n' % ', '.join(sl(self.stations.keys()))
s += 'locations: %s\n' % ', '.join(sl(self.locations.keys()))
s += 'channels: %s\n' % ', '.join(sl(self.channels.keys()))
s += 'deltats: %s\n' % ', '.join(sl(self.deltats.keys()))
else:
s = 'empty Pile'
return s
def snuffle(self, **kwargs):
'''Visualize it.
:param stations: list of `pyrocko.model.Station` objects or ``None``
:param events: list of `pyrocko.model.Event` objects or ``None``
:param markers: list of `pyrocko.gui_util.Marker` objects or ``None``
:param ntracks: float, number of tracks to be shown initially
(default: 12)
:param follow: time interval (in seconds) for real time follow mode or
``None``
:param controls: bool, whether to show the main controls (default:
``True``)
:param opengl: bool, whether to use opengl (default: ``False``)
'''
from pyrocko.gui.snuffler import snuffle
snuffle(self, **kwargs)
def make_pile(
paths=None, selector=None, regex=None,
fileformat='mseed',
cachedirname=None, show_progress=True):
'''Create pile from given file and directory names.
:param paths: filenames and/or directories to look for traces. If paths is
``None`` ``sys.argv[1:]`` is used.
:param selector: lambda expression taking group dict of regex match object
as a single argument and which returns true or false to keep or reject
a file
:param regex: regular expression which filenames have to match
:param fileformat: format of the files ('mseed', 'sac', 'kan',
'from_extension', 'detect')
:param cachedirname: loader cache is stored under this directory. It is
created as neccessary.
:param show_progress: show progress bar and other progress information
'''
if show_progress_force_off:
show_progress = False
if isinstance(paths, str):
paths = [paths]
if paths is None:
paths = sys.argv[1:]
if cachedirname is None:
cachedirname = config.config().cache_dir
fns = util.select_files(
paths, selector, regex, show_progress=show_progress)
cache = get_cache(cachedirname)
p = Pile()
p.load_files(
sorted(fns),
cache=cache,
fileformat=fileformat,
show_progress=show_progress)
return p
class Injector(trace.States):
def __init__(
self, pile,
fixation_length=None,
path=None,
format='from_extension',
forget_fixed=False):
trace.States.__init__(self)
self._pile = pile
self._fixation_length = fixation_length
self._format = format
self._path = path
self._forget_fixed = forget_fixed
def set_fixation_length(self, l):
'''Set length after which the fixation method is called on buffer traces.
The length should be given in seconds. Give None to disable.
'''
self.fixate_all()
self._fixation_length = l # in seconds
def set_save_path(
self,
path='dump_%(network)s.%(station)s.%(location)s.%(channel)s_'
'%(tmin)s_%(tmax)s.mseed'):
self.fixate_all()
self._path = path
def inject(self, trace):
logger.debug('Received a trace: %s' % trace)
buf = self.get(trace)
if buf is None:
trbuf = trace.copy()
buf = MemTracesFile(None, [trbuf])
self._pile.add_file(buf)
self.set(trace, buf)
else:
self._pile.remove_file(buf)
trbuf = buf.get_traces()[0]
buf.remove(trbuf)
trbuf.append(trace.ydata)
buf.add(trbuf)
self._pile.add_file(buf)
self.set(trace, buf)
trbuf = buf.get_traces()[0]
if self._fixation_length is not None:
if trbuf.tmax - trbuf.tmin > self._fixation_length:
self._fixate(buf, complete=False)
def fixate_all(self):
for state in list(self._states.values()):
self._fixate(state[-1])
self._states = {}
def free(self, buf):
self._fixate(buf)
def _fixate(self, buf, complete=True):
trbuf = buf.get_traces()[0]
del_state = True
if self._path:
if self._fixation_length is not None:
ttmin = trbuf.tmin
ytmin = util.year_start(ttmin)
n = int(math.floor((ttmin - ytmin) / self._fixation_length))
tmin = ytmin + n*self._fixation_length
traces = []
t = tmin
while t <= trbuf.tmax:
try:
traces.append(
trbuf.chop(
t,
t+self._fixation_length,
inplace=False,
snap=(math.ceil, math.ceil)))
except trace.NoData:
pass
t += self._fixation_length
if abs(traces[-1].tmax - (t - trbuf.deltat)) < \
trbuf.deltat/100. or complete:
self._pile.remove_file(buf)
else: # reinsert incomplete last part
new_trbuf = traces.pop()
self._pile.remove_file(buf)
buf.remove(trbuf)
buf.add(new_trbuf)
self._pile.add_file(buf)
del_state = False
else:
traces = [trbuf]
self._pile.remove_file(buf)
fns = io.save(traces, self._path, format=self._format)
if not self._forget_fixed:
self._pile.load_files(
fns, show_progress=False, fileformat=self._format)
if del_state:
del self._states[trbuf.nslc_id]
def __del__(self):
self.fixate_all()
|