PDS 2019.4 - RMS and PSD of an Electromyography Signal

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec 18 23:36:47 2019

@author: jeremiah
"""

#%% importings
import numpy as np
from scipy.io import wavfile
import matplotlib.pyplot as plt
import soundfile as sf
import wave
#%%

def rms_arr(arr, win_size, win_sobrep_proportion):
    def rms_win(win_arr):
        #return np.sqrt(np.sum(win_arr**2))/win_size
        return np.sqrt(np.mean(win_arr**2))

# =============================================================================
#     rest = (len(arr)%win_size)
#
#     if rest:
#         arr = arr[:-rest]
# =============================================================================
    ans = []
    for i in np.arange(0, len(arr), int(win_size/win_sobrep_proportion)):
        if i+win_size >= len(arr):
            break

        ans.append(rms_win(arr[i:i+win_size]))

    return ans

#%%
# =============================================================================
# Diego's function
# =============================================================================
def rms(arr, window_size):
    ans = []
    n = len(arr)
    if (window_size > n):
        return ans

    def compute_rms(x):
        return np.sqrt(x) / window_size

    cur = 0

    for i in range(window_size):
        cur += arr[i] * arr[i]
    ans.append(compute_rms(cur))
    for i in range(1, n - window_size + 1):
        cur -= arr[i - 1] * arr[i - 1]
        cur += arr[i + window_size - 1] * arr[i + window_size - 1]
        ans.append(compute_rms(cur))

    return ans

#%%
def psd(signal, fft_win_size, psd_win_size, win_sobrep_proportion):

    psd = []
    fft = []
    count = 0
    for i in np.arange(0, len(signal), int(fft_win_size/win_sobrep_proportion)):
        count += 1

        if i+fft_win_size >= len(signal):
            break
        fft.append(np.abs(np.fft.fft(signal[i:i+fft_win_size])))

        if count == psd_win_size:
            fft = np.array(fft, "float64")
            psd.append(np.sum(fft**2, axis=0)/psd_win_size)
            fft = []
            count = 0


    return np.array(psd, "float64")




#%%
file_name = "02.wav"
path = '/media/jeremiah/7E9BF5A34D96B6A4/2019.4/PDS/pds-2019.4-master/'+file_name
data, _ = sf.read(path)
#_, data = wavfile.read("/media/jeremiah/7E9BF5A34D96B6A4/2019.4/PDS/pds-2019.4-master/01.wav")
print(data[:10])
data = np.array(data, "float64")
win_sobrep_proportion = 4
win_size = 2**16
rms_data = rms_arr(data, win_size, win_sobrep_proportion)
plt.plot(rms_data)
plt.title("RMS with window size of "+str(win_size))
plt.savefig("/media/jeremiah/7E9BF5A34D96B6A4/2019.4/PDS/pds-2019.4-master/"+file_name+"rms-win-of"+str(win_size)+".png")
plt.show()

#%%
win_sobrep_proportion = 2
fft_win_size = 2**11
psd_win_size = 2**6
psd_data = psd(data, fft_win_size, psd_win_size, win_sobrep_proportion)
median_psd_data = np.median(psd_data, axis=1)
plt.plot(median_psd_data)
plt.title("PSD with fft window size of "+str(fft_win_size)+" and psd window size of "+str(psd_win_size))
plt.savefig("/media/jeremiah/7E9BF5A34D96B6A4/2019.4/PDS/pds-2019.4-master/"+file_name+"psd-fft_win-of"+str(fft_win_size)+"-psd_win-of"+str(psd_win_size)+".png")