numpy ndarray：僅當不為零時才打印虛部

Question

在這里繼續這個問題，我想問一下如何打印一個復雜的 numpy 數組，僅當它不為零時才打印虛部。 （這也適用於實際部分。）

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

預計 Output：

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

Answer 1

您可以使用np.isclose檢查數字是否接近於零，並使用.real和.imag屬性以不同方式訪問實部和復數部分，然后您可以編寫遞歸 function 進行打印：

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]]

這適用於任意嵌套的 arrays。

---

感謝@Koushik 提到內置的np.array2string ，使用它解決方案變得更簡單：

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}))

與相同的 output。

Answer 2

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}'}))