用于节省内存的 Numpy nditer？

Question

I'm lost when iterating over a ndarray with nditer.

Background

I am trying to compute the eigenvalues of 3x3 symmetric matrices for each point in a 3D array. My data is a 4D array of shape [6,x,y,z] with the 6 values being the values of matrix at point x,y,z, over a ~500x500x500 cube of float32. I first used numpy's eigvalsh, but it's optimized for large matrices, while I can use analytical simplification for 3x3 symmetric matrices.

I then implemented wikipedia's simplification , both as a function that takes a single matrix and computes eigenvalues (then iterating naively with nested for loops), and then vectorized using numpy.

The problem is that now inside my vectorization, each operation creates an internal array of my data's size, culminating in too much RAM used and PC freeze.

I tried using numexpr etc, it's still around 10G usage.

What I'm trying to do

I want to iterate (using numpy's nditer) through my array so that for each matrix, I compute my eigenvalues. This would remove the need to allocate huge intermediary arrays because we only calculate ~ 10 float numbers at a time. Basically trying to substitute nested for loops into one iterator.

I'm looking for something like this :

for a,b,c,d,e,f in np.nditer([symMatrix,eigenOut]): # for each matrix in x,y,z

    # computing my output for this matrix
    eigenOut[...] = myLovelyEigenvalue(a,b,c,d,e,f)

The best I have so far is this :

for i in np.nditer([derived],[],[['readonly']],op_axes=[[1,2,3]]):

But this means that i takes all values of the 4D array instead of being a tuple of 6 length. I can't seem to get the hang of the nditer documentation.

What am I doing wrong ? Do you have any tips and tricks as to iterating over "all but one" axis ?

The point is to have an nditer that would outperform regular nested loops on iteration (once this works i'll change function calls, buffer iteration ... but so far I just want it to work ^^)

Answer 1

You don't really need np.nditer for this. A simpler way of iterating over all but the first axis is just to reshape into a [6, 500 ** 3] array, transpose it to [500 ** 3, 6] , then iterate over the rows:

for (a, b, c, d, e, f) in (symMatrix.reshape(6, -1).T):
    # do something involving a, b, c, d, e, f...

If you really want to use np.nditer then you would do something like this:

for (a, b, c, d, e, f) in np.nditer(x, flags=['external_loop'], order='F'):
    # do something involving a, b, c, d, e, f...

A potentially important thing to consider is that if symMatrix is C-order (row-major) rather than Fortran-order (column-major) then iterating over the first dimension may be significantly faster than iterating over the last 3 dimensions, since then you will be accessing adjacent blocks of memory address. You might therefore want to consider switching to Fortran-order.

I wouldn't expect a massive performance gain from either of these, since at the end of the day you're still doing all of your looping in Python and operating only on scalars rather than taking advantage of vectorization.