# Fast square root of 3 x 3 matrix

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!

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

6 Likes

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

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

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

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

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

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.

Cholesky is definitely what you want here.

4 Likes