Skip to content

Commit 79305a9

Browse files
oestebaneffigies
andcommitted
enh: implement a memory limitation mechanism in loading data
Resolves: #218. Co-authored-by: Chris Markiewicz <[email protected]>
1 parent 7dcc78d commit 79305a9

File tree

1 file changed

+38
-1
lines changed

1 file changed

+38
-1
lines changed

nitransforms/resampling.py

+38-1
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ def apply(
9696
output_dtype: np.dtype = None,
9797
serialize_nvols: int = SERIALIZE_VOLUME_WINDOW_WIDTH,
9898
njobs: int = None,
99+
dtype_width: int = 8,
99100
) -> SpatialImage | np.ndarray:
100101
"""
101102
Apply a transformation to an image, resampling on the reference spatial object.
@@ -134,6 +135,10 @@ def apply(
134135
If ``reference`` is defined, then the return value is an image, with
135136
a data array of the effective dtype but with the on-disk dtype set to
136137
the input image's on-disk dtype.
138+
dtype_width: :obj:`int`
139+
Cap the width of the input data type to the given number of bytes.
140+
This argument is intended to work as a way to implement lower memory
141+
requirements in resampling.
137142
138143
Returns
139144
-------
@@ -157,7 +162,7 @@ def apply(
157162
spatialimage = _nbload(str(spatialimage))
158163

159164
# Avoid opening the data array just yet
160-
input_dtype = get_obj_dtype(spatialimage.dataobj)
165+
input_dtype = cap_dtype(get_obj_dtype(spatialimage.dataobj), dtype_width)
161166

162167
# Number of data volumes
163168
data_nvols = 1 if spatialimage.ndim < 4 else spatialimage.shape[-1]
@@ -277,3 +282,35 @@ def apply(
277282

278283
output_dtype = output_dtype or input_dtype
279284
return resampled.astype(output_dtype)
285+
286+
287+
def cap_dtype(dt, nbytes):
288+
"""
289+
Cap the datatype size to shave off memory requirements.
290+
291+
Examples
292+
--------
293+
>>> cap_dtype(np.dtype('f8'), 4)
294+
dtype('float32')
295+
296+
>>> cap_dtype(np.dtype('f8'), 16)
297+
dtype('float64')
298+
299+
>>> cap_dtype('float64', 4)
300+
dtype('float32')
301+
302+
>>> cap_dtype(np.dtype('i1'), 4)
303+
dtype('int8')
304+
305+
>>> cap_dtype('int8', 4)
306+
dtype('int8')
307+
308+
>>> cap_dtype('int32', 1)
309+
dtype('int8')
310+
311+
>>> cap_dtype(np.dtype('i8'), 4)
312+
dtype('int32')
313+
314+
"""
315+
dt = np.dtype(dt)
316+
return np.dtype(f"{dt.byteorder}{dt.kind}{min(nbytes, dt.itemsize)}")

0 commit comments

Comments
 (0)