Use Real as FC instead of AbstractFloat ?

[j-fu]

First, thank you for this package !

I am currently trying out its possibilities.
This includes differentiating through the calculations. In the type specification you always seem to require that FC is an AbstractFloat. However, this excludes ForwardDiff.Dual which is passed when performing forward mode autodiff. Their common supertype is Real.

So my question: Would it be possible replace AbstractFloat by Real ?

[amontoison]

Hi @j-fu!
Can you give a small example of a linear system where you use ForwardDiff.Dual?
I use AbstractFloat instead of Real to exclude Integer numbers but I didn't know that ForwardDiff.Dual was a subtype of Real.

[j-fu]

Hi, thanks for consideration, here is a simple example:

using LinearAlgebra
using ForwardDiff
using DiffResults
using Krylov

function makematrix(n,X)
    m=zeros(eltype(X),n,n)
    for i=1:n
        m[i,i]=10+X[1]^2
    end
    
    for i=1:n-2
        m[i,i+2]=-X[2]^2
        m[i+2,i]=-X[2]^2
    end
    m
end

function evalderiv(f,X)
    println("    direct evaluation:  f(X)=$(f(X))")
    
    dresult=DiffResults.JacobianResult(ones(1),X)
    ForwardDiff.jacobian!(dresult,f,X)
    println("ForwarDiff evaluation:  f(X)=$(DiffResults.value(dresult))")
    println(" ForwardDiff Jacobian: df(X)=$(DiffResults.jacobian(dresult))")
end

function testfwd(X=[1.0,0.1];n=20)

    fdirect(Y)=[sum(makematrix(n,Y)\ones(eltype(Y),n))]
    println("direct solver:")
    evalderiv(fdirect,X)

    println("cg iteration:")
    fcg(Y)=[sum( cg(makematrix(n,Y),ones(eltype(Y),n))[1])]
    evalderiv(fcg,X)
    
end

May be a bit contrived but it makes the point of matrix entries depending on parameters in a nonlinear way.

Generally I aim at differentiating through nonlinear parameter dependent PDE solves to do some interesting stuff (parameter ID, Bayes, Bifurcation).
With Sparspak.jl there is now a sparse direct solver alternative, but especially for 3D problems having iterative solvers up the sleeve would be good...

[j-fu]

As for the choice of Real as supertype in ForwardDiff, see these discussions: JuliaDiff/ForwardDiff.jl#157 (comment), JuliaDiff/ForwardDiff.jl#66

I am not really deeply in that, though, but as ForwardDiff seems to allow something like Dual{Int} I seem to understand that a Dual is not always an AbstractFloat.

[j-fu]

Bump this :)

[amontoison]

Sorry for the delay @j-fu, I will check again what I can do.

[j-fu]

... no problem, I appreciate your time!

I made a PR to LinearSolve where I propose have unit tests for general numbers. Would be cool to have Krylov.jl there with full functionality.

[michel2323]

@j-fu Would it help you if you had rules for ChainRules.jl such that one would not have to go through the linear solver with Dual numbers but solve the tangent system? This should also give you better stability for the sensitivities as the stopping criteria only applies to the values, not the duals, if you differentiate through it.

[j-fu]

Probably you are right. I need to think through this anyway, need to find some time though.

[michel2323]

@j-fu It We decided to create a new package DiffKrylov.jl to make Krylov.jl differentiable. Here is a snippet though that makes your example work (bugs maybe included).

using Krylov
using ChainRulesCore
using DiffResults
using ForwardDiff
using Krylov
using Test
using LinearAlgebra
using SparseArrays
using ForwardDiff
import ForwardDiff: Dual, Partials, partials, value

function Krylov.cg(_A::Matrix{Dual{T, V, N}}, _b::Vector{Dual{T, V, N}}; options...) where {T, V, N}
    A = Matrix{Float64}(value.(_A))
    dAs = Vector{Matrix{Float64}}(undef, N)
    for i in 1:N
        dAs[i] = Matrix(partials.(_A, i))
    end
    b = value.(_b)
    m = length(b)

    dbs = Matrix{V}(undef, m, N)
    for i in 1:m
        dbs[i,:] = partials(_b[i])
    end
    x, stats = cg(A,b)
    m = length(x)
    dxs = Matrix{V}(undef, m, N)
    px = Vector{Partials{N,V}}(undef, m)
    for i in 1:N
        nb = dbs[:,i] - dAs[i]*x
        dx, dstats = cg(dAs[i],nb)
        dxs[:, i] = dx
    end
    for i in 1:m
        px[i] = Partials{N,V}(Tuple(dxs[i,j] for j in 1:N))
    end
    duals = Dual{T,V,N}.(x, px)
    return (duals, stats)
end

function makematrix(n,X)
    m=zeros(eltype(X),n,n)
    for i=1:n
        m[i,i]=10+X[1]^2
    end

    for i=1:n-2
        m[i,i+2]=-X[2]^2
        m[i+2,i]=-X[2]^2
    end
    m
end

function evalderiv(f,X)
    println("    direct evaluation:  f(X)=$(f(X))")

    dresult=DiffResults.JacobianResult(ones(1),X)
    ForwardDiff.jacobian!(dresult,f,X)
    println("ForwarDiff evaluation:  f(X)=$(DiffResults.value(dresult))")
    println(" ForwardDiff Jacobian: df(X)=$(DiffResults.jacobian(dresult))")
end

function testfwd(X=[1.0,0.1];n=20)

    fdirect(Y)=[sum(makematrix(n,Y)\ones(eltype(Y),n))]
    println("direct solver:")
    evalderiv(fdirect,X)

    println("cg iteration:")
    fcg(Y)=[sum( cg(makematrix(n,Y),ones(eltype(Y),n))[1])]
    evalderiv(fcg,X)
end

testfwd()

[amontoison]

The solution is DiffKrylov.jl.