Python nested for loop on tensor-valued volume data

Question

Using ITK python wrapping, I'm reading a tensor-valued volume. On each tensor I perform an operation like so:

image      = reader.Execute()
dimensions = image.GetSize()

for z in range ( 0, dimensions[2] ):
  for y in range ( 0, dimensions[1] ):
    for x in range ( 0, dimensions[0] ):      
      image[x,y,z] = function( image[x,y,z] )

Obviously, the nested for loops are slow using Python. However, I can't seem to get this resolved using maps or list comprehension. Note that the function is performed on a list of 9 values, which represents the 3x3 tensor.

Anybody have a suggestion how to increase performance? Perhaps a flatten operation and manual indexing of the resulting list?

Thanks!

Answer 1

After some experiments, I'm now using a trade-off between readability and speed. The map function provides a neat solution, although it is still rather slow. In part, my question was addressed in: Numpy: Beginner nditer

Here is my current version:

reader     = sitk.ImageFileReader()
reader.SetFileName ( tensorVolumePath )  

image      = reader.Execute()
dimes      = image.GetSize()
origin     = image.GetOrigin()    

values     = sitk.GetArrayFromImage( image )
valuesFlat = values.reshape(dims[0]*dims[1]*dims[2], 3, 3)  

arrayInv   = np.array(map(np.linalg.inv, valuesFlat))  

imageInv = sitk.GetImageFromArray( valuesFlat.reshape(dims[0], dims[1], dims[2], -1) )
imageInv.SetOrigin(origin)

writer = sitk.ImageFileWriter()
writer.SetFileName ( tensorVolumeInversePath )
writer.Execute ( imageInv );

Answer 2

You can use itertools.product, but I'd profile before assuming that it's actually faster:

import itertools

xs = range(0, dimensions[0])
ys = range(0, dimensions[1])
zs = range(0, dimensions[2])

coords = itertools.product(xs, ys, zs)

for v in coords:
    image[v[0], v[1], v[2]] = function(image[v[0], v[1], v[2]])

My guess is that, sadly, any loop over the pixels of an image in Python will be incredibly slow and you'd need to somehow express your operation without loops using the imaging library.