Search code examples
pythonnumpycomplex-numbers

How to save and load an array of complex numbers using numpy.savetxt?

I want to use numpy.savetxt() to save an array of complex numbers to a text file. Problems:

  • If you save the complex array with the default format string, the imaginary part is discarded.
  • If you use fmt='%s', then numpy.loadtxt() can't load it unless you specify dtype=complex, converters={0: lambda s: complex(s)}. Even then, if there are NaN's in the array, loading still fails.

It looks like someone has inquired about this multiple times on the Numpy mailing list and even filed a bug, but has not gotten a response. Before I put something together myself, is there a canonical way to do this?

Solution

  • It's easier and saves a few temporary arrays to just reinterpret the array as a real array.

    Saving:

    numpy.savetxt('outfile.txt', array.view(float))

    Loading:

    array = numpy.loadtxt('outfile.txt').view(complex)

    If you prefer to have real and imaginary part on the same line in the file, you can use

    numpy.savetxt('outfile.txt', array.view(float).reshape(-1, 2))

    or

    array = numpy.loadtxt('outfile.txt').view(complex).reshape(-1)

    respectively.

    (Note that neither view() nor reshape() copies the array -- it will just reinterpret the same data in a different way.)

    Addendum from the question asker:

    If you want to save more than one complex array in the same file, you can do it like so:

    numpy.savetxt('outfile.txt', numpy.column_stack([
    array1.view(float).reshape(-1, 2),
    array2.view(float).reshape(-1, 2),
]))

array1, array2 = numpy.loadtxt('outfile.txt', unpack=True).view(complex)

    The reshaping is necessary because numpy.view() doesn't operate on strided arrays.