From 71845e53243beae5df462c33a9b2c28f4fb351c0 Mon Sep 17 00:00:00 2001
From: cneuromod_bot <cneuromod@gmail.com>
Date: Tue, 1 Aug 2023 15:52:32 -0400
Subject: [PATCH] gradunwarp to resampling: wip

---
 fmriprep/cli/parser.py                  |  4 +---
 fmriprep/workflows/base.py              |  1 +
 fmriprep/workflows/bold/base.py         | 16 ++++++++++++----
 fmriprep/workflows/bold/registration.py |  4 +++-
 fmriprep/workflows/bold/resampling.py   |  2 +-
 5 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/fmriprep/cli/parser.py b/fmriprep/cli/parser.py
index 119910a60..3b2e31ffd 100644
--- a/fmriprep/cli/parser.py
+++ b/fmriprep/cli/parser.py
@@ -282,9 +282,7 @@ def _slice_time_ref(value, parser):
     g_conf.add_argument(
         "--gradunwarp-file",
         metavar="PATH",
-        type=Path,
-        help="Path to vendor file for gradunwarp gradient distortion "
-        "correction.",
+        help="Path to vendor file for gradunwarp gradient distortion  correction.",
     )
     g_conf.add_argument(
         "--output-spaces",
diff --git a/fmriprep/workflows/base.py b/fmriprep/workflows/base.py
index ba2e9f037..9495d76b0 100644
--- a/fmriprep/workflows/base.py
+++ b/fmriprep/workflows/base.py
@@ -319,6 +319,7 @@ def init_single_subject_wf(subject_id: str):
         spaces=spaces,
         t1w=subject_data['t1w'],
         t2w=subject_data['t2w'],
+        flair=subject_data['flair'],
         gradunwarp_file=config.workflow.gradunwarp_file,
         cifti_output=config.workflow.cifti_output,
     )
diff --git a/fmriprep/workflows/bold/base.py b/fmriprep/workflows/bold/base.py
index 82a5e6cba..e3e671ca1 100644
--- a/fmriprep/workflows/bold/base.py
+++ b/fmriprep/workflows/bold/base.py
@@ -484,10 +484,13 @@ def init_func_preproc_wf(bold_file, has_fieldmap=False):
     # Gradient unwarping
     if config.workflow.gradunwarp_file:
         gradunwarp_wf = init_gradunwarp_wf()
-        gradunwarp_wf.inputs.inputnode.grad_file = gradunwarp_file
-        workflow.connect(
-            [(initial_boldref_wf, gradunwarp_wf, [('raw_ref_image', 'in_file')])]
-        )
+        gradunwarp_wf.inputs.inputnode.grad_file = config.workflow.gradunwarp_file
+        # input corrected file, raw not necessary, only uses acquisition matrix size
+        # unless mask/ref are used for sdcflows
+        workflow.connect([
+            (initial_boldref_wf, gradunwarp_wf, [
+                ('outputnode.raw_ref_image', 'inputnode.input_file')]),
+        ])
 
     # HMC on the BOLD
     bold_hmc_wf = init_bold_hmc_wf(
@@ -649,6 +652,7 @@ def init_func_preproc_wf(bold_file, has_fieldmap=False):
             ("mask", "inputnode.ref_bold_mask"),
             ("boldref", "inputnode.ref_bold_brain"),
         ]),
+        (gradunwarp_wf, bold_t1_trans_wf, [('outputnode.warp_file', 'inputnode.gradient_warp')]),
         (bold_t1_trans_wf, outputnode, [
             ("outputnode.bold_t1", "bold_t1"),
             ("outputnode.bold_t1_ref", "bold_t1_ref"),
@@ -814,6 +818,9 @@ def init_func_preproc_wf(bold_file, has_fieldmap=False):
                 ("outputnode.bold_std_ref", "bold_std_ref"),
                 ("outputnode.bold_mask_std", "bold_mask_std"),
             ]),
+            (gradunwarp_wf, bold_std_trans_wf, [
+                ('outputnode.warp_file', 'inputnode.gradient_warp')
+            ]),
         ])
         # fmt:on
 
@@ -1033,6 +1040,7 @@ def _last(inlist):
             (bold_hmc_wf, bold_bold_trans_wf, [
                 ("outputnode.xforms", "inputnode.hmc_xforms"),
             ]),
+            (gradunwarp_wf, bold_bold_trans_wf, [('outputnode.warp_file', 'inputnode.gradient_warp')]),
         ])
 
         workflow.connect([
diff --git a/fmriprep/workflows/bold/registration.py b/fmriprep/workflows/bold/registration.py
index c643bc7cf..bdb86edf2 100644
--- a/fmriprep/workflows/bold/registration.py
+++ b/fmriprep/workflows/bold/registration.py
@@ -325,6 +325,7 @@ def init_bold_t1_trans_wf(
                 't1w_aparc',
                 'bold_split',
                 'fieldwarp',
+                'gradient_warp',
                 'hmc_xforms',
                 'itk_bold_to_t1',
             ]
@@ -400,7 +401,7 @@ def init_bold_t1_trans_wf(
 
     # Merge transforms placing the head motion correction last
     merge_xforms = pe.Node(
-        niu.Merge(3),
+        niu.Merge(4),
         name='merge_xforms',
         run_without_submitting=True,
         mem_gb=DEFAULT_MEMORY_MIN_GB,
@@ -409,6 +410,7 @@ def init_bold_t1_trans_wf(
     workflow.connect([
         (inputnode, merge, [('name_source', 'header_source')]),
         (inputnode, merge_xforms, [
+            ('fieldwarp', 'in4'),  # May be 'identity' if no gradunwarp applied
             ('hmc_xforms', 'in3'),  # May be 'identity' if HMC already applied
             ('fieldwarp', 'in2'),   # May be 'identity' if SDC already applied
             ('itk_bold_to_t1', 'in1')]),
diff --git a/fmriprep/workflows/bold/resampling.py b/fmriprep/workflows/bold/resampling.py
index de665fb7a..91fc284b0 100644
--- a/fmriprep/workflows/bold/resampling.py
+++ b/fmriprep/workflows/bold/resampling.py
@@ -1184,7 +1184,7 @@ def init_bold_preproc_trans_wf(
     )
 
     merge_xforms = pe.Node(
-        niu.Merge(2),
+        niu.Merge(3),
         name="merge_xforms",
         run_without_submitting=True,
         mem_gb=DEFAULT_MEMORY_MIN_GB,