{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "32275606",
   "metadata": {},
   "source": [
    "# Using TIRL to map between MRI and microscopy data in BigMac\n",
    "\n",
    "The BigMac MRI and microscopy data is co-registered using the Tensor Image Registration Library (TIRL) in FSL (Huszar2019). This tutorial provides a basic overview of how to map between MRI and microscopy data in BigMac.\n",
    "\n",
    "More information on TIRL can be found at https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/TIRL and https://git.fmrib.ox.ac.uk/ihuszar/tirl. The git also includes a tutorial. The relevant paper is https://www.biorxiv.org/content/10.1101/849570v1.full.pdf.\n",
    "\n",
    "In BigMac, TIRL has been run on every microscopy slide in turn, to create a warpfield between the microscopy and structural MRI. This warpfield is stored as a TIRL image or .timg file.\n",
    "\n",
    "Here the structural MRI is used as reference due to its high spatial resolution and good image contrast (clear gm/wm delineation).\n",
    "\n",
    "MRI-MRI warpfields generated using FSLs FLIRT/FNIRT can be used to map between the microscopy and other MRI contrasts such as dMRI.\n",
    "\n",
    "Note, TIRL is all python-based. It is included in recent versions of FSL ann can be run using fslpython. Alternatively tirl can be installed via github above. If  doing the latter, please use version 2.1.3b1 or above\n",
    "\n",
    "The TIRL transformations in each timg describe going from the native/voxel space to the tirl/physical space where the data are co-aligned.\n",
    "\n",
    "Scripts for TIRL mapping in BigMac can be found at: https://git.fmrib.ox.ac.uk/amyh/bigmacanalysis/-/tree/main/TIRL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5005dcc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Libraries\n",
    "import tirl\n",
    "import numpy as np\n",
    "from tirl import fsl\n",
    "from tirl.timage import TImage\n",
    "import os\n",
    "import sys\n",
    "from PIL import Image\n",
    "import nibabel as nib\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42cbf258",
   "metadata": {},
   "source": [
    "## Load data and warpfields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "93ed7d86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# BigMac directory\n",
    "ddir = \"/vols/Data/km/ahoward/BigMac-clean/\"\n",
    "\n",
    "# Folder containing microscopy data and tirl warpfields\n",
    "base = ddir + \"/Microscopy/PLI/Anterior/P089x/\"\n",
    "\n",
    "# timg containing information to transform microscopy to physical space \n",
    "# (i.e. tirl intemediary space) where the data are coregistered\n",
    "micro_timg = base + \"/Reg2MRI/Full/slice.timg\"\n",
    "\n",
    "# timg containing information to transfrom mri to physical space\n",
    "mri_timg = base + \"/Reg2MRI/Full/volume.timg\"\n",
    "\n",
    "# High resolution PLI images\n",
    "pli_inplane = base+\"/In_plane.tif\" # The data we want to warp\n",
    "pli_inclination = base+\"/Inclination.tif\" # The image used to drive the registration\n",
    "\n",
    "# FNIRT transformations, dwi -> struct (opposite to what you think you need)\n",
    "# b10k 1mm data\n",
    "dwi = ddir + \"/MRI/Postmortem/dwi/b10k/1.0mm/data/S0.nii.gz\"\n",
    "field = ddir + \"/MRI/Postmortem/dwi/b10k/1.0mm/reg/dwi_1.0mm_2_struct_fnirt/dwi_1.0mm_2_struct_fnirt_field.nii.gz\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c3802fae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load microscopy timg\n",
    "micro = tirl.load(micro_timg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f5abdf3",
   "metadata": {},
   "source": [
    "## Inbuilt functions\n",
    "TIRL Timages contain lots of useful inbuilt functions. Use autocomplete to see full list.\\\n",
    "\n",
    "#### Timage structure\n",
    "- micro.domain.chain = transformation chain native => physical space where data is coregistered\n",
    "- micro.data = access / manipulate data\n",
    "\n",
    "#### Useful functions\n",
    "- micro.evaluate(domain) # => interpolate data at coords of some other domain\n",
    "- micro.snapshot(fout) # => save snapshot image of data to fout\n",
    "- micro.domain.chain # => access transformation chain. Output from tirl defined as from native/voxel space -> physical/tirl space\n",
    "- micro.domain.chain.inverse() # => invert chain\n",
    "- micro.domain.copy_transformations(domain) # => copy transformations from one timg to another\n",
    "- micro.domain.get_physical_coordinates() # => get microscopy coordinates in physical space where data is coregistered\n",
    "- micro.domain.get_voxel_coordinates() # => get microscopy pixels / mri voxels\n",
    "- micro.domain.chain.map_vectors(vectors,coordinates) # => map vectors (at voxel coordinates) through transformation chain\n",
    "- micro.domain.chain.map(coordinates) # => map coordinates through transformation chain\n",
    "- micro.domain.chain.map_vector_and_coordinates(vectors,coordinates) # => map vectors and voxel coordinates\n",
    "- micro.domain.map_physical_coordinates() # => map coordinates in physical space to voxel/native space\n",
    "- micro.data.flatten() # => returns flattened data\n",
    "- micro.shape # => shape of data (after offset transformations)\n",
    "- micro.header # => header information\n",
    "- micro.interpolator(coords) # => interpolate data at coordinate positions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "84608d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note we can load any data as a TImage to access these functions\n",
    "# from tirl.timage import TImage\n",
    "fa = TImage('/vols/Data/BigMac/MRI/Postmortem/dwi/b10k/1.0mm/data.dtifit/dti_FA.nii.gz')\n",
    "inplane = TImage(pli_inplane)\n",
    "inplane.domain.copy_transformations(micro.domain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8fa570f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# accessing help\n",
    "help(micro.domain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a20eaf",
   "metadata": {},
   "source": [
    "## Combining warpfields from tirl and flirt/fnirt\n",
    "\n",
    "#### example: create warpfield microscopy to dwi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6627e47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Load microscopy slide that we have registered to structural MRI space using TIRL\n",
    "# (microscopy -> physical)\n",
    "micro = tirl.load(micro_timg)\n",
    "chain = micro.domain.chain.copy()\n",
    "\n",
    "# 2. Create warpfield from microscopy to structural\n",
    "# (microscopy -> physical -> mri)\n",
    "struct = tirl.load(mri_timg)\n",
    "# Invert warpfield in mri_timg and add to warpfield from micro_timg \n",
    "tirl2struct = struct.domain.chain.inverse()\n",
    "chain += tirl2struct\n",
    "\n",
    "# 3. Import FNIRT transformation & add to transformation chain\n",
    "# (microscopy -> physical -> structural -> dwi)\n",
    "# Note, opposite transform to what you might expect due to how fsl defines warpfields\n",
    "warp = fsl.load_warp(field, moving=dwi)\n",
    "chain += warp\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d89833e",
   "metadata": {},
   "source": [
    "## Mapping coordinates \n",
    "\n",
    "#### a) map pixel coordinates to structural  MRI voxel locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "202492a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[229.49050031 107.73321454 218.01824585]\n",
      " [229.47817781 107.73373977 218.01531657]\n",
      " [229.46585674 107.73426504 218.01238556]\n",
      " [229.45353709 107.73479037 218.00945283]\n",
      " [229.44121889 107.73531575 218.00651839]\n",
      " [229.42890211 107.73584118 218.00358222]\n",
      " [229.41658677 107.73636666 218.00064433]\n",
      " [229.40427287 107.73689219 217.99770471]\n",
      " [229.3919604  107.73741778 217.99476338]\n",
      " [229.37964936 107.73794342 217.99182032]]\n"
     ]
    }
   ],
   "source": [
    "# Map coordinates of high resolution microscopy image through tirl chain to dwi space\n",
    "micro_hr = TImage(pli_inclination)\n",
    "micro_hr.domain.copy_transformations(micro.domain)\n",
    "    \n",
    "# Find hr pixel coordinates in physical (tirl) space\n",
    "pixel_physical_coords = micro_hr.domain.get_physical_coordinates()\n",
    "\n",
    "# Map to structural space\n",
    "pixel_coords_in_struct = tirl2struct.map(pixel_physical_coords[0:10,:])\n",
    "# Here only mapping first 10 pixels as an example\n",
    "print(pixel_coords_in_struct)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "177db55e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[229.49050031 107.73321454 218.01824585]\n",
      " [229.47817781 107.73373977 218.01531657]\n",
      " [229.46585674 107.73426504 218.01238556]\n",
      " [229.45353709 107.73479037 218.00945283]\n",
      " [229.44121889 107.73531575 218.00651839]\n",
      " [229.42890211 107.73584118 218.00358222]\n",
      " [229.41658677 107.73636666 218.00064433]\n",
      " [229.40427287 107.73689219 217.99770471]\n",
      " [229.3919604  107.73741778 217.99476338]\n",
      " [229.37964936 107.73794342 217.99182032]]\n"
     ]
    }
   ],
   "source": [
    "# Note, this is the same as:\n",
    "micro = tirl.load(micro_timg)\n",
    "micro2struct = micro.domain.chain + tirl2struct\n",
    "pixel_native_coords = micro_hr.domain.get_voxel_coordinates()  # coordinates in native space\n",
    "pixel_coords_in_struct = micro2struct.map(pixel_native_coords[0:10,:])\n",
    "print(pixel_coords_in_struct)\n",
    "# This version is less memory intensive as we are only ever mapping 10 pixels. In comparison,\n",
    "# the function get_pixel_coordinates() maps all microscopy pixels to real space."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d7ee43d",
   "metadata": {},
   "source": [
    "#### b) map pixel coordinates to dwi voxels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dde304d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[68.44825822 31.06628612 66.21822589]\n",
      " [68.44448256 31.06629764 66.21740675]\n",
      " [68.44070676 31.06630796 66.21658779]\n",
      " [68.43693085 31.0663171  66.21576896]\n",
      " [68.43315482 31.06632511 66.21495027]\n",
      " [68.42937871 31.06633201 66.21413169]\n",
      " [68.42560253 31.06633783 66.2133132 ]\n",
      " [68.42182628 31.06634262 66.21249479]\n",
      " [68.41804999 31.0663464  66.21167644]\n",
      " [68.41427366 31.0663492  66.21085813]]\n"
     ]
    }
   ],
   "source": [
    "# Define transformation chain from tirl/physical space to dwi space\n",
    "tirl2dwi = tirl2struct + warp\n",
    "\n",
    "# Map to dwi\n",
    "pixel_coords_in_dwi = tirl2dwi.map(pixel_physical_coords[0:10,:])\n",
    "print(pixel_coords_in_dwi)\n",
    "\n",
    "# Again this is equivalent to (where the following is less memory intensive)\n",
    "# pixel_coords_in_dwi = chain.map(pixel_native_coords[0:10,:])\n",
    "# where chain maps micro->dwi (see chain construction above)\n",
    "# print(pixel_coords_in_dwi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a531f7f5",
   "metadata": {},
   "source": [
    "## Mapping vectors\n",
    "TIRL can also map vectors and account for the rotations in the transformation chain.\\\n",
    "Mapping the normal from the microscopy plane is slightly more complex as the TIRL transformation chain has some 2D and some 3D transformations.\\ \n",
    "In general, we must take care to make sure we have a good microscopy -> mri mapping.\n",
    "\n",
    "#### a) mapping the orientation angle from PLI through to dwi space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed7b981e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966752  0.00262689  0.14335721]\n",
      " [-0.98966046  0.00263049  0.14340588]\n",
      " [-0.98966046  0.00263049  0.14340588]\n",
      " [-0.98966046  0.00263049  0.14340588]]\n"
     ]
    }
   ],
   "source": [
    "# Map 2D PLI vector (orientation angle) through tirl chain into diffusion space\n",
    "\n",
    "# Load in-plane angle\n",
    "inplane = TImage(pli_inplane)\n",
    "# Change scale 0-pi\n",
    "inplane.normalise()\n",
    "inplane = inplane*np.pi\n",
    "\n",
    "# Convert inplane angle to x/y \n",
    "\n",
    "x = -np.sin(inplane[:]).ravel()\n",
    "y = np.cos(inplane[:]).ravel()\n",
    "xy = np.c_[x, y]\n",
    "\n",
    "# map pli vector through img.domain.chain -> dwi space\n",
    "pixel_coords = inplane.domain.get_voxel_coordinates()\n",
    "inplane_in_dwi = chain.map_vector(xy[0:10,:], pixel_coords[0:10,:], rule=\"fs\")\n",
    "\n",
    "# flip vectors LR so they are stored in a similar manner to FSL vectors (e.g. the output from dtifit)\n",
    "inplane_in_dwi = inplane_in_dwi*[-1,1,1]\n",
    "print(inplane_in_dwi)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a036c13",
   "metadata": {},
   "source": [
    "#### b) mapping microscopy normal to dwi space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0aa40d19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01226941 -0.99771633 -0.06641971]\n",
      " [-0.01227313 -0.99771779 -0.06639714]\n",
      " [-0.01227313 -0.99771779 -0.06639714]\n",
      " [-0.01227313 -0.99771779 -0.06639714]]\n"
     ]
    }
   ],
   "source": [
    "# Need to separate 2D and 3D transformations in tirl\n",
    "chain2D = chain[:7]\n",
    "chain3D = chain[7:]\n",
    "\n",
    "# Define normal after 2D transformations\n",
    "norm = np.zeros(inplane_in_dwi.shape)\n",
    "norm[:,2] = 1\n",
    "pixel_coords_2D = chain2D.map(pixel_coords)\n",
    "norm_in_dwi = chain3D.map_vector(norm[0:10,:], pixel_coords_2D[0:10,:], rule=\"fs\")\n",
    "\n",
    "# flip vectors LR so they are stored in a similar manner to FSL vectorss (e.g. the output from dtifit)\n",
    "norm_in_dwi = norm_in_dwi*[-1,1,1]\n",
    "print(norm_in_dwi)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7db57b2",
   "metadata": {},
   "source": [
    "## Interpolating / resampling data\n",
    "\n",
    "#### e.g. resampling dwi data on microscopy slide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9db68486",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load FA image as TImage and compile correct transformation chain (dwi -> real)\n",
    "fa = TImage(ddir + '/MRI/Postmortem/dwi/b10k/1.0mm/data.dtifit/dti_FA.nii.gz')\n",
    "\n",
    "# Since mapping to HR microscopy space can be time consuming, can check mapping by instead interpolating \n",
    "# in the downsampled microscopy space, as defined by micro (instead of micro_hr)\n",
    "# e.g\n",
    "micro = TImage(micro_timg)\n",
    "micro.domain.chain = chain\n",
    "fa_in_micro = fa.evaluate(micro.domain)\n",
    "fa_in_micro.snapshot('./tmp_FA.tif',overwrite=True)\n",
    "micro.snapshot('./tmp_pli.tif',overwrite=True)\n",
    "\n",
    "# OR\n",
    "micro = TImage(micro_timg)\n",
    "fa_in_micro = fa.interpolator(tirl2dwi.map(micro.domain.get_physical_coordinates()))\n",
    "fa_in_micro = TImage(fa_in_micro.reshape(micro.shape))\n",
    "fa_in_micro.snapshot('./tmp_FA2.tif',overwrite=True)\n",
    "\n",
    "# OR can change interpolator\n",
    "from tirl.interpolators.scipyinterpolator import ScipyNearestNeighbours\n",
    "fa.interpolator = ScipyNearestNeighbours\n",
    "micro.domain.chain = chain\n",
    "fa_in_micro_nn = fa.evaluate(micro.domain)\n",
    "fa_in_micro_nn.snapshot('./tmp_FA3.tif',overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eb78cc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# OR to map to high resolution microscopy (computationally expensive):\n",
    "\n",
    "# Load FA image as TImage and compile correct transformation chain (dwi -> real)\n",
    "fa = TImage('/vols/Data/BigMac/MRI/Postmortem/dwi/b10k/1.0mm/data.dtifit/dti_FA.nii.gz')\n",
    "\n",
    "# Resample FA at microscopy pixel coordinates\n",
    "micro_hr.domain.chain = chain # chain defined above & transforms from microscopy -> dwi\n",
    "fa_in_micro = fa.evaluate(micro_hr.domain)\n",
    "fa_in_micro.snapshot('./tmp.tif',overwrite=True)\n",
    "\n",
    "# This is equivalent to\n",
    "fa_in_micro = fa.interpolator(chain.map(micro_hr.domain.get_voxel_coordinates()))\n",
    "fa_in_micro = TImage(fa_in_micro.reshape(micro_hr.shape))\n",
    "fa_in_micro.snapshot('./tmp2.tif')\n",
    "\n",
    "# OR\n",
    "fa.domain.chain = tirl2dwi.inverse() # This will take a long time to compute. Better to compile from fnirt/flirt & struct2tirl\n",
    "micro = TImage(micro_timg)\n",
    "micro_hr.domain.copy_transformations(micro.domain)\n",
    "fa_in_micro = fa.evaluate(micro_hr.domain)\n",
    "fa_in_micro.snapshot('./tmp3.tif',overwrite=True)\n",
    "\n",
    "# OR\n",
    "fa_in_micro = fa.interpolator(tirl2dwi.map(micro_hr.domain.get_physical_coordinates()))\n",
    "fa_in_micro = TImage(fa_in_micro.reshape(micro_hr.shape))\n",
    "fa_in_micro.snapshot('./tmp4.tif')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e899b6b",
   "metadata": {},
   "source": [
    "## Manipulating nifti volumes \n",
    "\n",
    "#### e.g. to save some output as a nifti\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc062e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy header information from dwi image and change resolution\n",
    "dwi_im = nib.load(dwi)\n",
    "new_hdr = dwi_im.header\n",
    "new_hdr.set_zooms([0.004,0.004,0.35])\n",
    "\n",
    "# Save pixel coordinates as a nifti volume\n",
    "niftiobj = nib.Nifti1Image(pixel_coords_in_dwi, new_hdr.get_best_affine(), new_hdr)\n",
    "nib.save(niftiobj, './tmp.nii.gz')"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}