# DiffEqFlux: SOSRA2() algorithm + BacksolveAdjoint() gives an error

Hi everyone. I have a question related to fitting a Neural SDE using DiffEqFlux.jl.

Description of the problem that is solved:
Given a specific SDE system, I want to steer that system (with a particular control) to the all-ones vector. In order to do so, I parametrize the control input as a Neural Network and try to fit the Neural Network. This NN takes as input both the position and the time. In order to train the Neural Network, I decided to use the sciml_train function of DiffEqFlux.jl in combination with the BacksolveAdjoint() method for calculating the gradient of the SDE.

Code:

``````using Flux, DiffEqFlux, DiffEqSensitivity, DifferentialEquations, Statistics

x_size = 6 # Size of the spatial dimensions in the SDE
v_size = 6 # Output size of the control

# Define Neural Network for the control input
input_size = x_size + 1 # size of the spatial dimensions PLUS one time dimensions
nn_initial = Chain(Dense(input_size,v_size)) # The actual neural network
p_nn, model = Flux.destructure(nn_initial)
nn(x,p) = model(p)(x)

# Define the right hand side of the SDE
const_mat = zeros(Float64, (x_size, v_size))
for i = 1:max(x_size,v_size)
const_mat[i,i] = 1
end

function f!(du,u,p,t)
du .= 2.0.*u + const_mat*nn([u;t],p)
end

function g!(du,u,p,t)
du .= 0.0.*u .+ sqrt(2*0.001)
end

# Define SDE problem
u0 = vec(rand(Float64, (x_size,1)))
tspan = (0.0, 1.0)
prob = SDEProblem{true}(f!, g!, u0, tspan, p_nn)

# Defining the loss function
function loss(pars, prob)

function prob_func(prob, i, repeat)
# Prepare new initial state and remake the problem
u0tmp = vec(rand(Float64,(x_size,1)))

remake(prob, p = pars, u0 = u0tmp)
end

ensembleprob = EnsembleProblem(prob, prob_func = prob_func)

adaptive = true, trajectories = 30)

A = convert(Array, _sol)

return sum(abs2, A .- 1)
end

# Actually training/fitting the model
result = DiffEqFlux.sciml_train((p) -> loss(p,prob), p_nn, ADAM(0.1), maxiters=5)
``````

Error:

