Avoiding memory allocations when simulating combinations of parameter sets in ParallelEnsemble

I have an array of arrays p = [a, b] for p in ode_model!(du, u, p, t). I am running a ParallelEnsemble where I want to simulate all combinations of a in A and b in B where A and B are arrays of arrays. I noticed that using the remake function to set p=[a, b] allocates some memory. I confirmed that [a, b] does indeed allocate e.g.

a=rand(5);
b=rand(5);
@benchmark [$a, $b]
BenchmarkTools.Trial:
  memory estimate:  96 bytes
  allocs estimate:  1
  --------------
  minimum time:     31.975 ns (0.00% GC)
  median time:      32.585 ns (0.00% GC)
  mean time:        35.434 ns (4.99% GC)
  maximum time:     1.148 ÎĽs (92.69% GC)
  --------------
  samples:          10000
  evals/sample:     993

It’s a small amount but with a ParallelEnsemble and a growing size of A or B, it ends up causing a lot of memory allocation.

Is there anyway to structure my parameter set p to avoid this allocation?

you can do prob.p[:,1] .= a; prob.p[:,2] .= b and that’ll reduce the allocations to zero on your end. But note shuttling to make things ship across machines and be thread-safe does require memory usage, probably more than the amount of memory you’re using.

Thanks. Is this safe? I am currently using remake(prob, p=[a,b]). I noticed that this will not affect p for the original prob but using prob.p[:,1] .= a; prob.p[:,2] .= b will also change the original prob.

Yeah it’s safe because we deepcopy the problem before we give it to the prob_func. This is something that I want to put an option on (and it wouldn’t be hard to), but we introduced that default a few years ago because when multithreading was first introduced we saw a lot of people doing exactly this (modifying the fields, the problem structs were also mutable back then though) and as you’d expect that was non-thread-safe chaos. We should probably reconsider some of the decisions though since now it’s safe by the immutability of the problem struct, so the deepcopy is probably still necessary by default (for example, if people are using closures for caches, that’s a nice trap), but it’s probably something we can add an option to remove.

Don’t worry, I have your MWEs open and plan to get to them soon. Maybe I can tackle this at the same time.

1 Like