numpy: converting array of vectors to array of symmetric matrices

Question

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?

Answer 1

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