[ENH] Stabilize autotrack performance

.circleci/AutoTrackTest.sh

@@ -25,7 +25,7 @@ CFG=${TESTDIR}/data/nipype.cfg
 EDDY_CFG=${TESTDIR}/data/eddy_config.json
 export FS_LICENSE=${TESTDIR}/data/license.txt
 
-# Test MRtrix3 multishell msmt with ACT
+# Test AutoTrack
 TESTNAME=autotrack
 setup_dir ${TESTDIR}/${TESTNAME}
 TEMPDIR=${TESTDIR}/${TESTNAME}/work
@@ -37,7 +37,7 @@ ${QSIPREP_CMD} \
 	 -w ${TEMPDIR} \
 	 --recon-input ${BIDS_INPUT_DIR} \
 	 --sloppy \
-     --stop-on-first-crash \
+         --stop-on-first-crash \
 	 --recon-spec dsi_studio_autotrack \
 	 --recon-only \
 	 -vv

.circleci/config.yml

@@ -2,7 +2,7 @@ version: 2.1
 
 .dockersetup: &dockersetup
   docker:
-    - image: pennbbl/qsiprep_build:23.5.6
+    - image: pennbbl/qsiprep_build:23.7.2
   working_directory: /src/qsiprep
 
 runinstall: &runinstall
@@ -218,7 +218,7 @@ jobs:
       - checkout
       - run: *runinstall
       - run:
-          name: Test the CSD recon workflows
+          name: Test the AutoTrack workflow
           no_output_timeout: 1h
           command: |
             cd .circleci

Dockerfile

@@ -1,4 +1,4 @@
-FROM pennbbl/qsiprep_build:23.5.6
+FROM pennbbl/qsiprep_build:23.7.2
 
 # WORKDIR /root/
 # # Installing qsiprep

qsiprep/interfaces/dsi_studio.py

@@ -759,7 +759,6 @@ class _AutoTrackInputSpec(DSIStudioCommandLineInputSpec):
                     copyfile=False,
                     argstr="--source=%s")
     map_file = File(exists=True,
-                    mandatory=True,
                     copyfile=False)
     track_id = traits.Str(
         "Fasciculus,Cingulum,Aslant,Corticos,Thalamic_R,Reticular,Optic,Fornix,Corpus",
@@ -812,9 +811,10 @@ class _AutoTrackInputSpec(DSIStudioCommandLineInputSpec):
     _boilerplate_traits = ["track_id", "track_voxel_ratio", "tolerance", "yield_rate"]
 
 
-class _AutoTrackOutputSpec(DSIStudioCommandLineInputSpec):
+class _AutoTrackOutputSpec(TraitedSpec):
     native_trk_files = OutputMultiObject(File(exists=True))
     stat_files = OutputMultiObject(File(exists=True))
+    map_file = File(exists=True)
 
 
 class AutoTrack(CommandLine):
@@ -834,28 +834,8 @@ def _list_outputs(self):
         stat_files = [str(fname.absolute()) for fname in cwd.rglob("*.stat.txt")]
         outputs["native_trk_files"] = trk_files
         outputs["stat_files"] = stat_files
-        return outputs
-
-
-class _AutoTrackInitInputSpec(_AutoTrackInputSpec):
-    num_threads = 1  # Force this so it doesn't use a ton of threads
-    map_file = File(mandatory=False)
-    track_id = "0"
-    num_retries = traits.Int(3, usedefault=True)
-    seconds_between_retries = traits.Int(500, usedefault=True)
-
-
-class _AutoTrackInitOutputSpec(_AutoTrackOutputSpec):
-    map_file = File(exists=True)
-
 
-class AutoTrackInit(AutoTrack):
-    input_spec = _AutoTrackInitInputSpec
-    output_spec = _AutoTrackInitOutputSpec
-
-    def _list_outputs(self):
-        outputs = self.output_spec().get()
-        cwd = Path(".")
+        # Find any mappings
         map_files = list(cwd.glob("*.map.gz"))
         if len(map_files) > 1:
             raise Exception("Too many map files generated")
@@ -864,27 +844,6 @@ def _list_outputs(self):
         outputs["map_file"] = str(map_files[0].absolute())
         return outputs
 
-    def _run_interface(self, runtime, correct_return_codes=(0,)):
-        if self.inputs.num_retries < 2:
-            return super(AutoTrackInit, self)._run_interface(
-                runtime, correct_return_codes)
-
-        for try_num in range(self.inputs.num_retries):
-            runtime = super(AutoTrackInit, self)._run_interface(
-                runtime, correct_return_codes)
-            succeeded = False
-            try:
-                self._list_outputs()
-                succeeded = True
-            except Exception:
-                LOGGER.info("Autotrack init has failed attempt {}".format(try_num + 1))
-                time.sleep(self.inputs.seconds_between_retries)
-
-            if succeeded:
-                return runtime
-
-        raise Exception("Autotrack was unable to run successfully.")
-
 
 class _AggregateAutoTrackResultsInputSpec(BaseInterfaceInputSpec):
     expected_bundles = InputMultiObject(traits.Str())

qsiprep/utils/sloppy_recon.py

@@ -16,7 +16,9 @@ def make_sloppy(spec):
         ("DSI Studio", "tractography"): {
             "fiber_count": 5000},
         ("DSI Studio", "autotrack"): {
-            "track_id": "0,1,2"},
+            "track_id": "Arcuate_Fasciculus_L,Arcuate_Fasciculus_R",
+            "tolerance": "30,40",
+            "track_voxel_ratio": 0.8},
 
         ("MRTrix3", "tractography"): {
             "tckgen": {

qsiprep/workflows/recon/dsi_studio.py

@@ -13,7 +13,7 @@
                                            DSIStudioAtlasGraph, DSIStudioExport,
                                            DSIStudioTracking, AutoTrack, _get_dsi_studio_bundles,
                                            FixDSIStudioExportHeader, AggregateAutoTrackResults,
-                                           AutoTrackInit, DSI_STUDIO_VERSION)
+                                           DSI_STUDIO_VERSION)
 
 import logging
 from ...interfaces.bids import ReconDerivativesDataSink
