It's a bit more complicated than just adding support for complex. The data we have is a structural array of complex

>>> np.dtype(( np.dtype(('complex128',(4,4))), (3,3)))
dtype((('<c16', (4, 4)), (3, 3)))

This gets expanded by fft_file and its dimensions are appended to the spatial dimensions. When get_fft is called, it only uses the last 3 dimensions,

out = get_fft(arr, cuda=cuda, axes=(-1, -2, -3))

which are now dimensions in the structural array rather than spatial indices.

The max_momentum option also assumes the spatial indices are last, so it will cut in the wrong dimension and possibly get an out of bounds error:

for axis, key in enumerate(["x", "y", "z"]):
    axis = -1 * (axis + 1)
    LOGGER.debug(
        "\t\t Axis %d: %s -> %s[%s]", axis, key, key, slice_index
    )
    out = np.take(out, slice_index, axis=axis)

Thanks @asmeyer2012 for the PR!

I have not looked at it in detail yet but I was wondering if this addition should be more general (@walkloud, do we expect other files to also require this dynamic shape/complex reshuffling). E.g., should this reshuffling be incorporated into the dset reader and kept outside of the FFT?

E.g., conceptually the reader converts hdf5 data types to python types and here the python type should obviously be complex.