Fast square root of 3 x 3 matrix

General Usage
,
1

Hi,

I’m currently using sqrt() from LinearAlgebra.jl to compute the square-root of 3 x 3 matrices repeatedly. But this has created a performance bottleneck. So I was wondering if there an implementation of a matrix square root that specializes for 3 x 3 matrices like the ones in StaticArrays.jl(which are limited to 2 x 2 matrices). FYI I’m dealing with real, symmetric and positive-definite matrices.

Thanks!

2

I think the best result will be building your own from the eigen decomposition which StaticArrays has a specialized method for 3x3

6 Likes
3

Depending on exactly what you need, you might be able to use the Cholesky decomposition, which computes a lower triangular matrix L or upper triangular matrix U such that L*L' == U'*U == A. I believe StaticArrays has specialized implementations for cholesky(A::SMatrix). This is similar to a sqrt in some ways, although not in the sense that S = sqrt(A) such that S*S == A (no transpose needed here and S is PSD if A is).

3 Likes
4

You will get some StaticArrays speedup by default, as long as you remember to wrap your array in a Symmetric:

r_ = rand(3, 3)
r = r_ + r_'

R = r^2
Rh = Symmetric(R)
Rs = SMatrix{3,3}(R)
Rsh = Symmetric(Rs)

@btime sqrt($R);  # normal Matrix
  3.538 μs (14 allocations: 2.17 KiB)

@btime sqrt($Rh);  # Symmetric Matrix
  3.413 μs (14 allocations: 2.17 KiB)

@btime sqrt($Rs);  # SMatrix 
  3.487 μs (16 allocations: 2.31 KiB)

@btime sqrt($Rsh);  # Symmetric SMatrix
  726.316 ns (0 allocations: 0 bytes)

So you don’t get any benefit of StaticArrays unless you also specify Symmetric as well.

8 Likes
5

Right, it would be good to know the context here. In many problems where you conceptually might want a matrix square root, there is a way to reformulate it to use Cholesky, which is way cheaper to compute.

2 Likes
6

The Symmetric LinearAlgebra sqrt is implemented with eigen, so with a Symmetric SMatrix you don’t have to build it yourself.

If speed is more critical than precision you can do considerably better than that, however. This function only works for sufficiently distinct eigenvalues, so needs to be complemented with some special cases for colliding eigenvalues.

function f(x)
    l = eigvals(x)
    J = Symmetric(SMatrix{3, 3}(I))
    A1 = Symmetric(x - l[1] * J)
    A2 = Symmetric(x - l[2] * J)
    A3 = Symmetric(x - l[3] * J)
    B3 = Symmetric(A1 * A2) / ((l[3] - l[1]) * (l[3] - l[2]))
    B2 = Symmetric(A1 * A3) / ((l[2] - l[1]) * (l[2] - l[3]))
    B1 = Symmetric(A2 * A3) / ((l[1] - l[2]) * (l[1] - l[3]))
    return sqrt(l[1]) * B1 + sqrt(l[2]) * B2 + sqrt(l[3]) * B3
end

With a Symmetric{Smatrix} I measure it 5-6 times faster than sqrt and it probably leaves a fair amount of micro-optimizations on the table.

3 Likes
7

Thanks this works, even when R is a MMatrix{3, 3}.

8

I need to factorize a 3 x 3 matrix A such that B*B’ = A and then B is used elsewhere in the code. The problem with the Cholesky decomposition is that now I’m not sure that my A will always be positive definite. But it is true that Cholesky is much faster than computing the square-root even with StaticArrays and Symmetric.

9

Cholesky is definitely what you want here.

4 Likes