new try with convolve

mlair/helpers/filter.py

@@ -3,6 +3,7 @@ import warnings
 from typing import Union, Callable
 import logging
 import os
+import time
 
 import datetime
 import numpy as np
@@ -266,12 +267,19 @@ class ClimateFIRFilter:
         # apply vectorized fir filter along the tmp dimension
         kwargs = {"fs": fs, "cutoff_high": cutoff_high, "order": order,
                   "causal": False, "padlen": int(min(padlen_factor, 1) * length), "h": h}
-        with TimeTracking(name="numpy_vec"):
-            filt = fir_filter_numpy_vectorized(filter_input_data, var_dim, new_dim, kwargs)
-        with TimeTracking(name="xr_apply_ufunc"):
-            filt = xr.apply_ufunc(fir_filter_vectorized, filter_input_data, time_axis,
-                                  input_core_dims=[[new_dim], []], output_core_dims=[[new_dim]], vectorize=True,
-                                  kwargs=kwargs)
+        # with TimeTracking(name="numpy_vec"):
+        #     filt = fir_filter_numpy_vectorized(filter_input_data, var_dim, new_dim, kwargs)
+        # with TimeTracking(name="xr_apply_ufunc"):
+        #     filt = xr.apply_ufunc(fir_filter_vectorized, filter_input_data, time_axis,
+        #                           input_core_dims=[[new_dim], []], output_core_dims=[[new_dim]], vectorize=True,
+        #                           kwargs=kwargs)
+        with TimeTracking(name="convolve"):
+            slicer = slice(int(-(length - 1) / 2), int((length - 1) / 2))
+            filt = xr.apply_ufunc(fir_filter_convolve_vectorized, filter_input_data.sel(window=slicer),
+                                  input_core_dims=[["window"]],
+                                  output_core_dims=[["window"]],
+                                  vectorize=True,
+                                  kwargs={"h": h})
 
         # plot
         if self.plot_path is not None:
@@ -404,6 +412,10 @@ def fir_filter_numpy_vectorized(filter_input_data, var_dim, new_dim, kwargs):
     return filt_np
 
 
+def fir_filter_convolve_vectorized(data, h):
+    return signal.convolve(data, h, mode='same', method="direct") / sum(h)
+
+
 def fir_filter_vectorized(data, time_stamp=None, fs=1, order=5, cutoff_low=None, cutoff_high=None, window="hamming",
                           h=None,
                           causal=True,
@@ -435,10 +447,10 @@ def fir_filter_vectorized(data, time_stamp=None, fs=1, order=5, cutoff_low=None,
         y = signal.lfilter(h, 1., data)
     else:
         padlen = padlen if padlen is not None else 3 * len(h)
+        # print(sum(sel))
         # if sum(sel) <= padlen:
         #     y = np.empty_like(data[sel])
         # else:
-        #     with TimeTracking():
         #     y = signal.filtfilt(h, 1., data[sel], padlen=padlen)
         y = signal.filtfilt(h, 1., data, padlen=padlen)
     # res[sel] = y
@@ -446,20 +458,10 @@ def fir_filter_vectorized(data, time_stamp=None, fs=1, order=5, cutoff_low=None,
     return y
 
 
-def fir_filter_vectorized_short(data, time_stamp=None, fs=1, order=5, cutoff_low=None, cutoff_high=None,
-                                window="hamming",
-                                h=None,
-                                causal=True,
-                                padlen=None):
+def fir_filter_vectorized_short(data, h=None, padlen=None):
     """Expects numpy array."""
-    sel = ~np.isnan(data)
-    res = np.empty_like(data)
-    if sum(sel) <= padlen:
-        y = np.empty_like(data[sel])
-    else:
-        y = signal.filtfilt(h, 1., data[sel], padlen=padlen)
-    res[sel] = y
-    return res
+    y = signal.filtfilt(h, 1., data, padlen=padlen)
+    return y
 
 
 class KolmogorovZurbenkoBaseClass: