From 5c047446795fe2ab84855bcd0d29f52f858b156b Mon Sep 17 00:00:00 2001
From: oesteban <code@oscaresteban.es>
Date: Tue, 19 Nov 2019 12:52:43 -0800
Subject: [PATCH] ENH: Deduplicating magnitude handling and fieldmap
 postprocessing workflows

This PR splits the magnitude processing and the fieldmap postprocessing
that is shared between phasediff (and phase1/2) and direct B0 mapping
approaches.

It also addresses some issues on the direct B0 mapping approach, where
the fieldmap was converted between units and subject to PRELUDE (which
is not correct as these fieldmaps are not phase-wrapped).

Adding a test case of SE fieldmaps would be great for a future PR.

Close #17
---
 sdcflows/workflows/fmap.py   |  85 +++-----------
 sdcflows/workflows/gre.py    | 220 +++++++++++++++++++++++++++++++++++
 sdcflows/workflows/phdiff.py | 109 +++--------------
 3 files changed, 254 insertions(+), 160 deletions(-)
 create mode 100644 sdcflows/workflows/gre.py

diff --git a/sdcflows/workflows/fmap.py b/sdcflows/workflows/fmap.py
index d7f4de8e21..9da412a463 100644
--- a/sdcflows/workflows/fmap.py
+++ b/sdcflows/workflows/fmap.py
@@ -19,16 +19,11 @@
 """
 
 from nipype.pipeline import engine as pe
-from nipype.interfaces import utility as niu, fsl, ants
-from niflow.nipype1.workflows.dmri.fsl.utils import demean_image, cleanup_edge_pipeline
+from nipype.interfaces import utility as niu, fsl
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
-from niworkflows.interfaces.bids import DerivativesDataSink
 from niworkflows.interfaces.images import IntraModalMerge
-from niworkflows.interfaces.masks import BETRPT
 
-from ..interfaces.fmap import (
-    FieldEnhance, FieldToRadS, FieldToHz
-)
+from .gre import init_fmap_postproc_wf, init_magnitude_wf
 
 
 def init_fmap_wf(omp_nthreads, fmap_bspline, name='fmap_wf'):
@@ -81,70 +76,28 @@ def init_fmap_wf(omp_nthreads, fmap_bspline, name='fmap_wf'):
     outputnode = pe.Node(niu.IdentityInterface(fields=['fmap', 'fmap_ref', 'fmap_mask']),
                          name='outputnode')
 
-    # Merge input magnitude images
-    magmrg = pe.Node(IntraModalMerge(), name='magmrg')
+    magnitude_wf = init_magnitude_wf(omp_nthreads=omp_nthreads)
+    workflow.connect([
+        (inputnode, magnitude_wf, [('magnitude', 'inputnode.magnitude')]),
+        (magnitude_wf, outputnode, [('outputnode.fmap_mask', 'fmap_mask'),
+                                    ('outputnode.fmap_ref', 'fmap_ref')]),
+    ])
+
     # Merge input fieldmap images
     fmapmrg = pe.Node(IntraModalMerge(zero_based_avg=False, hmc=False),
                       name='fmapmrg')
-
-    # de-gradient the fields ("bias/illumination artifact")
-    n4_correct = pe.Node(ants.N4BiasFieldCorrection(dimension=3, copy_header=True),
-                         name='n4_correct', n_procs=omp_nthreads)
-    bet = pe.Node(BETRPT(generate_report=True, frac=0.6, mask=True),
-                  name='bet')
-    ds_report_fmap_mask = pe.Node(DerivativesDataSink(
-        desc='brain', suffix='mask'), name='ds_report_fmap_mask',
-        run_without_submitting=True)
+    applymsk = pe.Node(fsl.ApplyMask(), name='applymsk')
+    fmap_postproc_wf = init_fmap_postproc_wf(omp_nthreads=omp_nthreads,
+                                             fmap_bspline=fmap_bspline)
 
     workflow.connect([
-        (inputnode, magmrg, [('magnitude', 'in_files')]),
         (inputnode, fmapmrg, [('fieldmap', 'in_files')]),
-        (magmrg, n4_correct, [('out_file', 'input_image')]),
-        (n4_correct, bet, [('output_image', 'in_file')]),
-        (bet, outputnode, [('mask_file', 'fmap_mask'),
-                           ('out_file', 'fmap_ref')]),
-        (inputnode, ds_report_fmap_mask, [('fieldmap', 'source_file')]),
-        (bet, ds_report_fmap_mask, [('out_report', 'in_file')]),
+        (fmapmrg, applymsk, [('out_file', 'in_file')]),
+        (magnitude_wf, applymsk, [('outputnode.fmap_mask', 'mask_file')]),
+        (applymsk, fmap_postproc_wf, [('out_file', 'inputnode.fmap')]),
+        (magnitude_wf, fmap_postproc_wf, [
+            ('outputnode.fmap_mask', 'inputnode.fmap_mask'),
+            ('outputnode.fmap_ref', 'inputnode.fmap_ref')]),
+        (fmap_postproc_wf, outputnode, [('outputnode.out_fmap', 'fmap')]),
     ])
-
-    if fmap_bspline:
-        # despike_threshold=1.0, mask_erode=1),
-        fmapenh = pe.Node(FieldEnhance(unwrap=False, despike=False),
-                          name='fmapenh', mem_gb=4, n_procs=omp_nthreads)
-
-        workflow.connect([
-            (bet, fmapenh, [('mask_file', 'in_mask'),
-                            ('out_file', 'in_magnitude')]),
-            (fmapmrg, fmapenh, [('out_file', 'in_file')]),
-            (fmapenh, outputnode, [('out_file', 'fmap')]),
-        ])
-
-    else:
-        torads = pe.Node(FieldToRadS(), name='torads')
-        prelude = pe.Node(fsl.PRELUDE(), name='prelude')
-        tohz = pe.Node(FieldToHz(), name='tohz')
-
-        denoise = pe.Node(fsl.SpatialFilter(operation='median', kernel_shape='sphere',
-                                            kernel_size=3), name='denoise')
-        demean = pe.Node(niu.Function(function=demean_image), name='demean')
-        cleanup_wf = cleanup_edge_pipeline(name='cleanup_wf')
-
-        applymsk = pe.Node(fsl.ApplyMask(), name='applymsk')
-
-        workflow.connect([
-            (bet, prelude, [('mask_file', 'mask_file'),
-                            ('out_file', 'magnitude_file')]),
-            (fmapmrg, torads, [('out_file', 'in_file')]),
-            (torads, tohz, [('fmap_range', 'range_hz')]),
-            (torads, prelude, [('out_file', 'phase_file')]),
-            (prelude, tohz, [('unwrapped_phase_file', 'in_file')]),
-            (tohz, denoise, [('out_file', 'in_file')]),
-            (denoise, demean, [('out_file', 'in_file')]),
-            (demean, cleanup_wf, [('out', 'inputnode.in_file')]),
-            (bet, cleanup_wf, [('mask_file', 'inputnode.in_mask')]),
-            (cleanup_wf, applymsk, [('outputnode.out_file', 'in_file')]),
-            (bet, applymsk, [('mask_file', 'mask_file')]),
-            (applymsk, outputnode, [('out_file', 'fmap')]),
-        ])
-
     return workflow
diff --git a/sdcflows/workflows/gre.py b/sdcflows/workflows/gre.py
new file mode 100644
index 0000000000..6e6cd59661
--- /dev/null
+++ b/sdcflows/workflows/gre.py
@@ -0,0 +1,220 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+"""Processing phase-difference (aka :abbr:`GRE (gradient-recalled echo)`) fieldmaps.
+
+.. _gre-fieldmaps:
+
+Workflows for processing :abbr:`GRE (gradient recalled echo)` fieldmaps
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Workflows for preparing the magnitude part of :abbr:`GRE (gradient-recalled echo)` fieldmap
+images and cleaning up the fieldmaps created from the phases or phasediff.
+
+"""
+
+from nipype.pipeline import engine as pe
+from nipype.interfaces import utility as niu, fsl, ants
+from niflow.nipype1.workflows.dmri.fsl.utils import cleanup_edge_pipeline
+from niworkflows.engine.workflows import LiterateWorkflow as Workflow
+from niworkflows.interfaces.images import IntraModalMerge
+from niworkflows.interfaces.masks import BETRPT
+
+
+def init_magnitude_wf(omp_nthreads, name='magnitude_wf'):
+    """
+    Prepare the magnitude part of :abbr:`GRE (gradient-recalled echo)` fieldmaps.
+
+    Average (if not done already) the magnitude part of the
+    :abbr:`GRE (gradient recalled echo)` images, run N4 to
+    correct for B1 field nonuniformity, and skull-strip the
+    preprocessed magnitude.
+
+    Workflow Graph
+        .. workflow ::
+            :graph2use: orig
+            :simple_form: yes
+
+            from sdcflows.workflows.fmap import init_magnitude_wf
+            wf = init_magnitude_wf(omp_nthreads=6)
+
+    Parameters
+    ----------
+    omp_nthreads : int
+        Maximum number of threads an individual process may use
+    name : str
+        Name of workflow (default: ``prepare_magnitude_w``)
+
+    Inputs
+    ------
+    magnitude : pathlike
+        Path to the corresponding magnitude path(s).
+
+    Outputs
+    -------
+    fmap_ref : pathlike
+        Path to the fieldmap reference calculated in this workflow.
+    fmap_mask : pathlike
+        Path to a binary brain mask corresponding to the reference above.
+
+    """
+    workflow = Workflow(name=name)
+    inputnode = pe.Node(
+        niu.IdentityInterface(fields=['magnitude']), name='inputnode')
+    outputnode = pe.Node(
+        niu.IdentityInterface(fields=['fmap_ref', 'fmap_mask', 'mask_report']),
+        name='outputnode')
+
+    # Merge input magnitude images
+    magmrg = pe.Node(IntraModalMerge(), name='magmrg')
+
+    # de-gradient the fields ("bias/illumination artifact")
+    n4_correct = pe.Node(ants.N4BiasFieldCorrection(dimension=3, copy_header=True),
+                         name='n4_correct', n_procs=omp_nthreads)
+    bet = pe.Node(BETRPT(generate_report=True, frac=0.6, mask=True),
+                  name='bet')
+
+    workflow.connect([
+        (inputnode, magmrg, [('magnitude', 'in_files')]),
+        (magmrg, n4_correct, [('out_file', 'input_image')]),
+        (n4_correct, bet, [('output_image', 'in_file')]),
+        (bet, outputnode, [('mask_file', 'fmap_mask'),
+                           ('out_file', 'fmap_ref'),
+                           ('out_report', 'mask_report')]),
+    ])
+    return workflow
+
+
+def init_fmap_postproc_wf(omp_nthreads, fmap_bspline, median_kernel_size=3,
+                          name='fmap_postproc_wf'):
+    """
+    Postprocess a B0 map estimated elsewhere.
+
+    This workflow denoises (mostly via smoothing) a B0 fieldmap.
+
+    Workflow Graph
+        .. workflow ::
+            :graph2use: orig
+            :simple_form: yes
+
+            from sdcflows.workflows.fmap import init_fmap_postproc_wf
+            wf = init_fmap_postproc_wf(omp_nthreads=6, fmap_bspline=False)
+
+    Parameters
+    ----------
+    omp_nthreads : int
+        Maximum number of threads an individual process may use
+    fmap_bspline : bool
+        Whether the fieldmap should be smoothed and extrapolated to off-brain regions
+        using B-Spline basis.
+    median_kernel_size : int
+        Size of the kernel when smoothing is done with a median filter.
+    name : str
+        Name of workflow (default: ``fmap_postproc_wf``)
+
+    Inputs
+    ------
+    fmap_mask : pathlike
+        A brain binary mask corresponding to this fieldmap.
+    fmap_ref : pathlike
+        A preprocessed magnitude/reference image for the fieldmap.
+    fmap : pathlike
+        A B0-field nonuniformity map (aka fieldmap) estimated elsewhere.
+
+    Outputs
+    -------
+    out_fmap : pathlike
+        Postprocessed fieldmap.
+
+    """
+    workflow = Workflow(name=name)
+    inputnode = pe.Node(niu.IdentityInterface(
+        fields=['fmap_mask', 'fmap_ref', 'fmap']), name='inputnode')
+    outputnode = pe.Node(niu.IdentityInterface(fields=['out_fmap']),
+                         name='outputnode')
+    if fmap_bspline:
+        from ..interfaces.fmap import FieldEnhance
+        # despike_threshold=1.0, mask_erode=1),
+        fmapenh = pe.Node(
+            FieldEnhance(unwrap=False, despike=False),
+            name='fmapenh', mem_gb=4, n_procs=omp_nthreads)
+
+        workflow.connect([
+            (inputnode, fmapenh, [('fmap_mask', 'in_mask'),
+                                  ('fmap_ref', 'in_magnitude'),
+                                  ('fmap_hz', 'in_file')]),
+            (fmapenh, outputnode, [('out_file', 'out_fmap')]),
+        ])
+
+    else:
+        recenter = pe.Node(niu.Function(function=_recenter),
+                           name='recenter', run_without_submitting=True)
+        denoise = pe.Node(fsl.SpatialFilter(
+                          operation='median', kernel_shape='sphere',
+                          kernel_size=median_kernel_size), name='denoise')
+        demean = pe.Node(niu.Function(function=_demean), name='demean')
+        cleanup_wf = cleanup_edge_pipeline(name="cleanup_wf")
+
+        workflow.connect([
+            (inputnode, cleanup_wf, [('fmap_mask', 'inputnode.in_mask')]),
+            (inputnode, recenter, [('fmap', 'in_file')]),
+            (recenter, denoise, [('out', 'in_file')]),
+            (denoise, demean, [('out_file', 'in_file')]),
+            (demean, cleanup_wf, [('out', 'inputnode.in_file')]),
+            (cleanup_wf, outputnode, [('outputnode.out_file', 'out_fmap')]),
+        ])
+
+    return workflow
+
+
+def _recenter(in_file):
+    """Recenter the phase-map distribution to the -pi..pi range."""
+    from os import getcwd
+    import numpy as np
+    import nibabel as nb
+    from nipype.utils.filemanip import fname_presuffix
+
+    nii = nb.load(in_file)
+    data = nii.get_fdata(dtype='float32')
+    msk = data != 0
+    msk[data == 0] = False
+    data[msk] -= np.median(data[msk])
+
+    out_file = fname_presuffix(in_file, suffix='_recentered',
+                               newpath=getcwd())
+    nb.Nifti1Image(data, nii.affine, nii.header).to_filename(out_file)
+    return out_file
+
+
+def _demean(in_file, in_mask=None, usemode=True):
+    """
+    Subtract the median (since it is robuster than the mean) from a map.
+
+    Parameters
+    ----------
+    usemode : bool
+        Use the mode instead of the median (should be even more robust
+        against outliers).
+
+    """
+    from os import getcwd
+    import numpy as np
+    import nibabel as nb
+    from nipype.utils.filemanip import fname_presuffix
+
+    nii = nb.load(in_file)
+    data = nii.get_fdata(dtype='float32')
+
+    msk = np.ones_like(data, dtype=bool)
+    if in_mask is not None:
+        msk[nb.load(in_mask).get_fdata(dtype='float32') < 1e-4] = False
+
+    if usemode:
+        from scipy.stats import mode
+        data[msk] -= mode(data[msk], axis=None)[0][0]
+    else:
+        data[msk] -= np.median(data[msk], axis=None)
+
+    out_file = fname_presuffix(in_file, suffix='_demean',
+                               newpath=getcwd())
+    nb.Nifti1Image(data, nii.affine, nii.header).to_filename(out_file)
+    return out_file
diff --git a/sdcflows/workflows/phdiff.py b/sdcflows/workflows/phdiff.py
index 4651e34688..fbd909e726 100644
--- a/sdcflows/workflows/phdiff.py
+++ b/sdcflows/workflows/phdiff.py
@@ -17,14 +17,12 @@
 
 """
 
