numpy: converting array of vectors to array of symmetric matrices-CodePudding

I receive an array of 1e6 <= n <= 1e10 vectors, which are the symmetric components of a matrix. Each row of my array contains the components in the following order:

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

Thus, the shape of the array is (n, 6)

In order to efficiently calculate some things, I need the original matrix, thus I want to get an array of shape (n, 3, 3).

So far, I was not able to find a faster solution than rewriting the array row by row:

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

I wonder if there is a more efficient way to transform this kind of data? vec is not needed anymore afterwards and could be overwritten in the transformation process.

edit: After writing this question, it came to me that I can do the lookup and a reshape in one 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)

This is already much faster:

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

is this already the most optimal way to do it?

CodePudding user response：

Using the tuple based indexing and numpy reshaping probably is as fast as you can get here:

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