<commit_before>import numpy as np

from scipy.interpolate import interp1d, splrep, splev
import scipy.interpolate as interpolate
from scipy.interpolate import interp1d
import scipy.ndimage as nd
import scipy.stats
import scipy

from scipy.signal import butter, bessel, lfilter

import matplotlib as mpl
from matplotlib import mlab

__all__ = ['baseline', 'normalize', 'normalizeAndBaseline', \
           'findLevels', 'findLevels1d', 'findLevelsNd', \
           'boxcar', 'smooth', 'lowess', \
           'fir_filter', 'butter_bandpass_filter', 'psd', 'specgram',\
           'mask_deviations', 'baseline_splines']

def baseline(A, baseRange, baseAxis):
    """Baseline a numpy array using a given range over a specfied axis.

    :param A: numpy array of arbitrary dimension
    :param baseRange: list of 2 numbers, specifying the range over which to compute the average for baselining.
    :param baseAxis: the axis of the np to baseline over
    :returns: basedlined array
    """

    shape = A.shape
    # make a slice to take the mean in the right dimension
    baseSlice = [slice(baseRange[0], baseRange[1]) if a is baseAxis else slice(None) for a in range(len(shape))]
    base = np.mean(A[baseSlice], axis=baseAxis)

    # make a slice to pad the numbers to make the broadcasting work
    try:
        return A - base
    except:
        subSlice=[slice(None) if axis == base.shape[0] else None for axis in shape]
        return A - base[subSlice]
    finally:
        pass
    
def normalize(A, normRange, normAxis):
    """Normalize a numpy array using a given range over a specfied axis.

    :param A: numpy array of arbitrary dimension
    :param normRange: list of 2 numbers, specifying the range over which to compute the average for normalization.
    :param normAxis: the axis of the np to normalize over
    :returns: normalized array
    """

    shape = A.shape
    # make a slice to take the mean in the right dimension
    # slice(None) effectively means ':', or all the elements
    normSlice = [slice(normRange[0], normRange[1]) if a is normAxis else slice(None) for a in range(len(shape))]
    norm = np.mean(A[normSlice], axis=normAxis)

    # make a slice to pad the numbers to make the broadcasting work
    # again, slice(None) means ':' and None means an empty dimension (note difference!)

    try:
        return A/norm
    except:
        subSlice=[slice(None) if axis == norm.shape[0] else None for axis in shape]
        return A / norm[subSlice]
    finally:
        pass
    
def normalizeAndBaseline(A, baseRange, baseAxis):
    """Normalize, then baseline a numpy array over a given range and on a specfied axis.

    Calls normalize, and then baseline

    :param A: numpy array of arbitrary dimension
    :param baseRange: list of 2 numbers, specifying the range over which to compute the average for baselining and normalization
    :param baseAxis: the axis of the np to baseline and normalize over
    :returns: normalized, baselined array
    """
    return baseline(normalize(A, baseRange, baseAxis), baseRange, baseAxis)

# -------------------- LEVEL FINDING ROUTINES------------------------------------------

def findLevels(A, level, mode='rising', boxWidth=0, rangeSubset=None):
    """Function to find level crossings in an 1d numpy array.  Based on the Igor
    function FindLevel. 

    Can find rising and/or falling crossings, control with the 'mode' paramter.

    Returns an numpy array of all crossings found and the number of crossings

    :param A: 1d numpy array
    :param level: floating point to search for in A
    :param mode: optional string: mode specfication. one of 'rising', 'falling' or 'both'
    :param boxWidth: optional int for local boxcar smoothing
    :param rangeSubset: optional list of ints to limit the search
    :returns: tuple, a numpy array of level crossings and the number of crossings
    """
    assert mode in ('rising', 'falling', 'both'), 'traceManip.findLevels: Unknown mode \'%s\'' % mode

    if boxWidth is not 0:
        A = np.convolve(A, np.array([1]*boxWidth)/float(boxWidth))

    crossings = np.diff(np.sign(A-level), axis=0)
    
    if mode is 'rising':
        rising_points = np.where(crossings > 0)
        return rising_points[0], len(rising_points[0])
    elif mode is 'falling':
        falling_points = np.where(crossings < 0)
        return falling_points[0], len(falling_points[0])
    else:
        all_crossing_points = np.where(np.abs(crossings) > 0)
        return all_crossing_points, len(all_crossing_points)

