{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Nonrigid Bundle Registration with BundleWarp\n\nThis example explains how you can nonlinearly register two bundles from two\ndifferent subjects directly in the space of streamlines [Chandio23]_, [Chandio20]_.\n\nTo show the concept, we will use two pre-saved uncinate fasciculus bundles. The\nalgorithm used here is called BundleWarp, streamline-based nonlinear\nregistration of white matter tracts [Chandio23]_.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from os.path import join as pjoin\n\nfrom dipy.align.streamwarp import (bundlewarp, bundlewarp_vector_filed,\n                                   bundlewarp_shape_analysis)\nfrom dipy.data import fetch_bundle_warp_dataset\nfrom dipy.io.stateful_tractogram import Space, StatefulTractogram\nfrom dipy.io.streamline import save_tractogram, load_trk\nfrom dipy.tracking.streamline import (set_number_of_points, unlist_streamlines,\n                                      Streamlines)\nfrom dipy.viz.streamline import (viz_two_bundles, viz_vector_field,\n                                 viz_displacement_mag)\nfrom time import time"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's download and loaf two uncinate fasciculus bundles in the left hemisphere\nof the brain (UF_L) available here:\nhttps://figshare.com/articles/dataset/Test_Bundles_for_DIPY/22557733\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "bundle_warp_files = fetch_bundle_warp_dataset()\ns_UF_L_path = pjoin(bundle_warp_files[1], 's_UF_L.trk')\nm_UF_L_path = pjoin(bundle_warp_files[1], 'm_UF_L.trk')\n\nuf_subj1 = load_trk(s_UF_L_path, reference=\"same\",\n                    bbox_valid_check=False).streamlines\nuf_subj2 = load_trk(m_UF_L_path, reference=\"same\",\n                    bbox_valid_check=False).streamlines"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's resample the streamlines so that they both have the same number of points\nper streamline. Here we will use 20 points.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "static = Streamlines(set_number_of_points(uf_subj1, 20))\nmoving = Streamlines(set_number_of_points(uf_subj2, 20))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We call ``uf_subj2`` a moving bundle as it will be nonlinearly aligned with\n``uf_subj1`` (static) bundle. Here is how this is done.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize static bundle in red and moving in green before registration.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "viz_two_bundles(static, moving, fname=\"static_and_moving.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "BundleWarp method provides a unique ability to either partially or fully deform\na moving bundle by the use of a single regularization parameter alpha.\nalpha controls the trade-off between regularizing the deformation and having\npoints match very closely. The lower the value of alpha, the more closely the\nbundles would match.\n\nLet's partially deform bundle by setting alpha=0.5.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "start = time()\ndeformed_bundle, moving_aligned, distances, match_pairs, warp_map = bundlewarp(\n                               static, moving, alpha=0.5, beta=20, max_iter=15)\nend = time()\n\nprint(\"time taken by BundleWarp registration in seconds = \", end-start)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize static bundle in red and moved (warped) in green. Note: You can\nset interactive=True in visualization functions throughout this tutorial if you\nprefer to get interactive visualization window.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "viz_two_bundles(static, deformed_bundle,\n                fname=\"static_and_partially_deformed.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize linearly moved bundle in blue and nonlinearly moved bundle in\ngreen to see BundleWarp registration improvement over linear SLR registration.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "viz_two_bundles(moving_aligned, deformed_bundle,\n                fname=\"linearly_and_nonlinearly_moved.png\", c1=(0, 0, 1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, let's visualize deformation vector field generated by BundleWarp.\nThis shows us visually where and how much and in what directions deformations\nwere added by BundleWarp.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "offsets, directions, colors = bundlewarp_vector_filed(moving_aligned,\n                                                      deformed_bundle)\n\npoints_aligned, _ = unlist_streamlines(moving_aligned)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualizing just the vector field.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"partially_vectorfield.png\"\nviz_vector_field(points_aligned, directions, colors, offsets, fname)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize vector field over linearly moved bundle. This will show how\nmuch deformations were introduced after linear registration.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"partially_vectorfield_over_linearly_moved.png\"\nviz_vector_field(points_aligned, directions, colors, offsets, fname,\n                 moving_aligned)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also visualize the magnitude of deformations in mm mapped over affinely\nmoved bundle. It shows which streamlines were deformed the most after affine\nregistration.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"partially_deformation_magnitude_over_linearly_moved.png\"\nviz_displacement_mag(moving_aligned, offsets, fname, interactive=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Saving partially warped bundle.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "new_tractogram = StatefulTractogram(deformed_bundle, m_UF_L_path, Space.RASMM)\nsave_tractogram(new_tractogram, \"partially_deformed_bundle.trk\",\n                bbox_valid_check=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's fully deform the moving bundle by setting alpha <= 0.01\n\nWe will use MDF distances computed and returned by previous run of BundleWarp\nmethod. This will save computation time.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "start = time()\ndeformed_bundle2, moving_aligned, distances, match_pairs, warp_map = bundlewarp(\n        static, moving, dist=distances, alpha=0.001, beta=20)\nend = time()\n\nprint(\"time taken by BundleWarp registration in seconds = \", end-start)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize static bundle in red and moved (completely warped) in green.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "viz_two_bundles(static, deformed_bundle2,\n                fname=\"static_and_fully_deformed.png\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, let's visualize the deformation vector field generated by BundleWarp.\nThis shows us visually where and how much and in what directions deformations\nwere added by BundleWarp to perfectly warp moving bundle to look like static.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "offsets, directions, colors = bundlewarp_vector_filed(moving_aligned,\n                                                      deformed_bundle2)\n\npoints_aligned, _ = unlist_streamlines(moving_aligned)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualizing just the vector field.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"fully_vectorfield.png\"\nviz_vector_field(points_aligned, directions, colors, offsets, fname)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize vector field over linearly moved bundle. This will show how\nmuch deformations were introduced after linear registration by fully deforming\nthe moving bundle.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"fully_vectorfield_over_linearly_moved.png\"\nviz_vector_field(points_aligned, directions, colors, offsets, fname,\n                 moving_aligned)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's visualize the magnitude of deformations in mm mapped over affinely moved\nbundle. It shows which streamlines were deformed the most after affine\nregistration.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fname = \"fully_deformation_magnitude_over_linearly_moved.png\"\nviz_displacement_mag(moving_aligned, offsets, fname, interactive=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also perform bundle shape difference analysis using the displacement\nfield generated by fully warping the moving bundle to look exactly like static\nbundle. Here, we plot bundle shape profile using BUAN. Bundle shape profile\nshows the average magnitude of deformations along the length of the bundle.\nSegments where we observe higher deformations are the areas where two bundles\ndiffer the most in shape.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "_, _ = bundlewarp_shape_analysis(moving_aligned, deformed_bundle, no_disks=10,\n                                 plotting=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Saving fully warped bundle.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "new_tractogram = StatefulTractogram(deformed_bundle2, m_UF_L_path,\n                                    Space.RASMM)\nsave_tractogram(new_tractogram, \"fully_deformed_bundle.trk\",\n                bbox_valid_check=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## References\n\n.. [Chandio23] Chandio et al., \"BundleWarp, streamline-based nonlinear\n                    registration of white matter tracts.\"\n                    bioRxiv (2023): 2023-01.\n.. [Chandio20] Chandio and Garyfallidis., \"StND: Streamline-based non-rigid\n                    partial-deformation tractography registration.\" Medical\n                    Imaging Meets NeurIPS (2020).\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}