BC type issue when used with GalacticOptim

[toollu]

As pointed out here, DiffEqOperators is in need of type expansion to work with GalaticOptim/AD.

Test case is a two phase problem along a plug flow reactor with a source term learned by a NN:

using DiffEqFlux
using DiffEqOperators
using GalacticOptim
using LinearAlgebra


## Parameters
#reactor inlet/ outlet concentration
c_in = 0.8
c_out = [2.22045e-16, 2.27763e-08, 8.9472e-08, 2.70662e-07, 7.18169e-07,1.75921e-06,4.07057e-06,9.02313e-06,1.93618e-05, 4.04692e-05,8.2552e-05,0.000165478,0.000326393,0.00063378,0.001214748,0.002295572,0.004281109,0.007839757,0.014035864,0.024269374,0.040172196,0.062886906,0.092739507,0.128716691,0.169202868,0.212285338,0.256330553,0.300039083,0.342489288,0.383087847,0.421482991,0.457497512,0.49109219,0.522325134,0.551298498,0.578142481,0.60300913,0.626074079,0.647493582,0.66739479,0.685911977]

#reactor & discretization parameters
ϵ = 0.36
ϕ =  (1 -ϵ)/ϵ
L = 10.0
nknots = 20

#fiffusion and velocity
D = 1.66e-9
v = 0.157

#ode related
c0 = fill(0.0,2*nknots)
t0 = 0.0
te = 200.0
t = collect(t0:5:te)
tspan = (t0,te)

## System Build Functions

function  discretize_pde(L,nknots,c_in,v,D)
    #discretization distance
    Δx = L/(nknots+1)
    #derivative & approximation orders
    D_derivative_order = 2
    v_derivative_order = 1
    order_approx = 2
    #Diffusion and Advection discretization
    DBand = CenteredDifference(D_derivative_order, order_approx, Δx, nknots,D)
    VBand = UpwindDifference(v_derivative_order, order_approx, Δx, nknots,-v)
    #boundary conditions
    bc  = RobinBC((1.0, -D/v, c_in),(0.0, 1.0, 0.0),Δx,1)
    #build pde system matrix
    ode_matrix = (DBand+VBand)*bc

    return ode_matrix
end


#define NN
ann = FastChain(FastDense(2,7,tanh), FastDense(7,1))
 function rhs(c,q,annpars)
      return  ann([c'; q'],annpars)
 end

#define ode system
 x2 = Float64.(zeros(nknots))
 x1 = Float64.(zeros(nknots))
 lc = Float64.(zeros(nknots))
 sc = Float64.(zeros(nknots))
 M  =  Float64.(zeros(nknots))

function ode_system!(du,u,p,t,ode_matrix)
    lc = @view u[1:nknots]
    sc = @view u[nknots+1:2*nknots]
    mul!(M,ode_matrix,lc)
    x2 = vec(rhs(lc,sc,p))
    x1 = M .- (ϕ.*[0; x2[2:end]])
    du .= vcat(x1,x2)
end


## Prediction and loss function
#(closures for later extension to train on multiple datasets )

function predict(θ,L,nknots,c_in,v,D,tspan,u0)
    ode_matrix = discretize_pde(L,nknots,c_in,v,D)
    ode_problem = ODEProblem((du,u,p,t) -> ode_system!(du, u, p,t,ode_matrix), u0,tspan)
    sol = solve(
            ode_problem,
            TRBDF2(autodiff=false),
            p = θ,
            saveat = t
          )
    return sol
end


function loss(θ)
    sol = predict(θ,L,nknots,c_in,v,D,tspan,c0)
    if any((s.retcode != :Success for s in sol))
        println("System integration failed.")
        return Inf
    else
        return sum(abs2, sol[end,:].-c_out)
    end
end


## Optimization

#Test loss function without optimization
θ = Float64.(initial_params(ann))
loss(θ) #works

