Search code examples
pythonnumpyfftwstride

Force numpy array to physically match striding in memory


For a cross-correlation routine I want to take some 2-D matrices (greyscale images), rotate half of them by 90 degrees, and Fourier transform them all. I am cross-correlating a huge number of frames so I am trying to use pyFFTW with the FFTW object interface, which I have used successfully in the past.

However, here with using numpy.rot90() I am running into the problem that numpy is not physically rotating the array in memory but simply changing the striding, whereas FFTW requires the array in physical memory to be actually rotated.

# Import a 2k x 2k image
mage = my_image_import_function( (2048,2048) )
# mage striding is (16384,8)
temp = np.rot90( mage, k=-1 )
# temp striding is (8, -16384 )
temp2 = np.copy( temp )
# temp2 striding is (8, 16384)
mage2 = np.lib.stride_tricks.as_strided( temp2, (2048,2048), (16384,8) )
# mage2 striding is (16384,8)
pyFFTWobj.update_arrays( mage2, mageFFT )
pyFFTWobj.execute()

The use of .as_strided() restores the original striding, so that it can be fed into pyFFTW. However, after application of the .as_strided() function, the mage2 is no longer rotated relative to mage. The .as_strided() has undone the rotation operation, so the above code does nothing.

How can a programmer physically force a numpy array to match its striding in memory?


Solution

  • You can supply np.copy an order kwarg that controls memory layout of the copied array. You seem to want a C contiguous array, so you want to do:

    temp2 = np.copy(temp, order='C')
    

    You could alternatively rely on the fact that, while the default value of order for the np.copy function is 'K', for the corresponding method of ndarray it is 'C', so this will also work:

    temp2 = temp.copy()
    

    Of course "explicit is better than implicit" and all that, so even if you go with the method, it is much better to explicitly ask for what you want:

    temp2 = temp.copy(order='C')
    

    Some fake data so see that it works:

    In [36]: a = np.random.randint(256, size=(2048, 2048)).astype(np.uint8)
    
    In [37]: a.strides
    Out[37]: (2048, 1)
    
    In [38]: np.rot90(a, k=-1).strides
    Out[38]: (1, -2048)
    # The method default works...
    In [39]: np.rot90(a, k=-1).copy().strides
    Out[39]: (2048, 1)
    # ...but explicit is better than implicit
    In [40]: np.rot90(a, k=-1).copy(order='C').strides
    Out[40]: (2048, 1)
    # The function default does not work
    In [41]: np.rot90(a, k=-1).copy(order='K').strides
    Out[41]: (1, 2048)