{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Student version!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Block coordinate minimization (BCM) for 1D dictionary learning of wave+spike data\n",
    "\n",
    "In-class group activity, EECS 598-006 W19\n",
    "\n",
    "Given a set of training signals $X = [x_1,\\ldots,x_L]$,\n",
    "this notebook performs *dictionary learning* by solving\n",
    "$$\n",
    "\\hat{D} = \\arg \\min_D \\min_Z \\frac{1}{2} \\| X - D Z \\|^2_F + \\beta \\|Z\\|_0\n",
    ",$$\n",
    "subject to the constraint that\n",
    "each column of $D$ must have unit norm.\n",
    "\n",
    "This notebook considers signals that are generated\n",
    "as the sum of\n",
    "a smooth component\n",
    "plus a few Kronecker impulses.\n",
    "\n",
    "This notebook applies BCM to update one column of $D$\n",
    "at a time,\n",
    "i.e., one atom,\n",
    "and then one corresponding column of $C = Z'$,\n",
    "looping over all columns of $D$\n",
    "to make one iteration.\n",
    "This is the \"sum of outer products\" (SOUP) approach\n",
    "to dictionary learning,\n",
    "developed by Sai Ravishankar et al.,\n",
    "at the University of Michigan.\n",
    "\n",
    "2019-04-10 Jeff Fessler, University of Michigan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Packages needed\n",
    "using MIRT: jim\n",
    "using FFTW: idct\n",
    "using LinearAlgebra: I, norm, pinv, svd\n",
    "using Plots; default(markerstrokecolor=nothing)\n",
    "using LaTeXStrings\n",
    "using Random: seed!, randperm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define basis functions for generating 1D training signals.\n",
    "# This 1D problem is small enough that we use the full DCT matrix and identity matrix.\n",
    "# For a large-scale problem, we would make D1 and D2 a LinearMap instead.\n",
    "#\n",
    "N = 2^7\n",
    "K1 = 10 # number of dct terms\n",
    "B1 = idct(Matrix(I,N,N)[:,1:K1], 1) # inverse DCT basis\n",
    "K2 = 12 # number of impulse terms\n",
    "#K2 = N # this would be the ideal value but it runs too slow for in-class work\n",
    "seed!(0); tmp = sort(randperm(N)[1:K2])\n",
    "B2 = Matrix(I,N,N)[:,tmp]\n",
    "jim([B1 B2]')\n",
    ";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scatter(B1, label=\"\", xlabel=L\"n\", ylabel=L\"D[:,k]\", title=\"$K1 iDCT basis atoms\")\n",
    "plot!(B1, label=\"\", color=:black)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate ideal training signals using the basis with a couple spikes\n",
    "L = 10^4\n",
    "seed!(0)\n",
    "tmp1 = rand(K1,L) .< 0.05 # 1/ndct\n",
    "Z1 = 3randn(K1,L) .* tmp1\n",
    "Xtrue = B1 * Z1\n",
    "tmp2 = rand(K2,L) .< 0.05 # 1/N\n",
    "Z2 = randn(K2,L) .* tmp2\n",
    "Xtrue .+= B2 * Z2\n",
    "Dtrue = [B1 B2]\n",
    "@assert isapprox(sum(Dtrue.^2,dims=1)[:], ones(K1+K2))\n",
    "Ztrue = [Z1; Z2]\n",
    "pctnz = Z -> round(count(Z .!= 0) / length(Z) * 100, digits=3)\n",
    "@show pctnz(Ztrue)\n",
    "@assert isapprox(Xtrue, Dtrue*Ztrue)\n",
    "jim(Dtrue')\n",
    "scatter(Xtrue[:,1:5], label=\"\", xlabel=L\"n\", ylabel=L\"x_l[n] \\mbox{ true}\", title=\"Xtrue\")\n",
    "plot!(Xtrue[:,1:5], label=\"\", color=:black)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make noisy training signals for DL\n",
    "sig = 0.01\n",
    "seed!(0); X = Xtrue + sig * randn(N,L)\n",
    "scatter(X[:,1:5], label=\"\", xlabel=L\"n\", ylabel=L\"x_l[n] \\mbox{ noisy}\")\n",
    "plot!(X[:,1:5], label=\"\", color=:black)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#savefig(\"dl-1d-soup-x.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility functions for dictionaries\n",
    "#\n",
    "# sort dictionary D to best align with true dictionary for nice display\n",
    "function sorter(D, Dtrue)\n",
    "    #Dout = copy(D)\n",
    "    @assert size(D,1) == size(Dtrue,1)\n",
    "    @assert size(D,2) >= size(Dtrue,2)\n",
    "    K = size(Dtrue,2)\n",
    "    klist = 1:size(D,2)\n",
    "    korder = zeros(Int, K)\n",
    "    for k=1:K\n",
    "        # unnnormalized inner product, assume atoms have unit norm\n",
    "        tmp = (Dtrue[:,k]' * D[:,klist])[:]\n",
    "        tmp = findmax(abs.(tmp))[2]\n",
    "        kbest = klist[tmp]\n",
    "        korder[k] = kbest\n",
    "        klist = klist[klist .!= kbest]\n",
    "    end\n",
    "    return [korder; klist]\n",
    "end\n",
    "\n",
    "# functions to make max value positive and normalize\n",
    "signs = D -> sign.(D[findmax(abs.(D), dims=1)[2]])\n",
    "normalize = D -> D ./ sqrt.(sum(abs.(D).^2, dims=1)) # normalize!\n",
    "\n",
    "seed!(0); tmp = normalize(rand(5,6))\n",
    "@assert sorter([tmp normalize(rand(5,3))], tmp) == 1:9 # test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initial guess of the dictionary D and coefficients Z, via PCA\n",
    "seed!(0)\n",
    "#K = 2*N # overcomplete\n",
    "K = K1 + K2 # cheat for now\n",
    "#D0 = randn(N,K) # initial dictionary - random\n",
    "D0 = svd(X).U[:,1:min(K,N)] # PCA\n",
    "if K > N\n",
    "    D0 = normalize([D0 randn(N,K-N)])\n",
    "end\n",
    "D0 = D0[:,sorter(D0, Dtrue)]\n",
    "D0 = D0 .* signs(D0)\n",
    "\n",
    "Z0 = pinv(D0) * X # think about why we use pinv here!\n",
    "# extrema(Z0)\n",
    "plot(jim(Dtrue', \"Dtrue\"), jim(D0', \"D0\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform BCM for DL via SOUP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# at least one of these shrinkage functions should be useful\n",
    "soft = (z,b) -> sign(z) * max(abs(z) - b, 0)\n",
    "hard = (z,b) -> z * (abs(z) > sqrt(2b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loop over 50 iterations of BCM here\n",
    "# update D[:,1] then C[:,1] then D[:,2], ... C[:,K]\n",
    "# where C = Z'\n",
    "\n",
    "reg = 0.05 # beta\n",
    "cost = (D,Z) -> 0.5 * norm(X - D*Z)^2 + reg * count(Z .!= 0) # \"0-norm\"\n",
    "\n",
    "niter = 60\n",
    "costs = zeros(niter + 1)\n",
    "D = copy(D0)\n",
    "Z = copy(Z0)\n",
    "costs[1] = cost(D, Z)\n",
    "R = X - D * Z # residual state\n",
    "C = Z'\n",
    "for iter = 1:niter\n",
    "    for k = 1:K\n",
    "        ck = @view C[:,k]\n",
    "        norm_ck = norm(ck)\n",
    "        \n",
    "        # put your dictionary atom update code here:\n",
    "\n",
    "        # ...\n",
    "\n",
    "        # be sure to keep the residual \"R\" up-to-date!\n",
    "        \n",
    "        # put your coefficient vector update here:\n",
    "\n",
    "    end\n",
    "    costs[iter+1] = cost(D, C')\n",
    "\n",
    "    if mod(iter, 5) == 0\n",
    "        @show iter, pctnz(C) # monitor progress of sparsity\n",
    "    end\n",
    "end\n",
    "Z = C'\n",
    "scatter(0:niter, costs, label=\"cost\") # plot cost vs iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort and sign \"correct\" the estimated dictionary for nice display\n",
    "# you sould see that the spike part of the dictionary looks very good\n",
    "# but the smooth part is hard to tell, even when sorted using sorter()\n",
    "tmp = sorter(D, Dtrue)\n",
    "D = D[:,tmp]\n",
    "Z = Z[tmp,:]\n",
    "tmp = signs(D)\n",
    "D = D .* tmp\n",
    "Z = Z .* tmp'\n",
    "plot(jim(Dtrue', \"Dtrue\"), jim(D0', L\"D_0\"), jim(D', L\"\\hat{D}\"), layout=(1,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#savefig(\"dl-1d-soup-d.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the estimated dictionary\n",
    "scatter(D, title=L\"\\hat{D}\", xlabel=L\"n\", ylabel=L\"\\hat{d}_k[n]\", label=\"\")\n",
    "#scatter!(legend=:bottomright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# examine the representation error and the cost function\n",
    "# think about these values!\n",
    "@show norm(X - D0*Z0)^2/(N*L*sig^2), cost(D0,Z0)\n",
    "@show norm(X - D*Z)^2/(N*L*sig^2), cost(D,Z);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Count the number of non-zero coefficients and report that value to Canvas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count(Z .!= 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test the learned dictionary by using it for sparse coding some new test data\n",
    "# first generate the data:\n",
    "seed!(7)\n",
    "Ltest = 10^3\n",
    "tmp1 = rand(K1,Ltest) .< 0.05 # 1/ndct\n",
    "Z1test = 3randn(K1,Ltest) .* tmp1\n",
    "Xtest = B1 * Z1test\n",
    "tmp2 = rand(K2,Ltest) .< 0.05 # 1/N\n",
    "Z2test = randn(K2,Ltest) .* tmp2\n",
    "Xtest .+= B2 * Z2test\n",
    "Ytest = Xtest + sig * randn(N, Ltest) # noise\n",
    ";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BCM sparse coding routine\n",
    "# study this function!\n",
    "# argmin_z 0.5*|D Z - X|_F^2 + reg |Z|_0 where columns of D have unit norm\n",
    "function sparse_code_cd(X, D, reg::Number; Z0 = pinv(D) * X, niter = 100)\n",
    "\t(N,K) = size(D)\n",
    "\tR = X - D * Z0 # residual vectors\n",
    "\tZ = copy(Z0)\n",
    "\n",
    "\tfor iter=1:niter # outer loop over iteration\n",
    "\t\tfor k=1:K # inner loop over atoms\n",
    "\t\t\tdk = @view D[:,k]\n",
    "\t\t\tR .+= dk * Z[k,:]'\n",
    "\t\t\tZ[k,:] = hard.(dk'*R, reg)\n",
    "\t\t\tR .-= dk * Z[k,:]'\n",
    "\t\tend\n",
    "\tend\n",
    "\treturn Z\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply BCM sparse coding to test data using the learned dictionary\n",
    "Ztest = sparse_code_cd(Ytest, D, reg/2)\n",
    ";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see how well the sparse coding works on test data with learned D\n",
    "@show norm(D * Ztest - Xtest) / norm(Xtest), sig\n",
    "@show count(Z .!= 0) / length(Z)\n",
    "ls = [3, 6, 8, 12]\n",
    "scatter(Xtest[:,ls], color=:blue, marker=:square, label=\"true\")\n",
    "scatter!(D * Ztest[:,ls], color=:red, label=\"fit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#jim(Xtest[:,1:100]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional extensions if time permits\n",
    "\n",
    "* Experiment with the regularization parameter(s) to try to improve the results.\n",
    "* Try other ways of initializing the dictionary D0.\n",
    "* Try larger values of K since in practice the appropriate K is unknown.\n",
    "* Compare hard thresholding (0-norm) to soft thresholding (1-norm).\n",
    "* Compare the 2-block BCM to the 2N-block BCM used here."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.1.0",
   "language": "julia",
   "name": "julia-1.1"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.1.0"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}