SciMLExpectations.jl icon indicating copy to clipboard operation
SciMLExpectations.jl copied to clipboard

Published 20 hours ago • verified publisher icon SciML

Problems evaluating Koopman Expectations over large parameter space

Open joegilkes opened this issue 2 years ago • 16 comments

Hi, first time using the package so I apologise if I'm missing something obvious.

I'm trying to propagate parametric uncertainty through a large, stiff system of ODEs representing a chemical reaction network (previous discussion here) using Koopman Expectations, such that I have a normal distribution over each of my ODESystem parameters and I am able to observe the mean and standard deviation over my state vector at every timestep specified by saveat. Here is my setup:

# `rates` represents reaction rate constants, the parameters of the model.
rates, rate_uncertainties = calc_rate_constants(...)
p_dist = [truncated(Normal(r, ru), r-(4*u), r+(4*u)) for (r, u) in zip(rates, rate_uncertainties)]

# `ODEProblem` constructed via Catalyst.jl, shouldn't be relevant.
oprob = ODEProblem(...)

# Want mean and variance for all saved timesteps.
g(sol) = [sol[1,:]; sol[1,:].^2]
gd = GenericDistribution(p_dist...)
# I believe this is the correct expression for `h(x, u, p)` when only applying uncertainty to parameters?
h(x, u, p) = u, x
sm = SystemMap(oprob, Rosenbrock23(); abstol=1e-14, reltol=1e-14, dtmin=1e-16, saveat=0.0:1e6:2e7)
exprob = ExpectationProblem(sm, g, h, gd; nout=2)
sol = solve(exprob, Koopman())

As far as I can tell this is set up correctly, but please correct me if not.

When I run this, it gets into the solve() line and then errors with the following:

