From 02797a42d3870738e8eab0864379376aec1b3767 Mon Sep 17 00:00:00 2001
From: Matt Cieslak <mattcieslak@gmail.com>
Date: Mon, 13 Jan 2025 09:43:42 -0500
Subject: [PATCH] add back chen

---
 .circleci/config.yml                   |  2 +-
 Dockerfile                             |  2 +-
 qsirecon/interfaces/dsi_studio.py      | 14 ++++++++++++++
 qsirecon/workflows/recon/dsi_studio.py | 13 +++++++++++--
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index c10b3a8a..a9c6dee6 100644
--- a/.circleci/config.yml
+++ b/.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
diff --git a/Dockerfile b/Dockerfile
index 7fc57dbd..651f7ba5 100644
--- a/Dockerfile
+++ b/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
diff --git a/qsirecon/interfaces/dsi_studio.py b/qsirecon/interfaces/dsi_studio.py
index fae236ed..5afe3079 100644
--- a/qsirecon/interfaces/dsi_studio.py
+++ b/qsirecon/interfaces/dsi_studio.py
@@ -775,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))
diff --git a/qsirecon/workflows/recon/dsi_studio.py b/qsirecon/workflows/recon/dsi_studio.py
index 0c21f2db..fea7bc84 100644
--- a/qsirecon/workflows/recon/dsi_studio.py
+++ b/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