script/sig_process.py

@@ -1,17 +1,18 @@
 # Signal processing functions
 import copy
-from mathutils import * 
+import numpy as np
+import scipy.optimize
 
 
 # Noise evaluation ####
-def noise_evaluation(noise):
+def noise_evaluation(noise: np.ndarray):
     """
     Returns offset and standard deviation of the noise signal.
     :param noise: noise signal
     :return: Tuple of (noise_offset, noise_sigma))
     """
 
-    return mean(noise), stdev(noise)
+    return np.mean(noise), np.std(noise)
 
 
 # Processing of BLM signal ####
@@ -21,22 +22,20 @@ def blm_remove_spikes(x):
     :param x: input array
     :return: output array
     """
-
     x = copy.copy(x)
     if x.size > 5:  # Must have enough sample points
         d_x = x[1:] - x[:-1]
         maximum = x.max() - x.min()
-
         for i in range(x.size-3):
             if d_x[i+1] > 0.5 * maximum and d_x[i+2] < -0.5 * maximum:
                 # 1 point spike
                 x[i+2] = (x[i:i+5].sum() - x[i+2])/4
-
-            if d_x[i+1] > 0.5 * maximum and d_x[i+2] >= -0.5 * maximum:
+            if d_x[i+1] > 0.5 * maximum and d_x[i+2] >= -0.5 * maximum:                
                 # 2 point spikes
-                x[i+2] = (x[i:i+2].sum() + x[i+4])/3
-                x[i+3] = (x[i+1] + x[i+4:i+6].sum())/3
-
+                if i < x.size-4:    #TODO: FIX BY AG, CHECK 
+                    x[i+2] = (x[i:i+2].sum() + x[i+4])/3
+                if i < x.size-6:    #TODO: FIX BY AG, CHECK 
+                    x[i+3] = (x[i+1] + x[i+4:i+6].sum())/3
         # Handle edge points
         if d_x[-1] > 0.5 * maximum:
             x[-1] = x[-3:-1].sum() / 2
@@ -76,12 +75,12 @@ def motor_to_wire_cs(pos, type: str = 'u', center_pos: float = 0.0):
     # u: normal positions, like GARAGE, POSITION, or any wires under 45 degrees to pipe horizon
     # x: which crosses pipe vertically (scans in x coordinate)
     # y: which crosses pipe horizontally (scans in y coordinate)
-    w_scale = {'u': 1, 'x': -math.sqrt(2).sqrt(2), 'y': math.sqrt(2).sqrt(2)}
+    w_scale = {'u': 1, 'x': -np.sqrt(2), 'y': np.sqrt(2)}
 
     return (pos-center_pos)/w_scale[type]
 
 
-def remove_beam_jitter(pos, bpm1, bpm2, d_b1_w=1, d_w_b2=1):
+def remove_beam_jitter(pos: np.ndarray, bpm1: np.ndarray, bpm2: np.ndarray, d_b1_w=1, d_w_b2=1):
     """
     This is a temporary solution which works only for first WSC on SwissFEL.
     !!TODO: In future this method must implement correction using online model.
@@ -109,23 +108,18 @@ def profile_gauss_stats(x, y, off=None, amp=None, com=None, sigma=None):
         off = y.min()  # good enough starting point for offset
 
     if com is None:
-        #com = x[y.argmax()]
-        com = x[y.index(max(y))]
+        com = x[y.argmax()]
 
     if amp is None:
         amp = y.max() - off
 
     # For normalised gauss curve sigma=1/(amp*sqrt(2*pi))
     if sigma is None:
-        surface = integrate((y-off), xdata=x)
-        sigma = surface / (amp * math.sqrt(2 * math.pi))
+        surface = np.trapz((y-off), x=x)
+        sigma = surface / (amp * np.sqrt(2 * np.pi))
     try:
-        #popt, pcov = scipy.optimize.curve_fit(gauss_fn, x, y, p0=[off, amp, com, sigma], method='lm')
-        #popt[3] = abs(popt[3])  # sigma should be returned as positive
-        pars = fit_gaussian(y, x, start_point = [amp, com, sigma])
-        popt = [0.0, pars[0], pars[1], abs([pars[2])]
-        
-        
+        popt, pcov = scipy.optimize.curve_fit(gauss_fn, x, y, p0=[off, amp, com, sigma], method='lm')
+        popt[3] = abs(popt[3])  # sigma should be returned as positive
     except Exception as e:
         print("Gauss fitting not successful.\n" + str(e))
         popt = [off, amp, com, abs(sigma)]
@@ -134,7 +128,7 @@ def profile_gauss_stats(x, y, off=None, amp=None, com=None, sigma=None):
 
 
 def gauss_fn(x, a, b, c, d):
-    return a + b * math.exp(-(math.pow((x - c), 2) / (2 * math.pow(d, 2))))
+    return a + b * np.exp(-(np.power((x - c), 2) / (2 * np.power(d, 2))))
 
 
 def profile_rms_stats(x, y, noise_std=0, n_sigma=3.5):
@@ -221,7 +215,7 @@ def profile_rms_stats(x, y, noise_std=0, n_sigma=3.5):
 
 def com_rms(x, y):
     # Centre of mass and rms
-    com = sum([a*b for a,b in zip(x,y)]) / sum(y)
-    com2 = sum([a*a*b for a,b in zip(x,y)]) / sum(y) 
-    rms = math.sqrt(math.abs(com2 - com * com))
+    com = (x * y).sum() / y.sum()
+    com2 = (x * x * y).sum() / y.sum()
+    rms = np.sqrt(np.abs(com2 - com * com))
     return com, rms