You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
contrib-snufflings/notch.py

113 lines
3.1 KiB
Python

from pyrocko import trace
import math
import numpy as num
from pyrocko.gui.snuffling import Param, Snuffling, Switch
import scipy.signal as S
import scipy.stats as SS
def detrend_data(x_axis,data):
slope,offset,dummy,dummy2,dumy3 = SS.linregress(x_axis, data)
detrended_data = S.detrend(data)
return detrended_data, slope, offset
def retrend_data(x_axis,data,slope, offset):
retrended_data = x_axis * slope + data + offset
return retrended_data
class Notch(Snuffling):
'''Apply gaussian notch filter to data in Snuffler
The filter is defined by a central frequency and a window width. If possible,
harmonics are removed as well. Filtering carried out in frequency domain.
Data trace is detrended, FFT calculated, respective notch coefficients
calculated and multiplied onto the FFT section. Then back-transform to time
domain, adding the underlying trendline.'''
def setup(self):
'''Customization of Notch Snuffling.'''
self.set_name('Notch Filter')
self.add_parameter(Param('Center Frequency [Hz]',
'centerfreq', 50., 0.001, 1000.0, low_is_none=False))
self.add_parameter(Param('Notch width (FWHM) [Hz]',
'notchwidth', 1., 0.001, 1000.0, low_is_none=False))
self.add_parameter(Switch('Also filter harmonics',
'filter_harmonics', False))
self.set_have_post_process_hook(True)
self.FFF = {}
def call(self):
# snuffling parameters have changed
# release cached frequency responses
self.FFF = {}
# released cached processed traces in snuffler internals
viewer = self.get_viewer()
viewer.old_processed_traces = None
def post_process_hook(self, traces):
FFF = self.FFF
for tr in traces:
indi = num.arange(tr.data_len(), dtype=num.float)
detr_data, m, b = detrend_data(indi, tr.get_ydata())
tr.set_ydata(detr_data)
freqs, FT_data = tr.spectrum(pad_to_pow2=True, tfade=None)
fff_length = FT_data.size
key = (tr.deltat, fff_length)
if key not in FFF:
FFF[key] = num.ones((fff_length), dtype=num.complex)
fcenter = self.centerfreq
fwidth = self.notchwidth
iharm = 1
while iharm*fcenter - fwidth*2 < freqs[-1]:
FFF[key] *= GaussNotch(iharm*fcenter, fwidth).evaluate(freqs)
iharm += 1
if not self.filter_harmonics:
break
filtered_data = num.fft.irfft(FT_data*FFF[key])[:tr.data_len()]
retrended_data = retrend_data(indi, filtered_data, m, b)
tr.set_ydata(retrended_data)
return traces
class GaussNotch(trace.FrequencyResponse):
def __init__(self, center, fwhm):
self.center = center
self.fwhm = fwhm
def evaluate(self, freqs):
denom = self.fwhm / (2.*math.sqrt(math.log(2.)))
return 1.0 - num.exp(-((freqs-self.center)/denom)**2)
def __snufflings__():
return [ Notch()]