1

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]
Caridorc
  • 6,222
  • 2
  • 31
  • 46
GoldenDremora
  • 167
  • 9

2 Answers2

2

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.

Caridorc
  • 6,222
  • 2
  • 31
  • 46
0
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}'}))
MatBailie
  • 83,401
  • 18
  • 103
  • 137
Koushik Reddy
  • 9
  • 1
  • Welcome to StackOverflow, please us the `{}` button in the editor to format your code and include an example output – Caridorc Jan 28 '23 at 10:52
  • Returns `0 + 2.j` where the op wants only `2.j` – MatBailie Jan 28 '23 at 10:54
  • You can use `np.array2string` with my `complex_to_string` function to get the desired output also for `2j` – Caridorc Jan 28 '23 at 10:55
  • 2
    Please read [answer] and [edit] your answer to contain an explanation as to why this code would actually solve the problem at hand. Always remember that you're not only solving the problem, but are also educating the OP and any future readers of this post – chrslg Jan 29 '23 at 00:19
  • Good answer! I agree with @chrslg: It would be _even better_ if it came with an explanation! Please see [answer]. Welcome to Stack Overflow. – Pranav Hosangadi Jan 29 '23 at 07:20
  • FWIW, it can be even shorter if you reduce the lambda by multiplying a boolean -- the `if x.imag != 0` -- by the string `f' + {x.imag}'`. i.e. `lambda x: f'{x.real}' + bool(x.imag) * f' + {x.imag}j'` – Pranav Hosangadi Jan 29 '23 at 07:23