enh: port from process pool into asyncio concurrent

Co-authored-by: Chris Markiewicz <[email protected]>

nitransforms/resampling.py

@@ -8,10 +8,11 @@
 ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
 """Resampling utilities."""
 
+import asyncio
 from os import cpu_count
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from functools import partial
 from pathlib import Path
-from typing import Tuple
+from typing import Callable, TypeVar
 
 import numpy as np
 from nibabel.loadsave import load as _nbload
@@ -27,30 +28,58 @@
  _as_homogeneous,
 )
 
+R = TypeVar("R")
+
 SERIALIZE_VOLUME_WINDOW_WIDTH: int = 8
 """Minimum number of volumes to automatically serialize 4D transforms."""
 
 
-def _apply_volume(
- index: int,
+async def worker(job: Callable[[], R], semaphore) -> R:
+ async with semaphore:
+ loop = asyncio.get_running_loop()
+ return await loop.run_in_executor(None, job)
+
+
+async def _apply_serial(
  data: np.ndarray,
+ spatialimage: SpatialImage,
  targets: np.ndarray,
+ transform: TransformBase,
+ ref_ndim: int,
+ ref_ndcoords: np.ndarray,
+ n_resamplings: int,
+ output: np.ndarray,
+ input_dtype: np.dtype,
  order: int = 3,
  mode: str = "constant",
  cval: float = 0.0,
  prefilter: bool = True,
-) -> Tuple[int, np.ndarray]:
+ max_concurrent: int = min(cpu_count(), 12),
+):
  """
- Decorate :obj:`~scipy.ndimage.map_coordinates` to return an order index for parallelization.
+ Resample through a given transform serially, in a 3D+t setting.
 
  Parameters
  ----------
- index : :obj:`int`
- The index of the volume to apply the interpolation to.
  data : :obj:`~numpy.ndarray`
  The input data array.
+ spatialimage : :obj:`~nibabel.spatialimages.SpatialImage` or `os.pathlike`
+ The image object containing the data to be resampled in reference
+ space
  targets : :obj:`~numpy.ndarray`
  The target coordinates for mapping.
+ transform : :obj:`~nitransforms.base.TransformBase`
+ The 3D, 3D+t, or 4D transform through which data will be resampled.
+ ref_ndim : :obj:`int`
+ Dimensionality of the resampling target (reference image).
+ ref_ndcoords : :obj:`~numpy.ndarray`
+ Physical coordinates (RAS+) where data will be interpolated, if the resampling
+ target is a grid, the scanner coordinates of all voxels.
+ n_resamplings : :obj:`int`
+ Total number of 3D resamplings (can be defined by the input image, the transform,
+ or be matched, that is, same number of volumes in the input and number of transforms).
+ output : :obj:`~numpy.ndarray`
+ The output data array where resampled values will be stored volume-by-volume.
  order : :obj:`int`, optional
  The order of the spline interpolation, default is 3.
  The order has to be in the range 0-5.
@@ -71,18 +100,46 @@ def _apply_volume(
 
  Returns
  -------
- (:obj:`int`, :obj:`~numpy.ndarray`)
- The index and the array resulting from the interpolation.
+ np.ndarray
+ Data resampled on the 3D+t array of input coordinates.
 
  """
