将复杂的NumPy数组转换为实部和虚部的(n,2)数组

Question

我有一个复数NumPy数组，我想将其转换为连续的NumPy数组，并将实部和虚部分开。

这个

import numpy

u = numpy.array([
    1.0 + 2.0j,
    2.0 + 4.0j,
    3.0 + 6.0j,
    4.0 + 8.0j
    ])

u2 = numpy.ascontiguousarray(numpy.vstack((u.real, u.imag)).T)

虽然这个方法可行，但转置、竖向堆叠和转换为连续数组可能是多余的步骤。

是否有原生的NumPy函数可以帮我完成这个操作？

- Nico Schlömer

Answer 1

我已经使用perfplot比较了建议的解决方案。对于大向量大小，常数时间视图。

u.view("(2,)float")

确实是最快的变体。对于较小的n，

numpy.column_stack([u.real, u.imag])

可能是最好的（最干净，最快）选择。

重现绘图的代码：

import numpy
import perfplot


def setup(n):
    return numpy.random.rand(n) + 1j * numpy.random.rand(n)


def column_stack(u):
    return numpy.column_stack([u.real, u.imag])


def dstack(u):
    return numpy.dstack((u.real, u.imag))[0]


def concatenate(u):
    return numpy.concatenate(
        (numpy.array(u.real, ndmin=2).T, numpy.array(u.imag, ndmin=2).T), axis=1
    )


def stack(u):
    return numpy.stack((u.real, u.imag), -1)


def view(u):
    return u.view("(2,)float")


perfplot.save(
    "out.png",
    setup=setup,
    kernels=[column_stack, dstack, concatenate, stack, view],
    n_range=[2 ** k for k in range(22)],
)

- Nico Schlömer

Answer 2

您可以使用column_stack来将两个1维数组作为列，组合成一个2维数组。

In [9]: np.column_stack((u.real,u.imag))
Out[9]: 
array([[ 1.,  2.],
       [ 2.,  4.],
       [ 3.,  6.],
       [ 4.,  8.]])

- Nickil Maveli

---

我刚刚检查了一下 :) 它避免了丑陋的索引！ - Colonel Beauvel

Answer 3

没有任何一种替代方案是本地的或可以避免重塑、转置等操作。

例如，在内部，column_stack将其输入转换为2D“列”数组。实际上它正在执行以下操作：

In [1171]: np.concatenate((np.array(u.real,ndmin=2).T,np.array(u.imag,ndmin=2).T),axis=1)
Out[1171]: 
array([[ 1.,  2.],
       [ 2.,  4.],
       [ 3.,  6.],
       [ 4.,  8.]])

vstack将其输入通过atleast_2d(m)，确保每个输入都是1行2D数组。而np.dstack则使用atleast_3d(m)。

新增加的函数是np.stack

In [1174]: np.stack((u.real,u.imag),-1)
Out[1174]: 
array([[ 1.,  2.],
       [ 2.,  4.],
       [ 3.,  6.],
       [ 4.,  8.]])

它使用None索引来纠正连接的维度；有效地：

np.concatenate((u.real[:,None],u.imag[:,None]),axis=1)

所有的都最终使用np.concatenate；它和np.array是仅有的编译连接函数。

另一个技巧是使用view

In [1179]: u.view('(2,)float')
Out[1179]: 
array([[ 1.,  2.],
       [ 2.,  4.],
       [ 3.,  6.],
       [ 4.,  8.]])

复数值被保存为相邻的两个浮点数。因此，同一数据缓冲区既可以视为纯浮点数，也可以使用此视图作为浮点数的二维数组。与“concatenate”函数相比，这里没有复制。另一个对替代方案的测试是询问当u为2d或更高时会发生什么？

- hpaulj

---

能否对一个 N x M 复杂数组执行相同的操作以避免复制？将其转换为 N x M x 2 的 float64 数组？ - wprins

Answer 4

你可以使用 dstack：

np.dstack((u.real, u.imag))[0]
#Out[210]:
#array([[ 1.,  2.],
#       [ 2.,  4.],
#       [ 3.,  6.],
#       [ 4.,  8.]])

- Colonel Beauvel