# Correct sensealg for an ensemble problem of stochastic UODEs with diagonal noise

**URL:** <https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490>\
**Category:** Machine Learning\
**Tags:** diffeq, flux\
**Created:** [November 20, 2020, 11:37am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490 "2020-11-20T11:37:20Z")\
**Posts on this page:** 18\
**Page:** 1

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 20, 2020, 11:37am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/1 "2020-11-20T11:37:20Z")

</div>

Hi all,

I can’t seem to get my stochastic UODE running. Here’s the setup and generation of synthetic data:

```julia
using Plots, DiffEqFlux, Flux, Optim, LinearAlgebra, DifferentialEquations
using Statistics
using DiffEqSensitivity
using DifferentialEquations.EnsembleAnalysis

function sys!(du, u, p, t)
    r, e, μ, h, ph, z, i = p
    du[1] = e * 0.5 * (5μ - u[1]) # nutrient input time series
    du[2] = e * 0.05 * (10μ - u[2]) # grazer density time series
    du[3] = 0.2 * exp(u[1]) - 0.05 * u[3] - r * u[3] / (h + u[3]) * u[4] # nutrient concentration
    du[4] =
        r * u[3] / (h + u[3]) * u[4] - 0.1 * u[4] -
        0.02 * u[4]^z / (ph^z + u[4]^z) * exp(u[2] / 2.0) + i #Algae density
end

function noise!(du, u, p, t)
    du[1] = p[end] # n
    du[2] = p[end] # n
    du[3] = 0.0
    du[4] = 0.0
end

# I've dropped the time frame down quite a bit here just to try to get it to run.
#datasize = 30
#tspan = (0.0f0, 10.0f0)
datasize = 10
tspan = (0.0f0, 3.0f0)
tsteps = range(tspan[1], tspan[2], length = datasize)
u0 = Float32[1.0, 1.0, 1.0, 1.0]

p_ = Float32[1.1, 1.0, 0.0, 2.0, 1.0, 1.0, 1e-6, 1.0]

prob = SDEProblem(sys!, noise!, u0, tspan, p_)
ensembleprob = EnsembleProblem(prob)

solution = solve(
    ensembleprob,
    EnsembleThreads();
    trajectories = 1000,
    abstol = 1e-5,
    reltol = 1e-5,
    maxiters = 1e8,
    saveat = tsteps,
)

(truemean, truevar) = Array.(timeseries_steps_meanvar(solution))

```

I’ve got 14 hardware threads, so this doesn’t take too long to run (`@btime` on the solve step is 23 seconds).

