Speed up get_stokes() by using scipy convolve

nu-radio · Nov 21, 2024 · ebfead7 · ebfead7
1 parent f77a38b
commit ebfead7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/NuRadioReco/utilities/trace_utilities.py b/NuRadioReco/utilities/trace_utilities.py
@@ -179,10 +179,10 @@ def get_stokes(trace_u, trace_v, window_samples=128, squeeze=True):
     elif window_samples is None:
         window_samples = len(h1)
 
-    stokes = np.array([[
-            np.sum(i[j:j+window_samples]) for j in range(len(h1)-window_samples+1)]
-        for i in stokes]
-    ) / window_samples
+    stokes = np.asarray([
+        scipy.signal.convolve(i, np.ones(window_samples), mode='valid', method='fft') for i in stokes
+    ])
+    stokes /= window_samples
 
     if squeeze:
         return np.squeeze(stokes)