source: fact/tools/pyscripts/pyfact/drs_spikes.py@ 13566

#!/usr/bin/python -tt
#
# Werner Lustermann
# ETH Zurich
#
import numpy as np

import fir_filter as fir

class DRSSpikes(object):
    """ remove spikes (single or double false readings) from DRS4 data
    Strategy:
    * filter the data, removing the signal, thus spike(s) are clearly visible
    * search single and double spikes
    * replace the spike by a value derived from the neighbors
    
    """
    
    def __init__(self, threshold=7., 
                 single_pattern=np.array( [-0.5, 1.0, -0.5]) ,
                 double_pattern=np.array([-1., 1., 1., -1.]), 
                 user_action=lambda candidates, singles, doubles, data, ind: None,
                 debug = False):
        """ initialize spike filter 
        template_single: template of a single slice spike
        template_double: template of a two slice spike

        """

        self.threshold = threshold
        self.single_pattern = single_pattern * threshold
        self.double_pattern = double_pattern * threshold
        
        self.remove_signal = fir.RemoveSignal()

        self.user_action = user_action
        self.debug = debug

    def __call__(self, data):

        self.row, self.col = data.shape
        indicator = self.remove_signal(data)
        a = indicator.flatten()
        singles = []
        doubles = []

        # a spike in the first or last channel is considered as a filter artefact
        candidates = np.where(a[1:-2] > self.threshold)
        # candidates = np.where(a[1:1022] > self.threshold)
        cc = candidates[0]
        #print 'cc: ', cc
        #: find single spikes
        p = self.single_pattern * np.sign( self.single_pattern )
        for i, can in enumerate( zip(a[cc], a[cc+1], a[cc+2]) ):
            #print 'can : p', can, p
            can = can * np.sign(self.single_pattern)
            if np.all(can > p):
                singles.append(cc[i])

        #: find double spikes
        p = self.double_pattern * np.sign( self.double_pattern )
        for i, can in enumerate( zip(a[cc], a[cc+1], a[cc+2], a[cc+3]) ):
            #print 'data:  ', [data[0,cc[i]+k] for k in range(3)]
            #print 'can : p', can, p
            can = can * np.sign(self.double_pattern)
            if np.all(can > p):
                doubles.append(cc[i])

        if self.debug:
            print 'singles: ', singles
            print 'doubles: ', doubles

        self.user_action(cc, singles, doubles, data, a)

        data = self.remove_single_spikes(singles, data)
        data = self.remove_double_spikes(doubles, data)
        return data

    def remove_single_spikes(self, singles, data):
        data = data.flatten() 
        for spike in singles:
            data[spike] = (data[spike-1] + data[spike+1]) / 2.
        return data.reshape(self.row, self.col)
    
    def remove_double_spikes(self, doubles, data):
        data = data.flatten() 
        for spike in doubles:
            data[spike:spike+2] = (data[spike-1] + data[spike+2]) / 2.
        return data.reshape(self.row, self.col)

