So I just had time to go back to a project I was working in the past (related to this). Basically I want to solve a Diffusion-Convection problem tracking two phases along a reactor with a source term that is learned by a neural network.
Since I started writing Julia last summer I read/watched and hopefully learned a lot about Julia and writing faster code. Therefore wanted to rewrite the existing code “properly”, e.g. using GalacticOptim instead of sciml_train , generating the PDE with DiffEqOperators and utilising some performance tips from the DiffEq tutorials. I however do not manage to get my new code running and after two frustrating days I thought maybe someone here might point me towards a solution.
What I have working so far is this:
using DiffEqFlux
using DiffEqOperators
using GalacticOptim
using LinearAlgebra
using OrdinaryDiffEq
## 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
This works well and is already faster what I had before. However when I try to learn the NN parameters with GalacticOptim:
## 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
the AD somewhat seems to trip over the boundary conditions in DiffEqOperators.
Therefore I have two questions:
- What am I doing wrong? How can I get the correct type
::T2inDiffEqOperators.BoundaryPaddedVector(::T, ::T, ::T2)? - Am I on the right track to make the first (working part) of my code faster or am I missing something essential?
The stack trace of the error is:
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