Proof of concept for buffered iterator on nibabel dataobj

iter_vox.py

from __future__ import division
import sys
import numpy as np
import nibabel as nb

try:
    range = xrange
except NameError:
    pass

def update_buf(dataobj, buf_start, buf_shape):
    buf_end = [min(buf_start[i] + buf_shape[i], dataobj.shape[i]) 
               for i in range(3)]
    buf = dataobj[buf_start[0]:buf_end[0], 
                  buf_start[1]:buf_end[1], 
                  buf_start[2]:buf_end[2], 
                  :
                 ]
    
    return buf, buf_end
    

def iter_vox(img, max_buf=20*(1024**2)):
    '''Produce the time curve for each voxel in the image'''
    #Figure out how many curves we can buffer
    curve_len = img.shape[-1]
    dtype = img.get_data_dtype()
    curve_bytes = curve_len * dtype.itemsize
    n_buf_curves = max_buf // curve_bytes
    
    #Figure out the shape of our buffer
    buf_shape = [1, 1, 1]
    size = 1
    for i in range(3):
        dim_size = min(n_buf_curves // size, img.shape[i])
        if dim_size != img.shape[i] and i != 2:
            buf_shape[i] = 1
            break
        buf_shape[i] = dim_size
        size *= img.shape[i]
    
    #Initialize the buffer
    buf_start = [0, 0, 0]
    buf, buf_end = update_buf(img.dataobj, buf_start, buf_shape)
    
    #Iterate through voxels yielding results from the buffer and updating 
    #the buffer contents as needed
    for z in range(img.shape[2]):
        for y in range(img.shape[1]):
            for x in range(img.shape[0]):
                yield buf[x - buf_start[0], 
                          y - buf_start[1], 
                          z - buf_start[2], 
                          :
                         ]
                             
                #If the buffer is exhausted refill it
                if (x == buf_end[0] - 1 and 
                    y == buf_end[1] - 1 and 
                    z == buf_end[2] - 1):
                    all_full = True
                    for idx, end_val in enumerate(buf_end):
                        if all_full and end_val == img.shape[idx]:
                            buf_start[idx] = 0
                        else:
                            if all_full:
                                buf_start[idx] = end_val
                            all_full = False
                    buf, buf_end = update_buf(img.dataobj, 
                                              buf_start, 
                                              buf_shape)


def test_vox_iter(img, test_iter):
    arr = img.get_data()
    for z in range(img.shape[2]):
        for y in range(img.shape[1]):
            for x in range(img.shape[0]):
                v = arr[x, y, z, :]
                assert np.all(v == test_iter.next())

if __name__ == '__main__':
    img = nb.load(sys.argv[1])
    test_iter = iter_vox(img)
    test_vox_iter(img, test_iter)

If I understand correctly, the difference in performance between my slice by slice iterator and yours is that you allow an arbitrary sized buffer to be pre-loaded, whereas I am using the slice size as a buffer size?