-from nipype.interfaces import ants, fsl, utility as niu
+from nipype.interfaces import fsl, utility as niu
 from nipype.pipeline import engine as pe
-from niflow.nipype1.workflows.dmri.fsl.utils import cleanup_edge_pipeline
 from niworkflows.engine.workflows import LiterateWorkflow as Workflow
-from niworkflows.interfaces.images import IntraModalMerge
-from niworkflows.interfaces.masks import BETRPT
 
 from ..interfaces.fmap import Phasediff2Fieldmap, PhaseMap2rads
+from .gre import init_fmap_postproc_wf, init_magnitude_wf
 
 
 def init_phdiff_wf(omp_nthreads, name='phdiff_wf'):
@@ -98,110 +96,33 @@ def init_phdiff_wf(omp_nthreads, name='phdiff_wf'):
     outputnode = pe.Node(niu.IdentityInterface(
         fields=['fmap', 'fmap_ref', 'fmap_mask']), name='outputnode')
 
-    # Merge input magnitude images
-    magmrg = pe.Node(IntraModalMerge(), name='magmrg')
-
-    # de-gradient the fields ("bias/illumination artifact")
-    n4 = pe.Node(ants.N4BiasFieldCorrection(dimension=3, copy_header=True),
-                 name='n4', n_procs=omp_nthreads)
-    bet = pe.Node(BETRPT(generate_report=True, frac=0.6, mask=True),
-                  name='bet')
-
-    # uses mask from bet; outputs a mask
-    # dilate = pe.Node(fsl.maths.MathsCommand(
-    #     nan2zeros=True, args='-kernel sphere 5 -dilM'), name='MskDilate')
+    magnitude_wf = init_magnitude_wf(omp_nthreads=omp_nthreads)
 
     # phase diff -> radians
     phmap2rads = pe.Node(PhaseMap2rads(), name='phmap2rads',
                          run_without_submitting=True)