The plan is to solve for the Michaelis-Menton kinetics portions of eqs 3 and 4, so I set up my UODE following a combination of documentation pages (I’m not 100% sure there’s an example that matches my case): [neural\_sde](https://diffeqflux.sciml.ai/dev/examples/neural_sde/), [PINNs](https://diffeqflux.sciml.ai/dev/examples/tensor_layer/), [zygote example](https://github.com/MikeInnes/zygote-paper/blob/master/neural_sde/neural_sde.jl), [uode examples](https://github.com/ChrisRackauckas/universal_differential_equations).

```julia
# Generate a NN with four inputs [L, G, N, A] and two outputs [NN1, NN2]
ann = FastChain(FastDense(4, 32, tanh), FastDense(32, 32, tanh), FastDense(32, 2))
α = initial_params(ann)

function dudt_(du, u, p, t)
    r, e, μ, h, ph, z, i = p_

    # we replace each of our MM values with the result of the NN.
    MM = ann(u, p)

    du[1] = e * 0.5 * (5μ - u[1]) # nutrient input time series
    du[2] = e * 0.05 * (10μ - u[2]) # grazer density time series
    du[3] = 0.2 * exp(u[1]) - 0.05 * u[3] - MM[1] # nutrient concentration
    du[4] = MM[2] - 0.1 * u[4] - 0.02 * u[4]^z / (ph^z + u[4]^z) * exp(u[2] / 2.0) + i #Algae density
end
function noise_(du, u, p, t)
    du[1] = p_[end]
    du[2] = p_[end]
    du[3] = 0.0
    du[4] = 0.0
end

prob_nn = SDEProblem(dudt_, noise_, u0, tspan, p = nothing)

function loss(θ)
    tmp_prob = remake(prob_nn, p = θ)
    ensembleprob = EnsembleProblem(tmp_prob)
    tmp_sol = solve(
        ensembleprob,
        EnsembleThreads();
        saveat = tsteps,
        trajectories = 100, # Drop down a little for troubleshooting
       # sensealg = ReverseDiffAdjoint(),
    # sensealg = ForwardDiffSensitivity(),
    )
    (tmp_mean, tmp_var) = Array.(timeseries_steps_meanvar(tmp_sol))
    sum(abs2, truemean - tmp_mean) + 0.1 * sum(abs2, truevar - tmp_var), tmp_mean
end

# Test
loss(α) # @btime 360 ms

const losses = []

callback(θ, l, pred) = begin
    push!(losses, l)
    println("Current loss after $(length(losses)) iterations: $(losses[end])")
    false
end

res1 = DiffEqFlux.sciml_train(
    loss,
    α,
    ADAM(0.1),
    cb = callback,
    maxiters = 200,
   # allow_f_increases = true,
)

```

If I’ve got this set up correctly, I cannot figure out what `sensealg` works here. If I choose the default (i.e. omit a `sensealg` keyword which is happening in the `loss` function above), I get an unhelpful error stating:

```julia
Training 100%|█████████████████████████████████████████████████████████████| Time: 0:00:54
ERROR: Mutating arrays is not supported

```

> **Stacktrace**
>
> ```julia
> Stacktrace:
> [1] error(::String) at ./error.jl:33
> [2] (::Zygote.var"#372#373")(::Nothing) at /home/tim/.julia/packages/Zygote/c0awc/src/lib/array.jl:65
> [3] (::Zygote.var"#2265#back#374"{Zygote.var"#372#373"})(::Nothing) at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
> [4] timeseries_steps_meanvar at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/ensemble_analysis.jl:71 [inlined]
> [5] (::typeof(∂(timeseries_steps_meanvar)))(::Tuple{Array{Float64,2},Array{Float64,2}}) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [6] loss at ./REPL[22]:12 [inlined]
> [7] (::typeof(∂(loss)))(::Tuple{Float64,Nothing}) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [8] #150 at /home/tim/.julia/packages/Zygote/c0awc/src/lib/lib.jl:191 [inlined]
> [9] (::Zygote.var"#1693#back#152"{Zygote.var"#150#151"{typeof(∂(loss)),Tuple{Tuple{Nothing},Tuple{}}}})(::Tuple{Float64,Nothing}) at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
> [10] #74 at /home/tim/.julia/packages/DiffEqFlux/8UHw5/src/train.jl:120 [inlined]
> [11] (::typeof(∂(λ)))(::Float64) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [12] (::Zygote.var"#54#55"{Zygote.Params,Zygote.Context,typeof(∂(λ))})(::Float64) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface.jl:177
> [13] gradient(::Function, ::Zygote.Params) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface.jl:54
> [14] macro expansion at /home/tim/.julia/packages/DiffEqFlux/8UHw5/src/train.jl:119 [inlined]
> [15] macro expansion at /home/tim/.julia/packages/ProgressLogging/BBN0b/src/ProgressLogging.jl:328 [inlined]
> [16] (::DiffEqFlux.var"#73#78"{typeof(callback),Int64,Bool,Bool,typeof(loss),Array{Float32,1},Zygote.Params})() at /home/tim/.julia/packages/DiffEqFlux/8UHw5/src/train.jl:64
> [17] with_logstate(::Function, ::Any) at ./logging.jl:408
> [18] with_logger at ./logging.jl:514 [inlined]
> [19] maybe_with_logger(::DiffEqFlux.var"#73#78"{typeof(callback),Int64,Bool,Bool,typeof(loss),Array{Float32,1},Zygote.Params}, ::LoggingExtras.TeeLogger{Tuple{LoggingExtras.EarlyFilteredLogger{TerminalLoggers.TerminalLogger,DiffEqFlux.var"#68#70"},LoggingExtras.EarlyFilteredLogger{Logging.ConsoleLogger,DiffEqFlux.var"#69#71"}}}) at /home/tim/.julia/packages/DiffEqFlux/8UHw5/src/train.jl:39
> [20] sciml_train(::Function, ::Array{Float32,1}, ::ADAM, ::Base.Iterators.Cycle{Tuple{DiffEqFlux.NullData}}; cb::Function, maxiters::Int64, progress::Bool, save_best::Bool) at /home/tim/.julia/packages/DiffEqFlux/8UHw5/src/train.jl:63
> [21] top-level scope at REPL[26]:2
> 
> ```

According to the [DiffEqSensitivity docs](https://diffeq.sciml.ai/latest/analysis/sensitivity/#Choosing-a-Sensitivity-Algorithm), I should be using `QuadratureAdjoint` for my system size, which if I’m not mistaken is what the defaults actually picked, since I get the `Mutating arrays not supported error` using it. At what point am I mutating arrays?

`ReverseDiffAdjoint` gives me an equally opaque error somewhere deep:

```julia
Training 100%|█████████████████████████████████████████████████████████████| Time: 0:01:12
ERROR: TaskFailedException:
MethodError: no method matching randn(::RandomNumbers.Xorshifts.Xoroshiro128Plus, ::Type{ReverseDiff.TrackedReal{Float32,Float32,Nothing}})                                                           
Closest candidates are:
  randn(::Random.AbstractRNG, ::Type{T}, ::Tuple{Vararg{Int64,N}} where N) where T at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:201
  randn(::Random.AbstractRNG, ::Type{T}, ::Integer, ::Integer...) where T at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:204
  randn(::Random.AbstractRNG) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:38
  ...

```

> **Stacktrace**
>
> ```julia
> Stacktrace:
> [1] randn!(::RandomNumbers.Xorshifts.Xoroshiro128Plus, ::Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:178
> [2] wiener_randn! at /home/tim/.julia/packages/DiffEqNoiseProcess/KaHvM/src/wiener.jl:9 [inlined]
> [3] INPLACE_WHITE_NOISE_DIST(::Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1}, ::NoiseProcess{ReverseDiff.TrackedReal{Float32,Float32,Nothing},2,Float32,Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1},Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1},Array{Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float32,Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1},Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1}},true},ResettableStacks.ResettableStack{Tuple{Float32,Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1},Array{ReverseDiff.TrackedReal{Float32,Float32,Nothing},1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus}, ::Float32, ::Array{ReverseDiff.TrackedReal{Float32,Float32,ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}},1}, ::ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}, ::Float32, ::RandomNumbers.Xorshifts.Xoroshiro128Plus) at /home/tim/.julia/packages/DiffEqNoiseProcess/KaHvM/src/wiener.jl:42
> [4] setup_next_step! at /home/tim/.julia/packages/DiffEqNoiseProcess/KaHvM/src/noise_interfaces/noise_process_interface.jl:145 [inlined]                                                             
> [5] setup_next_step! at /home/tim/.julia/packages/StochasticDiffEq/Abmgl/src/integrators/integrator_utils.jl:2 [inlined]                                                                             
> [6] __init(::SDEProblem{Array{ReverseDiff.TrackedReal{Float32,Float32,ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}},1},Tuple{Float32,Float32},true,ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}, ::SOSRI, ::Array{Any,1}, ::Array{Any,1}, ::Type{T} where T, ::Type{Val{true}}; saveat::StepRangeLen{Float32,Float64,Float64}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_noise::Bool, save_on::Bool, save_start::Bool, save_end::Bool, callback::Nothing, dense::Bool, calck::Bool, dt::Float32, adaptive::Bool, gamma::Rational{Int64}, abstol::Nothing, reltol::Nothing, qmax::Rational{Int64}, qmin::Rational{Int64}, qoldinit::Rational{Int64}, fullnormalize::Bool, failfactor::Int64, beta2::Rational{Int64}, beta1::Rational{Int64}, delta::Rational{Int64}, maxiters::Int64, dtmax::Float32, dtmin::Float32, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), 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{Symbol,Bool,Tuple{Symbol},NamedTuple{(:default_set,),Tuple{Bool}}}) at /home/tim/.julia/packages/StochasticDiffEq/Abmgl/src/solve.jl:562
> [7] #__solve#97 at /home/tim/.julia/packages/StochasticDiffEq/Abmgl/src/solve.jl:6 [inlined]
> [8] __solve(::SDEProblem{Array{ReverseDiff.TrackedReal{Float32,Float32,ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}},1},Tuple{Float32,Float32},true,ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}, ::Nothing; default_set::Bool, kwargs::Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}) at /home/tim/.julia/packages/DifferentialEquations/fpohE/src/default_solve.jl:7
> [9] #solve_call#456 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:65 [inlined]
> [10] #solve_up#458 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:92 [inlined]
> [11] #solve#457 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:74 [inlined]
> [12] (::DiffEqSensitivity.var"#reversediff_adjoint_forwardpass#168"{Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}},SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},Nothing,Tuple{}})(::ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}, ::ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}) at /home/tim/.julia/packages/DiffEqSensitivity/WiCRA/src/local_sensitivity/concrete_solve.jl:315
> [13] ReverseDiff.GradientTape(::Function, ::Tuple{Array{Float32,1},Array{Float32,1}}, ::ReverseDiff.GradientConfig{Tuple{ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}},ReverseDiff.TrackedArray{Float32,Float32,1,Array{Float32,1},Array{Float32,1}}}}) at /home/tim/.julia/packages/ReverseDiff/jFRo1/src/api/tape.jl:207
> [14] GradientTape at /home/tim/.julia/packages/ReverseDiff/jFRo1/src/api/tape.jl:204 [inlined]
> [15] _concrete_solve_adjoint(::SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}, ::Nothing, ::ReverseDiffAdjoint, ::Array{Float32,1}, ::Array{Float32,1}; kwargs::Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}) at /home/tim/.julia/packages/DiffEqSensitivity/WiCRA/src/local_sensitivity/concrete_solve.jl:326
> [16] #_solve_adjoint#478 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:277 [inlined]
> [17] #adjoint#469 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:241 [inlined]
> [18] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:53 [inlined]
> [19] adjoint at /home/tim/.julia/packages/Zygote/c0awc/src/lib/lib.jl:188 [inlined]
> [20] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
> [21] #solve#457 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:74 [inlined]
> [22] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve#457", ::ReverseDiffAdjoint, ::Nothing, ::Nothing, ::Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}, ::typeof(solve), ::SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}, ::Nothing) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [23] adjoint at /home/tim/.julia/packages/Zygote/c0awc/src/lib/lib.jl:188 [inlined]
> [24] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
> [25] _pullback(::Zygote.Context, ::DiffEqBase.var"#solve##kw", ::NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}, ::typeof(solve), ::SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing}, ::Nothing) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [26] #batch_func#360 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:143 [inlined]                                                                                
> [27] _pullback(::Zygote.Context, ::DiffEqBase.var"##batch_func#360", ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}}, ::typeof(DiffEqBase.batch_func), ::Int64, ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::Nothing) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [28] #365 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:187 [inlined]                                                                                           
> [29] #499 at /home/tim/.julia/packages/Zygote/c0awc/src/lib/array.jl:180 [inlined]
> [30] iterate at ./generator.jl:47 [inlined]
> [31] _collect(::UnitRange{Int64}, ::Base.Generator{UnitRange{Int64},Zygote.var"#499#503"{Zygote.Context,DiffEqBase.var"#365#366"{Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}},EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing},Nothing}}}, ::Base.EltypeUnknown, ::Base.HasShape{1}) at ./array.jl:699
> [32] collect_similar(::UnitRange{Int64}, ::Base.Generator{UnitRange{Int64},Zygote.var"#499#503"{Zygote.Context,DiffEqBase.var"#365#366"{Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}},EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing},Nothing}}}) at ./array.jl:628
> [33] map(::Function, ::UnitRange{Int64}) at ./abstractarray.jl:2162
> [34] ∇map(::Zygote.Context, ::Function, ::UnitRange{Int64}) at /home/tim/.julia/packages/Zygote/c0awc/src/lib/array.jl:180
> [35] adjoint at /home/tim/.julia/packages/Zygote/c0awc/src/lib/array.jl:196 [inlined]
> [36] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
> [37] #solve_batch#364 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:186 [inlined]                                                                               
> [38] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve_batch#364", ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}}, ::typeof(DiffEqBase.solve_batch), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::Nothing, ::EnsembleSerial, ::UnitRange{Int64}, ::Int64) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0 (repeats 2 times)                                                                                                  
> [39] #369 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:213 [inlined]                                                                                           
> [40] _pullback(::Zygote.Context, ::DiffEqBase.var"#369#371"{Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}},EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing},Nothing,UnitRange{Int64},Int64,Int64}, ::Int64) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [41] #639 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/init.jl:233 [inlined]
> [42] macro expansion at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:221 [inlined]                                                                                
> [43] (::DiffEqBase.var"#510#threadsfor_fun#372"{DiffEqBase.var"#639#644"{Zygote.Context,DiffEqBase.var"#369#371"{Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}},EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing},Nothing,UnitRange{Int64},Int64,Int64}},Tuple{UnitRange{Int64}},Array{Tuple{Any,typeof(∂(λ))},1},UnitRange{Int64}})(::Bool) at ./threadingconstructs.jl:81
> [44] (::DiffEqBase.var"#510#threadsfor_fun#372"{DiffEqBase.var"#639#644"{Zygote.Context,DiffEqBase.var"#369#371"{Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}},EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing},Nothing,UnitRange{Int64},Int64,Int64}},Tuple{UnitRange{Int64}},Array{Tuple{Any,typeof(∂(λ))},1},UnitRange{Int64}})() at ./threadingconstructs.jl:48
> Stacktrace:
> [1] wait at ./task.jl:267 [inlined]
> [2] threading_run(::Function) at ./threadingconstructs.jl:34
> [3] macro expansion at ./threadingconstructs.jl:93 [inlined]
> [4] tmap(::Function, ::UnitRange{Int64}) at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:220
> [5] ∇tmap(::Zygote.Context, ::Function, ::UnitRange{Int64}) at /home/tim/.julia/packages/DiffEqBase/3iigH/src/init.jl:233
> [6] adjoint at /home/tim/.julia/packages/DiffEqBase/3iigH/src/zygote.jl:55 [inlined]
> [7] _pullback(::Zygote.Context, ::typeof(DiffEqBase.tmap), ::Function, ::UnitRange{Int64}) at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47
> [8] #solve_batch#367 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:207 [inlined]
> [9] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve_batch#367", ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}}, ::typeof(DiffEqBase.solve_batch), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::Nothing, ::EnsembleThreads, ::UnitRange{Int64}, ::Int64) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [10] macro expansion at ./timing.jl:233 [inlined]
> [11] #__solve#359 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:108 [inlined]
> [12] _pullback(::Zygote.Context, ::DiffEqBase.var"##__solve#359", ::Int64, ::Int64, ::Int64, ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:saveat, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},ReverseDiffAdjoint}}}, ::typeof(DiffEqBase.__solve), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::Nothing, ::EnsembleThreads) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0 (repeats 2 times)
> [13] #__solve#358 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/ensemble/basic_ensemble_solve.jl:87 [inlined]
> [14] _pullback(::Zygote.Context, ::DiffEqBase.var"##__solve#358", ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol,Symbol},NamedTuple{(:saveat, :trajectories, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},Int64,ReverseDiffAdjoint}}}, ::typeof(DiffEqBase.__solve), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::EnsembleThreads) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [15] adjoint at /home/tim/.julia/packages/Zygote/c0awc/src/lib/lib.jl:188 [inlined]
> [16] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
> [17] #solve#459 at /home/tim/.julia/packages/DiffEqBase/3iigH/src/solve.jl:100 [inlined]
> [18] _pullback(::Zygote.Context, ::DiffEqBase.var"##solve#459", ::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol,Symbol},NamedTuple{(:saveat, :trajectories, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},Int64,ReverseDiffAdjoint}}}, ::typeof(solve), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::EnsembleThreads) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> [19] adjoint at /home/tim/.julia/packages/Zygote/c0awc/src/lib/lib.jl:188 [inlined]
> [20] _pullback at /home/tim/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
> [21] _pullback(::Zygote.Context, ::DiffEqBase.var"#solve##kw", ::NamedTuple{(:saveat, :trajectories, :sensealg),Tuple{StepRangeLen{Float32,Float64,Float64},Int64,ReverseDiffAdjoint}}, ::typeof(solve), ::EnsembleProblem{SDEProblem{Array{Float32,1},Tuple{Float32,Float32},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(dudt_),typeof(noise_),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(noise_),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},typeof(DiffEqBase.DEFAULT_PROB_FUNC),typeof(DiffEqBase.DEFAULT_OUTPUT_FUNC),typeof(DiffEqBase.DEFAULT_REDUCTION),Nothing}, ::EnsembleThreads) at /home/tim/.julia/packages/Zygote/c0awc/src/compiler/interface2.jl:0
> 
> ```
> 
> Had to cut this due to the text limit, I can paste the rest in a subsequent post if needed.

`ForwardDiffSensitivity` does not fail, but also does not solve anything. I’ve left the run going overnight and don’t even manage to get to the first callback. It just sits and spools forever.

So if all of these are failing, have I set up this problem incorrectly? Or is there a better strategy to solve this particular system, or a better `sensealg` setup here?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 20, 2020, 1:26pm UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/2 "2020-11-20T13:26:07Z")

</div>

> [@Libbum](#):
>
> According to the [DiffEqSensitivity docs](https://diffeq.sciml.ai/latest/analysis/sensitivity/#Choosing-a-Sensitivity-Algorithm), I should be using `QuadratureAdjoint` for my system size, which if I’m not mistaken is what the defaults actually picked, since I get the `Mutating arrays not supported error` using it. At what point am I mutating arrays?

Not for SDEs. For diagonal noise SDEs likely the right thing is `TrackerAdjoint`. `QuadratureAdjoint` doesn’t exist for SDEs given the weird rules of SDE integration, so we should throw a better error for that case. You can try `BacksolveAdjoint` or `InterpolatingAdjoint`, but you’ll need to switch to a commutative noise SDE solver because the reverse pass is no longer diagonal. We haven’t done all of the timings yet but @frankschae might have more comments.

> [@Libbum](#):
>
> `ReverseDiffAdjoint` gives me an equally opaque error somewhere deep

`ReverseDiffAdjoint` is really what we should be using, so @mohamed82008 could we get help fixing that?

> [@Libbum](#):
>
> I get the `Mutating arrays not supported error` using it. At what point am I mutating arrays?

`timeseries_steps_meanvar` mutates as I found out yesterday: [EnsembleProblem with timepoint\_meanvar errors due to internal mutation · Issue #446 · SciML/DiffEqFlux.jl · GitHub](https://github.com/SciML/DiffEqFlux.jl/issues/446) . I’ll probably patch this over the weekend.

---

<div class="post-metadata">

**Author:** ![frankschae](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/frankschae/32/8382_2.png) [@frankschae](https://discourse.julialang.org/u/frankschae)\
**Post date:** [November 20, 2020, 4:21pm UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/3 "2020-11-20T16:21:45Z")

</div>

Hi! You could try replacing the `timeseries_steps_meanvar` function by

```julia
tmp_sol = Array(solve(
        ensembleprob,
        EM(),
        dt = 0.001f0,
        EnsembleThreads();
      ...)
(tmp_mean, tmp_var) = mean(tmp_sol,dims=3)[:,:], var(tmp_sol,dims=3)[:,:]

```

With that change, Zygote can compute the gradients:

```julia
Zygote.gradient(p->loss(p)[1], α)

```

So this should solve the mutation issues until `timeseries_steps_meanvar` is patched.

Please note that I have chosen the 0.5 order `EM()` solver for Ito SDEs. This has essentially two reasons:  
(1) As Chris said, the reverse process will have non-diagonal noise, therefore not all solvers are applicable.  
(2) We have some problems with the reconstruction of the noise process with higher order solvers for Ito SDEs. (see [https://github.com/SciML/DiffEqNoiseProcess.jl/pull/62](https://github.com/SciML/DiffEqNoiseProcess.jl/pull/62)). So, if you want to use  
`BacksolveAdjoint` or `InterpolatingAdjoint` for SDEs in the Ito sense, one currently should choose `EM()` to be on the safe side…

Note also that the default sensitivity algorithm will be `InterpolatingAdjoint`. For SDEs written in the Ito sense, we have just implemented the tests for `BacksolveAdjoint` so far ([https://github.com/SciML/DiffEqSensitivity.jl/pull/317](https://github.com/SciML/DiffEqSensitivity.jl/pull/317))… More docs and tests should come soon 🙂

I am not sure what’s going wrong with `ReverseDiffAdjoint` and `ForwardDiffSensitivity` – Except that `ForwardDiffSensitivity` scales badly with the numbers of parameters… Does it work for fewer network parameters?

Regarding performance, it’s probably best if you do your own benchmarks for the different adjoint choices (I’d be interested to see them as well). Some small suggestions:

– setting return values to `nothing`, i.e.,

```julia
function noise_(du, u, p, t)
    du[1] = p_[end]
    du[2] = p_[end]
    du[3] = 0.0
    du[4] = 0.0
    return nothing
end

```

– make sure not to mix up different float types like in

```julia
du[2] = e * 0.05 * (10μ - u[2]) -> 0.05 is Float64

```

– localize all parameters in the drift and diffusion function

```julia
r, e, μ, h, ph, z, i = p[1:7]
p_nn = @view p[8:end]

```

(while of course just optimizing the network parameters.)

---

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 20, 2020, 4:24pm UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/4 "2020-11-20T16:24:24Z")

</div>

Thanks both! I’ll take a look at these suggestions and report back.

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [November 21, 2020, 2:51am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/5 "2020-11-21T02:51:11Z")

</div>

> [@ChrisRackauckas](#):
>
> `ReverseDiffAdjoint` is really what we should be using, so @mohamed82008 could we get help fixing that?

It’s not immediately obvious to me what’s wrong. @Libbum could you please open an issue to track this?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 21, 2020, 3:01am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/6 "2020-11-21T03:01:24Z")

</div>

do we have a good way to make it skip tracking of randoms?

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [November 21, 2020, 3:03am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/7 "2020-11-21T03:03:10Z")

</div>

Yes. We can define a custom gradient of `rand` or its equivalent to be nothing.

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [November 21, 2020, 3:07am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/8 "2020-11-21T03:07:21Z")

</div>

However, there has been a discussion before about whether this makes sense in general or not. [https://github.com/JuliaDiff/ChainRules.jl/issues/262](https://github.com/JuliaDiff/ChainRules.jl/issues/262) and [https://github.com/TuringLang/DistributionsAD.jl/issues/123](https://github.com/TuringLang/DistributionsAD.jl/issues/123)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 21, 2020, 3:19am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/9 "2020-11-21T03:19:16Z")

</div>

I would’ve thought I handled it via [https://github.com/SciML/DiffEqNoiseProcess.jl/blob/master/src/init.jl#L12-L16](https://github.com/SciML/DiffEqNoiseProcess.jl/blob/master/src/init.jl#L12-L16) , but it seems that’s not being picked up?

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [November 21, 2020, 3:41am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/10 "2020-11-21T03:41:29Z")

</div>

You would need one for `TrackedReal` as well. The user is using `inplace = true` so it’s using `Array{<:TrackedReal}` which is probably a bigger problem performance-wise but it should work.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 21, 2020, 7:42am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/11 "2020-11-21T07:42:16Z")

</div>

@Libbum try this PR branch:

[https://github.com/SciML/DiffEqNoiseProcess.jl/pull/73](https://github.com/SciML/DiffEqNoiseProcess.jl/pull/73)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 21, 2020, 6:04pm UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/12 "2020-11-21T18:04:05Z")

</div>

I released a patch with this.

---

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 23, 2020, 9:06am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/13 "2020-11-23T09:06:27Z")

</div>

Thanks for the assistance everyone!

I’m using the patched version of `DiffEqNoiseProcess.jl` @ChrisRackauckas, although since we seem to be tracking more than one problem in this thread I can’t say for sure if that’s assisted or not.

Here’s what’s working:

```julia
 function loss(θ)
    tmp_prob = remake(prob_nn, p = θ)
    ensembleprob = EnsembleProblem(tmp_prob)
    tmp_sol = Array(solve(
        ensembleprob,
        EM(); 
        dt = tsteps.step,
        trajectories = 100,
        sensealg = InterpolatingAdjoint(),
       ))                                  
    tmp_mean = mean(tmp_sol,dims=3)[:,:]   
    tmp_var = var(tmp_sol,dims=3)[:,:]     
    sum(abs2, truemean - tmp_mean) + 0.1 * sum(abs2, truevar - tmp_var), tmp_mean
end

```

Which is essentially what @frankschae has suggested above (remove the mutating functions, use `EM()` and `InterpolatingAdjoint()`).

If I follow correctly, the `TrackedReals` addition should assist when using `ReverseDiffAdjoint` correct?  
The error I’m getting there is now:

```julia
 ERROR: TaskFailedException:
MethodError: no method matching randn(::RandomNumbers.Xorshifts.Xoroshiro128Plus, ::Type{ReverseDiff.TrackedReal{Float32,Float32,Nothing}})                                                           
Closest candidates are:
  randn(::Random.AbstractRNG, ::Type{T}, ::Tuple{Vararg{Int64,N}} where N) where T at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:201
  randn(::Random.AbstractRNG, ::Type{T}, ::Integer, ::Integer...) where T at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:204
  randn(::Random.AbstractRNG) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/Random/src/normal.jl:38
  ...

```

@mohamed82008, concerning your latest comment: would it be better to _not_ use inplace in this instance? I don’t quite follow if Chris’ changes has resolved that worry or not. If you think this problem requires an issue to track it - which package are we looking at here? DiffEqSensitivity? Or should I just put it on DifferentialEquations?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 23, 2020, 9:11am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/14 "2020-11-23T09:11:14Z")

</div>

This looks like something that wouldn’t happen in my fix? Are you on DiffEqNoiseProcess v5.4.2? If so, please give a code that I can just copy paste to see the error, otherwise I’m not going to go through the effort to find a reproducer.

---

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 23, 2020, 9:13am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/15 "2020-11-23T09:13:34Z")

</div>

Yeah, I believe so. But I’ll add it explicitly and try again. If not I’ll put up a full MWE.

---

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 23, 2020, 9:26am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/16 "2020-11-23T09:26:00Z")

</div>

Fails with 5.4.2—yes.

```julia
(uode) pkg> st
Status `/mnt/turtle/scratch/uode/Project.toml`
  [ff3c4d4f] AutoOptimize v0.1.0 `https://github.com/SciML/AutoOptimize.jl.git#master`
  [052768ef] CUDA v1.3.3
  [2445eb08] DataDrivenDiffEq v0.4.1
  [a93c6f00] DataFrames v0.22.1
  [7806a523] DecisionTree v0.10.10
  [2b5f629d] DiffEqBase v6.48.2
  [aae7a2af] DiffEqFlux v1.24.0
  [071ae1c0] DiffEqGPU v1.8.0
  [77a26b50] DiffEqNoiseProcess v5.4.2
  [41bf760c] DiffEqSensitivity v6.34.0
  [587475ba] Flux v0.11.1
  [033835bb] JLD2 v0.3.0
  [2fda8390] LsqFit v0.11.0
  [961ee093] ModelingToolkit v3.21.0
  [429524aa] Optim v1.2.0
  [1dea7af3] OrdinaryDiffEq v5.45.1
  [789caeaf] StochasticDiffEq v6.26.0
  [5d786b92] TerminalLoggers v0.1.2
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [2f01184e] SparseArrays
  [10745b16] Statistics

```

Current script:

```julia
using DiffEqFlux, Flux, Optim, LinearAlgebra
using DiffEqNoiseProcess
using StochasticDiffEq
using Statistics
using DiffEqSensitivity
using DiffEqBase.EnsembleAnalysis

function sys!(du, u, p, t)
    r, e, μ, h, ph, z, i = p
    du[1] = e * 0.5 * (5μ - u[1]) # nutrient input time series
    du[2] = e * 0.05 * (10μ - u[2]) # grazer density time series
    du[3] = 0.2 * exp(u[1]) - 0.05 * u[3] - r * u[3] / (h + u[3]) * u[4] # nutrient concentration
    du[4] =
        r * u[3] / (h + u[3]) * u[4] - 0.1 * u[4] -
        0.02 * u[4]^z / (ph^z + u[4]^z) * exp(u[2] / 2.0) + i #Algae density
end

function noise!(du, u, p, t)
    du[1] = p[end] # n
    du[2] = p[end] # n
    du[3] = 0.0
    du[4] = 0.0
end

datasize = 10
tspan = (0.0f0, 3.0f0)
tsteps = range(tspan[1], tspan[2], length = datasize)
u0 = Float32[1.0, 1.0, 1.0, 1.0]

p_ = Float32[1.1, 1.0, 0.0, 2.0, 1.0, 1.0, 1e-6, 1.0]

prob = SDEProblem(sys!, noise!, u0, tspan, p_)
ensembleprob = EnsembleProblem(prob)

solution = solve(
    ensembleprob,
    SOSRI(),
    EnsembleThreads();  
    trajectories = 1000,
    abstol = 1e-5,
    reltol = 1e-5, 
    maxiters = 1e8, 
    saveat = tsteps,
)

(truemean, truevar) = Array.(timeseries_steps_meanvar(solution))

ann = FastChain(FastDense(4, 32, tanh), FastDense(32, 32, tanh), FastDense(32, 2))
α = initial_params(ann)

function dudt_(du, u, p, t)
    r, e, μ, h, ph, z, i = p_

    MM = ann(u, p)

    du[1] = e * 0.5 * (5μ - u[1]) # nutrient input time series
    du[2] = e * 0.05 * (10μ - u[2]) # grazer density time series
    du[3] = 0.2 * exp(u[1]) - 0.05 * u[3] - MM[1] # nutrient concentration
    du[4] = MM[2] - 0.1 * u[4] - 0.02 * u[4]^z / (ph^z + u[4]^z) * exp(u[2] / 2.0) + i #Algae density
    return nothing
end
function noise_(du, u, p, t)
    du[1] = p_[end]
    du[2] = p_[end]
    du[3] = 0.0
    du[4] = 0.0
    return nothing
end

prob_nn = SDEProblem(dudt_, noise_, u0, tspan, p = nothing)

function loss(θ)
    tmp_prob = remake(prob_nn, p = θ)
    ensembleprob = EnsembleProblem(tmp_prob)
    tmp_sol = Array(solve(
        ensembleprob,
        EM();
        dt = tsteps.step,
        trajectories = 100,
        sensealg = ReverseDiffAdjoint(),
       ))
    tmp_mean = mean(tmp_sol,dims=3)[:,:]
    tmp_var = var(tmp_sol,dims=3)[:,:]
    sum(abs2, truemean - tmp_mean) + 0.1 * sum(abs2, truevar - tmp_var), tmp_mean
end

const losses = []
callback(θ, l, pred) = begin
    push!(losses, l)
    if length(losses)%50 == 0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

res1 = DiffEqFlux.sciml_train(
    loss,
    α,
    ADAM(0.1),
    cb = callback,
    maxiters = 200,
)

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [November 23, 2020, 9:56am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/17 "2020-11-23T09:56:11Z")

</div>

This is the overload you need: [https://github.com/SciML/DiffEqNoiseProcess.jl/pull/75](https://github.com/SciML/DiffEqNoiseProcess.jl/pull/75) . But now I’m going to punt it over to @mohamed82008 since it’s erroring because the output of the function is an Array{TrackedReal} so it can’t `increment_deriv!`, and `reduce(vcat,output)` isn’t giving a TrackedArray.

---

<div class="post-metadata">

**Author:** ![Libbum](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/libbum/32/3935_2.png) [@Libbum](https://discourse.julialang.org/u/Libbum)\
**Post date:** [November 23, 2020, 10:03am UTC](https://discourse.julialang.org/t/correct-sensealg-for-an-ensemble-problem-of-stochastic-uodes-with-diagonal-noise/50490/18 "2020-11-23T10:03:04Z")

</div>

Thanks Chris!. @mohamed82008, I’m tracking the issue Chris just mentioned, so happy to provide details there if needed.