HandsOn_ISMRM_MoCo_Workshop_solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NjQ_TUENOqus"
   },
   "source": [
    "# **Hands-on Motion Estimation and Correction**\n",
    "*by [Gastao Cruz](https://twitter.com/gastaocruz) and [Thomas Kuestner](https://twitter.com/kuestnerthomas)*\n",
    "\n",
    "In this hands-on of the [ISMRM Workshop on Motion Detection and Correction 2022](https://www.ismrm.org/workshops/2022/Motion/), we aim to convey the principles of motion artifacts, their appearance in the MR image, means of estimating motion (conventional and deep learning) and correcting for the induced motion artifacts during reconstruction.\n",
    "\n",
    "This notebook is based on a [Github repository](https://github.com/lab-midas/ismrm-moco-workshop).\n",
    "\n",
    "## Getting started\n",
    "First you need to run the\n",
    "- **[0. Prerequisites](#Prerequisites)** which [installs the Python packages and libraries](#Installation), [imports the modules](#Imports) and [loads the data](#DataLoading) (see Provided material). In this hands-on, we will only focus on 2D single-slice brain MR image. If you want to try out the other datasets, you may need to adapt some functions, as only 2D processing is supported.\n",
    "\n",
    "Afterwards you can decide to run either (or all) of the following three topics (in any order):\n",
    "- **[1. Motion artifact appearance](#MotionArtifactAppearance)**: Learn something about the spatial and temporal behaviour of motion and how it manifests in the image for Cartesian and non-Cartesian imaging\n",
    "- **[2. Motion estimation / image registration](#MotionEstimation)**: Learn how to perform an image registration using conventional and deep learning based solutions\n",
    "- **[3. Motion-compensated image reconstruction](#MotionCompensated)**: Learn to incorporate motion estimation into image reconstruction to obtain a motion-corrected image\n",
    "\n",
    "## Provided material\n",
    "Two datasets (brain and heart) are distributed with this hands-on tutorial:\n",
    "- `brain_large.npz` (GIT LFS): multi-coil 3D T1w brain MR image; matrix size = 256 x 216 x 65 x 7 (X x Y x Slices x Channels)\n",
    "- `brain_slice.npz`: single central slice of `brain_large.npz`; matrix size = 256 x 216 x 7 (X x Y x Channels)\n",
    "- `heart_large.npz` (GIT LFS): multi-coil, multi-slice 2D cardiac CINE MR image; matrix size = 176 x 132 x 12 x 25 x 10 (X x Y x Slices x Cardiac Phases x Channels)\n",
    "\n",
    "**Attention:** The following codes are only defined for 2D processing. If you plan to use them for a 3D dataset, you would first need to extend some of the functions below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MuDbhyqpbRg7"
   },
   "source": [
    "# 0. Prerequisites\n",
    "<a name=\"Prerequisites\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7nhY_ySkO5UB"
   },
   "source": [
    "**Attention:** If you are running this code in Google Colab, at the end of the installation you should receive a message that the kernel has crashed and restarted. This kernel restart is required so that all installed libraries work correctly.\n",
    "\n",
    "Afterwards you can continue with [Imports](#Imports)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ey41-4MfKGs3"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "export USE_OPTOX=false  # [false | true]: optox is only needed if you want to work with the MERLIN provided functions for conjugate gradient etc.\n",
    "export MIN_INSTALL=true  # [false | true]: minimal installation for a faster install on Google Colab, but no radial imaging examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5XF7IHU9n35Z"
   },
   "source": [
    "## Installation\n",
    "<a name=\"Installation\"></a>\n",
    "Ensure Python 3.8 is used for best compatibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AbiYJgnTHXrj",
    "outputId": "b7a9cb29-11f7-41ab-b833-507937039fbe"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# apt-get install doxygen\n",
    "apt-get update -y\n",
    "apt-get install python3.8 python3.8-distutils\n",
    "update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 1\n",
    "update-alternatives --config python3\n",
    "apt-get install python3-pip\n",
    "python3 -m pip install --upgrade pip --user\n",
    "apt-get install libnfft3-dev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8p-WnaMaa_ks"
   },
   "source": [
    "### Python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "89Qi4m_2eqRq",
    "outputId": "8b765770-361a-4156-8376-4eec5035b790"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# package itself\n",
    "if [ -d \"/content/ismrm-moco-workshop\" ]; then\n",
    "  cd /content/ismrm-moco-workshop\n",
    "  git pull\n",
    "else\n",
    "  git clone https://github.com/lab-midas/ismrm-moco-workshop.git\n",
    "fi\n",
    "cp -r /content/ismrm-moco-workshop/utils /content/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "G8SNZO9LaSUy",
    "outputId": "4c4f78a5-684b-4ca5-e5d8-9ebc93484764"
   },
   "outputs": [],
   "source": [
    "# install dependencies\n",
    "! pip install --upgrade pip\n",
    "! pip install simpleitk medutils-mri scikit-image voxelmorph # imjoy\n",
    "! git clone https://github.com/voxelmorph/voxelmorph.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kvgc4nXKc22A",
    "outputId": "314fbbc3-8c54-4320-ee03-24fd4fde566f"
   },
   "outputs": [],
   "source": [
    "# Download ESPIRiT code for coil sensitivity map estimation\n",
    "!git clone https://github.com/mikgroup/espirit-python.git\n",
    "!cp /content/espirit-python/espirit.py ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qdq3_xRJiTj1"
   },
   "source": [
    "### GPUNUFFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "92POvtQ_f25x"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# setup gpunufft\n",
    "if [ \"$MIN_INSTALL\" = false ] ; then\n",
    "  git clone --branch cuda_streams https://github.com/khammernik/gpuNUFFT.git\n",
    "  export CUDA_ROOT_DIR=/usr/local/cuda\n",
    "  #!export CUDA_SDK_ROOT_DIR=/usr/local/cuda-samples\n",
    "  export LDFLAGS=-L/usr/local/cuda/lib64\n",
    "  mkdir -p gpuNUFFT/CUDA/build\n",
    "  cd gpuNUFFT/CUDA/build\n",
    "  cmake .. -DGEN_MEX_FILES=OFF\n",
    "  make\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_ess6AuioOl"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if [ \"$MIN_INSTALL\" = false ] ; then\n",
    "  apt install libnfft3-dev\n",
    "  pip install gpuNUFFT pynfft2\n",
    "  pip install python-pysap\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fgfwrJBljrNR"
   },
   "source": [
    "### Optox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qqLcXkCp2r7u",
    "outputId": "62d437aa-26f4-468c-826d-9dea8ff6c7d3"
   },
   "outputs": [],
   "source": [
    "# Ensure only tensorflow with GPU support is installed\n",
    "! pip uninstall --yes tensorflow\n",
    "! pip install tensorflow-gpu==2.8  # bug in cuDNN 8 + TF2.9 for Colab (crashes voxelmorph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-g8b9HahzuIE",
    "outputId": "c7529ac6-6b18-48e6-a642-9073356d7149"
   },
   "outputs": [],
   "source": [
    "#!pip install tensorflow-gpu==2.8\n",
    "!apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q5GXbPLzBOuw"
   },
   "outputs": [],
   "source": [
    "# set some environment variables in Google Colab that are not automatically handled\n",
    "import os\n",
    "import tensorflow as tf\n",
    "\n",
    "os.environ[\"Tensorflow_INCLUDE_DIRS\"] = str(tf.sysconfig.get_include())\n",
    "os.environ[\"Tensorflow_LIB_DIRS\"] = str(tf.sysconfig.get_lib())\n",
    "os.environ[\"Tensorflow_LIB\"] = str(tf.sysconfig.get_link_flags()[1])\n",
    "os.environ[\"Tensorflow_FLAGS\"] = str(tf.sysconfig.get_compile_flags()[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q_0o6k19fgXJ"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "# setup optox\n",
    "if [ \"$USE_OPTOX\" = true ] ; then\n",
    "  git clone https://github.com/midas-tum/optox.git\n",
    "  export GPUNUFFT_ROOT_DIR=/content/gpuNUFFT\n",
    "  export LDFLAGS=-L/usr/local/cuda/lib64\n",
    "  cd /content/optox\n",
    "  pip install -r requirements.txt\n",
    "  mkdir -p /content/optox/build\n",
    "  cd /content/optox/build\n",
    "  cmake .. -DWITH_PYTHON=ON -DWITH_PYTORCH=OFF -DWITH_TENSORFLOW=ON -DCUDA_SDK_ROOT_DIR=/usr/local/cuda/samples -DCUDA_ROOT=/usr/local/cuda/\n",
    "  make install\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sz5c6ly_ixdd"
   },
   "source": [
    "### MERLIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "a99l-8hxmj51",
    "outputId": "c6f60d75-e22b-41d5-a871-c35a095aa0f7"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip install merlinpy-mri merlinth-mri merlintf-mri"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c9QAsFoSSgG6"
   },
   "source": [
    "### BART\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M286PAF-lc3X"
   },
   "source": [
    "We install dependencies for BART. Make sure that you have installed the requirements if you run locally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fQEVVn59O2eP"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if [ \"$MIN_INSTALL\" = false ] ; then\n",
    "  [ $COLAB -ne 1 ] && echo \"Skipp cell (not on Colab)\" && exit 0\n",
    "\n",
    "  # Install BARTs dependencies\n",
    "  apt-get install -y make gcc libfftw3-dev liblapacke-dev libpng-dev libopenblas-dev &> /dev/null\n",
    "\n",
    "  # Install additional dependencies for converting ISMRMRD files\n",
    "  apt-get install -y libismrmrd-dev libboost-all-dev libhdf5-serial-dev &> /dev/null\n",
    "fi\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fFy511-hSgG8"
   },
   "source": [
    "#### Clone and Compile BART\n",
    "\n",
    "We clone BART into the current working directory of this notebook and delete any previous installation in this directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UG4RDIhQSgG8"
   },
   "outputs": [],
   "source": [
    "# Clone Bart\n",
    "! [ -d bart ] && rm -r bart\n",
    "! git clone https://github.com/mrirecon/bart/ bart &> /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YiN-TWaf39OB"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if [ \"$MIN_INSTALL\" = false ] ; then\n",
    "  cd bart\n",
    "\n",
    "  # Define compile options\n",
    "  COMPILE_SPECS=\" PARALLEL=1\n",
    "                  CUDA=1\n",
    "                  NON_DETERMINISTIC=1\n",
    "                  ISMRMRD=1\n",
    "                  \"\n",
    "\n",
    "  printf \"%s\\n\" $COMPILE_SPECS > Makefiles/Makefile.local\n",
    "\n",
    "\n",
    "  #if [ $COLAB -eq \"1\" ];\n",
    "  #then\n",
    "  # set path to cuda for Colab\n",
    "  echo \"CUDA_BASE=/usr/local/cuda\" >> Makefiles/Makefile.local\n",
    "  echo \"CUDA_LIB=lib64\" >> Makefiles/Makefile.local\n",
    "  #fi\n",
    "\n",
    "\n",
    "  make &> /dev/null\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iCwkW7baO2eQ"
   },
   "source": [
    "#### Add BART to PATH variable\n",
    "\n",
    "We add the BART directory to the PATH variable and include the python wrapper for reading *.cfl files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9WCi0TsVO2eQ"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "os.environ['TOOLBOX_PATH']=os.getcwd()+\"/bart/\"\n",
    "os.environ['PATH'] = os.environ['TOOLBOX_PATH'] + \":\" + os.environ['PATH']\n",
    "sys.path.append(os.environ['TOOLBOX_PATH'] + \"/python/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z6UbLX2Ki7Ey"
   },
   "source": [
    "Check BART setup:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3WwBQWEUi7Ey"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "  echo \"# The BART used in this notebook:\"\n",
    "  which bart\n",
    "  echo \"# BART version: \"\n",
    "  bart version\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TyNHPaMCn22A"
   },
   "source": [
    "Some helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5-UDh5HahDGH"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "def plotcfl(files, title=None):\n",
    "  import numpy as np\n",
    "  from matplotlib import pyplot as plt\n",
    "  import cfl\n",
    "  import os\n",
    "\n",
    "  from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "  import ipywidgets as widgets\n",
    "\n",
    "  def update(Range=(0.,1.),Coil=0, Map=0, Slice=0, Batch=0):\n",
    "\n",
    "    if(title==None):\n",
    "      wtitle=[]\n",
    "      for file in files:\n",
    "        head, tail = os.path.split(file)\n",
    "        wtitle.append(tail)\n",
    "    else:\n",
    "      wtitle=title\n",
    "\n",
    "    data=cfl.readcfl(files[0])\n",
    "    rat=data.shape[0]/data.shape[1]\n",
    "    width=16\n",
    "\n",
    "    ncols=len(files)\n",
    "    nrows=1\n",
    "\n",
    "    rat=rat\n",
    "    fig, axs = plt.subplots(nrows=nrows, ncols=ncols, constrained_layout=True, squeeze=False, figsize=(width, width*rat))\n",
    "\n",
    "    for i in range(len(files)):\n",
    "\n",
    "        data=cfl.readcfl(files[i])\n",
    "        nshp = [1] * 16\n",
    "\n",
    "        for k in range(len(data.shape)):\n",
    "          nshp[k] = data.shape[k]\n",
    "\n",
    "        idx=[0]*16\n",
    "        idx[3]=min(Coil,nshp[3]-1)\n",
    "        idx[4]=min(Map,nshp[4]-1)\n",
    "        idx[13]=min(Slice,nshp[13]-1)\n",
    "        idx[15]=min(Batch,nshp[15]-1)\n",
    "\n",
    "        idx[0]=slice(None, None, -1)\n",
    "        idx[1]=slice(None, None, None)\n",
    "\n",
    "        dat=np.abs(data.reshape(nshp)[tuple(idx)])\n",
    "        if 0 < np.max(dat):\n",
    "          dat=dat/np.max(dat)\n",
    "\n",
    "        axs.flatten()[i].imshow(dat,cmap=\"gray\",vmin=Range[0], vmax=Range[1])\n",
    "        axs.flatten()[i].set_title(wtitle[i])\n",
    "    for ax in axs.flatten():\n",
    "      ax.axis(\"off\")\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "  nshp = [1] * 16\n",
    "  for file in files:\n",
    "    data=cfl.readcfl(file)\n",
    "    for i in range(len(data.shape)):\n",
    "      nshp[i] = max(data.shape[i], nshp[i])\n",
    "\n",
    "  interact(update,\n",
    "           Range=widgets.FloatRangeSlider(value=[0.,1.], min=0, max=1.),\n",
    "           Coil=widgets.IntSlider(min=0, max=nshp[3]-1, step=1, value=0),\n",
    "           Map=widgets.IntSlider(min=0, max=nshp[4]-1, step=1, value=0),\n",
    "           Slice=widgets.IntSlider(min=0, max=nshp[13]-1, step=1, value=0),\n",
    "           Batch=widgets.IntSlider(min=0, max=nshp[15]-1, step=1, value=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yEyVnQ9GK4eu"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "function imshow () {\n",
    "    bart toimg -W $1 $1.png > /dev/null\n",
    "    cat $1.png | display\n",
    "    rm $1.png\n",
    "}\n",
    "export -f imshow\n",
    "\n",
    "# version without automatic windowing\n",
    "function imshow_nowin () {\n",
    "    bart toimg $1 $1.png > /dev/null\n",
    "    cat $1.png | display\n",
    "    rm $1.png\n",
    "}\n",
    "export -f imshow_nowin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IHCV5_OYqKuP"
   },
   "source": [
    "### Matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UxHdAwUyqTKU",
    "outputId": "ae2063a4-67d7-482a-e993-7760ac79af73"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip install matplotlib==3.2.2  # correct plotting bug"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K3YztvzE6zVD"
   },
   "source": [
    "### Final checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-NCG4SMDGmhG",
    "outputId": "9ef4c9d0-9249-4617-f583-4d550f649c83"
   },
   "outputs": [],
   "source": [
    "print(\"Restarting of kernel...\")  # needed to load installed libraries\n",
    "#get_ipython().kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2CaqkBhFhj_T"
   },
   "outputs": [],
   "source": [
    "# inspect the available GPU hardware\n",
    "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dYV08Ia5BP-k"
   },
   "outputs": [],
   "source": [
    "pip list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9rrR_V_2bOvX"
   },
   "source": [
    "## Imports\n",
    "<a name=\"Imports\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ThDfjT4Hqxms"
   },
   "source": [
    "**ATTENTION:** Please make sure that you have restarted the Jupyter kernel after installation of the above libraries (error message of unexpected crash and restart). Otherwise, the installed libraries will not be found. You can restart the kernel with:\n",
    "\n",
    "```\n",
    "get_ipython().kernel.do_shutdown(True)\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jRs3msSma5G-"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "os.environ['TOOLBOX_PATH']=os.getcwd()+\"/bart/\"\n",
    "os.environ['PATH'] = os.environ['TOOLBOX_PATH'] + \":\" + os.environ['PATH']\n",
    "sys.path.append(os.environ['TOOLBOX_PATH'] + \"/python/\")\n",
    "\n",
    "import numpy as np\n",
    "import medutils\n",
    "import SimpleITK as sitk\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "try:\n",
    "  import optotf\n",
    "except:\n",
    "  print('Optox not installed')\n",
    "import tensorflow as tf\n",
    "import espirit\n",
    "\n",
    "import voxelmorph as vxm\n",
    "#import neurite as ne\n",
    "\n",
    "# own utils\n",
    "from utils.flowplotting import *\n",
    "from utils.imageplotting import *\n",
    "from utils.mri import *\n",
    "from utils.padding import *\n",
    "from utils.radialsampling import *\n",
    "from utils.cartesiansampling import *\n",
    "from utils.warping import *\n",
    "from utils.motionsim import *\n",
    "from utils.motioncomp import *\n",
    "\n",
    "from mri.operators import NonCartesianFFT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LcNRw4_ebMOw"
   },
   "source": [
    "## Data loading\n",
    "<a name=\"DataLoading\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rKW697I8Z_Rs"
   },
   "outputs": [],
   "source": [
    "# Load the numpy\n",
    "import numpy as np\n",
    "from utils.padding import zpad\n",
    "\n",
    "datapath = '/content/ismrm-moco-workshop/data/brain_slice.npz'  # brain_slice\n",
    "data = np.load(datapath)\n",
    "img = data['arr_0']\n",
    "nRead = np.amax(np.shape(img)[0:2])\n",
    "img = zpad(img, (nRead, nRead, np.shape(img)[2])).astype(np.complex64)  # make image quadratic for\n",
    "\n",
    "nX, nY, nCha = np.shape(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nWy_UX3UobqE"
   },
   "outputs": [],
   "source": [
    "plot(img, title='Coil images')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "36tAS_UAfpry"
   },
   "source": [
    "# 1. Motion artifact appearance\n",
    "<a name=\"MotionArtifactAppearance\"></a>\n",
    "In this part, we will examine the impact of the motion artifacts on the image. Motion during an MR scan introduces blurring and aliasing (along the phase-encoding directions). Artifact appereance depends on the imaging trajectory: Cartesian and non-Cartesian. In the following we will first investigate the motion artifact appeareance on fully-sampled cases and then move towards undersampled/accelerated imaging trajectories.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sLi8D0z9dksl"
   },
   "source": [
    "## **Spatial description of motion**\n",
    "For translational, rigid and affine motion exists a direct linear relationship between the motion-corrected/free k-space $\\nu(k)$ and the motion-affected k-space $\\nu'(k')$.\n",
    "\n",
    "$$\n",
    "\\nu(k) = \\nu'(k') \\frac{\\operatorname{exp}(2\\pi i k'\\cdot t)}{|\\det (A)|}\n",
    "$$\n",
    "where $k'$ is the motion-affected k-space trajectory, $k$ is the motion-corrected/free k-space trajectory, $t=\\lbrack t_x, t_y, t_z \\rbrack^T$ is the translational component of the motion and $A$ is the affine motion matrix in 3D\n",
    "\n",
    "$\n",
    "\\begin{align}\n",
    "A = RGS = R_x R_y R_z G S = \\begin{bmatrix} 1 & 0 & 0 \\\\ 0 & \\cos \\phi & -\\sin \\phi \\\\ 0 & \\sin \\phi & \\cos \\phi \\end{bmatrix} \\begin{bmatrix} \\cos \\theta & 0 & \\sin \\theta \\\\ 0 & 1 & 0 \\\\ -\\sin \\theta & 0 & \\cos \\theta \\end{bmatrix} \\begin{bmatrix} \\cos \\psi & -\\sin \\psi & 0 \\\\ \\sin \\psi & \\cos \\psi & 0 \\\\ 0 & 0 & 1 \\end{bmatrix} \\begin{bmatrix} 1 & G_{xy} & G_{xz} \\\\ 0 & 1 & G_{yz} \\\\ 0 & 0 & 1 \\end{bmatrix} \\begin{bmatrix} S_x & 0 & 0 \\\\ 0 & S_y & 0 \\\\ 0 & 0 & S_z \\end{bmatrix}\n",
    "\\end{align}\n",
    "$\n",
    "\n",
    "consisting of the rotation matrices $R_x, R_y$ and $R_z$ (with angles $\\lbrace \\phi, \\theta, \\psi \\rbrace$) along the respective axis $x, y$ and $z$, shearing matrix $G$ (with 3D shear parameters $\\lbrace G_{xy}, G_{xz}, G_{yz} \\rbrace$) and scaling matrix $S$ (with 3D scale factors $\\lbrace S_x, S_y, S_z \\rbrace$).\n",
    "\n",
    "In 2D, the affine motion matrix simplifies to\n",
    "\n",
    "$\n",
    "A=RGS = \\begin{bmatrix} \\cos \\phi & \\sin \\phi \\\\-\\sin \\phi & \\cos \\phi \\end{bmatrix} \\begin{bmatrix} 1 & G_{xy} \\\\ 0 & 1 \\end{bmatrix} \\begin{bmatrix} S_x & 0 \\\\ 0 & S_y \\end{bmatrix}\n",
    "$\n",
    "\n",
    "### Affine motion\n",
    "The 3D affine transformation has 12 degrees of freedom consisting of 3D translations, rotations, shears and scaling. It maps the k-space point $k=\\lbrack k_x, k_y, k_z \\rbrack^T$ to the corrupted k-space point $k'=\\lbrack k'_x, k'_y, k'_z \\rbrack^T$ following $k'=A^{-T} k$\n",
    "\n",
    "in k-space: $\n",
    "\\begin{bmatrix} k'_x \\\\ k'_y \\\\ k'_z \\end{bmatrix} = R \\begin{bmatrix} 1 & 0 & 0 \\\\ -G_{xy} & 1 & 0 \\\\ (G_{xy}G_{yz}-G_{xz}) & -G_{yz} & 1\\end{bmatrix} \\begin{bmatrix} S_x^{-1} & 0 & 0 \\\\ 0 & S_y^{-1} & 0 \\\\ 0 & 0 & S_z^{-1} \\end{bmatrix} \\begin{bmatrix} k_x \\\\ k_y \\\\ k_z \\end{bmatrix}\n",
    "$\n",
    "\n",
    "or equivalently\n",
    "\n",
    "in image domain: $\n",
    "\\begin{bmatrix} x' \\\\ y' \\\\ z' \\\\ 1 \\end{bmatrix} = \\begin{bmatrix} & & & t_x \\\\ & R & & t_y \\\\ & & & t_z \\\\ 0 & 0 & 0 & 1 \\end{bmatrix} GS \\begin{bmatrix} x \\\\ y \\\\ z \\\\ 1 \\end{bmatrix}\n",
    "$\n",
    "\n",
    "In this case, we are considering the augmented matrix form (i.e. 4D for 3D spatial), which conveniently allows us to write translations and the other affine components as a series of matrix multiplications.\n",
    "\n",
    "### Rigid motion\n",
    "In case of rigid motion, we only have translation $t$ and rotation $R$, i.e. $G=S=I$, where $I$ is the identity matrix.\n",
    "\n",
    "### Translational motion\n",
    "For translational motion $t$, all other components are neglected, i.e. $R=G=S=I$.\n",
    "\n",
    "### Elastic (non-rigid) motion\n",
    "For elastic motion, the image $\\rho \\in \\mathbb{R}^N$ is deformed with a deformation field/motion model $u$ which is a vector field containing the motion vector at every voxel $u \\in \\mathbb{R}^{N \\times D}$ for a $D$-dimensional registration. The deformed image is given as\n",
    "$\n",
    "\\rho_d(x') = \\rho(u(x))\n",
    "$.\n",
    "\n",
    "## **Temporal description of motion**\n",
    "Motion represents a displacement over time and this can be addressed most directly by stating the spatial variation (i.e. the affine motion parameters) as a function of time. In practice, however, this belies the complexity of real physiologic motion, where limits on real-world acceleration, shear, etc. limit how displacement can practically vary over time. Moreover, some choice of discretization for the time domain is necessary.\n",
    "\n",
    "### Regular motion\n",
    "Many types of physiologic motion exhibit temporal regularity. Respiration has an inspiration and expiration phase; the heartbeat can be characterized by an approximately repeating waveform, identifiable by the QRS complex in the ECG. A key aspect of these types of motion is that while they are generally periodic, they do not necessarily have a fixed frequency through the duration of a scan — a subject may hold their breath; a heartbeat may skip. Therefore, these types of motion are often described as cyclic rather than periodic.\n",
    "\n",
    "### Irregular motion\n",
    "Irregular motions (e.g., head motion) are generally represented using a discretized time series of serial displacements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TqAj0XaoqWhs"
   },
   "source": [
    "## Motion in Cartesian imaging\n",
    "We will first examine the impact of motion on Cartesian imaging. We can apply the motion in the image domain by the function:\n",
    "\n",
    "```\n",
    "def transform_img(img, p):\n",
    "# img      input image to be transformed\n",
    "# p        affine transformation parameters\n",
    "#          3D (rank(img) == 3): t_x, t_y, t_z, \\phi [°], \\theta [°], \\psi [°], G_{xy}, G_{xz}, G_{yz}, S_x, S_y, S_z\n",
    "#          2D (rank(img) == 2): t_x, t_y, \\phi [°], G_{xy}, S_x, S_y\n",
    "# return   transformed image\n",
    "```\n",
    "takes care of creating and applying the transformation to the input `img`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lwtNNF9VNuuS"
   },
   "source": [
    "### Motion in fully-sampled imaging\n",
    "**Task 1:** Apply a 2D translational motion with $t_x=10, t_y=5$ to the RSS coil-combined image (`img_rss = rss(img)`) and plot the original and transformed image side-by-side using the function `plot([img_a, img_b])`.\n",
    "\n",
    "\n",
    "```\n",
    "def plot(img, flow=None, ...):\n",
    "# img        to be plotted image or list of 2D images\n",
    "#            2D: single 2D image\n",
    "#            3D: third dimension is plotted side-by-side\n",
    "# flow       2D: [x, y, flowDir]\n",
    "#            3D: [x, y, slices, flowDir]\n",
    "#            flowDir being the x and y components of the flow\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KcVe2BCgnreI"
   },
   "outputs": [],
   "source": [
    "img_rss = rss(img)\n",
    "img_trans = transform_img(img_rss, [30, 5, 0, 0, 1, 1])\n",
    "plot([img_rss, img_trans])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c9Ohkz8i-ZM3"
   },
   "source": [
    "**Task 2:** Apply an in-plane rotation of $\\phi=30^\\circ$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i8dKYVO6-tNR"
   },
   "outputs": [],
   "source": [
    "img_rot = transform_img(img_rss, [0, 0, 30, 0, 1, 1])\n",
    "plot([img_rss, img_rot])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eKoRMyI9Alet"
   },
   "source": [
    "**Task 3:** Try different combinations of translation, rotation, scalings and shearings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y7vM2jkhW69m"
   },
   "outputs": [],
   "source": [
    "img_affine = transform_img(img_rss, [4, 0, 10, 0.1, 1, 1])\n",
    "plot([img_rss, img_affine])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "94LMToIK37QH"
   },
   "source": [
    "**Task 4:** Write the equivalent translational operation of the image domain transformation in k-space. Plot the translational offsetted image performed in image domain, k-space domain and the absolute error between image domain and k-space domain transformation.<br/>\n",
    "*Optional:* Write the equivalent affine operation in k-space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BhGMtATzi6CP"
   },
   "outputs": [],
   "source": [
    "kspace = fft2c(img)\n",
    "p = [30, 5, 0, 0, 1, 1]\n",
    "\n",
    "mask_motion = np.ones(np.shape(kspace)[1])\n",
    "kspace_motion = np.zeros_like(kspace)\n",
    "y,x=np.meshgrid(np.linspace(-0.5,0.5-1/np.shape(kspace)[0], num=np.shape(kspace)[0]), np.linspace(-0.5,0.5-1/np.shape(kspace)[1], num=np.shape(kspace)[1]), indexing='ij')\n",
    "dy = 0\n",
    "for iL in np.arange(np.shape(kspace)[1]):\n",
    "  if mask_motion[iL]==1:\n",
    "    dPhase = np.exp(2*np.pi*1j*(x[:, dy+iL]*p[0]\\\n",
    "                    + y[:, dy+iL]*p[1]))\n",
    "\n",
    "    kspace_motion[:,iL,:] = kspace[:,iL,:]*dPhase[:,np.newaxis]\n",
    "  else:\n",
    "    kspace_motion[:,iL,:] = kspace[:,iL,:]\n",
    "\n",
    "img_motion_k = ifft2c(kspace_motion)\n",
    "\n",
    "plot([img_rss, img_trans, rss(img_motion_k), np.abs(img_trans - rss(img_motion_k))], title='motion-free | translational motion (image domain) | translational motion (k-space domain) | error (image <> k-space)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zeCH3nFNXbAw"
   },
   "source": [
    "Let us now simulate a simple case of a motion-affected MR scan. Therefore, we will impact every $n$-th phase-encoding line with motion using the function `simulate_motion()`. In order to simulate it, we first need to set up the functions for a zero-filled reconstruction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RdYyciw1Pd-0"
   },
   "source": [
    "### Zero-filled reconstruction\n",
    "In order to simulate motion and ultimately reconstruct an image, we need the multi-coil forward operator $A$ in `mriForwardOp(image, mask, smaps)` and adjoint operator $A^∗$ in `mriAdjointOp(kspace, mask, smaps)`, as well as an estimation of the coil sensitivities.\n",
    "\n",
    "*Hint*: We perform a centered FFT with `ortho` normalization in order to provide adjoint operators.\n",
    "Suggested Readings:<br/>\n",
    "Pruessmann et al. [SENSE: Sensitivity encoding for fast MRI](https://onlinelibrary.wiley.com/doi/abs/10.1002/%28SICI%291522-2594%28199911%2942%3A5%3C952%3A%3AAID-MRM16%3E3.0.CO%3B2-S) Magnetic Resonance in Medicine, 43(5):952-962, 1999."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WYuiGIQuVdHT"
   },
   "source": [
    "#### Forward and adjoint operator\n",
    "The operators are already defined (and imported) in `utils/mri.py`:\n",
    "\n",
    "\n",
    "```\n",
    "def mriAdjointOp(kspace, mask, smaps):\n",
    "  return np.sum(ifft2c(kspace * mask)*np.conj(smaps), axis=-1)\n",
    "\n",
    "def mriForwardOp(image, mask, smaps):\n",
    "  return fft2c(smaps * image[:,:,np.newaxis]) * mask\n",
    "\n",
    "def fft2c(image, axes=(0,1)):\n",
    "  return np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(image, axes=axes), norm='ortho', axes=axes), axes=axes)\n",
    "\n",
    "def ifft2c(kspace, axes=(0,1)):\n",
    "  return np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(kspace, axes=axes), norm='ortho', axes=axes), axes=axes)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZJEAml2QQ1NC"
   },
   "source": [
    "#### Coil sensitivity estimation\n",
    "The coil sensitivity maps (`smaps`) are smooth maps that show us in which parts the individual coil elements are sensitive. We will need these information for our multi-coil MRI forward and adjoint operators. We use the [python implementation](https://github.com/mikgroup/espirit-python) for ESPIRiT [1,2] to estimate these coil sensitivity maps.\n",
    "\n",
    "[1] Uecker et al. [ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: Where SENSE meets GRAPPA](https://onlinelibrary.wiley.com/doi/10.1002/mrm.24751). Magn Reson Med 71(3):990-1001, 2014.\n",
    "\n",
    "[2] https://github.com/mikgroup/espirit-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yzEGLw2bRBYf"
   },
   "outputs": [],
   "source": [
    "import espirit\n",
    "kspace = fft2c(img)\n",
    "kspace_espirit = kspace[:,:,np.newaxis,:]\n",
    "smaps_espirit = espirit.espirit(kspace_espirit, 8, 20, 0.05, 0)\n",
    "\n",
    "smaps = smaps_espirit[:,:,0,:,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9VYTTVKmSPxr"
   },
   "source": [
    "Let us visualize the sensitivity maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D-eqq4GCT6qh"
   },
   "outputs": [],
   "source": [
    "plot(smaps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vqDaBvGSVmoU"
   },
   "source": [
    "Now, you should check if the adjoint operator is working as expected. The result should be a coil-combined image. Right now, there is no undersampling mask involved, i.e., it is set to all ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8urAkcCtVugH"
   },
   "outputs": [],
   "source": [
    "sampling_mask_fs = np.ones_like(kspace)  # fully-sampled sampling mask\n",
    "img_cc = mriAdjointOp(kspace, sampling_mask_fs, smaps)\n",
    "plot(img_cc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c7HAJM6CaM9I"
   },
   "source": [
    "### Motion simulation\n",
    "Now we can simulate motion in the acquired k-space. The temporal behaviour, i.e. when the motion happens can be described by motion signals, which we will look into more detail in the next section. First, we will simulate an abrupt and periodic motion. We will make use of the function `simulate_motion()` in `utils/motionsim.py`.\n",
    "\n",
    "\n",
    "```\n",
    "def simulate_motion(img_cc, smaps, mask, p):\n",
    "  # img_cc      motion-free coil-combined image\n",
    "  # smaps       coil sensitivity maps\n",
    "  # mask        k-space sampling mask\n",
    "  # p           a) affine motion parameters (constant over time), 1x6\n",
    "  #             b) affine motion course (time-dependent or time-constant), nPE x 6\n",
    "  #             np.abs(p[:, :5]) > 0 = mask_motion; time points of phase-encoding steps\n",
    "  #             at which motion parameters > 0 are defined, i.e. motion is happening\n",
    "  # return:     motion-affected k-space, motion mask\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sizFqyYc2xFb"
   },
   "source": [
    "**Task 5:** Simulate a motion-affected image, such that every 4-th phase-encoding line (i.e. abrupt motion) is affected by a translational motion with $t_x=10, t_y=5$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hKXNgVVTtZuO"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(sampling_mask_fs)[1])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, sampling_mask_fs, p)\n",
    "img_fs_const_motion = mriAdjointOp(kspace_motion, sampling_mask_fs, smaps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rDHSVaYH0aey"
   },
   "source": [
    "**Task 6:** Let us now visualize the motion-free and the motion-affected image side-by-side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8AWlOxtt0fl0"
   },
   "outputs": [],
   "source": [
    "plot([img_cc, img_fs_const_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hvNKz4PI3UYk"
   },
   "source": [
    "### Temporal variation of motion\n",
    "Motion represents a displacement over time and this can be addressed most directly by stating the spatial variation (i.e. the affine motion parameters) as a function of time.\n",
    "\n",
    "The simulated motion from the previous example affected every $n$-th phase-encoding line, i.e. every $n$ TRs (assuming one echo line per TR) an abrupt motion was happening. We can depict the temporal behaviour of the motion as a function of affine motion parameters over time with the help of `plot_motion_course()` in `utils/motionsim.py`.\n",
    "\n",
    "\n",
    "```\n",
    "def plot_motion_course(motion_course, TR=1):\n",
    "  # motion_course     temporal variation of motion, i.e. temporal course of motion parameters, shape: [motion_parameters, total_time]\n",
    "  # TR                repetition time [ms]\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ylnlguz_78VM"
   },
   "source": [
    "**Task 7:** Plot the affine motion parameters for an abrupt translational motion ($t_x=10, t_y=5$) as a function over time for a TR=5 ms and one echo (phase-encoding line) per TR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oFBYYTqm8bJp"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "TR = 5E-3\n",
    "\n",
    "# abrupt motion\n",
    "motion_course = np.zeros((np.shape(mask_motion)[1], len(p)))\n",
    "for idx, cp in enumerate(p):\n",
    "  motion_course[:,idx] = np.squeeze(np.abs(mask_motion[0,:,0] * cp))\n",
    "\n",
    "plot_motion_course(motion_course, TR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1aRKRK09G0tM"
   },
   "source": [
    "**Task 8:** Create a periodic motion whose translational motion $t_x$ follows a sinus rhythm with frequency $\\omega_x=5 \\text{s}^{-1}$ , maximal amplitude 15, and $t_y$ follows a cosine rhythm with a frequency $\\omega_y=6 \\text{s}^{-1}$, maximal amplitude 8. Both motions happen at every 4-th phase-encoding line. Plot the affine motion parameters over time (TR=5 ms)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vpex4CH1HqLY"
   },
   "outputs": [],
   "source": [
    "p = [15, 8, 0, 0, 1, 1]\n",
    "TR = 5E-3\n",
    "\n",
    "# periodic motion\n",
    "time = np.arange(np.shape(mask_motion)[1]) * TR\n",
    "motion_signal = np.concatenate((np.stack([p[0] * np.sin(5 * time), p[1] * np.cos(6 * time)], axis=1), np.tile(p[2:], (len(time), 1))), axis=1)\n",
    "motion_course = np.zeros((np.shape(mask_motion)[1], len(p)))\n",
    "mask_motion_curr = np.tile(np.squeeze(np.abs(mask_motion[0,:,0]))[:, np.newaxis], (1, len(p)))  # when is motion happening over time for each motion parameter\n",
    "for idx, cp in enumerate(p):\n",
    "  motion_course[:,idx] = mask_motion_curr[:, idx] * motion_signal[:,idx]\n",
    "\n",
    "plot_motion_course(motion_course, TR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lexAHV_INgvW"
   },
   "source": [
    "**Task 9:** Corrupt the k-space now with the time-dependent motion parameters from **Task 8** and visualize the motion-affected image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u-9W8jXeNrXQ"
   },
   "outputs": [],
   "source": [
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, sampling_mask_fs, motion_course)\n",
    "img_motion = mriAdjointOp(kspace_motion, sampling_mask_fs, smaps)\n",
    "\n",
    "plot([img_cc, img_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NP02rnEYNe7h"
   },
   "source": [
    "It is also conceivable that different motion parameters are happening at different time points, i.e. we will have a different `mask_motion` for each motion parameter."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ntIpvdd-NyC6"
   },
   "source": [
    "**Task 10:** Create a periodic motion whose translational motion $t_x$ follows a sinus rhythm with frequency $\\omega_x=5 \\text{s}^{-1}$ , maximal amplitude 15, and $t_y$ follows a cosine rhythm with a frequency $\\omega_y=6 \\text{s}^{-1}$, maximal amplitude 8. Translation $t_x$ happens every 4-th phase-encoding line, whereas translation $t_y$ happens at every 8-th phase-encoding line. Plot the affine motion parameters over time (TR=5 ms)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ftS6FMWTOKrf"
   },
   "outputs": [],
   "source": [
    "p = [15, 8, 0, 0, 1, 1]\n",
    "TR = 5E-3\n",
    "\n",
    "# periodic motion\n",
    "time = np.arange(np.shape(mask_motion)[1]) * TR\n",
    "motion_signal = np.concatenate((np.stack([p[0] * np.sin(5 * time), p[1] * np.cos(6 * time)], axis=1), np.tile(p[2:], (len(time), 1))), axis=1)\n",
    "motion_course = np.zeros((np.shape(mask_motion)[1], len(p)))\n",
    "mask_motion_curr = np.tile(np.squeeze(np.abs(mask_motion[0,:,0]))[:, np.newaxis], (1, len(p)))  # when is motion happening over time for each motion parameter\n",
    "for idx, cp in enumerate(p):\n",
    "  motion_course[:,idx] = mask_motion_curr[:, idx] * motion_signal[:,idx]\n",
    "\n",
    "plot_motion_course(motion_course, TR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sATD_04LNKWf"
   },
   "source": [
    "### Motion in undersampled/accelerated imaging\n",
    "**Task 11:** Perform a Cartesian undersampling (regular and random) using the provided function `generate_mask()` from `utils/cartesiansampling.py`:\n",
    "\n",
    "\n",
    "```\n",
    "def generate_mask(R, nPE, nFE, nRef=20, mode='regular'):\n",
    "  # R     desired acceleration factor\n",
    "  # nPE   amount of phase-encoding lines\n",
    "  # nFE   amount of frequency-encoding lines\n",
    "  # nRef  amount of fully-sampled center lines\n",
    "  # mode  'regular': Parallel-Imaging-like/regular undersampling, 'random': Compressed-Sensing-like/random undersampling\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UVh940pQOrte"
   },
   "source": [
    "Now, generate a regular undersampling mask and visualize it (we visualize only a fraction in frequency encoding direction).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HLJ93gZqOmxP"
   },
   "outputs": [],
   "source": [
    "np.random.seed(1001)\n",
    "sampling_mask = generate_mask(R=3, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "\n",
    "medutils.visualization.imshow(sampling_mask[:40,:,0], 'Undersampling mask', figsize=(20,20))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "73_6M_N7V5XR"
   },
   "source": [
    "Next you can examine the impact of the undersampling masks on the reconstructed image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GhvK7vYGQhcE"
   },
   "source": [
    "**Task 12:** Examine the impact of a `regular` undersampling on the image for changing acceleration factors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rS_ojqqvWH28"
   },
   "outputs": [],
   "source": [
    "sampling_mask = generate_mask(R=3, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "img_cc_us = mriAdjointOp(kspace, sampling_mask, smaps)\n",
    "plot(img_cc_us)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z4GTsVBCrtDE"
   },
   "source": [
    "**Task 13:** Examine the impact of a `random` undersampling on the image for changing acceleration factors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lsM-SkFVrydu"
   },
   "outputs": [],
   "source": [
    "sampling_mask_r = generate_mask(R=3, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='random')\n",
    "img_cc_us_r = mriAdjointOp(kspace, sampling_mask_r, smaps)\n",
    "plot(img_cc_us_r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7fqv8QTD4vpJ"
   },
   "source": [
    "Let us now combine accelerated imaging with the impact of motion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kO7Hl3Qg41PP"
   },
   "source": [
    "**Task 14:** Induce motion in the `regular` undersampled data, assuming every 4-th phase-encoding line is affected by translational motion $t_x=10, t_y=5$. Plot the motion-free, fully-sampled motion-affected, and undersampled motion-affected image side-by-side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tB2JD87trWiL"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(sampling_mask_fs)[1])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, sampling_mask, p)\n",
    "img_us_motion = mriAdjointOp(kspace_motion, sampling_mask, smaps)\n",
    "plot([img_cc, img_fs_const_motion, img_us_motion], title=['motion-free', 'fully-sampled motion-affected', 'regular undersampled motion-affected'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j_6cCmr04Bqa"
   },
   "source": [
    "**Task 15:** Induce motion in the `random` undersampled data, assuming every 4-th phase-encoding line is affected by translational motion $t_x=10, t_y=5$. Plot the motion-free, fully-sampled motion-affected, and undersampled motion-affected image side-by-side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DmVKNJrN4Dyl"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(sampling_mask_fs)[1])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, sampling_mask_r, p)\n",
    "img_us_r_motion = mriAdjointOp(kspace_motion, sampling_mask_r, smaps)\n",
    "plot([img_cc, img_fs_const_motion, img_us_r_motion], title=['motion-free', 'fully-sampled motion-affected', 'random undersampled motion-affected'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MUGEGu5N8orn",
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Motion in radial imaging\n",
    "Let's examine the impact of motion in radial imaging. We will therefore first need to radially sample the data. For reconstruction, we use the [gpuNUFFT](https://github.com/andyschwarzl/gpuNUFFT) package with the Python interface provided by [pySAP-mri](https://github.com/CEA-COSMIC/pysap-mri/). Alternatively, one could also use the [BART](https://mrirecon.github.io/bart/) toolbox."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DpITuL0R4hv9"
   },
   "source": [
    "For simplicity, we first zero-pad the image to a quadrativ field of view. Next, we recalculate the coil sensitivity maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cDgxBrxTZgY4"
   },
   "outputs": [],
   "source": [
    "import espirit\n",
    "\n",
    "# getting the image into quadrativ FOV\n",
    "#img_qfov = np.pad(img, ((0,), (int(np.abs(nX-nY)/2),), (0,)))\n",
    "nRead = np.amax(np.shape(img)[0:2])\n",
    "img_qfov = zpad(img, (nRead, nRead, np.shape(img)[2])).astype(np.complex64)\n",
    "img_qfov = img_qfov[:,:,np.newaxis,:]\n",
    "kspace_qfov = fft2c(img_qfov)\n",
    "\n",
    "smaps_rad_espirit = espirit.espirit(kspace_qfov, 8, 20, 0.05, 0)\n",
    "smaps_rad = smaps_rad_espirit[:,:,0,:,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PW6qgbbPrAai"
   },
   "source": [
    "### Option A) gpuNUFFT\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3OSjbBEmEEuU"
   },
   "source": [
    "We can use the [gpuNUFFT](https://github.com/andyschwarzl/gpuNUFFT) package to reconstruct an image from a retrospective radial sampling. Please see **Task 16** for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PbA2KPSxq8Y9"
   },
   "source": [
    "### Option B) BART"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BJqaAzHRDfuG"
   },
   "source": [
    "We can use the [BART](https://mrirecon.github.io/bart/) toolbox to perform a radial sampling and NUFFT reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZWfRWS9fBlIS"
   },
   "outputs": [],
   "source": [
    "import cfl\n",
    "cfl.writecfl('imgrad', img_qfov)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aZOMksD5yJmi"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "SAMPLES=256\n",
    "SPOKES=402  # pi/2 * SAMPLES\n",
    "bart traj -x$SAMPLES -y$SPOKES -r tmp_traj\n",
    "bart scale 0.5 tmp_traj traj\n",
    "bart nufft traj imgrad kspace_rad\n",
    "bart show -m kspace_rad\n",
    "bart nufft -a traj kspace_rad imgreco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zjmxu2deOXg4"
   },
   "outputs": [],
   "source": [
    "#plotcfl(['imgreco'])\n",
    "imgreco = cfl.readcfl('imgreco')\n",
    "plot(np.squeeze(imgreco))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OIPVPTiLDuKz"
   },
   "source": [
    "### Motion in fully-sampled and undersampled/accelerated imaging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aC9dEmbc5P7m"
   },
   "source": [
    "**Task 16:** Define the non-uniform FFT (NUFFT) operator `NonCartesianFFT(samples, shape, n_coils, density_comp, smaps, implementation='gpuNUFFT')` using the [gpuNUFFT](https://github.com/andyschwarzl/gpuNUFFT) package and which requires the following inputs:\n",
    "- `samples`: radial sampling trajectory (shape: `[nSpokes, 2]`)\n",
    "- `shape`: (list) of quadrative field of view size: `[nRead, nRead]`\n",
    "- `n_coils`: amount of MR coils\n",
    "- `density_comp`: density compensation function (shape: `[nSpokes, 1]`)\n",
    "- `smaps`: coil sensitivity maps (shape: `[coils, nRead, nRead]`)\n",
    "\n",
    "Plot the radially sampled image.\n",
    "\n",
    "*Hint:* You can use the function `prepare_radial()` in `utils/radialsampling.py` to obtain the radial sampling trajectory and the density compensation function:\n",
    "\n",
    "```\n",
    "def prepare_radial(acc, nRead, nSlices=1):\n",
    "\"\"\"\n",
    "acc:         acceleration factor\n",
    "nRead:       number of readout steps\n",
    "nSlices:     number of slices\n",
    ":return:     radial trajectory, density compensation function\n",
    "\"\"\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lh_4EWNhrCW8"
   },
   "outputs": [],
   "source": [
    "from mri.operators.utils import convert_locations_to_mask\n",
    "from mri.operators import NonCartesianFFT\n",
    "\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "csm = np.transpose(smaps_rad, (2, 0, 1))\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "# NUFFT operator\n",
    "nufft = NonCartesianFFT(samples=kpos, shape=[nRead, nRead], n_coils=np.shape(csm)[0], density_comp=dcf, smaps=csm, implementation='gpuNUFFT')\n",
    "kspace_radial = nufft.op(img_rad)\n",
    "img_radial = nufft.adj_op(kspace_radial)\n",
    "plot(img_radial)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PHrCIFpX72d3"
   },
   "source": [
    "**Task 17:** Perform a retrospectively accelerated radial imaging with various acceleration factors and plot the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_YJxl0iw8GZh"
   },
   "outputs": [],
   "source": [
    "kpos, dcf = prepare_radial(acc=8, nRead=nRead)\n",
    "\n",
    "# NUFFT operator\n",
    "nufft = NonCartesianFFT(samples=kpos, shape=[nRead, nRead], n_coils=np.shape(csm)[0], density_comp=dcf, smaps=csm, implementation='gpuNUFFT')\n",
    "kspace_radial = nufft.op(img_rad)\n",
    "img_radial_acc = nufft.adj_op(kspace_radial)\n",
    "plot(img_radial_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ei4i62knDBxu"
   },
   "source": [
    "**Task 18:** Similar to the Cartesian function `simulate_motion()` in `utils/motionsim.py` create a function that simulates the motion on a radially sampled k-space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rIJkmoTzFl63"
   },
   "outputs": [],
   "source": [
    "def simulate_motion_radial(img_rad, smaps, mask, dcf, p):\n",
    "    # img_cc      motion-free coil-combined image (quadrativ FOV)\n",
    "    # smaps       coil sensitivity maps: [nCoils, nRead, nRead]\n",
    "    # mask        radial sampling mask\n",
    "    # dcf         density compensation function\n",
    "    # p           a) affine motion parameters (constant over time), 1x6\n",
    "    #             b) affine motion course (time-dependent or time-constant), nPE x 6\n",
    "    #             np.abs(p[:, :5]) > 0 = mask_motion; time points of phase-encoding steps\n",
    "    #             at which motion parameters > 0 are defined, i.e. motion is happening\n",
    "\n",
    "    nRead = np.shape(img_rad)[0]\n",
    "    nufft = NonCartesianFFT(samples=mask, shape=[nRead, nRead], n_coils=np.shape(smaps)[0], density_comp=dcf, smaps=smaps, implementation='gpuNUFFT')\n",
    "    kspace = nufft.op(img_rad)\n",
    "    p = np.asarray(p)\n",
    "    tmp = np.unique(p, axis=0)\n",
    "    if len(np.shape(p)) == 1:\n",
    "        mask_motion = np.ones_like(kspace)\n",
    "    else:\n",
    "        mask_motion = np.abs(np.sum(p[:, :5], axis=1)) > 0\n",
    "        mask_motion = np.tile(mask_motion[np.newaxis, :], (np.shape(kspace)[0], 1))\n",
    "\n",
    "    if len(np.shape(p)) == 1 or np.shape(tmp[~np.all(tmp == 0, axis=1)])[0] == 1:  # constant motion over time\n",
    "        if len(np.shape(p)) != 1:\n",
    "            p = np.squeeze(tmp[~np.all(tmp == 0, axis=1)])\n",
    "        kspace_motion = nufft.op(transform_img(img_rad, p))\n",
    "        return kspace * ( 1 -mask_motion) + kspace_motion * mask_motion, mask_motion\n",
    "    else:  # time-dependent motion\n",
    "        kspace_aff = np.zeros_like(kspace)\n",
    "        for ky in np.arange(np.shape(img_rad)[1]):\n",
    "            kspace_aff[:, ky, :] = nufft.op(transform_img(img_rad, p[ky, :]))[:, ky, :]\n",
    "        return kspace_aff, mask_motion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KXelIoGwJlHa"
   },
   "source": [
    "**Task 19:** Simulate that every 4-th k-space radial readout is affected by a translational motion with $t_x=10, t_y=5$ for a fully-sampled acquisition (`acc=1`). Plot the motion-free and the motion-affected image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_BuUy98KG6le"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "mask_motion = np.zeros(np.shape(kpos)[0])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "nufft = NonCartesianFFT(samples=kpos, shape=[nRead, nRead], n_coils=np.shape(csm)[0], density_comp=dcf, smaps=csm, implementation='gpuNUFFT')\n",
    "kspace_rad_motion, mask_rad_motion = simulate_motion_radial(img_rad, csm, kpos, dcf, p)\n",
    "img_rad_motion = nufft.adj_op(kspace_rad_motion)\n",
    "plot([img_radial, img_rad_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tao-ZYKUKfDm"
   },
   "source": [
    "**Task 20:** Plot and compare the fully-sampled (Cartesian/non-Cartesian), fully-sampled motion-affected Cartesian and fully-sampled motion-affected non-Cartesian (radial) images. Which one is less sensitive to motion?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8fVqHlwdKxEw"
   },
   "outputs": [],
   "source": [
    "plot([img_radial, img_fs_const_motion, img_rad_motion], title='fully-sampled motion-free | fully-sampled Cartesian motion-affected | fully-sampled non-Cartesian motion-affected')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M89ih6tTNbWx"
   },
   "source": [
    "**Task 21:** Simulate that every 4-th k-space radial readout is affected by a translational motion with $t_x=10, t_y=5$ for an accelerated acquisition (`acc=8`). Plot the motion-free and the motion-affected image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1uYOWgj6Nl6I"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 0, 0, 1, 1]\n",
    "kpos, dcf = prepare_radial(acc=8, nRead=nRead)\n",
    "\n",
    "mask_motion = np.zeros(np.shape(kpos)[0])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "nufft = NonCartesianFFT(samples=kpos, shape=[nRead, nRead], n_coils=np.shape(csm)[0], density_comp=dcf, smaps=csm, implementation='gpuNUFFT')\n",
    "kspace_rad_acc_motion, mask_rad_acc_motion = simulate_motion_radial(img_rad, csm, kpos, dcf, p)\n",
    "img_rad_acc_motion = nufft.adj_op(kspace_rad_acc_motion)\n",
    "plot([img_radial, img_rad_acc_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nny3DdWKPxpL"
   },
   "source": [
    "**Task 22:** Repeat the above examples for varying affine motion parameters $p$, motion courses (constant or changing over time), motion time points ($n$-th readout) and acceleration factors of the radial undersampling. Plot and compare the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "21DHXCmyQQhB"
   },
   "outputs": [],
   "source": [
    "p = [10, 5, 30, 0, 1, 1]\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "mask_motion = np.zeros(np.shape(kpos)[0])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "nufft = NonCartesianFFT(samples=kpos, shape=[nRead, nRead], n_coils=np.shape(csm)[0], density_comp=dcf, smaps=csm, implementation='gpuNUFFT')\n",
    "kspace_rad_rigid_motion, _ = simulate_motion_radial(img_rad, csm, kpos, dcf, p)\n",
    "img_rad_rigid_motion = nufft.adj_op(kspace_rad_rigid_motion)\n",
    "\n",
    "p = [10, 5, 20, 0.1, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(kpos)[0])\n",
    "mask_motion[::4] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "kspace_rad_affine_motion, _ = simulate_motion_radial(img_rad, csm, kpos, dcf, p)\n",
    "img_rad_affine_motion = nufft.adj_op(kspace_rad_affine_motion)\n",
    "\n",
    "plot([img_radial, img_rad_motion, img_rad_rigid_motion, img_rad_affine_motion], title='fs, translational fs, rigid fs, affine fs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e-Cqiuy587am"
   },
   "source": [
    "# 2. Motion estimation / image registration\n",
    "<a name=\"MotionEstimation\"></a>\n",
    "In this part we will learn about algorithms to register and estimate the motion. We will illustrate conventional (optical flow, B-Spline) and deep-learning methods for this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TAJJ1fkdQ7bD"
   },
   "source": [
    "In this part of the tutorial we will have a closer look at image registration using conventional and deep learning-based solutions. We will only focus on pairwise registrations, i.e. registering a moving image $\\rho_m(r)$ to a fixed image $\\rho_f(r)$. The fixed and moving image are of dimension $N$ and are each defined on their own spatial domain: $\\Omega_f \\subset \\mathbb{R}^N$ and $\\Omega_m \\subset \\mathbb{R}^N$. The task of registration is to find a displacement $u(r)$ that makes $\\rho_m(r+u(r))$ spatially aligned to $\\rho_f(r)$. Equivalently we can say that registration is the process to find the transformation/displacement $T(r) = r+u(r)$ that makes $\\rho_m(T(r))$ spatially aligned to $\\rho_f(x)$. The transformation is defined as a mapping from the fixed image to the moving image, i.e. $T: \\Omega_f \\subset \\mathbb{R}^N \\to \\Omega_m \\subset \\mathbb{R}^N$. The flow/displacement/motion field $u(x)$ contains the flows in the respective directions $u(r) = [u_x(r), u_y(r), \\dots]$.\n",
    "\n",
    "This transformation $T$, respectively its flow/displacement/motion field $u$ is also known as the *forward* motion $u^{(F)}$ (for simplicity we will refer to the forward motion field as $u$). Consequently, we also have a *backward* motion $u^{(B)}$. Ideally, the forward and backward motion should be isomorph, i.e. $u(u^{(B)}(r)) = r$ which allows to obtain either of the flows from the other one. For some cases, it might also be beneficial if the flows are diffeomorph, i.e. forward and backward motion are continously differentiable. During optimization, we are actually solving for the backward motion and later invert this flow to obtain the forward motion. This is usally more numerically stable as we expect that there is a matching voxel in the fixed image to which we can morph the moving image's voxel.   \n",
    "\n",
    "The main properties of a registration algorithm are:\n",
    "- dimensionality: 2D to 2D, 2D to 3D, 3D to 3D, ...\n",
    "- spatial motion: rigid, affine, elastic (non-rigid)\n",
    "- input modality: mono-modal, multi-modal\n",
    "- similarity measure: landmark-based, image-intensity based\n",
    "- cost function: mean-squared-error, mutual information, cross-correlation, ...\n",
    "- regularization: pyramidal resolution, smoothness or energy constraints, Jacobian constraints, isomorph or diffeomorph constraints, ...\n",
    "- algorithm: optical flow, diffusion-based, splines (B-splines, thin-plate splines, ...), deep learning, ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2YFPr3Bw9FQ3"
   },
   "source": [
    "## Conventional image registration\n",
    "For the conventional image registration methods, we will have a look at the differences between optical flow and B-Spline based registrations, as these methods are the most commonly used ones. Moreover, deep learning-based registrations are often motivated or based on optical flow methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7DJEvMReVP-B"
   },
   "source": [
    "### Optical Flow\n",
    "We will first examine an optical flow-based image registration. Therefore we create a pair of motion-free and motion-affected image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iH3fhugwrMlF"
   },
   "source": [
    "**Task 23:** Create a pair of motion-free (fixed) and moving image using `transform_img()` and `plot()` them side-by-side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oa3Va4boLcy8"
   },
   "outputs": [],
   "source": [
    "img_fixed = rss(img)\n",
    "p = [10, 0, 0, 0, 1, 1]\n",
    "#p = [4, 8, 10, 0.05, 1, 1]\n",
    "img_moving = transform_img(img_fixed, p)\n",
    "img_fixed = minmaxscale(img_fixed, [0, 1])\n",
    "img_moving = minmaxscale(img_moving, [0, 1])\n",
    "plot([img_fixed, img_moving])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VPDQmMTFMWeD"
   },
   "source": [
    "**Task 24:** Let us define the optical flow registration using the scikit-image (or any other library of your liking) for the Lucas-Kanade algorithm.\n",
    "\n",
    "*Hint:* Check the orientation of the flow for the $x$ and $y$ component, i.e. verify the displacement for a translational motion (`[10, 0, 0, 0, 0, 0]`) in $x$ direction first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FismVEK8Mokc"
   },
   "outputs": [],
   "source": [
    "from skimage.registration import optical_flow_ilk\n",
    "uy, ux = optical_flow_ilk(img_fixed, img_moving)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QEH2-a4q0MfW"
   },
   "source": [
    "**Task 25**: First stack the $x$ and $y$ components of the flow $u_x$, $u_y$ along the last dimension to obtain a flow field $u$ of shape $X \\times Y \\times 2$. Plot the obtained deformation field as quiver and color-plots using `plot` (in `utils/imageplotting.py`) and `plot_flow` (in `utils/flowplotting.py`), respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hKmfJtUQN4pD"
   },
   "source": [
    "For displaying the flow fields, we can show them as a quiver plot (vector field overlaid to image) using `plot([image1, image2], [flow_of_image1, flow_of_image2])` or as a color-encoded image (color depicts flow angle $\\arg u$ and hue depicts magnitude $|u|$) using `plot_flow(flow)`.\n",
    "\n",
    "The color-coded image has the following intepretation: <br/>\n",
    "![](https://github.com/lab-midas/ismrm-moco-workshop/blob/master/doc/colorwheel.png?raw=true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XRlRmvdWVOU1"
   },
   "outputs": [],
   "source": [
    "flows = np.stack([ux, uy], -1)\n",
    "\n",
    "# quiver plots\n",
    "plot([img_fixed, img_moving], [np.zeros_like(flows), flows], scale=5, figsize=(20,20))\n",
    "\n",
    "# color-wheel plot\n",
    "plot_flow(flows)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_U_fOxRFT8Yy"
   },
   "source": [
    "**Task 26:** Inspect and display the deformed/warped image using `warp_2D()` in `utils/warping.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LlcDgJzmSM8n"
   },
   "outputs": [],
   "source": [
    "img_warped = warp_2D(img_moving, np.stack([ux, uy], -1))\n",
    "plot([img_fixed, img_moving, img_warped, np.abs(img_warped-img_fixed)], title='fixed | moving | warped | error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4iHXM6yQSL0i"
   },
   "source": [
    "**Task 27:** Perform the registration using the Total-variation L1-regularized optical flow algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CO5jdmb5T0lU"
   },
   "outputs": [],
   "source": [
    "from skimage.registration import optical_flow_tvl1\n",
    "uy, ux = optical_flow_tvl1(img_fixed, img_moving)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sd5QRforQ5kX"
   },
   "source": [
    "**Task 28:** Repeat the **Tasks 25** and **26** with the TV L1-regularized optical flow from **Task 27**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "whxVlmUm10jK"
   },
   "outputs": [],
   "source": [
    "flows = np.stack([ux, uy], -1)\n",
    "\n",
    "# quiver plots\n",
    "plot([img_fixed, img_moving], [np.zeros_like(flows), flows], scale=5, figsize=(20,20))\n",
    "\n",
    "# color-wheel plot\n",
    "plot_flow(flows)\n",
    "\n",
    "img_warped = warp_2D(img_moving, np.stack([ux, uy], -1))\n",
    "plot([img_fixed, img_moving, img_warped, np.abs(img_warped-img_fixed)], title='fixed | moving | warped | error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tUTEz1frUGS2"
   },
   "source": [
    "**Task 29:** Evaluate the obtained registration performance in the warped image (i.e. warped to fixed image) and compare against before registration (i.e. moving to fixed image). Compute the photometric loss = $\\sum (\\rho_\\text{warped} - \\rho_f)^2$, structural similarity (SSIM) and peak signal to noise ratio (PSNR)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vybVu45ZX7mG"
   },
   "outputs": [],
   "source": [
    "from skimage.metrics import structural_similarity, peak_signal_noise_ratio, mean_squared_error\n",
    "\n",
    "photo = mean_squared_error(img_fixed, img_warped)\n",
    "ssim = structural_similarity(img_fixed, img_warped, data_range=img_fixed.max() - img_fixed.min())\n",
    "psnr = peak_signal_noise_ratio(img_fixed, img_warped)\n",
    "\n",
    "print('Warped to fixed image')\n",
    "print(f'Photometric: {photo}')\n",
    "print(f'SSIM: {ssim}')\n",
    "print(f'PSNR: {psnr}')\n",
    "\n",
    "photo_m = mean_squared_error(img_moving, img_warped)\n",
    "ssim_m = structural_similarity(img_moving, img_warped, data_range=img_fixed.max() - img_fixed.min())\n",
    "psnr_m = peak_signal_noise_ratio(img_moving, img_warped)\n",
    "\n",
    "print('\\nMoving to fixed image')\n",
    "print(f'Photometric: {photo_m}')\n",
    "print(f'SSIM: {ssim_m}')\n",
    "print(f'PSNR: {psnr_m}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "--C73gtfMpLg"
   },
   "source": [
    "**Task 30:** Now we evaluate the performance of the registration on the obtained motion/displacement field. Calculate\n",
    "- the determinant of the Jacobian of the transformation <br/>\n",
    "$\\text{Jac} = \\operatorname{det}(\\nabla T(r)) = \\left| \\begin{bmatrix} \\frac{\\partial u_x}{\\partial x} + 1 & \\frac{\\partial u_x}{\\partial y} \\\\ \\frac{\\partial u_y}{\\partial x} & \\frac{\\partial u_y}{\\partial y} + 1 \\end{bmatrix} \\right|$\n",
    "- the divergence of the displacement field <br/>\n",
    "$\\operatorname{div}(u) = \\nabla \\cdot u = \\frac{\\partial u_x}{\\partial x} + \\frac{\\partial u_y}{\\partial y}$\n",
    "- absolute displacement field change in selected direction $p=\\lbrace x,y \\rbrace$ <br/>\n",
    "$| \\nabla u | = \\sqrt{ \\left(\\frac{\\partial u_x}{\\partial p}\\right)^2 + \\left(\\frac{\\partial u_y}{\\partial p}\\right)^2}$\n",
    "\n",
    "and plot the obtained metrics.\n",
    "\n",
    "*Optional:* You may also try to implement these functions with SimpleITK ([1](https://itk.org/Doxygen/html/classitk_1_1DisplacementFieldJacobianDeterminantFilter.html), [2](https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1GradientImageFilter.html))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OeqELKWzWMv5"
   },
   "outputs": [],
   "source": [
    "# direct\n",
    "uxdx = np.gradient(ux, axis=0)\n",
    "uxdy = np.gradient(ux, axis=1)\n",
    "uydx = np.gradient(uy, axis=0)\n",
    "uydy = np.gradient(uy, axis=1)\n",
    "\n",
    "jac = (uxdx+1)*(uydy+1) - uxdy*uydx\n",
    "div = uxdx + uydy\n",
    "absdisp_x = np.sqrt(np.power(uxdx, 2) + np.power(uydx, 2))\n",
    "absdisp_y = np.sqrt(np.power(uxdy, 2) + np.power(uydy, 2))\n",
    "\n",
    "plot([jac, div, absdisp_x, absdisp_y], title='Jac | div | AbsDispX | AbsDispY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qgzs7hgM9Bku"
   },
   "outputs": [],
   "source": [
    "# SimpleITK\n",
    "flows = np.stack([ux, uy], -1)\n",
    "sitk_displacement_field = sitk.GetImageFromArray(flows, isVector=True)\n",
    "jacobian_det_volume = sitk.DisplacementFieldJacobianDeterminant(sitk_displacement_field)\n",
    "jac = sitk.GetArrayViewFromImage(jacobian_det_volume)\n",
    "sitk_ux = sitk.GetImageFromArray(ux) # , isVector=True)\n",
    "sitk_uy = sitk.GetImageFromArray(uy) # , isVector=True)\n",
    "filter = sitk.GradientImageFilter()\n",
    "#filter.SetUseImageDirection(True)\n",
    "uxd = filter.Execute(sitk_ux)\n",
    "uyd = filter.Execute(sitk_uy)\n",
    "uxd_np = sitk.GetArrayViewFromImage(uxd)\n",
    "uyd_np = sitk.GetArrayViewFromImage(uyd)\n",
    "div = uxd_np[..., 0] + uyd_np[..., 1]\n",
    "absdisp_x = np.sqrt(np.power(uxd_np[..., 0], 2) + np.power(uyd_np[..., 0], 2))\n",
    "absdisp_y = np.sqrt(np.power(uxd_np[..., 1], 2) + np.power(uyd_np[..., 1], 2))\n",
    "\n",
    "plot([jac, div, absdisp_x, absdisp_y], title='Jac | div | AbsDispX | AbsDispY')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j77zl576Wmll"
   },
   "source": [
    "Depending on your simulated motion you can identify with these metrics, directional changes or contraction/expansion of the moved tissues."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U3pP2tPN-u87"
   },
   "source": [
    "**Task 31:** Invert the flow fields by changing the registration direction (i.e. swapping moving and fixed image) and compare against direct inversion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ETIaeQlr-uqd"
   },
   "outputs": [],
   "source": [
    "# redo registration\n",
    "uBy, uBx = optical_flow_tvl1(img_moving, img_fixed)\n",
    "img_warped_B = warp_2D(img_fixed, np.stack([uBx, uBy], -1))\n",
    "\n",
    "# direct inversion\n",
    "uBx_direct = -ux\n",
    "uBy_direct = -uy\n",
    "img_warped_B_direct = warp_2D(img_fixed, np.stack([uBx_direct, uBy_direct], -1))\n",
    "\n",
    "# plot the results\n",
    "plot([img_warped, img_warped_B, img_warped_B_direct], title='Warped: mov->fix | fix->mov (inverted reg.) | fix->mov (direct inv.)')\n",
    "\n",
    "# plot the flow fields\n",
    "plot_flow([uBx, uBy], 'inverted registration')\n",
    "plot_flow([uBx_direct, uBy_direct], 'direct inversion')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0P9Zh58yFmS4"
   },
   "source": [
    "**Task 32:** Calculate the:\n",
    "- end-point-error (EPE): $\\text{EPE} = \\|u^{(F)} - u^{(B)}\\|_2 = \\sqrt{|u_x^{(F)}-u_x^{(B)}|^2 + |u_y^{(F)}-u_y^{(B)}|^2}$\n",
    "- end-angulation-error (EAE): $\\text{EAE} = \\arg(u^{(F)}, u^{(B)})$\n",
    "\n",
    "between the backward flow fields (direct inversion or inverse registration) and the forward flow fields. Plot the EPE and EAE as flow fields: $\\text{EPE} \\cdot \\exp{(1i \\cdot \\text{EAE})}$ to have a color-coded error map (using `plot_flow()`).\n",
    "\n",
    "*Note:* EPE and EAE can be calculated between any two flows to investigate their similarity/dissimilarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_mcPFCS_Hppf"
   },
   "outputs": [],
   "source": [
    "flow = np.stack([ux, uy], -1)\n",
    "flowB = np.stack([uBx, uBy], -1)\n",
    "flowB_direct = np.stack([uBx_direct, uBy_direct], -1)\n",
    "\n",
    "def epe(flow1, flow2):\n",
    "  return np.sqrt(np.sum(np.power(np.abs(flow1 - flow2),2), axis=-1))\n",
    "\n",
    "def eae(flow1, flow2):\n",
    "  return np.real(np.arccos((1 + np.sum(flow1 * flow2, axis=-1)) / (np.sqrt(1 + np.sum(np.power(flow1, 2), axis=-1)) * np.sqrt(1 + np.sum(np.power(flow2, 2), axis=-1))))) * 180 / np.pi\n",
    "\n",
    "EPE = epe(flowB, flow)\n",
    "EPE_direct = epe(flowB_direct, flow)\n",
    "\n",
    "EAE = eae(flowB, flow)\n",
    "EAE_direct = eae(flowB_direct, flow)\n",
    "\n",
    "#plot([EPE, EAE], title='inverted registration: EPE | EAE')\n",
    "#plot([EPE_direct, EAE_direct], title='direct inversion: EPE | EAE')\n",
    "flowtest = EPE * np.exp(1j * EAE)\n",
    "testx = np.real(flowtest)\n",
    "testy = np.imag(flowtest)\n",
    "plot_flow(np.stack([testx, testy], -1), title='Isomorphism: inverted registration')\n",
    "flowtest = EPE_direct * np.exp(1j * EAE_direct)\n",
    "testx = np.real(flowtest)\n",
    "testy = np.imag(flowtest)\n",
    "plot_flow(np.stack([testx, testy], -1), title='Isomorphism: direct inversion')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4E8OgIVxjc-V"
   },
   "source": [
    "**Task 33:** Check if the flows are isomorph by warping a random image. Plot the respective errors for the backward flows (inverse registration and direct inversion)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S7gN6tisXBE_"
   },
   "outputs": [],
   "source": [
    "tmp = np.random.normal(size=np.shape(img_fixed))\n",
    "iso = warp_2D(warp_2D(tmp, np.stack([uBx, uBy], -1)), np.stack([ux, uy], -1))\n",
    "\n",
    "iso_direct = warp_2D(warp_2D(tmp, np.stack([uBx_direct, uBy_direct], -1)), np.stack([ux, uy], -1))\n",
    "\n",
    "plot([tmp, iso, iso_direct, np.abs(tmp-iso), np.abs(tmp-iso_direct), np.abs(iso-iso_direct)], title='random fixed | inverted reg. | direct inversion | error (inverted reg.) | error (direct inversion) | error (inverted reg. <> direct inversion)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YgryboxXNWCi"
   },
   "source": [
    "Alternatively for an affine motion, we can also analytically calculate the motion field from the affine motion matrix. We provide the function `get_transform()` to return the flow fields:\n",
    "\n",
    "```\n",
    "def get_transform(img, p):\n",
    "    # img      input image to be transformed\n",
    "    # p        affine transformation parameters\n",
    "    #          3D (rank(img) == 3): t_x, t_y, t_z, \\phi [°], \\theta [°], \\psi [°], G_{xy}, G_{xz}, G_{yz}, S_x, S_y, S_z\n",
    "    #          2D (rank(img) == 2): t_x, t_y, \\phi [°], G_{xy}, S_x, S_y\n",
    "    # return   deformation field\n",
    "```\n",
    "\n",
    "**Task 34:** You may rerun the previous tasks using the analytically derived motion field and compare it against the one obtained from registration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4W0qngndetG6"
   },
   "outputs": [],
   "source": [
    "flow_ana = get_transform(img_rad, p)\n",
    "flow_reg = np.stack([ux, uy], -1)\n",
    "\n",
    "EPE = epe(flow_ana, flow_reg)\n",
    "EAE = eae(flow_ana, flow_reg)\n",
    "\n",
    "plot_flow(flow_ana, title='analytical flow')\n",
    "plot_flow(flow_reg, title='registered flow')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8p7CUayIJr8Z"
   },
   "source": [
    "### B-Spline\n",
    "We will now perform a B-Spline based registration on the same pair of motion images as above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ymyeZhlYxoC1"
   },
   "source": [
    "**Task 35:** Perform a cubic B-Spline registration, aligning the moving image to the fixed image. Please use a Mean-Squared-Error metric, gradient descent optimizer, pyramidal multi-resolution registration with B-Spline mesh grid sizes of $[10, 6, 2, 1]$.\n",
    "\n",
    "*Hint:* You may use [SimpleITK](https://simpleitk.readthedocs.io/en/master/link_ImageRegistrationMethodBSpline1_docs.html) for this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yvvMQYx79lm6"
   },
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "\n",
    "def command_iteration(method):\n",
    "    print(\n",
    "        f\"{method.GetOptimizerIteration():3} \"\n",
    "        + f\"= {method.GetMetricValue():10.5f} \"\n",
    "        + f\": {method.GetOptimizerPosition()}\"\n",
    "    )\n",
    "\n",
    "def command_multi_iteration(method):\n",
    "    print(\"--------- Resolution Changing ---------\")\n",
    "\n",
    "fixed = sitk.GetImageFromArray(img_fixed)\n",
    "moving = sitk.GetImageFromArray(img_moving)\n",
    "\n",
    "transformDomainMeshSize = [10] * moving.GetDimension()\n",
    "tx = sitk.BSplineTransformInitializer(fixed, transformDomainMeshSize)\n",
    "\n",
    "R = sitk.ImageRegistrationMethod()\n",
    "R.SetMetricAsMeanSquares()\n",
    "R.SetOptimizerAsRegularStepGradientDescent(4.0, 0.01, 200)\n",
    "R.SetOptimizerScalesFromPhysicalShift()\n",
    "R.SetInitialTransform(tx) # sitk.TranslationTransform(fixed.GetDimension())\n",
    "R.SetInterpolator(sitk.sitkLinear)\n",
    "R.SetShrinkFactorsPerLevel([6, 2, 1])\n",
    "R.SetSmoothingSigmasPerLevel([6, 2, 1])\n",
    "\n",
    "R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))\n",
    "R.AddCommand(\n",
    "    sitk.sitkMultiResolutionIterationEvent, lambda: command_multi_iteration(R)\n",
    ")\n",
    "\n",
    "outTx = R.Execute(fixed, moving)\n",
    "\n",
    "print(\"-------\")\n",
    "print(outTx)\n",
    "print(f\"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}\")\n",
    "print(f\" Iteration: {R.GetOptimizerIteration()}\")\n",
    "print(f\" Metric value: {R.GetMetricValue()}\")\n",
    "\n",
    "outFlow = sitk.TransformToDisplacementField(outTx,\n",
    "                                  sitk.sitkVectorFloat64,\n",
    "                                  fixed.GetSize(),\n",
    "                                  fixed.GetOrigin(),\n",
    "                                  fixed.GetSpacing(),\n",
    "                                  fixed.GetDirection())\n",
    "\n",
    "outFlow_np = sitk.GetArrayFromImage(outFlow)\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "resampler.SetReferenceImage(fixed)\n",
    "resampler.SetInterpolator(sitk.sitkLinear)\n",
    "resampler.SetDefaultPixelValue(1)\n",
    "resampler.SetTransform(outTx)\n",
    "warped = resampler.Execute(moving)\n",
    "warped_np = sitk.GetArrayFromImage(warped)\n",
    "\n",
    "# composing in sitk\n",
    "simg1 = sitk.Cast(sitk.RescaleIntensity(fixed), sitk.sitkUInt8)\n",
    "simg2 = sitk.Cast(sitk.RescaleIntensity(warped), sitk.sitkUInt8)\n",
    "cimg = sitk.Compose(simg1, simg2, simg1 // 2.0 + simg2 // 2.0)\n",
    "cimg_np = sitk.GetArrayFromImage(cimg)\n",
    "\n",
    "# plotting the results\n",
    "plot_flow(outFlow_np)\n",
    "plot([img_fixed, img_moving, warped_np, np.abs(warped_np - img_fixed)], title='fixed | moving | warped | error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FYjiXbQv7pRt"
   },
   "source": [
    "**Task 36:** Obtain the backward motion field by [displacement field inversion](https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1InverseDisplacementFieldImageFilter.html) and compare it against a registration with swapped moving and fixed image. Plot the flow fields and the warped images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VJLmGXbZC50Q"
   },
   "outputs": [],
   "source": [
    "# flow field inversion\n",
    "R = sitk.InverseDisplacementFieldImageFilter()\n",
    "R.SetReferenceImage(moving)\n",
    "outFlowB = R.Execute(outFlow)\n",
    "outTxB = sitk.DisplacementFieldTransform(outFlowB)\n",
    "outFlowB = sitk.TransformToDisplacementField(outTxB,\n",
    "                                  sitk.sitkVectorFloat64,\n",
    "                                  fixed.GetSize(),\n",
    "                                  fixed.GetOrigin(),\n",
    "                                  fixed.GetSpacing(),\n",
    "                                  fixed.GetDirection())\n",
    "outFlowB_np = sitk.GetArrayFromImage(outFlowB)\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "resampler.SetReferenceImage(moving)\n",
    "resampler.SetInterpolator(sitk.sitkLinear)\n",
    "resampler.SetDefaultPixelValue(1)\n",
    "resampler.SetTransform(outTxB)\n",
    "warpedB = resampler.Execute(fixed)\n",
    "warpedB_np = sitk.GetArrayFromImage(warpedB)\n",
    "\n",
    "# inverse registration\n",
    "def command_iteration(method):\n",
    "    print(\n",
    "        f\"{method.GetOptimizerIteration():3} \"\n",
    "        + f\"= {method.GetMetricValue():10.5f} \"\n",
    "        + f\": {method.GetOptimizerPosition()}\"\n",
    "    )\n",
    "\n",
    "def command_multi_iteration(method):\n",
    "    print(\"--------- Resolution Changing ---------\")\n",
    "\n",
    "fixed = sitk.GetImageFromArray(img_fixed)\n",
    "moving = sitk.GetImageFromArray(img_moving)\n",
    "\n",
    "transformDomainMeshSize = [10] * moving.GetDimension()\n",
    "tx = sitk.BSplineTransformInitializer(fixed, transformDomainMeshSize)\n",
    "\n",
    "R = sitk.ImageRegistrationMethod()\n",
    "R.SetMetricAsMeanSquares()\n",
    "R.SetOptimizerAsRegularStepGradientDescent(4.0, 0.01, 200)\n",
    "R.SetOptimizerScalesFromPhysicalShift()\n",
    "R.SetInitialTransform(tx) # sitk.TranslationTransform(fixed.GetDimension())\n",
    "R.SetInterpolator(sitk.sitkLinear)\n",
    "R.SetShrinkFactorsPerLevel([6, 2, 1])\n",
    "R.SetSmoothingSigmasPerLevel([6, 2, 1])\n",
    "\n",
    "R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))\n",
    "R.AddCommand(\n",
    "    sitk.sitkMultiResolutionIterationEvent, lambda: command_multi_iteration(R)\n",
    ")\n",
    "\n",
    "outTxBreg = R.Execute(moving, fixed)\n",
    "\n",
    "print(\"-------\")\n",
    "print(outTxBreg)\n",
    "print(f\"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}\")\n",
    "print(f\" Iteration: {R.GetOptimizerIteration()}\")\n",
    "print(f\" Metric value: {R.GetMetricValue()}\")\n",
    "\n",
    "outFlowBreg = sitk.TransformToDisplacementField(outTxBreg,\n",
    "                                  sitk.sitkVectorFloat64,\n",
    "                                  fixed.GetSize(),\n",
    "                                  fixed.GetOrigin(),\n",
    "                                  fixed.GetSpacing(),\n",
    "                                  fixed.GetDirection())\n",
    "\n",
    "outFlowBreg_np = sitk.GetArrayFromImage(outFlowBreg)\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "resampler.SetReferenceImage(moving)\n",
    "resampler.SetInterpolator(sitk.sitkLinear)\n",
    "resampler.SetDefaultPixelValue(1)\n",
    "resampler.SetTransform(outTxBreg)\n",
    "warpedBreg = resampler.Execute(fixed)\n",
    "warpedBreg_np = sitk.GetArrayFromImage(warpedBreg)\n",
    "\n",
    "# plotting the results\n",
    "plot_flow(outFlowB_np, title='inversion')\n",
    "plot([img_fixed, img_moving, warpedB_np, np.abs(warpedB_np - img_moving)], title='inversion: fixed | moving | warped | error')\n",
    "\n",
    "plot_flow(outFlowBreg_np, title='registration')\n",
    "plot([img_fixed, img_moving, warpedBreg_np, np.abs(warpedBreg_np - img_moving)], title='registration: fixed | moving | warped | error')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MvZdw0J0GHkZ"
   },
   "source": [
    "Did you spot any difference? They are most likely neglectable, but computational demand is higher for a swapped registration than inverting the flow field."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ixxGpceQ8Rk7"
   },
   "source": [
    "**Task 37:** Evaluate the isomorphism of the flow fields: $u^{(F)}(u^{(B)}(r)) = r$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OLvDdp5BGeKf"
   },
   "outputs": [],
   "source": [
    "iso_inv_w = warp_2D(warp_2D(np.ones_like(img_fixed), outFlow_np), outFlowB_np)\n",
    "iso_reg_w = warp_2D(warp_2D(np.ones_like(img_fixed), outFlow_np), outFlowBreg_np)\n",
    "plot([iso_inv_w, iso_reg_w], title=['inversion', 'registration'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_lKqUqFHXg2S"
   },
   "source": [
    "**Task 38:** Change the metric to a Mattes Mutual Information metric and compare the registration with the previous one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v-gr3a4ptoGM"
   },
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "\n",
    "def command_iteration(method):\n",
    "    print(\n",
    "        f\"{method.GetOptimizerIteration():3} \"\n",
    "        + f\"= {method.GetMetricValue():10.5f} \"\n",
    "        + f\": {method.GetOptimizerPosition()}\"\n",
    "    )\n",
    "\n",
    "def command_multi_iteration(method):\n",
    "    print(\"--------- Resolution Changing ---------\")\n",
    "\n",
    "fixed = sitk.GetImageFromArray(img_fixed)\n",
    "moving = sitk.GetImageFromArray(img_moving)\n",
    "\n",
    "transformDomainMeshSize = [10] * moving.GetDimension()\n",
    "tx = sitk.BSplineTransformInitializer(fixed, transformDomainMeshSize)\n",
    "\n",
    "R = sitk.ImageRegistrationMethod()\n",
    "R.SetMetricAsMattesMutualInformation(24)\n",
    "R.SetMetricSamplingPercentage(0.10, sitk.sitkWallClock)\n",
    "R.SetMetricSamplingStrategy(R.RANDOM)\n",
    "R.SetOptimizerAsRegularStepGradientDescent(4.0, 0.01, 200)\n",
    "R.SetOptimizerScalesFromPhysicalShift()\n",
    "R.SetInitialTransform(tx) # sitk.TranslationTransform(fixed.GetDimension())\n",
    "R.SetInterpolator(sitk.sitkLinear)\n",
    "R.SetShrinkFactorsPerLevel([6, 2, 1])\n",
    "R.SetSmoothingSigmasPerLevel([6, 2, 1])\n",
    "\n",
    "R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))\n",
    "R.AddCommand(\n",
    "    sitk.sitkMultiResolutionIterationEvent, lambda: command_multi_iteration(R)\n",
    ")\n",
    "\n",
    "outTx = R.Execute(fixed, moving)\n",
    "\n",
    "print(\"-------\")\n",
    "print(outTx)\n",
    "print(f\"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}\")\n",
    "print(f\" Iteration: {R.GetOptimizerIteration()}\")\n",
    "print(f\" Metric value: {R.GetMetricValue()}\")\n",
    "\n",
    "outFlow = sitk.TransformToDisplacementField(outTx,\n",
    "                                  sitk.sitkVectorFloat64,\n",
    "                                  fixed.GetSize(),\n",
    "                                  fixed.GetOrigin(),\n",
    "                                  fixed.GetSpacing(),\n",
    "                                  fixed.GetDirection())\n",
    "\n",
    "outFlow_mi_np = sitk.GetArrayFromImage(outFlow)\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "resampler.SetReferenceImage(fixed)\n",
    "resampler.SetInterpolator(sitk.sitkLinear)\n",
    "resampler.SetDefaultPixelValue(1)\n",
    "resampler.SetTransform(outTx)\n",
    "warped = resampler.Execute(moving)\n",
    "warped_mi_np = sitk.GetArrayFromImage(warped)\n",
    "\n",
    "# composing in sitk\n",
    "#simg1 = sitk.Cast(sitk.RescaleIntensity(fixed), sitk.sitkUInt8)\n",
    "#simg2 = sitk.Cast(sitk.RescaleIntensity(warped), sitk.sitkUInt8)\n",
    "#cimg = sitk.Compose(simg1, simg2, simg1 // 2.0 + simg2 // 2.0)\n",
    "#cimg_np = sitk.GetArrayFromImage(cimg)\n",
    "\n",
    "# plotting the results\n",
    "plot_flow(outFlow_np)\n",
    "plot([img_fixed, img_moving, warped_mi_np, np.abs(warped_mi_np - img_fixed)], title='fixed | moving | warped | error')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KNW_0mUUuctj"
   },
   "outputs": [],
   "source": [
    "# comparison between MSE and MI metric\n",
    "plot_flow([outFlow_np, outFlow_mi_np, np.abs(outFlow_np - outFlow_mi_np)], title=['MSE', 'MI', 'diff: MSE <> MI'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xZ4dc8bcx8Aa"
   },
   "source": [
    "You have most likely observed some incorrect motion estimates at the image boundaries. These are occuring in case you have too many degrees of freedom in your motion model (e.g. B-Spline) when applying to an affine motion (6 degrees of freedom in 2D).\n",
    "\n",
    "**Task 39:** Adjust the used transformation that better suits your underlying motion with which you simulated your image pair `img_fixed` and `img_moving`.\n",
    "\n",
    "*Hint:* You can find an overview of the available transformations [here](https://simpleitk.readthedocs.io/en/master/registrationOverview.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7b_Q73BL7gae"
   },
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "\n",
    "def command_iteration(method):\n",
    "    print(\n",
    "        f\"{method.GetOptimizerIteration():3} \"\n",
    "        + f\"= {method.GetMetricValue():10.5f} \"\n",
    "        + f\": {method.GetOptimizerPosition()}\"\n",
    "    )\n",
    "\n",
    "def command_multi_iteration(method):\n",
    "    print(\"--------- Resolution Changing ---------\")\n",
    "\n",
    "fixed = sitk.GetImageFromArray(img_fixed)\n",
    "moving = sitk.GetImageFromArray(img_moving)\n",
    "\n",
    "R = sitk.ImageRegistrationMethod()\n",
    "R.SetMetricAsMeanSquares()\n",
    "R.SetOptimizerAsRegularStepGradientDescent(4.0, 0.01, 200)\n",
    "R.SetInitialTransform(sitk.TranslationTransform(fixed.GetDimension()))\n",
    "R.SetInterpolator(sitk.sitkLinear)\n",
    "\n",
    "\n",
    "R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))\n",
    "R.AddCommand(\n",
    "    sitk.sitkMultiResolutionIterationEvent, lambda: command_multi_iteration(R)\n",
    ")\n",
    "\n",
    "outTx = R.Execute(fixed, moving)\n",
    "\n",
    "print(\"-------\")\n",
    "print(outTx)\n",
    "print(f\"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}\")\n",
    "print(f\" Iteration: {R.GetOptimizerIteration()}\")\n",
    "print(f\" Metric value: {R.GetMetricValue()}\")\n",
    "\n",
    "outFlow = sitk.TransformToDisplacementField(outTx,\n",
    "                                  sitk.sitkVectorFloat64,\n",
    "                                  fixed.GetSize(),\n",
    "                                  fixed.GetOrigin(),\n",
    "                                  fixed.GetSpacing(),\n",
    "                                  fixed.GetDirection())\n",
    "\n",
    "outFlow_trans_np = sitk.GetArrayFromImage(outFlow)\n",
    "\n",
    "resampler = sitk.ResampleImageFilter()\n",
    "resampler.SetReferenceImage(fixed)\n",
    "resampler.SetInterpolator(sitk.sitkLinear)\n",
    "resampler.SetDefaultPixelValue(1)\n",
    "resampler.SetTransform(outTx)\n",
    "warped = resampler.Execute(moving)\n",
    "warped_trans_np = sitk.GetArrayFromImage(warped)\n",
    "\n",
    "# composing in sitk\n",
    "simg1 = sitk.Cast(sitk.RescaleIntensity(fixed), sitk.sitkUInt8)\n",
    "simg2 = sitk.Cast(sitk.RescaleIntensity(warped), sitk.sitkUInt8)\n",
    "cimg = sitk.Compose(simg1, simg2, simg1 // 2.0 + simg2 // 2.0)\n",
    "cimg_np = sitk.GetArrayFromImage(cimg)\n",
    "\n",
    "# plotting the results\n",
    "plot_flow(outFlow_trans_np)\n",
    "plot([img_fixed, img_moving, warped_trans_np, np.abs(warped_trans_np - img_fixed)], title='fixed | moving | warped | error')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bYQDF9h-HZFq"
   },
   "source": [
    "**Task 40:** Compare the flows obtained from the optical flow with the SimpleITK flows (B-Spline and other transformation as set in **Tasks 24/27, 35, 38** and **39**)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UiPA_xUnHlO2"
   },
   "outputs": [],
   "source": [
    "plot_flow([outFlow_np, outFlow_mi_np, outFlow_trans_np, flows], title=['B-Spline MSE', 'B-Spline MI', 'Translational MSE', 'Optical Flow'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fFNiRTtT9K9K"
   },
   "source": [
    "## Deep learning-based image registration\n",
    "In recent years, deep learning solutions have also been proposed to address the image registration. A short (and by far complete) overview of deep learning image registrations for (single-modality) MRI is provided here:\n",
    "- [LAPNet](https://ieeexplore.ieee.org/document/9478906)\n",
    "- [CARMEN](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6677286/)\n",
    "- [Voxelmorph](https://ieeexplore.ieee.org/document/8633930)\n",
    "- [FlowNet](https://arxiv.org/abs/1504.06852)\n",
    "- [RAFT](https://arxiv.org/abs/2003.12039)\n",
    "- [MRAFT](https://link.springer.com/chapter/10.1007/978-3-030-88552-6_2)\n",
    "- [MoCo-MoDL](https://onlinelibrary.wiley.com/doi/10.1002/mrm.28851)\n",
    "- [GRN](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9095964/)\n",
    "- [DDR](https://www.sciencedirect.com/science/article/abs/pii/S0167865517301708)\n",
    "- [DIRNet](https://link.springer.com/chapter/10.1007/978-3-319-67558-9_24)\n",
    "- [AIRNet](https://arxiv.org/abs/1810.02583)\n",
    "- [BIRNet](https://www.sciencedirect.com/science/article/abs/pii/S1361841519300805)\n",
    "- [CVAE](https://arxiv.org/abs/1804.07172)\n",
    "- [FCN](https://arxiv.org/abs/1801.04012)\n",
    "- [DeformableReg](https://openreview.net/forum?id=HkmkmW2jM)\n",
    "- [ICNet](https://arxiv.org/abs/1809.03443)\n",
    "- ...\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1qBqr5henDEr"
   },
   "source": [
    "We will employ in this example the pre-trained Voxelmorph architecture - as we do not have the time and computing ressources to properly train a neural network from scratch here 😉 But we will dive into a simple exemplary training of an image registration network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t-TGuzBBnFba"
   },
   "source": [
    "Download and extract the example data and pre-trained networks of Voxelmorph [here](https://surfer.nmr.mgh.harvard.edu/pub/data/voxelmorph/tutorial_data.tar.gz).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cfUuIsdKwkgY"
   },
   "outputs": [],
   "source": [
    "!wget https://surfer.nmr.mgh.harvard.edu/pub/data/voxelmorph/tutorial_data.tar.gz -O data.tar.gz\n",
    "!tar -xzvf data.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y5-7u5uDodcl"
   },
   "source": [
    "**Task 41:** Register the moving and the fixed image using the pre-trained Voxelmorph architecture `brain_2d_smooth.h5`. You frist need to specify the network architecture and hyperparameters as detailed here:\n",
    "\n",
    "\n",
    "```\n",
    "import voxelmorph as vxm\n",
    "nb_features = [\n",
    "    [32, 32, 32, 32],         # encoder features\n",
    "    [32, 32, 32, 32, 32, 16]  # decoder features\n",
    "]\n",
    "vol_shape = tuple(np.shape(img_fixed))\n",
    "vxm_model = vxm.networks.VxmDense(vol_shape, nb_features, int_steps=0)\n",
    "\n",
    "losses = ['mse', vxm.losses.Grad('l2').loss]\n",
    "loss_weights = [1, 0.01]\n",
    "\n",
    "vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)\n",
    "\n",
    "vxm_model.load_weights('brain_2d_smooth.h5')\n",
    "```\n",
    "\n",
    "The input to the network is defined as `inp = [img_moving, img_fixed]` and the prediction can be carried out with `out = vxm_model.predict(inp)` which returns a list containing the warped image and the motion field. Plot the warped image (using `plot()`) and the motion field (using `plot_flow()`).\n",
    "\n",
    "\n",
    "*Hint:* Do not forget to add a batch and channel dimension (additional first and last dimension) to the inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LnT6GLwlwDY8"
   },
   "outputs": [],
   "source": [
    "import voxelmorph as vxm\n",
    "nb_features = [\n",
    "    [32, 32, 32, 32],         # encoder features\n",
    "    [32, 32, 32, 32, 32, 16]  # decoder features\n",
    "]\n",
    "vol_shape = tuple(np.shape(img_fixed))\n",
    "vxm_model = vxm.networks.VxmDense(vol_shape, nb_features, int_steps=0)\n",
    "\n",
    "losses = ['mse', vxm.losses.Grad('l2').loss]\n",
    "loss_weights = [1, 0.01]\n",
    "\n",
    "vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)\n",
    "\n",
    "vxm_model.load_weights('brain_2d_smooth.h5')\n",
    "\n",
    "val_input = [img_moving[np.newaxis, :, :, np.newaxis], img_fixed[np.newaxis, :, :, np.newaxis]]\n",
    "val_pred = vxm_model.predict(val_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "35EGsYFp2Mm3"
   },
   "outputs": [],
   "source": [
    "plot([img_fixed, img_moving, np.squeeze(val_pred[0], (0, -1))], title='fixed | moving | warped')\n",
    "plot_flow(np.squeeze(val_pred[1], 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fHTfFUuNEhnY"
   },
   "source": [
    "**Task 42:** Perform the above registration using the pre-trained network `brain_2d_no_smooth.h5`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RK4WXvPBE_X_"
   },
   "outputs": [],
   "source": [
    "import voxelmorph as vxm\n",
    "nb_features = [\n",
    "    [32, 32, 32, 32],         # encoder features\n",
    "    [32, 32, 32, 32, 32, 16]  # decoder features\n",
    "]\n",
    "vol_shape = tuple(np.shape(img_fixed))\n",
    "vxm_model = vxm.networks.VxmDense(vol_shape, nb_features, int_steps=0)\n",
    "\n",
    "losses = ['mse', vxm.losses.Grad('l2').loss]\n",
    "loss_weights = [1, 0.01]\n",
    "\n",
    "vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)\n",
    "\n",
    "vxm_model.load_weights('brain_2d_no_smooth.h5')\n",
    "\n",
    "val_input = [img_moving[np.newaxis, :, :, np.newaxis], img_fixed[np.newaxis, :, :, np.newaxis]]\n",
    "val_pred = vxm_model.predict(val_input)\n",
    "plot([img_fixed, img_moving, np.squeeze(val_pred[0], (0, -1))], title='fixed | moving | warped')\n",
    "plot_flow(np.squeeze(val_pred[1], 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FxfJXLGo5STf"
   },
   "source": [
    "Most likely you have obtained in both cases a quite artistic looking warped image together with a smooth motion field. The reason for this is, that the pre-trained architectures were trained to match sagittal slices between different subjects (inter-subject registration) and different slices from the same subject to enable an atlas-style mapping.\n",
    "\n",
    "Since we want to register the same moving slices onto each other, we need to retrain the architecture. We will use the axial slices of the two validation volumes in `subj1.npz` and `subj2.npz` provided with the Voxelmorph example data. The subsqeuent steps follow the [Voxelmorph tutorial](http://tutorial.voxelmorph.net/) loosely which can be used as guidance or inspiration.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yrO8EGIQIY6q"
   },
   "source": [
    "We first need to rotate the images into axial orientation (`np.rot90(..., 1)`) and subsequently zeropad the images to a size of $X \\times Y \\times \\text{slices} = 256 \\times 256 \\times 224$ (`zpad(..., (256, 256, 224))`). Move the axial slices now in the batch dimension (first dimension).\n",
    "\n",
    "**Task 43:** Prepare the training data taken from the `subj1.npz` and `subj2.npz` as described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2uFcO9HH6bqI"
   },
   "outputs": [],
   "source": [
    "# prepare data from 3D val_volumes\n",
    "val_volume_1 = np.load('subj1.npz')['vol']\n",
    "val_volume_2 = np.load('subj2.npz')['vol']\n",
    "traindata = np.concatenate((np.transpose(zpad(np.rot90(val_volume_1, 1), (256, 256, 224)), (2, 0, 1)),\n",
    "                            np.transpose(zpad(np.rot90(val_volume_2, 1), (256, 256, 224)), (2, 0, 1))),0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KOP0MvxyI-IZ"
   },
   "source": [
    "**Task 44:** Implement a data generator for neural network training that inputs to the network the `[moving_image, fixed_image]` in which the `moving_image` is obtained by an affine transformation (using `transform_img()`) of the `fixed_image` with randomly selected affine transformation parameters `p`. Decide for a suitable range of transformation parameters `p`.\n",
    "\n",
    "*Hint:* You may refer to the function `vxm_data_generator()` in the [Voxelmorph tutorial](http://tutorial.voxelmorph.net/) for guidance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XrKFkz12Kgjv"
   },
   "outputs": [],
   "source": [
    "def vxm_data_generator_new(x_data, batch_size=32, p_min=[0, 0, 0, 0, 80, 80], p_max=[100, 100, 90, 50, 100, 100]):\n",
    "    \"\"\"\n",
    "    Generator that takes in data of size [N, H, W], and yields data for\n",
    "    our custom vxm model.\n",
    "\n",
    "    inputs:  moving [bs, H, W, 1], fixed image [bs, H, W, 1]\n",
    "    outputs: moved image [bs, H, W, 1], zero-gradient [bs, H, W, 2]\n",
    "    \"\"\"\n",
    "\n",
    "    # preliminary sizing\n",
    "    vol_shape = x_data.shape[1:] # extract data shape\n",
    "    ndims = len(vol_shape)\n",
    "\n",
    "    # prepare a zero array the size of the deformation\n",
    "    # to enable regularization\n",
    "    zero_phi = np.zeros([batch_size, *vol_shape, ndims])\n",
    "\n",
    "    while True:\n",
    "        # prepare inputs:\n",
    "        # images need to be of the size [batch_size, H, W, 1]\n",
    "        idx1 = np.random.randint(0, x_data.shape[0], size=batch_size)\n",
    "        fixed_images = x_data[idx1, ..., np.newaxis]\n",
    "        moving_images = np.zeros_like(fixed_images)\n",
    "        for b in range(np.shape(fixed_images)[0]):\n",
    "          p_rand = np.random.randint(p_min, p_max).astype(np.float32)\n",
    "          p_rand[-3:] = p_rand[-3:]/100\n",
    "          moving_images[b, ...] = transform_img(fixed_images[b, ..., 0], p_rand)[..., np.newaxis]\n",
    "\n",
    "        inputs = [moving_images, fixed_images]\n",
    "\n",
    "        # prepare outputs (the 'true' moved image):\n",
    "        # of course, we don't have this, but we know we want to compare\n",
    "        # the resulting moved image with the fixed image.\n",
    "        # we also wish to penalize the deformation field.\n",
    "        outputs = [fixed_images, zero_phi]\n",
    "\n",
    "        yield (inputs, outputs)\n",
    "\n",
    "\n",
    "train_generator = vxm_data_generator_new(traindata, batch_size=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VFsXKYBY_tK1"
   },
   "source": [
    "**Task 45:** Now train the Voxelmorph architecture using your prepared training data generator: `vxm_model.fit_generator(train_generator, epochs=nb_epochs, steps_per_epoch=steps_per_epoch, verbose=2)`. You may decide for the amount of epochs to train `nb_epochs` and the `steps_per_epoch` ($\\leq 448$). Depending on your selection, training may take up to several minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PMTp-1Kg-a1Y"
   },
   "outputs": [],
   "source": [
    "nb_features = [\n",
    "    [32, 32, 32, 32],         # encoder features\n",
    "    [32, 32, 32, 32, 32, 16]  # decoder features\n",
    "]\n",
    "vol_shape = tuple(np.shape(img_fixed))\n",
    "vxm_model = vxm.networks.VxmDense(vol_shape, nb_features, int_steps=0)\n",
    "\n",
    "losses = ['mse', vxm.losses.Grad('l2').loss]\n",
    "loss_weights = [1, 0.01]\n",
    "\n",
    "vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)\n",
    "hist = vxm_model.fit_generator(train_generator, epochs=20, steps_per_epoch=200, verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r8_vuaR-ML50"
   },
   "source": [
    "**Task 46:** Register the images using your trained network. Plot the warped image and the obtained flow.\n",
    "\n",
    "*Hint:* `model.predict()` provides you the warped image and the flow field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SKhz8fH2A5mk"
   },
   "outputs": [],
   "source": [
    "val_input = [img_moving[np.newaxis, :, :, np.newaxis], img_fixed[np.newaxis, :, :, np.newaxis]]\n",
    "val_pred = vxm_model.predict(val_input)\n",
    "\n",
    "plot([img_fixed, img_moving, np.squeeze(val_pred[0], (0, -1))], title='fixed | moving | warped')\n",
    "plot_flow(np.squeeze(val_pred[1], 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Iw5T-utLMpUH"
   },
   "source": [
    "Were you able to train a good image registration network? A few points to consider when training your registration network:\n",
    "- the complexity of your motion (rigid, affine, elastic)\n",
    "- availability of enough diverse training data\n",
    "- network architecture parametrization (kernel sizes, feature maps/channels, ...)\n",
    "- hyperparameters (loss, learning rate, ...)\n",
    "\n",
    "**Task 47:** You can play around now by changing some of the above points to see if you are able to obtain a better result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P9g7EGBP9REy"
   },
   "source": [
    "# 3. Motion-compensated image reconstruction\n",
    "<a name=\"MotionCompensated\"></a>\n",
    "In this part, we will embed the estimated motion in the image reconstruction task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ShBjSpjVpLet"
   },
   "source": [
    "Before diving into the motion-compensated reconstruction task, we will shortly recap an iterative SENSE image reconstruction.\n",
    "\n",
    "The goal is to recover the clean image $x$, which is obtained by (undersampled) k-space data $y$ and corrupted by additive Gaussian white noise $n$,\n",
    "$$ y = Ax + n. $$\n",
    "The rawdata $y$ was aquired with multiple receive coils. The linear operator $A$ denotes the mapping from image space to k-space.\n",
    "\n",
    "Consequently we can define and implement the multi-coil forward operator $A$ in `mriForwardOp(image, mask, smaps)` mapping from `image` to k-space and the adjoint operator $A^*$ in `mriAdjointOp(kspace, mask, smaps)` mapping from `kspace` to image domain. `smaps` describes the coil sensitivity maps, and `mask` depicts the sampling trajectory.\n",
    "\n",
    "We consider now the following minimization problem:\n",
    "\n",
    "$$ \\min_x  E(x,y) = \\min_x \\frac{1}{2} \\Vert Ax - y \\Vert_2^2 .$$\n",
    "\n",
    "While in image denoising we are still able to compute a closed-form solution for this problem, this is not feasible for the task of MRI reconstruction anymore. We instead use first-order optimization methods and solve this by Gradient Descent:\n",
    "$$ x^{t+1} = x^{t} - \\alpha \\nabla_x E(x,y) $$\n",
    "$$ x^{t+1} = x^{t} - \\alpha A^* (Ax^t - y) $$\n",
    "\n",
    "\n",
    "**Suggested Readings:**\n",
    "\n",
    "Pruessmann et al. [SENSE: Sensitivity encoding for fast MRI](https://onlinelibrary.wiley.com/doi/abs/10.1002/%28SICI%291522-2594%28199911%2942%3A5%3C952%3A%3AAID-MRM16%3E3.0.CO%3B2-S) Magnetic Resonance in Medicine, 43(5):952-962, 1999."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f9T8kwH8FBen"
   },
   "source": [
    "## iterative SENSE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6u6QFiNVrC1U"
   },
   "source": [
    "Below is the implementation of the iterative SENSE reconstruction that follows the above formalism. The interface can be used to reconstruction Cartesian and non-Cartesian data - only determined by the definition of the respective forward and adjoint operators.\n",
    "\n",
    "The Cartesian operators are given by:\n",
    "- numpy: `mriForwardOp()` and `mriAdjointOp()` in `utils/mri.py`\n",
    "- Tensorflow: `MulticoilForwardOp` and `MulticoilAdjointOp` in `merlintf.keras.layers.mri`.\n",
    "\n",
    "The non-Cartesian operators are defined as:\n",
    "- numpy: `GPUNUFFTOp.forward()` and `GPUNUFFTOp.adjoint()` in class `GPUNUFFTOp` of `utils/mri.py`\n",
    "- Tensorflow: `GPUNUFFTFwd` and `GPUNUFFTAdj` in `utils/mri.py`.\n",
    "\n",
    "The iterative SENSE reconstruction `iterativeSENSE()` in `utisl/mri.py` is automatically executed on GPU (if available and optox is used) and has the following interface:\n",
    "\n",
    "```\n",
    "def iterativeSENSE(kspace, smap=None, mask=None, noisy=None, dcf=None, flow=None,\n",
    "            fwdop=MulticoilForwardOp, adjop=MulticoilAdjointOp,\n",
    "            add_batch_dim=True, max_iter=10, tol=1e-12, weight_init=1.0, weight_scale=1.0, use_optox=False):\n",
    "    # kspace        raw k-space data as [X, Y, coils] which will be converted to:\n",
    "    #               Cartesian + no-motion compensation: [batch, coils, X, Y] or [batch, coils, X, Y, Z] or [batch, coils, time, X, Y] or [batch, coils, time, X, Y, Z] (numpy array)\n",
    "    #               Cartesian + motion-compensation / non-Cartesian + no-motion/motion-comp.: [batch, X, Y, coils]\n",
    "    # smap          coil sensitivity maps with same shape as kspace (or singleton dimension for time) (numpy array)\n",
    "    # mask          subsampling including/excluding soft-weights with same shape as kspace (no-motion-comp) and shape: X, Y, coils, T (motion-comp) (numpy array)\n",
    "    # noisy         initialiaztion for reconstructed image, if None it is created from A^H(kspace) (numpy array)\n",
    "    # dcf           density compensation function (only non-Cartesian) (numpy array)\n",
    "    # flow          flow field (numpy array)\n",
    "    # fwdop         forward operator A\n",
    "    # adjop         adjoint operator A^H\n",
    "    # add_batch_dim automatically append batch dimension for GPU execution\n",
    "    # max_iter      maximum number of iterations for CG/iterative SENSE\n",
    "    # tol           tolerance for stopping condition for CG/iterative SENSE\n",
    "    # weight_init   initial weighting for lambda regularization parameter\n",
    "    # weight_scale  scaling factor for lambda regularization parameter\n",
    "    # return:       reconstructed image (numpy array)\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uRvXenuDgA5N"
   },
   "source": [
    "In the following, we will only focus on the numpy implementation (due to RAM limitations and GPU availability), but codes for GPU implementation are also provided (Attention: The GPU implementations require further improvement as currently handovers between CPU and GPU, i.e. type conversion between numpy and Tensorflow are happening).\n",
    "\n",
    "\n",
    "We will see two examples below how to perform an iterative SENSE reconstruction for a 2D Cartesian and a 2D non-Cartesian imaging."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7WEd96fNs5Bd"
   },
   "source": [
    "### 2D Cartesian reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RzVh2Z0a4EF1"
   },
   "outputs": [],
   "source": [
    "from merlintf.keras.layers.mri import MulticoilForwardOp, MulticoilAdjointOp\n",
    "\n",
    "# prepare data\n",
    "kspace = fft2c(img)\n",
    "kspace_espirit = kspace[:,:,np.newaxis,:]\n",
    "smaps_espirit = espirit.espirit(kspace_espirit, 8, 20, 0.05, 0)\n",
    "smaps = smaps_espirit[:,:,0,:,0]\n",
    "\n",
    "# Cartesian\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = MulticoilForwardOp(center=True, coil_axis=-3, channel_dim_defined=False)\n",
    "  AH = MulticoilAdjointOp(center=True, coil_axis=-3, channel_dim_defined=False)\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = mriForwardOp\n",
    "  AH = mriAdjointOp\n",
    "  use_optox = False\n",
    "\n",
    "img_recon = iterativeSENSE(kspace, smaps, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "plot(img_recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-RzTwUHys8r7"
   },
   "source": [
    "### 2D non-Cartesian reconstruction\n",
    "<a name=\"2DnonCart\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3t9ZEFe20GL2"
   },
   "outputs": [],
   "source": [
    "from mri.operators import NonCartesianFFT\n",
    "\n",
    "# prepare data\n",
    "nRead = np.amax(np.shape(img)[0:2])\n",
    "img_qfov = zpad(img, (nRead, nRead, np.shape(img)[2])).astype(np.complex64)\n",
    "img_qfov = img_qfov[:,:,np.newaxis,:]\n",
    "kspace_qfov = fft2c(img_qfov)\n",
    "smaps_rad_espirit = espirit.espirit(kspace_qfov, 8, 20, 0.05, 0)\n",
    "smaps_rad = smaps_rad_espirit[:,:,0,:,0]\n",
    "csm = np.transpose(smaps_rad, (2, 0, 1))\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "# non-Cartesian\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = GPUNUFFTFwd(nRead, kpos, csm, dcf)\n",
    "  AH = GPUNUFFTAdj(nRead, kpos, csm, dcf)\n",
    "  #kspace_radial = np.squeeze(A(numpy2tensor(rss(img_qfov[:,:,0,:]), add_batch_dim=True, add_channel_dim=False)).numpy())\n",
    "  kspace_radial = A.op.op(rss(img_qfov[:,:,0,:]))\n",
    "  use_optox = True\n",
    "else:\n",
    "  nufft = GPUNUFFTOp(kpos, csm, dcf, nRead)\n",
    "  A = nufft.forward\n",
    "  AH = nufft.adjoint\n",
    "  kspace_radial = nufft.forward(rss(img_qfov[:,:,0,:]))\n",
    "  use_optox = False\n",
    "\n",
    "img_recon_rad = iterativeSENSE(kspace_radial, csm, kpos, dcf=dcf, fwdop=A, adjop=AH, weight_scale=0.1, max_iter=2, use_optox=use_optox)\n",
    "plot(img_recon_rad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FaifvQyGeHMp"
   },
   "source": [
    "## Motion-compensated Batchelor-type reconstruction\n",
    "\n",
    "Affine motion has a well-understood representation in k-space, however generalized elastic motion does not. An approach to correct for this type of motion is to consider the forward model of an acquisition corrupted by motion and solve the corresponding reconstruction problem [1]:\n",
    "\n",
    "$$ \\hat{x} = \\arg\\min_x \\frac{1}{2} \\Vert Ex - y \\Vert_2^2 $$\n",
    "\n",
    "where the forward model is given by:\n",
    "\n",
    "$$ E = \\sum_n A_n F C U_n $$\n",
    "\n",
    "Under this representation, $A_n$ are the k-space sampling matrices for each motion state $n$, $F$ is the Fourier transform, $C$ are the coil sensitivities and $U_n$ are the motion fields for each motion state. The motion fields $U_n$ may be considered as sparse motion matrices where each row contains the interpolation weights $w_i^j$ associated with a given displacement. If nearest neighbour interpolation is considered, then $U_n$ become permutation matrices, i.e. only one interpolant per row and with value $w_i^j=1$. If linear interpolation is considered (as it is in this example implementation), then each row contains four interpolants (in 2D), such that $ \\sum_j w_i^j = 1 $.\n",
    "\n",
    "![](https://github.com/lab-midas/ismrm-moco-workshop/blob/master/doc/sparse_motion_mat.png?raw=true)\n",
    "\n",
    "For a motion-compensated Batchelor-type reconstruction, we first need a simple representation of the motion as a sparse matrix.\n",
    "\n",
    "This is implemented in `get_sparse_motion_matrix()` (converts dense flow field to sparse motion matrix) and `apply_sparse_motion()` (warps an image via matrix-vector multiplication) from `utils/motioncomp.py`. In the Cartesian case, these functions are invoked from the forward `BatchForwardOp()` $E$ and adjoint `BatchAdjointOp()` $E^H$ operators which are handed over to the iterative SENSE reconstruction `iterativeSENSE()`.\n",
    "\n",
    "The forward and adjoint operators that take into account the motion model are defined as `BatchForwardOp()` and `BatchAdjointOp()` in `utils/mri.py`.\n",
    "\n",
    "<br>\n",
    "\n",
    "[1] Batchelor PG, Atkinson D, Irarrazaval P, Hill DL, Hajnal J, Larkman D. Matrix description of general motion correction applied to multishot images. Magnetic Resonance in Medicine: An Official Journal of the International Society for Magnetic Resonance in Medicine. 2005 Nov;54(5):1273-80."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TTQq4QkBbuCe"
   },
   "source": [
    "**Task 48:** Apply an affine transformation to the coil-combined image (`img_cc = mriAdjointOp(kspace, np.ones_like(kspace), smaps`) and plot the motion-free and motion-affected image side-by-side (using `plot()`).\n",
    "\n",
    "*Hint:* Please see **Task 1** for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kM9t24FrcisE"
   },
   "outputs": [],
   "source": [
    "kspace = fft2c(img)\n",
    "kspace_espirit = kspace[:,:,np.newaxis,:]\n",
    "smaps_espirit = espirit.espirit(kspace_espirit, 8, 20, 0.05, 0)\n",
    "smaps = smaps_espirit[:,:,0,:,0]\n",
    "smaps = smaps / rss(smaps)[:,:,np.newaxis]  # normalize coil sensitivities\n",
    "img_cc = mriAdjointOp(kspace, np.ones_like(kspace), smaps)\n",
    "\n",
    "img_affine = transform_img(img_cc, [-10, 30, 45, 0.1, 1, 1])\n",
    "plot([img_cc, img_affine])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V1AhUCoefloM"
   },
   "source": [
    "**Task 49:** Simulate a motion-corrupted acquisition such that every 2nd phase-encoding step is affected by an abrupt (and time-constant) affine motion. You can use the function `simulate_motion()` in `utils/motionsim.py`. Plot the motion-free and motion-affected scans.\n",
    "\n",
    "*Hint:* Please see **Task 5** for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zBIPREIRf2sD"
   },
   "outputs": [],
   "source": [
    "p = [-10, 30, 45, 0.1, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(kspace)[1])\n",
    "mask_motion[::2] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, np.ones_like(kspace), p)\n",
    "img_motion = mriAdjointOp(kspace_motion, np.ones_like(kspace), smaps)\n",
    "\n",
    "plot([img_cc, img_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e5cyCDHTjqJk"
   },
   "source": [
    "**Task 50:** Simulate a translational motion-affected scan (every 2nd phase-encoding line is affected by motion) using `simulate_motion()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oqz_VrF_jptj"
   },
   "outputs": [],
   "source": [
    "p = [-10, 30, 0, 0, 1, 1]\n",
    "mask_motion = np.zeros(np.shape(kspace)[1])\n",
    "mask_motion[::2] = 1\n",
    "p = mask_motion[:, np.newaxis] * np.asarray(p)[np.newaxis, :]\n",
    "\n",
    "kspace_motion, mask_motion = simulate_motion(img_cc, smaps, np.ones_like(kspace), p)\n",
    "img_motion = mriAdjointOp(kspace_motion, np.ones_like(kspace), smaps)\n",
    "\n",
    "plot([img_cc, img_motion])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TB-WQdxRPk11"
   },
   "source": [
    "Before we can perform a motion-compensated image reconstruction, we need the deformation fields which are used in the forward `BatchForwardOp()` and adjoint operators `BatchAdjointOp()`.\n",
    "\n",
    "We can either perform an image registration between the motion-free and motion-affected scan, please see **Tasks 23** to **47** in [2. Motion estimation / image registration](#MotionEstimation), or since we are applying an explicit (and known) affine motion to the image, we can analytically calculate the deformation field for it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d0yNYmX5jztX"
   },
   "source": [
    "**Task 51:** Write the functions `get_affine_matrix()` which gives you the affine motion matrix for the motion parameters (see [Motion artifact appearance](#MotionArtifactAppearance)) and plug this into `get_deformation_field_from_affine()` which applies the affine motion matrix at each voxel location of the image to return the deformation field of shape $X \\times Y \\times 2$.\n",
    "\n",
    "*Attention:* For `get_deformation_field_from_affine()` (and also in the Matlab code), the deformation field is usually overlaid with a grid which is used in the creation of the sparse matrix. For this, a grid is created via `np.mgrid` to which the affine motion matrix from `get_affine_matrix()` is multiplied. However, the Python code does not use this grid in the creation of the sparse matrix. We therefore provide here functions that add/remove the mesh from the flow field.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Na2E_WWyoXxC"
   },
   "outputs": [],
   "source": [
    "def get_affine_matrix(img, p):\n",
    "  if len(np.shape(img)) > 2:\n",
    "    raise \"Only 2D processing implemented\"\n",
    "  x, y = np.shape(img)\n",
    "  x /= 2\n",
    "  y /= 2\n",
    "\n",
    "  # translation from origin to point (x,y)\n",
    "  P1 = np.array([[1, 0, -x], [0, 1, -y], [0, 0, 1]])\n",
    "  # translation back to origin from point (x,y)\n",
    "  P2 = np.array([[1, 0, x], [0, 1, y], [0, 0, 1]])\n",
    "\n",
    "  # translation\n",
    "  T = np.array([[0, 0, p[0]], [0, 0, p[1]], [0, 0, 0]])\n",
    "\n",
    "  # rotation\n",
    "  radians = -np.pi * np.asarray(p[2]) / 180.\n",
    "  R = np.array([[np.cos(radians), -np.sin(radians), 0],[np.sin(radians), np.cos(radians), 0], [0, 0, 1]])\n",
    "\n",
    "  # shearing\n",
    "  G = np.array([[1, p[3], 0], [0, 1, 0], [0, 0, 1]])\n",
    "\n",
    "  # scaling\n",
    "  S = np.array([[p[4], 0, 0], [0, p[5], 0], [0, 0, 1]])\n",
    "\n",
    "  # affine matrix\n",
    "  return np.linalg.multi_dot([P2, G, S, R, P1]) + T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jM2wLal0k-4J"
   },
   "outputs": [],
   "source": [
    "def get_deformation_field_from_affine(img, affine_mat):\n",
    "  height, width = np.shape(img)\n",
    "  gridY, gridX = np.mgrid[1:width+1, 1:height+1]\n",
    "  positions = np.concatenate((np.transpose(gridX.flatten()[:, np.newaxis], (1, 0)), np.transpose(gridY.flatten()[:, np.newaxis], (1, 0)), np.ones((1, img.size))), axis=0)  # with overlaid grid\n",
    "  #positions = np.concatenate((np.ones((1, img.size)), np.ones((1, img.size)), np.ones((1, img.size))), axis=0)\n",
    "  u = affine_mat @ positions\n",
    "  u = np.transpose(np.reshape(np.transpose(u[0:2, :], (1, 0)), (height, width, 2)), (1, 0, 2))\n",
    "  return u\n",
    "\n",
    "def add_mesh_to_def(u):\n",
    "  height, width, nd = np.shape(u)\n",
    "  gridY, gridX = np.mgrid[1:width+1, 1:height+1]\n",
    "  q = np.zeros(np.shape(u))\n",
    "  q[:,:,0] = u[:,:,0] + gridY\n",
    "  q[:,:,1] = u[:,:,1] + gridX\n",
    "  return q\n",
    "\n",
    "def remove_mesh_from_def(u):\n",
    "  height, width, nd = np.shape(u)\n",
    "  gridY, gridX = np.mgrid[1:width+1, 1:height+1]\n",
    "  q = np.zeros(np.shape(u))\n",
    "  q[:,:,0] = u[:,:,0] - gridY\n",
    "  q[:,:,1] = u[:,:,1] - gridX\n",
    "  return q\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FYACGpvgk8MI"
   },
   "source": [
    "**Task 52:** Obtain the deformation field via `get_flow()` in `utils/motionsim.py` for a translational displacement $t_x=10$ and plot the flow field using `plot_flow()`. Compare the warped image using the obtained flow field with the directly transformed image from `transform_img()`.\n",
    "\n",
    "*Hint:* Please see **Task 1** for more information on transforming an image, and **Task 25** for more information on flow plotting. Do not forget to scale the images to a common range before comparing them.<br/>\n",
    "The function `get_flow(img, p)` is a wrapper around `get_affine_matrix()`, `get_deformation_field_from_affine()` and `remove_mesh_from_def()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BUDNxr1IlS95"
   },
   "outputs": [],
   "source": [
    "# Converting affine matrices to dense motion fields will include an underlying\n",
    "# coordinates mesh, which can be removed if needed\n",
    "p = [-10, 0, 0, 0, 1, 1]\n",
    "aff = get_affine_matrix(img_cc, p)\n",
    "u = get_deformation_field_from_affine(img_cc, aff)\n",
    "#q = remove_mesh_from_def(u)\n",
    "q = get_flow(img_cc, p)\n",
    "\n",
    "img_in = minmaxscale(np.abs(img_cc), [0, 1])\n",
    "img_trans = transform_img(img_in, p)\n",
    "img_warped = warp_2D(img_in, q)\n",
    "mae = np.sum(np.sum(np.abs(img_warped-img_trans)))\n",
    "plot([img_in, img_trans, np.abs(img_trans - img_in), img_warped, np.abs(img_warped-img_in), np.abs(img_warped-img_trans)])\n",
    "plot_flow(u) # with mesh\n",
    "plot_flow(q) # without mesh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dDAVjVO3MOjx"
   },
   "source": [
    "#### Cartesian imaging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "adnmpXaclwi1"
   },
   "source": [
    "**Task 53:** Apply the forward operator `BatchForwardOp()` to the fully-sampled motion-free image using the flow field $u$ from **Task 52** to obtain a motion-corrupted acquisition. Subsequently apply the adjoint operator `BatchAdjointOp()` on this motion-corrupted data to obtain the zero-filled motion-corrected reconstruction. Compare this against the non-motion-corrected zero-filled reconstruction (using `mriAdjointOp()`). Simulate that every 2nd phase-encoding line is affected by motion. Plot the images.\n",
    "\n",
    "*Hint:* Since we need a motion field for all time frames, we also require the identity flow field for the first time point which we can stack into a rank 4 matrix: $X \\times Y \\times 2 \\ (u_x, u_y) \\times 2 \\ (\\text{Time})$ <br/>\n",
    "In a similar fashion we need to stack the masks so that they have a shape: $X \\times Y \\times \\text{Channels} \\times \\text{Time}$.<br/>\n",
    "Please use `get_flow()` to obtain the flow field from the affine motion parameters and `get_sparse_motion_matrix()` in `utils/motioncomp.py` to obtain the sparse motion matrices. A stack of vertical sparse tensors\n",
    "\n",
    "```\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1, ...])\n",
    "```\n",
    "\n",
    "is used in the Batchelor forward and adjoint operators.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kq0G2iCzi-tz"
   },
   "outputs": [],
   "source": [
    "mask_fs = np.ones_like(kspace)\n",
    "p = [-20, 40, 0, 0, 1, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "mask_motion1 = np.zeros_like(kspace)\n",
    "mask_motion1[::2] = 1\n",
    "mask_motion2 = np.ones_like(kspace)\n",
    "mask_motion2 -= mask_motion1\n",
    "masks_motion = np.stack([mask_motion1, mask_motion2], -1)\n",
    "masks = np.stack([mask_fs, mask_fs], -1)\n",
    "\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NA-a6vGslKVq"
   },
   "outputs": [],
   "source": [
    "kspace_motion = BatchForwardOp(img_cc, masks_motion, smaps, smm)\n",
    "\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks_motion, smaps, smm)\n",
    "img_no_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smm)\n",
    "plot([img_cc, img_no_motion_corr_zf, img_motion_corr_zf], title='motion-free | motion-affected | zero-filled motion-comp.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "408sSylVTkMQ"
   },
   "source": [
    "Residual artefacts can be observed. This is because just applying the adjoint operator leads to a naive inversion, which corresponds to an image-based motion correction approach (i.e. reconstruct separate motion-resolved aliased images and then warp them to a common reference). This only works if the motion and the (undersampled) Fourier transform for each motion state commute, which is generally not the case.\n",
    "\n",
    "Before going to an iterative SENSE reconstruction, we will solve the linear inverse problem by a simple gradient descent algorithm: compute the local gradient and take a small step in the direction that minimizes it. It iteratively solves for the solution.\n",
    "\n",
    "**Attention**: The iterative SENSE reconstruction (on the GPU)requires some time. The gradient descent provides a faster approximate solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y-DTGT6UMo9J"
   },
   "outputs": [],
   "source": [
    "# Gradient descent solver for linear inverse problem\n",
    "\n",
    "def grad_descent_batch(kspace, masks_motion, smaps, smm, max_it):\n",
    "\n",
    "  Ny, Nx, Nc = np.shape(kspace)\n",
    "\n",
    "  #Right hand side\n",
    "  m = BatchAdjointOp(kspace_motion, masks_motion, smaps, smm)\n",
    "  ExEHxm = BatchAdjointOp(BatchForwardOp(m, masks_motion, smaps, smm), masks_motion, smaps, smm)\n",
    "\n",
    "  # initial residual\n",
    "  r = m - ExEHxm\n",
    "\n",
    "  rho = np.zeros((Ny,Nx,max_it+1)) + 1j * np.zeros((Ny,Nx,max_it+1))\n",
    "\n",
    "  # 0-th iteration, i.e., the \"zero-filled\" recon\n",
    "  rho[:,:,0] = m\n",
    "\n",
    "  for it in range(max_it):\n",
    "\n",
    "      # Compute alpha\n",
    "      aux = BatchAdjointOp(BatchForwardOp(r, masks_motion, smaps, smm), masks_motion, smaps, smm)\n",
    "      top = np.sum(np.transpose(np.matrix.flatten(np.conj(r))) * np.matrix.flatten(r))\n",
    "      bot = np.sum(np.transpose(np.matrix.flatten(np.conj(r))) * np.matrix.flatten(aux))\n",
    "      alpha = (top / bot)\n",
    "\n",
    "      # Update solution and residual\n",
    "      m = m + alpha*r\n",
    "      r = r - alpha*BatchAdjointOp(BatchForwardOp(r, masks_motion, smaps, smm), masks_motion, smaps, smm)\n",
    "\n",
    "      # Save output\n",
    "      rho[:,:,it+1] = m\n",
    "\n",
    "  return rho"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rkcY4VRgSYJL"
   },
   "outputs": [],
   "source": [
    "# run batch recon iteratively\n",
    "recon_its = grad_descent_batch(kspace_motion, masks_motion, smaps, smm, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sA5maMHrTIX2"
   },
   "outputs": [],
   "source": [
    "# Now we can see how aliasing artefacts are completely removed when the reconstruction is properly solved\n",
    "print(np.shape(recon_its))\n",
    "plot(recon_its, title='gradient descent iterations')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4FjAn_kOoJJl"
   },
   "source": [
    "**Task 54:** Perform a motion-compensated image reconstruction in `iterativeSENSE()` using the ingredients from **Task 52** and **Task 53**, and a non-motion compensated image reconstruction. Compare the reconstructed images with the zero-filled reconstructions.\n",
    "\n",
    "*Hint:* You need to use the functions `BatchelorFwd` and `BatchelorAdj` inside the iterative SENSE which are Tensorflow layers wrapping around the `BatchForwardOp()` and `BatchAdjointOp()`. The iterative SENSE requires the non-sparse motion flows in order to wrap them into tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gU29Qey-kzIh"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks_motion, smaps, smm, 3)\n",
    "\n",
    "plot([img_cc, img_motion_corr_zf, img_no_motion_corr_zf, recon_its[...,-1]], title='motion-free | zero-filled motion-comp. | zero-filled no motion-comp. | motion-comp. recon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Rvnu7NAypZm8"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "img_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, mask=masks_motion, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = MulticoilForwardOp(center=True, coil_axis=-3, channel_dim_defined=False)\n",
    "  AH = MulticoilAdjointOp(center=True, coil_axis=-3, channel_dim_defined=False)\n",
    "else:\n",
    "  A = mriForwardOp\n",
    "  AH = mriAdjointOp\n",
    "\n",
    "img_no_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_zf, img_no_motion_corr_zf, img_motion_corr_recon, img_no_motion_corr_recon], title='motion-free | zero-filled motion-comp. | zero-filled no motion-comp. | motion-comp. recon | no-motion-comp. recon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AsICg_5mGf9C"
   },
   "outputs": [],
   "source": [
    "A = BatchForwardOp\n",
    "AH = BatchAdjointOp\n",
    "noisy = None\n",
    "dcf = None\n",
    "mask = None\n",
    "sm,a\n",
    "\n",
    "if dcf is not None:\n",
    "    bradial = True\n",
    "else:\n",
    "    bradial = False\n",
    "\n",
    "if flow is not None:\n",
    "    motioncomp = True\n",
    "else:\n",
    "    motioncomp = False\n",
    "\n",
    "if bradial:\n",
    "    Nx, Nspokes = np.shape(kspace)\n",
    "    Ny = Nx\n",
    "    Nc = 1\n",
    "else:\n",
    "    Nx, Ny, Nc = np.shape(kspace)\n",
    "if motioncomp:\n",
    "    Nt = np.shape(flow)[-1]\n",
    "else:\n",
    "    Nt = 1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uQ3nTfJz34yQ"
   },
   "source": [
    "**Task 55:** Repeat the **Tasks 52** to **54** for a rotational motion, i.e. perform a motion-compensated image reconstruction on rotational motion-affected data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FolzZPs04hFl"
   },
   "outputs": [],
   "source": [
    "# Fully-sampled sampling masks\n",
    "mask_fs = np.ones_like(kspace)\n",
    "\n",
    "# using meshless flow\n",
    "p = [0, 0, 45, 0, 1, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GdWBeFvmSrur"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Residual aliasing artefacts present with the naive inversion, but these are\n",
    "# corrected with the iterative solution\n",
    "\n",
    "kspace_motion = BatchForwardOp(img_cc, masks_motion, smaps, smm)\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks_motion, smaps, smm)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks_motion, smaps, smm, 3)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf,img_motion_corr_zf], title='motion-free | zero-filled no motion compensation | zero-filled motion compensation')\n",
    "plot(recon_its, title='motion-compensated recon: gradient descent iterations')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m8Ry3TRkXnn_"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "\n",
    "img_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, mask=masks_motion, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf,img_motion_corr_zf, img_motion_corr_recon], title='motion-free | zero-filled no motion compensation | zero-filled motion compensation | motion-comp. recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3cGOpKRh4uCF"
   },
   "source": [
    "**Task 56:** Now we want to examine the impact of the coil sensitivity maps on a motion-compensated reconstruction for a rotational motion (as in **Task 55**). Therefore, perform a motion-compensated reconstruction with and without coil sensitivity maps. Plot and compare the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p4q93l0JUIgr"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Considerable artefacts visible for recon without coil sensitivies. This is the actual\n",
    "# image with the highest data consistency. It has removed all of the motion, however it is\n",
    "# riddled with aliasing artefacts due to the k-space gaps opened by the rotational motion.\n",
    "recon_its_no_csm = grad_descent_batch(kspace_motion, masks_motion, np.ones_like(smaps), smm, 3)\n",
    "\n",
    "plot(recon_its, title='recon w/ csm')\n",
    "plot(recon_its_no_csm, title='recon w/o csm')\n",
    "#plot(recon_its_no_csm[:,:,0])\n",
    "#plot(recon_its_no_csm[:,:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "TQfjwfsb5D9g"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "# with csm\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "#img_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, fwdop=A, adjop=AH)  # result from previous cell\n",
    "\n",
    "# without csm\n",
    "img_motion_corr_recon_no_csm = iterativeSENSE(kspace_motion, np.ones_like(smaps), mask=masks, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_motion_corr_recon, img_motion_corr_recon_no_csm], title='with CSM | without CSM')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-kuDz6H95kt2"
   },
   "source": [
    "**Task 57:** Perform a motion-compensated image reconstruction for rotational and translational motion in a (regular and random) undersampled Cartesian acquisition. Each motion state has a different undersampling mask (but same acceleration factor).\n",
    "\n",
    "*Hint:* Please see **Task 11** for creating an undersampling mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ow4XIUHd52Sn"
   },
   "outputs": [],
   "source": [
    "p = [0, 10, 30, 0, 1, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "\n",
    "#sampling_mask1 = generate_mask(R=3, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='random')\n",
    "#sampling_mask2 = generate_mask(R=3, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='random')\n",
    "\n",
    "mask_fs = np.ones_like(kspace)\n",
    "sampling_mask1 = np.zeros_like(kspace)\n",
    "sampling_mask1[:,1::3,:] = 1\n",
    "sampling_mask2 = np.zeros_like(kspace)\n",
    "sampling_mask2[:,2::3,:] = 1\n",
    "\n",
    "masks = np.stack([sampling_mask1, sampling_mask2], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DQbQuseJWMrH"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# More challenging case, with undersampling and rotational motion. Coil sensitivities\n",
    "# help, but residual aliasing remains\n",
    "kspace_motion = BatchForwardOp(img_cc, masks, smaps, smm)\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smm)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks, smaps, smm, 10)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf,img_motion_corr_zf], title='motion-free | zero-filled no motion compensation | zero-filled motion compensation')\n",
    "plot(recon_its[:,:,0::3], title='motion-compensated recon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QbhhKHwXVvlD"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "masks = np.stack([sampling_mask1, sampling_mask2], -1)\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smm)\n",
    "img_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_zf, img_motion_corr_recon], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CTqVLS3F_Z23"
   },
   "source": [
    "**Task 58:** Perform a motion-compensated image reconstruction for an affine motion ($S_x=0.9$) on regular undersampled Cartesian data with acceleration factor 2. The k-space is additionally corrupted by additive white Gaussian noise. Coil sensitivity maps are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JhjfyF88_wzp"
   },
   "outputs": [],
   "source": [
    "p = [0, 0, 0, 0, 0.9, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "\n",
    "sampling_mask = generate_mask(R=2, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "masks = np.stack([np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1])), np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1]))], -1)\n",
    "\n",
    "# motion-affected k-space\n",
    "kspace_motion = BatchForwardOp(img_cc, masks, smaps, smm)\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 1E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6tKEDjgqMIYR"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Reconstruct undersampled and noisy data corrupted by scaling motion. Residual\n",
    "# aliasing and noise amplification is visible in the reconstructed images.\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smm)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks, smaps, smm, 9)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf, img_motion_corr_zf], title='motion-free | zero-filled no motion-comp. | zero-filled motion-comp.')\n",
    "plot(recon_its[:,:,0::3], title='motion-compensated recon')\n",
    "#plot(recon_its[:,:,0])\n",
    "#plot(recon_its[:,:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3Y4Ohsy_L4pj"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smm)\n",
    "img_motion_corr_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_zf, img_motion_corr_recon], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yfFRh92KBKzG"
   },
   "source": [
    "**Task 59:** We want to see now what happens if we do not have access to an accurate motion model, i.e. if the image registration was inaccurate. For this purpose, we will adjust the affine motion parameters by an error term (e.g. 10% error), so that the actual obtained deformation field is not exactly reflecting the underlying applied motion. Perform a motion-compensated image reconstruction for an affine motion on regular undersampled Cartesian data. Additionally the k-space is affected by additive white Gaussian noise (refer to **Task 58**). Coil sensitivity maps are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-pnIyNmGCQo2"
   },
   "outputs": [],
   "source": [
    "# sampling\n",
    "sampling_mask = generate_mask(R=2, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "masks = np.stack([np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1])), np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1]))], -1)\n",
    "\n",
    "# using get_def_field and meshless\n",
    "p = [0, 0, 0, 0, 0.9, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "\n",
    "# motion-affected k-space\n",
    "kspace_motion = BatchForwardOp(img_cc, masks, smaps, smm)\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 1E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n",
    "\n",
    "# add error on motion field\n",
    "p_error = p\n",
    "p_error[4] *= 1.1  # 10% error\n",
    "\n",
    "ue = get_flow(img_cc, p_error)\n",
    "smm1e = get_sparse_motion_matrix(ue)\n",
    "\n",
    "smme = vstack([smm0, smm1e])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eSldFffPUaok"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Reconstruct undersampled and noisy data corrupted by scaling motion with\n",
    "# errors in the motion model. Errors in the motion model propagate into scaling motion\n",
    "# artefacts in the reconstructed images.\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks, smaps, smme, 9)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf,img_motion_corr_zf], title='motion-free | zero-filled no motion-comp. | zero-filled motion-comp.')\n",
    "plot(recon_its[:,:,0::3], title='motion-compensated recon')\n",
    "#plot(recon_its[:,:,0])\n",
    "#plot(recon_its[:,:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5RqdeHRBUaFq"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow_error\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smme\n",
    "  use_optox = False\n",
    "img_motion_corr_error_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_motion_corr_error_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_zf, img_motion_corr_error_recon], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wvvREolyBf3O"
   },
   "source": [
    "**Task 60:** In a next step we will regularize the reconstruction by providing a prior to help with the challenging case of motion-compensated image reconstruction in undersampled Cartesian data with k-space noise. The setting is otherwise similar to **Task 59**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aiSk3zrtcrfx"
   },
   "outputs": [],
   "source": [
    "# A simple version of a regularized gradient descent for problems of the type\n",
    "# ||Ex-k||_2 + lambda ||x-x_reg||_2, which is equivalent to solving the problem\n",
    "# Fx = b, where F = E*E^H + I*lambda and b = E^H*k + lambda*x_reg\n",
    "\n",
    "def reg_grad_descent_batch(kspace, masks_motion, smaps, smm, max_it,x_reg,lambda_reg):\n",
    "\n",
    "  Ny, Nx, Nc = np.shape(kspace)\n",
    "\n",
    "  #Right hand side\n",
    "  m = BatchAdjointOp(kspace_motion, masks_motion, smaps, smm) + lambda_reg*x_reg\n",
    "  ExEHxm = BatchAdjointOp(BatchForwardOp(m, masks_motion, smaps, smm), masks_motion, smaps, smm) + lambda_reg*m\n",
    "\n",
    "  # initial residual\n",
    "  r = m - ExEHxm\n",
    "\n",
    "  rho = np.zeros((Ny,Nx,max_it+1)) + 1j * np.zeros((Ny,Nx,max_it+1))\n",
    "\n",
    "  # 0-th iteration, i.e., the \"zero-filled\" recon\n",
    "  rho[:,:,0] = m\n",
    "\n",
    "  for it in range(max_it):\n",
    "\n",
    "      # Compute alpha\n",
    "      aux = BatchAdjointOp(BatchForwardOp(r, masks_motion, smaps, smm), masks_motion, smaps, smm) + lambda_reg*r\n",
    "      top = np.sum(np.transpose(np.matrix.flatten(np.conj(r))) * np.matrix.flatten(r))\n",
    "      bot = np.sum(np.transpose(np.matrix.flatten(np.conj(r))) * np.matrix.flatten(aux))\n",
    "      alpha = (top / bot)\n",
    "\n",
    "      # Update solution and residual\n",
    "      m = m + alpha*r\n",
    "      r = r - alpha*BatchAdjointOp(BatchForwardOp(r, masks_motion, smaps, smm), masks_motion, smaps, smm) + lambda_reg*r\n",
    "\n",
    "      # Save output\n",
    "      rho[:,:,it+1] = m\n",
    "\n",
    "  return rho\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4c5YxjhzDdpi"
   },
   "outputs": [],
   "source": [
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "# sampling\n",
    "sampling_mask = generate_mask(R=2, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "masks = np.stack([np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1])), np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1]))], -1)\n",
    "\n",
    "# using meshless u\n",
    "p = [0, 0, 0, 0, 0.9, 1]\n",
    "u = get_flow(img_cc, p)\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "\n",
    "# motion-affected k-space\n",
    "kspace_motion = BatchForwardOp(img_cc, masks, smaps, smm)\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 2E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n",
    "\n",
    "# add error on motion field\n",
    "p_error = p\n",
    "p_error[4] *= 1.1  # 10% error\n",
    "\n",
    "ue = get_flow(img_cc, p_error)\n",
    "smm1e = get_sparse_motion_matrix(ue)\n",
    "\n",
    "smme = vstack([smm0, smm1e])\n",
    "\n",
    "img_prior = gaussian_filter(img_cc, sigma=5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bDc2bkwJfkX9"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Reconstruct undersampled and noisy data corrupted by scaling motion with\n",
    "# errors in the motion model, but now regularized by a low resolution prior.\n",
    "# Some suppression of noise and aliasing is achieved, but superior performance\n",
    "# is achieved with a regularized conjugate gradient\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "maxit = 9\n",
    "lambda_reg = 1\n",
    "\n",
    "recon_its = grad_descent_batch(kspace_motion, masks, smaps, smme, maxit)\n",
    "recon_its_reg = reg_grad_descent_batch(kspace_motion, masks, smaps, smme, maxit,img_prior,lambda_reg)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf, img_motion_corr_zf], title='motion-free | zero-filled no motion-comp. | zero-filled motion-comp.')\n",
    "plot(recon_its[:,:,0::3], title='motion-compensated reconstruction')\n",
    "plot(recon_its_reg[:,:,0::3], title='regularized motion-compensated reconstruction')\n",
    "#plot(recon_its[:,:,-1])\n",
    "#plot(recon_its_reg[:,:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bStKWEJpfSPC"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow_error\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smme\n",
    "  use_optox = False\n",
    "img_motion_corr_error_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_motion_corr_error_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, noisy=img_prior, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_zf, img_motion_corr_error_recon], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oSztldtVLPtq"
   },
   "source": [
    "**Task 61:** Perform a similar motion-compensated image reconstruction as in **Task 60**, but now having three motion states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6x0D6mRZLPC9"
   },
   "outputs": [],
   "source": [
    "from scipy.ndimage import gaussian_filter\n",
    "\n",
    "p = [3, 6, 10, 0, 1, 1.05]\n",
    "u = get_flow(img_cc, p)\n",
    "\n",
    "p2 = [-4, -8, -20, 0, 0.95, 1]\n",
    "u2 = get_flow(img_cc, p2)\n",
    "\n",
    "u_id = get_flow(img_cc, [0, 0, 0, 0, 1, 1])\n",
    "\n",
    "flow = np.stack([u_id, u, u2], -1)\n",
    "\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "smm2 = get_sparse_motion_matrix(flow[:,:,:,2])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1, smm2])\n",
    "\n",
    "# sampling\n",
    "sampling_mask = generate_mask(R=2, nPE=np.shape(img_cc)[1], nFE=np.shape(img_cc)[0], mode='regular')\n",
    "masks = np.stack([np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1])),\n",
    "                  np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1])),\n",
    "                  np.tile(sampling_mask, (1, 1, np.shape(smaps)[-1]))], -1)\n",
    "\n",
    "\n",
    "kspace_motion = BatchForwardOp(img_cc, masks, smaps, smm)\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 2E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iQTIYBoyNpgV"
   },
   "outputs": [],
   "source": [
    "# add error on motion field\n",
    "p_error = np.multiply(p,[0.98, 0.98, 0.98, 0.98, 0.98, 0.98])\n",
    "u_error = get_flow(img_cc, p_error)\n",
    "\n",
    "p2_error = np.multiply(p2,[0.98, 0.98, 0.98, 0.98, 0.98, 0.98])\n",
    "u2_error = get_flow(img_cc, p2_error)\n",
    "\n",
    "flow_error = np.stack([u_id, u_error, u2_error], -1)\n",
    "\n",
    "smm1e = get_sparse_motion_matrix(u_error)\n",
    "smm2e = get_sparse_motion_matrix(u2_error)\n",
    "\n",
    "smme = vstack([smm0, smm1e, smm2e])\n",
    "\n",
    "# image prior\n",
    "img_prior = gaussian_filter(img_cc, sigma=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aPFB6S6Y4pmh"
   },
   "outputs": [],
   "source": [
    "# gradient descent\n",
    "\n",
    "# Reconstruct undersampled and noisy data corrupted by scaling motion with\n",
    "# errors in the motion model and regularized by a low resolution prior but with\n",
    "# 3 motion states (and less motion amplitude per state).\n",
    "# Some suppression of noise and aliasing is achieved, but superior performance\n",
    "# is achieved with a regularized conjugate gradient\n",
    "img_motion_corr_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_no_motion_corr_zf = mriAdjointOp(kspace_motion, mask_fs, smaps)\n",
    "\n",
    "maxit = 9\n",
    "lambda_reg = 1.5\n",
    "\n",
    "recon_true_motion = grad_descent_batch(kspace_motion, masks, smaps, smm, maxit)\n",
    "recon_its = grad_descent_batch(kspace_motion, masks, smaps, smme, maxit)\n",
    "recon_its_reg = reg_grad_descent_batch(kspace_motion, masks, smaps, smme, maxit,img_prior,lambda_reg)\n",
    "\n",
    "plot([img_cc, img_no_motion_corr_zf, img_motion_corr_zf], title='motion-free | zero-filled no motion-comp. | zero-filled motion-comp.')\n",
    "plot(recon_true_motion[:,:,0::3], title='motion-compensated reconstruction (true motion)')\n",
    "plot(recon_its[:,:,0::3], title='motion-compensated reconstruction (error motion)')\n",
    "plot(recon_its_reg[:,:,0::3], title='motion-compensated reconstruction (error motion) with regularization')\n",
    "#plot(recon_true_motion[:,:,-1])\n",
    "#plot(recon_its[:,:,-1])\n",
    "#plot(recon_its_reg[:,:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-bOBuwic4pFy"
   },
   "outputs": [],
   "source": [
    "# iterative SENSE\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorFwd()\n",
    "  AH = BatchelorAdj()\n",
    "  flowin = flow_error\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchForwardOp\n",
    "  AH = BatchAdjointOp\n",
    "  flowin = smme\n",
    "  use_optox = False\n",
    "img_motion_corr_error_zf = BatchAdjointOp(kspace_motion, masks, smaps, smme)\n",
    "img_motion_corr_error_recon = iterativeSENSE(kspace_motion, smaps, mask=masks, flow=flowin, noisy=img_prior, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_zf, img_motion_corr_error_recon], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SgVcUAbYVNk-"
   },
   "source": [
    "**Task 62:** Repeat the above **Task 61** with changing sampling pattern per motion state (acceleration factor and/or sampled points), motion errors on motion fields, number of motion states or changing additive noise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AY_0qIplMZyc"
   },
   "source": [
    "#### Non-Cartesian imaging\n",
    "**Attention:** If you are running this in Google Colab, you may not have enough RAM to execute the motion-compensated reconstruction together with GPUNUFFT. You need to run this locally, upgrade to a paid tier in Google Colab or have a look at the Matlab code which showcases these examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_pmQOdsa6_Sa"
   },
   "source": [
    "**Task 63:** Perform a motion-compensated image reconstruction for an affine motion on fully-sampled non-Cartesian data. The respective forward and adjoint operators are given in `utils/mri.py` as `BatchGPUNUFFTForwardOp()` and `BatchGPUNUFFTAdjointOp()`. Inside the iterative SENSE reconstruction, you need to use for Tensorflow: `BatchelorGPUNUFFTFwd` and `BatchelorGPUNUFFTAdj` (Tensorflow wrappers similar to the Cartesian case around the previously mentioned motion-compensated GPUNUFFT operators); for numpy: `BatchGPUNUFFTForwardOp()` and `BatchGPUNUFFTAdjointOp()`.\n",
    "\n",
    "*Hint:* Please refer to the [above](#2DnonCart) example for non-Cartesian iterative SENSE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tirqcUF57S1x"
   },
   "outputs": [],
   "source": [
    "# prepare radial image\n",
    "nRead = np.amax(np.shape(img)[0:2])\n",
    "img_qfov = zpad(img, (nRead, nRead, np.shape(img)[2])).astype(np.complex64)\n",
    "img_qfov = img_qfov[:,:,np.newaxis,:]\n",
    "kspace_qfov = fft2c(img_qfov)\n",
    "smaps_rad_espirit = espirit.espirit(kspace_qfov, 8, 20, 0.05, 0)\n",
    "smaps_rad = smaps_rad_espirit[:,:,0,:,0]\n",
    "csm = np.transpose(smaps_rad, (2, 0, 1))\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "\n",
    "# affine motion\n",
    "p = [15, 10, 30, 0.1, 1, 1]\n",
    "u = get_transform(img_rad, p)\n",
    "u_id_rad = get_transform(img_rad, [0, 0, 0, 0, 1, 1])\n",
    "#u = get_deformation_field_from_affine(img_rad, get_affine_matrix(img_rad, p))\n",
    "#u_id_rad = get_deformation_field_from_affine(img_rad, get_affine_matrix(img_rad, [0, 0, 0, 0, 1, 1]))\n",
    "flow = np.stack([u_id_rad, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "# split kpos and dcf into motion states\n",
    "kpos1 = kpos[::2,:]\n",
    "dcf1 = dcf[::2,:]\n",
    "kpos2 = kpos[1::2,:]\n",
    "dcf2 = dcf[1::2,:]\n",
    "kpos_all = np.stack([kpos1, kpos2], -1)\n",
    "dcf_all = np.stack([dcf1, dcf2], -1)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, csm, dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, csm, dcf_all)\n",
    "  kspace_motion_rad = A.op(img_rad, kpos_all, csm, dcf_all, flow)\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  kspace_motion_rad = A(img_rad, kpos_all, csm, dcf_all, smm)\n",
    "  use_optox = False\n",
    "\n",
    "# no motion compensation\n",
    "flow_no = np.stack([u_id_rad, u_id_rad], -1)\n",
    "smm_no = vstack([smm0, smm0])\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  flowin = flow_no\n",
    "  img_no_motion_corr_zf_rad = AH.op(kspace_motion_rad, kpos_all, csm, dcf_all, flow_no)\n",
    "else:\n",
    "  flowin = smm_no\n",
    "  img_no_motion_corr_zf_rad = AH(kspace_motion_rad, kpos_all, csm, dcf_all, smm_no)\n",
    "\n",
    "img_no_motion_corr_recon_rad = iterativeSENSE(kspace_radial, csm, kpos_all, dcf=dcf_all, flow=flowin, fwdop=A, adjop=AH, weight_scale=0.1, use_optox=use_optox)\n",
    "\n",
    "# motion-compensated\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  flowin = flow\n",
    "  img_motion_corr_zf_rad = AH.op(kspace_motion_rad, kpos_all, csm, dcf_all, flow)\n",
    "else:\n",
    "  flowin = smm\n",
    "  img_motion_corr_zf_rad = AH(kspace_motion_rad, kpos_all, csm, dcf_all, smm)\n",
    "\n",
    "img_motion_corr_recon_rad = iterativeSENSE(kspace_radial, csm, kpos_all, dcf=dcf_all, flow=flowin, fwdop=A, adjop=AH, weight_scale=0.1, use_optox=use_optox)\n",
    "plot([img_rad, img_no_motion_corr_zf_rad, img_no_motion_corr_recon_rad, img_motion_corr_zf_rad, img_motion_corr_recon_rad], title='motion-free | zero-filled no motion-comp. | no motion-comp. recon | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iXjVh4Y6-NXS"
   },
   "source": [
    "**Task 64:** Perform a motion-compensated image reconstruction for an affine motion on accelerated non-Cartesian data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JzlbtOD7-fXv"
   },
   "outputs": [],
   "source": [
    "p = [15, 10, 30, 0.1, 1, 1]\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "u = get_transform(img_rad, p)\n",
    "u_id_rad = get_transform(img_rad, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id_rad, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "kpos, dcf = prepare_radial(acc=4, nRead=nRead)\n",
    "\n",
    "# split kpos and dcf into motion states\n",
    "kpos1 = kpos[::2,:]\n",
    "dcf1 = dcf[::2,:]\n",
    "kpos2 = kpos[1::2,:]\n",
    "dcf2 = dcf[1::2,:]\n",
    "kpos_all = np.stack([kpos1, kpos2], -1)\n",
    "dcf_all = np.stack([dcf1, dcf2], -1)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, csm, dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, csm, dcf_all)\n",
    "  kspace_motion_rad = A.op(img_rad, kpos_all, csm, dcf_all, flow)\n",
    "  img_motion_corr_zf_rad = AH.op(kspace_motion_rad, csm, kpos_all, dcf_all, flow)\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  kspace_motion_rad = A(img_rad, kpos_all, csm, dcf_all, smm)\n",
    "  img_motion_corr_zf_rad = AH(kspace_motion_rad, csm, kpos_all, dcf_all, smm)\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "\n",
    "img_motion_corr_recon_rad = iterativeSENSE(kspace_radial, csm, kpos_all, dcf=dcf_all, flow=flowin, fwdop=A, adjop=AH, weight_scale=0.1, use_optox=use_optox)\n",
    "plot([img_rad, img_motion_corr_zf_rad, img_motion_corr_recon_rad], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MIUI8V2U-lbQ"
   },
   "source": [
    "**Task 65:** Examine again the impact of using the coil sensitivity maps (with and without) for non-Cartesian imaging and compare to the reconstructions in **Task 64**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wjs5utnY-z1R"
   },
   "outputs": [],
   "source": [
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, np.ones_like(csm), dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, np.ones_like(csm), dcf_all)\n",
    "  flowin = flow\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  flowin = smm\n",
    "  use_optox = False\n",
    "\n",
    "img_motion_corr_recon_no_csm_rad = iterativeSENSE(kspace_radial, np.ones_like(csm), kpos_all, dcf=dcf_all, flow=flowin, fwdop=A, adjop=AH, weight_scale=0.1, use_optox=use_optox)\n",
    "\n",
    "plot([img_motion_corr_recon_rad, img_motion_corr_recon_no_csm_rad], title='with CSM | without CSM')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_5x0Cxp29EF5"
   },
   "source": [
    "**Task 66:** In a similar fashion, we want to examine the impact of a motion field error (i.e. inaccurate image registration) for the fully-sampled and undersampled non-Cartesian imaging. Perform a motion-compensated image reconstruction for an affine motion on regular undersampled Cartesian data. Additionally the k-space is affected by additive white Gaussian noise. Coil sensitivity maps are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z9aClerx9jLS"
   },
   "outputs": [],
   "source": [
    "p = [15, 10, 30, 0.1, 1, 1]\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "u = get_transform(img_rad, p)\n",
    "u_id_rad = get_transform(img_rad, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id_rad, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "kpos, dcf = prepare_radial(acc=1, nRead=nRead)\n",
    "\n",
    "# split kpos and dcf into motion states\n",
    "kpos1 = kpos[::2,:]\n",
    "dcf1 = dcf[::2,:]\n",
    "kpos2 = kpos[1::2,:]\n",
    "dcf2 = dcf[1::2,:]\n",
    "kpos_all = np.stack([kpos1, kpos2], -1)\n",
    "dcf_all = np.stack([dcf1, dcf2], -1)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, csm, dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, csm, dcf_all)\n",
    "  kspace_motion_rad = A.op(img_rad, kpos_all, csm, dcf_all, flow)\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  kspace_motion_rad = A(img_rad, kpos_all, csm, dcf_all, smm)\n",
    "  use_optox = False\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 2E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n",
    "\n",
    "# add error on motion field\n",
    "p_error = np.multiply(p,[0.9, 0.9, 0.9, 0.9, 1, 1])\n",
    "u_error = get_transform(img_rad, p_error)\n",
    "flow_error = np.stack([u_id_rad, u_error], -1)\n",
    "smm_error0 = get_sparse_motion_matrix(flow_error[:,:,:,0])\n",
    "smm_error1 = get_sparse_motion_matrix(flow_error[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm_error = vstack([smm_error0, smm_error1])\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  flowin = flow_error\n",
    "  img_motion_corr_error_rad_zf = AH.op(kspace_motion_rad, kpos_all, csm, dcf_all, flow_error)\n",
    "else:\n",
    "  flowin = smm_error\n",
    "  img_motion_corr_error_rad_zf = AH(kspace_motion_rad, kpos_all, csm, dcf_all, smm_error)\n",
    "\n",
    "img_motion_corr_error_recon_rad = iterativeSENSE(kspace_motion_rad, smaps, mask=kpos_all, flow=flowin, dcf=dcf_all, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_rad_zf, img_motion_corr_error_recon_rad], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AFWLVJIiJOi3"
   },
   "source": [
    "**Task 67:** Provide the same image prior for the motion-compensated image reconstruction on undersampled non-Cartesian data. The setting is otherwise similar to **Task 66**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PjFeQgZ6JgPh"
   },
   "outputs": [],
   "source": [
    "p = [15, 10, 30, 0.1, 1, 1]\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "u = get_transform(img_rad, p)\n",
    "u_id_rad = get_transform(img_rad, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id_rad, u], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "kpos, dcf = prepare_radial(acc=4, nRead=nRead)\n",
    "\n",
    "# split kpos and dcf into motion states\n",
    "kpos1 = kpos[::2,:]\n",
    "dcf1 = dcf[::2,:]\n",
    "kpos2 = kpos[1::2,:]\n",
    "dcf2 = dcf[1::2,:]\n",
    "kpos_all = np.stack([kpos1, kpos2], -1)\n",
    "dcf_all = np.stack([dcf1, dcf2], -1)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, csm, dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, csm, dcf_all)\n",
    "  kspace_motion_rad = A.op(img_rad, kpos_all, csm, dcf_all, flow)\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  kspace_motion_rad = A(img_rad, kpos_all, csm, dcf_all, smm)\n",
    "  use_optox = False\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 2E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n",
    "\n",
    "# add error on motion field\n",
    "p_error = np.multiply(p,[0.9, 0.9, 0.9, 0.9, 1, 1]) # 10% error\n",
    "u_error = get_transform(img_rad, p_error)\n",
    "flow_error = np.stack([u_id_rad, u_error], -1)\n",
    "smm_error0 = get_sparse_motion_matrix(flow_error[:,:,:,0])\n",
    "smm_error1 = get_sparse_motion_matrix(flow_error[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm_error = vstack([smm_error0, smm_error1])\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  flowin = flow_error\n",
    "  img_motion_corr_error_rad_zf = AH.op(kspace_motion_rad, kpos_all, csm, dcf_all, flow_error)\n",
    "else:\n",
    "  flowin = smm_error\n",
    "  img_motion_corr_error_rad_zf = AH(kspace_motion_rad, kpos_all, csm, dcf_all, smm_error)\n",
    "\n",
    "img_prior = gaussian_filter(img_cc, sigma=11)\n",
    "\n",
    "img_motion_corr_error_recon_rad = iterativeSENSE(kspace_motion_rad, smaps, mask=kpos_all, flow=flowin, dcf=dcf, noisy=img_prior, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_rad_zf, img_motion_corr_error_recon_rad], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZAP3wdIiLd0I"
   },
   "source": [
    "**Task 68:** Perform a similar motion-compensated image reconstruction as in **Task 67**, but now having three motion states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "84jQ4cqfTabj"
   },
   "outputs": [],
   "source": [
    "p = [15, 10, 30, 0.1, 1, 1]\n",
    "img_rad = rss(img_qfov[:,:,0,:])\n",
    "u = get_transform(img_rad, p)\n",
    "p2 = [5, -10, 20, 0, 1, 1]\n",
    "u2 = get_transform(img_rad, p2)\n",
    "u_id_rad = get_transform(img_rad, [0, 0, 0, 0, 1, 1])\n",
    "flow = np.stack([u_id_rad, u, u2], -1)\n",
    "smm0 = get_sparse_motion_matrix(flow[:,:,:,0])\n",
    "smm1 = get_sparse_motion_matrix(flow[:,:,:,1])\n",
    "from scipy.sparse import vstack\n",
    "smm = vstack([smm0, smm1])\n",
    "kpos, dcf = prepare_radial(acc=2, nRead=nRead)\n",
    "\n",
    "# split kpos and dcf into motion states\n",
    "kpos1 = kpos[::3,:]\n",
    "dcf1 = dcf[::3,:]\n",
    "kpos2 = kpos[1::3,:]\n",
    "dcf2 = dcf[1::3,:]\n",
    "kpos3 = kpos[2::3,:]\n",
    "dcf3 = dcf[2::3,:]\n",
    "kpos_all = np.stack([kpos1, kpos2, kpos3], -1)\n",
    "dcf_all = np.stack([dcf1, dcf2, dcf3], -1)\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  A = BatchelorGPUNUFFTFwd(nRead, kpos_all, csm, dcf_all)\n",
    "  AH = BatchelorGPUNUFFTAdj(nRead, kpos_all, csm, dcf_all)\n",
    "  kspace_motion_rad = A.op(img_rad, kpos_all, csm, dcf_all, flow)\n",
    "  use_optox = True\n",
    "else:\n",
    "  A = BatchGPUNUFFTForwardOp\n",
    "  AH = BatchGPUNUFFTAdjointOp\n",
    "  kspace_motion_rad = A(img_rad, kpos_all, csm, dcf_all, smm)\n",
    "  use_optox = False\n",
    "\n",
    "# add noise to k-space\n",
    "max_s = np.amax(np.abs(kspace_motion))\n",
    "noisestd = 2E-3\n",
    "kspace_motion += max_s * noisestd * (np.random.randn(*np.shape(kspace_motion)) + 1j * np.random.randn(*np.shape(kspace_motion)))\n",
    "\n",
    "# add error on motion field\n",
    "p_error = np.multiply(p,[0.9, 0.9, 0.9, 0.9, 1, 1]) # 10% error\n",
    "u_error = get_transform(img_rad, p_error)\n",
    "p_error = np.multiply(p2,[0.9, 0.9, 0.9, 0.9, 1, 1])\n",
    "u2_error = get_transform(img_rad, p2_error)\n",
    "flow_error = np.stack([u_id_rad, u_error, u2_error], -1)\n",
    "smm_error0 = get_sparse_motion_matrix(flow_error[:,:,:,0])\n",
    "smm_error1 = get_sparse_motion_matrix(flow_error[:,:,:,1])\n",
    "smm_error2 = get_sparse_motion_matrix(flow_error[:,:,:,2])\n",
    "from scipy.sparse import vstack\n",
    "smm_error = vstack([smm_error0, smm_error1, smm_error2])\n",
    "\n",
    "if os.getenv(\"USE_OPTOX\") and os.environ[\"USE_OPTOX\"] == \"true\":  # Tensorflow\n",
    "  flowin = flow_error\n",
    "  img_motion_corr_error_rad_zf = AH.op(kspace_motion_rad, kpos_all, csm, dcf_all, flow_error)\n",
    "else:\n",
    "  flowin = smm_error\n",
    "  img_motion_corr_error_rad_zf = AH(kspace_motion_rad, kpos_all, csm, dcf_all, smm_error)\n",
    "\n",
    "img_prior = gaussian_filter(img_cc, sigma=11)\n",
    "\n",
    "img_motion_corr_error_recon_rad = iterativeSENSE(kspace_motion_rad, smaps, mask=kpos_all, flow=flowin, dcf=dcf_all, noisy=img_prior, fwdop=A, adjop=AH, use_optox=use_optox)\n",
    "\n",
    "plot([img_cc, img_motion_corr_error_rad_zf, img_motion_corr_error_recon_rad], title='motion-free | zero-filled motion-comp. | motion-compensated recon')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O_2MMxj8WDMJ"
   },
   "source": [
    "**Task 69:** Repeat the above **Task 68** with changing sampling pattern per motion state (acceleration factor and/or sampled points), motion errors on motion fields, number of motion states or changing additive noise."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H5NFIndbWKPB"
   },
   "source": [
    "**Task 70:** Repeat the **Tasks 53** to **69** using a different dataset, e.g. cardiac CINE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qu_-QeT6WYFV"
   },
   "source": [
    "### Further exercises\n",
    "To be performed on numerical and/or in-vivo data:\n",
    "1. Evaluate the performance of both reconstructions in the presence of noise, different degrees of undersampling, model errors (coils and/or motion), etc.\n",
    "\n",
    "2. Implement a motion corrected reconstruction using nearest neighbour or cubic spline interpolations (current implementation uses linear). Evaluate the performance of the recon, particularly in the presence of large scalings. <br/>\n",
    "Notice how the motion matrix reduces to a permutation matrix in the case of nearest neighbour interpolation.\n",
    "\n",
    "3. Inspect the convergence of the motion corrected reconstruction without normalizing the interpolants when the tranpose motion is applied (i.e. set the \"motion_norm\" variable inside \"mtimes.m\" to identity). Experiment with different interpolation strategies, rotations, scalings and shearings.\n",
    "\n",
    "4. Combine both itSENSE recon and the motion correction recon with compressed sensing and evalute their performances as in \"1.\".\n",
    "\n",
    "5. Evaluate the performance of all the previous exercises in terms of SNR, SSIM, RMSE, etc\n",
    "\n",
    "6. Simulate a case with through-plane motion and observe the performance of a (in-plane) motion corrected reconstruction.\n",
    "\n",
    "7. Evaluate the g-factor of the motion corrected reconstruction for different types of motion and sampling.\n",
    "\n",
    "8. Motion correction can be combined with soft-weighting to alleviate the inherent noise amplification. With soft-weighting, the sampling matrices are real-valued [0,1] and each k-space data can partially belong to multiple motion states. Implement such a soft-weighted motion corrected reconstruction and evaluate its' performance.\n",
    "\n",
    "9. The current implementation assumes the object is acquired during a steady state. Contrast resolved reconstructions can be easily solved using low rank subspace forward models. Combine a low rank constrained reconstruction with motion correction to enable contrast resolved motion correction.\n",
    "\n",
    "10. Plenty of other parameters can go into the forward model. For example, B0 and B1 can vary with motion. Implement a variant of the motion corrected reconstruction that accounts for B0 and B1 variations due to motion.\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "IHCV5_OYqKuP",
    "K3YztvzE6zVD",
    "LcNRw4_ebMOw"
   ],
   "machine_shape": "hm",
   "name": "HandsOn_ISMRM_MoCo_Workshop_solution.ipynb",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}