-
     # FSL PRELUDE will perform phase-unwrapping
     prelude = pe.Node(fsl.PRELUDE(), name='prelude')
 
-    recenter = pe.Node(niu.Function(function=_recenter),
-                       name='recenter', run_without_submitting=True)
-
-    denoise = pe.Node(fsl.SpatialFilter(operation='median', kernel_shape='sphere',
-                                        kernel_size=3), name='denoise')
-
-    demean = pe.Node(niu.Function(function=_demean), name='demean')
-
-    cleanup_wf = cleanup_edge_pipeline(name="cleanup_wf")
-
+    fmap_postproc_wf = init_fmap_postproc_wf(omp_nthreads=omp_nthreads,
+                                             fmap_bspline=False)
     compfmap = pe.Node(Phasediff2Fieldmap(), name='compfmap')
 
     workflow.connect([
         (inputnode, compfmap, [('metadata', 'metadata')]),
-        (inputnode, magmrg, [('magnitude', 'in_files')]),
-        (magmrg, n4, [('out_avg', 'input_image')]),
-        (n4, prelude, [('output_image', 'magnitude_file')]),
-        (n4, bet, [('output_image', 'in_file')]),
-        (bet, prelude, [('mask_file', 'mask_file')]),
+        (inputnode, magnitude_wf, [('magnitude', 'inputnode.magnitude')]),
+        (magnitude_wf, prelude, [('outputnode.fmap_ref', 'magnitude_file'),
+                                 ('outputnode.fmap_mask', 'mask_file')]),
         (inputnode, phmap2rads, [('phasediff', 'in_file')]),
         (phmap2rads, prelude, [('out_file', 'phase_file')]),
-        (prelude, recenter, [('unwrapped_phase_file', 'in_file')]),
-        (recenter, denoise, [('out', 'in_file')]),
-        (denoise, demean, [('out_file', 'in_file')]),
-        (demean, cleanup_wf, [('out', 'inputnode.in_file')]),
-        (bet, cleanup_wf, [('mask_file', 'inputnode.in_mask')]),
-        (cleanup_wf, compfmap, [('outputnode.out_file', 'in_file')]),
+        (prelude, fmap_postproc_wf, [('unwrapped_phase_file', 'inputnode.fmap')]),
+        (magnitude_wf, fmap_postproc_wf, [
+            ('outputnode.fmap_mask', 'inputnode.fmap_mask'),
+            ('outputnode.fmap_ref', 'inputnode.fmap_ref')]),
+        (fmap_postproc_wf, compfmap, [('outputnode.out_fmap', 'in_file')]),
         (compfmap, outputnode, [('out_file', 'fmap')]),
-        (bet, outputnode, [('mask_file', 'fmap_mask'),
-                           ('out_file', 'fmap_ref')]),
+        (magnitude_wf, outputnode, [('outputnode.fmap_ref', 'fmap_ref'),
+                                    ('outputnode.fmap_mask', 'fmap_mask')]),
     ])
 
     return workflow
