deepcopy error on sparse analytic Jacobian EnsembleProblems

[joegilkes]

Sorry if this is in the wrong place, it covers a fairly broad set of packages so let me know if there's a better place to move it to.

Whenever I try to run an EnsembleProblem on an ODEProblem that has been created from a ModelingToolkit ODESystem with an analytic Jacobian jac=true and sparsity sparse=true (a bit of a mouthful, I know), I get the following error:

ERROR: deepcopy of Modules not supported
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] deepcopy_internal(x::Module, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:33
  [3] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
  [4] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
  [5] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
  [6] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
  [7] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
  [8] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
  [9] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [10] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [11] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [12] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [13] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [14] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [15] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [16] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [17] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [18] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [19] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [20] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [21] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [22] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [23] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [24] _deepcopy_array_t(x::Array, T::Type, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:105
 [25] deepcopy_internal(x::Vector{Any}, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:92
 [26] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any})
    @ Base ./deepcopy.jl:65
 [27] deepcopy_internal(x::Any, stackdict::IdDict{Any, Any}) (repeats 4 times)
    @ Base ./deepcopy.jl:76
 [28] deepcopy
    @ ./deepcopy.jl:26 [inlined]
 [29] batch_func(i::Int64, prob::EnsembleProblem{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, SciMLBase.NullParameters, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xcb760beb, 0xef6a2ddc, 0xf02b42e4, 0x861e2e82, 0x7639da69)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x5024cc6a, 0x6af89867, 0x0ad7dad3, 0xcc69ec74, 0x21a02068)}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x477f85ee, 0xd0e8c452, 0xbd3625c2, 0xd167f375, 0xbb4b8d4e)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xca37cb04, 0x0af10484, 0x3c689b9e, 0xbcb1957d, 0x6e5acb8b)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Any}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ODESystem, Dict{Any, Any}}, Nothing, ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, typeof(prob_func), typeof(SciMLBase.DEFAULT_OUTPUT_FUNC), typeof(SciMLBase.DEFAULT_REDUCTION), Nothing}, alg::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLBase ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:89
 [30] batch_func
    @ ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:87 [inlined]
 [31] (::SciMLBase.var"#527#528"{Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, EnsembleProblem{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, SciMLBase.NullParameters, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xcb760beb, 0xef6a2ddc, 0xf02b42e4, 0x861e2e82, 0x7639da69)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x5024cc6a, 0x6af89867, 0x0ad7dad3, 0xcc69ec74, 0x21a02068)}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x477f85ee, 0xd0e8c452, 0xbd3625c2, 0xd167f375, 0xbb4b8d4e)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xca37cb04, 0x0af10484, 0x3c689b9e, 0xbcb1957d, 0x6e5acb8b)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Any}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ODESystem, Dict{Any, Any}}, Nothing, ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, typeof(prob_func), typeof(SciMLBase.DEFAULT_OUTPUT_FUNC), typeof(SciMLBase.DEFAULT_REDUCTION), Nothing}, Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}})(i::Int64)
    @ SciMLBase ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:146
 [32] responsible_map(f::Function, II::UnitRange{Int64})
    @ SciMLBase ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:139
 [33] #solve_batch#526
    @ ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:145 [inlined]
 [34] solve_batch
    @ ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:144 [inlined]
 [35] macro expansion
    @ ./timing.jl:382 [inlined]
 [36] __solve(prob::EnsembleProblem{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, SciMLBase.NullParameters, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xcb760beb, 0xef6a2ddc, 0xf02b42e4, 0x861e2e82, 0x7639da69)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x5024cc6a, 0x6af89867, 0x0ad7dad3, 0xcc69ec74, 0x21a02068)}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x477f85ee, 0xd0e8c452, 0xbd3625c2, 0xd167f375, 0xbb4b8d4e)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xca37cb04, 0x0af10484, 0x3c689b9e, 0xbcb1957d, 0x6e5acb8b)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Any}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ODESystem, Dict{Any, Any}}, Nothing, ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, typeof(prob_func), typeof(SciMLBase.DEFAULT_OUTPUT_FUNC), typeof(SciMLBase.DEFAULT_REDUCTION), Nothing}, alg::Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, ensemblealg::EnsembleSerial; trajectories::Int64, batch_size::Int64, pmap_batch_size::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLBase ~/.julia/packages/SciMLBase/QqtZA/src/ensemble/basic_ensemble_solve.jl:56
 [37] #solve#28
    @ ~/.julia/packages/DiffEqBase/ENGnO/src/solve.jl:872 [inlined]

