Add Chen version as an option for autotrack (#201)

.circleci/config.yml

@@ -4,7 +4,7 @@ orbs:
 
 .dockersetup: &dockersetup
   docker:
-    - image: pennlinc/qsirecon_build:24.10.3
+    - image: pennlinc/qsirecon_build:25.1.0
   working_directory: /src/qsirecon
 
 runinstall: &runinstall

Dockerfile

@@ -5,7 +5,7 @@ RUN pip install build
 RUN apt-get update && \
     apt-get install -y --no-install-recommends git
 
-FROM pennlinc/qsirecon_build:24.10.3
+FROM pennlinc/qsirecon_build:25.1.0
 
 # Install qsirecon
 COPY . /src/qsirecon

qsirecon/interfaces/dsi_studio.py

@@ -713,7 +713,22 @@ class _AutoTrackInputSpec(DSIStudioCommandLineInputSpec):
     output_dir = traits.Str(
         "cwd", argstr="%s", usedefault=True, desc="Forces DSI Studio to write results in cwd"
     )
-    _boilerplate_traits = ["track_id", "track_voxel_ratio", "tolerance", "yield_rate"]
+    tip_iterations = traits.Int(
+        16,
+        usedefault=False,
+        desc="Topologically-informed pruning iterations",
+        argstr="--tip_iteration=%d",
+    )
+    template = traits.Int(
+        0, usedefault=True, argstr="--template=%d", desc="Must be 0 for autotrack"
+    )
+    _boilerplate_traits = [
+        "track_id",
+        "track_voxel_ratio",
+        "tolerance",
+        "yield_rate",
+        "tip_iteration",
+    ]
 
 
 class _AutoTrackOutputSpec(TraitedSpec):
@@ -760,6 +775,20 @@ def _list_outputs(self):
         return outputs
 
 
+class _ChenAutoTrackInputSpec(_AutoTrackInputSpec):
+    pass
+
+
+class _ChenAutoTrackOutputSpec(_AutoTrackOutputSpec):
+    pass
+
+
+class ChenAutoTrack(AutoTrack):
+    input_spec = _ChenAutoTrackInputSpec
+    output_spec = _ChenAutoTrackOutputSpec
+    _cmd = "dsi_studio_chen --action=atk"
+
+
 class _AggregateAutoTrackResultsInputSpec(BaseInterfaceInputSpec):
     expected_bundles = InputMultiObject(traits.Str())
     trk_files = InputMultiObject(File(exists=True))

qsirecon/workflows/recon/dsi_studio.py

@@ -25,6 +25,7 @@
     DSI_STUDIO_VERSION,
     AggregateAutoTrackResults,
     AutoTrack,
+    ChenAutoTrack,
     DSIStudioAtlasGraph,
     DSIStudioCreateSrc,
     DSIStudioExport,
@@ -306,9 +307,13 @@ def init_dsi_studio_autotrack_registration_wf(
     omp_nthreads = config.nipype.omp_nthreads
     workflow = Workflow(name=name)
 
+    dsi_studio_version = params.pop("dsi_studio_version", "hou")
+
+    _AutoTrack = AutoTrack if dsi_studio_version == "hou" else ChenAutoTrack
+
     # Run autotrack on only one bundle. The important part is getting the map.gz
     registration_atk = pe.Node(
-        AutoTrack(num_threads=omp_nthreads, track_id="Association_ArcuateFasciculusL"),
+        _AutoTrack(num_threads=omp_nthreads, track_id="Association_ArcuateFasciculusL"),
         name="registration_atk",
         n_procs=omp_nthreads,
     )
@@ -402,9 +407,13 @@ def init_dsi_studio_autotrack_wf(
     workflow = Workflow(name=name)
     workflow.__desc__ = desc + bundle_desc
 
+    dsi_studio_version = params.pop("dsi_studio_version", "hou")
+
+    _AutoTrack = AutoTrack if dsi_studio_version == "hou" else ChenAutoTrack
+
     # Run autotrack!
     actual_trk = pe.Node(
-        AutoTrack(num_threads=omp_nthreads, **params), name="actual_trk", n_procs=omp_nthreads
+        _AutoTrack(num_threads=omp_nthreads, **params), name="actual_trk", n_procs=omp_nthreads
     )  # An extra thread is needed
 
     # Create a single output