utils.py

# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
# -*- coding: utf-8 -*-

import os.path as op

from ..base import (
    CommandLineInputSpec,
    CommandLine,
    traits,
    TraitedSpec,
    File,
    InputMultiPath,
    isdefined,
)
from .base import MRTrix3BaseInputSpec, MRTrix3Base


class BrainMaskInputSpec(MRTrix3BaseInputSpec):
    in_file = File(
        exists=True,
        argstr="%s",
        mandatory=True,
        position=-2,
        desc="input diffusion weighted images",
    )
    out_file = File(
        "brainmask.mif",
        argstr="%s",
        mandatory=True,
        position=-1,
        usedefault=True,
        desc="output brain mask",
    )


class BrainMaskOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="the output response file")


class BrainMask(CommandLine):
    """
    Convert a mesh surface to a partial volume estimation image


    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> bmsk = mrt.BrainMask()
    >>> bmsk.inputs.in_file = 'dwi.mif'
    >>> bmsk.cmdline                               # doctest: +ELLIPSIS
    'dwi2mask dwi.mif brainmask.mif'
    >>> bmsk.run()                                 # doctest: +SKIP
    """

    _cmd = "dwi2mask"
    input_spec = BrainMaskInputSpec
    output_spec = BrainMaskOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class Mesh2PVEInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-3, desc="input mesh"
    )
    reference = File(
        exists=True,
        argstr="%s",
        mandatory=True,
        position=-2,
        desc="input reference image",
    )
    in_first = File(
        exists=True,
        argstr="-first %s",
        desc="indicates that the mesh file is provided by FSL FIRST",
    )

    out_file = File(
        "mesh2volume.nii.gz",
        argstr="%s",
        mandatory=True,
        position=-1,
        usedefault=True,
        desc="output file containing SH coefficients",
    )


class Mesh2PVEOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="the output response file")


class Mesh2PVE(CommandLine):
    """
    Convert a mesh surface to a partial volume estimation image


    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> m2p = mrt.Mesh2PVE()
    >>> m2p.inputs.in_file = 'surf1.vtk'
    >>> m2p.inputs.reference = 'dwi.mif'
    >>> m2p.inputs.in_first = 'T1.nii.gz'
    >>> m2p.cmdline                               # doctest: +ELLIPSIS
    'mesh2pve -first T1.nii.gz surf1.vtk dwi.mif mesh2volume.nii.gz'
    >>> m2p.run()                                 # doctest: +SKIP
    """

    _cmd = "mesh2pve"
    input_spec = Mesh2PVEInputSpec
    output_spec = Mesh2PVEOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class Generate5ttInputSpec(MRTrix3BaseInputSpec):
    algorithm = traits.Enum(
        "fsl",
        "gif",
        "freesurfer",
        argstr="%s",
        position=-3,
        mandatory=True,
        desc="tissue segmentation algorithm",
    )
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-2, desc="input image"
    )
    out_file = File(argstr="%s", mandatory=True, position=-1, desc="output image")


class Generate5ttOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="output image")


class Generate5tt(MRTrix3Base):
    """
    Generate a 5TT image suitable for ACT using the selected algorithm


    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> gen5tt = mrt.Generate5tt()
    >>> gen5tt.inputs.in_file = 'T1.nii.gz'
    >>> gen5tt.inputs.algorithm = 'fsl'
    >>> gen5tt.inputs.out_file = '5tt.mif'
    >>> gen5tt.cmdline                             # doctest: +ELLIPSIS
    '5ttgen fsl T1.nii.gz 5tt.mif'
    >>> gen5tt.run()                               # doctest: +SKIP
    """

    _cmd = "5ttgen"
    input_spec = Generate5ttInputSpec
    output_spec = Generate5ttOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class TensorMetricsInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-1, desc="input DTI image"
    )

    out_fa = File(argstr="-fa %s", desc="output FA file")
    out_adc = File(argstr="-adc %s", desc="output ADC file")
    out_evec = File(argstr="-vector %s", desc="output selected eigenvector(s) file")
    out_eval = File(argstr="-value %s", desc="output selected eigenvalue(s) file")
    component = traits.List(
        [1],
        usedefault=True,
        argstr="-num %s",
        sep=",",
        desc=(
            "specify the desired eigenvalue/eigenvector(s). Note that "
            "several eigenvalues can be specified as a number sequence"
        ),
    )
    in_mask = File(
        exists=True,
        argstr="-mask %s",
        desc=(
            "only perform computation within the specified binary" " brain mask image"
        ),
    )
    modulate = traits.Enum(
        "FA",
        "none",
        "eval",
        argstr="-modulate %s",
        desc=("how to modulate the magnitude of the" " eigenvectors"),
    )


class TensorMetricsOutputSpec(TraitedSpec):
    out_fa = File(desc="output FA file")
    out_adc = File(desc="output ADC file")
    out_evec = File(desc="output selected eigenvector(s) file")
    out_eval = File(desc="output selected eigenvalue(s) file")


class TensorMetrics(CommandLine):
    """
    Compute metrics from tensors


    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> comp = mrt.TensorMetrics()
    >>> comp.inputs.in_file = 'dti.mif'
    >>> comp.inputs.out_fa = 'fa.mif'
    >>> comp.cmdline                               # doctest: +ELLIPSIS
    'tensor2metric -num 1 -fa fa.mif dti.mif'
    >>> comp.run()                                 # doctest: +SKIP
    """

    _cmd = "tensor2metric"
    input_spec = TensorMetricsInputSpec
    output_spec = TensorMetricsOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()

        for k in list(outputs.keys()):
            if isdefined(getattr(self.inputs, k)):
                outputs[k] = op.abspath(getattr(self.inputs, k))

        return outputs


class ComputeTDIInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-2, desc="input tractography"
    )
    out_file = File(
        "tdi.mif", argstr="%s", usedefault=True, position=-1, desc="output TDI file"
    )
    reference = File(
        exists=True,
        argstr="-template %s",
        desc="a reference" "image to be used as template",
    )
    vox_size = traits.List(
        traits.Int, argstr="-vox %s", sep=",", desc="voxel dimensions"
    )
    data_type = traits.Enum(
        "float",
        "unsigned int",
        argstr="-datatype %s",
        desc="specify output image data type",
    )
    use_dec = traits.Bool(argstr="-dec", desc="perform mapping in DEC space")
    dixel = File(
        argstr="-dixel %s",
        desc="map streamlines to"
        "dixels within each voxel. Directions are stored as"
        "azimuth elevation pairs.",
    )
    max_tod = traits.Int(
        argstr="-tod %d",
        desc="generate a Track Orientation " "Distribution (TOD) in each voxel.",
    )

    contrast = traits.Enum(
        "tdi",
        "length",
        "invlength",
        "scalar_map",
        "scalar_map_conut",
        "fod_amp",
        "curvature",
        argstr="-constrast %s",
        desc="define the desired " "form of contrast for the output image",
    )
    in_map = File(
        exists=True,
        argstr="-image %s",
        desc="provide the"
        "scalar image map for generating images with "
        "'scalar_map' contrasts, or the SHs image for fod_amp",
    )

    stat_vox = traits.Enum(
        "sum",
        "min",
        "mean",
        "max",
        argstr="-stat_vox %s",
        desc="define the statistic for choosing the final"
        "voxel intesities for a given contrast",
    )
    stat_tck = traits.Enum(
        "mean",
        "sum",
        "min",
        "max",
        "median",
        "mean_nonzero",
        "gaussian",
        "ends_min",
        "ends_mean",
        "ends_max",
        "ends_prod",
        argstr="-stat_tck %s",
        desc="define the statistic for choosing "
        "the contribution to be made by each streamline as a function of"
        " the samples taken along their lengths.",
    )

    fwhm_tck = traits.Float(
        argstr="-fwhm_tck %f",
        desc="define the statistic for choosing the"
        " contribution to be made by each streamline as a function of the "
        "samples taken along their lengths",
    )

    map_zero = traits.Bool(
        argstr="-map_zero",
        desc="if a streamline has zero contribution based "
        "on the contrast & statistic, typically it is not mapped; use this "
        "option to still contribute to the map even if this is the case "
        "(these non-contributing voxels can then influence the mean value in "
        "each voxel of the map)",
    )

    upsample = traits.Int(
        argstr="-upsample %d",
        desc="upsample the tracks by"
        " some ratio using Hermite interpolation before "
        "mappping",
    )

    precise = traits.Bool(
        argstr="-precise",
        desc="use a more precise streamline mapping "
        "strategy, that accurately quantifies the length through each voxel "
        "(these lengths are then taken into account during TWI calculation)",
    )
    ends_only = traits.Bool(
        argstr="-ends_only", desc="only map the streamline" " endpoints to the image"
    )

    tck_weights = File(
        exists=True,
        argstr="-tck_weights_in %s",
        desc="specify" " a text scalar file containing the streamline weights",
    )
    nthreads = traits.Int(
        argstr="-nthreads %d",
        desc="number of threads. if zero, the number" " of available cpus will be used",
        nohash=True,
    )


