diff options
Diffstat (limited to 'numpy/lib/npyio.py')
| -rw-r--r-- | numpy/lib/npyio.py | 42 |
1 files changed, 32 insertions, 10 deletions
diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py index 3e3d00c57..3a575a048 100644 --- a/numpy/lib/npyio.py +++ b/numpy/lib/npyio.py @@ -880,7 +880,15 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', fmt : str or sequence of strs, optional A single format (%10.5f), a sequence of formats, or a multi-format string, e.g. 'Iteration %d -- %10.5f', in which - case `delimiter` is ignored. + case `delimiter` is ignored. For complex `X`, the legal options + for `fmt` are: + a) a single specifier, `fmt='%.4e'`, resulting in numbers formatted + like `' (%s+%sj)' % (fmt, fmt)` + b) a full string specifying every real and imaginary part, e.g. + `' %.4e %+.4j %.4e %+.4j %.4e %+.4j'` for 3 columns + c) a list of specifiers, one per column - in this case, the real + and imaginary part must have separate specifiers, + e.g. `['%.3e + %.3ej', '(%.15e%+.15ej)']` for 2 columns delimiter : str, optional Character separating columns. newline : str, optional @@ -895,11 +903,10 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', String that will be prepended to the ``header`` and ``footer`` strings, to mark them as comments. Default: '# ', as expected by e.g. ``numpy.loadtxt``. - .. versionadded:: 2.0.0 + .. versionadded:: 1.7.0 Character separating lines. - See Also -------- save : Save an array to a binary file in NumPy ``.npy`` format @@ -1003,6 +1010,7 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', else: ncol = X.shape[1] + iscomplex_X = np.iscomplexobj(X) # `fmt` can be a string with multiple insertion points or a # list of formats. E.g. '%10.5f\t%10d' or ('%10.5f', '$10d') if type(fmt) in (list, tuple): @@ -1010,20 +1018,34 @@ def savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', raise AttributeError('fmt has wrong shape. %s' % str(fmt)) format = asstr(delimiter).join(map(asstr, fmt)) elif type(fmt) is str: - if fmt.count('%') == 1: - fmt = [fmt, ]*ncol + n_fmt_chars = fmt.count('%') + error = ValueError('fmt has wrong number of %% formats: %s' % fmt) + if n_fmt_chars == 1: + if iscomplex_X: + fmt = [' (%s+%sj)' % (fmt, fmt),] * ncol + else: + fmt = [fmt, ] * ncol format = delimiter.join(fmt) - elif fmt.count('%') != ncol: - raise AttributeError('fmt has wrong number of %% formats. %s' - % fmt) + elif iscomplex_X and n_fmt_chars != (2 * ncol): + raise error + elif ((not iscomplex_X) and n_fmt_chars != ncol): + raise error else: format = fmt if len(header) > 0: header = header.replace('\n', '\n' + comments) fh.write(asbytes(comments + header + newline)) - for row in X: - fh.write(asbytes(format % tuple(row) + newline)) + if iscomplex_X: + for row in X: + row2 = [] + for number in row: + row2.append(number.real) + row2.append(number.imag) + fh.write(asbytes(format % tuple(row2) + newline)) + else: + for row in X: + fh.write(asbytes(format % tuple(row) + newline)) if len(footer) > 0: footer = footer.replace('\n', '\n' + comments) fh.write(asbytes(comments + footer + newline)) |