-
-
-def _recenter(in_file):
-    """Recenter the phase-map distribution to the -pi..pi range."""
-    from os import getcwd
-    import numpy as np
-    import nibabel as nb
-    from nipype.utils.filemanip import fname_presuffix
-
-    nii = nb.load(in_file)
-    data = nii.get_fdata(dtype='float32')
-    msk = data != 0
-    msk[data == 0] = False
-    data[msk] -= np.median(data[msk])
-
-    out_file = fname_presuffix(in_file, suffix='_recentered',
-                               newpath=getcwd())
-    nb.Nifti1Image(data, nii.affine, nii.header).to_filename(out_file)
-    return out_file
-
-
-def _demean(in_file, in_mask=None, usemode=True):
-    """
-    Subtract the median (since it is robuster than the mean) from a map.
-
-    Parameters
-    ----------
-    usemode : bool
-        Use the mode instead of the median (should be even more robust
-        against outliers).
-
-    """
-    from os import getcwd
-    import numpy as np
-    import nibabel as nb
-    from nipype.utils.filemanip import fname_presuffix
-
-    nii = nb.load(in_file)
-    data = nii.get_fdata(dtype='float32')
-
-    msk = np.ones_like(data, dtype=bool)
-    if in_mask is not None:
-        msk[nb.load(in_mask).get_fdata(dtype='float32') < 1e-4] = False
-
-    if usemode:
-        from scipy.stats import mode
-        data[msk] -= mode(data[msk], axis=None)[0][0]
-    else:
-        data[msk] -= np.median(data[msk], axis=None)
-
-    out_file = fname_presuffix(in_file, suffix='_demean',
-                               newpath=getcwd())
-    nb.Nifti1Image(data, nii.affine, nii.header).to_filename(out_file)
-    return out_file