FFT to work on complex data vs [RE,IM] final shape of data

[walkloud]

Chroma sometimes writes complex data structures to H5 format with

[<array_structure>, re, im]

and sometimes as

[<array_structure>, complex]

So - we want either FFT to detect or a flag for the user to switch which data structure to act on

[asmeyer2012]

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)

[ckoerber]

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.