{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallel Magnetic Resonance imaging (MRI)\n",
    "\n",
    "Standard MRI is described by the Fourier transform $\\mathcal{F}$ as forward operator (here in two dimensions). \n",
    "To accelerate data acquisition, parallel MRI uses simultaneous measurements by $N$ receiver coils. This allows undersampling \n",
    "of the Fourier domain leading to speed-ups. Parallel MRI is described by the forward operator\n",
    "$$F\\left(\\begin{array}{c}\\rho \\\\ c_1\\\\ \\vdots \\\\ c_N\\end{array}\\right) \n",
    "= \\left(\\begin{array}{c}M\\cdot\\mathcal{F}(c_1 \\cdot \\rho)\\\\ \\vdots \\\\\n",
    "M\\cdot\\mathcal{F}(c_N \\cdot \\rho)\\end{array}\\right).$$\n",
    "Here $\\rho$ describes the hydrogen density and is the main quantity of interest. To take into account effects such as motion artifacts, $\\rho$ has to be modeled as a complex-valued function. $c_1,\\dots, c_N$ describe complex-valued coil profile, which may be assumed to be smooth. As they depend on the sample $\\rho$, they must be reconstructed together with $\\rho$. $M$ is a 0-1-mask describing the undersampling pattern.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mplib\n",
    "from matplotlib.colors import hsv_to_rgb\n",
    "import numpy as np\n",
    "from scipy.io import loadmat\n",
    "\n",
    "import regpy.stoprules as rules\n",
    "import regpy.util as uti\n",
    "from examples.mri.mri import cartesian_sampling, normalize, parallel_mri, sobolev_smoother, estimate_sampling_pattern\n",
    "from regpy.operators import PtwMultiplication\n",
    "from regpy.solvers import RegularizationSetting\n",
    "from regpy.solvers.nonlinear.irgnm import IrgnmCG\n",
    "from regpy.vecsps import UniformGridFcts\n",
    "from regpy.hilbert import L2\n",
    "import os\n",
    "\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format='%(asctime)s %(levelname)s %(name)-40s :: %(message)s'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complex to rgb conversion\n",
    "\n",
    "Converts array of complex numbers into array of RGB color values for plotting. The hue corresponds to the argument.\n",
    "The brighntess corresponds to the absolute value.  \n",
    "\n",
    "`Parameters`\n",
    ">z : `numpy.ndarray`\n",
    ">array of complex numbers\n",
    "\n",
    "`Returns`\n",
    "> `numpy.ndarray`\n",
    "> Array that contains three values for each value in z containing the RGB representation of this value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def complex_to_rgb(z):\n",
    "    HSV = np.dstack( (np.mod(np.angle(z)/(2.*np.pi),1), 1.0*np.ones(z.shape), np.abs(z)/np.max((np.abs(z[:]))), ))\n",
    "    return hsv_to_rgb(HSV)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data from file and estimate sampling pattern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = loadmat('data/ksp3x2.mat')['Y']\n",
    "data = np.transpose(data,(2,0,1))*(100/np.linalg.norm(data))\n",
    "# normalize and transpose data \n",
    "nrcoils,n1,n2 = data.shape\n",
    "grid = UniformGridFcts((-1, 1, n1), (-1, 1, n2), dtype=complex)\n",
    "mask = estimate_sampling_pattern(data)\n",
    "plt.imshow(mask.T); plt.title('Undersampling pattern of data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up forward operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sobolev_index = 32\n",
    "\n",
    "full_mri_op = parallel_mri(grid=grid, ncoils=nrcoils,centered=True)\n",
    "sampling = PtwMultiplication(full_mri_op.codomain,(1.+0j)* mask)\n",
    "#cartesian_sampling(full_mri_op.codomain, mask=mask)\n",
    "smoother = sobolev_smoother(full_mri_op.domain, sobolev_index, factor=220.)\n",
    "\n",
    "parallel_mri_op = sampling * full_mri_op * smoother"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up initial guess\n",
    "We use constant density and zero coil profiles as initial guess."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init = parallel_mri_op.domain.zeros()\n",
    "init_density, _ = parallel_mri_op.domain.split(init)\n",
    "init_density[...] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up regularization method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "setting = RegularizationSetting(op=parallel_mri_op, penalty=L2, data_fid=L2)\n",
    "\n",
    "solver = IrgnmCG(\n",
    "    setting=setting,\n",
    "    data=data,\n",
    "    regpar=1,\n",
    "    regpar_step=1/3.,\n",
    "    init=init\n",
    ")\n",
    "\n",
    "stoprule = rules.CountIterations(max_iterations=5) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run solver by hand and plot iterates\n",
    "Get an iterator from the solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "it = iter(solver)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate this cell by hand about 5 times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "reco, reco_data = next(it)\n",
    "\n",
    "rho, coils = smoother.codomain.split(smoother(reco))\n",
    "#rho, coils = normalize(rho,coils)\n",
    "\n",
    "fig = plt.figure(figsize = (15,9))\n",
    "\n",
    "gs = fig.add_gridspec(3,7)\n",
    "axs = [fig.add_subplot(gs[0:3, 0:3])]\n",
    "im0 = axs[0].imshow(np.abs(rho),cmap=mplib.colormaps['Greys_r'],origin='lower')\n",
    "axs[0].xaxis.set_ticklabels([])\n",
    "axs[0].yaxis.set_ticklabels([])\n",
    "for j in range(3):\n",
    "    for k in range(3,7):\n",
    "        axs.append(fig.add_subplot(gs[j,k]))\n",
    "        axs[-1].xaxis.set_ticklabels([])\n",
    "        axs[-1].yaxis.set_ticklabels([])\n",
    "for j in range(nrcoils):\n",
    "    axs[1+j].imshow(complex_to_rgb(coils[j,:,:]),origin='lower')\n",
    "plt.show()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ... or run solver by stopping rule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for reco, reco_data in solver.while_(stoprule):\n",
    "    rho, coils = smoother.codomain.split(smoother(reco))\n",
    "    #rho, coils = normalize(rho,coils)\n",
    "\n",
    "    fig = plt.figure(figsize = (15,9))\n",
    "\n",
    "    gs = fig.add_gridspec(3,7)\n",
    "    axs = [fig.add_subplot(gs[0:3, 0:3])]\n",
    "    axs[0].imshow(np.abs(rho),cmap=mplib.colormaps['Greys_r'],origin='lower')\n",
    "    axs[0].xaxis.set_ticklabels([])\n",
    "    axs[0].yaxis.set_ticklabels([])\n",
    "    for j in range(3):\n",
    "        for k in range(3,7):\n",
    "            axs.append(fig.add_subplot(gs[j,k]))\n",
    "            axs[-1].xaxis.set_ticklabels([])\n",
    "            axs[-1].yaxis.set_ticklabels([])\n",
    "    for j in range(nrcoils):\n",
    "        axs[1+j].imshow(complex_to_rgb(coils[j,:,:]),origin='lower')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ngsolve",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}