def findLevels1d(A, level, mode='rising', boxWidth=0):
    return findLevelsNd(A, level, mode=mode, axis=0, boxWidth=boxWidth)

def findLevelsNd(A, level, mode='rising', axis=0, boxWidth=0):
    """Function to find level crossings in an Nd numpy array. 

    Can find rising and/or falling crossings, control with the 'mode' paramter.

    Returns a binary array of level crossings, with true elements right AFTER a crossing.

    NOTE THAT THIS RETURNS DIFFERENT VALUES THAN findLevels().  if you want to get a list of
    locations where the crossings occurs, then use the following syntax:

    levels = findLevelsNd(array, level)
    level_crossings_locations = levels.nonzero()
    number_of_level_crossings = len(level_crossing_locations[0])

    Often, the crossings are noisy.  You can use np.diff() and findLevelsNd() again to help yourself out.

    :param A: 1d numpy array
    :param level: floating point to search for in A
    :param mode: optional string: mode specfication. one of 'rising', 'falling' or 'both'
    :param axis: optional integer, specifies dimension
    :param boxWidth: optional int for local boxcar smoothing
    :returns: binary array of level crossing locations
    """
    assert mode in ('rising', 'falling', 'both'), 'traceManip.findLevels: Unknown mode \'%s\'' % mode

    if boxWidth is not 0:
        A = nd.convolve1d(A, np.array([1]*boxWidth)/float(boxWidth), axis=axis)

    crossings = np.diff(np.sign(A-level), axis=axis)
    
    if mode is 'rising':
        return crossings>0
    elif mode is 'falling':
        return crossings<0
    else:
        return np.abs(crossings>0)

# -------------------- SMOOTHING ROUTINES------------------------------------------

def lowess(x, y, f=2./3., iters=3): 
    """Lowess smoother: Robust locally weighted regression. 
    The lowess function fits a nonparametric regression curve to a scatterplot. 
    The arrays x and y contain an equal number of elements; each pair 
    (x[i], y[i]) defines a data point in the scatterplot. The function returns 
    the estimated (smooth) values of y. 

    The smoothing span is given by f. A larger value for f will result in a 
    smoother curve. The number of robustifying iterations is given by iter. The 
    function will run faster with a smaller number of iterations. 

    x and y should be numpy float arrays of equal length.  The return value is 
    also a numpy float array of that length. 

    :param: x - x values
    :param: y - y values
    :param: f - span (region size to weight for smoothing)
    :param: iters - number of times to apply smoothing
    :returns: yest - a smoothed version of y
    """ 
    n = len(x) 
    r = int(np.ceil(f*n)) 
    h = [np.sort(abs(x-x[i]))[r] for i in range(n)] 
    w = np.clip(abs(([x]-np.transpose([x]))/h),0.0,1.0) 
    w = 1-w*w*w 
    w = w*w*w 
    yest = np.zeros(n) 
    delta = np.ones(n) 
    for iteration in range(iters): 
        for i in xrange(n): 
            weights = delta * w[:,i] 
            weights_mul_x = weights * x 
            b1 = np.dot(weights,y) 
            b2 = np.dot(weights_mul_x,y) 
            A11 = sum(weights) 
            A12 = sum(weights_mul_x) 
            A21 = A12 
            A22 = np.dot(weights_mul_x,x) 
            determinant = A11*A22 - A12*A21 
            beta1 = (A22*b1-A12*b2) / determinant 
            beta2 = (A11*b2-A21*b1) / determinant 
            yest[i] = beta1 + beta2*x[i] 
        residuals = y-yest 
        s = np.median(abs(residuals)) 
        delta[:] = np.clip(residuals/(6*s),-1,1) 
        delta[:] = 1-delta*delta 
        delta[:] = delta*delta 
    return yest 

