speed comparison

mlair/helpers/filter.py

@@ -265,13 +265,13 @@ class ClimateFIRFilter:
         time_axis = filter_input_data.coords["datetime"]
         # 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)}
-        with TimeTracking():
+                  "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():
-        #     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="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)
 
         # plot
         if self.plot_path is not None:
@@ -413,8 +413,8 @@ def fir_filter_vectorized(data, time_stamp=None, fs=1, order=5, cutoff_low=None,
         pd_date = pd.to_datetime(time_stamp)
         if pd_date.day == 1 and pd_date.month in [1, 7]:
             logging.info(time_stamp)
-    sel = ~np.isnan(data)
-    res = np.empty_like(data)
+    # sel = ~np.isnan(data)
+    # res = np.empty_like(data)
     if h is None:
         cutoff = []
         if cutoff_low is not None:
@@ -431,9 +431,29 @@ def fir_filter_vectorized(data, time_stamp=None, fs=1, order=5, cutoff_low=None,
             raise ValueError("Please provide either cutoff_low or cutoff_high.")
         h = signal.firwin(order, cutoff, pass_zero=filter_type, fs=fs, window=window)
     if causal:
-        y = signal.lfilter(h, 1., data[sel])
+        # y = signal.lfilter(h, 1., data[sel])
+        y = signal.lfilter(h, 1., data)
     else:
         padlen = padlen if padlen is not None else 3 * len(h)
+        # 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
+    # return res
+    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):
+    """Expects numpy array."""
+    sel = ~np.isnan(data)
+    res = np.empty_like(data)
     if sum(sel) <= padlen:
         y = np.empty_like(data[sel])
     else: