Add drop_outliers_node and consolidate outputs

Author: Adam Richie-Halford

dmriprep/run.py

@@ -266,7 +266,7 @@ def run_dmriprep_pe(dwi_file, dwi_file_AP, dwi_file_PA, bvec_file, bval_file,
     eddy.inputs.repol = True
     eddy.inputs.niter = 1  # TODO: make this a parameter to the function with default 5
 
-    def drop_outliers_fn(outlier_report, threshold, dwi_file):
+    def id_outliers_fn(outlier_report, threshold, dwi_file):
         """Get list of scans that exceed threshold for number of outliers
 
         Parameters
@@ -309,75 +309,116 @@ def num_outliers(scan, outliers):
             return len([d for d in outliers if d['scan'] == scan])
 
         if 0 < threshold < 1:
-            threshold *= threshold * nib.load(dwi_file).shape[2]
+            threshold *= nib.load(dwi_file).shape[2]
 
         drop_scans = np.array([
             s for s in scans
             if num_outliers(s, outliers) > threshold
         ])
 
-        return drop_scans
+        outpath = op.abspath("dropped_scans.txt")
+        np.savetxt(outpath, drop_scans, fmt="%d")
 
-    drop_outliers_node = pe.Node(niu.Function(
+        return drop_scans, outpath
+
+    id_outliers_node = pe.Node(niu.Function(
         input_names=["outlier_report", "threshold", "dwi_file"],
-        output_names=["drop_scans"],
-        function=drop_outliers_fn),
-        name="drop_outliers_node"
+        output_names=["drop_scans", "outpath"],
+        function=id_outliers_fn),
+        name="id_outliers_node"
     )
 
-    drop_outliers_node.inputs.threshold = 0.02
-    drop_outliers_node.inputs.dwi_file = dwi_file
+    # TODO: make this a parameter to the function
+    id_outliers_node.inputs.threshold = 0.02
+    id_outliers_node.inputs.dwi_file = dwi_file
     wf.connect(prep, "fsl_eddy.out_outlier_report",
-               drop_outliers_node, "outlier_report")
+               id_outliers_node, "outlier_report")
+
+    merge = pe.Node(fsl.Merge(dimension='t'), name="mergeAPPA")
+    merge.inputs.in_files = [dwi_file_AP, dwi_file_PA]
+    wf.connect(merge, 'merged_file', prep, 'inputnode.alt_file')
+
+    fslroi = pe.Node(fsl.ExtractROI(t_min=0, t_size=1), name="fslroi")
+    wf.connect(prep, "outputnode.out_file", fslroi, "in_file")
 
-    def save_outlier_list_fn(drop_scans):
-        """Save list of outlier scans to file
+    bbreg = pe.Node(fs.BBRegister(contrast_type="t2", init="coreg",
+                                  out_fsl_file=True,
+                                  subjects_dir=subjects_dir,
+                                  epi_mask=True),
+                    name="bbreg")
+    bbreg.inputs.subject_id = 'freesurfer'  # bids_sub_name
+    wf.connect(fslroi, "roi_file", bbreg, "source_file")
+
+    def drop_outliers_fn(in_file, in_bval, in_bvec, drop_scans):
+        """Drop outlier volumes from dwi file
 
         Parameters
         ----------
+        in_file: string
+            Path to nii dwi file to drop outlier volumes from
+
+        in_bval: string
+            Path to bval file to drop outlier volumes from
+
+        in_bvec: string
+            Path to bvec file to drop outlier volumes from
+
         drop_scans: numpy.ndarray
-            Path to the fsl_eddy outlier report
+            List of scan indices to drop
 
         Returns
         -------
-        outpath: string
-            Path to output file where list is saved
+        out_file: string
+            Path to "thinned" output dwi file
+
+        out_bval: string
+            Path to "thinned" output bval file
+
+        out_bvec: string
+            Path to "thinned" output bvec file
         """
+        import nibabel as nib
         import numpy as np
         import os.path as op
-        outpath = op.abspath("dropped_scans.txt")
-        np.savetxt(outpath, drop_scans, fmt="%d")
-        return outpath
 
-    save_outlier_list_node = pe.Node(niu.Function(
-        input_names=["drop_scans"],
-        output_names=["outpath"],
-        function=save_outlier_list_fn),
-        name="save_outlier_list_node"
-    )
+        img = nib.load(op.abspath(in_file))
+        img_data = img.get_fdata()
+        img_data_thinned = np.delete(img_data, drop_scans, axis=3)
+        img_thinned = nib.Nifti1Image(img_data_thinned.astype(np.float64), img.affine)
 
-    wf.connect(drop_outliers_node, "drop_scans",
-               save_outlier_list_node, "drop_scans")
+        root, ext1 = op.splitext(in_file)
+        root, ext0 = op.splitext(root)
+        out_file = ''.join([root + "_thinned", ext0, ext1])
+        nib.save(img_thinned, op.abspath(out_file))
 
-    merge = pe.Node(fsl.Merge(dimension='t'), name="mergeAPPA")
-    merge.inputs.in_files = [dwi_file_AP, dwi_file_PA]
-    wf.connect(merge, 'merged_file', prep, 'inputnode.alt_file')
+        bval = np.loadtxt(in_bval)
+        bval_thinned = np.delete(bval, drop_scans, axis=0)
+        out_bval = '_thinned'.join(op.splitext(in_bval))
+        np.savetxt(out_bval, bval_thinned)
 
-    fslroi = pe.Node(fsl.ExtractROI(t_min=0, t_size=1), name="fslroi")
-    wf.connect(prep, "outputnode.out_file", fslroi, "in_file")
+        bvec = np.loadtxt(in_bvec)
+        bvec_thinned = np.delete(bvec, drop_scans, axis=1)
+        out_bvec = '_thinned'.join(op.splitext(in_bvec))
+        np.savetxt(out_bvec, bvec_thinned)
 
-    bbreg = pe.Node(fs.BBRegister(contrast_type="t2", init="coreg",
-                                  out_fsl_file=True,
-                                  subjects_dir=subjects_dir,
-                                  epi_mask=True),
-                    name="bbreg")
-    bbreg.inputs.subject_id = 'freesurfer'  # bids_sub_name
-    wf.connect(fslroi, "roi_file", bbreg, "source_file")
+        return out_file, out_bval, out_bvec
+
+    drop_outliers_node = pe.Node(niu.Function(
+        input_names=["in_file", "in_bval", "in_bvec", "drop_scans"],
+        output_names=["out_file", "out_bval", "out_bvec"],
+        function=drop_outliers_fn),
+        name="drop_outliers_node"
+    )
+
+    wf.connect(id_outliers_node, "drop_scans", drop_outliers_node, "drop_scans")
+    wf.connect(prep, "outputnode.out_file", drop_outliers_node, "in_file")
+    drop_outliers_node.inputs.in_bval = bval_file
+    wf.connect(prep, "outputnode.out_bvec", drop_outliers_node, "in_bvec")
 
     get_tensor = pe.Node(DTI(), name="dipy_tensor")
-    wf.connect(prep, "outputnode.out_file", get_tensor, "in_file")
-    get_tensor.inputs.in_bval = bval_file
-    wf.connect(prep, "outputnode.out_bvec", get_tensor, "in_bvec")
+    wf.connect(drop_outliers_node, "out_file", get_tensor, "in_file")
+    wf.connect(drop_outliers_node, "out_bval", get_tensor, "in_bval")
+    wf.connect(drop_outliers_node, "out_bvec", get_tensor, "in_bvec")
 
     scale_tensor = pe.Node(name='scale_tensor', interface=fsl.BinaryMaths())
     scale_tensor.inputs.operation = 'mul'
@@ -464,6 +505,10 @@ def get_orig(subjects_dir, sub):
         ("derivatives/dmriprep", "derivatives/{}/dmriprep".format(bids_sub_name))
     ]
 
+    wf.connect(drop_outliers_node, "out_file", datasink, "dmriprep.dwi.@thinned")
+    wf.connect(drop_outliers_node, "out_bval", datasink, "dmriprep.dwi.@bval_thinned")
+    wf.connect(drop_outliers_node, "out_bvec", datasink, "dmriprep.dwi.@bvec_thinned")
+
     wf.connect(prep, "outputnode.out_file", datasink, "dmriprep.dwi.@corrected")
     wf.connect(prep, "outputnode.out_bvec", datasink, "dmriprep.dwi.@rotated")
     wf.connect(prep, "fsl_eddy.out_parameter",
@@ -478,7 +523,7 @@ def get_orig(subjects_dir, sub):
     wf.connect(prep, "fsl_eddy.out_shell_alignment_parameters",
                datasink, "dmriprep.qc.@eddyparamsshellalign")
 
-    wf.connect(save_outlier_list_node, "outpath", datasink, "dmriprep.qc.@droppedscans")
+    wf.connect(id_outliers_node, "outpath", datasink, "dmriprep.qc.@droppedscans")
 
     wf.connect(get_tensor, "out_file", datasink, "dmriprep.dti.@tensor")
     wf.connect(get_tensor, "fa_file", datasink, "dmriprep.dti.@fa")