pyscripts-common/signal_processing.py at master · bmotevalli/pyscripts-common · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from scipy.fftpack import fft
import numpy as np


def get_fft(t,x):

    """
    This function returns fast-fourier transform. fftpack in scipy is used for this purpose, however,
    numpy.fft.fft can also be used in similar manner.

    t:       time
    x:       signal

    fr:      frequency axis
    X_fr:    amplitude complex number
    amp:     amplitude
    angle:   phase angle
    """

    # Sampling Frequency
    Fs = 1 / (t[1] - t[0])

    # Generate Frequency Axis
    n = np.size(t)
    n_half = int(n / 2)
    fr = (Fs / 2) * np.linspace(0, 1, n_half)

    # Compute FFT
    X     = fft(x)
    X_fr  = (1 / n_half) * X[0:n_half]
    amp   = np.absolute(X_fr)
    angle = np.angle(X_fr)

    return fr, X_fr, amp, angle

def get_ifft(X):

    """
    X:            amplitude complex number
    """
    return np.fft.ifft(X).real * np.size(X) / 2


def fft_comps_cutoff(t, x, cutoff = 0.01):

    fr, X_fr, amp, angle = get_fft(t, x)

    # filter
    amp_norm = amp / amp.max()
    mask = amp_norm > cutoff

    return fr[mask], X_fr[mask], amp[mask], angle[mask]

def construct_time_signal_from_comps(t_vec, freq, amplitude, angle):

    y_vec = np.zeros_like(t_vec)
    for fr, amp, a in zip(freq, amplitude, angle):
        # y_vec += amp * np.sin(2 * np.pi * fr * t_vec + a)
        y_vec += amp * np.sin(2 * np.pi * fr * t_vec) # + a)

    return y_vec