class ComputeTDIOutputSpec(TraitedSpec):
    out_file = File(desc="output TDI file")


class ComputeTDI(MRTrix3Base):
    """
    Use track data as a form of contrast for producing a high-resolution
    image.

    .. admonition:: References

      * For TDI or DEC TDI: Calamante, F.; Tournier, J.-D.; Jackson, G. D. &
        Connelly, A. Track-density imaging (TDI): Super-resolution white
        matter imaging using whole-brain track-density mapping. NeuroImage,
        2010, 53, 1233-1243

      * If using -contrast length and -stat_vox mean: Pannek, K.; Mathias,
        J. L.; Bigler, E. D.; Brown, G.; Taylor, J. D. & Rose, S. E. The
        average pathlength map: A diffusion MRI tractography-derived index
        for studying brain pathology. NeuroImage, 2011, 55, 133-141

      * If using -dixel option with TDI contrast only: Smith, R.E., Tournier,
        J-D., Calamante, F., Connelly, A. A novel paradigm for automated
        segmentation of very large whole-brain probabilistic tractography
        data sets. In proc. ISMRM, 2011, 19, 673

      * If using -dixel option with any other contrast: Pannek, K., Raffelt,
        D., Salvado, O., Rose, S. Incorporating directional information in
        diffusion tractography derived maps: angular track imaging (ATI).
        In Proc. ISMRM, 2012, 20, 1912

      * If using -tod option: Dhollander, T., Emsell, L., Van Hecke, W., Maes,
        F., Sunaert, S., Suetens, P. Track Orientation Density Imaging (TODI)
        and Track Orientation Distribution (TOD) based tractography.
        NeuroImage, 2014, 94, 312-336

      * If using other contrasts / statistics: Calamante, F.; Tournier, J.-D.;
        Smith, R. E. & Connelly, A. A generalised framework for
        super-resolution track-weighted imaging. NeuroImage, 2012, 59,
        2494-2503

      * If using -precise mapping option: Smith, R. E.; Tournier, J.-D.;
        Calamante, F. & Connelly, A. SIFT: Spherical-deconvolution informed
        filtering of tractograms. NeuroImage, 2013, 67, 298-312 (Appendix 3)



    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> tdi = mrt.ComputeTDI()
    >>> tdi.inputs.in_file = 'dti.mif'
    >>> tdi.cmdline                               # doctest: +ELLIPSIS
    'tckmap dti.mif tdi.mif'
    >>> tdi.run()                                 # doctest: +SKIP
    """

    _cmd = "tckmap"
    input_spec = ComputeTDIInputSpec
    output_spec = ComputeTDIOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class TCK2VTKInputSpec(CommandLineInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-2, desc="input tractography"
    )
    out_file = File(
        "tracks.vtk", argstr="%s", usedefault=True, position=-1, desc="output VTK file"
    )
    reference = File(
        exists=True,
        argstr="-image %s",
        desc="if specified, the properties of"
        " this image will be used to convert track point positions from real "
        "(scanner) coordinates into image coordinates (in mm).",
    )
    voxel = File(
        exists=True,
        argstr="-image %s",
        desc="if specified, the properties of"
        " this image will be used to convert track point positions from real "
        "(scanner) coordinates into image coordinates.",
    )

    nthreads = traits.Int(
        argstr="-nthreads %d",
        desc="number of threads. if zero, the number" " of available cpus will be used",
        nohash=True,
    )


class TCK2VTKOutputSpec(TraitedSpec):
    out_file = File(desc="output VTK file")