def boxcar(A, boxWidth=3, axis=1):
    """Boxcar smoothes a matrix of 1d traces with a boxcar of a specified width.
    Does this by convolving the traces with another flat array.

    :param A: a 1d (time) or 2d numpy array (traces by time)
    :param boxWidth: an optional int, specifying the width of the boxcar
    :returns: 2d numpy array, a smoothed version of A.
    """
    if A.ndim is 1:
        axis=0
    
    return nd.convolve1d(A, np.array([1]*boxWidth)/float(boxWidth), axis=axis)

def smooth(A, window_len=11, window='hanning'):
    """smooth the data (1D numpy array) using a window with requested size.
    
    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal 
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.

    ripped off from http://www.scipy.org/Cookbook/SignalSmooth
  
    :param A: the input signal 
    :param window_len: the dimension of the smoothing window; should be an odd integer
    :param window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' flat window will produce a moving average smoothing.
    :returns: the smoothed signal
        
    See also: numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
    scipy.signal.lfilter
    """

    if A.ndim != 1:
        raise ValueError, "smooth only accepts 1 dimension arrays."

    if A.size < window_len:
        raise ValueError, "Input vector needs to be bigger than window size."

    if window_len<3:
        return A

    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError, "Window is one of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"

    s=np.r_[A[window_len-1:0:-1],A,A[-1:-window_len:-1]]
    if window == 'flat': #moving average
        w=np.ones(window_len,'d')
    else:
        w=eval('np.'+window+'(window_len)')

    y=np.convolve(w/w.sum(),s,mode='valid')
    return y

# -------------------- Filtering Routines ------------------------------------------
# from http://code.google.com/p/python-neural-analysis-scripts/source/browse/trunk/scripts/Filtering/Fir.py