@@ -260,9 +260,6 @@ def init_dsi_studio_autotrack_wf(omp_nthreads, available_anatomical_data,
             This rate will be used to terminate tracking early if DSI Studio finds that the
             fiber tracking is not generating results. (default: 0.00001)
     """
-    if omp_nthreads < 3:
-        LOGGER.warn("Please set --omp-nthreads to 3 or greater and ensure you have enough CPU "
-                    "resources. Without at least 3 cpus this workflow will likely fail!!")
     inputnode = pe.Node(
         niu.IdentityInterface(
             fields=recon_workflow_input_fields + ['fibgz']),
@@ -282,19 +279,13 @@ def init_dsi_studio_autotrack_wf(omp_nthreads, available_anatomical_data,
     workflow = Workflow(name=name)
     workflow.__desc__ = desc + bundle_desc
 
-    # This initial autotrack is used to calculate the registration to the template
-    register_bundles = pe.Node(
-        AutoTrackInit(),
-        name="register_bundles",
-        n_procs=omp_nthreads)  # This process intrinsically uses ~3 cpus if you specify thread_count=1
-
-    # Real autotrack is run here on all the bundles
+    # Run autotrack!
     actual_trk = pe.Node(
-        AutoTrack(num_threads=max(1, omp_nthreads-1), **params),
+        AutoTrack(num_threads=omp_nthreads, **params),
         name="actual_trk",
         n_procs=omp_nthreads)  # An extra thread is needed
 
-    # Create a single
+    # Create a single output
     aggregate_atk_results = pe.Node(
         AggregateAutoTrackResults(
             expected_bundles=bundle_names),
@@ -303,8 +294,7 @@ def init_dsi_studio_autotrack_wf(omp_nthreads, available_anatomical_data,
     convert_to_tck = pe.MapNode(
         DSIStudioTrkToTck(),
         name="convert_to_tck",
-        iterfield="trk_file"
-    )
+        iterfield="trk_file")
 
     # Save tck files of the bundles into the outputs
     ds_tckfiles = pe.MapNode(
@@ -328,12 +318,10 @@ def init_dsi_studio_autotrack_wf(omp_nthreads, available_anatomical_data,
         run_without_submitting=True)
 
     workflow.connect([
-        (inputnode, register_bundles, [('fibgz', 'fib_file')]),
         (inputnode, actual_trk, [('fibgz', 'fib_file')]),
         (inputnode, aggregate_atk_results, [('dwi_file', 'source_file')]),
         (inputnode, convert_to_tck, [('dwi_file', 'reference_nifti')]),
-        (register_bundles, actual_trk, [
-            ('map_file', 'map_file')]),
+        (actual_trk, ds_mapping, [('map_file', 'in_file')]),
         (actual_trk, aggregate_atk_results, [
             ("native_trk_files", "trk_files"),
             ("stat_files", "stat_files")]),