Memory usage in the ensemble problems

[henry2004y]

According to my preliminary observations, the memory usage in the multi-particle tracing problem goes up quickly. On my 16 GB laptop, it cannot handle 300 particles simultaneously for an EM field of 80 MB in total. More detailed results to be followed.

[henry2004y]

Platform: Ubuntu 18, memory 15.3 GiB, swap memory 15.9 GiB
Before launching Julia: memory ~1.8 GiB, swap memory ~3.6 GiB
After starting Julia: memory ~2.0 GiB
Size of field data (E+M): 84 MiB
Launch mode: serial

Imported packages (allocations of which are included in the table below):

using TestParticle, OrdinaryDiffEq
using JLD2
using Vlasiator: RE

# of particles	Time span	Allocations	Memory usage	System memory	Swap memory
1	[0, 1]s	46.23M	3.715 GiB	4.1 GiB	3.1 GiB
50	[0, 1]s	46.26M	7.882 GiB	7.9 GiB	3.1 GiB
100	[0, 1]s	46.28M	12.134 GiB	11.3 GiB	3.6 GiB
150	[0, 1]s	46.30M	16.386 GiB	15.1 GiB	4.2 GiB
200	[0, 1]s	46.33M	20.638 GiB	15.1 GiB	8.6 GiB
250	[0, 1]s	46.35M	24.890 GiB	15.1 GiB	12.0 GiB
275	[0, 1]s	46.36M	27.016 GiB	15.1 GiB	14.9 GiB
300	[0, 1]s	N/A	OOM	N/A	N/A
100	[0, 10]s	46.32M	12.137 GiB	11.3 GiB	3.6 GiB
100	[0, 100]s	46.86M	12.181 GiB	11.3 GiB	3.6 GiB
100	[0, 1000]s	53.72M	12.731 GiB	11.7 GiB	3.6 GiB

If we save the solutions of 50 particles into a JLD2 file via

using JLD2
@save "particle.jld2" sols

it takes 4.3 GB, so ~ 86 MB per particle. Why is this so large?

As a comparison, without building the EnsembleProblem,

# of particles	Time span	Allocations	Memory usage	System memory	Swap memory
1	[0, 1]s	44.20M	3.523 GiB	3.8 GiB	3.1 GiB

Setting dense=false has no effect on memory usage.

[TCLiuu]

Platform: Ubuntu 18, memory 15.3 GiB, swap memory 15.9 GiB Before launching Julia: memory ~1.8 GiB, swap memory ~3.6 GiB After starting Julia: memory ~2.0 GiB Size of field data (E+M): 84 MiB Launch mode: serial

Imported packages (allocations of which are included in the table below):

using TestParticle, OrdinaryDiffEq
using JLD2
using Vlasiator: RE

of particles Time span Allocations Memory usage System memory Swap memory

1 [0, 1]s 46.23M 3.715 GiB 4.1 GiB 3.1 GiB
50 [0, 1]s 46.26M 7.882 GiB 7.9 GiB 3.1 GiB
100 [0, 1]s 46.28M 12.134 GiB 11.3 GiB 3.6 GiB
150 [0, 1]s 46.30M 16.386 GiB 15.1 GiB 4.2 GiB
200 [0, 1]s 46.33M 20.638 GiB 15.1 GiB 8.6 GiB
250 [0, 1]s 46.35M 24.890 GiB 15.1 GiB 12.0 GiB
275 [0, 1]s 46.36M 27.016 GiB 15.1 GiB 14.9 GiB
300 [0, 1]s N/A OOM N/A N/A
100 [0, 10]s 46.32M 12.137 GiB 11.3 GiB 3.6 GiB
100 [0, 100]s 46.86M 12.181 GiB 11.3 GiB 3.6 GiB
100 [0, 1000]s 53.72M 12.731 GiB 11.7 GiB 3.6 GiB
If we save the solutions of 50 particles into a JLD2 file via

using JLD2
@save "particle.jld2" sols

it takes 4.3 GB, so ~ 86 MB per particle. Why is this so large?