ERROR: LoadError: OutOfMemoryError()
Stacktrace:
  [1] Array
    @ ./boot.jl:459 [inlined]
  [2] signcombos(k::Int64, λ::Float64, #unused#::Val{36})
    @ HCubature ~/.julia/packages/HCubature/QvyJW/src/genz-malik.jl:34
  [3] HCubature.GenzMalik(v::Val{36}, ::Type{Float64})
    @ HCubature ~/.julia/packages/HCubature/QvyJW/src/genz-malik.jl:99
  [4] cubrule
    @ ~/.julia/packages/HCubature/QvyJW/src/HCubature.jl:37 [inlined]
  [5] hcubature_(f::Integrals.var"#19#21"{IntegralProblem{false, RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, SciMLExpectations.var"#14#15"{SciMLExpectations.GenericDistribution{SciMLExpectations.var"#pdf_func#4"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, SciMLExpectations.var"#rand_func#6"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}}, AutoBreakdown.var"#h#45", AutoBreakdown.var"#g#44", SciMLExpectations.SystemMap{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11954d43, 0xfeb7cb16, 0x36117848, 0xb62dc5d6, 0x86beea97)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xf7d03d71, 0x4a56f8e0, 0x6b4c3b18, 0xa7e9f6b2, 0x184e3749)}}, UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xeb9be0d9, 0x805797f6, 0x35efa81e, 0xf7ba3bad, 0x26a49a9b)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x6e53ff16, 0x22f787c8, 0xc15faf6f, 0x2bf5f454, 0x2964e826)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Symbol}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ModelingToolkit.ODESystem, Dict{Any, Any}}, Nothing, ModelingToolkit.ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Rosenbrock23{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing}}, Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:saveat, :reltol, :progress_steps, :dtmin, :abstol, :progress, :kwargshandle), Tuple{Vector{Float64}, Float64, Int64, Float64, Float64, Bool, DataType}}}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}}, a::StaticArraysCore.SVector{36, Float64}, b::StaticArraysCore.SVector{36, Float64}, norm::typeof(norm), rtol_::Float64, atol::Float64, maxevals::Int64, initdiv::Int64)
    @ HCubature ~/.julia/packages/HCubature/QvyJW/src/HCubature.jl:56
  [6] #hcubature#3
    @ ~/.julia/packages/HCubature/QvyJW/src/HCubature.jl:179 [inlined]
  [7] __solvebp_call(::IntegralProblem{false, RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, SciMLExpectations.var"#14#15"{SciMLExpectations.GenericDistribution{SciMLExpectations.var"#pdf_func#4"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, SciMLExpectations.var"#rand_func#6"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}}, AutoBreakdown.var"#h#45", AutoBreakdown.var"#g#44", SciMLExpectations.SystemMap{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11954d43, 0xfeb7cb16, 0x36117848, 0xb62dc5d6, 0x86beea97)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xf7d03d71, 0x4a56f8e0, 0x6b4c3b18, 0xa7e9f6b2, 0x184e3749)}}, UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xeb9be0d9, 0x805797f6, 0x35efa81e, 0xf7ba3bad, 0x26a49a9b)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x6e53ff16, 0x22f787c8, 0xc15faf6f, 0x2bf5f454, 0x2964e826)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Symbol}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ModelingToolkit.ODESystem, Dict{Any, Any}}, Nothing, ModelingToolkit.ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Rosenbrock23{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing}}, Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:saveat, :reltol, :progress_steps, :dtmin, :abstol, :progress, :kwargshandle), Tuple{Vector{Float64}, Float64, Int64, Float64, Float64, Bool, DataType}}}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::Integrals.HCubatureJL, ::Integrals.ReCallVJP{Integrals.ZygoteVJP}, ::StaticArraysCore.SVector{36, Float64}, ::StaticArraysCore.SVector{36, Float64}, ::RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}; reltol::Float64, abstol::Float64, maxiters::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Integrals ~/.julia/packages/Integrals/KYdjn/src/Integrals.jl:260
  [8] #__solvebp#13
    @ ~/.julia/packages/Integrals/KYdjn/src/Integrals.jl:221 [inlined]
  [9] #solve#12
    @ ~/.julia/packages/Integrals/KYdjn/src/Integrals.jl:217 [inlined]
 [10] integrate(quadalg::Integrals.HCubatureJL, adalg::SciMLExpectations.NonfusedAD, f::Function, lb::StaticArraysCore.SVector{36, Float64}, ub::StaticArraysCore.SVector{36, Float64}, p::RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}; nout::Int64, batch::Int64, kwargs::Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:reltol, :abstol, :maxiters), Tuple{Float64, Float64, Int64}}})
    @ SciMLExpectations ~/.julia/packages/SciMLExpectations/DWytJ/src/expectation.jl:148
 [11] solve(::SciMLExpectations.ExpectationProblem{SciMLExpectations.SystemMap{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11954d43, 0xfeb7cb16, 0x36117848, 0xb62dc5d6, 0x86beea97)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xf7d03d71, 0x4a56f8e0, 0x6b4c3b18, 0xa7e9f6b2, 0x184e3749)}}, UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xeb9be0d9, 0x805797f6, 0x35efa81e, 0xf7ba3bad, 0x26a49a9b)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x6e53ff16, 0x22f787c8, 0xc15faf6f, 0x2bf5f454, 0x2964e826)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Symbol}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ModelingToolkit.ODESystem, Dict{Any, Any}}, Nothing, ModelingToolkit.ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Rosenbrock23{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing}}, Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:saveat, :reltol, :progress_steps, :dtmin, :abstol, :progress, :kwargshandle), Tuple{Vector{Float64}, Float64, Int64, Float64, Float64, Bool, DataType}}}}, AutoBreakdown.var"#g#44", AutoBreakdown.var"#h#45", SciMLExpectations.GenericDistribution{SciMLExpectations.var"#pdf_func#4"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, SciMLExpectations.var"#rand_func#6"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}}, RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}}, ::SciMLExpectations.Koopman{SciMLExpectations.NonfusedAD}; maxiters::Int64, batch::Int64, quadalg::Integrals.HCubatureJL, ireltol::Float64, iabstol::Float64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLExpectations ~/.julia/packages/SciMLExpectations/DWytJ/src/expectation.jl:134
 [12] solve(::SciMLExpectations.ExpectationProblem{SciMLExpectations.SystemMap{ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Vector{Float64}, ODEFunction{true, SciMLBase.AutoSpecialize, ModelingToolkit.var"#f#487"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x11954d43, 0xfeb7cb16, 0x36117848, 0xb62dc5d6, 0x86beea97)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xf7d03d71, 0x4a56f8e0, 0x6b4c3b18, 0xa7e9f6b2, 0x184e3749)}}, UniformScaling{Bool}, Nothing, Nothing, ModelingToolkit.var"#_jac#491"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0xeb9be0d9, 0x805797f6, 0x35efa81e, 0xf7ba3bad, 0x26a49a9b)}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", (0x6e53ff16, 0x22f787c8, 0xc15faf6f, 0x2bf5f454, 0x2964e826)}}, Nothing, Nothing, SparseArrays.SparseMatrixCSC{Float64, Int64}, Nothing, Nothing, Nothing, Nothing, Vector{Symbol}, Symbol, Vector{Symbol}, ModelingToolkit.var"#503#generated_observed#495"{Bool, ModelingToolkit.ODESystem, Dict{Any, Any}}, Nothing, ModelingToolkit.ODESystem}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Tuple{Rosenbrock23{0, true, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing}}, Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:saveat, :reltol, :progress_steps, :dtmin, :abstol, :progress, :kwargshandle), Tuple{Vector{Float64}, Float64, Int64, Float64, Float64, Bool, DataType}}}}, AutoBreakdown.var"#g#44", AutoBreakdown.var"#h#45", SciMLExpectations.GenericDistribution{SciMLExpectations.var"#pdf_func#4"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, SciMLExpectations.var"#rand_func#6"{NTuple{36, Distributions.Truncated{Distributions.Normal{Float64}, Distributions.Continuous, Float64}}}, StaticArraysCore.SVector{36, Float64}, StaticArraysCore.SVector{36, Float64}}, RecursiveArrayTools.ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}}, ::SciMLExpectations.Koopman{SciMLExpectations.NonfusedAD})
    @ SciMLExpectations ~/.julia/packages/SciMLExpectations/DWytJ/src/expectation.jl:125

Watching my system memory usage at the time indicates that I really shouldn't be running out of memory (~7 GB / 32 GB used when it errors), so I'm really not sure what to make of the issue. I'd like this to work on systems of >2000 ODES with > 7000 parameters, but right now this crash is happening on a much smaller system of 20 ODEs with only 36 parameters (hence the Rosenbrock23 solver, I usually use QNDF).

joegilkes avatar Jan 12 '23 11:01 joegilkes