mohamed82008 / differentiablefactorizations.jl Goto Github PK

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]