source: SHX/trunk/SeismicHandler/patches/ObsPy.py @ 529

# -*- coding: utf-8 -*-
#
# Copyright (C) 2008-2011 Marcus Walther (walther@szgrf.bgr.de)
#
# This file is part of Seismic Handler eXtended (SHX)
# Full details can be found at project website http://www.seismic-handler.org/
#
# SHX is free software; you can redistribute it and/or modify
# it under the terms of the GNU LESSER GENERAL PUBLIC LICENSE as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# SHX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with SHX (see license.txt); if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA

"""
This modules defines extensions to obspy. They are patched into the
original obspy modules at a very basic level in SeismicHandler/__init__.py

This module is named ObsPy.py to prevent import loops.
"""

import numpy as np
import math
from obspy.core import Stream, Trace
from obspy.core.util import AttribDict
from obspy.sh.core import fromUTCDateTime
from SeismicHandler.basics.messages import logMessage, setStatus

__all__ = []

def _streamAutoMerge(self, *args, **kwargs):
    """
    Call obspy's Stream merge method to fill up gaps using the latest pre-gap
    sample value. This destroys the spectra of the data but it's common.
    """

    # Change value data type to float32 as in former versions of SH.
    # This will be changed in later versions.
    for t in self.traces:
        if t.data.dtype != np.dtype('float32'):
            t.data = t.data.astype('float32')

    if self.traces:
        self.merge(method=1, fill_value='latest')

def streamAutoMergeInit(self, *args, **kwargs):
    self.__shx_init__(*args, **kwargs)
    _streamAutoMerge(self)

def streamAutoMergeAppend(self, *args, **kwargs):
    self.__shx_append(*args, **kwargs)
    _streamAutoMerge(self, *args, **kwargs)

def streamAutoMergeExtend(self, *args, **kwargs):
    self.__shx_extend(*args, **kwargs)
    _streamAutoMerge(self, *args, **kwargs)

def streamAutoMergeInsert(self, *args, **kwargs):
    self.__shx_insert(*args, **kwargs)
    _streamAutoMerge(self, *args, **kwargs)

def streamAutoMergeAdd(self, *args, **kwargs):
    self.__shx_add__(*args, **kwargs)
    _streamAutoMerge(self, *args, **kwargs)

def tracePrepareDataForImage(self, width, height, timescale=None, zoom=1, norm="window"):
    """
    Preparing trace data for fast plotting. Using numpy's minmax feature - this
    will not just downsample the data!

    width : total width of plotting window in pixels
    height : hight of plotting window in pixels
    zoom : global zoom factor (y extension)
    timescale : extension of current timescale (tuple of min and max)
    norm : norm method: "window", "total" or value (effects y extension)
    """
    logMessage('debug.patches', "Preparing plotting data...", self.id)

    if not timescale:
        # use trace's total dimensions
        windowstart, windowend = (self.stats.starttime, self.stats.endtime)
    else:
        windowstart, windowend = timescale

    # get slice for time window
    pt = self.slice(windowstart, windowend)

    duration_total = windowend - windowstart
    duration_trace = pt.stats.endtime - pt.stats.starttime

    # width of the total screen that is covered by trace data (pixels)
    pixel_width = duration_trace * width / duration_total

    # remember values for caching
    CacheTraceID = \
              (windowstart, windowend, "%.5f" % pixel_width, height, zoom, norm)
    try:
        cachedX = CacheTraceID[:3] == self.shx.CacheTraceID[:3]
        cachedY = cachedX and CacheTraceID[3:] == self.shx.CacheTraceID[3:]
    except:
        cachedX = False
        cachedY = False

    if not cachedX:
        # save new ID
        self.shx.CacheTraceID = CacheTraceID

        # data points per pixel
        npts = pt.stats.npts
        dpp = int(npts / pixel_width)

        # get copy of data
        self.shx.PlotData = pt.data.copy()
    else:
        logMessage('debug.patches.devel',
                                    "using cachedX data for", CacheTraceID[:3])

    # use minmax approach if more than 4 data points per pixel
    # self.shx_plotdata will be changed
    if not cachedX and dpp >= 4:
        # do integer division since array shape values cannot be floats
        dimx = npts // dpp
        covered = dpp * dimx

        # get data for reshape operation
        data = self.shx.PlotData[:covered]

        # extra treatment for remaining values
        remaining = self.shx.PlotData[covered:]

        # reshape data and get min/max values
        data.shape = dimx, dpp
        _min = data.min(axis=1)
        _max = data.max(axis=1)
        # combine data
        joined = np.append([_min], [_max], axis=0).T.flatten()
        # handle remaining
        if len(remaining):
            joined = np.append(joined, [remaining.min(), remaining.max()])

        logMessage(
            'debug.patches.devel',
             "pixel_width", pixel_width, "minmax", "npts", npts, "dpp",
             dpp, "dimx", dimx, "covered", covered, "len_rest",
             len(remaining), "data_len", len(data), "joined", len(joined),
             "width", width
        )

        # At this stage the x-transformed data can be cached! If pixel_width
        # doesn't change, the data can be reused to save time.
        self.shx.PlotData = joined

    logMessage('debug.patches.devel',
                                  width, height, norm, type(self.shx.PlotData))

    if cachedY:
        logMessage('debug.patches.devel',
                                    "using cachedY data for", CacheTraceID[3:])
        return
    
    # renew cache id
    self.shx.CacheTraceID = CacheTraceID

    # get basis for normation
    if hasattr(norm, "lower"):
        norm = norm.lower()
        if norm == "window":
            norm = abs(pt.max())
        elif norm == "total":
            norm = self.max()
        else:
            raise ValueError("Invalid input for normation!")

    # Calculate y extension, leads to raw "image" data.
    # Use a copy to save time if only y changes occur (e.g. zooming).
    # Apply normation, calibration and total height factor
    y = self.shx.PlotData.copy()
    y *= self.stats.calib / norm * height / 2 * zoom

