Skip to content

Commit

Permalink
Merge pull request #462 from effigies/fix/optional-jacobian
Browse files Browse the repository at this point in the history
ENH: Allow Jacobian correction to be toggled on/off
  • Loading branch information
effigies authored Oct 4, 2024
2 parents d953979 + 9c48489 commit 3f4bbf7
Show file tree
Hide file tree
Showing 5 changed files with 40 additions and 15 deletions.
6 changes: 6 additions & 0 deletions sdcflows/interfaces/bspline.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,11 @@ class _ApplyCoeffsFieldInputSpec(BaseInterfaceInputSpec):
mandatory=True,
desc="the phase-encoding direction corresponding to in_data",
)
jacobian = traits.Bool(
False,
usedefault=True,
desc="apply Jacobian determinant correction after unwarping",
)
num_threads = traits.Int(nohash=True, desc="number of threads")
approx = traits.Bool(
True,
Expand Down Expand Up @@ -421,6 +426,7 @@ def _run_interface(self, runtime):
self.inputs.in_data,
self.inputs.pe_dir,
self.inputs.ro_time,
jacobian=self.inputs.jacobian,
xfms=self.inputs.in_xfms or None,
xfm_data2fmap=data2fmap_xfm,
approx=self.inputs.approx,
Expand Down
26 changes: 18 additions & 8 deletions sdcflows/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def _sdc_unwarp(
coordinates: np.ndarray,
pe_info: Tuple[int, float],
hmc_xfm: np.ndarray | None,
jacobian: bool,
fmap_hz: np.ndarray,
output_dtype: str | np.dtype | None = None,
order: int = 3,
Expand All @@ -95,13 +96,6 @@ def _sdc_unwarp(
vsm = fmap_hz * pe_info[1]
coordinates[pe_info[0], ...] += vsm

# The Jacobian determinant image is the amount of stretching in the PE direction.
# Using central differences accounts for the shift in neighboring voxels.
# The full Jacobian at each voxel would be a 3x3 matrix, but because there is
# only warping in one direction, we end up with a diagonal matrix with two 1s.
# The following is the other entry at each voxel, and hence the determinant.
jacobian = 1 + np.gradient(vsm, axis=pe_info[0])

resampled = ndi.map_coordinates(
data,
coordinates,
Expand All @@ -110,7 +104,15 @@ def _sdc_unwarp(
mode=mode,
cval=cval,
prefilter=prefilter,
) * jacobian
)

# The Jacobian determinant image is the amount of stretching in the PE direction.
# Using central differences accounts for the shift in neighboring voxels.
# The full Jacobian at each voxel would be a 3x3 matrix, but because there is
# only warping in one direction, we end up with a diagonal matrix with two 1s.
# The following is the other entry at each voxel, and hence the determinant.
if jacobian:
resampled *= 1 + np.gradient(vsm, axis=pe_info[0])

return resampled

Expand Down Expand Up @@ -138,6 +140,7 @@ async def unwarp_parallel(
fmap_hz: np.ndarray,
pe_info: Sequence[Tuple[int, float]],
xfms: Sequence[np.ndarray],
jacobian: bool,
order: int = 3,
mode: str = "constant",
cval: float = 0.0,
Expand All @@ -162,6 +165,8 @@ async def unwarp_parallel(
pe_info : :obj:`tuple` of (:obj:`int`, :obj:`float`)
A tuple containing the index of the phase-encoding axis in the data array and
the readout time (including sign, if displacements must be reversed)
jacobian : :class:`bool`
If :obj:`True`, apply Jacobian determinant correction after unwarping.
xfms : :obj:`list` of obj:`~numpy.ndarray`
A list of 4×4 matrices, each one formalizing
the estimated head motion alignment to the scan's reference.
Expand Down Expand Up @@ -200,6 +205,7 @@ async def unwarp_parallel(

func = partial(
_sdc_unwarp,
jacobian=jacobian,
fmap_hz=fmap_hz,
output_dtype=output_dtype,
order=order,
Expand Down Expand Up @@ -370,6 +376,7 @@ def apply(
pe_dir: str | Sequence[str],
ro_time: float | Sequence[float],
xfms: Sequence[np.ndarray] | None = None,
jacobian: bool = True,
xfm_data2fmap: np.ndarray | None = None,
approx: bool = True,
order: int = 3,
Expand All @@ -394,6 +401,8 @@ def apply(
A valid ``PhaseEncodingDirection`` metadata value.
ro_time : :obj:`float` or :obj:`list` of :obj:`float`
The total readout time in seconds.
jacobian : :class:`bool`
If :obj:`True`, apply Jacobian determinant correction after unwarping.
xfms : :obj:`None` or :obj:`list`
A list of 4×4 matrices, each one formalizing
the estimated head motion alignment to the scan's reference.
Expand Down Expand Up @@ -528,6 +537,7 @@ def apply(
self.mapped.get_fdata(dtype="float32"), # fieldmap in Hz
pe_info,
xfms,
jacobian,
output_dtype='float32',
order=order,
mode=mode,
Expand Down
19 changes: 14 additions & 5 deletions sdcflows/workflows/apply/correction.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,14 @@
from niworkflows.engine.workflows import LiterateWorkflow as Workflow


def init_unwarp_wf(*, free_mem=None, omp_nthreads=1, debug=False, name="unwarp_wf"):
def init_unwarp_wf(
*,
jacobian=True,
free_mem=None,
omp_nthreads=1,
debug=False,
name="unwarp_wf",
):
r"""
Set up a workflow that unwarps the input :abbr:`EPI (echo-planar imaging)` dataset.
Expand All @@ -40,11 +47,13 @@ def init_unwarp_wf(*, free_mem=None, omp_nthreads=1, debug=False, name="unwarp_w
Parameters
----------
omp_nthreads : :obj:`int`
jacobian : :class:`bool`
If :obj:`True`, apply Jacobian determinant correction after unwarping.
omp_nthreads : :class:`int`
Maximum number of threads an individual process may use.
name : :obj:`str`
name : :class:`str`
Unique name of this workflow.
debug : :obj:`bool`
debug : :class:`bool`
Whether to run in *sloppy* mode.
Inputs
Expand Down Expand Up @@ -123,7 +132,7 @@ def init_unwarp_wf(*, free_mem=None, omp_nthreads=1, debug=False, name="unwarp_w
num_threads = omp_nthreads

resample = pe.Node(
ApplyCoeffsField(num_threads=num_threads),
ApplyCoeffsField(jacobian=jacobian, num_threads=num_threads),
mem_gb=mem_per_thread * num_threads,
name="resample",
)
Expand Down
2 changes: 1 addition & 1 deletion sdcflows/workflows/fit/pepolar.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,7 @@ def init_topup_wf(
from sdcflows.interfaces.bspline import ApplyCoeffsField

# Separate the runs again, as our ApplyCoeffsField corrects them separately
unwarp = pe.Node(ApplyCoeffsField(), name="unwarp")
unwarp = pe.Node(ApplyCoeffsField(jacobian=True), name="unwarp")
unwarp.interface._always_run = True
concat_corrected = pe.Node(MergeSeries(), name="concat_corrected")

Expand Down
2 changes: 1 addition & 1 deletion sdcflows/workflows/fit/syn.py
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,7 @@ def init_syn_sdc_wf(
DisplacementsField2Fieldmap(), name="extract_field"
)

unwarp = pe.Node(ApplyCoeffsField(), name="unwarp")
unwarp = pe.Node(ApplyCoeffsField(jacobian=False), name="unwarp")

# Check zooms (avoid very expensive B-Splines fitting)
zooms_field = pe.Node(
Expand Down

0 comments on commit 3f4bbf7

Please sign in to comment.