numpy：将向量数组转换为对称矩阵数组

Question

我收到一个1e6 <= n <= 1e10向量的数组，它们是矩阵的对称分量。 我的数组的每一行都按以下顺序包含组件：

[11, 22, 33, 23, 13, 12]

因此，数组的形状为(n, 6)

为了有效地计算一些东西，我需要原始矩阵，因此我想得到一个形状为(n, 3, 3)的数组。

到目前为止，我找不到比逐行重写数组更快的解决方案：

def vec_to_mat(vec):
    """Convert an array of shape (N, 6) to (N, 3, 3)"""
    mat = np.empty((vec.shape[0], 3, 3))
    for i, s in enumerate(vec):
        mat[i] = np.array(s[[0, 5, 4, 5, 1, 3, 4, 3, 2]]).reshape((3, 3))
    return mat


# Example usage:
>>> x = np.array([[1,2,3,4,5,6], [4,6,8,2,4,6]])
>>> vec_to_mat(x)
array([[[1., 6., 5.],
        [6., 2., 4.],
        [5., 4., 3.]],

       [[4., 6., 4.],
        [6., 6., 2.],
        [4., 2., 8.]]])

我想知道是否有更有效的方法来转换这种数据？ vec之后不再需要，并且可以在转换过程中被覆盖。

编辑：写完这个问题后，我想到我可以在一个 go 中进行查找和重塑：

def vec_to_mat_2(vec):
    """Convert an array of shape (N, 6) to (N, 3, 3)"""
    return vec[:, [0, 5, 4, 5, 1, 3, 4, 3, 2]].reshape(-1, 3, 3)

这已经快得多了：

%timeit vec_to_mat(np.random.normal(4, 17, (10000, 6)))
43.3 ms ± 234 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit vec_to_mat_2(np.random.normal(4, 17, (10000, 6)))
2.42 ms ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这已经是最好的方法了吗？

Answer 1

使用基于元组的索引和 numpy 重塑可能与您在这里一样快：

def vec_to_mat(vec):
    """Convert an array of shape (N, 6) to (N, 3, 3)"""
    mat = vec[:, (0, 5, 4, 5, 1, 3, 4, 3, 2)].reshape(-1, 3, 3)
    return mat

x = np.array([[1,2,3,4,5,6], [4,6,8,2,4,6]])
vec_to_mat(x)
>>> array([[[1, 6, 5],
        [6, 2, 4],
        [5, 4, 3]],

       [[4, 6, 4],
        [6, 6, 2],
        [4, 2, 8]]])