Coverage for /usr/local/lib/python3.7/dist-packages/pyrocko/pile.py : 74%

# 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: