Skip to content

Commit 754785f

Browse files
committed
enh: prepare code for easy parallelization with a process pool executor
Resolves: #214.
1 parent 4c06174 commit 754785f

File tree

1 file changed

+23
-18
lines changed

1 file changed

+23
-18
lines changed

nitransforms/resampling.py

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
99
"""Resampling utilities."""
1010

11+
from functools import partial
1112
from pathlib import Path
1213
import numpy as np
1314
from nibabel.loadsave import load as _nbload
@@ -135,33 +136,37 @@ def apply(
135136
else None
136137
)
137138

138-
# Order F ensures individual volumes are contiguous in memory
139-
# Also matches NIfTI, making final save more efficient
140-
resampled = np.zeros(
141-
(len(ref_ndcoords), len(transform)), dtype=input_dtype, order="F"
139+
map_coordinates = partial(
140+
ndi.map_coordinates,
141+
order=order,
142+
mode=mode,
143+
cval=cval,
144+
prefilter=prefilter,
142145
)
143146

144-
for t in range(n_resamplings):
145-
xfm_t = transform if n_resamplings == 1 else transform[t]
147+
def _apply_volume(index, data, transform, targets=None):
148+
xfm_t = transform if n_resamplings == 1 else transform[index]
146149

147150
if targets is None:
148151
targets = ImageGrid(spatialimage).index( # data should be an image
149152
_as_homogeneous(xfm_t.map(ref_ndcoords), dim=_ref.ndim)
150153
)
151154

152-
# Interpolate
153-
resampled[..., t] = ndi.map_coordinates(
154-
(
155-
data
156-
if data is not None
157-
else spatialimage.dataobj[..., t].astype(input_dtype, copy=False)
158-
),
159-
targets,
160-
order=order,
161-
mode=mode,
162-
cval=cval,
163-
prefilter=prefilter,
155+
data_t = (
156+
data
157+
if data is not None
158+
else spatialimage.dataobj[..., index].astype(input_dtype, copy=False)
164159
)
160+
return map_coordinates(data_t, targets)
161+
162+
# Order F ensures individual volumes are contiguous in memory
163+
# Also matches NIfTI, making final save more efficient
164+
resampled = np.zeros(
165+
(len(ref_ndcoords), len(transform)), dtype=input_dtype, order="F"
166+
)
167+
for t in range(n_resamplings):
168+
# Interpolate
169+
resampled[..., t] = _apply_volume(t, data, transform, targets=targets)
165170

166171
else:
167172
data = np.asanyarray(spatialimage.dataobj, dtype=input_dtype)

0 commit comments

Comments
 (0)