#optimize
optprob = OptimizationFunction((x,p) -> loss(x), GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(optprob, θ)
result = GalacticOptim.solve(prob, ADAM(), maxiters = 300) #fails

StackTrace:

LoadError: MethodError: no method matching DiffEqOperators.BoundaryPaddedVector(::ReverseDiff.TrackedReal{Float64,Float64,Nothing}, ::ReverseDiff.TrackedReal{Float64,Float64,Nothing}, ::SubArray{ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}},1,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},Tuple{UnitRange{Int64}},true})

Closest candidates are:

DiffEqOperators.BoundaryPaddedVector(::T, ::T, ::T2) where {T, T2<:AbstractArray{T,1}} at /Users/me/.julia/packages/DiffEqOperators/viN6A/src/boundary_padded_arrays.jl:14

Stacktrace:
[1] *(::RobinBC{Float64,Array{Float64,1}}, ::SubArray{ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}},1,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},Tuple{UnitRange{Int64}},true}) at /Users/me/.julia/packages/DiffEqOperators/viN6A/src/derivative_operators/bc_operators.jl:153

[2] mul!(::Array{Float64,1}, ::GhostDerivativeOperator{Float64,DerivativeOperator{Float64,1,false,Float64,StaticArrays.SArray{Tuple{3},Float64,1,3},StaticArrays.SArray{Tuple{0},StaticArrays.SArray{Tuple{4},Float64,1,4},1,0},Array{Float64,1},Float64},RobinBC{Float64,Array{Float64,1}}}, ::SubArray{ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}},1,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},Tuple{UnitRange{Int64}},true}) at /Users/me/.julia/packages/DiffEqOperators/viN6A/src/derivative_operators/ghost_derivative_operator.jl:15

[3] mul!(::Array{Float64,1}, ::DiffEqOperators.DiffEqOperatorCombination{Float64,Tuple{GhostDerivativeOperator{Float64,DerivativeOperator{Float64,1,false,Float64,StaticArrays.SArray{Tuple{3},Float64,1,3},StaticArrays.SArray{Tuple{0},StaticArrays.SArray{Tuple{4},Float64,1,4},1,0},Array{Float64,1},Float64},RobinBC{Float64,Array{Float64,1}}},GhostDerivativeOperator{Float64,DerivativeOperator{Float64,1,true,Float64,StaticArrays.SArray{Tuple{3},Float64,1,3},StaticArrays.SArray{Tuple{1},StaticArrays.SArray{Tuple{3},Float64,1,3},1,1},Array{Float64,1},Float64},RobinBC{Float64,Array{Float64,1}}}},Array{Float64,1}}, ::SubArray{ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}},1,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},Tuple{UnitRange{Int64}},true}) at /Users/me/.julia/packages/DiffEqOperators/viN6A/src/composite_operators.jl:86

[4] ode_system!(::Array{ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}},1}, ::ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}, ::ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}, ::ReverseDiff.TrackedReal{Float64,Float64,ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}}, ::DiffEqOperators.DiffEqOperatorCombination{Float64,Tuple{GhostDerivativeOperator{Float64,DerivativeOperator{Float64,1,false,Float64,StaticArrays.SArray{Tuple{3},Float64,1,3},StaticArrays.SArray{Tuple{0},StaticArrays.SArray{Tuple{4},Float64,1,4},1,0},Array{Float64,1},Float64},RobinBC{Float64,Array{Float64,1}}},GhostDerivativeOperator{Float64,DerivativeOperator{Float64,1,true,Float64,StaticArrays.SArray{Tuple{3},Float64,1,3},StaticArrays.SArray{Tuple{1},StaticArrays.SArray{Tuple{3},Float64,1,3},1,1},Array{Float64,1},Float64},RobinBC{Float64,Array{Float64,1}}}},Array{Float64,1}}) at /Users/me/projectSRC/JL/testfile.jl:67

Since I'm quite new to Julia I'm not really comfortable to do this (correctly) on my own. I'm however happy to learn/help once I have a clearer understanding of what needs to be done.