class DRSSpikes_2D(object):
    """ remove spikes (single or double false readings) from DRS4 data
    Strategy:
    * filter the data, removing the signal, thus spike(s) are clearly visible
    * search single and double spikes
    * replace the spike by a value derived from the neighbors
    
    """
    
    def __init__(self, threshold=7., 
                 single_pattern=np.array( [-0.5, 1.0, -0.5]) ,
                 double_pattern=np.array([-1., 1., 1., -1.]), 
                 user_action=lambda self, data: None,
                 debug = False):
        """ initialize spike filter 
        template_single: template of a single slice spike
        template_double: template of a two slice spike
        """

        self.threshold = threshold
        self.single_pattern = single_pattern * threshold
        self.double_pattern = double_pattern * threshold
        
        self.remove_signal = fir.RemoveSignal()

        self.user_action = user_action
        self.debug = debug

    def __call__(self, data):
        # shortcuts
        row, col = data.shape
        thr = self.threshold
        
        # these are: lists, which will contain positiones of spikes
        # lets see if this is feasible
        self.singles = []
        self.doubles = []
        singles = self.singles
        doubles = self.doubles

        # indi means indicator, i.e. a filter output, which indicates, where spikes
        # are positioned in the unfiltered data
        # indi is delayed w.r.t. data by 1
        self.indi = self.remove_signal(data)
        indi = self.indi

        # cand (candidates), is a tuple of two equal length np.arrays
        # each pair of array elements can be understood as coordinates, pointing out
        # where the condition was fullfilled in indi. 
        # e.g. in pixel = cand[0][0] around slice = cand[1][0]
        # there is probably a spike
        #
        # for double spikes, two neighboring slices fulfill the condition
        # which lead to something like: 
        # cand =( array([ ... 3, 3, ... ]) , array([ ... 102, 103 ...]) )
        cand = np.where(indi[:,1:-2] > thr)
        self.cand = cand
        # in order to verify, that the candidate is really a single or double
        # spike, we compare the spike with a 3 or 4 slices pattern.
        # therefor we want to slice out 4 slices out of indi, where ever the 
        # condition was fullfilled
        #
        # note: since indi was sliced in the np.where statement, 
        # the resulting cand coordinates are reduced by 1 in the slice coordinate
        # this is actually what we want, since a spike has a distinctive low-high-low 
        # pattern in the indicator. So we *want* the indicator slices to be shifted 1
        # to the left
        # and in addition, by pure chance, the coordinates in cand[1] point directly 
        # to the real spike in data, since indi was delayed by one anyway.
        cand_slices = np.empty( (len(cand[0]), 4), dtype=np.int )
        self.cand_slices = cand_slices
        for i in range(4):
            cand_slices[i] = indi[ (cand[0], cand[1]+i )
            
        # search for single spikes
        sp = self.single_pattern * np.sign( self.single_pattern )
        for i, can in enumerate(cand_slices[:-1]):
            can *= np.sign(sp)
            if np.all( can > sp):
                singles.append( (cand[0][i],can[1][i]) )

        # I guess in principle it is possible, that a candidate looks like a 
        # single and like a double, ... nut with the current patters
        # but in case one changes the patterns ... then it might happen.
        # In addition the treatment of double spikes is maybed not smart:
        #   In case both parts of a double spike fulfill the 1st condition
        #       only the first candidate will fulfill the 2nd condition
        #   In case only the first part fulfilled the 1st conition, then
        #       we are fine
        #   In case only the second part triggered the first time, then
        #       we sliced out the wrong piece and it wouldn't fulfull the 2nd anyway.
        # 
        #   This means, in case there are neighboring hits in cand,
        #   The 2nd neighbor will never fulfill the 2nd condition.

        # search for double spikes
        dp = self.double_pattern * np.sign( self.double_pattern )
        for i, can in enumerate( cand_slices ):
            can *= np.sign(dp)
            if np.all(can > dp):
                doubles.append( (cand[0][i],can[1][i]) )

        self.user_action(self, data)

        data = self.remove_single_spikes(singles, data)
        data = self.remove_double_spikes(doubles, data)
        return data

    def remove_single_spikes(self, singles, data):
        for spike in singles:
            data[spike[0],spike[1]] = (data[spike[0],spike[1]-1] + data[spike[0],spike[1]+1]) / 2.
        return data
    
    def remove_double_spikes(self, doubles, data):
        for spike in doubles:
            data[spike[0],spike[1]:spike[1]+2] = (data[spike[0],spike[1]-1] + data[spike[0],spike[1]+2]) / 2.
        return data



def _test():
  
    a = np.ones((3,12)) * 3.
    a[0,3] = 7.
    a[1,7] = 14.
    a[1,8] = 14.
    a[2,4] = 50.
    a[2,5] = 45.
    a[2,8] = 20.
    
    print a 

    SpikeRemover = DRSSpikes(3., debug=True)
    print 'single spike pattern ', SpikeRemover.single_pattern
    print 'double spike pattern ', SpikeRemover.double_pattern
    afilt = SpikeRemover(a)
    print afilt
    
if __name__ == '__main__':
    """ test the class """
    _test()