Memory issues when resampling large 4D files

[oesteban]

The resampling of transform-mappings is typically killed in standard settings when files are large (we've seen it with DWIs of ~150 volumes).

[oesteban]

@effigies - would you agree that:

1. We want to open only one volume of the reference image and then make sure it is cleared out from memory
2. We want to open the moving image with memory mapping
3. We want to resample volume-by-volume (or by N-length 4D chunks), so that the full file is not required in memory at once.

[oesteban]

For 1, I think if we open memory-maping it should be sufficient as we don't access the data array, correct?

nitransforms/nitransforms/linear.py

Line 440 in 1674e86

reference = _nbload(str(reference))

For 2, the moving image is loaded here:

nitransforms/nitransforms/linear.py

Line 447 in 1674e86

spatialimage = _nbload(str(spatialimage))

and then this is the relevant loop where we should make sure only one volume is loaded on to memory at a time:

nitransforms/nitransforms/linear.py

Lines 466 to 476 in 1674e86

	ndi.map_coordinates(
	data[..., t],
	targets[t, ..., : _ref.ndim].T,
	output=output_dtype,
	order=order,
	mode=mode,
	cval=cval,
	prefilter=prefilter,
	)
	for t in range(data.shape[-1])
	],

does it make sense?

[effigies]

1. You do not need to load a volume of the reference image at all, unless you're using a mask to narrow the field-of-view. All you need is the shape, affine and dtype.

2. If you switch to using dataobj, then it doesn't matter whether you have a mmap or not. You can create an iterator that returns one volume at a time:

   (dataobj[:, :, :, i].astype(dtype, copy=False) for i in range(dataobj.shape[3]))

   This could be made into a generator function if you need more flexibiilty.

3. Volume-by-volume definitely makes sense. Let the thread manager choose how many to load, and then set your threads based on available cores and memory.

[oesteban]

Thanks SO MUCH :)

[effigies]

I think you have a serious inefficiency here:

nitransforms/nitransforms/linear.py

Lines 464 to 478 in 1674e86

	resampled = np.stack(
	[
	ndi.map_coordinates(
	data[..., t],
	targets[t, ..., : _ref.ndim].T,
	output=output_dtype,
	order=order,
	mode=mode,
	cval=cval,
	prefilter=prefilter,
	)
	for t in range(data.shape[-1])
	],
	axis=0,
	)

You're creating a bunch of individual volumes, and then concatenating them in C order. There are going to be unnecessary copies under the hood, and another one when saving to NIfTI.

In fMRIPrep, I've moved to

    # Order F ensures individual volumes are contiguous in memory
    # Also matches NIfTI, making final save more efficient
    out_array = np.zeros(coordinates.shape[1:] + data.shape[-1:], dtype=output_dtype, order='F')

    tasks = [
        asyncio.create_task(
            worker(
                partial(
                    resample_vol,
                    data=volume,
                    coordinates=coordinates,
                    pe_info=pe_info[volid],
                    hmc_xfm=hmc_xfms[volid] if hmc_xfms else None,
                    fmap_hz=fmap_hz,
                    output=out_array[..., volid],
                    order=order,
                    mode=mode,
                    cval=cval,
                    prefilter=prefilter,
                ),
                semaphore,
            )
        )
        for volid, volume in enumerate(np.rollaxis(data, -1, 0))
    ]

The equivalent here would be:

            resampled = np.zeros(coordinates.shape[1:] + data.shape[-1:], dtype=output_dtype, order='F')
            for t in range(data.shape[-1]):
                ndi.map_coordinates(
                    data[..., t],
                    targets[t, ..., : _ref.ndim].T,
                    output=resampled[..., t],
                    order=order,
                    mode=mode,
                    cval=cval,
                    prefilter=prefilter,
                )

[oesteban]

Okay, I'm testing your suggestion, and before even getting there, we hit these lines:

nitransforms/nitransforms/linear.py

Lines 452 to 455 in 1674e86

	ycoords = self.map(_ref.ndcoords.T)
	targets = ImageGrid(spatialimage).index( # data should be an image
	_as_homogeneous(np.vstack(ycoords), dim=_ref.ndim)
	)

These push RSS from ~10GB to ~30GB.

[effigies]

I don't fully understand that line. We don't use it in sdcflows/fmriprep. We would just use:

targets = nt.base.SpatialReference.factory(spatialimage).ndcoords.astype('f4').T

[effigies]

Ah, hold on. Nevermind, I see the equivalents. Need to think it through a bit.

[oesteban]

I think that's somehow equivalent to _ref.ndcoords.T of line 452.

Lines 453-455 converts those physical coordinates into voxel coordinates - seemingly in a very inefficient way

[effigies]

What's the dimensionality of your target image?

[effigies]

And here's the equivalent of what we have in fMRIPrep:

ycoords = self.map(_ref.ndcoords.astype('f4').T)
ras2vox = ~nt.Affine(spatialimage.affine)
targets = ras2vox.map(ycoords)

If switching to float32 and Affine.map don't help, then you're just working with much bigger data. In which case it might make sense to start chunking arrays by slice as well as volume.

[oesteban]

>>> _ref.shape
(76, 102, 60)

Not so big, although it's 750 timepoints, so

>>> targets.shape
(750, 465120, 3)

[effigies]

Timepoints shouldn't matter at this point. You're just calculating the spatial coordinates, not iinterpolating values. Even at float64, targets should only be 10MB, so something's going wildly wrong.

I would skip the index() thing anyway. I don't think it buys you any clarity.

[oesteban]

Okay, timepoints do matter because here:

targets = ras2vox.map(ycoords)

you have to ravel the two first dimensions of targets.

[oesteban]

Looks like cramming all the coordinates in a single mapping is really not memory efficient.