class TCK2VTK(MRTrix3Base):
    """
    Convert a track file to a vtk format, cave: coordinates are in XYZ
    coordinates not reference

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> vtk = mrt.TCK2VTK()
    >>> vtk.inputs.in_file = 'tracks.tck'
    >>> vtk.inputs.reference = 'b0.nii'
    >>> vtk.cmdline                               # doctest: +ELLIPSIS
    'tck2vtk -image b0.nii tracks.tck tracks.vtk'
    >>> vtk.run()                                 # doctest: +SKIP
    """

    _cmd = "tck2vtk"
    input_spec = TCK2VTKInputSpec
    output_spec = TCK2VTKOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class DWIExtractInputSpec(MRTrix3BaseInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-2, desc="input image"
    )
    out_file = File(argstr="%s", mandatory=True, position=-1, desc="output image")
    bzero = traits.Bool(argstr="-bzero", desc="extract b=0 volumes")
    nobzero = traits.Bool(argstr="-no_bzero", desc="extract non b=0 volumes")
    singleshell = traits.Bool(
        argstr="-singleshell", desc="extract volumes with a specific shell"
    )
    shell = traits.List(
        traits.Float,
        sep=",",
        argstr="-shell %s",
        desc="specify one or more gradient shells",
    )


class DWIExtractOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="output image")


class DWIExtract(MRTrix3Base):
    """
    Extract diffusion-weighted volumes, b=0 volumes, or certain shells from a
    DWI dataset

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> dwiextract = mrt.DWIExtract()
    >>> dwiextract.inputs.in_file = 'dwi.mif'
    >>> dwiextract.inputs.bzero = True
    >>> dwiextract.inputs.out_file = 'b0vols.mif'
    >>> dwiextract.inputs.grad_fsl = ('bvecs', 'bvals')
    >>> dwiextract.cmdline                             # doctest: +ELLIPSIS
    'dwiextract -bzero -fslgrad bvecs bvals dwi.mif b0vols.mif'
    >>> dwiextract.run()                               # doctest: +SKIP
    """

    _cmd = "dwiextract"
    input_spec = DWIExtractInputSpec
    output_spec = DWIExtractOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class MRConvertInputSpec(MRTrix3BaseInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-2, desc="input image"
    )
    out_file = File(
        "dwi.mif",
        argstr="%s",
        mandatory=True,
        position=-1,
        usedefault=True,
        desc="output image",
    )
    coord = traits.List(
        traits.Float,
        sep=" ",
        argstr="-coord %s",
        desc="extract data at the specified coordinates",
    )
    vox = traits.List(
        traits.Float, sep=",", argstr="-vox %s", desc="change the voxel dimensions"
    )
    axes = traits.List(
        traits.Int,
        sep=",",
        argstr="-axes %s",
        desc="specify the axes that will be used",
    )
    scaling = traits.List(
        traits.Float,
        sep=",",
        argstr="-scaling %s",
        desc="specify the data scaling parameter",
    )


class MRConvertOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="output image")


class MRConvert(MRTrix3Base):
    """
    Perform conversion between different file types and optionally extract a
    subset of the input image

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> mrconvert = mrt.MRConvert()
    >>> mrconvert.inputs.in_file = 'dwi.nii.gz'
    >>> mrconvert.inputs.grad_fsl = ('bvecs', 'bvals')
    >>> mrconvert.cmdline                             # doctest: +ELLIPSIS
    'mrconvert -fslgrad bvecs bvals dwi.nii.gz dwi.mif'
    >>> mrconvert.run()                               # doctest: +SKIP
    """

    _cmd = "mrconvert"
    input_spec = MRConvertInputSpec
    output_spec = MRConvertOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class MRMathInputSpec(MRTrix3BaseInputSpec):
    in_file = File(
        exists=True, argstr="%s", mandatory=True, position=-3, desc="input image"
    )
    out_file = File(argstr="%s", mandatory=True, position=-1, desc="output image")
    operation = traits.Enum(
        "mean",
        "median",
        "sum",
        "product",
        "rms",
        "norm",
        "var",
        "std",
        "min",
        "max",
        "absmax",
        "magmax",
        argstr="%s",
        position=-2,
        mandatory=True,
        desc="operation to computer along a specified axis",
    )
    axis = traits.Int(
        0, argstr="-axis %d", desc="specfied axis to perform the operation along"
    )


class MRMathOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="output image")


