## Ce programme affiche un signal à partir de son spectre, puis
## lui applique un filtrage et affiche le signal filtré
## Armel MARTIN - Septembre 2021
## distribué sous licence CC-by-nc-sa
import matplotlib.pyplot as plt
import numpy as np
import random

## définition du spectre via la Transformée de Fourier (TF)
## ou la Série de Fourier (SF) si suite de fréquences multiples entiers de
## mm fréquence fondamentale

## Construction de la fonction du temps. 
def signal(tt, ff, tfs):
    """Calcule la fonction s(t) pour les instants contenus dans le vecteur tt
    à partir du spectre donné sous la forme du vecteur des fréquences ff et
    du vecteur tfs dont le module est l'amplitude de l'harmonique associée
    à chaque fréquence (en cosinus) et l'argument sa phase à l'origine:
            tfs[k] = A_k * exp ( i*phi_k ) pour la fréquence ff[k].
    """
    n = len(ff) # Nb de composantes spectrales (composante continue comprise)
    ss = np.zeros(len(tt))  # Initialisation de la somme 
    for k in range(n): # boucle d'addition des composantes spectrales présentes
        amp_k, phi_k = np.abs(tfs[k]), np.angle(tfs[k]) # amplitude, phase
        ss += amp_k * np.cos( 2*np.pi*ff[k]*tt + phi_k )
    return ss # vecteur des valeurs de s(t) pour les instants t dans tt

## Spectres
def spectre_2_comp(f1,A1,phi1,f2,A2,phi2):
    """Créée un spectre tfs sous la forme d'un vecteur de valeurs
    tfs[k] = A_k * exp ( i*phi_k ) et d'un vecteur ff des fréquences associées.
    Signal constitué d'une superposition d'une composante lente et d'une rapide:
        s(t) = A1 cos(2pi*f1*t + phi1) + A2 cos(2pi*f2*t + phi2)
    """
    ff = np.array([f1,f2])
    tfs = np.array([ A1*np.exp(1j*phi1) , A2*np.exp(1j*phi2) ])
    return ff, tfs

def spectre_creneau(f,A0,Sm,p):
    """Créée un spectre tfs sous la forme d'un vecteur de valeurs
    tfs[k] = A_k * exp ( i*phi_k ) et d'un vecteur ff des fréquences associées.
    Signal créneau pair (ou impair) de fondamentale f, de valeur moyenne A0,
    d'amplitude Sm, tronqué à l'harmonique de rang n = 2*p-1
    (il y a donc p harmoniques non nulles).
    """
    n = 2*p-1 # rang de la dernière harmonique, n+1 composantes avec la moyenne
    ff = np.arange(n+1) * f # série de Fourier (multiples de la fondamentale)
    tfs = np.zeros(n+1, dtype=complex) # initialisation par des zéros (type complexe !)
    tfs[0] = A0 # composante continue = valeur moyenne
    for i in range(p):
        k = 2*i+1 # rang de l'harmonique présente
        tfs[k] = (-1)**i / k # forme en cosinus, Paire
        #tfs[k] = 1 / k * np.exp(1j * (-np.pi/2)) # forme en sinus, Impaire
    tfs = tfs * 4*Sm/np.pi
    return ff, tfs

def spectre_triangle(f,A0,Sm,p):
    """Créée un spectre tfs sous la forme d'un vecteur de valeurs
    tfs[k] = A_k * exp ( i*phi_k ) et d'un vecteur ff des fréquences associées.
    Signal triangulaire pair (ou impair) de fondamentale f, de valeur moyenne A0,
    d'amplitude Sm, tronqué à l'harmonique de rang n = 2*p-1
    (il y a donc p harmoniques non nulles).
    """
    n = 2*p-1 # rang de la dernière harmonique, n+1 composantes avec la moyenne
    ff = np.arange(n+1) * f # série de Fourier (multiples de la fondamentale)
    tfs = np.zeros(n+1, dtype=complex) # initialisation par des zéros (type complexe !)
    tfs[0] = A0 # composante continue = valeur moyenne
    for i in range(p):
        k = 2*i+1 # rang de l'harmonique présente
        tfs[k] = 1 / k**2 # forme en cosinus, Paire
        #tfs[k] = (-1)**i / k**2 * np.exp(1j * (-np.pi/2)) # forme en sinus, Impaire
    tfs = tfs * 8*Sm/np.pi**2
    return ff, tfs

def spectre_bruit_colore(f,p,Amp,n):
    """Créée un spectre tfs sous la forme d'un vecteur de valeurs
    tfs[k] = A_k * exp ( i*phi_k ) et d'un vecteur ff des fréquences associées.
    Signal aléatoire correspondant à un bruit coloré, i.e.:
        A_k = Amp / (k*f)**p
    p = 0 -> bruit blanc
    p = 1 -> bruit rose
    p = 2 -> bruit rouge
    p = -1 -> bruit bleu
    p = -2 -> bruit violet
    ATTENTION: limiter la fréquence maximale à f_n = fe/2 = N/2T au max
    (T = durée simulation, N = nombre d'échantillons temporels)
    pour respecter le critère de Shannon et éviter le repliement spectral.
    Par exemple: f = 1/T et n = int(N/2)
    """
    ff = np.arange(n+1) * f # série de Fourier (multiples de la fondamentale)
    tfs = np.zeros(n+1, dtype=complex) # initialisation par des zéros (type complexe !)
    tfs[0] = 0 # composante continue = valeur moyenne
    for k in range(n+1):
        phi_k = random.uniform(-np.pi,np.pi) # phase aléatoire
        tfs[k] = 1/k**p * np.exp( 1j*phi_k ) 
    tfs = tfs * Amp
    return ff, tfs

