Reduce memory usage for ensemble problems

Specific Domains Modelling & Simulations
1

Hi,

I have a test particle tracer which is built on top of OrdinaryDiffEq.jl. This model can have 3D numerical fields as input parameters and solve ODEs for each set of initial conditions.

However, when testing on the multiple particles setup, i.e. an ensemble problem, I encountered a memory bottleneck where the field arrays (wrapped by Interpolations.jl here) are replicated for each solution in the ensemble problem. For example, given a test EM field of ~ 80MB, each solution is also ~ 80MB in memory. More details can be found here.

In my case, the field parameter would be the same for all particles, but each test particle may have different initial condition. Is there a way to reuse the same interpolated field information for the ensemble problem to save memory?

1 Like
2

So to put the original question shorter and more precise: is it possible to share parameters for an ensemble problem?

For example,

using DifferentialEquations

function lorenz!(du,u,p,t)
    σ,ρ,β = p
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
end

"Initial state perturbation for EnsembleProblem."
function prob_func(prob, i, repeat)
   prob.u0[1] += i
   prob
end

u0 = [1.0,0.0,0.0]

p_tuple = (10,28,8/3)
#p_vector = Any[10,28, 8/3]

tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan,p_tuple)
#sol = solve(prob, Tsit5())
ensemble_prob = EnsembleProblem(prob; prob_func)

trajectories = 1

@time sols = solve(ensemble_prob, Tsit5(), EnsembleThreads(); trajectories,
   save_on=false, save_start=false, save_end=false,);

In this case, the parameter p = p_tuple = (10,28,8/3) is not shared: each trajectory will generate a copy. This may not be a problem for a parameter with 3 scalars, but in my original problem, two of the parameters are arrays, and I cannot afford to copy field arrays for each trajectory.

I guess if I use a vector type for the parameter, p = Any[10, 28, 8/3], then probably the parameters are referring to the same memory? Even this is true, there is another warning from DifferentialEquations

┌ Warning: Using arrays or dicts to store parameters of different types can hurt performance.
│ Consider using tuples instead.
└ @ SciMLBase C:\Users\hyzho\.julia\packages\SciMLBase\ZUwcW\src\performance_warnings.jl:32

I added Any on purpose, because in my original problem, param is of type

Tuple{Float64, Float64, TestParticle.Field{false, TestParticle.var"#get_field#3"{Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}, Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}, Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}}}, TestParticle.Field{false, TestParticle.var"#get_field#3"{Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}, Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}, Interpolations.ScaledInterpolation{Float64, 3, Interpolations.FilledExtrapolation{Float64, 3, Interpolations.BSplineInterpolation{Float64, 3, Array{Float64, 3}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}, Base.OneTo{Int64}}}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Float64}, Interpolations.BSpline{Interpolations.Linear{Interpolations.Throw{Interpolations.OnGrid}}}, Tuple{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}}}}}

and when I tried to convert that into an vector via collect(param), Julia gave up in type inference.

I got stuck here.

3

ensemble_prob = EnsembleProblem(prob; prob_func, safetycopy = false)

Use p = [10, 28, 8/3] instead.

2 Likes
4

I now have an additional question on the safetycopy keyword argument. The current version of the code let safetycopy control whether or not prob is deep copied or not, including u0 and param. Is it possible to have a fine-grain control such that u0 is deep copied whereas param is referenced? The reason I wish to have this is that in my specific problem, param is memory consuming but u0 is cheap. To avoid race conditions when running ensemble problems, I need different u0 for different threads but shared param.

5

You can always just set safetycopy=false and do the copying yourself.

1 Like