def fir_filter(sig, sampling_freq, critical_freq, kernel_window = 'hamming', taps = 101, kind = 'band', **kwargs):
    """This is a wrapper around scipy.signal.lfilter(), which is for finite impulse response filters.<commit_msg>Update traceRoutines.py<commit_after>import numpy as np

from scipy.interpolate import interp1d, splrep, splev
import scipy.interpolate as interpolate
from scipy.interpolate import interp1d
import scipy.ndimage as nd
import scipy.stats
import scipy

from scipy.signal import butter, bessel, lfilter

import matplotlib as mpl
from matplotlib import mlab

__all__ = ['baseline', 'normalize', 'normalizeAndBaseline', \
           'findLevels', 'findLevels1d', 'findLevelsNd', \
           'boxcar', 'smooth', 'lowess', \
           'fir_filter', 'butter_bandpass_filter', 'psd', 'specgram',\
           'mask_deviations', 'baseline_splines']

def baseline(A, baseRange, baseAxis):
    """Baseline a numpy array using a given range over a specfied axis.

    :param A: numpy array of arbitrary dimension
    :param baseRange: list of 2 numbers, specifying the range over which to compute the average for baselining.
    :param baseAxis: the axis of the np to baseline over
    :returns: basedlined array
    """

    shape = A.shape
    # make a slice to take the mean in the right dimension
    baseSlice = [slice(baseRange[0], baseRange[1]) if a is baseAxis else slice(None) for a in range(len(shape))]
    base = np.mean(A[baseSlice], axis=baseAxis)

    # make a slice to pad the numbers to make the broadcasting work
    try:
        return A - base
    except:
        subSlice=[slice(None) if axis == base.shape[0] else None for axis in shape]
        return A - base[subSlice]
    finally:
        pass
    
def normalize(A, normRange, normAxis):
    """Normalize a numpy array using a given range over a specfied axis.

    :param A: numpy array of arbitrary dimension
    :param normRange: list of 2 numbers, specifying the range over which to compute the average for normalization.
    :param normAxis: the axis of the np to normalize over
    :returns: normalized array
    """

    shape = A.shape
    # make a slice to take the mean in the right dimension
    # slice(None) effectively means ':', or all the elements
    normSlice = [slice(normRange[0], normRange[1]) if a is normAxis else slice(None) for a in range(len(shape))]
    norm = np.mean(A[normSlice], axis=normAxis)

    # make a slice to pad the numbers to make the broadcasting work
    # again, slice(None) means ':' and None means an empty dimension (note difference!)

    try:
        return A/norm
    except:
        subSlice=[slice(None) if axis == norm.shape[0] else None for axis in shape]
        return A / norm[subSlice]
    finally:
        pass
    
def normalizeAndBaseline(A, baseRange, baseAxis):
    """Normalize, then baseline a numpy array over a given range and on a specfied axis.

    Calls normalize, and then baseline

    :param A: numpy array of arbitrary dimension
    :param baseRange: list of 2 numbers, specifying the range over which to compute the average for baselining and normalization
    :param baseAxis: the axis of the np to baseline and normalize over
    :returns: normalized, baselined array
    """
    return baseline(normalize(A, baseRange, baseAxis), baseRange, baseAxis)

# -------------------- LEVEL FINDING ROUTINES------------------------------------------

def findLevels(A, level, mode='rising', boxWidth=0, rangeSubset=None):
    """Function to find level crossings in an 1d numpy array.  Based on the Igor
    function FindLevel. 

    Can find rising and/or falling crossings, control with the 'mode' paramter.

    Returns an numpy array of all crossings found and the number of crossings

    :param A: 1d numpy array
    :param level: floating point to search for in A
    :param mode: optional string: mode specfication. one of 'rising', 'falling' or 'both'
    :param boxWidth: optional int for local boxcar smoothing
    :param rangeSubset: optional list of ints to limit the search
    :returns: tuple, a numpy array of level crossings and the number of crossings
    """
    assert mode in ('rising', 'falling', 'both'), 'traceManip.findLevels: Unknown mode \'%s\'' % mode

    if boxWidth is not 0:
        A = np.convolve(A, np.array([1]*boxWidth)/float(boxWidth))

    crossings = np.diff(np.sign(A-level), axis=0)
    
    if mode is 'rising':
        rising_points = np.where(crossings > 0)
        return rising_points[0], len(rising_points[0])
    elif mode is 'falling':
        falling_points = np.where(crossings < 0)
        return falling_points[0], len(falling_points[0])
    else:
        all_crossing_points = np.where(np.abs(crossings) > 0)
        return all_crossing_points, len(all_crossing_points)

def findLevels1d(A, level, mode='rising', boxWidth=0):
    return findLevelsNd(A, level, mode=mode, axis=0, boxWidth=boxWidth)

def findLevelsNd(A, level, mode='rising', axis=0, boxWidth=0):
    """Function to find level crossings in an Nd numpy array. 

    Can find rising and/or falling crossings, control with the 'mode' paramter.

    Returns a binary array of level crossings, with true elements right AFTER a crossing.

    NOTE THAT THIS RETURNS DIFFERENT VALUES THAN findLevels().  if you want to get a list of
    locations where the crossings occurs, then use the following syntax:

    levels = findLevelsNd(array, level)
    level_crossings_locations = levels.nonzero()
    number_of_level_crossings = len(level_crossing_locations[0])

    Often, the crossings are noisy.  You can use np.diff() and findLevelsNd() again to help yourself out.

    :param A: 1d numpy array
    :param level: floating point to search for in A
    :param mode: optional string: mode specfication. one of 'rising', 'falling' or 'both'
    :param axis: optional integer, specifies dimension
    :param boxWidth: optional int for local boxcar smoothing
    :returns: binary array of level crossing locations
    """
    assert mode in ('rising', 'falling', 'both'), 'traceManip.findLevels: Unknown mode \'%s\'' % mode

    if boxWidth is not 0:
        A = nd.convolve1d(A, np.array([1]*boxWidth)/float(boxWidth), axis=axis)

    crossings = np.diff(np.sign(A-level), axis=axis)
    
    if mode is 'rising':
        return crossings>0
    elif mode is 'falling':
        return crossings<0
    else:
        return np.abs(crossings>0)

# -------------------- SMOOTHING ROUTINES------------------------------------------

def lowess(x, y, f=2./3., iters=3): 
    """Lowess smoother: Robust locally weighted regression. 
    The lowess function fits a nonparametric regression curve to a scatterplot. 
    The arrays x and y contain an equal number of elements; each pair 
    (x[i], y[i]) defines a data point in the scatterplot. The function returns 
    the estimated (smooth) values of y. 

    The smoothing span is given by f. A larger value for f will result in a 
    smoother curve. The number of robustifying iterations is given by iter. The 
    function will run faster with a smaller number of iterations. 

    x and y should be numpy float arrays of equal length.  The return value is 
    also a numpy float array of that length. 

    :param: x - x values
    :param: y - y values
    :param: f - span (region size to weight for smoothing)
    :param: iters - number of times to apply smoothing
    :returns: yest - a smoothed version of y
    """ 
    n = len(x) 
    r = int(np.ceil(f*n)) 
    h = [np.sort(abs(x-x[i]))[r] for i in range(n)] 
    w = np.clip(abs(([x]-np.transpose([x]))/h),0.0,1.0) 
    w = 1-w*w*w 
    w = w*w*w 
    yest = np.zeros(n) 
    delta = np.ones(n) 
    for iteration in range(iters): 
        for i in xrange(n): 
            weights = delta * w[:,i] 
            weights_mul_x = weights * x 
            b1 = np.dot(weights,y) 
            b2 = np.dot(weights_mul_x,y) 
            A11 = sum(weights) 
            A12 = sum(weights_mul_x) 
            A21 = A12 
            A22 = np.dot(weights_mul_x,x) 
            determinant = A11*A22 - A12*A21 
            beta1 = (A22*b1-A12*b2) / determinant 
            beta2 = (A11*b2-A21*b1) / determinant 
            yest[i] = beta1 + beta2*x[i] 
        residuals = y-yest 
        s = np.median(abs(residuals)) 
        delta[:] = np.clip(residuals/(6*s),-1,1) 
        delta[:] = 1-delta*delta 
        delta[:] = delta*delta 
    return yest 

def boxcar(A, boxWidth=3, axis=1):
    """Boxcar smoothes a matrix of 1d traces with a boxcar of a specified width.
    Does this by convolving the traces with another flat array.

    :param A: a 1d (time) or 2d numpy array (traces by time)
    :param boxWidth: an optional int, specifying the width of the boxcar
    :returns: 2d numpy array, a smoothed version of A.
    """
    if A.ndim is 1:
        axis=0
    
    return nd.convolve1d(A, np.array([1]*boxWidth)/float(boxWidth), axis=axis)

def smooth(A, window_len=11, window='hanning'):
    """smooth the data (1D numpy array) using a window with requested size.
    
    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal 
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.

    ripped off from http://www.scipy.org/Cookbook/SignalSmooth
  
    :param A: the input signal 
    :param window_len: the dimension of the smoothing window; should be an odd integer
    :param window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' flat window will produce a moving average smoothing.
    :returns: the smoothed signal
        
    See also: numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
    scipy.signal.lfilter
    """

    if A.ndim != 1:
        raise ValueError, "smooth only accepts 1 dimension arrays."

    if A.size < window_len:
        raise ValueError, "Input vector needs to be bigger than window size."

    if window_len<3:
        return A

    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError, "Window is one of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'"

    s=np.r_[A[window_len-1:0:-1],A,A[-1:-window_len:-1]]
    if window == 'flat': #moving average
        w=np.ones(window_len,'d')
    else:
        w=eval('np.'+window+'(window_len)')

    y=np.convolve(w/w.sum(),s,mode='same')
    return y

# -------------------- Filtering Routines ------------------------------------------
# from http://code.google.com/p/python-neural-analysis-scripts/source/browse/trunk/scripts/Filtering/Fir.py

def fir_filter(sig, sampling_freq, critical_freq, kernel_window = 'hamming', taps = 101, kind = 'band', **kwargs):
    """This is a wrapper around scipy.signal.lfilter(), which is for finite impulse response filters.
{'fraction_non_alphanumeric': 0.11250644299704794, 'fraction_numerical': 0.010308795276697437, 'mean_word_length': 2.841944194419442, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2938088', 'n_tokens_mistral': 6772, 'n_tokens_neox': 6304, 'n_words': 2416}