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

# http://pyrocko.org - GPLv3 

# 

# The Pyrocko Developers, 21st Century 

# ---|P------/S----------~Lg---------- 

from __future__ import absolute_import 

import os 

import logging 

from . import pile, io 

from . import trace as tracemod 

logger = logging.getLogger('pyrocko.hamster_pile') 

class Processor(object): 

def __init__(self): 

self._buffers = {} 

def process(self, trace): 

return [trace] 

def get_buffer(self, trace): 

nslc = trace.nslc_id 

if nslc not in self._buffers: 

return None 

trbuf = self._buffers[nslc] 

if (abs((trbuf.tmax+trbuf.deltat) - trace.tmin) < 1.0e-1*trbuf.deltat 

and trbuf.ydata.dtype == trace.ydata.dtype 

and trbuf.deltat == trace.deltat): 

return trbuf 

else: 

return None 

def set_buffer(self, trace): 

nslc = trace.nslc_id 

self._buffers[nslc] = trace 

def empty_buffer(self, trace): 

nslc = trace.nslc_id 

del self._buffers[nslc] 

def flush_buffers(self): 

traces = list(self._buffers.values()) 

self._buffers = {} 

return traces 

class Renamer(object): 

def __init__(self, mapping): 

self._mapping = mapping 

def process(self, trace): 

target_id = self._mapping(trace) 

if target_id is None: 

return [] 

out = trace.copy() 

out.set_codes(*target_id) 

return [out] 

class Chain(Processor): 

def __init__(self, *processors): 

self._processors = processors 

def process(self, trace): 

traces = [trace] 

xtraces = [] 

for p in self._processors: 

for tr in traces: 

xtraces.extend(p.process(traces)) 

return xtraces 

class Downsampler(Processor): 

def __init__(self, mapping, deltat): 

Processor.__init__(self) 

self._mapping = mapping 

self._downsampler = tracemod.co_downsample_to(deltat) 

def process(self, trace): 

target_id = self._mapping(trace) 

if target_id is None: 

return [] 

ds_trace = self._downsampler.send(trace) 

ds_trace.set_codes(*target_id) 

if ds_trace.data_len() == 0: 

return [] 

return [ds_trace] 

def __del__(self): 

self._downsampler.close() 

class Grower(Processor): 

def __init__(self, tflush=None): 

Processor.__init__(self) 

self._tflush = tflush 

def process(self, trace): 

buffer = self.get_buffer(trace) 

if buffer is None: 

buffer = trace 

self.set_buffer(buffer) 

else: 

buffer.append(trace.ydata) 

if buffer.tmax - buffer.tmin >= self._tflush: 

self.empty_buffer(buffer) 

return [buffer] 

else: 

return [] 

class HamsterPile(pile.Pile): 

def __init__( 

self, 

fixation_length=None, 

path=None, 

format='from_extension', 

forget_fixed=False, 

processors=None): 

pile.Pile.__init__(self) 

self._buffers = {} # keys: nslc, values: MemTracesFile 

self._fixation_length = fixation_length 

self._format = format 

self._path = path 

self._forget_fixed = forget_fixed 

if processors is None: 

self._processors = [Processor()] 

else: 

self._processors = [] 

for p in processors: 

self.add_processor(p) 

def set_fixation_length(self, length): 

'''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 = length # 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 add_processor(self, processor): 

self.fixate_all() 

self._processors.append(processor) 

def insert_trace(self, trace): 

logger.debug('Received a trace: %s' % trace) 

for p in self._processors: 

for tr in p.process(trace): 

self._insert_trace(tr) 

def _insert_trace(self, trace): 

buf = self._append_to_buffer(trace) 

nslc = trace.nslc_id 

if buf is None: # create new buffer trace 

if nslc in self._buffers: 

self._fixate(self._buffers[nslc]) 

trbuf = trace.copy() 

buf = pile.MemTracesFile(None, [trbuf]) 

self.add_file(buf) 

self._buffers[nslc] = buf 

buf.recursive_grow_update([trace]) 

trbuf = buf.get_traces()[0] 

if self._fixation_length is not None: 

if trbuf.tmax - trbuf.tmin > self._fixation_length: 

self._fixate(buf) 

del self._buffers[nslc] 

def _append_to_buffer(self, trace): 

'''Try to append the trace to the active buffer traces. 

Returns the current buffer trace or None if unsuccessful. 

''' 

nslc = trace.nslc_id 

if nslc not in self._buffers: 

return None 

buf = self._buffers[nslc] 

trbuf = buf.get_traces()[0] 

if (abs((trbuf.tmax+trbuf.deltat) - trace.tmin) < 1.0e-1*trbuf.deltat 

and trbuf.ydata.dtype == trace.ydata.dtype 

and trbuf.deltat == trace.deltat): 

trbuf.append(trace.ydata) 

return buf 

return None 

def fixate_all(self): 

for buf in self._buffers.values(): 

self._fixate(buf) 

self._buffers = {} 

def _fixate(self, buf): 

if self._path: 

trbuf = buf.get_traces()[0] 

fns = io.save([trbuf], self._path, format=self._format) 

self.remove_file(buf) 

if not self._forget_fixed: 

self.load_files( 

fns, show_progress=False, fileformat=self._format) 

def drop_older(self, tmax, delete_disk_files=False): 

self.drop( 

condition=lambda file: file.tmax < tmax, 

delete_disk_files=delete_disk_files) 

def drop(self, condition, delete_disk_files=False): 

candidates = [] 

buffers = list(self._buffers.values()) 

for file in self.iter_files(): 

if condition(file) and file not in buffers: 

candidates.append(file) 

self.remove_files(candidates) 

if delete_disk_files: 

for file in candidates: 

if file.abspath and os.path.exists(file.abspath): 

os.unlink(file.abspath) 

def __del__(self): 

self.fixate_all()