- return index, ndi.map_coordinates(
- data,
- targets,
- order=order,
- mode=mode,
- cval=cval,
- prefilter=prefilter,
- )
+ tasks = []
+ semaphore = asyncio.Semaphore(max_concurrent)
+
+ for t in range(n_resamplings):
+ xfm_t = transform if n_resamplings == 1 else transform[t]
+
+ if targets is None:
+ targets = ImageGrid(spatialimage).index( # data should be an image
+ _as_homogeneous(xfm_t.map(ref_ndcoords), dim=ref_ndim)
+ )
+
+ data_t = (
+ data
+ if data is not None
+ else spatialimage.dataobj[..., t].astype(input_dtype, copy=False)
+ )
+
+ tasks.append(
+ asyncio.create_task(
+ worker(
+ partial(
+ ndi.map_coordinates,
+ data_t,
+ targets,
+ output=output[..., t],
+ order=order,
+ mode=mode,
+ cval=cval,
+ prefilter=prefilter,
+ ),
+ semaphore,
+ )
+ )
+ )
+ await asyncio.gather(*tasks)
+ return output
 
 
 def apply(
@@ -94,15 +151,17 @@ def apply(
  cval: float = 0.0,
  prefilter: bool = True,
  output_dtype: np.dtype = None,
- serialize_nvols: int = SERIALIZE_VOLUME_WINDOW_WIDTH,
- njobs: int = None,
  dtype_width: int = 8,
+ serialize_nvols: int = SERIALIZE_VOLUME_WINDOW_WIDTH,
+ max_concurrent: int = min(cpu_count(), 12),
 ) -> SpatialImage | np.ndarray:
  """
  Apply a transformation to an image, resampling on the reference spatial object.
 
  Parameters
  ----------
+ transform: :obj:`~nitransforms.base.TransformBase`
+ The 3D, 3D+t, or 4D transform through which data will be resampled.
  spatialimage : :obj:`~nibabel.spatialimages.SpatialImage` or `os.pathlike`
  The image object containing the data to be resampled in reference
  space
@@ -118,15 +177,15 @@ def apply(
  or ``'wrap'``. Default is ``'constant'``.
  cval : :obj:`float`, optional
  Constant value for ``mode='constant'``. Default is 0.0.
- prefilter: :obj:`bool`, optional
+ prefilter : :obj:`bool`, optional
  Determines if the image's data array is prefiltered with
  a spline filter before interpolation. The default is ``True``,
  which will create a temporary *float64* array of filtered values
  if *order > 1*. If setting this to ``False``, the output will be
  slightly blurred if *order > 1*, unless the input is prefiltered,
  i.e. it is the result of calling the spline filter on the original
  input.
- output_dtype: :obj:`~numpy.dtype`, optional
+ output_dtype : :obj:`~numpy.dtype`, optional
  The dtype of the returned array or image, if specified.
  If ``None``, the default behavior is to use the effective dtype of
  the input image. If slope and/or intercept are defined, the effective
@@ -135,10 +194,17 @@ def apply(
  If ``reference`` is defined, then the return value is an image, with
  a data array of the effective dtype but with the on-disk dtype set to
  the input image's on-disk dtype.
- dtype_width: :obj:`int`
+ dtype_width : :obj:`int`
  Cap the width of the input data type to the given number of bytes.
  This argument is intended to work as a way to implement lower memory
  requirements in resampling.
+ serialize_nvols : :obj:`int`
+ Minimum number of volumes in a 3D+t (that is, a series of 3D transformations
+ independent in time) to resample on a one-by-one basis.
+ Serialized resampling can be executed concurrently (parallelized) with
+ the argument ``max_concurrent``.
+ max_concurrent : :obj:`int`
+ Maximum number of 3D resamplings to be executed concurrently.
 
  Returns
  -------
@@ -201,46 +267,30 @@ def apply(
  else None
  )
 
- njobs = cpu_count() if njobs is None or njobs < 1 else njobs
-
- with ProcessPoolExecutor(max_workers=min(njobs, n_resamplings)) as executor:
- results = []
- for t in range(n_resamplings):
- xfm_t = transform if n_resamplings == 1 else transform[t]
-
- if targets is None:
- targets = ImageGrid(spatialimage).index( # data should be an image
- _as_homogeneous(xfm_t.map(ref_ndcoords), dim=_ref.ndim)
- )
-
- data_t = (
- data
- if data is not None
- else spatialimage.dataobj[..., t].astype(input_dtype, copy=False)
- )
-
- results.append(
- executor.submit(
- _apply_volume,
- t,
- data_t,
- targets,
- order=order,
- mode=mode,
- cval=cval,
- prefilter=prefilter,
- )
- )
+ # Order F ensures individual volumes are contiguous in memory
+ # Also matches NIfTI, making final save more efficient
+ resampled = np.zeros(
+ (len(ref_ndcoords), len(transform)), dtype=input_dtype, order="F"
+ )
 
- # Order F ensures individual volumes are contiguous in memory
- # Also matches NIfTI, making final save more efficient
- resampled = np.zeros(
- (len(ref_ndcoords), len(transform)), dtype=input_dtype, order="F"
+ resampled = asyncio.run(
+ _apply_serial(
+ data,
+ spatialimage,
+ targets,
+ transform,
+ _ref.ndim,
+ ref_ndcoords,
+ n_resamplings,
+ resampled,
+ input_dtype,
+ order=order,
+ mode=mode,
+ cval=cval,
+ prefilter=prefilter,
+ max_concurrent=max_concurrent,
  )
-
- for future in as_completed(results):
- t, resampled_t = future.result()
- resampled[..., t] = resampled_t
+ )
  else:
  data = np.asanyarray(spatialimage.dataobj, dtype=input_dtype)
 

nitransforms/tests/test_resampling.py

@@ -15,7 +15,7 @@
 from nitransforms import nonlinear as nitnl
 from nitransforms import manip as nitm
 from nitransforms import io
-from nitransforms.resampling import apply, _apply_volume
+from nitransforms.resampling import apply
 
 RMSE_TOL_LINEAR = 0.09
 RMSE_TOL_NONLINEAR = 0.05
@@ -363,16 +363,3 @@ def test_LinearTransformsMapping_apply(
  reference=testdata_path / "sbref.nii.gz",
  serialize_nvols=2 if serialize_4d else np.inf,
  )
-
-
-@pytest.mark.parametrize("t", list(range(4)))
-def test_apply_helper(monkeypatch, t):
- """Ensure the apply helper function correctly just decorates with index."""
- from nitransforms.resampling import ndi
-
- def _retval(*args, **kwargs):
- return 1
-
- monkeypatch.setattr(ndi, "map_coordinates", _retval)
-
- assert _apply_volume(t, None, None) == (t, 1)