mohamed82008 / DifferentiableFactorizations.jl

Differentiable matrix factorizations using ImplicitDifferentiation.jl.
MIT License
30 stars 1 forks source link

add Schur(A,B) decomposition #5

Closed thorek1 closed 2 years ago

thorek1 commented 2 years ago

Hi,

I tried to implement schur(A,B) the way you did for the eigen decomposition but the solver doesn't find a solution.

My guess is that numerical stability in the decomposition might be an issue here. What's your experience with the other decompositions? Any hints?

Here is some code to show what I mean:

using LinearAlgebra, SparseArrays, ChainRulesCore, ImplicitDifferentiation, ComponentArrays, Zygote, Krylov, FiniteDifferences

D = [
 0.0  -1.0         -0.213967   0.0        0.0       0.0
 0.0   0.0          0.0        0.0        0.0       0.0
 0.0   0.0         -0.77138    0.0        0.335701  0.404972
 0.0   0.0          0.636374   0.0        0.40692   0.490887
 0.0   0.00173608   0.0       -0.0194976  0.683867  0.0
 1.0   0.0          0.0        0.0        0.0       0.0
]

E =[
 -0.0  -1.001  -0.0       -1.42321   1.0       -0.0
 -0.9  -0.0    -0.0        1.0      -0.0       -0.0
 -0.0  -0.0    -0.694242  -0.0       0.335701   0.607458
 -0.0  -0.0     0.572737  -0.0       0.40692    0.73633
 -0.0  -0.0    -0.0       -0.0       0.683867  -0.0
  0.0   0.0     0.0        1.0       0.0        0.0
]

comp_vec(A, B) = ComponentVector((; A, B))

function ChainRulesCore.rrule(::typeof(comp_vec), A, B)
    out = comp_vec(A, B)
    T = typeof(out)
    return out, Δ -> begin
        _Δ = convert(T, Δ)
        (NoTangent(), _Δ.A, _Δ.B) 
    end
end

function schur_conditions(ab, F)
    (; left, right, S, T) = F
    (; A, B) = ab
    return vcat(
        vec(A - left * S * right'),
        vec(B - left * T * right')
    )
end

function schur_forward(ab)
    (; A, B) = ab

    SCH = schur(A, B)

    left = SCH.left
    right = SCH.right
    S = SCH.S
    T = SCH.T

    # left = sparse(SCH.left)
    # right = sparse(SCH.right)
    # S = sparse(SCH.S)
    # T = sparse(SCH.T)

    # droptol!(left,eps(Float32))
    # droptol!(right,eps(Float32))
    # droptol!(S,eps(Float32))
    # droptol!(T,eps(Float32))

    return ComponentVector(; left, right, S, T)
end

_diff_schur = ImplicitFunction(schur_forward, schur_conditions, Krylov.bicgstab)

#function diff_schur(AB)
#    (;A,B) = AB
#    (; left, right, S, T) = _diff_schur(comp_vec(A, B))
#    return ComponentVector(; left, right , S, T)
#end

AB = ComponentVector(; A = D, B = E)

#f1(AB) = begin
#    A = AB.A
#    B = AB.B
#    diff_schur(A, B)
#end

zjac1 = Zygote.jacobian(_diff_schur, AB)[1]
fjac1 = FiniteDifferences.jacobian(central_fdm(5, 1), _diff_schur, AB)[1]
mohamed82008 commented 2 years ago

Sorry for the delay. I implemented Schur and generalized Schur in #6. It was a bit tricky because of the quasi upper triangular matrix S and how the structure of T and S are related but I think I got it right.

thorek1 commented 2 years ago

thanks a lot. works as intended so far. last piece missing to get my difference equation solver diffable is ordschur, I'll open another issue for that one