## Graphes
def graphe_spectre(spectre):
    plt.figure(1) # ouverture de la figure 1
    plt.subplot(1,2,1)  # Graphe Spectre en amplitude
    plt.title("Spectre en amplitude")  # Titre
    plt.xlabel("f(Hz)")  # Légende des x
    plt.ylabel("Amplitude (SI)")  # Légende des y
    plt.bar(spectre[0],np.abs(spectre[1]), width=0.2)  # Tracé du spectre en amplitude
    plt.subplot(1,2,2)  # Graphe Spectre en phase
    plt.title("Spectre en phase")  # Titre
    plt.xlabel("f(Hz)")  # Légende des x
    plt.ylabel("rad")  # Légende des y
    plt.bar(spectre[0],np.angle(spectre[1]), width=0.2)  # Tracé du spectre en phase
    #plt.show()
    return 

def graphe_signal(tt,ss,style='-k',epaisseur=1,etiquette="s(t)"):
    plt.figure(2) # ouverture de la figure 1
    plt.title("Signal")  # Titre
    plt.xlabel("t (s)")  # Légende des x
    plt.ylabel("s(t)")  # Légende des y
    plt.plot(tt, ss, style, linewidth=epaisseur, label=etiquette)  # Tracé temporel
    plt.grid()
    plt.legend()
    #plt.show()
    return 

def FT_pbas_1erO(f0,H0,ff):
    xx = ff / f0 # fréquence réduite
    return H0 / ( 1 + 1j * xx ) # forme canonique

def FT_pbas_2ndO(f0,Q,H0,ff):
    xx = ff / f0 # fréquence réduite
    return H0 / ( 1 + 1j * xx / Q - xx**2) # forme canonique

def FT_phaut_1erO(f0,H0,ff):
    xx = ff / f0 # fréquence réduite
    return H0 * 1j * xx / ( 1 + 1j * xx ) # forme canonique

def FT_phaut_2ndO(f0,Q,H0,ff):
    xx = ff / f0 # fréquence réduite
    return H0 * (-xx**2) / ( 1 + 1j * xx / Q - xx**2) # forme canonique

def FT_pbande_2ndO(f0,Q,H0,ff):
    xx = ff / f0 # fréquence réduite
    return H0 * 1j * xx / Q / ( 1 + 1j * xx / Q - xx**2) # forme canonique

## Simulations
## Spectre
#spectre = spectre_2_comp(0.2,1.,0.,10.,0.1,0.)
spectre = spectre_creneau(0.5,1.,2.,30)
#spectre = spectre_triangle(1.,1.,2.,10)
#spectre = spectre_bruit_colore(1.,0,1.,int(len(instants)/2))
graphe_spectre(spectre)

## Signal 
instants = np.linspace(0,5,500)
ff, tfs = spectre
s = signal(instants,ff,tfs)
graphe_signal(instants,s)
## Signal filtré
# Filtrage 1er Ordre
sf = signal(instants,ff,tfs*FT_pbas_1erO(0.5,1.,ff))
graphe_signal(instants,sf,'-c',2,"filtré PBas 1erO - fc = 0.5Hz")
sf = signal(instants,ff,tfs*FT_pbas_1erO(0.05,1.,ff))
graphe_signal(instants,sf,'-y',2,"filtré PBas 1erO - fc = 0.05Hz")
#sf = signal(instants,ff,tfs*FT_pbas_1erO(2,1.,ff))
#sf = signal(instants,ff,tfs*FT_phaut_1erO(5.,1.,ff))
#graphe_signal(instants,sf,'-c',2,"filtré PBas 1erO")
#graphe_signal(instants,sf,'-c',2,"filtré PHaut 1erO")
# Filtrage 2nd Ordre
#sf = signal(instants,ff,tfs*FT_pbas_2ndO(0.1,1/np.sqrt(2),1.,ff))
#sf = signal(instants,ff,tfs*FT_pbas_2ndO(5,1/np.sqrt(2),1.,ff))
#sf = signal(instants,ff,tfs*FT_phaut_2ndO(5,1/np.sqrt(2),1.,ff))
#graphe_signal(instants,sf,'-y',2,"filtré PBas 2ndO")
#graphe_signal(instants,sf,'-y',2,"filtré PHaut 2ndO")
#sf = signal(instants,ff,tfs*FT_pbande_2ndO(5,10,1.,ff))
#sf = signal(instants,ff,tfs*FT_phaut_2ndO(5,1/np.sqrt(2),1.,ff))
#graphe_signal(instants,sf,'-r',2,"filtré PBande 2ndO")
plt.show()