diff --git a/src/wela/butterworth_filter.py b/src/wela/butterworth_filter.py index a9e100b98343e76032f2acbce78479fa718209fe..3c8a903546283a9062e3139f9b84578e03a36bf7 100644 --- a/src/wela/butterworth_filter.py +++ b/src/wela/butterworth_filter.py @@ -1,8 +1,9 @@ import numpy as np from scipy import signal +import pandas as pd -def butterworth_filter(timeseries, params=None): +def butterworth_filter(data, params=None): """Apply Butterworth filter to time series arranged in rows.""" # second-order-sections output # by default, using a digital filter @@ -20,18 +21,18 @@ def butterworth_filter(timeseries, params=None): fs=params["sampling_freq"], output="sos", ) + # fill NaNs between data points + data = pd.DataFrame(data).interpolate(axis=1).to_numpy() # subtract time series by mean - timeseries_norm = timeseries - np.nanmean(timeseries, axis=1).reshape( - timeseries.shape[0], 1 - ) - if np.any(np.isnan(timeseries_norm)): + data_norm = data - np.nanmean(data, axis=1).reshape(data.shape[0], 1) + if np.any(np.isnan(data_norm)): # filter one at a time because of NaNs at different times - filtered_timeseries = np.nan * np.ones(timeseries_norm.shape) - for i in range(timeseries_norm.shape[0]): - data = timeseries_norm[i, :] - filtered = signal.sosfiltfilt(sos, data[~np.isnan(data)]) - filtered_timeseries[i, ~np.isnan(data)] = filtered + filtered_timeseries = np.nan * np.ones(data_norm.shape) + for i in range(data_norm.shape[0]): + s_data = data_norm[i, :] + filtered = signal.sosfiltfilt(sos, s_data[~np.isnan(s_data)]) + filtered_timeseries[i, ~np.isnan(s_data)] = filtered else: # filter all at once - filtered_timeseries = signal.sosfiltfilt(sos, timeseries_norm, axis=1) + filtered_timeseries = signal.sosfiltfilt(sos, data_norm, axis=1) return filtered_timeseries