6_Deconvolution/clark_clean_assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "### Clark's Method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "This is a notebook to implement Clark's `CLEAN` method. See $\\S$ 6.3 for details on the method. The pseudo-code for the algorithm is presented below. The goal of this assignment is to implement the major cycle and PSF selection function. The minor cycle uses Högbom's method which has all ready been written. An example of Högbom's `CLEAN` method can be found in Chapter 6. Code needs to be added to the functions `clark()` and `selectSubPSF()` in order for the notebook to work. Marks are given based completeness and generality.\n",
    "\n",
    "Pseudo-code:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "$\\textbf{input: } I^{D}(l,m), \\ \\textrm{PSF}(l,m), \\ \\gamma, \\ f_{\\textrm{thresh}}, \\ N$\n",
    "\n",
    "$\\textbf{initialize: } S^{\\textrm{model}} \\leftarrow \\{\\}, \\ I^{\\textrm{res}} \\leftarrow I^{D}, \\ i \\leftarrow 0, \\ (\\textrm{PSF}_{\\textrm{sub}}(l,m), \\ R_{\\textrm{PSF}}) \\leftarrow g(\\textrm{PSF}(l,m))$\n",
    "\n",
    "$\\textbf{while} \\ \\textrm{any}(I^{\\textrm{res}} > f_{\\textrm{thresh}}) \\ \\textrm{or} \\ i \\leq N \\ \\textbf{do:} \\quad [\\textrm{Major Cycle}]$\n",
    "\n",
    "$\\qquad l_{\\textrm{max}}, m_{\\textrm{max}} \\leftarrow \\underset{l,m}{\\operatorname{argmax}} I^{\\textrm{res}}(l,m)$\n",
    "\n",
    "$\\qquad f_{\\textrm{max}} \\leftarrow I^{D}(l_{\\textrm{max}}, m_{\\textrm{max}})$\n",
    "\n",
    "$\\qquad S^{\\textrm{model}}_{\\textrm{partial}} \\leftarrow \\textrm{Hogbom}(I^{\\textrm{res}}, \\  \\textrm{PSF}_{\\textrm{sub}}, \\ \\gamma, \\ f_{\\textrm{max}} \\cdot R_{\\textrm{PSF}}) \\quad [\\textrm{Minor Cycle}]$\n",
    "\n",
    "$\\qquad V^{\\textrm{model}}_{\\textrm{partial}} \\leftarrow \\mathscr{F}\\{S^{\\textrm{model}}_{\\textrm{partial}}\\}, V^S \\leftarrow \\mathscr{F}\\{\\textrm{PSF}\\}$\n",
    "\n",
    "$\\qquad I^{\\textrm{res}} \\leftarrow I^{\\textrm{res}} - \\mathscr{F}^{-1}\\{V^S \\cdot V^{\\textrm{model}}_{\\textrm{partial}}\\}$\n",
    "\n",
    "$\\qquad S^{\\textrm{model}} \\leftarrow S^{\\textrm{model}} + S^{\\textrm{model}}_{\\textrm{partial}}$\n",
    "\n",
    "$\\qquad i \\leftarrow i +1$\n",
    "\n",
    "$\\textbf{ouput: } S^{\\textrm{model}}, I^{\\textrm{res}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from scipy import optimize\n",
    "\n",
    "from astropy.io import fits\n",
    "import aplpy\n",
    "\n",
    "#Disable astropy/aplpy logging\n",
    "import logging\n",
    "logger0 = logging.getLogger('astropy')\n",
    "logger0.setLevel(logging.CRITICAL)\n",
    "logger1 = logging.getLogger('aplpy')\n",
    "logger1.setLevel(logging.CRITICAL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def subtractPSF(img, psf, l, m, flux, gain):\n",
    "    \"\"\"Subtract the PSF (attenuated by the gain and flux) centred at (l,m) from an image\"\"\"\n",
    "    \n",
    "    #get the half lengths of the PSF\n",
    "    if (psf.shape[0] % 2 == 0): psfMidL = psf.shape[0]/2 #even\n",
    "    else: psfMidL = (psf.shape[0]+1)/2 #odd\n",
    "    if (psf.shape[1] % 2 == 0): psfMidM = psf.shape[1]/2 #even\n",
    "    else: psfMidM = (psf.shape[1]+1)/2 #odd\n",
    "    \n",
    "    #determine limits of sub-images\n",
    "    #starting m\n",
    "    if m-psfMidM < 0:\n",
    "        subM0 = 0\n",
    "        subPSFM0 = psfMidM-m\n",
    "    else:\n",
    "        subM0 = m-psfMidM\n",
    "        subPSFM0 = 0\n",
    "    #starting l\n",
    "    if l-psfMidL < 0:\n",
    "        subL0 = 0\n",
    "        subPSFL0 = psfMidL-l\n",
    "    else:\n",
    "        subL0 = l-psfMidL\n",
    "        subPSFL0 = 0\n",
    "    #ending m\n",
    "    if img.shape[1] > m+psfMidM:\n",
    "        subM1 = m+psfMidM\n",
    "        subPSFM1 = psf.shape[1]\n",
    "    else:\n",
    "        subM1 = img.shape[1]\n",
    "        subPSFM1 = psfMidM + (img.shape[1]-m)\n",
    "    #ending l\n",
    "    if img.shape[0] > l+psfMidL:\n",
    "        subL1 = l+psfMidL\n",
    "        subPSFL1 = psf.shape[0]\n",
    "    else:\n",
    "        subL1 = img.shape[0]\n",
    "        subPSFL1 = psfMidL + (img.shape[0]-l)\n",
    "    \n",
    "    #select subset of image\n",
    "    #subImg = img[subL0:subL1, subM0:subM1]\n",
    "    #select subset of PSF\n",
    "    subPSF = psf[subPSFL0:subPSFL1, subPSFM0:subPSFM1]\n",
    "    \n",
    "    #subtract PSF centred on (l,m) position\n",
    "    img[subL0:subL1, subM0:subM1] -= flux * gain * psf[subPSFL0:subPSFL1, subPSFM0:subPSFM1]\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def gauss2D(x, y, amp, meanx, meany, sigmax, sigmay):\n",
    "    \"\"\"2D Gaussian Function\"\"\"\n",
    "    gx = -(x - meanx)**2/(2*sigmax**2)\n",
    "    gy = -(y - meany)**2/(2*sigmay**2)\n",
    "    \n",
    "    return amp * np.exp( gx + gy)\n",
    "\n",
    "def err(p, xx, yy, data):\n",
    "    \"\"\"Error function for least-squares fitting\"\"\"\n",
    "    return gauss2D(xx.flatten(), yy.flatten(), *p) - data.flatten()\n",
    "\n",
    "def idealPSF(psfImg):\n",
    "    \"\"\"Determine the ideal PSF size based on the observing PSF doing a simple 2D Gaussian least-squares fit\"\"\"\n",
    "    xx, yy = np.meshgrid(np.arange(0, psfImg.shape[0]), np.arange(0, psfImg.shape[1]))\n",
    "    # Initial estimate: PSF should be amplitude 1, and usually imaging over sample the PSF by 3-5 pixels\n",
    "    params0 = 1., psfImg.shape[0]/2., psfImg.shape[1]/2., 3., 3.\n",
    "    params, pcov, infoDict, errmsg, sucess = optimize.leastsq(err, params0, \\\n",
    "                            args=(xx.flatten(), yy.flatten(), psfImg.flatten()), full_output=1)\n",
    "    #fwhm = [2.*np.sqrt(2.*np.log(2.)) * params[3], 2.*np.sqrt(2.*np.log(2.)) * params[4]]\n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def restoreImg(skyModel, residImg, params):\n",
    "    \"\"\"Generate a restored image from a deconvolved sky model, residual image, ideal PSF params\"\"\"\n",
    "    mdlImg = np.zeros_like(residImg)\n",
    "    for l,m, flux in skyModel:\n",
    "        mdlImg[l,m] += flux\n",
    "    \n",
    "    #generate an ideal PSF image\n",
    "    psfImg = np.zeros_like(residImg)\n",
    "    xx, yy = np.meshgrid(np.arange(0, psfImg.shape[0]), np.arange(0, psfImg.shape[1]))\n",
    "    psfImg = gauss2D(xx, yy, params[0], params[1], params[2], params[3], params[4])\n",
    "    \n",
    "    #convolve ideal PSF with model image\n",
    "    sampFunc = np.fft.fft2(psfImg) #sampling function\n",
    "    mdlVis = np.fft.fft2(mdlImg) #sky model visibilities\n",
    "    sampMdlVis = sampFunc * mdlVis #sampled sky model visibilities\n",
    "    convImg = np.abs(np.fft.fftshift(np.fft.ifft2(sampMdlVis))) + residImg #sky model convolved with PSF\n",
    "    \n",
    "    #return mdlImg + residImg\n",
    "    return convImg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def hogbom(dirtyImg, psfImg, gain, niter, fthresh):\n",
    "    \"\"\"Implementation of Hogbom's CLEAN method\n",
    "    inputs:\n",
    "    dirtyImg: 2-D numpy array, dirty image\n",
    "    psfImg: 2-D numpy array, PSF\n",
    "    gain: float, gain factor, 0 < gain < 1\n",
    "    niter: int, maximum number of iterations to halt deconvolution\n",
    "    fthresh: flaot, maximum flux threshold to halt deconvolution\n",
    "    \n",
    "    outputs:\n",
    "    residImg: 2-D numpy array, residual image\n",
    "    skyModel: list of sky model components, each component is a list [l, m, flux]\n",
    "    \"\"\"\n",
    "    #Initalization\n",
    "    skyModel = [] #initialize empty model\n",
    "    residImg = np.copy(dirtyImg) #copy the dirty image to the initial residual image\n",
    "    i = 0 #number of iterations\n",
    "    print '\\tMinor Cycle: fthresh: %f'%fthresh\n",
    "    \n",
    "    #CLEAN iterative loop\n",
    "    while np.max(residImg) > fthresh and i < niter:\n",
    "        lmax, mmax = np.unravel_index(residImg.argmax(), residImg.shape) #get pixel position of maximum value\n",
    "        fmax = residImg[lmax, mmax] #flux value of maximum pixel\n",
    "        #print 'iter %i, (l,m):(%i, %i), flux: %f'%(i, lmax, mmax, fmax)\n",
    "        residImg = subtractPSF(residImg, psfImg, lmax, mmax, fmax, gain)\n",
    "        skyModel.append([lmax, mmax, gain*fmax])\n",
    "        i += 1\n",
    "    \n",
    "    return residImg, skyModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def selectSubPSF(psfImg):\n",
    "    \"\"\"Select out a central region of the PSF for the Hogbom minor cycle,\n",
    "    and compute the ratio between the main lobe and highest sidelobe\n",
    "    \n",
    "    There are a number of ways to implement this function. A general method\n",
    "    would determine where the first sidelobe is, and select out a region of\n",
    "    the PSF which includes those sidelobes. A simple method would be to hard-code\n",
    "    the size of the sub-image. Marks will be given based on how general the\n",
    "    function is.\n",
    "    \n",
    "    inputs:\n",
    "    psfImg: 2-D numpy array, PSF\n",
    "    \n",
    "    outputs:\n",
    "    subPsfImg: 2-D numpy array, subset of the PSF image which contains the main lobe out to the \n",
    "        largest sidelobe.\n",
    "    peakRatio: float, peakRatio < 1, ratio of the largest sidelobe to the main lobe\n",
    "    \"\"\"\n",
    "    \n",
    "    #ADD CODE HERE\n",
    "    \n",
    "    #return sub-PSF and ratio\n",
    "    return subPsfImg, peakRatio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def clark(dirtyImg, psfImg, gain, niter, fthresh):\n",
    "    \"\"\"Implementation of Clark's CLEAN method\n",
    "    inputs:\n",
    "    dirtyImg: 2-D numpy array, dirty image\n",
    "    psfImg: 2-D numpy array, PSF\n",
    "    gain: float, gain factor, 0 < gain < 1\n",
    "    niter: int, maximum number of iterations to halt deconvolution\n",
    "    fthresh: flaot, maximum flux threshold to halt deconvolution\n",
    "    \n",
    "    outputs:\n",
    "    residImg: 2-D numpy array, residual image\n",
    "    skyModel: list of sky model components, each component is a list [l, m, flux]\n",
    "    \"\"\"\n",
    "    \n",
    "    #Initalization\n",
    "    skyModel = [] #initialize empty model\n",
    "    residImg = np.copy(dirtyImg) #copy the dirty image to the initial residual image\n",
    "    i = 0 #number of iterations\n",
    "    \n",
    "    #select subset of PSF\n",
    "    subPsfImg, peakRatio = selectSubPSF(psfImg) #ADD CODE: this function needs to be written\n",
    "    \n",
    "    #CLEAN iterative Major cycle\n",
    "    while np.max(residImg) > fthresh and i < niter:\n",
    "        #ADD CODE: get flux and pixel position of maximum value\n",
    "        \n",
    "        #ADD CODE: minor cycle, run partial Hogbom, return partial sky model\n",
    "        \n",
    "        #ADD CODE: Fourier transform partial sky model to model visibilities\n",
    "        \n",
    "        #ADD CODE: compute sampling function from PSF\n",
    "        \n",
    "        #ADD CODE: subtract sampled model image from the residual image\n",
    "\n",
    "        #ADD CODE: update full sky model\n",
    "\n",
    "        i += 1\n",
    "    \n",
    "    return residImg, skyModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#input parameters\n",
    "gain = 0.1 #loop gain, range: 0 < gain < 1\n",
    "niter = 20 #number of loop iterations\n",
    "fthresh = 2.5 #minimum flux threshold to deconvolve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#input images: dirty, PSF\n",
    "#assuming unpolarized, single frequency image\n",
    "fh = fits.open('data/KAT-7_6h60s_dec-30_10MHz_10chans_uniform_n100-dirty.fits')\n",
    "dirtyImg = fh[0].data[0,0]\n",
    "fh = fits.open('data/KAT-7_6h60s_dec-30_10MHz_10chans_uniform_n100-psf.fits')\n",
    "psfImg = fh[0].data[0,0]\n",
    "idealPSFparams = idealPSF(psfImg) #compute ideal PSF parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#run deconvolution\n",
    "residImg, skyModel = clark(dirtyImg, psfImg, gain, niter, fthresh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#plot the dirty image\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "plt.imshow(dirtyImg)\n",
    "plt.title('Dirty Image')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#plot the residual image\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "plt.imshow(residImg)\n",
    "plt.title('Residual Image')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#plot the restored image\n",
    "restImg = restoreImg(skyModel, residImg, idealPSFparams)\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "plt.imshow(restImg)\n",
    "plt.title('Restored Image')\n",
    "plt.colorbar()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}