02_nodes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module nodes\n",
    "\n",
    "> Modules containing all Nodes used in workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/bsipes/opt/anaconda3/envs/tracts/lib/python3.9/site-packages/bids/config.py:39: FutureWarning: Setting 'extension_initial_dot' will be removed in pybids 0.16.\n",
      "  warnings.warn(\"Setting 'extension_initial_dot' will be removed in pybids 0.16.\",\n",
      "/Users/bsipes/opt/anaconda3/envs/tracts/lib/python3.9/site-packages/nilearn/datasets/__init__.py:93: FutureWarning: Fetchers from the nilearn.datasets module will be updated in version 0.9 to return python strings instead of bytes and Pandas dataframes instead of Numpy arrays.\n",
      "  warn(\"Fetchers from the nilearn.datasets module will be \"\n"
     ]
    }
   ],
   "source": [
    "#exporti\n",
    "import os\n",
    "import bids\n",
    "\n",
    "bids.config.set_option(\"extension_initial_dot\", True)\n",
    "\n",
    "from pathlib import Path\n",
    "from itertools import product\n",
    "\n",
    "import pipetography.core as ppt\n",
    "\n",
    "from nipype import IdentityInterface, Function\n",
    "from nipype.interfaces.io import SelectFiles, DataSink\n",
    "from nipype.pipeline import Node, MapNode, Workflow\n",
    "from nipype.interfaces.mrtrix3.utils import (\n",
    "    BrainMask,\n",
    "    TensorMetrics,\n",
    "    DWIExtract,\n",
    "    MRMath,\n",
    ")\n",
    "from nipype.interfaces.mrtrix3.preprocess import MRDeGibbs, DWIBiasCorrect, ResponseSD\n",
    "from nipype.interfaces.mrtrix3.reconst import (\n",
    "    FitTensor,\n",
    "    EstimateFOD,\n",
    "    ConstrainedSphericalDeconvolution,\n",
    ")\n",
    "from nipype.interfaces import ants\n",
    "from nipype.interfaces import fsl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class PreProcNodes:\n",
    "    \"\"\"\n",
    "    Initiate DWI preprocessing pipeline nodes. This doesn't connect the nodes into a pipeline.\n",
    "\n",
    "    Inputs:\n",
    "        - bids_dir (str) - \n",
    "        - bids_path_template (dict) - template for file naming conventions\n",
    "        - bids_ext (str) - \n",
    "        - rpe_design (str) - This is necessary to initiate nodes for single DWI volume, or additional nodes for reverse volume.\n",
    "        - mrtrix_nthreads (int)\n",
    "        - img_resol (str)\n",
    "        - sub_list (List)\n",
    "        - ses_list (List)\n",
    "        - exclude_list (tuple)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        bids_dir,\n",
    "        bids_path_template,\n",
    "        bids_ext,\n",
    "        rpe_design,\n",
    "        mrtrix_nthreads,\n",
    "        img_resol,\n",
    "        sub_list,\n",
    "        ses_list,\n",
    "        exclude_list=[()],\n",
    "    ):\n",
    "        # filtter & create sub-graphs for subjects and sessions combos\n",
    "        sub_iter, ses_iter = ppt.filter_workflow(bids_dir, sub_list, ses_list, exclude_list)\n",
    "        # Create BIDS nodes:\n",
    "        BIDSFolders = [\n",
    "            (\n",
    "                \"preprocessed/ses-%ssub-%s\" % (session, subject),\n",
    "                \"sub-%s/ses-%s/preprocessed\" % (subject, session),\n",
    "            )\n",
    "            for session in ses_list\n",
    "            for subject in sub_list\n",
    "        ]\n",
    "        # IdentityInterface for file input:\n",
    "        self.subject_source = Node(\n",
    "            IdentityInterface(fields=[\"subject_id\", \"session_id\"]),\n",
    "            iterables=[(\"subject_id\", sub_iter), (\"session_id\", ses_iter)],\n",
    "            synchronize=True,\n",
    "            name=\"SubjectDataSource\",\n",
    "        )\n",
    "        self.subject_source.inputs.ext = bids_ext\n",
    "\n",
    "        # reverse phase encoding design selection\n",
    "        # If only one DWI volume:\n",
    "        if rpe_design == \"-rpe_none\":\n",
    "            # Gradient files input:\n",
    "            self.sub_grad_files = Node(\n",
    "                Function(\n",
    "                    input_names=[\"sub_dwi\", \"ext\"],\n",
    "                    output_names=[\"fslgrad\"],\n",
    "                    function=ppt.get_sub_gradfiles,\n",
    "                    ext=bids_ext,\n",
    "                ),\n",
    "                name=\"SubjectGradientSource\",\n",
    "            )\n",
    "            self.sub_grad_files.inputs.ext = bids_ext\n",
    "            # Reorient DWI to standard orientation for FSL:\n",
    "            self.DWIReorient = Node(\n",
    "                fsl.utils.Reorient2Std(\n",
    "                    args = '-m mnitransformation.mat',\n",
    "                    out_file = 'dwi_reorient.nii.gz'\n",
    "                ),\n",
    "                name = \"DWIReorient\"\n",
    "            )\n",
    "            # Convert from nifti to mrtrix3's mif format:\n",
    "            self.mrconvert = Node(\n",
    "                ppt.Convert(\n",
    "                    out_file=\"raw_dwi.mif\",\n",
    "                    export_grad=\"raw_dwi.b\",\n",
    "                    nthreads=mrtrix_nthreads,\n",
    "                ),\n",
    "                name=\"Convert2Mif\",\n",
    "            )\n",
    "        # If there are two DWI images with second volume being reverse direction:\n",
    "        elif rpe_design == \"-rpe_all\":\n",
    "            # Forward direction gradient file\n",
    "            self.sub_grad_files1 = Node(\n",
    "                Function(\n",
    "                    input_names=[\"sub_dwi\", \"ext\"],\n",
    "                    output_names=[\"fslgrad\"],\n",
    "                    function=ppt.get_sub_gradfiles,\n",
    "                    ext=bids_ext,\n",
    "                ),\n",
    "                name=\"SubjectGradientForward\",\n",
    "            )\n",
    "            # Reverse direction gradient file\n",
    "            self.sub_grad_files2 = Node(\n",
    "                Function(\n",
    "                    input_names=[\"sub_dwi\", \"ext\"],\n",
    "                    output_names=[\"fslgrad\"],\n",
    "                    function=ppt.get_sub_gradfiles,\n",
    "                    ext=bids_ext,\n",
    "                ),\n",
    "                name=\"SubjectGradientReverse\",\n",
    "            )\n",
    "            self.sub_grad_files1.inputs.ext = bids_ext\n",
    "            self.sub_grad_files2.inputs.ext = bids_ext\n",
    "            # Forward and reverses direction reorientation to FSL standard\n",
    "            self.DWIReorientForward = Node(\n",
    "                fsl.utils.Reorient2Std(\n",
    "                    args = '-m forward_mni_transformation.mat',\n",
    "                    out_file = 'dwi_reorient_forward.nii.gz'\n",
    "                ),\n",
    "                name = \"DWIReorientForward\",\n",
    "            )\n",
    "            self.DWIReorientReverse = Node(\n",
    "                fsl.utils.Reorient2Std(\n",
    "                    args = '-m reverse_mni_transformation.mat',\n",
    "                    out_file = 'dwi_reorient_reverse.nii.gz'\n",
    "                ),\n",
    "                name = \"DWIReorientReverse\",\n",
    "            )\n",
    "            # Conversion from nifti to mif\n",
    "            self.mrconvert1 = Node(\n",
    "                ppt.Convert(\n",
    "                    out_file=\"raw_dwi1.mif\",\n",
    "                    export_grad=\"raw_dwi1.b\",\n",
    "                    nthreads=mrtrix_nthreads,\n",
    "                ),\n",
    "                name=\"Convert2MifForward\",\n",
    "            )\n",
    "            self.mrconvert2 = Node(\n",
    "                ppt.Convert(\n",
    "                    out_file=\"raw_dwi2.mif\",\n",
    "                    export_grad=\"raw_dwi2.b\",\n",
    "                    nthreads=mrtrix_nthreads,\n",
    "                ),\n",
    "                name=\"Convert2MifReverse\",\n",
    "            )\n",
    "            # concatenate the two images and their gradient files.\n",
    "            self.mrconcat = Node(ppt.MRCat(out_file=\"raw_dwi.mif\"), name=\"concat_dwi\",)\n",
    "            self.gradcat = Node(ppt.GradCat(out_file=\"raw_dwi.b\"), name=\"concat_grad\",)\n",
    "\n",
    "        self.select_files = Node(\n",
    "            SelectFiles(bids_path_template, base_directory=bids_dir),\n",
    "            name=\"SelectFiles\",\n",
    "        )\n",
    "\n",
    "        self.get_metadata = Node(\n",
    "            Function(\n",
    "                input_names=[\"path\", \"bids_dir\"],\n",
    "                output_names=[\"ReadoutTime\", \"PE_DIR\"],\n",
    "                function=ppt.BIDS_metadata,\n",
    "            ),\n",
    "            name=\"GetMetaData\",\n",
    "        )\n",
    "        self.get_metadata.inputs.bids_dir = bids_dir\n",
    "\n",
    "        self.createMask = Node(\n",
    "            BrainMask(out_file=\"b0_brain_mask.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"Dwi2Mask\",\n",
    "        )\n",
    "\n",
    "        self.GradCheck = Node(\n",
    "            ppt.GradCheck(export_grad=\"corrected.b\", nthreads=mrtrix_nthreads),\n",
    "            name=\"GradientCheck\",\n",
    "        )\n",
    "\n",
    "        self.NewGradMR = Node(\n",
    "            ppt.Convert(out_file=\"corrected_dwi.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"ConvertDWI2Mif\",\n",
    "        )\n",
    "\n",
    "        self.denoise = Node(\n",
    "            ppt.dwidenoise(\n",
    "                out_file=\"denoised.mif\", noise=\"noise_map.mif\", nthreads=mrtrix_nthreads\n",
    "            ),\n",
    "            name=\"Denoise\",\n",
    "        )\n",
    "\n",
    "        self.degibbs = Node(\n",
    "            MRDeGibbs(out_file=\"unring.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"RingingRemoval\",\n",
    "        )\n",
    "\n",
    "        self.fslpreproc = Node(\n",
    "            ppt.dwipreproc(\n",
    "                out_file=\"preproc.mif\",\n",
    "                rpe_options=rpe_design,\n",
    "                eddy_options='\"--slm=linear --repol \"',\n",
    "                nthreads=mrtrix_nthreads,\n",
    "                export_grad=\"eddy_dwi.b\",\n",
    "            ),\n",
    "            name=\"DWIFSLPreproc\",\n",
    "        )\n",
    "\n",
    "        self.GradUpdate = Node(\n",
    "            ppt.GradCheck(export_grad=\"tmp.b\"), name=\"AlteredGradient\"\n",
    "        )\n",
    "\n",
    "        self.ModGrad = Node(\n",
    "            ppt.MRInfo(export_grad=\"modified.b\"), name=\"ModifyGradient\"\n",
    "        )\n",
    "\n",
    "        self.UpdateMif = Node(ppt.Convert(), name=\"UpdateMif\")\n",
    "\n",
    "        self.NewMask = Node(BrainMask(), name=\"RecreateMask\")\n",
    "\n",
    "        self.biascorrect = Node(\n",
    "            ppt.BiasCorrect(\n",
    "                use_ants=True,\n",
    "                out_file=\"dwi_bias.mif\",\n",
    "                bias=\"biasfield.mif\",\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"BiasCorrection\",\n",
    "        )\n",
    "\n",
    "        self.grad_info = Node(\n",
    "            ppt.MRInfo(export_grad=\"rician_tmp.b\", nthreads=mrtrix_nthreads),\n",
    "            name=\"NewGradient\",\n",
    "        )\n",
    "\n",
    "        self.low_noise_map = Node(\n",
    "            ppt.CheckNIZ(out_file=\"lownoisemap.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"LowNoiseMap\",\n",
    "        )\n",
    "\n",
    "        self.rician_noise = Node(\n",
    "            ppt.RicianNoise(\n",
    "                power=2,\n",
    "                denoise=2,\n",
    "                out_file=\"rician_removed_dwi.mif\",\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"RicianNoise\",\n",
    "        )\n",
    "\n",
    "        self.check_rician = Node(\n",
    "            ppt.CheckNIZ(out_file=\"rician_tmp.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"NoiseComparison\",\n",
    "        )\n",
    "\n",
    "        self.convert_rician = Node(\n",
    "            ppt.Convert(out_file=\"rician_corrected_dwi.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"ConvertRician\",\n",
    "        )\n",
    "        self.dwi_mask = Node(BrainMask(out_file=\"dwi_mask.mif\"), name=\"dwi2mask\",)\n",
    "        self.fit_tensor = Node(FitTensor(out_file=\"dti.mif\"), name=\"dwi2tensor\",)\n",
    "        self.tensor_FA = Node(TensorMetrics(out_fa=\"fa.mif\"), name=\"tensor2metrics\",)\n",
    "        self.wm_mask = Node(\n",
    "            ppt.MRThreshold(opt_abs=0.5, out_file=\"wm.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"mrthreshold\",\n",
    "        )\n",
    "        self.norm_intensity = Node(\n",
    "            ppt.DWINormalize(\n",
    "                opt_intensity=1000,\n",
    "                out_file=\"dwi_norm_intensity.mif\",\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"DWINormalise\",\n",
    "        )\n",
    "        self.sub_b0extract = Node(\n",
    "            DWIExtract(bzero=True, out_file=\"b0_volume.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"ExtractB0Image\",\n",
    "        )\n",
    "        self.sub_b0mean = Node(\n",
    "            MRMath(\n",
    "                operation=\"mean\",\n",
    "                axis=3,\n",
    "                out_file=\"b0_dwi.mif\",\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"MeanB0Volume\",\n",
    "        )\n",
    "        self.sub_b0mask = Node(\n",
    "            BrainMask(out_file=\"dwi_norm_mask.mif\", nthreads=mrtrix_nthreads),\n",
    "            name=\"DWI2Mask\",\n",
    "        )\n",
    "        self.sub_convert_dwi = Node(\n",
    "            ppt.Convert(out_file=\"b0_dwi.nii.gz\"), name=\"DWI2Nifti\",\n",
    "        )\n",
    "        self.sub_convert_mask = Node(\n",
    "            ppt.Convert(out_file=\"dwi_norm_mask.nii.gz\"), name=\"Mask2Nifti\",\n",
    "        )\n",
    "        self.sub_apply_mask = Node(\n",
    "            fsl.ApplyMask(out_file=\"b0_dwi_brain.nii.gz\"), name=\"ApplyMask\",\n",
    "        )\n",
    "        self.mni_b0extract = Node(\n",
    "            DWIExtract(\n",
    "                bzero=True,\n",
    "                out_file=\"dwi_space-acpc_res-{}_b0.mif\".format(img_resol),\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"MNIExtractB0Volume\",\n",
    "        )\n",
    "        self.mni_b0mean = Node(\n",
    "            MRMath(\n",
    "                operation=\"mean\",\n",
    "                axis=3,\n",
    "                out_file=\"dwi_space-acpc_res-{}_b0mean.mif\".format(img_resol),\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"MNIB0MeanVolume\",\n",
    "        )\n",
    "        self.mni_b0mask = Node(\n",
    "            BrainMask(\n",
    "                out_file=\"dwi_space-acpc_res-{}_mask.mif\".format(img_resol),\n",
    "                nthreads=mrtrix_nthreads,\n",
    "            ),\n",
    "            name=\"MNIB0BrainMask\",\n",
    "        )\n",
    "        self.mni_convert_dwi = Node(\n",
    "            ppt.Convert(\n",
    "                out_file=\"dwi_space-acpc_res-{}_b0mean.nii.gz\".format(img_resol)\n",
    "            ),\n",
    "            name=\"MNIDWI2Nifti\",\n",
    "        )\n",
    "        self.mni_convert_mask = Node(\n",
    "            ppt.Convert(\n",
    "                out_file=\"dwi_space-acpc_res-{}_seg-brain_mask.nii.gz\".format(img_resol)\n",
    "            ),\n",
    "            name=\"MNIMask2Nifti\",\n",
    "        )\n",
    "        self.mni_apply_mask = Node(\n",
    "            fsl.ApplyMask(\n",
    "                out_file=\"dwi_space-acpc_res-{}_seg-brain.nii.gz\".format(img_resol)\n",
    "            ),\n",
    "            name=\"MNIApplyMask\",\n",
    "        )\n",
    "        self.mni_dwi = Node(\n",
    "            ppt.Convert(\n",
    "                out_file=\"dwi_space-acpc_res-{}.nii.gz\".format(img_resol),\n",
    "                export_grad=\"dwi_space-acpc_res-{}.b\".format(img_resol),\n",
    "                export_fslgrad=(\n",
    "                    \"dwi_space-acpc_res-{}.bvecs\".format(img_resol),\n",
    "                    \"dwi_space-acpc_res-{}.bvals\".format(img_resol),\n",
    "                ),\n",
    "                export_json=True,\n",
    "                nthreads=mrtrix_nthreads,\n",
    "                out_json=\"dwi_space-acpc_res-{}.json\".format(img_resol),\n",
    "            ),\n",
    "            name=\"MNI_Outputs\",\n",
    "        )\n",
    "\n",
    "        self.datasink = Node(\n",
    "            DataSink(\n",
    "                base_directory=os.path.join(bids_dir, \"derivatives\", \"pipetography\")\n",
    "            ),\n",
    "            name=\"DataSink\",\n",
    "        )\n",
    "        substitutions = [\n",
    "            (\"_subject_id_\", \"sub-\"),\n",
    "            (\"_session_id_\", \"ses-\"),\n",
    "            *BIDSFolders,\n",
    "        ]\n",
    "\n",
    "        self.datasink.inputs.substitutions = substitutions\n",
    "        print(\n",
    "            \"Data sink (output folder) is set to {}\".format(\n",
    "                os.path.join(bids_dir, \"derivatives\", \"pipetography\")\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating layout of data directory, might take a while if there are a lot of subjects\n",
      "Data sink (output folder) is set to ./testing/BIDS_dir/derivatives/pipetography\n",
      "Data sink (output folder) is set to ./testing/BIDS_dir/derivatives/pipetography\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "#test\n",
    "bidsdir = \"./testing/BIDS_dir\"\n",
    "sub_list, ses_list, layout = ppt.get_subs(bidsdir)\n",
    "sub_template = {\n",
    "    \"anat\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"anat\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_T1w.{ext}\",\n",
    "    ),\n",
    "    \"dwi\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_dwi.{ext}\",\n",
    "    ),\n",
    "    \"b_files\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_dwi.bv*\",\n",
    "    ),\n",
    "}\n",
    "test_nodes = PreProcNodes(\n",
    "    bids_dir=bidsdir,\n",
    "    bids_path_template=sub_template,\n",
    "    bids_ext=\".nii.gz\",\n",
    "    rpe_design=\"-rpe_none\",\n",
    "    mrtrix_nthreads=6,\n",
    "    img_resol='1mm',\n",
    "    sub_list=sub_list,\n",
    "    ses_list=ses_list,\n",
    ")\n",
    "\n",
    "# for rpe_all\n",
    "sub_template = {\n",
    "    \"anat\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"anat\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_T1w.{ext}\",\n",
    "    ),\n",
    "    \"dwi\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_dwi.{ext}\",\n",
    "    ),\n",
    "    \"bfiles\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_dwi.bv*\",\n",
    "    ),\n",
    "    \"rdwi\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_run-2_dwi.{ext}\",\n",
    "    ),\n",
    "    \"rbfiles\": os.path.join(\n",
    "        \"sub-{subject_id}\",\n",
    "        \"ses-{session_id}\",\n",
    "        \"dwi\",\n",
    "        \"sub-{subject_id}_ses-{session_id}_run-2_dwi.bv*\",\n",
    "    ),\n",
    "}\n",
    "test_nodes = PreProcNodes(\n",
    "    bids_dir=bidsdir,\n",
    "    bids_path_template=sub_template,\n",
    "    bids_ext=\".nii.gz\",\n",
    "    rpe_design=\"-rpe_all\",\n",
    "    mrtrix_nthreads=6,\n",
    "    img_resol='1mm',\n",
    "    sub_list=sub_list,\n",
    "    ses_list=ses_list,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ACPCNodes:\n",
    "    \"\"\"\n",
    "    T1 anatomy image related nodes. Mainly ACPC alignment of T1 and DWI and extraction of white matter mask.\n",
    "    Inputs:\n",
    "        - MNI_template: path to MNI template provided by FSL. By default uses the environment variable FSLDIR to locate the reference templates for ACPC alignment.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, MNI_template):\n",
    "        self.T1Reorient = Node(\n",
    "            fsl.utils.Reorient2Std(\n",
    "                args = '-m mnitransformation.mat',\n",
    "                out_file = 'T1_reoriented.nii.gz'\n",
    "            ),\n",
    "            name = 'T1Reorient'\n",
    "        )\n",
    "        self.reduceFOV = Node(\n",
    "            fsl.utils.RobustFOV(\n",
    "                out_transform=\"roi2full.mat\", out_roi=\"robustfov.nii.gz\"\n",
    "            ),\n",
    "            name=\"ReduceFOV\",\n",
    "        )\n",
    "        self.xfminverse = Node(\n",
    "            fsl.utils.ConvertXFM(out_file=\"full2roi.mat\", invert_xfm=True),\n",
    "            name=\"InverseTransformation\",\n",
    "        )\n",
    "        self.flirt = Node(\n",
    "            fsl.preprocess.FLIRT(\n",
    "                reference=MNI_template,\n",
    "                interp=\"spline\",\n",
    "                out_matrix_file=\"roi2std.mat\",\n",
    "                out_file=\"acpc_mni.nii.gz\",\n",
    "            ),\n",
    "            name=\"FLIRT\",\n",
    "        )\n",
    "        self.concatxfm = Node(\n",
    "            fsl.utils.ConvertXFM(concat_xfm=True, out_file=\"full2std.mat\"),\n",
    "            name=\"ConcatTransform\",\n",
    "        )\n",
    "\n",
    "        self.alignxfm = Node(\n",
    "            ppt.fslaff2rigid(out_file=\"outputmatrix\"), name=\"aff2rigid\",\n",
    "        )\n",
    "\n",
    "        self.ACPC_warp = Node(\n",
    "            fsl.preprocess.ApplyWarp(\n",
    "                out_file=\"T1w_space-acpc.nii.gz\",\n",
    "                relwarp=True,\n",
    "                output_type=\"NIFTI_GZ\",\n",
    "                interp=\"spline\",\n",
    "                ref_file=MNI_template,\n",
    "            ),\n",
    "            name=\"ACPCApplyWarp\",\n",
    "        )\n",
    "\n",
    "        self.t1_bet = Node(\n",
    "            fsl.preprocess.BET(mask=True, robust=True, out_file=\"acpc_t1_brain.nii.gz\"),\n",
    "            name=\"FSLBet\",\n",
    "        )\n",
    "\n",
    "        self.epi_reg = Node(fsl.epi.EpiReg(out_base=\"dwi2acpc\"), name=\"fsl_epireg\",)\n",
    "\n",
    "        self.acpc_xfm = Node(\n",
    "            ppt.TransConvert(flirt=True, out_file=\"dwi2acpc_xfm.mat\", force=True),\n",
    "            name=\"ConvertTransformation\",\n",
    "        )\n",
    "\n",
    "        self.apply_xfm = Node(\n",
    "            ppt.MRTransform(out_file=\"dwi_acpc.mif\"), name=\"MRTransform\",\n",
    "        )\n",
    "\n",
    "        self.regrid = Node(\n",
    "            ppt.MRRegrid(\n",
    "                out_file=\"dwi_space-acpc_res-1mm.mif\",\n",
    "                regrid=MNI_template,\n",
    "            ),\n",
    "            name=\"Regrid\",\n",
    "        )\n",
    "\n",
    "        self.gen_5tt = Node(\n",
    "            ppt.Make5ttFSL(premasked=True, out_file=\"T1w_space-acpc_seg-5tt.mif\"),\n",
    "            name=\"Mrtrix5TTGen\",\n",
    "        )\n",
    "        self.gmwmi = Node(ppt.gmwmi(out_file=\"gmwmi.nii.gz\"), name=\"5tt2gmwmi\")\n",
    "        self.binarize_gmwmi = Node(\n",
    "            ppt.MRThreshold(\n",
    "                opt_abs=0.05, out_file=\"T1w_space-acpc_seg-gmwmi_mask.nii.gz\"\n",
    "            ),\n",
    "            name=\"GMWMI\",\n",
    "        )\n",
    "        self.convert2wm = Node(\n",
    "            ppt.Convert(\n",
    "                coord=[3, 2],\n",
    "                axes=[0, 1, 2],\n",
    "                out_file=\"T1w_space-acpc_seg-wm_mask.nii.gz\",\n",
    "            ),\n",
    "            name=\"GetWMMask\",\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "#test\n",
    "test_node = ACPCNodes(\n",
    "    MNI_template=os.path.expandvars(\"$FSLDIR/data/standard/MNI152_T1_1mm.nii.gz\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Post-processing Nodes\n",
    "\n",
    "Nodes for fetching streamlines, MRI, atlases and combine into connectomes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class PostProcNodes:\n",
    "    \"\"\"\n",
    "    Inputs:\n",
    "        BIDS_dir (str): Path to BIDS directory\n",
    "        subj_template (dict): template directory for tck, dwi, T1, mask files\n",
    "        sub_list (list): subjects IDs list generated from BIDS layout\n",
    "        ses_list (list): sessions IDs listgenerated from BIDS layout\n",
    "        skip_tuples (tuple): [('subject', 'session')] string pair to skip\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, BIDS_dir, subj_template, sub_list, ses_list, skip_tuples):\n",
    "        # filter & create sub-graphs for subjects and session combos:\n",
    "        sub_iter, ses_iter = ppt.filter_workflow(BIDS_dir, sub_list, ses_list, skip_tuples)\n",
    "        \n",
    "        # Create BIDS output folder list for datasink\n",
    "        BIDSFolders = [\n",
    "            (\n",
    "                \"connectomes/ses-%ssub-%s\" % (session, subject),\n",
    "                \"sub-%s/ses-%s/connectomes\" % (subject, session),\n",
    "            )\n",
    "            for session in ses_list\n",
    "            for subject in sub_list\n",
    "        ]\n",
    "        # BIDS derivatives directory containing preprocessed outputs and streamline outputs\n",
    "        preproc_dir = os.path.join(\n",
    "            BIDS_dir, \"derivatives\", \"pipetography\"\n",
    "        )  \n",
    "\n",
    "        # DWI input:\n",
    "        self.subject_source = Node(\n",
    "            IdentityInterface(fields=[\"subject_id\", \"session_id\"]),\n",
    "            iterables=[(\"subject_id\", sub_iter), (\"session_id\", ses_iter)],\n",
    "            synchronize=True,\n",
    "            name=\"subj_source\",\n",
    "        )\n",
    "        self.select_files = Node(\n",
    "            SelectFiles(subj_template), base_directory=BIDS_dir, name=\"select_subjects\"\n",
    "        )\n",
    "        self.linear_reg = Node(\n",
    "            ants.Registration(\n",
    "                output_transform_prefix=\"atlas_in_dwi_affine\",\n",
    "                dimension=3,\n",
    "                collapse_output_transforms=True,\n",
    "                transforms=[\"Affine\"],\n",
    "                transform_parameters=[(0.1,)],\n",
    "                metric=[\"MI\"],\n",
    "                metric_weight=[1],\n",
    "                radius_or_number_of_bins=[64],\n",
    "                number_of_iterations=[[500, 200, 200, 100]],\n",
    "                convergence_threshold=[1e-6],\n",
    "                convergence_window_size=[10],\n",
    "                smoothing_sigmas=[[4, 2, 1, 0]],\n",
    "                sigma_units=[\"vox\"],\n",
    "                shrink_factors=[[8, 4, 2, 1]],\n",
    "                use_histogram_matching=[True],\n",
    "                output_warped_image=\"atlas_in_dwi_affine.nii.gz\",\n",
    "                interpolation=\"GenericLabel\",\n",
    "            ),\n",
    "            name=\"LinearRegistration\",\n",
    "        )\n",
    "        self.nonlinear_reg = Node(\n",
    "            ants.Registration(\n",
    "                output_transform_prefix=\"atlas_in_dwi_syn\",\n",
    "                dimension=3,\n",
    "                collapse_output_transforms=True,\n",
    "                transforms=[\"SyN\"],\n",
    "                transform_parameters=[(0.1,)],\n",
    "                metric=[\"MI\"],\n",
    "                metric_weight=[1],\n",
    "                radius_or_number_of_bins=[64],\n",
    "                number_of_iterations=[[500, 200, 200, 100]],\n",
    "                convergence_threshold=[1e-06],\n",
    "                convergence_window_size=[10],\n",
    "                smoothing_sigmas=[[4, 2, 1, 0]],\n",
    "                sigma_units=[\"vox\"],\n",
    "                shrink_factors=[[8, 4, 2, 1]],\n",
    "                use_histogram_matching=[True],\n",
    "                output_warped_image=\"atlas_in_dwi_syn.nii.gz\",\n",
    "                interpolation=\"GenericLabel\",\n",
    "            ),\n",
    "            name=\"NonLinearRegistration\",\n",
    "        )\n",
    "        self.response = Node(\n",
    "            ResponseSD(\n",
    "                algorithm=\"dhollander\",\n",
    "                wm_file=\"wm.txt\",\n",
    "                gm_file=\"gm.txt\",\n",
    "                csf_file=\"csf.txt\",\n",
    "            ),\n",
    "            name=\"SDResponse\",\n",
    "        )\n",
    "        self.fod = Node(\n",
    "            ConstrainedSphericalDeconvolution(\n",
    "                algorithm=\"msmt_csd\",\n",
    "                wm_txt=\"wm.txt\",\n",
    "                gm_txt=\"gm.txt\",\n",
    "                gm_odf=\"gm.mif\",\n",
    "                csf_txt=\"csf.txt\",\n",
    "                csf_odf=\"csf.mif\",\n",
    "            ),\n",
    "            name=\"dwiFOD\",\n",
    "        )\n",
    "        self.sift2 = Node(\n",
    "            ppt.tckSIFT2(\n",
    "                out_file=\"sift2.txt\"\n",
    "            ),\n",
    "            name=\"SIFT2\",\n",
    "        )\n",
    "        self.connectome = Node(\n",
    "            ppt.MakeConnectome(\n",
    "                out_file=\"connectome.csv\", symmetric=True, zero_diag=True\n",
    "            ),\n",
    "            name=\"WeightConnectome\",\n",
    "        )\n",
    "        self.distance = Node(\n",
    "            ppt.MakeConnectome(\n",
    "                scale_length=True,\n",
    "                stat_edge=\"mean\",\n",
    "                symmetric=True,\n",
    "                zero_diag=True,\n",
    "                out_file=\"distances.csv\",\n",
    "            ),\n",
    "            name=\"WeightDistance\",\n",
    "        )\n",
    "        self.datasink = Node(DataSink(base_directory=preproc_dir), name=\"datasink\")\n",
    "        substitutions = [(\"_subject_id_\", \"sub-\"), (\"_session_id_\", \"ses-\")]\n",
    "        substitutions.extend(BIDSFolders)\n",
    "        self.datasink.inputs.substitutions = substitutions\n",
    "        self.datasink.inputs.regexp_substitutions = [\n",
    "            (r\"(_moving_image_.*\\.\\.)\", \"\"),\n",
    "            (r\"(\\.nii|\\.gz)\", \"\"),\n",
    "        ]\n",
    "        print(\"Data sink (output folder) is set to {}\".format(preproc_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data sink (output folder) is set to ./testing/BIDS_dir/derivatives/pipetography\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "#test\n",
    "data_dir = \"./testing/BIDS_dir\"\n",
    "\n",
    "subj_template = {\n",
    "    \"tck\": os.path.join(\n",
    "        data_dir,\n",
    "        \"cuda_tracking\",\n",
    "        \"_session_id_{session_id}_subject_id_{subject_id}\",\n",
    "        \"sub-{subject_id}_ses-{session_id}.tck\",\n",
    "    ),\n",
    "    \"dwi\": os.path.join(\n",
    "        data_dir,\n",
    "        \"derivatives\",\n",
    "        \"dwi_acpc_aligned_1mm\",\n",
    "        \"_session_id_{session_id}_subject_id_{subject_id}\",\n",
    "        \"dwi_acpc_1mm.mif\",\n",
    "    ),\n",
    "    \"T1A\": os.path.join(\n",
    "        data_dir,\n",
    "        \"derivatives\",\n",
    "        \"t1_acpc_aligned\",\n",
    "        \"_session_id_{session_id}_subject_id_{subject_id}\",\n",
    "        \"acpc_t1.nii\",\n",
    "    ),\n",
    "    \"mask\": os.path.join(\n",
    "        data_dir,\n",
    "        \"derivatives\",\n",
    "        \"t1_acpc_aligned\",\n",
    "        \"_session_id_{session_id}_subject_id_{subject_id}\",\n",
    "        \"dwi_acpc_1mm_mask.nii.gz\",\n",
    "    ),\n",
    "}\n",
    "\n",
    "test_node = PostProcNodes(\n",
    "    BIDS_dir=data_dir, subj_template=subj_template, sub_list=['11042', '11045'], ses_list=['01', '02'], skip_tuples=[()]\n",
    ")\n",
    "node_names = list(\n",
    "    [\n",
    "        \"subj_source\",\n",
    "        \"select_subjects\",\n",
    "        \"linear_registration\",\n",
    "        \"nonlinear_registration\",\n",
    "        \"SDResponse\",\n",
    "        \"dwiFOD\",\n",
    "        \"sift2_filtering\",\n",
    "        \"weight_connectome\",\n",
    "        \"weight_distance\",\n",
    "        \"datasink\",\n",
    "    ]\n",
    ")\n",
    "for key, value in test_node.__dict__.items():\n",
    "    assert str(value) in node_names"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}