utils.py - TB065 -

tp/scripts/utils.py (16304B)
Line numbers
   1 import matplotlib.pyplot as plt
   2 from matplotlib.ticker import AutoMinorLocator
   3 from matplotlib.ticker import MultipleLocator
   4 from matplotlib.ticker import MaxNLocator
   5 from matplotlib.ticker import FuncFormatter
   6 from cycler import cycler
   7 
   8 import matplotlib
   9 import numpy as np
  10 import os
  11 
  12 from scipy.io import wavfile
  13 from scipy.fft import fft, ifft, fftfreq
  14 from scipy.signal import spectrogram
  15 
  16 from data import *
  17 
  18 matplotlib.rcParams['font.family'] = 'Inter'
  19 matplotlib.rcParams['font.size'] = 12
  20 matplotlib.rcParams['axes.prop_cycle'] = cycler(
  21         color=['#1f77b4', '#ff0000', '#ff5f1f', 'green'])
  22 matplotlib.use("TkAgg")
  23 
  24 def ticks_label_format(x, pos):
  25     # 3 decimales, se eliminan los ceros y puntos
  26     return f"{x:.3f}".rstrip("0").rstrip(".")
  27 
  28 def time_graph_multiple_data(x, y_arr, y_lab, t=0, dt=0, a=0, da=0, show=True):
  29     figure, axis = plt.subplots(figsize=(8, 4))
  30 
  31     for i, y in enumerate(y_arr):
  32         axis.plot(x, y, label=y_lab[i], alpha=0.75)
  33 
  34     axis.set(xlabel='Tiempo [s]', ylabel='Amplitud normalizada')
  35 
  36     axis.minorticks_on()
  37     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
  38     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
  39 
  40     # configuracion de ticks del eje x
  41     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
  42     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
  43 
  44     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
  45     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
  46 
  47     plt.tight_layout()
  48 
  49     # max 3 decimales
  50     axis.xaxis.set_major_formatter(FuncFormatter(ticks_label_format))
  51 
  52     axis.set_xlim([t, t+dt if dt > 0 else x[-1]])
  53     axis.set_ylim(-1.1, 1.1)
  54 
  55     # resaltado de parte de la señal (solo si a != 0)
  56     axis.axvspan(a, a+da, color='skyblue',
  57                  alpha=0 if a == 0 else 0.50,
  58                  label=f"Un periodo T={da}s" if da != 0 else "")
  59     axis.legend(loc='upper left')
  60 
  61     if show:
  62         plt.show()
  63 
  64     return figure, axis
  65 
  66 # todos deben la misma cantidad de elementos que el primero
  67 def time_plot_multiple(fs, data_arr, leg_arr, save_name="", t=0, dt=0, a=0, da=0, show=False):
  68     x = np.arange(len(data_arr[0])) / fs
  69     fig, ax = time_graph_multiple_data(x, data_arr, leg_arr, t, dt, a=a, da=da, show=show)
  70 
  71     if show == False:
  72         save_plot(fig, save_name)
  73 
  74     return fig, ax
  75 
  76 def time_graph_data(x, y, t=0, dt=0, a=0, da=0, ylim=[], show=True):
  77     figure, axis = plt.subplots(figsize=(8, 4))
  78 
  79     axis.plot(x, y, label='Señal de audio')
  80     axis.set(xlabel='Tiempo [s]', ylabel='Amplitud normalizada')
  81 
  82     axis.minorticks_on()
  83     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
  84     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
  85 
  86     # configuracion de ticks del eje x
  87     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
  88     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
  89 
  90     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
  91     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
  92 
  93     plt.tight_layout()
  94 
  95     # max 3 decimales
  96     axis.xaxis.set_major_formatter(FuncFormatter(ticks_label_format))
  97 
  98     axis.set_xlim(t, t+dt if dt > 0 else x[-1])
  99 
 100     if len(ylim) != 0:
 101         axis.set_ylim(ylim[0], ylim[1])
 102     else:
 103         axis.set_ylim(-1.1, 1.1)
 104 
 105     # resaltado de parte de la señal (solo si a != 0)
 106     axis.axvspan(a, a+da, color='skyblue',
 107                  alpha=0 if a == 0 else 0.50,
 108                  label=f"Un periodo T={da}s" if da != 0 else "")
 109     axis.legend(loc='upper left')
 110 
 111     if show:
 112         plt.show()
 113 
 114     return figure, axis
 115 
 116 def normalize(data):
 117     data = data.astype(np.float32)
 118     data /= np.max(np.abs(data))
 119     return data
 120 
 121 def time_plot(fs, data, save_name="", t=0, dt=0, a=0, da=0):
 122     show = True if save_name == "" else False
 123 
 124     # normaliza la amplitud dividiendo por el valor maximo del tipo de dato
 125     data = normalize(data)
 126 
 127     x = np.arange(len(data)) / fs
 128     fig, ax = time_graph_data(x, data, t, dt, a, da, show)
 129 
 130     if show == False:
 131         save_plot(fig, save_name)
 132 
 133     return fig, ax
 134 
 135 def save_plot(fig, name, save_dir="", overwrite=True):
 136     base_name = os.path.basename(name)
 137     file_name, ext = os.path.splitext(base_name.replace('.', '_'))
 138 
 139     if save_dir == "":
 140         save_dir = plot_dir
 141 
 142     if save_dir != "" and not save_dir.endswith('/'):
 143         save_dir += "/"
 144 
 145     file_path_no_ext = f'{save_dir}{file_name}'
 146     save_name = f'{file_path_no_ext}.png'
 147 
 148     if overwrite == False:
 149         if os.path.exists(save_name):
 150             i = 1
 151             while True:
 152                 new_save_name = f'{file_path_no_ext}_{i:02d}'
 153                 if not os.path.exists(f'{new_save_name}.png'):
 154                     save_name = f'{new_save_name}.png'
 155                     break
 156                 i += 1
 157 
 158     print(f'[LOG] Guardando figura `{save_name}`')
 159 
 160      # crea carpeta para plots
 161     os.makedirs(plot_dir, exist_ok=True)
 162     fig.savefig(save_name, dpi=100, bbox_inches="tight")
 163     plt.close(fig) # liberar memoria
 164 
 165 def save_to_wav(fs, data, save_name):
 166     # normalizar para prevenir clipping
 167     data = data / np.max(np.abs(data))
 168 
 169     # convertir a 16-bit PCM para WAV
 170     data_as_int16 = np.int16(data * 32767)
 171 
 172     # crea carpeta para wavs
 173     os.makedirs(out_dir, exist_ok=True)
 174 
 175     file_path = f'{out_dir}{save_name}'
 176     print(f'[LOG] Guardando audio: `{file_path}')
 177     wavfile.write(file_path, fs, data_as_int16)
 178 
 179 # frecuencia
 180 
 181 # data = [[fft], [freqs], [legends]]
 182 def freq_graph_multiple_data(data, f_min=0, f_max=0, y_min=0, y_max=0, show=True):
 183     fig, axis = plt.subplots(figsize=(8, 4))
 184 
 185     for i, (fft, freqs, label) in enumerate(data):
 186         # print(label)
 187         N = len(freqs)
 188         x = freqs[:N // 2]
 189         y = np.abs(fft[:N // 2])
 190         axis.plot(x, y, label=label, alpha=0.90, linewidth=((len(data)-i-1)*0.5 + 1.5))
 191 
 192     axis.set(xlabel='Frecuencia [Hz]', ylabel='Magnitud [dB]')
 193 
 194     axis.minorticks_on()
 195     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 196     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 197 
 198     # configuracion de ticks del eje x
 199     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
 200     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
 201 
 202     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 203     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
 204 
 205     plt.tight_layout()
 206 
 207     # max 3 decimales
 208     axis.xaxis.set_major_formatter(FuncFormatter(ticks_label_format))
 209     plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
 210 
 211     axis.set_xlim([f_min, f_max if f_max != 0 else 20000])
 212     axis.set_ylim([y_min, y_max if y_max != 0 else 1.05*max(y)])
 213 
 214     axis.legend(loc='upper right')
 215 
 216     if show:
 217         plt.show()
 218 
 219     return fig, axis
 220 
 221 def freq_graph_data_norm(x, y, x_min=0, x_max=0, y_min=0, y_max=0, show=True,
 222                     xlabel="", ylabel=""):
 223 
 224     fig, axis = plt.subplots(figsize=(8, 4))
 225 
 226     axis.plot(x, y)
 227 
 228     if xlabel == "":
 229         axis.set(xlabel='Frecuencia normalizada')
 230 
 231     if ylabel == "":
 232         axis.set(ylabel='Magnitud [dB]')
 233 
 234     axis.minorticks_on()
 235     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 236     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 237 
 238     # configuracion de ticks del eje x
 239     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
 240     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
 241 
 242     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 243     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
 244 
 245     plt.tight_layout()
 246 
 247     axis.set_xlim([x_min, x_max if x_max != 0 else 1])
 248     axis.set_ylim([y_min, y_max if y_max != 0 else 1.05*max(y)])
 249 
 250     if show:
 251         plt.show()
 252 
 253     return fig, axis
 254 
 255 def freq_graph_multiple_data_norm(x, y_arr, labels,
 256                                   x_min=0, x_max=0, y_min=0, y_max=0,
 257                                   show=True, xlabel="", ylabel=""):
 258 
 259     fig, axis = plt.subplots(figsize=(8, 4))
 260 
 261     for i, y in enumerate(y_arr):
 262         axis.plot(x, y, label=labels[i], alpha=0.75)
 263 
 264     if xlabel == "":
 265         axis.set(xlabel='Frecuencia normalizada')
 266 
 267     if ylabel == "":
 268         axis.set(ylabel='Magnitud [dB]')
 269 
 270     axis.minorticks_on()
 271     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 272     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 273 
 274     # configuracion de ticks del eje x
 275     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
 276     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
 277 
 278     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 279     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
 280 
 281     plt.tight_layout()
 282     axis.legend(loc='upper right')
 283 
 284     axis.set_xlim([x_min, x_max if x_max != 0 else 1])
 285     axis.set_ylim([y_min, y_max if y_max != 0 else 1.05*max(y)])
 286 
 287     if show:
 288         plt.show()
 289 
 290     return fig, axis
 291 
 292 
 293 def freq_graph_data(x, y, f_min=0, f_max=0, y_min=0, y_max=0, show=True,
 294                     xlabel="", ylabel=""):
 295     fig, axis = plt.subplots(figsize=(8, 4))
 296 
 297     axis.plot(x, y)
 298 
 299     if xlabel == "":
 300         axis.set(xlabel='Frecuencia [Hz]')
 301 
 302     if ylabel == "":
 303         axis.set(ylabel='Magnitud')
 304 
 305     axis.minorticks_on()
 306     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 307     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 308 
 309     # configuracion de ticks del eje x
 310     axis.xaxis.set_major_locator(MaxNLocator(nbins=5))
 311     axis.xaxis.set_minor_locator(AutoMinorLocator(5))
 312 
 313     axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 314     axis.yaxis.set_minor_locator(AutoMinorLocator(4))
 315 
 316     plt.tight_layout()
 317 
 318     # max 3 decimales
 319     axis.xaxis.set_major_formatter(FuncFormatter(ticks_label_format))
 320     plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
 321 
 322     axis.set_xlim([f_min, f_max if f_max != 0 else 20000])
 323     axis.set_ylim([y_min, y_max if y_max != 0 else 1.05*max(y)])
 324 
 325     if show:
 326         plt.show()
 327 
 328     return fig, axis
 329 
 330 def freq_compute_fft(fs, data, t=0, dt=0, N=0):
 331     i = 0
 332     di = fs*len(data)
 333     if t != 0 or dt != 0:
 334         i = int(t*fs)
 335         di = int((t+dt)*fs)
 336 
 337     interval_data = data[i:di]
 338 
 339     # puntos de la fft
 340     if N == 0:
 341         N = len(interval_data)
 342 
 343     interval_fft = fft(interval_data, N)
 344     interval_freqs = fftfreq(N, d=1/fs)
 345 
 346     return interval_fft, interval_freqs
 347 
 348 # hace la transformacion a frecuencias y pasa lo transformado a `freq_graph_data`
 349 def freq_plot(fs, data, f_min=0, f_max=0, y_min=0, y_max=0,
 350               t=0, dt=0, a=0, da=0, show=False, N=0, save_name="", save_dir=""):
 351 
 352     interval_fft, interval_freqs = freq_compute_fft(fs, data, t, dt, N=N)
 353     N = len(interval_fft)
 354 
 355     # se toma la parte positiva en ambos casos (primer parte del arreglo)
 356     x = interval_freqs[:N // 2]
 357     y = np.abs(interval_fft[:N // 2])
 358 
 359     fig, ax = freq_graph_data(x, y, f_min, f_max, y_min, y_max, show=show)
 360 
 361     if save_name != "":
 362         save_plot(fig, name=save_name, save_dir=save_dir)
 363 
 364     return fig, ax
 365 
 366 # frecuencia de muestreo comun
 367 # computa y grafica en una figura la fft the los datos en `data_arr`
 368 def freq_plot_multiple(fs, data_arr, leg_arr, save_name="",
 369                        f_min=0, f_max=0, y_min=0, y_max=0, t=0, dt=0, show=True):
 370 
 371     fft_freqs_arr = []
 372     for i, data in enumerate(data_arr):
 373         fft, freqs = freq_compute_fft(fs, data, t, dt)
 374         fft_freqs_arr.append([fft, freqs, leg_arr[i]])
 375 
 376     fig, axis = freq_graph_multiple_data(fft_freqs_arr, f_min, f_max, y_min, y_max, show)
 377 
 378     if save_name != "":
 379         save_plot(fig, save_name)
 380 
 381 def generate_spectrogram(fs, data, t=0, dt=0, N=1024, overlp=16, win='hamm'):
 382     if dt == 0:
 383         dt = (len(data)/fs)-t
 384 
 385     i = int(t*fs)
 386     di = int((t+dt)*fs)
 387     interval_data = data[i:di]
 388 
 389     # `nperseg`  tamaño de ventana (número de muestras por segmento)
 390     # `noverlap` cantidad de solapamiento entre ventanas
 391     f, time, Sxx = spectrogram(interval_data, fs=fs, nperseg=N, noverlap=overlp,
 392                                window=win)
 393 
 394     return f, time, Sxx
 395 
 396 def spectrogram_plot(fs, data, save_name="", save_dir="", t=0, dt=0, N=1024,
 397                     overlp=16, win='hamm', xlim=[], ylim=[],
 398                     shading='gouraud', cmap='viridis', show=False):
 399 
 400     f, time, Sxx = generate_spectrogram(fs, data, t, dt, N, overlp, win)
 401     fig, axis = plt.subplots(figsize=(8, 4))
 402 
 403     # plt.pcolormesh(time, f, Sxx**0.10)
 404     plt.pcolormesh(time, f, 10*np.log10(Sxx + 1e-12), shading=shading, cmap=cmap)
 405 
 406     plt.ylabel('Frecuencia [Hz]')
 407     plt.xlabel('Tiempo [s]')
 408 
 409     if len(xlim) != 0:
 410         plt.xlim(xlim)
 411 
 412     if len(ylim) != 0:
 413         plt.ylim(ylim)
 414     else:
 415         plt.ylim(1, 20000)
 416 
 417     if show == True:
 418         plt.show()
 419     else:
 420         if save_name != "":
 421             save_plot(fig, name=save_name, save_dir=save_dir, overwrite=True)
 422 
 423     return fig, axis
 424 
 425 def bode_plot(w, H, show=True):
 426     figure, axis = plt.subplots(figsize=(8, 4))
 427 
 428     axis.plot(w, 20*np.log10(np.abs(H)))
 429 
 430     axis.set(xlabel='Frecuencia [Hz]', ylabel='Magnitud [dB]')
 431     axis.minorticks_on()
 432     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 433     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 434     plt.tight_layout()
 435 
 436     axis.set_xlim(0.0, 20e3)
 437 
 438     axis.legend(loc='upper left')
 439 
 440     if show:
 441         plt.show()
 442 
 443     return figure, axis
 444 
 445 def freq_response_plot(w, H, phase, show=True, fc=20e3):
 446     fig, ax1 = plt.subplots(figsize=(8, 4))
 447 
 448     H_db = 20*np.log10(np.abs(H))
 449 
 450     line1, = ax1.plot(w, H_db)
 451     ax1.set(xlabel='Frecuencia [Hz]', ylabel='Magnitud [dB]')
 452     ax1.minorticks_on()
 453     ax1.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 454     ax1.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 455     ax1.set_xlim(0.0, fc)
 456 
 457     ax2 = ax1.twinx()
 458     line2, = ax2.plot(w, phase, color="tab:red")
 459     ax2.set_ylabel("Fase [grados]", color="black")
 460     ax2.tick_params(axis='y', labelcolor="black")
 461 
 462     # axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 463     ax1.yaxis.set_minor_locator(AutoMinorLocator(2))
 464 
 465     # Add ONE point
 466 
 467     # Find index closest to -3 dB
 468     idx = np.argmin(np.abs(H_db + 3))    # H_db = -3 => H_db +3 = 0
 469     w_3db = w[idx]
 470     H_3db = H_db[idx]
 471     line3 = ax1.scatter(w_3db, H_3db, color='tab:green', s=50, zorder=10)
 472 
 473     # nyquist
 474     w_nyquist = 2756.25
 475     idx = np.argmin(np.abs(w - w_nyquist))    # H_db = -3 => H_db +3 = 0
 476     H_nyquist = H_db[idx]
 477     line4 = ax1.scatter(w_nyquist, H_nyquist, color='tab:orange', s=50, zorder=10)
 478 
 479     ax1.legend([line1, line2, line3, line4],
 480                ["Magnitud [dB]",
 481                 "Fase [grados]",
 482                 r'-3dB $\approx$ %0.0f Hz'%w_3db,
 483                 r'Nyquist $\approx$ %0.0f Hz'%w_nyquist],
 484                loc='upper right')
 485 
 486     plt.tight_layout()
 487 
 488     if show:
 489         plt.show()
 490 
 491     return fig, ax1, ax2
 492 
 493 def dtime_plot(N, f, save_name="", legend="", n=0, dn=0, a=0, da=0):
 494     show = True if save_name == "" else False
 495 
 496     n = np.arange(N + 1)
 497 
 498     fig, axis = plt.subplots(figsize=(8,4))
 499     axis.set(xlabel='Tiempo discreto', ylabel='Amplitud')
 500 
 501     markerline, stemlines, baseline = axis.stem(
 502         n, f,
 503         markerfmt='o',     # tipo de marcador en la cabeza
 504         basefmt="k-",
 505     )
 506 
 507     markerline.set_markersize(2.0)
 508     stemlines.set_linewidth(0.35)
 509     baseline.set_linewidth(0.5)
 510 
 511     stemlines.set_zorder(2)
 512     markerline.set_zorder(3)
 513     baseline.set_zorder(1)
 514 
 515     axis.grid(True, which='major', color='black', linestyle=':', linewidth=1.00)
 516     axis.grid(True, which='minor', color='black', linestyle=':', linewidth=0.50)
 517     axis.set_xlim(0, N+1)
 518     axis.set_ylim(-0.03, 0.15)
 519 
 520     # configuracion de ticks del eje x
 521     axis.xaxis.set_major_locator(MaxNLocator(nbins=15))
 522     axis.xaxis.set_minor_locator(AutoMinorLocator(2))
 523 
 524     axis.yaxis.set_minor_locator(AutoMinorLocator(2))
 525 
 526     # axis.yaxis.set_major_locator(MaxNLocator(nbins=5))
 527     # axis.yaxis.set_minor_locator(AutoMinorLocator(4))
 528 
 529     if legend != "":
 530         axis.legend([markerline], [legend], loc='upper right')
 531 
 532     if show == False:
 533         save_plot(fig, save_name)
 534     else:
 535         plt.show()
 536 
 537     return fig, axis
 538 
 539 # np.linspace(start, stop, num).astype(int)