fix indexing and don't use I

Author: ChrisRackauckas

src/ensemble/basic_ensemble_solve.jl

@@ -29,8 +29,8 @@ if (kwargs[:parallel_type] == != :none && kwargs[:parallel_type] == != :threads)
 end
 (batch_size != trajectories) && warn("batch_size and reductions are ignored when !collect_result")
 
-elapsed_time = @elapsed u = DArray((trajectories,)) do I
-    solve_batch(prob,alg,kwargs[:parallel_type] ==,I[1],pmap_batch_size,kwargs...)
+elapsed_time = @elapsed u = DArray((trajectories,)) do II
+    solve_batch(prob,alg,kwargs[:parallel_type] ==,II[1],pmap_batch_size,kwargs...)
 end
 return EnsembleSolution(u,elapsed_time,false)
 =#
@@ -88,8 +88,8 @@ function __solve(prob::AbstractEnsembleProblem,
   num_batches < 1 && error("trajectories ÷ batch_size cannot be less than 1, got $num_batches")
   num_batches * batch_size != trajectories && (num_batches += 1)
 
-  function batch_function(I)
-    batch_data = solve_batch(prob,alg,ensemblealg,I,pmap_batch_size;kwargs...)
+  function batch_function(II)
+    batch_data = solve_batch(prob,alg,ensemblealg,II,pmap_batch_size;kwargs...)
   end
 
   if num_batches == 1 && prob.reduction === DEFAULT_REDUCTION
@@ -99,19 +99,21 @@ function __solve(prob::AbstractEnsembleProblem,
   end
 
   converged::Bool = false
+  i = 1
+  II = (batch_size*(i-1)+1):batch_size*i
 
-  batch_data = batch_function(I)
+  batch_data = batch_function(II)
   u = prob.u_init === nothing ? similar(batch_data, 0) : prob.u_init
-  u,converged = prob.reduction(u,batch_data,I)
+  u,converged = prob.reduction(u,batch_data,II)
   elapsed_time = @elapsed for i in 2:num_batches
     converged && break
     if i == num_batches
-      I = (batch_size*(i-1)+1):trajectories
+      II = (batch_size*(i-1)+1):trajectories
     else
-      I = (batch_size*(i-1)+1):batch_size*i
+      II = (batch_size*(i-1)+1):batch_size*i
     end
-    batch_data = batch_function(I)
-    u,converged = prob.reduction(u,batch_data,I)
+    batch_data = batch_function(II)
+    u,converged = prob.reduction(u,batch_data,II)
   end
 
   u = reduce(vcat, u)
@@ -150,24 +152,24 @@ function batch_func(i,prob,alg;kwargs...)
   _x[1]
 end
 
-function solve_batch(prob,alg,ensemblealg::EnsembleDistributed,I,pmap_batch_size;kwargs...)
+function solve_batch(prob,alg,ensemblealg::EnsembleDistributed,II,pmap_batch_size;kwargs...)
   wp=CachingPool(workers())
-  batch_data = pmap(wp,I,batch_size=pmap_batch_size) do i
+  batch_data = pmap(wp,II,batch_size=pmap_batch_size) do i
     batch_func(i,prob,alg;kwargs...)
   end
   map(identity,batch_data)
 end
 
-function solve_batch(prob,alg,::EnsembleSerial,I,pmap_batch_size;kwargs...)
-  batch_data = map(I) do i
+function solve_batch(prob,alg,::EnsembleSerial,II,pmap_batch_size;kwargs...)
+  batch_data = map(II) do i
     batch_func(i,prob,alg;kwargs...)
   end
   batch_data
 end
 
-function solve_batch(prob,alg,ensemblealg::EnsembleThreads,I,pmap_batch_size;kwargs...)
+function solve_batch(prob,alg,ensemblealg::EnsembleThreads,II,pmap_batch_size;kwargs...)
   function multithreaded_batch(batch_idx)
-    i = I[batch_idx]
+    i = II[batch_idx]
     iter = 1
     _prob = prob.safetycopy ? deepcopy(prob.prob) : prob.prob
     new_prob = prob.prob_func(_prob,i,iter)
@@ -196,7 +198,7 @@ function solve_batch(prob,alg,ensemblealg::EnsembleThreads,I,pmap_batch_size;kwa
     _x[1]
   end
 
-  batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(I))})}(undef,length(I))
+  batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(II))})}(undef,length(II))
   let
     Threads.@threads for batch_idx in axes(batch_data, 1)
         batch_data[batch_idx] = multithreaded_batch(batch_idx)
@@ -205,26 +207,26 @@ function solve_batch(prob,alg,ensemblealg::EnsembleThreads,I,pmap_batch_size;kwa
   batch_data
 end
 
-function solve_batch(prob,alg,::EnsembleSplitThreads,I,pmap_batch_size;kwargs...)
+function solve_batch(prob,alg,::EnsembleSplitThreads,II,pmap_batch_size;kwargs...)
   wp=CachingPool(workers())
   N = nworkers()
-  batch_size = length(I)÷N
+  batch_size = length(II)÷N
   batch_data = let
     pmap(wp,1:N,batch_size=pmap_batch_size) do i
       if i == N
-        I_local = I[(batch_size*(i-1)+1):end]
+        I_local = II[(batch_size*(i-1)+1):end]
       else
-        I_local = I[(batch_size*(i-1)+1):(batch_size*i)]
+        I_local = II[(batch_size*(i-1)+1):(batch_size*i)]
       end
       thread_monte(prob,I_local,alg,i;kwargs...)
     end
   end
   reduce(vcat,batch_data)
 end
 
-function thread_monte(prob,I,alg,procid;kwargs...)
+function thread_monte(prob,II,alg,procid;kwargs...)
   function multithreaded_batch(j)
-    i = I[j]
+    i = II[j]
     iter = 1
     _prob = prob.safetycopy ? deepcopy(prob.prob) : prob.prob
     new_prob = prob.prob_func(_prob,i,iter)
@@ -253,10 +255,10 @@ function thread_monte(prob,I,alg,procid;kwargs...)
     _x[1]
   end
 
-  batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(I))})}(undef,length(I))
+  batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(II))})}(undef,length(II))
 
   let
-    Threads.@threads for j in 1:length(I)
+    Threads.@threads for j in 1:length(II)
       batch_data[j] = multithreaded_batch(j)
     end
   end