``````ERROR: LoadError: BoundsError: attempt to access 995-element Array{Array{Float64,1},1} at index [996]
Stacktrace:
[1] getindex at .\array.jl:809 [inlined]
[2] #interpolate!#70 at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_process_interface.jl:524 [inlined]
[3] #interpolate!#72 at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:15 [inlined]
[4] NoiseWrapper at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\types.jl:179 [inlined]
[5] calculate_step!(::NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true}, ::Float64, ::Array{Float64,1}, ::Array{Float32,1}) at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:23
[6] setup_next_step! at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:71 [inlined]
[7] setup_next_step! at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\integrators\integrator_utils.jl:2 [inlined]
[8] __init(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true},SDEFunction{true,DiffEqSensitivity.ODEBacksolveSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,ReverseDiff.GradientTape{DiffEqSensitivity.var"#89#100"{ODEFunction{true,StochasticTransformedFunction{SDEProblem{Array{Float64,1}, ............. verbose::Bool, force_dtmin::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, initialize_integrator::Bool, seed::UInt64, alias_u0::Bool, alias_jumps::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\solve.jl:571
[9] #__solve#100 at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\solve.jl:6 [inlined]
ypeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},SDEFunction{true,typeof(f!),typeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Array{Float64,1},Nothing},UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,typeof(SciMLBase.DEFAULT_OBSERVED),Nothing}},Tuple{ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},ReverseDiff.TrackedArray{Float32,Float64,1,Array{Float32,1},Array{Float64,1}},ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}}, .................................., ::SOSRA2; kwargs::Base.Iterators.Pairs{Symbol,Any,NTuple{7,Symbol},NamedTuple{(:save_everystep, :save_start, :saveat, :tstops, :abstol, :reltol, :adaptive),Tuple{Bool,Bool,Array{Float64,1},Array{Float64,1},Float64,Float64,Bool}}}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:82
[12] #solve#57 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:70 [inlined]
[17] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
[19] #solve#57 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:70 [inlined]
[20] (::typeof(∂(#solve#57)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[21] (::Zygote.var"#180#181"{typeof(∂(#solve#57)),Tuple{NTuple{6,Nothing},Tuple{Nothing}}})(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194
[23] (::typeof(∂(solve##kw)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[24] #batch_func#452 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:143 [inlined]
[25] (::typeof(∂(#batch_func#452)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
[26] #457 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:195 [inlined]
[27] (::typeof(∂(λ)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[28] (::DiffEqBase.var"#219#227")(::typeof(∂(λ)), ::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\init.jl:259
[29] responsible_map(::Function, ::Array{typeof(∂(λ)),1}, ::Vararg{Any,N} where N) at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:188
[30] (::DiffEqBase.var"#∇responsible_map_internal#226"{Array{typeof(∂(λ)),1}})(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\init.jl:259
[32] #solve_batch#456 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:194 [inlined]
[33] (::typeof(∂(#solve_batch#456)))(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
[34] #solve_batch#459 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:203 [inlined]
[35] (::typeof(∂(#solve_batch#459)))(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
[36] macro expansion at .\timing.jl:233 [inlined]
[37] #__solve#451 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:108 [inlined]
[38] (::typeof(∂(#__solve#451)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
[39] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
[41] #solve#59 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:96 [inlined]
[42] (::typeof(∂(#solve#59)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[43] (::Zygote.var"#180#181"{typeof(∂(#solve#59)),Tuple{Tuple{Nothing,Nothing,Nothing},Tuple{Nothing,Nothing}}})(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194
[45] (::typeof(∂(solve##kw)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[46] loss at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:43 [inlined]
[47] (::typeof(∂(loss)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[48] #39 at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:52 [inlined]
[49] (::typeof(∂(#39)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[50] #69 at C:\Users\Sven\.julia\packages\DiffEqFlux\alPQ3\src\train.jl:3 [inlined]
[51] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
[53] OptimizationFunction at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\problems\basic_problems.jl:107 [inlined]
... (the last 3 lines are repeated 1 more time)
[57] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
[59] #8 at C:\Users\Sven\.julia\packages\GalacticOptim\JnLwV\src\solve.jl:94 [inlined]
[60] (::typeof(∂(λ)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
[61] (::Zygote.var"#69#70"{Zygote.Params,Zygote.Context,typeof(∂(λ))})(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface.jl:252
[63] __solve(::OptimizationProblem{false,OptimizationFunction{false,GalacticOptim.AutoZygote,OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Nothing}},GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Nothing}},GalacticOptim.var"#154#164",Nothing,Nothing,Nothing},Array{Float32,1},SciMLBase.NullParameters,Nothing,Nothing,Nothing,Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxiters,),Tuple{Int64}}}}, ::ADAM, ::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; maxiters::Int64, cb::Function, progress::Bool, save_best::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at C:\Users\Sven\.julia\packages\GalacticOptim\JnLwV\src\solve.jl:93
[64] #solve#468 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\solve.jl:3 [inlined]
[65] sciml_train(::var"#39#40", ::Array{Float32,1}, ::ADAM, ::GalacticOptim.AutoZygote; lower_bounds::Nothing, upper_bounds::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxiters,),Tuple{Int64}}}) at C:\Users\Sven\.julia\packages\DiffEqFlux\alPQ3\src\train.jl:6
in expression starting at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:52
``````

Small remark
I removed a part of the error at lines [8], [10] and [11]. The reason is that it was very large and made me go over the character limit for the post. The parts that I removed are replaced by dots (so ‘…’).

Possible location of the error
At the top line of the error, I see `ERROR: LoadError: BoundsError: attempt to access 995-element Array{Array{Float64,1},1} at index [996]`. Here the values 995 and 996 change from run to run. Hence, most likely the error is related to the stochasticity and adaptivity when solving a SDE (either the forward SDE or the Backward SDE of BacksolveAdjoint()). Moreover, when changing the SOSRA2() algorithm to the LambaEulerHeun() algorithm, no errors occur. Hence, I suspect that the error is related to the SOSRA2() algorithm and related algorithms (for instance, with SRA2() the code gives the same error).

Question:
Is there any way I can change the code such that the SOSRA2() algorithm can be used in combination with BacksolveAdjoint()?

Thank you very much for your reply. After updating the package, I used SOSRA2() on some newly made code and it worked.