import numpy as np
from numba import jit
@jit
def VMD(signal, fs, T, alpha, tau, K, DC, init, tol):
f_mirror = np.zeros((2*T))
for i in range(T//2):
f_mirror[i] = signal[T//2-1-i]
for k in range(T//2, 3*T//2):
f_mirror[k] = signal[k - T // 2]
for l in range(3*T//2, 2*T):
f_mirror[l] = signal[5*T//2-l-1]
f = f_mirror
T = len(f)
t = np.arange(1, T+1, 1)/T
freqs = t - 0.5 - 1.0 / T
freqs = freqs + 0.j
N = 500
Alpha = alpha * np.ones((1, K))
Alpha = Alpha + 0.j
f_hat = np.fft.fftshift(np.fft.fft(f))
f_hat_plus = f_hat
for l in range(0, T//2):
f_hat_plus[l] = 0
u_hat_plus = np.zeros((N, T, K))
u_hat_plus = u_hat_plus + 0.j
omega_plus = np.zeros((N, K)) # 生成N行K列的数组,元素为0, 500*3
omega_plus = omega_plus + 0.j
if init == 1:
for i in range(0, K):
omega_plus[0][i] = 0.5/K*i
elif init == 2:
omega_plus[0, :] = np.sort(np.exp(np.log(fs)) + (np.log(0.5)-np.log(fs))*np.random.rand(0, K))
else:
omega_plus[0, :] = 0
if DC:
omega_plus[0, 0] = 0
else:
pass
# 从空双变量开始
lambda_hat = np.zeros((N, T)) #
lambda_hat = lambda_hat + 0.j
# 其他参数初始化
uDiff = tol + 2.77555756e-17 # 更新步长
n = 1 # 循环计数
sum_uk = 0.j # 器材累加
while uDiff > tol and n < N : # 表示没有收敛或低于迭代次数,其中,n和K都从1开始
# 更新第一阶模态累加器
k = 0
sum_uk = u_hat_plus[n-1, :, K-1] + sum_uk - u_hat_plus[n-1, :, 0] # sum_uk此时还是一维数组
# print(sum_uk) # ------------------------ ----检验专用-----------------------
# 通过残差维纳滤波器更新第一模态的频谱
u_hat_plus[n, :, k] = (f_hat_plus - sum_uk - lambda_hat[n - 1, :] / 2.0) / (1 + Alpha[0, k]*(freqs - omega_plus[n-1, k])**2)
if ~DC:
omega_plus[n, k] = np.dot(freqs[T//2: T], ((abs(u_hat_plus[n, T//2:T, k])).reshape(T//2, 1))**2) / sum((abs(u_hat_plus[n, T//2:T, k]))**2)
for k in range(1, K):