Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: Allow Jacobian correction to be toggled on/off #462

Merged
merged 2 commits into from
Oct 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Assuming you've checked all instances - inserting a new parameter in the middle of a non keyword only function makes me nervous

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's what testing's for. And yes, I've checked.

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 @@
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")

Check warning on line 222 in sdcflows/workflows/fit/pepolar.py

View check run for this annotation

Codecov / codecov/patch

sdcflows/workflows/fit/pepolar.py#L222

Added line #L222 was not covered by tests
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
Loading