02_12_2023 Python Implementation of Variational mode decomposition.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "012f17dc",
   "metadata": {},
   "source": [
    "# vmdpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f3455ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\"\"\" \n",
    "Created on Wed Feb 20 19:24:58 2019 \n",
    " \n",
    "@author: Vinícius Rezende Carvalho \n",
    "\"\"\"\n",
    "import numpy as np\n",
    "\n",
    "def  VMD(f, alpha, tau, K, DC, init, tol):\n",
    "    \"\"\" \n",
    "    u,u_hat,omega = VMD(f, alpha, tau, K, DC, init, tol) \n",
    "    Variational mode decomposition \n",
    "    Python implementation by Vinícius Rezende Carvalho - vrcarva@gmail.com \n",
    "    code based on Dominique Zosso's MATLAB code, available at: \n",
    "    https://www.mathworks.com/matlabcentral/fileexchange/44765-variational-mode-decomposition \n",
    "    Original paper: \n",
    "    Dragomiretskiy, K. and Zosso, D. (2014) ‘Variational Mode Decomposition’,  \n",
    "    IEEE Transactions on Signal Processing, 62(3), pp. 531–544. doi: 10.1109/TSP.2013.2288675. \n",
    "     \n",
    "     \n",
    "    Input and Parameters: \n",
    "    --------------------- \n",
    "    f       - the time domain signal (1D) to be decomposed \n",
    "    alpha   - the balancing parameter of the data-fidelity constraint \n",
    "    tau     - time-step of the dual ascent ( pick 0 for noise-slack ) \n",
    "    K       - the number of modes to be recovered \n",
    "    DC      - true if the first mode is put and kept at DC (0-freq) \n",
    "    init    - 0 = all omegas start at 0 \n",
    "                       1 = all omegas start uniformly distributed \n",
    "                      2 = all omegas initialized randomly \n",
    "    tol     - tolerance of convergence criterion; typically around 1e-6 \n",
    " \n",
    "    Output: \n",
    "    ------- \n",
    "    u       - the collection of decomposed modes \n",
    "    u_hat   - spectra of the modes \n",
    "    omega   - estimated mode center-frequencies \n",
    "    \"\"\"\n",
    "    \n",
    "    if len(f)%2:\n",
    "       f = f[:-1]\n",
    "\n",
    "    # Period and sampling frequency of input signal\n",
    "    fs = 1./len(f)\n",
    "    \n",
    "    ltemp = len(f)//2 \n",
    "    fMirr =  np.append(np.flip(f[:ltemp],axis = 0),f)  \n",
    "    fMirr = np.append(fMirr,np.flip(f[-ltemp:],axis = 0))\n",
    "\n",
    "    # Time Domain 0 to T (of mirrored signal)\n",
    "    T = len(fMirr)\n",
    "    t = np.arange(1,T+1)/T  \n",
    "    \n",
    "    # Spectral Domain discretization\n",
    "    freqs = t-0.5-(1/T)\n",
    "\n",
    "    # Maximum number of iterations (if not converged yet, then it won't anyway)\n",
    "    Niter = 500\n",
    "    # For future generalizations: individual alpha for each mode\n",
    "    Alpha = alpha*np.ones(K)\n",
    "    \n",
    "    # Construct and center f_hat\n",
    "    f_hat = np.fft.fftshift((np.fft.fft(fMirr)))\n",
    "    f_hat_plus = np.copy(f_hat) #copy f_hat\n",
    "    f_hat_plus[:T//2] = 0\n",
    "\n",
    "    # Initialization of omega_k\n",
    "    omega_plus = np.zeros([Niter, K])\n",
    "\n",
    "\n",
    "    if init == 1:\n",
    "        for i in range(K):\n",
    "            omega_plus[0,i] = (0.5/K)*(i)\n",
    "    elif init == 2:\n",
    "        omega_plus[0,:] = np.sort(np.exp(np.log(fs) + (np.log(0.5)-np.log(fs))*np.random.rand(1,K)))\n",
    "    else:\n",
    "        omega_plus[0,:] = 0\n",
    "            \n",
    "    # if DC mode imposed, set its omega to 0\n",
    "    if DC:\n",
    "        omega_plus[0,0] = 0\n",
    "    \n",
    "    # start with empty dual variables\n",
    "    lambda_hat = np.zeros([Niter, len(freqs)], dtype = complex)\n",
    "    \n",
    "    # other inits\n",
    "    uDiff = tol+np.spacing(1) # update step\n",
    "    n = 0 # loop counter\n",
    "    sum_uk = 0 # accumulator\n",
    "    # matrix keeping track of every iterant // could be discarded for mem\n",
    "    u_hat_plus = np.zeros([Niter, len(freqs), K],dtype=complex)    \n",
    "\n",
    "    #*** Main loop for iterative updates***\n",
    "\n",
    "    while ( uDiff > tol and  n < Niter-1 ): # not converged and below iterations limit\n",
    "        # update first mode accumulator\n",
    "        k = 0\n",
    "        sum_uk = u_hat_plus[n,:,K-1] + sum_uk - u_hat_plus[n,:,0]\n",
    "        \n",
    "        # update spectrum of first mode through Wiener filter of residuals\n",
    "        u_hat_plus[n+1,:,k] = (f_hat_plus - sum_uk - lambda_hat[n,:]/2)/(1.+Alpha[k]*(freqs - omega_plus[n,k])**2)\n",
    "        \n",
    "        # update first omega if not held at 0\n",
    "        if not(DC):\n",
    "            omega_plus[n+1,k] = np.dot(freqs[T//2:T],(abs(u_hat_plus[n+1, T//2:T, k])**2))/np.sum(abs(u_hat_plus[n+1,T//2:T,k])**2)\n",
    "\n",
    "        # update of any other mode\n",
    "        for k in np.arange(1,K):\n",
    "            #accumulator\n",
    "            sum_uk = u_hat_plus[n+1,:,k-1] + sum_uk - u_hat_plus[n,:,k]\n",
    "            # mode spectrum\n",
    "            u_hat_plus[n+1,:,k] = (f_hat_plus - sum_uk - lambda_hat[n,:]/2)/(1+Alpha[k]*(freqs - omega_plus[n,k])**2)\n",
    "            # center frequencies\n",
    "            omega_plus[n+1,k] = np.dot(freqs[T//2:T],(abs(u_hat_plus[n+1, T//2:T, k])**2))/np.sum(abs(u_hat_plus[n+1,T//2:T,k])**2)\n",
    "            \n",
    "        # Dual ascent\n",
    "        lambda_hat[n+1,:] = lambda_hat[n,:] + tau*(np.sum(u_hat_plus[n+1,:,:],axis = 1) - f_hat_plus)\n",
    "        \n",
    "        # loop counter\n",
    "        n = n+1\n",
    "        \n",
    "        # converged yet?\n",
    "        uDiff = np.spacing(1)\n",
    "        for i in range(K):\n",
    "            uDiff = uDiff + (1/T)*np.dot((u_hat_plus[n,:,i]-u_hat_plus[n-1,:,i]),np.conj((u_hat_plus[n,:,i]-u_hat_plus[n-1,:,i])))\n",
    "\n",
    "        uDiff = np.abs(uDiff)        \n",
    "            \n",
    "    #Postprocessing and cleanup\n",
    "    \n",
    "    #discard empty space if converged early\n",
    "    Niter = np.min([Niter,n])\n",
    "    omega = omega_plus[:Niter,:]\n",
    "    \n",
    "    idxs = np.flip(np.arange(1,T//2+1),axis = 0)\n",
    "    # Signal reconstruction\n",
    "    u_hat = np.zeros([T, K],dtype = complex)\n",
    "    u_hat[T//2:T,:] = u_hat_plus[Niter-1,T//2:T,:]\n",
    "    u_hat[idxs,:] = np.conj(u_hat_plus[Niter-1,T//2:T,:])\n",
    "    u_hat[0,:] = np.conj(u_hat[-1,:])    \n",
    "    \n",
    "    u = np.zeros([K,len(t)])\n",
    "    for k in range(K):\n",
    "        u[k,:] = np.real(np.fft.ifft(np.fft.ifftshift(u_hat[:,k])))\n",
    "        \n",
    "    # remove mirror part\n",
    "    u = u[:,T//4:3*T//4]\n",
    "\n",
    "    # recompute spectrum\n",
    "    u_hat = np.zeros([u.shape[1],K],dtype = complex)\n",
    "    for k in range(K):\n",
    "        u_hat[:,k]=np.fft.fftshift(np.fft.fft(u[k,:]))\n",
    "\n",
    "    return u, u_hat, omega"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}