Sparse_Dictionary_Learning_KSVD.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## K-SVD Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on the implementation of IshitaTakeshi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m registry at `~/.julia/registries/General`\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m git-repo `https://github.com/JuliaRegistries/General.git`\n",
      "\u001b[2K\u001b[?25h[1mFetching:\u001b[22m\u001b[39m [========================================>]  100.0 %.0 %          ]  75.0 %\u001b[36m\u001b[1mFetching:\u001b[22m\u001b[39m [================================>        ]  78.8 %\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m Resolving\u001b[22m\u001b[39m package versions...\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Project.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n",
      "\u001b[32m\u001b[1m  Updating\u001b[22m\u001b[39m `~/.julia/environments/v1.2/Manifest.toml`\n",
      "\u001b[90m [no changes]\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "using Pkg\n",
    "Pkg.add(\"PyCall\") \n",
    "Pkg.add(\"DataStructures\")\n",
    "Pkg.add(\"SparseArrays\")\n",
    "Pkg.add(\"ProgressBars\")\n",
    "Pkg.add(\"DelimitedFiles\")\n",
    "Pkg.add(\"Plots\")\n",
    "Pkg.add(\"Images\")\n",
    "using DataStructures\n",
    "using SparseArrays\n",
    "using LinearAlgebra\n",
    "using ProgressBars\n",
    "using PyCall\n",
    "using Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "init_dictionary (generic function with 1 method)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function init_dictionary(n::Int, K::Int)\n",
    "    \"\"\"\n",
    "    Initialize the dictionary.\n",
    "    Args: \n",
    "        n: dimension of input signal\n",
    "        k = number of atoms in the dictionary\n",
    "    \"\"\"\n",
    "    # D must be a full-rank matrix\n",
    "    D = rand(n, K)\n",
    "    while rank(D) != min(n, K)\n",
    "        D = rand(n, K)\n",
    "    end\n",
    "\n",
    "    @inbounds for k in 1:K\n",
    "        D[:, k] ./= norm(@view(D[:, k]))\n",
    "    end\n",
    "    return D\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "matching_pursuit"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The implementation is referencing the wikipedia page\n",
    "# https://en.wikipedia.org/wiki/Matching_pursuit#The_algorithm\n",
    "\n",
    "const default_max_iter = 20\n",
    "const default_tolerance = 1e-6\n",
    "\n",
    "\n",
    "function SparseArrays.sparsevec(d::DefaultDict, m::Int)\n",
    "    SparseArrays.sparsevec(collect(keys(d)), collect(values(d)), m)\n",
    "end\n",
    "\n",
    "\n",
    "function matching_pursuit_(data::AbstractVector, dictionary::AbstractMatrix,\n",
    "                           max_iter::Int, tolerance::Float64)\n",
    "    n_atoms = size(dictionary, 2)\n",
    "\n",
    "    residual = copy(data)\n",
    "\n",
    "    xdict = DefaultDict{Int, Float64}(0.)\n",
    "    for i in 1:max_iter\n",
    "        if norm(residual) < tolerance\n",
    "            return sparsevec(xdict, n_atoms)\n",
    "        end\n",
    "\n",
    "        # find an atom with maximum inner product\n",
    "        products = dictionary' * residual\n",
    "        _, maxindex = findmax(abs.(products))\n",
    "        maxval = products[maxindex]\n",
    "        atom = dictionary[:, maxindex]\n",
    "\n",
    "        # c is the length of the projection of data onto atom\n",
    "        a = maxval / sum(abs2, atom)  # equivalent to maxval / norm(atom)^2\n",
    "        residual -= atom * a\n",
    "\n",
    "        xdict[maxindex] += a\n",
    "    end\n",
    "    return sparsevec(xdict, n_atoms)\n",
    "end\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "    matching_pursuit(data::Vector, dictionary::AbstractMatrix;\n",
    "                     max_iter::Int = $default_max_iter,\n",
    "                     tolerance::Float64 = $default_tolerance)\n",
    "Find ``x`` such that ``Dx = y`` or ``Dx ≈ y`` where y is `data` and D is `dictionary`.\n",
    "```\n",
    "# Arguments\n",
    "* `max_iter`: Hard limit of iterations\n",
    "* `tolerance`: Exit when the norm of the residual < tolerance\n",
    "```\n",
    "\"\"\"\n",
    "function matching_pursuit(data::AbstractVector, dictionary::AbstractMatrix;\n",
    "                          max_iter::Int = default_max_iter,\n",
    "                          tolerance = default_tolerance)\n",
    "\n",
    "    if tolerance <= 0\n",
    "        throw(ArgumentError(\"`tolerance` must be > 0\"))\n",
    "    end\n",
    "\n",
    "    if max_iter <= 0\n",
    "        throw(ArgumentError(\"`max_iter` must be > 0\"))\n",
    "    end\n",
    "\n",
    "    if size(data, 1) != size(dictionary, 1)\n",
    "        throw(ArgumentError(\n",
    "            \"Dimensions must match: `size(data, 1)` and `size(dictionary, 1)`.\"\n",
    "        ))\n",
    "    end\n",
    "\n",
    "    matching_pursuit_(data, dictionary, max_iter, tolerance)\n",
    "end\n",
    "# The implementation is referencing the wikipedia page\n",
    "# https://en.wikipedia.org/wiki/Matching_pursuit#The_algorithm\n",
    "\n",
    "const default_max_iter = 20\n",
    "const default_tolerance = 1e-6\n",
    "\n",
    "\n",
    "function SparseArrays.sparsevec(d::DefaultDict, m::Int)\n",
    "    SparseArrays.sparsevec(collect(keys(d)), collect(values(d)), m)\n",
    "end\n",
    "\n",
    "\n",
    "function matching_pursuit_(data::AbstractVector, dictionary::AbstractMatrix,\n",
    "                           max_iter::Int, tolerance::Float64)\n",
    "    n_atoms = size(dictionary, 2)\n",
    "\n",
    "    residual = copy(data)\n",
    "\n",
    "    xdict = DefaultDict{Int, Float64}(0.)\n",
    "    for i in 1:max_iter\n",
    "        if norm(residual) < tolerance\n",
    "            return sparsevec(xdict, n_atoms)\n",
    "        end\n",
    "\n",
    "        # find an atom with maximum inner product\n",
    "        products = dictionary' * residual\n",
    "        _, maxindex = findmax(abs.(products))\n",
    "        maxval = products[maxindex]\n",
    "        atom = dictionary[:, maxindex]\n",
    "\n",
    "        # c is the length of the projection of data onto atom\n",
    "        a = maxval / sum(abs2, atom)  # equivalent to maxval / norm(atom)^2\n",
    "        residual -= atom * a\n",
    "\n",
    "        xdict[maxindex] += a\n",
    "    end\n",
    "    return sparsevec(xdict, n_atoms)\n",
    "end\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "    matching_pursuit(data::Vector, dictionary::AbstractMatrix;\n",
    "                     max_iter::Int = $default_max_iter,\n",
    "                     tolerance::Float64 = $default_tolerance)\n",
    "Find ``x`` such that ``Dx = y`` or ``Dx ≈ y`` where y is `data` and D is `dictionary`.\n",
    "```\n",
    "# Arguments\n",
    "* `max_iter`: Hard limit of iterations\n",
    "* `tolerance`: Exit when the norm of the residual < tolerance\n",
    "```\n",
    "\"\"\"\n",
    "function matching_pursuit(data::AbstractVector, dictionary::AbstractMatrix;\n",
    "                          max_iter::Int = default_max_iter,\n",
    "                          tolerance = default_tolerance)\n",
    "\n",
    "    if tolerance <= 0\n",
    "        throw(ArgumentError(\"`tolerance` must be > 0\"))\n",
    "    end\n",
    "\n",
    "    if max_iter <= 0\n",
    "        throw(ArgumentError(\"`max_iter` must be > 0\"))\n",
    "    end\n",
    "\n",
    "    if size(data, 1) != size(dictionary, 1)\n",
    "        throw(ArgumentError(\n",
    "            \"Dimensions must match: `size(data, 1)` and `size(dictionary, 1)`.\"\n",
    "        ))\n",
    "    end\n",
    "\n",
    "    matching_pursuit_(data, dictionary, max_iter, tolerance)\n",
    "end\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "    matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;\n",
    "                     max_iter::Int = $default_max_iter,\n",
    "                     tolerance::Float64 = $default_tolerance)\n",
    "Find ``X`` such that ``DX = Y`` or ``DX ≈ Y`` where Y is `data` and D is `dictionary`.\n",
    "```\n",
    "# Arguments\n",
    "* `max_iter`: Hard limit of iterations\n",
    "* `tolerance`: Exit when the norm of the residual < tolerance\n",
    "```\n",
    "\"\"\"\n",
    "function matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;\n",
    "                          max_iter::Int = default_max_iter,\n",
    "                          tolerance::Float64 = default_tolerance)\n",
    "    K = size(dictionary, 2)\n",
    "    N = size(data, 2)\n",
    "\n",
    "    X = spzeros(K, N)\n",
    "\n",
    "    for i in 1:N\n",
    "        X[:, i] = matching_pursuit(\n",
    "            vec(data[:, i]),\n",
    "            dictionary,\n",
    "            max_iter = max_iter,\n",
    "            tolerance = tolerance\n",
    "        )\n",
    "    end\n",
    "    return X\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "    matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;\n",
    "                     max_iter::Int = $default_max_iter,\n",
    "                     tolerance::Float64 = $default_tolerance)\n",
    "Find ``X`` such that ``DX = Y`` or ``DX ≈ Y`` where Y is `data` and D is `dictionary`.\n",
    "```\n",
    "# Arguments\n",
    "* `max_iter`: Hard limit of iterations\n",
    "* `tolerance`: Exit when the norm of the residual < tolerance\n",
    "```\n",
    "\"\"\"\n",
    "function matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;\n",
    "                          max_iter::Int = default_max_iter,\n",
    "                          tolerance::Float64 = default_tolerance)\n",
    "    K = size(dictionary, 2)\n",
    "    N = size(data, 2)\n",
    "\n",
    "    X = spzeros(K, N)\n",
    "\n",
    "    for i in 1:N\n",
    "        X[:, i] = matching_pursuit(\n",
    "            vec(data[:, i]),\n",
    "            dictionary,\n",
    "            max_iter = max_iter,\n",
    "            tolerance = tolerance\n",
    "        )\n",
    "    end\n",
    "    return X\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "K_SVD (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function K_SVD(Y,niter_KSVD,n_atoms)\n",
    "    \"\"\"\n",
    "    Computes the K-SVD algorithm.\n",
    "    \n",
    "    Args:\n",
    "    Y - \n",
    "    niter_KSVD: number of iterations for the algorithm\n",
    "    n_atoms - number of atoms in the dictionary\n",
    "    returns: Dictionary (D) and Sparse coefficients (X) \n",
    "    \n",
    "    \"\"\"\n",
    "    D = init_dictionary(size(Y,1),n_atoms)\n",
    "    X = matching_pursuit(Y,D)\n",
    "    for i in ProgressBar(1:niter_KSVD)\n",
    "        X = matching_pursuit(Y,D)\n",
    "        for k = 1:n_atoms\n",
    "            Xk = X[k,:]\n",
    "            all(iszero,Xk)&&continue\n",
    "            wk = findall(!iszero,Xk)\n",
    "            indices = [j for j=1:size(D,2) if j!=k]\n",
    "            Ek = Y - D[:,indices]*X[indices,:]\n",
    "            Ωk = sparse(wk,1:length(wk),ones(length(wk)),size(Y,2),length(wk))\n",
    "            U, S, V= svd(Ek*Ωk, full=true)\n",
    "            D[:,k]=U[:,1]\n",
    "            X[k,wk] = V[:,1]*S[1]\n",
    "        end\n",
    "    end\n",
    "    return D,X\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### First Task: Image Compression _(MNIST dataset)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = pyimport(\"sklearn.datasets\")\n",
    "digits = datasets.load_digits()\n",
    "Y = digits[\"data\"];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D,X = K_SVD(Y,200,256);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find D and X such that Y ≈ DX\n",
    "\n",
    "println(\"||Y - D * X|| = $(norm(Y - D * X))\")\n",
    "\n",
    "println(\"The ratio of zero elemnts in the matrix X: \",\n",
    "        sum(X .== 0) / length(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# saves result in a dlm file\n",
    "using DelimitedFiles\n",
    "\n",
    "writedlm(\"dictionary.dlm\", D)\n",
    "writedlm(\"reconstruction.dlm\", D*X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Second Analysis _(CIFAR10_dataset)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "using DelimitedFiles\n",
    "\n",
    "\n",
    "data = readdlm(\"../CIFAR10_data.dlm\");\n",
    "labels = readdlm(\"../CIFAR10_labels.dlm\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.0%┣████▉                                                        ┫ 8/100 [13:28<02:56:53, 0.0 it/s]"
     ]
    },
    {
     "ename": "InterruptException",
     "evalue": "InterruptException:",
     "output_type": "error",
     "traceback": [
      "InterruptException:",
      "",
      "Stacktrace:",
      " [1] gesdd!(::Char, ::Array{Float64,2}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.2/LinearAlgebra/src/lapack.jl:1628",
      " [2] #svd!#73(::Bool, ::typeof(svd!), ::Array{Float64,2}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.2/LinearAlgebra/src/svd.jl:63",
      " [3] #svd#74 at ./none:0 [inlined]",
      " [4] #svd at ./none:0 [inlined]",
      " [5] K_SVD(::Array{Float64,2}, ::Int64, ::Int64) at ./In[4]:23",
      " [6] top-level scope at In[15]:1"
     ]
    }
   ],
   "source": [
    "D,X = K_SVD(data[1:2000,:],100,256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "ename": "DimensionMismatch",
     "evalue": "DimensionMismatch(\"dimensions must match\")",
     "output_type": "error",
     "traceback": [
      "DimensionMismatch(\"dimensions must match\")",
      "",
      "Stacktrace:",
      " [1] promote_shape at ./indices.jl:154 [inlined]",
      " [2] promote_shape at ./indices.jl:145 [inlined]",
      " [3] -(::Array{Float64,2}, ::Array{Float64,2}) at ./arraymath.jl:38",
      " [4] top-level scope at In[17]:1"
     ]
    }
   ],
   "source": [
    "# Find D and X such that Y ≈ DX\n",
    "\n",
    "println(\"||Y - D * X|| = $(norm(data[1:2000,:] - D * X))\")\n",
    "\n",
    "println(\"The ratio of zero elemnts in the matrix X: \",\n",
    "        sum(X .== 0) / length(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writedlm(\"dictionary_CIFAR.dlm\", D)\n",
    "writedlm(\"reconstruction_CIFAR.dlm\", D*X)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.2.0",
   "language": "julia",
   "name": "julia-1.2"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}