demos/notebooks/demo_motion_correction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Motion Correction demo\n",
    "\n",
    "This notebook demonstrates the various routines for motion correction in the CaImAn package. It demonstrates the usage of rigid and piecewise rigid motion correction on a two-photon calcium imaging dataset using the NoRMCorre algorithm [[1]](#normcorre), as well as several measures for quality assessment. This notebook should be interpreted more as a tutorial of the various methods. In practice, you can use either rigid or piecewise rigid motion correction depending on the motion of the dataset.\n",
    "\n",
    "The dataset used in this notebook is provided by Sue Ann Koay and David Tank, Princeton University. This is a two photon calcium imaging dataset. For motion correction of one photon microendoscopic data the procedure is similar, with the difference, that the shifts are inferred on high pass spatially filtered version of the data. For more information check the demos for one photon data in the CaImAn package.\n",
    "\n",
    "More information about the NoRMCorre algorithm can be found in the following paper:\n",
    "\n",
    "<a name=\"normcorre\"></a>[1] Pnevmatikakis, E.A., and Giovannucci A. (2017). NoRMCorre: An online algorithm for piecewise rigid motion correction of calcium imaging data. Journal of Neuroscience Methods, 291:83-92 [[paper]](https://doi.org/10.1016/j.jneumeth.2017.07.031)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from IPython import get_ipython\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os.path\n",
    "import logging\n",
    "\n",
    "try:\n",
    "    cv2.setNumThreads(0)\n",
    "except:\n",
    "    pass\n",
    "\n",
    "try:\n",
    "    if __IPYTHON__:\n",
    "        get_ipython().run_line_magic('load_ext', 'autoreload')\n",
    "        get_ipython().run_line_magic('autoreload', '2')\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "logging.basicConfig(format=\n",
    "                          \"%(relativeCreated)12d [%(filename)s:%(funcName)20s():%(lineno)s] [%(process)d] %(message)s\",\n",
    "                    # filename=\"/tmp/caiman.log\",\n",
    "                    level=logging.DEBUG)\n",
    "\n",
    "import caiman as cm\n",
    "from caiman.motion_correction import MotionCorrect, tile_and_correct, motion_correction_piecewise\n",
    "from caiman.utils.utils import download_demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First download the file and load it in memory to view it. Note that it is not necessary to load the file in memory in order to perform motion correction. Here we load it to inspect it. Viewing the file occurs with OpenCV and will a open a new window. **To exit click on the video and press q.**\n",
    "\n",
    "The `download_demo` function will download the specific file for you and return the complete path to the file which will be stored in your `caiman_data` directory. If you adapt this demo for your data make sure to pass the complete path to your file(s). Remember to pass the `fname` variable as a list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fnames = 'Sue_2x_3000_40_-46.tif'\n",
    "fnames = [download_demo(fnames)]     # the file will be downloaded if it doesn't already exist\n",
    "m_orig = cm.load_movie_chain(fnames)\n",
    "downsample_ratio = .2  # motion can be perceived better when downsampling in time\n",
    "m_orig.resize(1, 1, downsample_ratio).play(q_max=99.5, fr=30, magnification=2)   # play movie (press q to exit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now set some parameters that are used for motion correction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_shifts = (6, 6)  # maximum allowed rigid shift in pixels (view the movie to get a sense of motion)\n",
    "strides =  (48, 48)  # create a new patch every x pixels for pw-rigid correction\n",
    "overlaps = (24, 24)  # overlap between patches (size of patch strides+overlaps)\n",
    "max_deviation_rigid = 3   # maximum deviation allowed for patch with respect to rigid shifts\n",
    "pw_rigid = False  # flag for performing rigid or piecewise rigid motion correction\n",
    "shifts_opencv = True  # flag for correcting motion using bicubic interpolation (otherwise FFT interpolation is used)\n",
    "border_nan = 'copy'  # replicate values along the boundary (if True, fill in with NaN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that here the data presented here has been downsampled in space by a factor of 2 to reduce the file size. As a result the spatial resolution is coarser here (around 2 microns per pixel). If we were operating at the original resolution, several of the parameters above, e.g., ```max_shifts, strides, overlaps, max_deviation_rigid```, could have been larger by a factor of 2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Motion correction is performed in parallel on chunks taken across times.\n",
    "\n",
    "We first need to start a cluster. The default backend mode for parallel processing is through the multiprocessing package. To make sure that this package is viewable from everywhere before starting the notebook these commands need to be executed from the terminal (in Linux and Windows):\n",
    "```bash\n",
    "   export MKL_NUM_THREADS=1\n",
    "   export OPENBLAS_NUM_THREADS=1 \n",
    "   ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% start the cluster (if a cluster already exists terminate it)\n",
    "if 'dview' in locals():\n",
    "    cm.stop_server(dview=dview)\n",
    "c, dview, n_processes = cm.cluster.setup_cluster(\n",
    "    backend='multiprocessing', n_processes=None, single_thread=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first need to create a motion correction object with the parameters specified above. We pass directly its input arguments in the constructor below. Alternatively, we can use the `params` object and construct it by passing the arguments of `params.motion`. See the notebook `demo_pipeline.ipynb` for an example of this usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a motion correction object\n",
    "mc = MotionCorrect(fnames, dview=dview, max_shifts=max_shifts,\n",
    "                  strides=strides, overlaps=overlaps,\n",
    "                  max_deviation_rigid=max_deviation_rigid, \n",
    "                  shifts_opencv=shifts_opencv, nonneg_movie=True,\n",
    "                  border_nan=border_nan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1> Rigid motion correction</h1>\n",
    "<p> The original file exhibits a lot of motion. In order to correct for it we are first trying a simple rigid motion correction algorithm. This has already been selected by setting the parameter `pw_rigid=False` during the construction of the `MotionCorrect` object. The algorithm first creates a template by averaging frames from the video. It then tries to match each frame to this template. In addition the template will get updated during the matching process, resulting in a single precise template that is used for subpixel registration.  </p>\n",
    "<img src=\"images/rigidcorrection.png\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# correct for rigid motion correction and save the file (in memory mapped form)\n",
    "mc.motion_correct(save_movie=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The motion corrected file is automatically save as memory mapped file in the location given by `mc.mmap_file`. The rigid shifts are also save in `mc.shifts_rig`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load motion corrected movie\n",
    "m_rig = cm.load(mc.mmap_file)\n",
    "bord_px_rig = np.ceil(np.max(mc.shifts_rig)).astype(int)\n",
    "#%% visualize templates\n",
    "plt.figure(figsize = (20,10))\n",
    "plt.imshow(mc.total_template_rig, cmap = 'gray');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#%% inspect movie\n",
    "m_rig.resize(1, 1, downsample_ratio).play(\n",
    "    q_max=99.5, fr=30, magnification=2, bord_px = 0*bord_px_rig) # press q to exit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plot the shifts computed by rigid registration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% plot rigid shifts\n",
    "plt.close()\n",
    "plt.figure(figsize = (20,10))\n",
    "plt.plot(mc.shifts_rig)\n",
    "plt.legend(['x shifts','y shifts'])\n",
    "plt.xlabel('frames')\n",
    "plt.ylabel('pixels');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Piecewise rigid registration\n",
    "While rigid registration corrected for a lot of the movement, there is still non-uniform motion present in the registered file. To correct for that we can use piece-wise rigid registration **directly in the original file** by setting `mc.pw_rigid=True`. As before the registered file is saved in a memory mapped format in the location given by `mc.mmap_file`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "#%% motion correct piecewise rigid\n",
    "mc.pw_rigid = True  # turn the flag to True for pw-rigid motion correction\n",
    "mc.template = mc.mmap_file  # use the template obtained before to save in computation (optional)\n",
    "\n",
    "mc.motion_correct(save_movie=True, template=mc.total_template_rig)\n",
    "m_els = cm.load(mc.fname_tot_els)\n",
    "m_els.resize(1, 1, downsample_ratio).play(\n",
    "    q_max=99.5, fr=30, magnification=2,bord_px = bord_px_rig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now concatenate all the movies (raw, rigid, and pw-rigid) for inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm.concatenate([m_orig.resize(1, 1, downsample_ratio) - mc.min_mov*mc.nonneg_movie,\n",
    "                m_rig.resize(1, 1, downsample_ratio), m_els.resize(\n",
    "            1, 1, downsample_ratio)], axis=2).play(fr=60, q_max=99.5, magnification=2, bord_px=bord_px_rig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the movie we can see that pw-rigid registration corrected for the non uniform motion of the data. This was done by estimating different displacement vectors for the different patches in the FOV. This can be visualized by plotting all the computed shifts were a dispersion in the shifts in the y direction is apparent. In this case, the shifts along the two axes are stored in `mc.x_shifts_els` and `mc.y_shifts_els`, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#%% visualize elastic shifts\n",
    "plt.close()\n",
    "plt.figure(figsize = (20,10))\n",
    "plt.subplot(2, 1, 1)\n",
    "plt.plot(mc.x_shifts_els)\n",
    "plt.ylabel('x shifts (pixels)')\n",
    "plt.subplot(2, 1, 2)\n",
    "plt.plot(mc.y_shifts_els)\n",
    "plt.ylabel('y_shifts (pixels)')\n",
    "plt.xlabel('frames')\n",
    "#%% compute borders to exclude\n",
    "bord_px_els = np.ceil(np.maximum(np.max(np.abs(mc.x_shifts_els)),\n",
    "                                 np.max(np.abs(mc.y_shifts_els)))).astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The improvement in performance can also be seen by a more crisp summary statistic image. Below we plot the correlation images for the three datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (20,10))\n",
    "plt.subplot(1,3,1); plt.imshow(m_orig.local_correlations(eight_neighbours=True, swap_dim=False))\n",
    "plt.subplot(1,3,2); plt.imshow(m_rig.local_correlations(eight_neighbours=True, swap_dim=False))\n",
    "plt.subplot(1,3,3); plt.imshow(m_els.local_correlations(eight_neighbours=True, swap_dim=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm.stop_server(dview=dview) # stop the server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quality assessment \n",
    "\n",
    "Apart from inspection, the performance of the registration methods can be quantified using several measures. Below we compute measures such as correlation of each frame with mean, crispness of summary image, and residual optical flow for all three cases. For more info see [[1]](#normcorre). Note that computation of the residual optical flow can be computationally intensive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "#% compute metrics for the results (TAKES TIME!!)\n",
    "final_size = np.subtract(mc.total_template_els.shape, 2 * bord_px_els) # remove pixels in the boundaries\n",
    "winsize = 100\n",
    "swap_dim = False\n",
    "resize_fact_flow = .2    # downsample for computing ROF\n",
    "\n",
    "tmpl_orig, correlations_orig, flows_orig, norms_orig, crispness_orig = cm.motion_correction.compute_metrics_motion_correction(\n",
    "    fnames[0], final_size[0], final_size[1], swap_dim, winsize=winsize, play_flow=False, resize_fact_flow=resize_fact_flow)\n",
    "\n",
    "tmpl_rig, correlations_rig, flows_rig, norms_rig, crispness_rig = cm.motion_correction.compute_metrics_motion_correction(\n",
    "    mc.fname_tot_rig[0], final_size[0], final_size[1],\n",
    "    swap_dim, winsize=winsize, play_flow=False, resize_fact_flow=resize_fact_flow)\n",
    "\n",
    "tmpl_els, correlations_els, flows_els, norms_els, crispness_els = cm.motion_correction.compute_metrics_motion_correction(\n",
    "    mc.fname_tot_els[0], final_size[0], final_size[1],\n",
    "    swap_dim, winsize=winsize, play_flow=False, resize_fact_flow=resize_fact_flow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot correlation with mean frame for each dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize = (20,10))\n",
    "plt.subplot(211); plt.plot(correlations_orig); plt.plot(correlations_rig); plt.plot(correlations_els)\n",
    "plt.legend(['Original','Rigid','PW-Rigid'])\n",
    "plt.subplot(223); plt.scatter(correlations_orig, correlations_rig); plt.xlabel('Original'); \n",
    "plt.ylabel('Rigid'); plt.plot([0.3,0.7],[0.3,0.7],'r--')\n",
    "axes = plt.gca(); axes.set_xlim([0.3,0.7]); axes.set_ylim([0.3,0.7]); plt.axis('square');\n",
    "plt.subplot(224); plt.scatter(correlations_rig, correlations_els); plt.xlabel('Rigid'); \n",
    "plt.ylabel('PW-Rigid'); plt.plot([0.3,0.7],[0.3,0.7],'r--')\n",
    "axes = plt.gca(); axes.set_xlim([0.3,0.7]); axes.set_ylim([0.3,0.7]); plt.axis('square');\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print crispness values\n",
    "print('Crispness original: ' + str(int(crispness_orig)))\n",
    "print('Crispness rigid: ' + str(int(crispness_rig)))\n",
    "print('Crispness elastic: ' + str(int(crispness_els)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% plot the results of Residual Optical Flow\n",
    "fls = [cm.paths.fname_derived_presuffix(mc.fname_tot_els[0], 'metrics', swapsuffix='npz'),\n",
    "       cm.paths.fname_derived_presuffix(mc.fname_tot_rig[0], 'metrics', swapsuffix='npz'),\n",
    "       cm.paths.fname_derived_presuffix(mc.fname[0],         'metrics', swapsuffix='npz'),\n",
    "      ]\n",
    "\n",
    "plt.figure(figsize = (20,10))\n",
    "for cnt, fl, metr in zip(range(len(fls)), fls, ['pw_rigid','rigid','raw']):\n",
    "    with np.load(fl) as ld:\n",
    "        print(ld.keys())\n",
    "        print(fl)\n",
    "        print(str(np.mean(ld['norms'])) + '+/-' + str(np.std(ld['norms'])) +\n",
    "              ' ; ' + str(ld['smoothness']) + ' ; ' + str(ld['smoothness_corr']))\n",
    "        \n",
    "        plt.subplot(len(fls), 3, 1 + 3 * cnt)\n",
    "        plt.ylabel(metr)\n",
    "        print(f\"Loading data with base {fl[:-12]}\")\n",
    "        try:\n",
    "            mean_img = np.mean(\n",
    "            cm.load(fl[:-12] + '.mmap'), 0)[12:-12, 12:-12]\n",
    "        except:\n",
    "            try:\n",
    "                mean_img = np.mean(\n",
    "                    cm.load(fl[:-12] + '.tif'), 0)[12:-12, 12:-12]\n",
    "            except:\n",
    "                mean_img = np.mean(\n",
    "                    cm.load(fl[:-12] + 'hdf5'), 0)[12:-12, 12:-12]\n",
    "                    \n",
    "        lq, hq = np.nanpercentile(mean_img, [.5, 99.5])\n",
    "        plt.imshow(mean_img, vmin=lq, vmax=hq)\n",
    "        plt.title('Mean')\n",
    "        plt.subplot(len(fls), 3, 3 * cnt + 2)\n",
    "        plt.imshow(ld['img_corr'], vmin=0, vmax=.35)\n",
    "        plt.title('Corr image')\n",
    "        plt.subplot(len(fls), 3, 3 * cnt + 3)\n",
    "        flows = ld['flows']\n",
    "        plt.imshow(np.mean(\n",
    "        np.sqrt(flows[:, :, :, 0]**2 + flows[:, :, :, 1]**2), 0), vmin=0, vmax=0.3)\n",
    "        plt.colorbar()\n",
    "        plt.title('Mean optical flow');  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}