将多维复杂 numpy 数组打印到字符串的 Pythonic 方法

Question

我有一个这样定义的 3D 复杂 numpy 数组:

> import numpy as np
> a = np.random.rand(2,3,4) + np.random.rand(2,3,4) * 1j
> a
array([[[ 0.40506245+0.68587874j,  0.74700976+0.73208816j,
      0.42010818+0.31124884j,  0.27181199+0.54599156j],
    [ 0.29457621+0.34057513j,  0.82490182+0.63943948j,
      0.46887722+0.12734375j,  0.77184637+0.21522095j],
    [ 0.67774944+0.8208908j ,  0.41476702+0.85332392j,
      0.10084665+0.56146324j,  0.71325041+0.77306548j]],

   [[ 0.77843387+0.23660274j,  0.23671262+0.63997834j,
      0.60831419+0.41741288j,  0.53870756+0.13747055j],
    [ 0.12477767+0.54603678j,  0.60537090+0.89208227j,
      0.16027151+0.17575777j,  0.18801875+0.27282324j],
    [ 0.82308271+0.97238411j,  0.47458327+0.75200695j,
      0.16085009+0.60620705j,  0.79766571+0.76470634j]]])

我需要将它打印成一个字符串 s，格式有点像 MATLAB，我发现最好的方法如下:(对我来说，描述格式与此代码)

> s = ''
> for k in range(a.shape[2]):
>   for j in range(a.shape[1]):
>     for i in range(a.shape[0]):
>       s += str(a[i,j,k].real) + ' '
>   for j in range(a.shape[1]):
>     for i in range(a.shape[0]):
>       s += str(a[i,j,k].imag) + ' '

我对这段代码不满意，它对我来说看起来不太“pythonic”(我来自 C++，对 Python 了解不多)。我确信 Python 提供了一些可以在这里使用的良好语法(例如列表理解)，但我对它不是很熟悉。

因此我的问题如下:如何改进此代码以使其更符合 Python 风格？

编辑:此 3 维数组被视为 2×3 复矩阵的数组。该格式包括打印第一个矩阵的实部，然后是它的虚部，然后以这种方式遍历每个矩阵。

这是在 MATLAB 中运行这段代码时得到的格式:

> a = rand(2,3,4) + rand(2,3,4) * 1i;
> s = sprintf('%g %g ', [real(a) imag(a)]);

我的主要目标是与这种格式兼容。

Answer 1

字符串连接通常使用 join 完成:

s += str(a[i,j,k].imag) + ' '

可以替换为

s += ' '.join(str(a[i,j,k].imag))

全局应用，1-liner 可以是:

s = ' '.join(' '.join(str(a[i,j,k].real) for j in range(a.shape[1]) for i in range(a.shape[0])) + ' ' + ' '.join(str(a[i,j,k].imag) for j in range(a.shape[1]) for i in range(a.shape[0])) for k in range(a.shape[2]))

不太清楚。我会保留 for ... k 循环并这样写:

s = ''
for k in range(a.shape[2]):
    s += ' '.join(str(a[i,j,k].real) for j in range(a.shape[1]) for i in range(a.shape[0]))
    s += ' '
    s += ' '.join(str(a[i,j,k].imag) for j in range(a.shape[1]) for i in range(a.shape[0]))
    s += ' '

编辑

这个比较重，numpy有很多工具。这是一个更简单的版本。 第一行重新格式化矩阵以简化第二行的工作:

b = [numpy.vstack((a.real.T[i], a.imag.T[i])) for i in range(a.shape[2])]
s = ' '.join(str(d) for x in b for d in x.flat)

编辑2

还可以简化

' '.join([str(x) for x in np.hstack((a.T.real, a.T.imag)).flat])