As a comparison, without building the EnsembleProblem,

of particles Time span Allocations Memory usage System memory Swap memory

1 [0, 1]s 44.20M 3.523 GiB 3.8 GiB 3.1 GiB
Setting dense=false has no effect on memory usage.

Can you provide your test file? My server has 64GB of RAM, so maybe we can compare the results.

[henry2004y]

Link to the EM file here

Note that if the X extent goes from [-26, 12] and the inner boundary sphere is at radius = 4.7. If you want more points, you may need to change the seeding procedure.

Test script

# Tracing particles in EGI.
#

using TestParticle, OrdinaryDiffEq
using JLD2
using Vlasiator: RE

function initial_conditions(i, prob)
   stateinit = prob.u0
   Δx = .05RE
   u0 = [stateinit[1] - (i-1)*Δx, stateinit[2:end]...]
end

"Initial state perturbation for EnsembleProblem."
function prob_func(prob, i, repeat)
   remake(prob, u0=initial_conditions(i, prob))
end

function isoutofdomain(u, p, t)
   if hypot(u[1], u[2], u[3]) < 4.7RE
      return true
   else
      return false
   end
end

function main(stateinit::AbstractVector; species::TestParticle.Species=Proton,
   tspan=(0.0, 1.0), trajectories::Int=1)
   file = "EM_EGI_t1000.jld2"
   data = load(file)
   E = data["E"]       # [V/m]
   B = data["B"]       # [T]
   x = data["x"] .* RE # [m]
   y = data["y"] .* RE # [m]
   z = data["z"] .* RE # [m]

   param = prepare(x, y, z, E, B; species)

   prob = ODEProblem(trace!, stateinit, tspan, param)
   #sol = solve(prob, Tsit5(); save_idxs=[1,2,3], isoutofdomain)
   ensemble_prob = EnsembleProblem(prob; prob_func)
   #dense=false
   sols = solve(ensemble_prob, Tsit5(), EnsembleThreads();
      trajectories, save_idxs=[1,2,3], isoutofdomain)

   return sols
end

stateinit = [-7RE, 0.0RE, 0RE, 0.0, 0.0, 0.0]

tspan = (0.0, 1.0)
trajectories = 10

# Electron
#@time sols = main(stateinit; species=Electron, tspan, trajectories);
# Proton
@time sols = main(stateinit; species=Proton, tspan, trajectories);
#=
using JLD2
@save "particle.jld2" sols
=#

Plotting script

# Test particle visualization

using GLMakie
using TestParticleMakie
using OrdinaryDiffEq
using JLD2
import Vlasiator: RE

#JLD2.@load "particle.jld2" sols

f = Figure(resolution = (1200, 800), fontsize = 18)
ax = Axis3(f[1, 1],
   title = "Particle trajectories, EGI, t=1000s",
   xlabel = L"X [$\text{R}_\text{E}$]",
   ylabel = L"Y [$\text{R}_\text{E}$]",
   zlabel = L"Z [$\text{R}_\text{E}$]",
   aspect = :data,
   elevation = 0.18π,
   azimuth = 0.85π,
)

invRE = 1/RE

xmin, xmax = -14.0, 8.0
ymin, ymax = -8.0, 8.0
zmin, zmax = -6.0, 6.0
# inner boundary
sphere = Sphere(Point3f(0.0), 4.7RE)
sph = mesh!(ax, sphere, color=(:white, 0.9), transparency=true, shading=true)
scale!(sph, invRE, invRE, invRE)
limits!(ax, xmin, xmax, ymin, ymax, zmin, zmax)

for i in eachindex(sols)
   l = lines!(sols[i], label="$i")
   scale!(l, invRE, invRE, invRE)
end
f

[TCLiuu]

# of particles	Time span	Allocations	Memory usage
1	[0, 1]s	46.51M	3.697 GiB
50	[0, 1]s	46.62M	7.868 GiB
100	[0, 1]s	46.67M	12.121 GiB
200	[0, 1]s	46.73M	20.626 GiB
250	[0, 1]s	46.74M	24.877 GiB
300	[0, 1]s	46.75M	29.129GiB
400	[0, 1]s	46.78M	37.632GiB
100	[0, 10]s	46.68M	12.123 GiB
100	[0, 100]s	46.86M	12.181 GiB

