filter_sgolay2d.py 5.49 KB
title = "Savitzky-Golay Filter"
tip = "applies Savitzky-Golay filter to data"
onein = True

import numpy as np
import scipy.signal as sps

from guidata.dataset.datatypes import DataSet
from guidata.dataset.dataitems import (IntItem, FloatArrayItem, StringItem,
                                       ChoiceItem, FloatItem, DictItem,
                                       BoolItem)
from guiqwt.config import _

def sgolay2d(z, window_size, order, derivative=None):
    """Smooth (and optionally differentiate) data with a Savitzky-Golay filter.
    The Savitzky-Golay filter removes high frequency noise from data.
    It has the advantage of preserving the original shape and
    features of the signal better than other types of filtering
    approaches, such as moving averages techniques.
    Parameters
    """
    # number of terms in the polynomial expression
    n_terms = ( order + 1 ) * ( order + 2)  / 2.0
    
    if window_size % 2 == 0:
       raise ValueError('window_size must be odd')
    
    if window_size**2 < n_terms:
       raise ValueError('order is too high for the window size')

    half_size = window_size // 2
     
    # exponents of the polynomial. 
    # p(x,y) = a0 + a1*x + a2*y + a3*x^2 + a4*y^2 + a5*x*y + ... 
    # this line gives a list of two item tuple. Each tuple contains 
    # the exponents of the k-th term. First element of tuple is for x
    # second element for y.
    # Ex. exps = [(0,0), (1,0), (0,1), (2,0), (1,1), (0,2), ...]
    exps = [ (k-n, n) for k in range(order+1) for n in range(k+1) ]
    # coordinates of points
    ind = np.arange(-half_size, half_size+1, dtype=np.float64)
    dx = np.repeat(ind, window_size)
    dy = np.tile(ind, [window_size, 1]).reshape(window_size**2,)

    # build matrix of system of equation
    A = np.empty( (window_size**2, len(exps)) )
    for i, exp in enumerate( exps ):
        A[:,i] = (dx**exp[0]) * (dy**exp[1])

    # pad input array with appropriate values at the four borders
    new_shape = z.shape[0] + 2*half_size, z.shape[1] + 2*half_size
    Z = np.zeros( (new_shape) )
    # top band
    band = z[0, :]
    Z[:half_size, half_size:-half_size] =  band -  np.abs( np.flipud( z[1:half_size+1, :] ) - band )
    # bottom band
    band = z[-1, :]
    Z[-half_size:, half_size:-half_size] = band  + np.abs( np.flipud( z[-half_size-1:-1, :] )  -band )
    # left band
    band = np.tile( z[:,0].reshape(-1,1), [1,half_size])
    Z[half_size:-half_size, :half_size] = band - np.abs( np.fliplr( z[:, 1:half_size+1] ) - band )
    # right band
    band = np.tile( z[:,-1].reshape(-1,1), [1,half_size] )
    Z[half_size:-half_size, -half_size:] =  band + np.abs( np.fliplr( z[:, -half_size-1:-1] ) - band )
    # central band
    Z[half_size:-half_size, half_size:-half_size] = z
    # top left corner
    band = z[0,0]
    Z[:half_size,:half_size] = band - np.abs( np.flipud(np.fliplr(z[1:half_size+1,1:half_size+1]) ) - band )
    # bottom right corner
    band = z[-1,-1]
    Z[-half_size:,-half_size:] = band + np.abs( np.flipud(np.fliplr(z[-half_size-1:-1,-half_size-1:-1]) ) - band )
    # top right corner
    band = Z[half_size,-half_size:]
    Z[:half_size,-half_size:] = band - np.abs( np.flipud(Z[half_size+1:2*half_size+1,-half_size:]) - band )
    # bottom left corner
    band = Z[-half_size:,half_size].reshape(-1,1)
    Z[-half_size:,:half_size] = band - np.abs( np.fliplr(Z[-half_size:, half_size+1:2*half_size+1]) - band )

    # solve system and convolve
    if type(derivative) == type(None):
       m = np.linalg.pinv(A)[0].reshape((window_size, -1))
       return sps.fftconvolve(Z, m, mode='valid')
    elif derivative == 'col':
       c = np.linalg.pinv(A)[1].reshape((window_size, -1))
       return sps.fftconvolve(Z, -c, mode='valid')
    elif derivative == 'row':
       r = np.linalg.pinv(A)[2].reshape((window_size, -1))
       return sps.fftconvolve(Z, -r, mode='valid')
    elif derivative == 'both':
       c = np.linalg.pinv(A)[1].reshape((window_size, -1))
       r = np.linalg.pinv(A)[2].reshape((window_size, -1))
       return sps.fftconvolve(Z, -r, mode='valid'), sps.fftconvolve(Z, -c, mode='valid')


class NOD3_App:

    def __init__(self, parent):
        self.parent = parent
        self.parent.activateWindow()

    def compute_app(self, **args):
        class FuncParam(DataSet):
            window = IntItem('Window_size:', default=5)
            order = IntItem('Order:', default=3)
            derivative = ChoiceItem('Deriv:', ((None, "None"), ("both", "Both")), default=None)
            #derivative = ChoiceItem('Deriv:', ((None, "None"), ("col", "Col"), ("row", "Row"), ("both", "Both")), default=None)
        name = title.replace(" ", "")
        if args == {}:
           param = FuncParam(_(title), "Fits n-order polynomials to surface")
        else:
           param = self.parent.ScriptParameter(name, args)
        # if no parameter needed set param to None. activate next line
        #param = None
        self.name = name
        self.parent.compute_11(name, lambda m, p: self.function(m, p), param, onein)

    def function(self, m, p):
        mask = np.isnan(m.data)
        data = np.where(mask, np.nanmin(m.data), m.data)
        if p.derivative == "both":
           data = sgolay2d(data, p.window, p.order, derivative=p.derivative)
           data1 = np.where(mask, np.nan, data[0]) 
           data2 = np.where(mask, np.nan, data[1]) 
           m.data = np.sqrt(data1**2 + data2**2)
           return m, p
        else:
           data = sgolay2d(data, p.window, p.order)
           m.data = np.where(mask, np.nan, data) 
           return m, p