#    print self.shxPlotData.min(), self.shxPlotData.max(), y.min(), y.max()

    # factor for x stretching (trade-off between pixel_width and array shape)
    stretch = pixel_width / len(y)
    width = int(math.ceil(pixel_width))

    # calculate line data
    # do not rely on framework axis' management -> shift/flip y-axis
    offset = height // 2
    lines = []
    oldx, oldy = 0, -y[0] + offset

    for i in xrange(1, len(y)):
        newx = i * stretch
        newy = -y[i] + offset
        lines.append([oldx, oldy, newx, newy])
        oldx, oldy = newx, newy

    self.shx.ImageData = lines
    self.shx.PlotPixels = pixel_width

def traceAddShxInit(self, *args, **kwargs):
    self.__shx_init__(*args, **kwargs)

    # container for SHX releated data
    self.shx = AttribDict()

def traceGetStatusInfo(self, name):
    """
    Mapping of Seismic Handler trace information entries to ObsPy status info.
    """
    name = name.lower()

    # identical name
    if name in self.stats:
        return self.stats[name]

    # mapping
    if name == "length":
        return self.stats.npts

    if name == "chan1":
        return self.stats.channel[0]

    if name == "chan2":
        return self.stats.channel[1]

    if name == "comp":
        return self.stats.channel[2]

    if name == "start":
        return fromUTCDateTime(self.stats.starttime)

    # finally look into self defined entries
    if name in self.shx:
        return self.shx[name]

    raise NameError

# monkey patching obspy stream class: traces are merge automatically if stream
# get's altered.
Stream.__shx_init__, Stream.__init__ = Stream.__init__, streamAutoMergeInit
#Stream.__shx_append, Stream.append = Stream.append, streamAutoMergeAppend
#Stream.__shx_extend, Stream.extend = Stream.extend, streamAutoMergeExtend
#Stream.__shx_insert, Stream.insert = Stream.insert, streamAutoMergeInsert
#Stream.__shx_add__, Stream.__add__ = Stream.__add__, streamAutoMergeAdd

# monkey patching obspy trace class: add shx container and method for fast
# plotting
Trace.__shx_init__, Trace.__init__ = Trace.__init__, traceAddShxInit
Trace._shxPrepareImageData = tracePrepareDataForImage
Trace._shxInfo = traceGetStatusInfo