Is there a quicker way to find matrix square root for matrices from a 3-dimensional array?

Viewed 863

I am New at Python and Numpy. Given a 3-dimensional array containing matrices I need to find a square root for each matrix in this array and store the result as another 3-dimensional array containing these roots.

I did the following (Example for shape (2,2,2)):

import numpy as np
import scipy.linalg

matrices = np.array([
    [
        [1, 0],
        [0, 1]
    ],
    [
        [33, 24],
        [48, 57]
    ],
])

# Finding roots
roots = []
for i in range(np.size(matrices[:,0, 0])): 
    root = scipy.linalg.sqrtm(matrices[i,:, :])
    roots.append(root)
roots = np.array(roots)

And got the output

roots

array([[[1., 0.],
        [0., 1.]],

       [[5., 2.],
        [4., 7.]]])

...which is perfectly fine, but slow if there are many matrices in the array (e.g. shape is (3000,n,n) ) . Is there a more efficient way of doing this?

0 Answers
Related