notebooks/macaque_template_warps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Registration across Macaque MRI templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Authors:** Nikoloz Sirmpilatze (nsirmpilatze@dpz.eu) & Chris Klink (c.klink@nin.knaw.nl)                \n",
    "\n",
    "**Last updated:** May 11, 2020     \n",
    "\n",
    "**Requirements:**    \n",
    "* _python_ >= 3.7\n",
    "* _nipype_ >= 1.2.0\n",
    "* _nilearn_ >= 0.5.2\n",
    "    * Used only for visualisation\n",
    "* _nibabel_ >= 2.3.3\n",
    "* _joblib_ >= 0.14.1\n",
    "    * Used only for parallel processing, not necessary when registrations are done serially\n",
    "* _ANTs_ >= 2.4.0\n",
    "    * _antsRegistration_,  _antsApplyTransforms_ and _antsAverageImages_ need to be in your path as executables    \n",
    "\n",
    "**Citation**: Sirmpilatze, Nikoloz and Klink, P. Christiaan (2020). RheMAP: Non-linear warps between common rhesus macaque brain templates (Version 1.2)[Data set]. Zenodo. https://doi.org/10.5281/zenodo.3786357     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following templates are used: \n",
    "1. [NMT v1.2](https://afni.nimh.nih.gov/pub/dist/doc/htmldoc/nonhuman/macaque_tempatl/template_nmtv1.html)\n",
    "2. [NMT v1.3](https://afni.nimh.nih.gov/pub/dist/doc/htmldoc/nonhuman/macaque_tempatl/template_nmtv1.html)\n",
    "3. [NMT v2.0](https://afni.nimh.nih.gov/pub/dist/doc/htmldoc/nonhuman/macaque_tempatl/template_nmtv2.html)\n",
    "4. [D99](https://afni.nimh.nih.gov/Macaque)\n",
    "5. [INIA19](https://www.nitrc.org/projects/inia19/https://www.nitrc.org/projects/inia19/)\n",
    "6. [MNI macaque](http://www.bic.mni.mcgill.ca/ServicesAtlases/Macaque)\n",
    "7. [Yerkes19](https://github.com/Washington-University/NHPPipelines)\n",
    "8. [ONPRC18](https://www.nitrc.org/projects/onprc18_atlas)     \n",
    "9. [F99](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/XTRACT)\n",
    "10. [MEBRAINS](https://doi.org/10.25493/5454-ZEA)       \n",
    "\n",
    "Within this notebook, they are abbreviated as *NMTv12*, *NMTv13*, *NMTv20_sym*, *NMTv20_asym*, *NMTv20_05mm_sym_brain*, and *NMTv20_05mm_asym*, *D99*, *INIA*, *MNI*, *YRK*, *ONPRC18*, *F99*, and *MEBRAINS*.     \n",
    "\n",
    "**NB!** We do not provide copies of the actual templates (licenses often forbids redistribution), but instead suggest you follow the links above and get them at the source. We do offer the warp files and warped templates that will be produced by this workflow. They can be downloaded from [Zenodo](https://doi.org/10.5281/zenodo.3786357) or [G-NODE GIN](https://gin.g-node.org/ChrisKlink/RheMAP). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to emulate this code, you can consider setting up your templates in the following folder structure:    \n",
    "\n",
    "|--- RheMAP   \n",
    "&emsp; |--- notebooks     \n",
    "&emsp; |--- templates     \n",
    "&emsp; &emsp; |--- D99     \n",
    "&emsp; &emsp; &emsp; |--- D99_atlas_1.2a.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_atlas_1.2a_in_MNI.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_template.nii.gz      \n",
    "&emsp; &emsp; |--- F99     \n",
    "&emsp; &emsp; &emsp; |--- struct.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- struct_brain.nii.gz \n",
    "&emsp; &emsp; |--- INIA   \n",
    "&emsp; &emsp; &emsp; |--- inia19-t1-brain_truncated.nii.gz      \n",
    "&emsp; &emsp; |--- MNI     \n",
    "&emsp; &emsp; &emsp; |--- macaque_25_model-MNI_brain.nii.gz      \n",
    "&emsp; &emsp; |--- MEBRAINS     \n",
    "&emsp; &emsp; &emsp; |--- MEBRAINS_T1_masked.nii.gz      \n",
    "&emsp; &emsp; |--- NMT     \n",
    "&emsp; &emsp; &emsp; |--- NMT_v1.2     \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_SS.nii.gz      \n",
    "&emsp; &emsp; &emsp; |--- NMT_v1.3     \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_SS.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- NMT_v2.0     \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_v2.0_asym_05mm_SS.nii.gz  \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_v2.0_asym_SS.nii.gz  \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_v2.0_sym_05mm_SS.nii.gz  \n",
    "&emsp; &emsp; &emsp; &emsp; |--- NMT_v2.0_sym_SS.nii.gz  \n",
    "&emsp; &emsp; |--- ONPRC18     \n",
    "&emsp; &emsp; &emsp; |--- ONPRC18_T1W.nii.gz  \n",
    "&emsp; &emsp; |--- YRK      \n",
    "&emsp; &emsp; &emsp; |--- MacaqueYerkes19_T1w_0.5mm_brain.nii.gz       \n",
    "\n",
    "After downloading the warp files and warped templates from [Zenodo](https://zenodo.org/record/3776856#.XqqfI3UzZjE), we suggest you include them like this:      \n",
    "\n",
    "|--- RheMAP   \n",
    "&emsp; |--- notebooks     \n",
    "&emsp; |--- templates     \n",
    "&emsp; |--- warps       \n",
    "&emsp; &emsp; |--- final     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_INIA_CompositeWarp.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_MNI_CompositeWarp.nii.gz      \n",
    "&emsp; &emsp; &emsp; |--- etc   \n",
    "&emsp; &emsp; |--- linear     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_INIA_affine_0GenericAffine.mat     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_MNI_affine_0GenericAffine.mat     \n",
    "&emsp; &emsp; &emsp; |--- etc   \n",
    "&emsp; &emsp; |--- nonlinear     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_INIA_1InverseWarp.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_INIA_1Warp.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_MNI_1InverseWarp.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_to_MNI_1Warp.nii.gz  \n",
    "&emsp; &emsp; &emsp; |--- etc   \n",
    "&emsp; |--- warped_templates       \n",
    "&emsp; &emsp; |--- final     \n",
    "&emsp; &emsp; &emsp; |--- D99_in_INIA_composite.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_in_MNI_composite.nii.gz       \n",
    "&emsp; &emsp; &emsp; |--- etc   \n",
    "&emsp; &emsp; |--- linear     \n",
    "&emsp; &emsp; &emsp; |--- D99_in_INIA_linear.nii.gz    \n",
    "&emsp; &emsp; &emsp; |--- D99_in_MNI_linear.nii.gz          \n",
    "&emsp; &emsp; &emsp; |--- etc   \n",
    "&emsp; &emsp; |--- nonlinear     \n",
    "&emsp; &emsp; &emsp; |--- D99_in_INIA_linear+SyN.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- D99_in_MNI_linear+SyN.nii.gz     \n",
    "&emsp; &emsp; &emsp; |--- etc  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on the warp files and warped templates you could of course reconstruct the original templates with something like the following.     \n",
    "\n",
    "On the command line:     \n",
    "```bash\n",
    "antsApplyTransforms -i <TEMPLATE1_in_TEMPLATE2_composite.nii.gz> \\\n",
    "                    -r <TEMPLATE2_in_TEMPLATE1.nii.gz> \\        \n",
    "                    -o <RECONSTRUCTED_ORIGINAL_TEMPLATE1.nii.gz> \\\n",
    "                    -t [<TEMPLATE1_to_TEMPLATE2_CompositeWarp>,1] \\\n",
    "                    -n Linear \\\n",
    "                    -d 3\n",
    "```    \n",
    "\n",
    "In NiPype:     \n",
    "```python\n",
    "import nipype.interfaces.ants as ants    \n",
    "ants.ApplyTransforms(\n",
    "            input_image=<TEMPLATE1_in_TEMPLATE2_composite.nii.gz>,\n",
    "            reference_image=<TEMPLATE2_in_TEMPLATE1.nii.gz>,        \n",
    "            output_image=<RECONSTRUCTED_ORIGINAL_TEMPLATE1.nii.gz>,\n",
    "            transforms=<TEMPLATE1_to_TEMPLATE2_CompositeWarp>,\n",
    "            invert_transform_flags=True,\n",
    "            interpolation='Linear',\n",
    "            dimension=3)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The end goal is to generate warps between each unique pair of the templates (forward and backward).\n",
    "* forwards (A to B, e.g. *NMTv1.2_to_D99*)\n",
    "* backwards (B to A, e.g. *D99_to_NMTv1.2*)\n",
    "\n",
    "![warps](RegisterTemplates.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Preparations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0a. Import required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import time\n",
    "import shutil as sh\n",
    "import nibabel as nb\n",
    "\n",
    "from itertools import combinations\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import nipype.interfaces.fsl as fsl # nipype interface for FSL\n",
    "import nipype.interfaces.ants as ants # nipype interface for ANTs\n",
    "from nilearn import plotting # Plotting function from nilearn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0b. Define relative paths to template files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The skull-stripped isotropic volumetric images are used for registration.\n",
    "* For **NMT**, **D99**, and **YRK** the provided skull-stripped brains were used\n",
    "* In **INIA** the brain stem extends further down the spinal cord compared to the other temlplates, so the braisn stem was truncated at a level similar to the others\n",
    "* No skull-stripped image was provided with **MNI**, so the brain was segmented semi-manually using ITK-SNAP\n",
    "* For **YRK** we used the version provided together with [NHPPipelines](https://github.com/Washington-University/NHPPipelines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================\n",
    "# NB! THIS COULD BE DIFFERENT FOR EVERY USER\n",
    "# ============================================\n",
    "BASE_path     = os.path.dirname(os.getcwd()) + '/'   # repo base folder\n",
    "print(BASE_path)\n",
    "\n",
    "# these follow the directory structure as outline above\n",
    "TEMPLATE_path = BASE_path + 'templates/'                   # templates base folder\n",
    "NMTv12_path   = TEMPLATE_path + 'NMT/NMT_v1.2/'\n",
    "NMTv13_path   = TEMPLATE_path + 'NMT/NMT_v1.3/'\n",
    "NMTv20_path   = TEMPLATE_path + 'NMT/NMT_v2.0/'\n",
    "D99_path      = TEMPLATE_path + 'D99/'\n",
    "F99_path      = TEMPLATE_path + 'F99/'\n",
    "INIA_path     = TEMPLATE_path + 'INIA/'\n",
    "MNI_path      = TEMPLATE_path + 'MNI/'\n",
    "YRK_path      = TEMPLATE_path + 'YRK/'\n",
    "ONPRC18_path  = TEMPLATE_path + 'ONPRC18/'\n",
    "MEBRAINS_path  = TEMPLATE_path + 'MEBRAINS/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NMTv12_brain = NMTv12_path + 'NMT_SS.nii.gz'\n",
    "NMTv13_brain = NMTv13_path + 'NMT_SS.nii.gz'\n",
    "NMTv20_sym_brain = NMTv20_path + 'NMT_v2.0_sym_SS.nii.gz'\n",
    "NMTv20_05mm_sym_brain = NMTv20_path + 'NMT_v2.0_sym_05mm_SS.nii.gz'\n",
    "NMTv20_asym_brain = NMTv20_path + 'NMT_v2.0_asym_SS.nii.gz'\n",
    "NMTv20_05mm_asym_brain = NMTv20_path + 'NMT_v2.0_asym_05mm_SS.nii.gz'\n",
    "D99_brain    = D99_path + 'D99_template.nii.gz'\n",
    "F99_brain    = F99_path + 'struct_brain.nii.gz'\n",
    "INIA_brain   = INIA_path + 'inia19-t1-brain_truncated.nii.gz'\n",
    "MNI_brain    = MNI_path + 'macaque_25_model-MNI_brain.nii.gz'\n",
    "YRK_brain    = YRK_path + 'MacaqueYerkes19_T1w_0.5mm_brain.nii.gz'\n",
    "ONPRC18_brain = ONPRC18_path + 'ONPRC18_T1W.nii.gz'\n",
    "MEBRAINS_brain = MEBRAINS_path + 'MEBRAINS_T1W.nii.gz'\n",
    "\n",
    "temp_names   = ['NMTv1.2', 'NMTv1.3',\n",
    "                'NMTv2.0-sym','NMTv2.0-0.5mm-sym',\n",
    "                'NMTv2.0-asym','NMTv2.0-0.5mm-asym',\n",
    "                'D99', 'F99', 'INIA', 'MNI',\n",
    "                'YRK', 'ONPRC18','MEBRAINS']\n",
    "temp_brains  = [NMTv12_brain, NMTv13_brain, \n",
    "                NMTv20_sym_brain, NMTv20_05mm_sym_brain,\n",
    "                NMTv20_asym_brain,NMTv20_05mm_asym_brain,\n",
    "                D99_brain, F99_brain, INIA_brain, MNI_brain,\n",
    "                YRK_brain, ONPRC18_brain, MEBRAINS_brain]\n",
    "\n",
    "# strength of contrast in the plot_anat function (\"dim\")\n",
    "contrasts = [0]*len(temp_brains) # create a list of zeros because this should exist\n",
    "contrasts = [-1, 0, 0, 0, 0, 0, 0, -1, -1, -0.5, 0, 0, 0] # specify to improve display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following step you can create a subset of the templates to include in creating the warps. This can be useful when you add a new template and want to warp it to templates that have pre-existing warps (to each other). This way you can generate the new unique temlates without redoing the prexisting ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "include_pairs_with = temp_names # all included\n",
    "include_pairs_with = ['MEBRAINS']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The headers of NMT > v1.3 has both qform_code and sform_code set to 5. This is a relatively new definition ([see NITRC](https://www.nitrc.org/forum/forum.php?thread_id=10029&forum_id=1942)). Older version of nibabel (< 2.4.0) only recognized values 0-4 and reset the 5 to 0 which causes weird plotting artefacts (registrations will still be ok). As a workaround, you can use the code below to change the header value from 5 to something like 2 ('aligned'). A better option would of course be to update nibabel.\n",
    "```\n",
    "if NMTv13_brain in locals():\n",
    "    sh.copyfile(NMTv13_brain,NMTv13_brain + '.bak'); # back up the original file\n",
    "    img = nb.load(NMTv13_brain)\n",
    "    hdr = img.header\n",
    "    hdr.set_qform(hdr.get_qform(),code=2)\n",
    "    hdr.set_sform(hdr.get_sform(),code=2)\n",
    "    nb.save(img, NMTv13_brain) \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0c. Plot template brains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for name, brain,con in zip(temp_names, temp_brains, contrasts):\n",
    "    if name in include_pairs_with:\n",
    "        display = plotting.plot_anat(brain, display_mode='ortho', title=name, \n",
    "                  draw_cross=False, annotate=True, dim=con);\n",
    "        plt.savefig(name + '_template.png'); # save png\n",
    "        plt.draw()\n",
    "        #display.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0d. Parallel or serial    \n",
    "Registration can be a time consuming process. If you have the resources, you may opt to parallelize some of the steps. You can configure that here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "do_parallel = True  # Set to True for parallel processing\n",
    "par_numjobs = 5     # How many parallel processes? This is only used when do_parallel is True\n",
    "if do_parallel is True:\n",
    "    from joblib import Parallel, delayed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Linear registration\n",
    "* This is performed using [ANTs](http://stnava.github.io/ANTs/) **Affine** registration (12 degrees of freedom: 3 translations + 3 rotations + 3 scalings + 3 shears)\n",
    "* The _antsRegistration_ function is called through its [nipype wrapper](https://nipype.readthedocs.io/en/latest/interfaces/generated/interfaces.ants/registration.html#registration).\n",
    "* [Anatomy of an antsRegistration call](https://github.com/ANTsX/ANTs/wiki/Anatomy-of-an-antsRegistration-call) does a good job of explaining what the various parameters mean."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1a. Define function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def register_linear(moving, target,\n",
    "                    moving_prefix, target_prefix,\n",
    "                    include_pairs_with):\n",
    "    \"\"\"Linearly registers moving image to target image\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    moving: 3D NIFTI image file\n",
    "    target: 3D NIFTI image file\n",
    "    moving_prefix: string\n",
    "        one of the template names \n",
    "        is used in the naming of output files\n",
    "    target_prefix: string\n",
    "        one of the template names \n",
    "        is used in the naming of output files\n",
    "    include_pairs_with: list\n",
    "        list of templates to include in execution\n",
    "        will only run registration if 'moving' or 'target' is on this list\n",
    "        \n",
    "    \n",
    "    Return\n",
    "    ------\n",
    "    Linear transformation matrix\n",
    "        file named as '{moving_prefix}_to_{target_prefix}_0GenericAffine.mat'\n",
    "    Transformed image\n",
    "        moving image transformed into target image space\n",
    "        file named as: {moving_prefix}_in_{target_prefix}_linear.nii.gz\n",
    "    \"\"\"\n",
    "    if moving_prefix in include_pairs_with or target_prefix in include_pairs_with: \n",
    "        rgstr = ants.registration.Registration(\n",
    "               dimension=3,\n",
    "               float=False,\n",
    "               fixed_image=target,\n",
    "               moving_image=moving,\n",
    "               output_transform_prefix='{0}_to_{1}_affine_'.format(moving_prefix, target_prefix),\n",
    "               output_warped_image='{0}_in_{1}_linear.nii.gz'.format(moving_prefix, target_prefix),\n",
    "               initial_moving_transform_com=0,\n",
    "               winsorize_lower_quantile=0.05,\n",
    "               winsorize_upper_quantile=0.95,\n",
    "               interpolation='Linear',\n",
    "               use_histogram_matching=[True],\n",
    "               transforms=['Affine'],\n",
    "               transform_parameters=[(0.1,)],\n",
    "               metric=['MI'],\n",
    "               metric_weight=[1],\n",
    "               radius_or_number_of_bins=[32],\n",
    "               sampling_strategy=['Regular'],\n",
    "               sampling_percentage=[0.2],\n",
    "               number_of_iterations=[[1000, 500, 250, 100]],\n",
    "               convergence_threshold=[1e-6],\n",
    "               convergence_window_size=[10],\n",
    "               shrink_factors=[[8, 4, 2, 1]],\n",
    "               smoothing_sigmas=[[3, 2, 1, 0]],\n",
    "               sigma_units=['vox'])\n",
    "    \n",
    "        rgstr.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1b. Apply function to register all unique template pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if do_parallel: # parallel\n",
    "    print('\\nPerforming linear registration in parallel...')\n",
    "    start_time = time.time()\n",
    "    Parallel(n_jobs=par_numjobs, backend='multiprocessing')(\n",
    "        delayed(register_linear)\n",
    "        (brains[0],brains[1],names[0],names[1],include_pairs_with) \n",
    "        for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)))\n",
    "    end_time = time.time()\n",
    "    print('Linear registration completed in  {0:.2f} minutes'.format((end_time - start_time)/60))\n",
    "        \n",
    "else: # serial\n",
    "    for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "        \n",
    "        # explicit here for clarity\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = brains[0]\n",
    "        targ_brain = brains[1]\n",
    "    \n",
    "        print('\\nLinearly registering {0} to {1} template...'.format(mov_pref, targ_pref))\n",
    "        start_time = time.time()\n",
    "        register_linear(mov_brain, targ_brain, mov_pref, targ_pref, include_pairs_with)\n",
    "        end_time = time.time()\n",
    "        print('Linear registration completed in  {0:.2f} minutes'.format((end_time - start_time)/60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1c. Visualize linear registration results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "    if names[0] in include_pairs_with or names[1] in include_pairs_with:\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = brains[0]\n",
    "        targ_brain = brains[1]\n",
    "        targ_index = temp_names.index(targ_pref)\n",
    "\n",
    "        display = plotting.plot_anat(targ_brain, display_mode='ortho',\n",
    "                            title='{0} edges on {1} volume'.format(mov_pref, targ_pref),\n",
    "                            draw_cross=False, annotate=False, dim=contrasts[targ_index])\n",
    "        brain = '{0}_in_{1}_linear.nii.gz'.format(mov_pref, targ_pref)\n",
    "        #display.add_edges(brain)\n",
    "        display.add_contours(brain)\n",
    "        plt.savefig('Linear_{0}_on_{1}.png'.format(mov_pref, targ_pref))\n",
    "        plt.draw()\n",
    "        display.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Non-linear registration\n",
    "* Symmetric Diffeomorphic ([SyN](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2276735/)) registration is performed using [ANTs](http://stnava.github.io/ANTs/).\n",
    "* The starting point for this step is the transformed image from the previous (linear) step\n",
    "* The _antsRegistration_ function is called through its [nipype wrapper](https://nipype.readthedocs.io/en/latest/interfaces/generated/interfaces.ants/registration.html#registration).\n",
    "* [Anatomy of an antsRegistration call](https://github.com/ANTsX/ANTs/wiki/Anatomy-of-an-antsRegistration-call) does a good job of explaining what the various parameters mean."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2a. Define function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def register_SyN(moving, target,\n",
    "                 moving_prefix, target_prefix,\n",
    "                 include_pairs_with):\n",
    "    \"\"\"Non-linearly registers moving image to target image using ANTs SyN\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    Same as register_linear function\n",
    "\n",
    "    Return\n",
    "    ------\n",
    "    Non-linear forward warp\n",
    "        file named as '{moving_prefix}_to_{target_prefix}_1Warp.nii.gz'\n",
    "    Non-linear backward warp (inverse of the above)\n",
    "        file named as '{moving_prefix}_to_{target_prefix}_1InverseWarp.nii.gz'\n",
    "    Transformed image\n",
    "        moving image transformed into target image space\n",
    "        file named as: {moving_prefix}_in_{target_prefix}_linear+SyN.nii.gz\n",
    "    \"\"\"\n",
    "    if moving_prefix in include_pairs_with or target_prefix in include_pairs_with: \n",
    "        rgstr = ants.registration.Registration(\n",
    "               dimension=3,\n",
    "               float=False,\n",
    "               fixed_image=target,\n",
    "               moving_image=moving,\n",
    "               output_transform_prefix='{0}_to_{1}_'.format(moving_prefix, target_prefix),\n",
    "               output_warped_image='{0}_in_{1}_linear+SyN.nii.gz'.format(moving_prefix, target_prefix),\n",
    "               initial_moving_transform_com=0,\n",
    "               winsorize_lower_quantile=0.05,\n",
    "               winsorize_upper_quantile=0.95,\n",
    "               interpolation='Linear',\n",
    "               use_histogram_matching=[True],\n",
    "               transforms=['SyN'],\n",
    "               transform_parameters=[(0.1,)],\n",
    "               metric=['CC'],\n",
    "               metric_weight=[1],\n",
    "               radius_or_number_of_bins=[4],\n",
    "               sampling_strategy=['Regular'],\n",
    "               sampling_percentage=[0.2],\n",
    "               number_of_iterations=[[500, 200, 50]],\n",
    "               convergence_threshold=[1e-6],\n",
    "               convergence_window_size=[10],\n",
    "               shrink_factors=[[8, 4, 2]],\n",
    "               smoothing_sigmas=[[3, 2, 1]],\n",
    "               sigma_units=['vox'])\n",
    "\n",
    "        rgstr.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2b. Apply function to register all unique template pairs    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if do_parallel: # parallel\n",
    "    print('\\nPerforming non-linear registration in parallel...')\n",
    "    start_time = time.time()\n",
    "    Parallel(n_jobs=par_numjobs, backend='multiprocessing')(\n",
    "        delayed(register_SyN)\n",
    "        ('{0}_in_{1}_linear.nii.gz'.format(names[0], names[1]), brains[1], names[0], names[1], include_pairs_with) \n",
    "        for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)))\n",
    "    end_time = time.time()\n",
    "    print('Non-linear registration completed in  {0:.2f} minutes'.format((end_time - start_time)/60))\n",
    "    \n",
    "else: # serial\n",
    "    for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "\n",
    "        # explicit here for clarity\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = '{0}_in_{1}_linear.nii.gz'.format(mov_pref, targ_pref)\n",
    "        targ_brain = brains[1]\n",
    "    \n",
    "        print('\\nNon-linearly registering {0} to {1} template...'.format(mov_pref, targ_pref))\n",
    "        start_time = time.time()\n",
    "        register_SyN(mov_brain, targ_brain, mov_pref, targ_pref, include_pairs_with)\n",
    "        end_time = time.time()\n",
    "        print('Non-linear registration completed in  {0:.2f} minutes'.format((end_time - start_time)/60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2c: Visualize non-linear registration results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "    if names[0] in include_pairs_with or names[1] in include_pairs_with:\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        targ_brain = brains[1]\n",
    "        targ_index = temp_names.index(targ_pref)\n",
    "\n",
    "        display = plotting.plot_anat(targ_brain, display_mode='ortho',\n",
    "                            title='{0} edges on {1} volume'.format(mov_pref, targ_pref),\n",
    "                            draw_cross=False, annotate=False, dim=contrasts[targ_index])\n",
    "        brain = '{0}_in_{1}_linear+SyN.nii.gz'.format(mov_pref, targ_pref)\n",
    "        #display.add_edges(brain)\n",
    "        display.add_contours(brain)\n",
    "        plt.savefig('Nonlinear_{0}_on_{1}.png'.format(mov_pref, targ_pref))\n",
    "        plt.draw()\n",
    "        display.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "toc-hr-collapsed": false
   },
   "source": [
    "## Step 3: Derive composite warps by combining linear and non-linear steps\n",
    "* This is done with the _antsApplyTransforms_ function, using its [nipype wrapper](https://nipype.readthedocs.io/en/latest/interfaces/generated/interfaces.ants/resampling.html#applytransforms).\n",
    "* **Caution:** the order with which transforms are passed is akin to linear algebra notation, so first transform is applied last\n",
    "* Both forward (templateA-to-templateB) and backward (templateB-to-templateA) warps are computed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3a. Define function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def combine_transforms(moving, target,\n",
    "                       moving_prefix, target_prefix,\n",
    "                       include_pairs_with):\n",
    "\n",
    "    \"\"\"Combines linear transform with non-linear warp to generate a composite warp\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    Same as register_linear function\n",
    "\n",
    "    Return\n",
    "    ------\n",
    "    Composite forward warp\n",
    "        file named as '{moving_prefix}_to_{target_prefix}_CompositeWarp.nii.gz'\n",
    "    Composite backward warp (inverse of the above)\n",
    "        file named as '{target_prefix}_to_{moving_prefix}_CompositeWarp.nii.gz'\n",
    "    \"\"\"\n",
    "    if moving_prefix in include_pairs_with or target_prefix in include_pairs_with: \n",
    "        forward = ants.ApplyTransforms(\n",
    "                input_image=moving,\n",
    "                output_image='{0}_to_{1}_CompositeWarp.nii.gz'.format(moving_prefix, target_prefix),\n",
    "                reference_image=target,\n",
    "                transforms=['{0}_to_{1}_1Warp.nii.gz'.format(moving_prefix, target_prefix),\n",
    "                            '{0}_to_{1}_affine_0GenericAffine.mat'.format(moving_prefix, target_prefix)],\n",
    "                dimension=3,\n",
    "                print_out_composite_warp_file=True)\n",
    "        forward.run()\n",
    "\n",
    "        backward = ants.ApplyTransforms(\n",
    "                input_image=target,\n",
    "                output_image='{1}_to_{0}_CompositeWarp.nii.gz'.format(moving_prefix, target_prefix),\n",
    "                reference_image=moving,\n",
    "                transforms=['{0}_to_{1}_affine_0GenericAffine.mat'.format(moving_prefix, target_prefix),\n",
    "                            '{0}_to_{1}_1InverseWarp.nii.gz'.format(moving_prefix, target_prefix)],\n",
    "                invert_transform_flags=[True, False], # invert the linear transform\n",
    "                dimension=3,\n",
    "                print_out_composite_warp_file=True)\n",
    "\n",
    "        backward.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3b. Apply function to generate all composite warp pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if do_parallel: # parallel\n",
    "    print('\\nCreating composite warps in parallel...')\n",
    "    start_time = time.time()\n",
    "    Parallel(n_jobs=par_numjobs, backend='multiprocessing')(\n",
    "        delayed(combine_transforms)\n",
    "        (brains[0], brains[1], names[0], names[1], include_pairs_with) \n",
    "        for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)))\n",
    "    end_time = time.time()\n",
    "    print('Composite warps computed in {0:.2f} minutes'.format((end_time - start_time)/60))\n",
    "    \n",
    "else: # serial\n",
    "    for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "\n",
    "        # make explicit here\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = brains[0]\n",
    "        targ_brain = brains[1]\n",
    "\n",
    "        print('\\nCreating composite warps between {0} and {1} templates...'.format(mov_pref, targ_pref))\n",
    "        start_time = time.time()\n",
    "        combine_transforms(mov_brain, targ_brain, mov_pref, targ_pref, include_pairs_with)\n",
    "        end_time = time.time()\n",
    "        print('Composite warps computed in {0:.2f} minutes'.format((end_time - start_time)/60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3c. Visually assess the derived composite warp pairs\n",
    "* First use the derived composite warps to transform brains in both directions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "    if names[0] in include_pairs_with or names[1] in include_pairs_with:\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = brains[0]\n",
    "        targ_brain = brains[1]\n",
    "\n",
    "        forward = ants.ApplyTransforms(\n",
    "                input_image=mov_brain,\n",
    "                output_image='{0}_in_{1}_composite.nii.gz'.format(mov_pref, targ_pref),\n",
    "                reference_image=targ_brain,\n",
    "                transforms=['{0}_to_{1}_CompositeWarp.nii.gz'.format(mov_pref, targ_pref)],\n",
    "                dimension=3)\n",
    "        backward = ants.ApplyTransforms(\n",
    "                input_image=targ_brain,\n",
    "                output_image='{1}_in_{0}_composite.nii.gz'.format(mov_pref, targ_pref),\n",
    "                reference_image=mov_brain,\n",
    "                transforms=['{1}_to_{0}_CompositeWarp.nii.gz'.format(mov_pref, targ_pref)],\n",
    "                dimension=3)\n",
    "\n",
    "        print('\\nTransforming {0} into {1} space...'.format(mov_pref, targ_pref))\n",
    "        start_time = time.time()\n",
    "        forward.run()\n",
    "        end_time = time.time()\n",
    "        print('Transformation complete in {0:.2f} minutes'.format((end_time - start_time)/60))\n",
    "        print('Transforming {1} into {0} space...'.format(mov_pref, targ_pref))\n",
    "        start_time = time.time()\n",
    "        backward.run()\n",
    "        end_time = time.time()\n",
    "        print('Transformation complete in {0:.2f} minutes'.format((end_time - start_time)/60))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Then, overlay the transformed (by composite warps) brains on target brains for visualization\n",
    "* two plots (forward/backward) will be generated for each template pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for names, brains in zip(combinations(temp_names, 2), combinations(temp_brains, 2)):\n",
    "    if names[0] in include_pairs_with or names[1] in include_pairs_with:\n",
    "        mov_pref = names[0]\n",
    "        targ_pref = names[1]\n",
    "        mov_brain = brains[0]\n",
    "        targ_brain = brains[1]\n",
    "        mov_index = temp_names.index(mov_pref)\n",
    "        targ_index = temp_names.index(targ_pref)\n",
    "\n",
    "        display = plotting.plot_anat(targ_brain, display_mode='ortho',\n",
    "                            title='{0} edges on {1} volume'.format(mov_pref, targ_pref),\n",
    "                            draw_cross=False, annotate=False, dim=contrasts[targ_index])\n",
    "        brain = '{0}_in_{1}_composite.nii.gz'.format(mov_pref, targ_pref)\n",
    "        #display.add_edges(brain)\n",
    "        display.add_contours(brain)\n",
    "        plt.savefig('Composite_{0}_on_{1}.png'.format(mov_pref, targ_pref))\n",
    "        plt.draw()\n",
    "        display.close()\n",
    "\n",
    "        display = plotting.plot_anat(mov_brain, display_mode='ortho',\n",
    "                            title='{1} edges on {0} volume'.format(mov_pref, targ_pref),\n",
    "                            draw_cross=False, annotate=False, dim=contrasts[mov_index])\n",
    "        brain = '{1}_in_{0}_composite.nii.gz'.format(mov_pref, targ_pref)\n",
    "        #display.add_edges(brain)\n",
    "        display.add_contours(brain)\n",
    "        plt.savefig('Composite_{1}_on_{0}.png'.format(mov_pref, targ_pref))\n",
    "        plt.draw()\n",
    "        plt.show()\n",
    "        display.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Clean-up\n",
    "* Create subfolders for linear and non-linear warps, and for the warped volumes\n",
    "* Put all files in sensible subfolders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(BASE_path + \"warps/linear\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"warps/nonlinear\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"warps/final\" , exist_ok=True)\n",
    "\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"warps/linear/\" + f) for f in glob.glob('*_affine_*Affine.mat')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"warps/nonlinear/\" + f) for f in glob.glob('*_to_*_1Warp.nii.gz')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"warps/nonlinear/\" + f) for f in glob.glob('*_to_*_1InverseWarp.nii.gz')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"warps/final/\" + f) for f in glob.glob('*_to_*_CompositeWarp.nii.gz')];\n",
    "[os.remove(f) for f in glob.glob('*Affine.mat')]; # spurious linear warps (generated by nonlinear)\n",
    "\n",
    "os.makedirs(BASE_path + \"warped_templates/linear\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"warped_templates/nonlinear\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"warped_templates/final\" , exist_ok=True)\n",
    "\n",
    "[os.rename(os.getcwd() + '/' +  f, BASE_path + \"warped_templates/linear/\" + f) for f in glob.glob('*_in_*linear.nii.gz')];\n",
    "[os.rename(os.getcwd() + '/' +  f, BASE_path + \"warped_templates/nonlinear/\" + f) for f in glob.glob('*_in_*linear+SyN.nii.gz')];\n",
    "[os.rename(os.getcwd() + '/' +  f, BASE_path + \"warped_templates/final/\" + f) for f in glob.glob('*_in_*composite.nii.gz')];\n",
    "\n",
    "os.makedirs(BASE_path + \"images/templates\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"images/linear_reg\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"images/nonlinear_reg\" , exist_ok=True)\n",
    "os.makedirs(BASE_path + \"images/all_warp_pairs\" , exist_ok=True)\n",
    "\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"images/templates/\" + f) for f in glob.glob('*_template.png')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"images/linear_reg/\" + f) for f in glob.glob('Linear*.png')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"images/nonlinear_reg/\" + f) for f in glob.glob('Nonlinear*.png')];\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"images/all_warp_pairs/\" + f) for f in glob.glob('Composite*.png')];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. We will now create the average 'meta'-template in NMTv2.0 symmetric space (just because we can)\n",
    "* Get all the templates (in NMTv2.0 sym space) & average them into a 'MetaTemplate'\n",
    "* Plot the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wtp = BASE_path + \"warped_templates/final/\"\n",
    "image_list = [NMTv20_sym_brain]\n",
    "for file in os.listdir(wtp):\n",
    "    if file.split('_')[2] == 'NMTv2.0-sym':\n",
    "        image_list.append(wtp + file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "avg = ants.AverageImages(\n",
    "    dimension=3,\n",
    "    output_average_image=\"MetaTemplate_in_NMTv2.0-sym.nii.gz\",\n",
    "    normalize=True,\n",
    "    images=image_list\n",
    ");\n",
    "avg.run();\n",
    "\n",
    "plotting.plot_anat(\"MetaTemplate_in_NMTv2.0-sym.nii.gz\", display_mode='ortho', \n",
    "          title=\"MetaTemplate in NMTv2.0-sym\", draw_cross=False, annotate=True, dim=con);\n",
    "plt.savefig('MetaTemplate_in_NMTv2.0-sym.png')\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(BASE_path + \"warped_templates/metatemplate\" , exist_ok=True)\n",
    "[os.rename(os.getcwd() + '/' +  f, BASE_path + \"warped_templates/metatemplate/\" + f) for f in glob.glob('MetaTemplate*.nii.gz')];\n",
    "os.makedirs(BASE_path + \"images/metatemplate\" , exist_ok=True)\n",
    "[os.rename(os.getcwd() + '/' + f, BASE_path + \"images/metatemplate/\" + f) for f in glob.glob('MetaTemplate*.png')];"
   ]
  },
  {
   "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.7.12"
  },
  "toc-showcode": false,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}