I got same results. And my server can handle 400 particles simultaneously, but I got lots of warnings from OrdinaryDiffEq.

[henry2004y]

I got lots of warnings from OrdinaryDiffEq.

The warnings can be turned off if we set verbose=false. I keep it there since during the testing phase I want to get a sense of how many particles reach the inner/outer boundary.

[henry2004y]

it takes 4.3 GB, so ~ 86 MB per particle. Why is this so large?

My current guess is that for each solution type, the numerical field information is replicated. This is consistent with the fact that the input EM field is ~ 84 MB. So the new question is, how can we reduce the memory usage?

[TCLiuu]

it takes 4.3 GB, so ~ 86 MB per particle. Why is this so large?

My current guess is that for each solution type, the numerical field information is replicated. This is consistent with the fact that the input EM field is ~ 84 MB. So the new question is, how can we reduce the memory usage?

Your guess sounds reasonable, maybe the ODEProblem will record the argument param. To verify this, simply change the EM field of a different size and see if the memory usage changes. But if this idea is right, I don't have any ideas to reduce this memory usage, perhaps starting with the design of ODEProblem or the implementation of AbstractField.

[henry2004y]

This is indeed because all the parameters are replicated for each particle trajectory. Here is a better way to handle the memory issue:

Running script

using TestParticle, OrdinaryDiffEq
using JLD2
using Vlasiator: RE

function isoutofdomain(u, p, t)
   if hypot(u[1], u[2], u[3]) < 4.7RE
      return true
   else
      return false
   end
end

struct Trajectory{T<:AbstractFloat}
   x::Vector{T}
   y::Vector{T}
   z::Vector{T}
end

function trace(stateinit::AbstractVector; species::TestParticle.Species=Proton,
   tspan=(0.0, 1.0), trajectories::Int=1, saveat::Float64=0.0)

   file = "EM_EGI_t1298.jld2"
   data = load(file)
   E = data["E"]       # [V/m]
   B = data["B"]       # [T]
   x = data["x"] .* RE # [m]
   y = data["y"] .* RE # [m]
   z = data["z"] .* RE # [m]

   param = prepare(x, y, z, E, B; species)

   prob = ODEProblem(trace!, stateinit, tspan, param)
   #
   sols = Vector{Trajectory}(undef, trajectories)
   Δx = 0.125RE / 8

   for i in 1:trajectories
      u0 = if i ≤ 40
         [stateinit[1] - (i-1)*Δx, stateinit[2], 0.4RE, stateinit[4:end]...]
      elseif i ≤ 80
         [stateinit[1] - (i-41)*Δx, stateinit[2], -0.4RE, stateinit[4:end]...]
      else
         [stateinit[1] - (i-81)*Δx, stateinit[2], 0.0, stateinit[4:end]...]
      end

      prob = remake(prob; u0, tspan)

      if saveat == 0.0
         sol = solve(prob, Tsit5(); save_idxs=[1,2,3], isoutofdomain)
      else
         sol = solve(prob, Tsit5(); save_idxs=[1,2,3], isoutofdomain, saveat)
      end
      x = getindex.(sol.u, 1)
      y = getindex.(sol.u, 2)
      z = getindex.(sol.u, 3)
      sols[i] = Trajectory(x, y, z)
   end

   return sols
end

stateinit = [-7RE, 0.0RE, 0RE, 0.0, 0.0, 0.0]

tspan = (0.0, 10.0)
trajectories = 100 * 8

@time sols = trace(stateinit; species=Proton, tspan, trajectories, saveat=0.05);

I need to update the part about ensemble problems.

[henry2004y]

I posted a question on Discourse.

[henry2004y]

Chris helped me solved this! Here is the trick:

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

[TCLiuu]

Chris helped me solved this! Here is the trick:

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

A nice trick. There are too many things buried in the documentation.