From 23c570e540bf0be9f1d1b97cf7883051753f02ac Mon Sep 17 00:00:00 2001
From: Oscar Esteban <code@oscaresteban.es>
Date: Sat, 28 Nov 2020 19:03:20 +0100
Subject: [PATCH] enh: add minimal unit test and upgrade from Function to
 fully-fledged interface

---
 sdcflows/interfaces/tests/test_utils.py | 22 +++++++++++++-
 sdcflows/interfaces/utils.py            | 39 +++++++++++++++++++++++++
 sdcflows/workflows/fit/pepolar.py       | 23 +++------------
 3 files changed, 64 insertions(+), 20 deletions(-)

diff --git a/sdcflows/interfaces/tests/test_utils.py b/sdcflows/interfaces/tests/test_utils.py
index 231827cba7..e5715e8490 100644
--- a/sdcflows/interfaces/tests/test_utils.py
+++ b/sdcflows/interfaces/tests/test_utils.py
@@ -1,8 +1,9 @@
 """Test utilites."""
+import pytest
 import numpy as np
 import nibabel as nb
 
-from ..utils import Flatten
+from ..utils import Flatten, ConvertWarp
 
 
 def test_Flatten(tmpdir):
@@ -31,3 +32,22 @@ def test_Flatten(tmpdir):
     assert out_meta[0] == {"a": 1}
     assert out_meta[1] == out_meta[2] == out_meta[3] == {"b": 2}
     assert out_meta[4] == out_meta[5] == {"c": 3}
+
+
+@pytest.mark.parametrize("shape", [
+    (10, 10, 10, 1, 3),
+    (10, 10, 10, 3)
+])
+def test_ConvertWarp(tmpdir, shape):
+    """Exercise the interface."""
+    tmpdir.chdir()
+
+    nb.Nifti1Image(np.zeros(shape, dtype="uint8"),
+                   np.eye(4), None).to_filename("3dQwarp.nii.gz")
+
+    out = ConvertWarp(in_file="3dQwarp.nii.gz").run()
+
+    nii = nb.load(out.outputs.out_file)
+    assert nii.header.get_data_dtype() == np.float32
+    assert nii.header.get_intent() == ("vector", (), "")
+    assert nii.shape == (10, 10, 10, 1, 3)
diff --git a/sdcflows/interfaces/utils.py b/sdcflows/interfaces/utils.py
index f1d1127aca..7d0200b6ed 100644
--- a/sdcflows/interfaces/utils.py
+++ b/sdcflows/interfaces/utils.py
@@ -50,6 +50,27 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class _ConvertWarpInputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc="output of 3dQwarp")
+
+
+class _ConvertWarpOutputSpec(TraitedSpec):
+    out_file = File(exists=True, desc="the warp converted into ANTs")
+
+
+class ConvertWarp(SimpleInterface):
+    """Convert a displacements field from ``3dQwarp`` to ANTS-compatible."""
+
+    input_spec = _ConvertWarpInputSpec
+    output_spec = _ConvertWarpOutputSpec
+
+    def _run_interface(self, runtime):
+        self._results["out_file"] = _qwarp2ants(
+            self.inputs.in_file, newpath=runtime.cwd
+        )
+        return runtime
+
+
 def _flatten(inlist, max_trs=50, out_dir=None):
     """
     Split the input EPIs and generate a flattened list with corresponding metadata.
@@ -83,3 +104,21 @@ def _flatten(inlist, max_trs=50, out_dir=None):
                 output.append((str(out_name), meta))
 
     return output
+
+
+def _qwarp2ants(in_file, newpath=None):
+    """Ensure the data type and intent of a warp is acceptable by ITK-based tools."""
+    import numpy as np
+    import nibabel as nb
+    from nipype.utils.filemanip import fname_presuffix
+
+    nii = nb.load(in_file)
+    hdr = nii.header.copy()
+    hdr.set_data_dtype("<f4")
+    hdr.set_intent("vector", (), "")
+    out_file = fname_presuffix(in_file, "_warpfield", newpath=newpath)
+    data = np.squeeze(nii.get_fdata(dtype="float32"))[..., np.newaxis, :]
+    nb.Nifti1Image(data, nii.affine, hdr).to_filename(
+        out_file
+    )
+    return out_file
diff --git a/sdcflows/workflows/fit/pepolar.py b/sdcflows/workflows/fit/pepolar.py
index 4a58be5f78..949a9d7ce3 100644
--- a/sdcflows/workflows/fit/pepolar.py
+++ b/sdcflows/workflows/fit/pepolar.py
@@ -160,7 +160,7 @@ def init_3dQwarp_wf(omp_nthreads=1, name="pepolar_estimate_wf"):
     from niworkflows.interfaces.freesurfer import StructuralReference
     from niworkflows.func.util import init_enhance_and_skullstrip_bold_wf
     from ...utils.misc import front as _front, last as _last
-    from ...interfaces.utils import Flatten
+    from ...interfaces.utils import Flatten, ConvertWarp
 
     workflow = Workflow(name=name)
     workflow.__desc__ = f"""{_PEPOLAR_DESC} \
@@ -226,7 +226,7 @@ def init_3dQwarp_wf(omp_nthreads=1, name="pepolar_estimate_wf"):
         n_procs=min(omp_nthreads, 4),
     )
 
-    to_ants = pe.Node(niu.Function(function=_fix_hdr), name="to_ants", mem_gb=0.01)
+    to_ants = pe.Node(ConvertWarp(), name="to_ants", mem_gb=0.01)
 
     cphdr_warp = pe.Node(CopyHeader(), name="cphdr_warp", mem_gb=0.01)
 
@@ -251,31 +251,16 @@ def init_3dQwarp_wf(omp_nthreads=1, name="pepolar_estimate_wf"):
         (inputnode, cphdr_warp, [(("in_data", _front), "hdr_file")]),
         (qwarp, cphdr_warp, [("source_warp", "in_file")]),
         (cphdr_warp, to_ants, [("out_file", "in_file")]),
-        (to_ants, unwarp_reference, [("out", "transforms")]),
+        (to_ants, unwarp_reference, [("out_file", "transforms")]),
         (inputnode, unwarp_reference, [("in_reference", "reference_image"),
                                        ("in_reference", "input_image")]),
         (unwarp_reference, outputnode, [("output_image", "fmap_ref")]),
-        (to_ants, outputnode, [("out", "fmap")]),
+        (to_ants, outputnode, [("out_file", "fmap")]),
     ])
     # fmt: on
     return workflow
 
 
-def _fix_hdr(in_file, newpath=None):
-    import nibabel as nb
-    from nipype.utils.filemanip import fname_presuffix
-
-    nii = nb.load(in_file)
-    hdr = nii.header.copy()
-    hdr.set_data_dtype("<f4")
-    hdr.set_intent("vector", (), "")
-    out_file = fname_presuffix(in_file, "_warpfield", newpath=newpath)
-    nb.Nifti1Image(nii.get_fdata(dtype="float32"), nii.affine, hdr).to_filename(
-        out_file
-    )
-    return out_file
-
-
 def _pe2fsl(metadata):
     """
     Convert ijk notation to xyz.