Implement more matrix functions

[jiahao]

Higham and Deadman have published a list of software implementing algorithms for matrix functions, showing that Julia's library for these functions could be improved greatly relative to the state of the art.

The purpose of this issue is to discuss and track the implementation of matrix functions.

From Higham and Deadman's list:

General function evaluations

• General matrix function with derivatives of the underlying scalar function available: Schur–Parlett algorithm (Davies and Higham, 2003)
• Function of a symmetric or Hermitian matrix by diagonalization
• Condition number estimate for general functions: cond(f,A)
• matrix function - vector product evaluation: f(A) * b

Specific functions

• expm: scaling and squaring algorithm (Al-Mohy and Higham, 2009)
• logm: inverse scaling and squaring algorithm (Al-Mohy, Higham, and Relton, 2012,
  2013)
• sqrtm:
  • Schur algorithm (Björck and Hammarling, 1983)
  • real version for real matrices (Higham, 1987)
  • blocking algorithm for performance improvements (Deadman, Higham, and Ralha, 2013)
• A^t for real matrix powers: Schur– Padé algorithm (Higham and Lin, 2013) (Fixed the algorithm for powers of a matrix. #21184)
• Matrix unwinding function

Fréchet derivatives

• Exponential: scaling and squaring algorithm (Al-Mohy and Higham, 2009)
• Logarithm: inverse scaling and squaring algorithm (Al-Mohy, Higham, and Relton, 2013)
• General matrix function: complex step algorithm (Al-Mohy and Higham, 2010) or using block 2 × 2 matrix formula.
• A^t for real matrix powers: Schur– Padé algorithm (Higham and Lin, 2013)

Miscellaneous functions

• inv and det for Tridiagonal and SymTridiagonal matrices (Usmani, 1994) (Provide inv and det of Tridiagonal and SymTridiagonal matrices #5358)

[StefanKarpinski]

It would be amazing to include all of these algorithms. Thanks for putting this issue together, @jiahao.

[jiahao]

It's not clear to me what interface should be implemented for the very generic things like "function of a symmetric or Hermitian matrix".

We could define methods for elementary functions until the cows come home, but that is both tedious and namespace polluting (we'd need sinm, cosm, atanhm, etc. to disambiguate from the ordinary forms which broadcast over the elements) , and furthermore wouldn't scale well for linear combinations of functions.

Maybe we want something like applym, so that a call like applym(A->exp(A)+sin(2A)-cosh(4A), A) would evaluate the matrix expm(A)+sinm(2A)-coshm(4A)?

[StefanKarpinski]

This is one of the reasons why having all these functions be vectorized is kind of shitty. I suspect that @JeffBezanson will vehemently agree. There could be a MatrixFunctions module that you have to import from to get the matrix variants with all the same names.

[jiahao]

Sure, but that won't help with the fact that a complicated matrix function like A->expm(A)+sinm(2A)-coshm(4A) can be computed more quickly all at once than computing (expm(A))+(sinm(2A))-(coshm(4A)) by working out the individual elementary functions separately, since the underlying Schur-Parlett algorithm is essentially the same.

[StefanKarpinski]

Yes, that's a very good point.

[stevengj]

Trefethen's contour-integration method is also neat for mostly analytic functions (with known singularities), though I don't know how it compares to the state of the art (in performance and generality).

[jiahao]

@stevengj it looks like that paper is Ref. 32 of the Higham and Deadman report, and isn't currently available in any major package (although it has Matlab code).

[simonbyrne]

This seems like a good candidate for a GSOC project.

As far as an interface goes, I would suggest a macro

B = @matfun exp(A)+sin(2A)-cosh(4A)

A naive approach could simply replace the functions (e.g. exp -> expm), a more advanced approach could try to identify the common decompositions.

[GunnarFarneback]

The item "[sqrtm] real version for real matrices (Higham, 1987)" is checked. I can't see that algorithm in base. I made an attempt on it about a year ago but concluded that the complexity just wasn't worth it. I've put up the code as a gist if someone wants to have a look.

[jiahao]

Thanks for catching that. I must have ticked it by accident when writing up this issue.

[dlfivefifty]

Just tried to call logm, but wasn't implemented. Had to go to Mathematica

[fabianlischka]

Whoever tackles this issue might find this paper helpful for testing:
Higham, Deadman: Testing Matrix Function Algorithms Using Identities, March 2014

[StefanKarpinski]

Great reference – thanks, @fabianlischka.

[mfasi]

The logm function has been added by pull request #12217.

[ViralBShah]

Also cc: @higham @alanedelman

[antoine-levitt]

Unless I'm missing something, not really: differentials of matrix functions are tricky due to non-commutativity. @sethaxen it would be strange to have this live in LinearAlgebra, can you maybe refactor LinearAlgebra to expose the functions that you need and use them in ChainRules?

[dlfivefifty]

I think @antoine-levitt is right. Looks like differentiating f(A(t)) using Taylor series is a mess. Cauchy integrals formula is a bit easier: if we differentiate

f(A(t)) = ∮ f(ζ) inv(ζ*I - A(t)) dζ

then we get

d/dt f(A(t)) = ∮ f(ζ) inv(ζ*I - A(t)) * dA/dt * inv(ζ*I - A(t))  dζ

which I don't believe can be reduced to matrix factorisations.

Really interested to know if anyone has a reference for differentiating matrix functions. Don't have a copy of Higham handy but nothing pops out of the table of contents.

[sethaxen]

Possibly but isn't it possible to define these in terms of existing matrix functions?

Sort of. For any matrix function f we can compute its Frechet derivative by calling f on a block matrix with a particular structure:

function frechet(f, A, ΔA)
    n = checksquare(A)
    Ablock = [A ΔA; zero(A) A]
    Yblock = f(Ablock)
    Y = Yblock[1:n,1:n]
    ΔY = Yblock[1:n,n+1:2n]
    return Y, ΔY
end

But this is ~8x the cost of f, whereas the algorithms by Higham in the OP are more efficient (I think ~3x or less).

@sethaxen it would be strange to have this live in LinearAlgebra

I agree, I only ask because given this issue it seemed there was (old) interest in them being in base, and the Frechet derivative of exp is included in scipy.

can you maybe refactor LinearAlgebra to expose the functions that you need and use them in ChainRules?

Perhaps. While the most efficient way to implement will be completely in parallel, we can get close by reusing all matrix products. So a nice refactor might be to move the work of exp! into a function that also returns all intermediate matrix products, which I could then reuse. But it would also be nice to reduce the number of allocations in exp!, in which case its Frechet derivative would need to be computed in parallel and duplicate a lot of code.

[antoine-levitt]

Really interested to know if anyone has a reference for differentiating matrix functions. Don't have a copy of Higham handy but nothing pops out of the table of contents.

Basically continue what you were doing with the integral formula, expand in an eigenbasis and integrate the poles explicitly, and you get divided differences (f(lambda_n) - f(lambda_m)) / (lambda_n - lambda_m), with the convention (f(x)-f(y))/(x-y) = f'(x) when x==y. It's probably discussed in Higham somewhere, I think it's called Daleskii-Krein there (I know it from quantum mechanics, where it's known under the generic name of "perturbation theory"). It's a bit tricky to implement in the case of close eigenvalues for generic functions f because of the divided differences (what I've been doing in https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/Smearing.jl#L44 is to use forward-diff to use f' and f'' for x and y close together, although that still doesn't get you to O(eps)). For exp/log you can use expm1 and log1p. All this is for normal matrices, @sethaxen I imagine the method you're looking at is more robust for non-normal?

It was suggested above to split matrix functions (more than exp/log/sqrt) to a separate MatrixFunctions.jl package; I think that would be a very natural place to put derivatives.

[dlfivefifty]

@sethaxen's method is probably a bad idea for normal matrices as Ablock will not be normal.

[sethaxen]

Really interested to know if anyone has a reference for differentiating matrix functions. Don't have a copy of Higham handy but nothing pops out of the table of contents.

Section 3.2 of Higham's Functions of Matrices gives several approaches for computing Frechet derivative:

• the block method given above (Eq 3.16)
• Daleckii-Krein using the eigendecomposition for normal matrices (Corollary 3.12)
• Using the Schur decomposition, if you know the Frechet derivative of f(T) for triangular T (Problem 3.2)
• Computing the Frechet derivatives of the primitive functions of the implementation of f and composing with the chain rule, though this needs to be checked for accuracy.

@sethaxen I imagine the method you're looking at is more robust for non-normal?

For automatic differentiation, we indeed only want to define a rule if it works for all matrices, and the latter approach should be fine. In fact, the Frechet derivative of the matrix exponential by Higham is precisely exactly this, just checked for accuracy. I don't think we can autodiff exp with any AD (Zygote can't handle in-place modification, and ForwardDiff and ReverseDiff require types other than BlasFloat), so a rule is needed. The reverse-mode derivative (pullback) is the Frechet derivative applied to the adjoint of the input.

It was suggested above to split matrix functions (more than exp/log/sqrt) to a separate MatrixFunctions.jl package; I think that would be a very natural place to put derivatives.

The issue with that, raised above, is type piracy. I guess that could be worked around by re-appending the m suffix or using an interface like matfun(sin, A), but sin(A) is cleaner. For Hermitian, almost all of the matrix functions share a common form (using Daleckii-Krein), so it would be odd to split a subset of them into a separate package. The advantage though is such a package could provide multiple algorithms for computing each f and also bundle the ChainRules rules side-by-side, which is always preferable to putting them in the ChainRules package. There are also cases like ExponentialUtilities.exp! where the code for LinearAlgebra.exp! was almost exactly duplicated just to use a cache.

[antoine-levitt]

Computing the Frechet derivatives of the primitive functions of the implementation of f and composing with the chain rule, though this needs to be checked for accuracy.

I see. In the applications I'm interested in the function f is often the Heaviside function, so that's not an option.

Regarding a separate package, clearly exp/sin/log/sqrt are special because they're in base, but there are other functions one might want to use, so a general matfun(f, A) is needed, and it would be natural to put the rules there. But indeed the question of where to put the rules for exp/sin/log/sqrt remains... I would say they belong to ChainRules, and code duplication is necessary in any case (as code reuse for the function itself and its derivatives looks hard from what you say).

[ViralBShah]

This is best done in external packages. Moving to a discussion so that we have the list.

[aravindh-krishnamoorthy]

Hey, I'm just done with cbrt in PR #50661 based on Smith, M. I. (2003). A Schur Algorithm for Computing Matrix pth Roots. In SIAM Journal on Matrix Analysis and Applications (Vol. 24, Issue 4, pp. 971–989). Comments and suggestions are welcome!

I'm new to Julia ($2\text{--}3$ months), but I've prior experience in vectorisation and numerical LA. I can get started with expm_frechet and follow through with the remainder of the "Frechet derivatives" list. Is there still interest in implementing these?

[ViralBShah]

Certainly! These packages can live in the JuliaMath org.

[antoine-levitt]

Yes, it'd be awesome to have something like a package that implements matrix_function(f, A) and has derivatives through compatibility with ChainRules/ForwardDiff.

[aravindh-krishnamoorthy]

Yes, it'd be awesome to have something like a package that implements matrix_function(f, A) and has derivatives through compatibility with ChainRules/ForwardDiff.

Thank you for the encouragement. Pipe dream is to implement something like MatrixCalculus.org in Julia where you could hand in a function along with the parameters and get back the value F0 and the Fréchet derivative L0 at X=X0:

F0, L0 = matrix_function(3*X + sin(X) - C, X, X0)

But, I haven't a clue on how to get started with this. Perhaps a good first step to understand this better is to just implement

F0, L0 = matrix_function1(f, X0)

based on [1] (function evaluation) and [2] (Fréchet derivative), and then see how to go on from there....

[1]	Davies, Philip I., and Nicholas J. Higham. "A Schur-Parlett algorithm for computing matrix functions." SIAM Journal on Matrix Anal1 ysis and Applications 25.2 (2003): 464-485.
[2]	Al-Mohy, Awad H., and Nicholas J. Higham. "The complex step approximation to the Fréchet derivative of a matrix function." Numerical Algorithms 53.1 (2010): 133-148.

[antoine-levitt]

I don't think getting back the full Frechet derivative is super helpful, as that's probably too big for most applications; much more useful is the jvp (L * delta_X) and the vjp (delta_X * L), which are what ChainRules needs for its frule and rrule. I'm not a big fan of complex step (it's cute, but it's just simpler to use the real one)

I implemented it for my own purposes for Symmetric matrices, see below. I imagine something similar holds for non-symmetric, but I've never worked with those directly.

function matrix_function(f, H)
    E = eigen(Symmetric(H))
    D = E.vectors * Diagonal(f.(E.values)) * E.vectors'
end

@views function dmatrix_function(f, H, δH)
    E = eigen(Symmetric(H))
    occ = f.(E.values)
    N = size(H, 1)
    δD = zeros(eltype(δH), N, N)
    # δD = ∑nm (fn-fm)/(en-em) |n><m| δHnm
    for n = 1:N
        for m = 1:N
            (occ[n] == occ[m]) && continue
            δD .+= ((occ[n]-occ[m])/(E.values[n]-E.values[m]) * dot(E.vectors[:, m], δH * E.vectors[:, n])) .* (E.vectors[:, m] .* E.vectors[:, n]')
        end
    end
    δD
end


# makes matrix_function compatible with ForwardDiff
function matrix_function(f, H::Matrix{T}) where {T <: ForwardDiff.Dual}
    H_value = ForwardDiff.value.(H)
    D = matrix_function(f, H_value)
    DT = ForwardDiff.Dual{ForwardDiff.tagtype(T)}
    res = [dmatrix_function(f, H_value, ForwardDiff.partials.(H, α)) for α = 1:ForwardDiff.npartials(T)]
    map((val, δval...) -> DT(val, δval...), D, res...)
end

EDIT: this function is wrong, see below

[aravindh-krishnamoorthy]

I don't think getting back the full Frechet derivative is super helpful, as that's probably too big for most applications; much more useful is the jvp (L * delta_X) and the vjp (delta_X * L), which are what ChainRules needs for its frule and rrule. I'm not a big fan of complex step (it's cute, but it's just simpler to use the real one)

Hello @antoine-levitt, I'm not sure I fully understand your comment. Do you mean the Jacobian? Perhaps we must discuss this in another discussion thread?

In fact, I couldn't also understand what is being computed in the function dmatrix_function above, which is equivalent to:

function dmatrix_function(f, H, dH)
    N = size(H,1)
    E = eigen(Symmetric(H))
    U = E.vectors
    occ = f.(E.values)
    F = [if m == n 0 else (occ[m] - occ[n])/(E.values[m]-E.values[n]) end for m=1:N, n=1:N]
    δD = U*(F.*(U'*dH*U))*U'
end

On the other hand, the matrix $L_f(\boldsymbol{H}, \boldsymbol{dH}),$ which I could quickly derive, can be utilized for symmetric matrices:

function Lf(f, H, dH)
    N = size(H,1)
    E = eigen(Symmetric(H))
    U = E.vectors
    fL = Diagonal(f.(E.values))
    dfL = Diagonal(ForwardDiff.derivative.(f, E.values))
    F = [if m == n 0 else 1/(E.values[n]-E.values[m]) end for m=1:N, n=1:N]
    L = U*dfL*Diagonal(U'*dH*U)*U' + U*(F.*(U'*dH*U))*fL*U' + U*fL*(F'.*(U'*dH*U))*U'
end

and works for symmetric dH and non-singular, symmetric H with non-repeating eigenvalues. Furthermore, the above function can be tested, based on [2] and [3] (below), as follows:

julia> A = [1 2; 2 5]/10 ;
julia> dA = [4 5; 5 7]*1e-6 ;
julia> A1 = [A dA; zeros(2,2) A] ;

julia> exp(A1)[1:2,3:4]
2×2 Matrix{Float64}:
 5.82694e-6  8.46229e-6
 8.46229e-6  1.3176e-5

julia> Lf(exp, A, dA)
2×2 Matrix{Float64}:
 5.82694e-6  8.46229e-6
 8.46229e-6  1.3176e-5

It also works for sin, sqrt, cbrt, log, and several other functions. (Please see also the symmetric functions linked in the post by @sethaxen below. Those functions essentially compute Lf(f, H, dH) in a computationally and numerically efficient way.)

I'm not sure how to use the functions in your comment above to reach the same output. This is because ForwardDiff doesn't seem to accept matrix inputs other than for jacobian. Would the Fréchet derivatives for arbitrary matrices still be useful?

[3]	Awad H. Al-Mohy and Nicholas J. Higham (2009) Computing the Frechet Derivative of the Matrix Exponential, with an application to Condition Number Estimation. SIAM Journal On Matrix Analysis and Applications., 30 (4). pp. 1639-1657. ISSN 1095-7162

[antoine-levitt]

Sorry about that, I took my dmatrix_function without looking from a code that used it with f piecewise constant, so there was no diagonal terms. The F should have a derivative on the diagonal as in your code, of course.

[sethaxen]

FYI, ChainRules implements forward- and reverse-mode rules for most of the matrix functions defined in base Julia.

• functions of symmetric matrices: https://github.com/JuliaDiff/ChainRules.jl/blob/df672c3b86aed6ec2638f9e56e279c26ebe92828/src/rulesets/LinearAlgebra/symmetric.jl#L302-L462
• functions of dense matrices: https://github.com/JuliaDiff/ChainRules.jl/blob/df672c3b86aed6ec2638f9e56e279c26ebe92828/src/rulesets/LinearAlgebra/matfun.jl

The latter uses the following signatures:

• _matfun(f, A) -> (Y, intermediates): compute the matrix function Y = f(A), returning intemediates that may be reused to compute the Frechet derivative
• _matfun_frechet(f, E, A, Y, intermediates): compute the Frechet derivative of the matrix function in the direction of E.

I don't think getting back the full Frechet derivative is super helpful, as that's probably too big for most applications; much more useful is the jvp (L * delta_X) and the vjp (delta_X * L), which are what ChainRules needs for its frule and rrule. I'm not a big fan of complex step (it's cute, but it's just simpler to use the real one)

The Frechet derivative is the pushforward/jvp used in forward-mode AD, i.e. the action of the Jacobian on the direction matrix. As a pushforward, it is a linear map whose adjoint is the pullback/vjp used in reverse-mode AD. A really convenient property of matrix functions is that if the pushforward (Frechet derivative) at A is (f_*)_A(ΔA), then the pullback at A is (f^*)_A(ΔY) = ((f_*)_A(ΔY'))'.

So the Frechet derivative of a matrix function is a really useful primitive that can be used for both forward- and reverse-mode AD. It might be worth starting from the implementations in ChainRules. My advice would be to keep the package very lightweight.

[aravindh-krishnamoorthy]

Hello @sethaxen, thank you for this really good input.

My advice would be to keep the package very lightweight.

Actually, matfun.jl seems like a good home for these functions to me. It's just an extension of the current exp implementations to arbitrary functions (hopefully).

Do you think a separate package is needed for this?

[antoine-levitt]

OK sorry I misunderstood what you meant by Fréchet derivative, I thought the full 4th order tensor.

LinearAlgebra does not depend on chainrules. Though I guess there could be a matrix_function and a frechet_matrix_function in Linear Algebra, and chainrules would just wrap that in its own API. Generally for non trivial changes it's easier to start with a separate package, collect feedback and use cases, iron out the bugs etc then move the code to the julia codebase. This is especially true here because there are many potential numerical instabilities (even for hermitian matrices, almost repeated eigenvalues are tricky) that have to be carefully checked before getting into base julia ; a separate package is simpler for this.

[antoine-levitt]

To clarify: the issue for numerical stability (in the hermitian case, the non-hermitian case is even worse...) is to compute (f(x)-f(y))/(x-y) stably (meaning, with a relative precision of order eps()). If done naively, the error is eps()/|x-y|. You can switch to f'(x), which has error |x-y|. So you should switch when |x-y| < sqrt(eps()). Actually, a better formula is to use (f'(x)+f'(y))/2, which is accurate to |x-y|^2, and therefore the switching point is cbrt(eps()), for an overall precision of eps()^(2/3). To do better, you need higher derivatives (or wider stencils). A special case is exp, where you can factor the exponential and use expm1. We use these tricks in https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/Smearing.jl (function occupation_divided_difference)

[sethaxen]

Actually, a better formula is to use (f'(x)+f'(y))/2, which is accurate to |x-y|^2, and therefore the switching point is cbrt(eps()), for an overall precision of eps()^(2/3).

Agreed, ChainRules uses the same approach: https://github.com/JuliaDiff/ChainRules.jl/blob/df672c3b86aed6ec2638f9e56e279c26ebe92828/src/rulesets/LinearAlgebra/symmetric.jl#L433-L444

I agree also that this would be better to start as its own package. If it's lightweight, then ChainRules could also depend on it instead.

The following API functions make sense:

• matrix_function(f, A)
• frechet_matrix_function(f, A, E)

but to be useful for reverse-mode AD, I suggest including the following as well:

matrix_function_with_cache(f, A) = (matrix_function(f, A), nothing)
frechet_matrix_function(f, A, cache) = frechet_matrix_function(f, A, E)

One could then overload matrix_function_with_cache for a specific function to return intermediates necessary for computing the Frechet derivative.

Similarly, to be really efficient for forward-mode AD, you might want to include the following:

function matrix_function_with_frechet(f, A, E)
    Y, cache = matrix_function_with_cache(f, A)
    return Y, frechet_matrix_function(f, A, cache)
end

This could then be overloaded for specific matrix functions to avoid storing all intermediates.

[aravindh-krishnamoorthy]

Thank you very much for making me feel welcome with this valuable discussion. I'll get started on this. I'd also like to add support for Fréchet derivatives of scalar functions of matrix variables, which can be obtained independently of $\boldsymbol{E}.$ It's a simple implementation and might be useful in some cases.