numpy ndarray: print imaginary part only if not zero-CodePudding

Continuing this question here, I'd like to ask how I can print a complex numpy array in a way that prints the imaginary part only if it's not zero. (This also goes for the real part.)

print(np.array([1.23456, 2j, 0, 3 4j], dtype='complex128'))

Expected Output:

[1.23    2.j   0.   3.   4.j]

CodePudding user response：

You can use np.isclose to check if a number is close to zero and the .real and .imag attributes to access the real and complex parts differently, then you can write a recursive function for printing:

import numpy as np

x = np.array([[1.23456, 2j, 0, 3 4j], [1,2,3,4]], dtype='complex128')

def complex_to_string(c):
    if c.imag == 0:
        return '{0.real:.2f}'.format(c)
    if c.real == 0:
        return '{0.imag:.2f}j'.format(c)
    return '{0.real:.2f}   {0.imag:.2f}j'.format(c)

def complex_arr_to_string(arr):
    if isinstance(arr, complex):
        return complex_to_string(arr)
    return "[" ' '.join(complex_arr_to_string(c) for c in arr) "]"

print(complex_arr_to_string(x))

Output:

[[1.23 0.00   2.00j 0.00 3.00   4.00j] [1.00 2.00 3.00 4.00]]

This works for arbitrarily nested arrays.

Thanks to @Koushik for mentioning the np.array2string builtin, using it the solution gets simpler:

import numpy as np
arr = np.array([[1.23456, 2j],[0, 3 4j]], dtype='complex128')

def complex_to_string(c):
    if c.imag == 0:
        return '{0.real:.2f}'.format(c)
    if c.real == 0:
        return '{0.imag:.2f}j'.format(c)
    return '{0.real:.2f}   {0.imag:.2f}j'.format(c)

print(np.array2string(arr, formatter={'complexfloat': complex_to_string}))

With the same output.

CodePudding user response：

import numpy as np
arr = np.array([1.23456, 2j, 0, 3 4j], dtype='complex128')
print(np.array2string(arr, formatter={'complexfloat':lambda x: f'{x.real}   {x.imag}j' if x.imag != 0 else f'{x.real}'}))