BespostX using Nipype

bedpostx.py

import argparse
import os

from nipype.pipeline.engine import Workflow, Node, MapNode
from nipype.interfaces.utility import IdentityInterface
from nipype.interfaces.fsl.utils import ImageMaths
from nipype.interfaces.fsl.utils import Split
from nipype.interfaces.fsl.dti import XFibres
from nipype.interfaces.fsl.utils import Merge
from nipype.interfaces.fsl.dti import MakeDyadicVectors
from nipype.interfaces.io import DataSink

# Functions
def transpose(samples_over_fibres):
    import numpy as np
    a = np.array(samples_over_fibres)
    if len(a.shape) == 1:
        a = a.reshape(-1, 1)
    return a.T.tolist()

def run_workflow(dwi, bvecs, bvals, mask, outputdir, outputname):
    # Workflow
    trac_bedp = Workflow("trac_bedp")
    # Node: trac_bedp.inputnode
    trac_bedp_inputnode = Node(IdentityInterface(fields=['dwi', 'mask', 'bvecs', 'bvals'], mandatory_inputs=True), name="trac_bedp_inputnode")
    trac_bedp_inputnode.inputs.bvals = os.path.abspath(bvals)
    trac_bedp_inputnode.inputs.bvecs = os.path.abspath(bvecs)
    trac_bedp_inputnode.inputs.dwi = os.path.abspath(dwi)
    trac_bedp_inputnode.inputs.mask = os.path.abspath(mask)
    # Node: trac_bedp.preproc.inputnode
    trac_bedp_preproc_inputnode = Node(IdentityInterface(fields=['dwi', 'mask'], mandatory_inputs=True), name="trac_bedp_preproc_inputnode")
    trac_bedp.connect(trac_bedp_inputnode, "mask", trac_bedp_preproc_inputnode, "mask")
    trac_bedp.connect(trac_bedp_inputnode, "dwi", trac_bedp_preproc_inputnode, "dwi")
    # Node: trac_bedp.preproc.mask_dwi
    trac_bedp_preproc_mask_dwi = Node(ImageMaths(), name="trac_bedp_preproc_mask_dwi")
    trac_bedp_preproc_mask_dwi.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_preproc_mask_dwi.inputs.ignore_exception = False
    trac_bedp_preproc_mask_dwi.inputs.op_string = '-mas'
    trac_bedp_preproc_mask_dwi.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_preproc_mask_dwi.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_preproc_inputnode, "dwi", trac_bedp_preproc_mask_dwi, "in_file")
    trac_bedp.connect(trac_bedp_preproc_inputnode, "mask", trac_bedp_preproc_mask_dwi, "in_file2")
    # Node: trac_bedp.preproc.slice_dwi
    trac_bedp_preproc_slice_dwi = Node(Split(), name="trac_bedp_preproc_slice_dwi")
    trac_bedp_preproc_slice_dwi.inputs.dimension = 'z'
    trac_bedp_preproc_slice_dwi.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_preproc_slice_dwi.inputs.ignore_exception = False
    trac_bedp_preproc_slice_dwi.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_preproc_slice_dwi.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_preproc_mask_dwi, "out_file", trac_bedp_preproc_slice_dwi, "in_file")
    # Node: trac_bedp.preproc.slice_mask
    trac_bedp_preproc_slice_mask = Node(Split(), name="trac_bedp_preproc_slice_mask")
    trac_bedp_preproc_slice_mask.inputs.dimension = 'z'
    trac_bedp_preproc_slice_mask.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_preproc_slice_mask.inputs.ignore_exception = False
    trac_bedp_preproc_slice_mask.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_preproc_slice_mask.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_preproc_inputnode, "mask", trac_bedp_preproc_slice_mask, "in_file")
    # Node: trac_bedp.xfibres
    trac_bedp_xfibres = MapNode(XFibres(), iterfield=['dwi', 'mask'], name="trac_bedp_xfibres")
    trac_bedp_xfibres.inputs.burn_in = 1000
    trac_bedp_xfibres.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_xfibres.inputs.force_dir = True
    trac_bedp_xfibres.inputs.fudge = 1
    trac_bedp_xfibres.inputs.ignore_exception = False
    trac_bedp_xfibres.inputs.logdir = 'logdir'
    trac_bedp_xfibres.inputs.model = 1
    trac_bedp_xfibres.inputs.n_fibres = 2
    trac_bedp_xfibres.inputs.n_jumps = 1250
    trac_bedp_xfibres.inputs.non_linear = True
    trac_bedp_xfibres.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_xfibres.inputs.sample_every = 25
    trac_bedp_xfibres.inputs.terminal_output = 'stream'
    trac_bedp_xfibres.inputs.update_proposal_every = 24
    trac_bedp.connect(trac_bedp_inputnode, "bvals", trac_bedp_xfibres, "bvals")
    trac_bedp.connect(trac_bedp_inputnode, "bvecs", trac_bedp_xfibres, "bvecs")
    trac_bedp.connect(trac_bedp_preproc_slice_mask, "out_files", trac_bedp_xfibres, "mask")
    trac_bedp.connect(trac_bedp_preproc_slice_dwi, "out_files", trac_bedp_xfibres, "dwi")
    # Node: trac_bedp.postproc.inputnode
    trac_bedp_postproc_inputnode = Node(IdentityInterface(fields=['thsamples', 'phsamples', 'fsamples', 'dyads', 'mean_dsamples', 'mask'], mandatory_inputs=True), name="trac_bedp_postproc_inputnode")
    trac_bedp.connect(trac_bedp_inputnode, "mask", trac_bedp_postproc_inputnode, "mask")
    trac_bedp.connect(trac_bedp_xfibres, "thsamples", trac_bedp_postproc_inputnode, "thsamples")
    trac_bedp.connect(trac_bedp_xfibres, "phsamples", trac_bedp_postproc_inputnode, "phsamples")
    trac_bedp.connect(trac_bedp_xfibres, "fsamples", trac_bedp_postproc_inputnode, "fsamples")
    trac_bedp.connect(trac_bedp_xfibres, "dyads", trac_bedp_postproc_inputnode, "dyads")
    trac_bedp.connect(trac_bedp_xfibres, "mean_dsamples", trac_bedp_postproc_inputnode, "mean_dsamples")
    # Node: trac_bedp.postproc.merge_mean_dsamples
    trac_bedp_postproc_merge_mean_dsamples = Node(Merge(), name="trac_bedp_postproc_merge_mean_dsamples")
    trac_bedp_postproc_merge_mean_dsamples.inputs.dimension = 'z'
    trac_bedp_postproc_merge_mean_dsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_merge_mean_dsamples.inputs.ignore_exception = False
    trac_bedp_postproc_merge_mean_dsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_merge_mean_dsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_inputnode, "mean_dsamples", trac_bedp_postproc_merge_mean_dsamples, "in_files")
    # Node: trac_bedp.postproc.merge_phsamples
    trac_bedp_postproc_merge_phsamples = MapNode(Merge(), iterfield=['in_files'], name="trac_bedp_postproc_merge_phsamples")
    trac_bedp_postproc_merge_phsamples.inputs.dimension = 'z'
    trac_bedp_postproc_merge_phsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_merge_phsamples.inputs.ignore_exception = False
    trac_bedp_postproc_merge_phsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_merge_phsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_inputnode, ('phsamples', transpose), trac_bedp_postproc_merge_phsamples, "in_files")
    # Node: trac_bedp.postproc.mean_phsamples
    trac_bedp_postproc_mean_phsamples = MapNode(ImageMaths(), iterfield=['in_file'], name="trac_bedp_postproc_mean_phsamples")
    trac_bedp_postproc_mean_phsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_mean_phsamples.inputs.ignore_exception = False
    trac_bedp_postproc_mean_phsamples.inputs.op_string = '-Tmean'
    trac_bedp_postproc_mean_phsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_mean_phsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_merge_phsamples, "merged_file", trac_bedp_postproc_mean_phsamples, "in_file")
    # Node: trac_bedp.postproc.merge_fsamples
    trac_bedp_postproc_merge_fsamples = MapNode(Merge(), iterfield=['in_files'], name="trac_bedp_postproc_merge_fsamples")
    trac_bedp_postproc_merge_fsamples.inputs.dimension = 'z'
    trac_bedp_postproc_merge_fsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_merge_fsamples.inputs.ignore_exception = False
    trac_bedp_postproc_merge_fsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_merge_fsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_inputnode, ('fsamples', transpose), trac_bedp_postproc_merge_fsamples, "in_files")
    # Node: trac_bedp.postproc.mean_fsamples
    trac_bedp_postproc_mean_fsamples = MapNode(ImageMaths(), iterfield=['in_file'], name="trac_bedp_postproc_mean_fsamples")
    trac_bedp_postproc_mean_fsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_mean_fsamples.inputs.ignore_exception = False
    trac_bedp_postproc_mean_fsamples.inputs.op_string = '-Tmean'
    trac_bedp_postproc_mean_fsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_mean_fsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_merge_fsamples, "merged_file", trac_bedp_postproc_mean_fsamples, "in_file")
    # Node: trac_bedp.postproc.merge_thsamples
    trac_bedp_postproc_merge_thsamples = MapNode(Merge(), iterfield=['in_files'], name="trac_bedp_postproc_merge_thsamples")
    trac_bedp_postproc_merge_thsamples.inputs.dimension = 'z'
    trac_bedp_postproc_merge_thsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_merge_thsamples.inputs.ignore_exception = False
    trac_bedp_postproc_merge_thsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_merge_thsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_inputnode, ('thsamples', transpose), trac_bedp_postproc_merge_thsamples, "in_files")
    # Node: trac_bedp.postproc.make_dyads
    trac_bedp_postproc_make_dyads = MapNode(MakeDyadicVectors(), iterfield=['theta_vol', 'phi_vol'], name="trac_bedp_postproc_make_dyads")
    trac_bedp_postproc_make_dyads.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_make_dyads.inputs.ignore_exception = False
    trac_bedp_postproc_make_dyads.inputs.output = 'dyads'
    trac_bedp_postproc_make_dyads.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_make_dyads.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_inputnode, "mask", trac_bedp_postproc_make_dyads, "mask")
    trac_bedp.connect(trac_bedp_postproc_merge_thsamples, "merged_file", trac_bedp_postproc_make_dyads, "theta_vol")
    trac_bedp.connect(trac_bedp_postproc_merge_phsamples, "merged_file", trac_bedp_postproc_make_dyads, "phi_vol")
    # Node: trac_bedp.postproc.mean_thsamples
    trac_bedp_postproc_mean_thsamples = MapNode(ImageMaths(), iterfield=['in_file'], name="trac_bedp_postproc_mean_thsamples")
    trac_bedp_postproc_mean_thsamples.inputs.environ = {'FSLOUTPUTTYPE': 'NIFTI_GZ'}
    trac_bedp_postproc_mean_thsamples.inputs.ignore_exception = False
    trac_bedp_postproc_mean_thsamples.inputs.op_string = '-Tmean'
    trac_bedp_postproc_mean_thsamples.inputs.output_type = 'NIFTI_GZ'
    trac_bedp_postproc_mean_thsamples.inputs.terminal_output = 'stream'
    trac_bedp.connect(trac_bedp_postproc_merge_thsamples, "merged_file", trac_bedp_postproc_mean_thsamples, "in_file")
    # Node: trac_bedp.outputnode
    trac_bedp_outputnode = Node(IdentityInterface(fields=['thsamples', 'phsamples', 'fsamples', 'mean_thsamples', 'mean_phsamples', 'mean_fsamples', 'dyads', 'dyads_dispersion'], mandatory_inputs=True), name="trac_bedp_outputnode")
    trac_bedp.connect(trac_bedp_postproc_merge_thsamples, "merged_file", trac_bedp_outputnode, "thsamples")
    trac_bedp.connect(trac_bedp_postproc_mean_fsamples, "out_file", trac_bedp_outputnode, "mean_fsamples")
    trac_bedp.connect(trac_bedp_postproc_mean_thsamples, "out_file", trac_bedp_outputnode, "mean_thsamples")
    trac_bedp.connect(trac_bedp_postproc_merge_phsamples, "merged_file", trac_bedp_outputnode, "phsamples")
    trac_bedp.connect(trac_bedp_postproc_merge_fsamples, "merged_file", trac_bedp_outputnode, "fsamples")
    trac_bedp.connect(trac_bedp_postproc_make_dyads, "dyads", trac_bedp_outputnode, "dyads")
    trac_bedp.connect(trac_bedp_postproc_make_dyads, "dispersion", trac_bedp_outputnode, "dyads_dispersion")
    trac_bedp.connect(trac_bedp_postproc_mean_phsamples, "out_file", trac_bedp_outputnode, "mean_phsamples")
    # Node: trac_bedp.sinker
    trac_bedp_sinker = Node(DataSink(infields=None), name="trac_bedp_sinker")
    trac_bedp_sinker.inputs._outputs = {}
    trac_bedp_sinker.inputs.base_directory = os.path.abspath(outputdir)
    trac_bedp_sinker.inputs.container = outputname
    trac_bedp_sinker.inputs.ignore_exception = False
    trac_bedp_sinker.inputs.parameterization = True
    trac_bedp_sinker.inputs.regexp_substitutions = [('_mean_.+samples./', 'mean_'),
                                                    ('_merged_maths', ''),
                                                    ('_merge_.+samples./', 'merged_'),
                                                    ('_merged', ''),
                                                    ('_trac_bedp_postproc', '')]
    trac_bedp_sinker.inputs.remove_dest_dir = False
    trac_bedp_sinker.inputs.substitutions = [('/dyads', ''), ('_make_dyads0', 'dyads1'),
                                             ('_make_dyads1', 'dyads2')]
    trac_bedp.connect(trac_bedp_outputnode, "dyads", trac_bedp_sinker, "@dyads")
    trac_bedp.connect(trac_bedp_outputnode, "dyads_dispersion", trac_bedp_sinker, "@disp")
    trac_bedp.connect(trac_bedp_outputnode, "fsamples", trac_bedp_sinker, "@fsamples")
    trac_bedp.connect(trac_bedp_outputnode, "mean_fsamples", trac_bedp_sinker, "@mean_fsamples")
    trac_bedp.connect(trac_bedp_outputnode, "mean_phsamples", trac_bedp_sinker, "@mean_phsamples")
    trac_bedp.connect(trac_bedp_outputnode, "mean_thsamples", trac_bedp_sinker, "@mean_thsamples")
    trac_bedp.connect(trac_bedp_outputnode, "phsamples", trac_bedp_sinker, "@phsamples")
    trac_bedp.connect(trac_bedp_outputnode, "thsamples", trac_bedp_sinker, "@thsamples")

    trac_bedp.base_dir = os.getcwd()
    trac_bedp.config['execution']['remove_unnecessary_outputs'] = False
    trac_bedp.run(plugin='PBS', plugin_args={'qsub_args': '-q many'})


from argparse import ArgumentParser
if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("--dwi", dest="dwi", help="dwi file", required=True)
     parser.add_argument("--bvecs", dest="bvecs", help="bvec file", required=True)
     parser.add_argument("--bvals", dest="bvals", help="bvals file", required=True)
     parser.add_argument("--mask", dest="mask", help="mask file", required=True)
     parser.add_argument("-o", "--outputdir", dest="outputdir", help="dwi file",
                         default=os.getcwd())
     parser.add_argument("--subject", dest="subject", help="subject id", required=True)
     args = parser.parse_args()

     run_workflow(args.dwi, args.bvecs, args.bvals, args.mask,
                  args.outputdir, args.subject)