{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -r3 -n1\n",
    "from numpy import empty,arange,exp,real,imag,pi,array,place,argwhere\n",
    "from numpy.fft import rfft,irfft,fft2,ifft2\n",
    "from scipy.fftpack import dct,idct\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.sparse import coo_matrix\n",
    "\n",
    "# Numpy does not have DCT directly, so we have to do it manually - note that scipy does have it, but this is a useful teaching example\n",
    "######################################################################\n",
    "# 1D DCT Type-II function\n",
    "\n",
    "# def dct(y):\n",
    "#     N = len(y)\n",
    "#     y2 = empty(2*N,float) # create an empty array for the function\n",
    "#     y2[:N] = y[:] # the first N elements of y2 are taken from y\n",
    "#     y2[N:] = y[::-1] # the elements after N are reversed giving us a symmetric, even function\n",
    "\n",
    "#     c = rfft(y2) # perform the fourier transform. This uses a numpy function (Real Fast Fourier Transform), otherwise everything takes too long - we will discuss it later.\n",
    "#     phi = exp(-1j*pi*arange(N)/(2*N)) # create an array with terms from (0-999) and use it to define the leading phase factor\n",
    "#     return real(phi*c[:N]) # return the real part of the transform\n",
    "\n",
    "######################################################################\n",
    "# 2D DCT function\n",
    "\n",
    "# def dct2(y):\n",
    "#     M = y.shape[0] # first dimension of input function, use to define limits\n",
    "#     N = y.shape[1] # second dimension of input function, use to define limits\n",
    "#     a = empty([M,N],float) # create arrays to store our result - a for the first transform and b for the second\n",
    "#     b = empty([M,N],float)\n",
    "\n",
    "# # run the 1D transforms\n",
    "#     for i in range(M):\n",
    "#         a[i,:] = dct(y[i,:])\n",
    "#     for j in range(N):\n",
    "#         b[:,j] = dct(a[:,j])\n",
    "\n",
    "#     return b\n",
    "\n",
    "######################################################################\n",
    "# 1D inverse DCT Type-II function\n",
    "\n",
    "# def idct(a):\n",
    "#     N = len(a)\n",
    "#     c = empty(N+1,complex)\n",
    "\n",
    "#     phi = exp(1j*pi*arange(N)/(2*N))\n",
    "#     c[:N] = phi*a\n",
    "#     c[N] = 0.0\n",
    "#     return irfft(c)[:N] # (Inverse Real Fast Fourier Transform)\n",
    "\n",
    "######################################################################\n",
    "# 2D inverse DCT function\n",
    "\n",
    "# def idct2(b):\n",
    "#     M = b.shape[0]\n",
    "#     N = b.shape[1]\n",
    "#     a = empty([M,N],float)\n",
    "#     y = empty([M,N],float)\n",
    "\n",
    "#     for i in range(M):\n",
    "#         a[i,:] = idct(b[i,:])\n",
    "#     for j in range(N):\n",
    "#         y[:,j] = idct(a[:,j])\n",
    "\n",
    "#     return y\n",
    "\n",
    "adam = plt.imread('images/photo.tiff') # Read in TIFF picture file. TIFF is lossless, so generally very large compared to compressed formats.\n",
    "\n",
    "singch_adam = adam[:, :, 0] # slice image into a single channel so the transform can handle it.\n",
    "plt.figure(figsize=(12,6)) # set the figsize\n",
    "\n",
    "# plot the original image\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(singch_adam, cmap=\"plasma\") # a default colourmap is applied, there are no colour channels in the data anymore\n",
    "plt.colorbar()\n",
    "adam_dct = fft2(singch_adam) # apply the 2D DCT to the single channel image\n",
    "\n",
    "# plot the DCT of the image, limiting the scale so we can actually see something\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(real(adam_dct), cmap=\"plasma\",clim=(0, 1e5))\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "# define a threshold for cutting the coefficients\n",
    "cutoff_amplitude = 1e6 # \n",
    "place(adam_dct, abs(adam_dct) < cutoff_amplitude, [0])\n",
    "original_size = adam_dct.shape[0]*adam_dct.shape[1]\n",
    "components_removed = argwhere(abs(adam_dct) < cutoff_amplitude).shape[0] # Replace values below cutoff with 0\n",
    "\n",
    "# give some nice output on how much we have cut from the DCT\n",
    "percent_removed = (1-components_removed/original_size)*100\n",
    "print(\"Removed {} out of {} components. Image contains {:.2f}% of the original components.\".format(components_removed, original_size, percent_removed))\n",
    "\n",
    "plt.figure(figsize=(12,6)) # set the figsize\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(real(adam_dct), cmap=\"plasma\",clim=(0, 1e5))\n",
    "plt.colorbar()\n",
    "\n",
    "adam_idct = ifft2(adam_dct) # run the inverse DCT on the coefficients to generate the \"compressed\" image\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(real(adam_idct), cmap=\"plasma\")\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "# now actually remove all the zeroes so that we are really handling a smaller array\n",
    "adam_compressed = coo_matrix(adam_dct)\n",
    "print(\"Now we have really reduced the array to {:.2f}% of its original size.\".format(100*size(adam_compressed)/size(adam_dct)))\n",
    "\n",
    "# but to actually do anything with it, we need to decompress it and see whether we lost any data.\n",
    "adam_decomp = adam_compressed.toarray()\n",
    "\n",
    "plt.figure(figsize=(12,6)) # set the figsize\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(real(adam_decomp), cmap=\"plasma\",clim=(0, 1e5))\n",
    "plt.colorbar()\n",
    "\n",
    "adam_idct_decomp = ifft2(adam_decomp)\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(real(adam_idct_decomp), cmap=\"plasma\")\n",
    "plt.colorbar()\n",
    "plt.show()\n",
    "\n",
    "plt.imsave('images/photo_compressed',real(adam_idct_decomp), format='tiff') # but this is still the same size, as we had to decompress it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3.76 s ± 152 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)\n",
    "3.45 s ± 27.9 ms per loop (mean ± std. dev. of 3 runs, 1 loop each) with scipy DCT\n",
    "2.57 s ± 37.9 ms per loop (mean ± std. dev. of 3 runs, 1 loop each) with scipy DCT and iDCT\n",
    "2.54 s ± 70.2 ms per loop (mean ± std. dev. of 3 runs, 1 loop each) with numpy FFT and iFFT using real components"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}