Draft MEDIC dynamic distortion correction method (second attempt)

Dockerfile

@@ -80,6 +80,11 @@ RUN mkdir /opt/convert3d && \
     curl -fsSL --retry 5 https://sourceforge.net/projects/c3d/files/c3d/Experimental/c3d-1.4.0-Linux-gcc64.tar.gz/download \
     | tar -xz -C /opt/convert3d --strip-components 1
 
+# Julia for MEDIC
+FROM downloader as julia
+RUN curl -fsSL https://install.julialang.org | sh -s -- --yes --default-channel 1.9.4 && \
+    mkdir -p /opt/julia/ && cp -r /root/.julia/juliaup/*/* /opt/julia/
+
 # Micromamba
 FROM downloader as micromamba
 WORKDIR /
@@ -156,12 +161,18 @@ RUN apt-get update -qq \
 # Install files from stages
 COPY --from=afni /opt/afni-latest /opt/afni-latest
 COPY --from=c3d /opt/convert3d/bin/c3d_affine_tool /usr/bin/c3d_affine_tool
+COPY --from=julia /opt/julia/ /opt/julia/
 
 # AFNI config
 ENV PATH="/opt/afni-latest:$PATH" \
     AFNI_IMSAVE_WARNINGS="NO" \
     AFNI_PLUGINPATH="/opt/afni-latest"
 
+# Julia config
+ENV PATH="/opt/julia/bin:${PATH}"
+# add libjulia to ldconfig
+RUN echo "/opt/julia/lib" >> /etc/ld.so.conf.d/julia.conf && ldconfig
+
 # Create a shared $HOME directory
 RUN useradd -m -s /bin/bash -G users sdcflows
 WORKDIR /home/sdcflows

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "scipy >= 1.8.1",
     "templateflow",
     "toml",
+    "warpkit == 0.1.1",
 ]
 dynamic = ["version"]
 

sdcflows/fieldmaps.py

@@ -49,6 +49,7 @@ class EstimatorType(Enum):
     PHASEDIFF = auto()
     MAPPED = auto()
     ANAT = auto()
+    MEDIC = auto()
 
 
 MODALITIES = {

sdcflows/interfaces/fmap.py

@@ -29,6 +29,8 @@
 from nipype import logging
 from nipype.interfaces.base import (
     BaseInterfaceInputSpec,
+    CommandLineInputSpec,
+    CommandLine,
     TraitedSpec,
     File,
     traits,
@@ -62,6 +64,27 @@ def _run_interface(self, runtime):
         return runtime
 
 
+class _PhaseMap2rads2InputSpec(BaseInterfaceInputSpec):
+    in_file = File(exists=True, mandatory=True, desc="input (wrapped) phase map")
+
+
+class _PhaseMap2rads2OutputSpec(TraitedSpec):
+    out_file = File(desc="the phase map in the range -3.14 - 3.14")
+
+
+class PhaseMap2rads2(SimpleInterface):
+    """Convert a phase map given in a.u. (e.g., 0-4096) to radians."""
+
+    input_spec = _PhaseMap2rads2InputSpec
+    output_spec = _PhaseMap2rads2OutputSpec
+
+    def _run_interface(self, runtime):
+        from ..utils.phasemanip import au2rads2
+
+        self._results["out_file"] = au2rads2(self.inputs.in_file, newpath=runtime.cwd)
+        return runtime
+
+
 class _SubtractPhasesInputSpec(BaseInterfaceInputSpec):
     in_phases = traits.List(File(exists=True), min=1, max=2, desc="input phase maps")
     in_meta = traits.List(
@@ -390,3 +413,94 @@ def _check_gross_geometry(
             f"{img1.get_filename()} {''.join(nb.aff2axcodes(img1.affine))}, "
             f"{img2.get_filename()} {''.join(nb.aff2axcodes(img2.affine))}"
         )
+
+
+class _MEDICInputSpec(CommandLineInputSpec):
+    mag_files = traits.List(
+        File(exists=True),
+        argstr="--magnitude %s",
+        mandatory=True,
+        minlen=2,
+        desc="Magnitude image(s) to verify registration",
+    )
+    phase_files = traits.List(
+        File(exists=True),
+        argstr="--phase %s",
+        mandatory=True,
+        minlen=2,
+        desc="Phase image(s) to verify registration",
+    )
+    metadata = traits.List(
+        File(exists=True),
+        argstr="--metadata %s",
+        mandatory=True,
+        minlen=2,
+        desc="Metadata corresponding to the inputs",
+    )
+    prefix = traits.Str(
+        "medic",
+        argstr="--out_prefix %s",
+        usedefault=True,
+        desc="Prefix for output files",
+    )
+    noise_frames = traits.Int(
+        0,
+        argstr="--noiseframes %d",
+        usedefault=True,
+        desc="Number of noise frames to remove",
+    )
+    n_cpus = traits.Int(
+        4,
+        argstr="--n_cpus %d",
+        usedefault=True,
+        desc="Number of CPUs to use",
+    )
+    debug = traits.Bool(
+        False,
+        argstr="--debug",
+        usedefault=True,
+        desc="Enable debugging output",
+    )
+    wrap_limit = traits.Bool(
+        False,
+        argstr="--wrap_limit",
+        usedefault=True,
+        desc="Turns off some heuristics for phase unwrapping",
+    )
+
+
+class _MEDICOutputSpec(TraitedSpec):
+    native_field_map = File(
+        exists=True,
+        desc="4D ative (distorted) space field map in Hertz",
+    )
+    displacement_map = File(
+        exists=True,
+        desc="4D displacement map in millimeters",
+    )
+    field_map = File(
+        exists=True,
+        desc="4D undistorted field map in Hertz",
+    )
+
+
+class MEDIC(CommandLine):
+    """Run MEDIC."""
+
+    _cmd = "medic"
+    input_spec = _MEDICInputSpec
+    output_spec = _MEDICOutputSpec
+
+    def _list_outputs(self):
+        outputs = self._outputs().get()
+        out_dir = os.getcwd()
+        outputs['native_field_map'] = os.path.join(
+            out_dir,
+            f'{self.inputs.prefix}_fieldmaps_native.nii',
+        )
+        outputs['displacement_map'] = os.path.join(
+            out_dir,
+            f'{self.inputs.prefix}_displacementmaps.nii',
+        )
+        outputs['field_map'] = os.path.join(out_dir, f'{self.inputs.prefix}_fieldmaps.nii')
+        return outputs

sdcflows/utils/phasemanip.py

@@ -46,6 +46,30 @@ def au2rads(in_file, newpath=None):
     return out_file
 
 
+def au2rads2(in_file, newpath=None):
+    """Convert the input phase map in arbitrary units (a.u.) to rads (-pi to pi)."""
+    import numpy as np
+    import nibabel as nb
+    from nipype.utils.filemanip import fname_presuffix
+
+    im = nb.load(in_file)
+    data = im.get_fdata(caching="unchanged")  # Read as float64 for safety
+    hdr = im.header.copy()
+
+    # Rescale to [0, 2*pi]
+    data = (data - data.min()) * (2 * np.pi / (data.max() - data.min()))
+    data = data - np.pi
+
+    # Round to float32 and clip
+    data = np.clip(np.float32(data), -np.pi, np.pi)
+
+    hdr.set_data_dtype(np.float32)
+    hdr.set_xyzt_units("mm")
+    out_file = fname_presuffix(str(in_file), suffix="_rads", newpath=newpath)
+    nb.Nifti1Image(data, None, hdr).to_filename(out_file)
+    return out_file
+
+
 def subtract_phases(in_phases, in_meta, newpath=None):
     """Calculate the phase-difference map, given two input phase maps."""
     import numpy as np

sdcflows/workflows/fit/medic.py

@@ -0,0 +1,134 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+#
+# Copyright 2021 The NiPreps Developers <[email protected]>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# We support and encourage derived works from this project, please read
+# about our expectations at
+#
+#     https://www.nipreps.org/community/licensing/
+#
+"""Processing of dynamic field maps from complex-valued multi-echo BOLD data."""
+
+from nipype.pipeline import engine as pe
+from nipype.interfaces import utility as niu
+from niworkflows.engine.workflows import LiterateWorkflow as Workflow
+
+from sdcflows.interfaces.fmap import MEDIC, PhaseMap2rads2
+
+INPUT_FIELDS = ("magnitude", "phase")
+
+
+def init_medic_wf(name="medic_wf"):
+    """Create the MEDIC dynamic field estimation workflow.
+
+    Workflow Graph
+        .. workflow ::
+            :graph2use: orig
+            :simple_form: yes
+
+            from sdcflows.workflows.fit.medic import init_medic_wf
+
+            wf = init_medic_wf()   # doctest: +SKIP
+
+    Parameters
+    ----------
+    name : :obj:`str`
+        Name for this workflow
+
+    Inputs
+    ------
+    magnitude : :obj:`list` of :obj:`str`
+        A list of echo-wise magnitude EPI files that will be fed into MEDIC.
+    phase : :obj:`list` of :obj:`str`
+        A list of echo-wise phase EPI files that will be fed into MEDIC.
+    metadata : :obj:`list` of :obj:`dict`
+        List of metadata dictionaries corresponding to each of the input magnitude files.
+
+    Outputs
+    -------
+    fieldmap : :obj:`str`
+        The path of the estimated fieldmap time series file. Units are Hertz.
+    displacement : :obj:`list` of :obj:`str`
+        Path to the displacement time series files. Units are mm.
+    method: :obj:`str`
+        Short description of the estimation method that was run.
+
+    Notes
+    -----
+    This is a translation of the MEDIC algorithm, as implemented in ``vandandrew/warpkit``
+    (specifically the function ``unwrap_and_compute_field_maps``), into a Nipype workflow.
+
+    """
+    workflow = Workflow(name=name)
+
+    workflow.__desc__ = """\
+A dynamic fieldmap was estimated from multi-echo EPI data using the MEDIC algorithm (@medic).
+"""
+
+    inputnode = pe.Node(niu.IdentityInterface(fields=INPUT_FIELDS), name="inputnode")
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=["fieldmap", "displacement", "method"]),
+        name="outputnode",
+    )
+    outputnode.inputs.method = "MEDIC"
+
+    # Write metadata dictionaries to JSON files
+    write_metadata = pe.MapNode(
+        niu.Function(
+            input_names=["metadata"],
+            output_names=["out_file"],
+            function=write_json,
+        ),
+        iterfield=["metadata"],
+        name="write_metadata",
+    )
+    workflow.connect([(inputnode, write_metadata, [("metadata", "metadata")])])
+
+    # Convert phase to radians (-pi to pi, not 0 to 2pi)
+    phase2rad = pe.MapNode(
+        PhaseMap2rads2(),
+        iterfield=["in_file"],
+        name="phase2rad",
+    )
+    workflow.connect([(inputnode, phase2rad, [("phase", "in_file")])])
+
+    medic = pe.Node(
+        MEDIC(),
+        name="medic",
+    )
+    workflow.connect([
+        (inputnode, medic, [("magnitude", "magnitude")]),
+        (phase2rad, medic, [("out_file", "metadata")]),
+        (phase2rad, medic, [("out_file", "phase")]),
+        (medic, outputnode, [
+            ("native_field_map", "fieldmap"),
+            ("displacement_map", "displacement"),
+        ]),
+    ])  # fmt:skip
+
+    return workflow
+
+
+def write_json(metadata):
+    """Write a dictionary to a JSON file."""
+    import json
+    import os
+
+    out_file = os.path.abspath("metadata.json")
+    with open(out_file, "w") as fobj:
+        json.dump(metadata, fobj, sort_keys=True, indent=4)
+
+    return out_file