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

# -*- 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 import UTCDateTime

__all__ = [
    "patched",
]

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.
    """
    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=None, 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)
    """
    if not timescale:
        # use trace's total dimensions
        windowstart, windowend = (self.stats.starttime, self.stats.endtime)
    else:
        windowstart, windowend = timescale

    # XXX
    if not height:
        height = 1000

    # 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.shxCacheTraceID[:3]
        cachedY = cachedX and CacheTraceID[3:] == self.shxCacheTraceID[3:]
    except:
        cachedX = False
        cachedY = False

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

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

        # get copy of data
        self.shxPlotData = pt.data.copy()
    else:
        print "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.shxPlotData[:covered]

        # extra treatment for remaining values
        remaining = self.shxPlotData[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()])

        print "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.shxPlotData = joined

    print width, height, norm, type(self.shxPlotData)

    if cachedY:
        print "using cachedY data for", CacheTraceID[3:]
        return

    # 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.shxPlotData.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.shxImageData = lines

def __traceX(parameters):
    """
    Only the x axis is adopted here. Zoom, norm etc. regarding the y extension
    is done ...
    """
    pass

# 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

# add method for faster plotting
Trace.shxPrepareImageData = tracePrepareDataForImage

# just an indicator
patched = True