add safetycopy argument

Author: ChrisRackauckas

src/ensemble/basic_ensemble_solve.jl

@@ -104,9 +104,10 @@ function __solve(prob::AbstractEnsembleProblem,
   return EnsembleSolution(_u,elapsed_time,converged)
 end
 
-function batch_func(i,prob,alg,I,kwargs...)
+function batch_func(i,prob,alg,I,safetycopy,kwargs...)
   iter = 1
-  new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+  _prob = safetycopy ? deepcopy(prob.prob) : prob.prob
+  new_prob = prob.prob_func(_prob,i,iter)
   rerun = true
   x = prob.output_func(solve(new_prob,alg;kwargs...),i)
   if !(typeof(x) <: Tuple)
@@ -118,7 +119,8 @@ function batch_func(i,prob,alg,I,kwargs...)
   rerun = _x[2]
   while rerun
       iter += 1
-      new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+      _prob = safetycopy ? deepcopy(prob.prob) : prob.prob
+      new_prob = prob.prob_func(_prob,i,iter)
       x = prob.output_func(solve(new_prob,alg;kwargs...),i)
       if !(typeof(x) <: Tuple)
           @warn("output_func should return (out,rerun). See docs for updated details")
@@ -131,11 +133,11 @@ function batch_func(i,prob,alg,I,kwargs...)
   _x[1]
 end
 
-function solve_batch(prob,alg,::EnsembleDistributed,I,pmap_batch_size,kwargs...)
+function solve_batch(prob,alg,ensemblealg::EnsembleDistributed,I,pmap_batch_size,kwargs...)
   wp=CachingPool(workers())
   batch_data = let
     pmap(wp,I,batch_size=pmap_batch_size) do i
-      batch_func(i,prob,alg,I,kwargs...)
+      batch_func(i,prob,alg,I,ensemblealg.safetycopy,kwargs...)
     end
   end
   map(i->batch_data[i],1:length(batch_data))
@@ -150,13 +152,14 @@ function solve_batch(prob,alg,::EnsembleSerial,I,pmap_batch_size,kwargs...)
   map(i->batch_data[i],1:length(batch_data))
 end
 
-function solve_batch(prob,alg,::EnsembleThreads,I,pmap_batch_size,kwargs...)
+function solve_batch(prob,alg,ensemblealg::EnsembleThreads,I,pmap_batch_size,kwargs...)
   batch_data = Vector{Any}(undef,length(I))
   let
     Threads.@threads for batch_idx in axes(batch_data, 1)
         i = I[batch_idx]
         iter = 1
-        new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+        _prob = ensemblealg.safetycopy ? deepcopy(prob.prob) : prob.prob
+        new_prob = prob.prob_func(_prob,i,iter)
         x = prob.output_func(solve(new_prob,alg;kwargs...),i)
         if !(typeof(x) <: Tuple)
             @warn("output_func should return (out,rerun). See docs for updated details")
@@ -168,7 +171,8 @@ function solve_batch(prob,alg,::EnsembleThreads,I,pmap_batch_size,kwargs...)
 
         while rerun
             iter += 1
-            new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+            _prob = ensemblealg.safetycopy ? deepcopy(prob.prob) : prob.prob
+            new_prob = prob.prob_func(_prob,i,iter)
             x = prob.output_func(solve(new_prob,alg;kwargs...),i)
             if !(typeof(x) <: Tuple)
                 @warn("output_func should return (out,rerun). See docs for updated details")
@@ -195,20 +199,21 @@ function solve_batch(prob,alg,::EnsembleSplitThreads,I,pmap_batch_size,kwargs...
       else
         I_local = I[(batch_size*(i-1)+1):(batch_size*i)]
       end
-      thread_monte(prob,I_local,alg,i,kwargs...)
+      thread_monte(prob,I_local,alg,i,ensemblealg.safetycopy,kwargs...)
     end
   end
   _batch_data = vector_batch_data_to_arr(batch_data)
 end
 
-function thread_monte(prob,I,alg,procid,kwargs...)
+function thread_monte(prob,I,alg,procid,safetycopy,kwargs...)
   batch_data = Vector{Any}(undef,length(I))
   let
     j = 0
     Threads.@threads for i in I
       j += 1
       iter = 1
-      new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+      _prob = safetycopy ? deepcopy(prob.prob) : prob.prob
+      new_prob = prob.prob_func(_prob,i,iter)
       rerun = true
       x = prob.output_func(solve(new_prob,alg;kwargs...),i)
       if !(typeof(x) <: Tuple)
@@ -220,7 +225,8 @@ function thread_monte(prob,I,alg,procid,kwargs...)
       rerun = _x[2]
       while rerun
           iter += 1
-          new_prob = prob.prob_func(deepcopy(prob.prob),i,iter)
+          _prob = safetycopy ? deepcopy(prob.prob) : prob.prob
+          new_prob = prob.prob_func(_prob,i,iter)
           x = prob.output_func(solve(new_prob,alg;kwargs...),i)
           if !(typeof(x) <: Tuple)
               @warn("output_func should return (out,rerun). See docs for updated details")

src/ensemble/ensemble_problems.jl

@@ -7,18 +7,23 @@ struct EnsembleProblem{T,T2,T3,T4,T5} <: AbstractEnsembleProblem
   output_func::T3
   reduction::T4
   u_init::T5
+  safetycopy::Bool
 end
 
+DEFAULT_PROB_FUNC(prob,i,repeat) = prob
+DEFAULT_OUTPUT_FUNC(sol,i) = (sol,false)
+DEFAULT_REDUCTION(u,data,I) = (append!(u,data),false)
 EnsembleProblem(prob;
-    output_func = (sol,i)-> (sol,false),
-    prob_func= (prob,i,repeat)->prob,
-    reduction = (u,data,I)->(append!(u,data),false),
-    u_init = []) =
-    EnsembleProblem(prob,prob_func,output_func,reduction,u_init)
+    output_func = DEFAULT_OUTPUT_FUNC,
+    prob_func= DEFAULT_PROB_FUNC,
+    reduction = DEFAULT_REDUCTION,
+    u_init = [],
+    safetycopy = prob_func !== DEFAULT_PROB_FUNC) =
+    EnsembleProblem(prob,prob_func,output_func,reduction,u_init,safetycopy)
 
 EnsembleProblem(;prob,
-    output_func = (sol,i)-> (sol,false),
-    prob_func= (prob,i,repeat)->prob,
-    reduction = (u,data,I)->(append!(u,data),false),
-    u_init = [], p = nothing) =
-    EnsembleProblem(prob,prob_func,output_func,reduction,u_init)
+    output_func = DEFAULT_OUTPUT_FUNC,
+    prob_func= DEFAULT_PROB_FUNC,
+    reduction = DEFAULT_REDUCTION,
+    u_init = [], p = nothing, safetycopy = prob_func !== DEFAULT_PROB_FUNC) =
+    EnsembleProblem(prob,prob_func,output_func,reduction,u_init,safetycopy)

test/downstream/ensemble.jl

@@ -73,7 +73,9 @@ reduction = function (u,batch,I)
   u,((var(u)/sqrt(last(I)))/mean(u)<0.5) ? true : false
 end
 
-prob2 = EnsembleProblem(prob,prob_func=prob_func,output_func=output_func,reduction=reduction,u_init=Vector{Float64}())
+prob2 = EnsembleProblem(prob,prob_func=prob_func,output_func=output_func,
+                        reduction=reduction,u_init=Vector{Float64}(),
+                        safetycopy=false)
 sim = solve(prob2,Tsit5(),trajectories=10000,batch_size=20)
 @test sim.converged == true