This can be replicated by the following, albeit rather contrived, MWE:

using DifferentialEquations, ModelingToolkit, Distributions

do_jac = true
do_sparse = true

function lorenz!(du,u,p,t)
 du[1] = 10.0(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)

prob = ODEProblem(lorenz!, u0, tspan)
sys = modelingtoolkitize(prob)
oprob = ODEProblem(sys, u0, tspan, jac=do_jac, sparse=do_sparse)

new_u0() = [rand(Normal(1.0, 0.1)), 0.0, 0.0]
prob_func(prob, i, repeat) = remake(prob, u0=new_u0())

ensemble_prob = EnsembleProblem(oprob, prob_func=prob_func)
sol = solve(ensemble_prob, Tsit5(), EnsembleSerial(), trajectories=10)

When do_jac or do_sparse are independently true, the EnsembleProblem runs fine. Only when the two are combined does it cause this issue. Overloading Base.deepcopy_internal(x::Module, stackdict::IdDict) to print out the module in question reveals that Setfield is causing the issue.

The MWE was a fairly pointless example, but this is currently causing problems for me when solving large reaction networks from Catalyst.jl, which produces ODEProblems from ModelingToolkit ODESystems, and under which I require sparse analytic Jacobians in order to solve in reasonable times. A Catalyst-specific MWE is below:

using Catalyst, DifferentialEquations, Distributions

repressilator = @reaction_network Repressilator begin
    hillr(P₃,α,K,n), ∅ --> m₁
    hillr(P₁,α,K,n), ∅ --> m₂
    hillr(P₂,α,K,n), ∅ --> m₃
    (δ,γ), m₁ <--> ∅
    (δ,γ), m₂ <--> ∅
    (δ,γ), m₃ <--> ∅
    β, m₁ --> m₁ + P₁
    β, m₂ --> m₂ + P₂
    β, m₃ --> m₃ + P₃
    μ, P₁ --> ∅
    μ, P₂ --> ∅
    μ, P₃ --> ∅
end α K n δ γ β μ

@parameters  α K n δ γ β μ
@variables t m₁(t) m₂(t) m₃(t) P₁(t) P₂(t) P₃(t)
pmap  = (α => .5, K => 40, n => 2, δ => log(2)/120,
         γ => 5e-3, β => 20*log(2)/120, μ => log(2)/60)
u0map = [m₁ => 0., m₂ => 0., m₃ => 0., P₁ => 20., P₂ => 0., P₃ => 0.]
tspan = (0., 10000.)
pdist = [
    Normal(0.5, 0.01),
    Normal(40.0, 0.5),
    Normal(2, 1e-2),
    Normal(log(2)/120, 1e-4),
    Normal(5e-3, 1e-5),
    Normal(20*log(2)/120, 1e-3),
    Normal(log(2)/60, 5e-4)
]

oprob = ODEProblem(repressilator, u0map, tspan, pmap, jac=true, sparse=true)

function prob_func(prob, i, repeat)
    remake(prob, p=rand.(pdist))
end
ensemble_prob = EnsembleProblem(oprob, prob_func=prob_func)
sol = solve(ensemble_prob, Tsit5(), EnsembleSerial(); trajectories=50, saveat=10.)