{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data\n",
    "\n",
    "Load complex image data from MATLAB mat-file (v7.3 or later), which is actually an HDF5 format\n",
    "\n",
    "Complex data is loaded as a (real, imag) tuple, so it neds to be explicitly converted to complex double\n",
    "\n",
    "In this section:\n",
    "- using h5py module\n",
    "- np.transpose\n",
    "- 1j as imaginary constant\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get hdf5 object for the mat-file\n",
    "h = h5py.File('data.mat','r')\n",
    "\n",
    "# get img variable from the mat-file\n",
    "dat = h.get('img')\n",
    "\n",
    "# turn array of (real, imag) tuples into an array of complex doubles\n",
    "# transpose to keep data in same orientation as MATLAB\n",
    "img = np.transpose(dat['real'] + 1j*dat['imag'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6/8 Partial Fourier sampling\n",
    "\n",
    "Fourier transform the image to get k-space data, and add complex Gaussian noise\n",
    "\n",
    "To simulate 6/8 Partial Fourier sampling, zero out the bottom 1/4 of k-space\n",
    "\n",
    "In this section:\n",
    "- np.random.randn\n",
    "- np.fft\n",
    "- 0-based indexing\n",
    "- image plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate normally-distributed complex noise\n",
    "n = np.random.randn(96,96) + 1j*np.random.randn(96,96)\n",
    "\n",
    "# Fourier transform the image and add noise\n",
    "y = np.fft.fftshift(np.fft.fft2(img), axes=0) + n\n",
    "\n",
    "# set bottom 24/96 lines to 0\n",
    "y[72:,:] = 0\n",
    "\n",
    "# show sampling\n",
    "_, ax = plt.subplots()\n",
    "ax.imshow(np.log(np.abs(np.fft.fftshift(y, axes=1))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Estimate phase\n",
    "\n",
    "Filter the k-space data and extract a low-resolution phase estimate\n",
    "\n",
    "Filtering can help reduce ringing in the phase image\n",
    "\n",
    "In this section:\n",
    "- np.pad\n",
    "- np.hanning\n",
    "- reshaping 1D array to 2D array using np.newaxis (or None)\n",
    "- subplots with titles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create zero-padded hanning filter for ky-filtering\n",
    "filt = np.pad(np.hanning(48),24,'constant')\n",
    "\n",
    "# reshape 1D array into 2D array\n",
    "filt = filt[:,np.newaxis]\n",
    "# or\n",
    "# filt = filt[:,None]\n",
    "\n",
    "# generate low-res image with inverse Fourier transform\n",
    "low = np.fft.ifft2(np.fft.ifftshift(y*filt, axes=0))\n",
    "\n",
    "# get phase image\n",
    "phs = np.exp(1j*np.angle(low))\n",
    "\n",
    "# show phase estimate alongside true phase\n",
    "_, ax = plt.subplots(1,2)\n",
    "ax[0].imshow(np.angle(img))\n",
    "ax[0].set_title('True image phase')\n",
    "ax[1].imshow(np.angle(phs))\n",
    "ax[1].set_title('Estimated phase')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# POCS reconstruction\n",
    "\n",
    "Perform the projection-onto-convex-sets (POCS) partial Fourier reconstruction method.\n",
    "\n",
    "POCS is an iterative scheme estimates the reconstructed image as any element in the intersection of the following two (convex) sets:\n",
    "1. Set of images consistent with the measured data\n",
    "2. Set of images that are non-negative real\n",
    "\n",
    "This requires prior knowledge of the image phase (hence the estimate above), and it works because although we have less than a full k-space of measurements, we're now only estimating half the number of free parameters (real values only, instead of real + imag), and we're no longer under-determined. Equivalently, consider the fact that real-valued images have conjugate symmetric k-spaces, so we only require half of k-space to reconstruct our image.\n",
    "\n",
    "In this section:\n",
    "- np.zeros\n",
    "- range() builtin\n",
    "- point-wise multiplication (*)\n",
    "- np.fft operations default to last axis, not first\n",
    "- np.maximum vs np.max\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise image estimate to be zeros\n",
    "est = np.zeros((96,96))\n",
    "\n",
    "# set the number of iterations \n",
    "iters = 10\n",
    "\n",
    "# each iteration cycles between projections\n",
    "for i in range(iters):\n",
    "# projection onto data-consistent set:\n",
    "    # use a-priori phase to get complex image\n",
    "    est = est*phs\n",
    "    \n",
    "    # Fourier transform to get k-space\n",
    "    est = np.fft.fftshift(np.fft.fft2(est), axes=0)\n",
    "    \n",
    "    # replace data with measured lines\n",
    "    est[:72,:] = y[:72,:]\n",
    "    \n",
    "    # inverse Fourier transform to get back to image space\n",
    "    est = np.fft.ifft2(np.fft.ifftshift(est, axes=0))\n",
    "\n",
    "# projection onto non-negative reals:\n",
    "    # remove a-priori phase\n",
    "    est = est*np.conj(phs)\n",
    "    \n",
    "    # get real part\n",
    "    est = np.real(est)\n",
    "\n",
    "    # ensure output is non-negative\n",
    "    est = np.maximum(est, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Display error and plot reconstruction\n",
    "\n",
    "The POCS reconstruction is compared to a zero-filled reconstruction (i.e., where the missing data is zeroed prior to inverse Fourier Transform)\n",
    "\n",
    "In this section:\n",
    "- print formatted strings to standard output\n",
    "- 2D subplots with min/max scales, figure size\n",
    "- np.sum sums over all elements by default"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute zero-filled recon\n",
    "zf = np.fft.ifft2(np.fft.ifftshift(y, axes=0))\n",
    "\n",
    "# compute rmse for zero-filled and POCS recon\n",
    "err_zf = np.sqrt(np.sum(np.abs(zf - img)**2))\n",
    "err_pocs = np.sqrt(np.sum(np.abs(est*phs - img)**2))\n",
    "\n",
    "# print errors\n",
    "print(f'RMSE for zero-filled recon: {err_zf}')\n",
    "print(f'RMSE for POCS recon: {err_pocs}')\n",
    "\n",
    "# plot both recons side-by-side\n",
    "_, ax = plt.subplots(2,2,figsize=(16,16))\n",
    "\n",
    "# plot zero-filled\n",
    "ax[0,0].imshow(np.abs(zf), vmin=0, vmax=1)\n",
    "ax[0,0].set_title('Zero-filled')\n",
    "ax[1,0].plot(np.abs(zf[:,47]))\n",
    "\n",
    "# plot POCS\n",
    "ax[0,1].imshow(est, vmin=0, vmax=1)\n",
    "ax[0,1].set_title('POCS recon')\n",
    "ax[1,1].plot(np.abs(est[:,47]))"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}