class MRMath(MRTrix3Base):
    """
    Compute summary statistic on image intensities
    along a specified axis of a single image

    Example
    -------

    >>> import nipype.interfaces.mrtrix3 as mrt
    >>> mrmath = mrt.MRMath()
    >>> mrmath.inputs.in_file = 'dwi.mif'
    >>> mrmath.inputs.operation = 'mean'
    >>> mrmath.inputs.axis = 3
    >>> mrmath.inputs.out_file = 'dwi_mean.mif'
    >>> mrmath.inputs.grad_fsl = ('bvecs', 'bvals')
    >>> mrmath.cmdline                             # doctest: +ELLIPSIS
    'mrmath -axis 3 -fslgrad bvecs bvals dwi.mif mean dwi_mean.mif'
    >>> mrmath.run()                               # doctest: +SKIP
    """

    _cmd = "mrmath"
    input_spec = MRMathInputSpec
    output_spec = MRMathOutputSpec

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["out_file"] = op.abspath(self.inputs.out_file)
        return outputs


class MRResizeInputSpec(MRTrix3BaseInputSpec):
    in_file = File(
        exists=True, argstr="%s", position=-2, mandatory=True, desc="input DWI image"
    )
    image_size = traits.Tuple(
        (traits.Int, traits.Int, traits.Int),
        argstr="-size %d,%d,%d",
        mandatory=True,
        desc="Number of voxels in each dimension of output image",
        xor=["voxel_size", "scale_factor"],
    )
    voxel_size = traits.Tuple(
        (traits.Float, traits.Float, traits.Float),
        argstr="-voxel %g,%g,%g",
        mandatory=True,
        desc="Desired voxel size in mm for the output image",
        xor=["image_size", "scale_factor"],
    )
    scale_factor = traits.Tuple(
        (traits.Float, traits.Float, traits.Float),
        argstr="-scale %g,%g,%g",
        mandatory=True,
        desc="Scale factors to rescale the image by in each dimension",
        xor=["image_size", "voxel_size"],
    )
    interpolation = traits.Enum(
        "cubic",
        "nearest",
        "linear",
        "sinc",
        argstr="-interp %s",
        usedefault=True,
        desc="set the interpolation method to use when resizing (choices: "
        "nearest, linear, cubic, sinc. Default: cubic).",
    )
    out_file = File(
        argstr="%s",
        name_template="%s_resized",
        name_source=["in_file"],
        keep_extension=True,
        position=-1,
        desc="the output resized DWI image",
    )


class MRResizeOutputSpec(TraitedSpec):
    out_file = File(desc="the output resized DWI image", exists=True)


class MRResize(MRTrix3Base):
    """
    Resize an image by defining the new image resolution, voxel size or a
    scale factor. If the image is 4D, then only the first 3 dimensions can be
    resized. Also, if the image is down-sampled, the appropriate smoothing is
    automatically applied using Gaussian smoothing.
    For more information, see
    <https://mrtrix.readthedocs.io/en/latest/reference/commands/mrresize.html>

    Example
    -------
    >>> import nipype.interfaces.mrtrix3 as mrt

    Defining the new image resolution:
    >>> image_resize = mrt.MRResize()
    >>> image_resize.inputs.in_file = 'dwi.mif'
    >>> image_resize.inputs.image_size = (256, 256, 144)
    >>> image_resize.cmdline                               # doctest: +ELLIPSIS
    'mrresize -size 256,256,144 -interp cubic dwi.mif dwi_resized.mif'
    >>> image_resize.run()                                 # doctest: +SKIP

    Defining the new image's voxel size:
    >>> voxel_resize = mrt.MRResize()
    >>> voxel_resize.inputs.in_file = 'dwi.mif'
    >>> voxel_resize.inputs.voxel_size = (1, 1, 1)
    >>> voxel_resize.cmdline                               # doctest: +ELLIPSIS
    'mrresize -interp cubic -voxel 1,1,1 dwi.mif dwi_resized.mif'
    >>> voxel_resize.run()                                 # doctest: +SKIP

    Defining the scale factor of each image dimension:
    >>> scale_resize = mrt.MRResize()
    >>> scale_resize.inputs.in_file = 'dwi.mif'
    >>> scale_resize.inputs.scale_factor = (0.5,0.5,0.5)
    >>> scale_resize.cmdline                               # doctest: +ELLIPSIS
    'mrresize -interp cubic -scale 0.5,0.5,0.5 dwi.mif dwi_resized.mif'
    >>> scale_resize.run()                                 # doctest: +SKIP
    """

    _cmd = "mrresize"
    input_spec = MRResizeInputSpec
    output_spec = MRResizeOutputSpec