pyrocko/src/hamster_pile.py

# 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()