Merge pull request #511 from SciML/myb/inference

Inference fix

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)
 =#
@@ -75,36 +75,55 @@ function __solve(prob::AbstractEnsembleProblem,
     __solve(prob,alg,ensemblealg;trajectories=trajectories,kwargs...)
 end
 
+tighten_container_eltype(u::Vector{Any}) = map(identity, u)
+tighten_container_eltype(u) = u
+
 function __solve(prob::AbstractEnsembleProblem,
                  alg::Union{DEAlgorithm,Nothing},
                  ensemblealg::BasicEnsembleAlgorithm;
                  trajectories, batch_size = trajectories,
                  pmap_batch_size = batch_size÷100 > 0 ? batch_size÷100 : 1, kwargs...)
 
   num_batches = trajectories ÷ batch_size
+  num_batches < 1 && error("trajectories ÷ batch_size cannot be less than 1, got $num_batches")
   num_batches * batch_size != trajectories && (num_batches += 1)
 
-  u = deepcopy(prob.u_init)
-  converged = false
-  elapsed_time = @elapsed for i in 1:num_batches
+  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
+    elapsed_time = @elapsed u = batch_function(1:trajectories)
+    _u = tighten_container_eltype(u)
+    return EnsembleSolution(_u,elapsed_time,true)
+  end
+
+  converged::Bool = false
+  i = 1
+  II = (batch_size*(i-1)+1):batch_size*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,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 = solve_batch(prob,alg,ensemblealg,I,pmap_batch_size,kwargs...)
-    u,converged = prob.reduction(u,batch_data,I)
-    converged && break
-  end
-  if typeof(u) <: Vector{Any}
-    _u = map(i->u[i],1:length(u))
-  else
-    _u = u
+    batch_data = batch_function(II)
+    u,converged = prob.reduction(u,batch_data,II)
   end
+
+  u = reduce(vcat, u)
+  _u = tighten_container_eltype(u)
+
   return EnsembleSolution(_u,elapsed_time,converged)
+
 end
 
-function batch_func(i,prob,alg,I,kwargs...)
+function batch_func(i,prob,alg;kwargs...)
   iter = 1
   _prob = prob.safetycopy ? deepcopy(prob.prob) : prob.prob
   new_prob = prob.prob_func(_prob,i,iter)
@@ -133,31 +152,38 @@ function batch_func(i,prob,alg,I,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 = let
-    pmap(wp,I,batch_size=pmap_batch_size) do i
-      batch_func(i,prob,alg,I,kwargs...)
-    end
+  batch_data = pmap(wp,II,batch_size=pmap_batch_size) do i
+    batch_func(i,prob,alg;kwargs...)
   end
-  map(i->batch_data[i],1:length(batch_data))
+  map(identity,batch_data)
 end
 
-function solve_batch(prob,alg,::EnsembleSerial,I,pmap_batch_size,kwargs...)
-  batch_data = let
-    map(I) do i
-      batch_func(i,prob,alg,I,kwargs...)
-    end
+function solve_batch(prob,alg,::EnsembleSerial,II,pmap_batch_size;kwargs...)
+  batch_data = map(II) do i
+    batch_func(i,prob,alg;kwargs...)
   end
-  map(i->batch_data[i],1:length(batch_data))
+  batch_data
 end
 
-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
+function solve_batch(prob,alg,ensemblealg::EnsembleThreads,II,pmap_batch_size;kwargs...)
+  function multithreaded_batch(batch_idx)
+    i = II[batch_idx]
+    iter = 1
+    _prob = 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")
+        _x = (x,false)
+    else
+      _x = x
+    end
+    rerun = _x[2]
+
+    while rerun
+        iter += 1
         _prob = 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)
@@ -168,87 +194,73 @@ function solve_batch(prob,alg,ensemblealg::EnsembleThreads,I,pmap_batch_size,kwa
           _x = x
         end
         rerun = _x[2]
+    end
+    _x[1]
+  end
 
-        while rerun
-            iter += 1
-            _prob = 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")
-                _x = (x,false)
-            else
-              _x = x
-            end
-            rerun = _x[2]
-        end
-        batch_data[batch_idx] = _x[1]
+  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)
     end
   end
-  map(i->batch_data[i],1:length(batch_data))
+  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...)
+      thread_monte(prob,I_local,alg,i;kwargs...)
     end
   end
-  _batch_data = vector_batch_data_to_arr(batch_data)
+  reduce(vcat,batch_data)
 end
 
-function thread_monte(prob,I,alg,procid,kwargs...)
-  batch_data = Vector{Any}(undef,length(I))
-  let
-    Threads.@threads for j in 1:length(I)
-      i = I[j]
-      iter = 1
-      _prob = 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)
-          @warn("output_func should return (out,rerun). See docs for updated details")
-          _x = (x,false)
-      else
-        _x = x
-      end
-      rerun = _x[2]
-      while rerun
-          iter += 1
-          _prob = 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")
-              _x = (x,false)
-          else
-            _x = x
-          end
-          rerun = _x[2]
-      end
-      batch_data[j] = _x[1]
+function thread_monte(prob,II,alg,procid;kwargs...)
+  function multithreaded_batch(j)
+    i = II[j]
+    iter = 1
+    _prob = 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)
+        @warn("output_func should return (out,rerun). See docs for updated details")
+        _x = (x,false)
+    else
+      _x = x
     end
+    rerun = _x[2]
+    while rerun
+        iter += 1
+        _prob = 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")
+            _x = (x,false)
+        else
+          _x = x
+        end
+        rerun = _x[2]
+    end
+    _x[1]
   end
-  batch_data
-end
 
-function vector_batch_data_to_arr(batch_data)
-  _batch_data = Vector{Any}(undef,sum((length(x) for x in batch_data)))
-  idx = 0
-  @inbounds for a in batch_data
-    for x in a
-      idx += 1
-      _batch_data[idx] = x
+  batch_data = Vector{Core.Compiler.return_type(multithreaded_batch,Tuple{typeof(first(II))})}(undef,length(II))
+
+  let
+    Threads.@threads for j in 1:length(II)
+      batch_data[j] = multithreaded_batch(j)
     end
   end
-  map(i->_batch_data[i],1:length(_batch_data))
+  batch_data
 end

src/ensemble/ensemble_problems.jl

@@ -12,18 +12,18 @@ end
 
 DEFAULT_PROB_FUNC(prob,i,repeat) = prob
 DEFAULT_OUTPUT_FUNC(sol,i) = (sol,false)
-DEFAULT_REDUCTION(u,data,I) = (append!(u,data),false)
+DEFAULT_REDUCTION(u,data,I) = push!(u, data), false
 EnsembleProblem(prob;
     output_func = DEFAULT_OUTPUT_FUNC,
     prob_func= DEFAULT_PROB_FUNC,
     reduction = DEFAULT_REDUCTION,
-    u_init = [],
+    u_init = nothing,
     safetycopy = prob_func !== DEFAULT_PROB_FUNC) =
     EnsembleProblem(prob,prob_func,output_func,reduction,u_init,safetycopy)
 
 EnsembleProblem(;prob,
     output_func = DEFAULT_OUTPUT_FUNC,
     prob_func= DEFAULT_PROB_FUNC,
     reduction = DEFAULT_REDUCTION,
-    u_init = [], p = nothing, safetycopy = prob_func !== DEFAULT_PROB_FUNC) =
+    u_init = nothing, p = nothing, safetycopy = prob_func !== DEFAULT_PROB_FUNC) =
     EnsembleProblem(prob,prob_func,output_func,reduction,u_init,safetycopy)

src/solve.jl

@@ -6,8 +6,11 @@ NO_TSPAN_PROBS = Union{AbstractLinearProblem, AbstractNonlinearProblem,
                        AbstractQuadratureProblem,
                        AbstractSteadyStateProblem,AbstractJumpProblem}
 
+has_kwargs(_prob::DEProblem) = has_kwargs(typeof(_prob))
+Base.@pure has_kwargs(::Type{T}) where T = :kwargs ∈ fieldnames(T)
+
 function init_call(_prob,args...;kwargs...)
-  if :kwargs ∈ propertynames(_prob)
+  if has_kwargs(_prob)
     __init(_prob,args...;_prob.kwargs...,kwargs...)
   else
     __init(_prob,args...;kwargs...)
@@ -34,7 +37,7 @@ function init(prob::DEProblem,args...;kwargs...)
 end
 
 function solve_call(_prob,args...;merge_callbacks = true, kwargs...)
-  if :kwargs ∈ propertynames(_prob)
+  if has_kwargs(_prob)
     if merge_callbacks && haskey(_prob.kwargs,:callback) && haskey(kwargs, :callback)
       kwargs_temp = NamedTuple{Base.diff_names(Base._nt_names(
       values(kwargs)), (:callback,))}(values(kwargs))
@@ -44,10 +47,19 @@ function solve_call(_prob,args...;merge_callbacks = true, kwargs...)
     kwargs = merge(values(_prob.kwargs), kwargs)
   end
 
-  logger = get(kwargs, :progress, false) ? default_logger(Logging.current_logger()) : nothing
-  maybe_with_logger(logger) do
-    __solve(_prob,args...; kwargs...)
+  T = Core.Compiler.return_type(__solve,Tuple{typeof(_prob),map(typeof, args)...})
+
+  progress = get(kwargs, :progress, false)
+  if progress
+    logger = default_logger(Logging.current_logger())
+    x = maybe_with_logger(logger) do
+      __solve(_prob,args...; kwargs...)
+    end
+    return x::T
+  else
+    __solve(_prob,args...; kwargs...)::T
   end
+
 end
 
 function solve(prob::DEProblem,args...;kwargs...)

test/downstream/ensemble.jl

@@ -80,7 +80,10 @@ sim = solve(prob2,Tsit5(),trajectories=10000,batch_size=20)
 @test sim.converged == true
 
 
-Random.seed!(100)
+prob_func = function (prob,i,repeat)
+  ODEProblem(prob.f,(1 + i/100)*prob.u0,prob.tspan,1.01)
+end
+
 reduction = function (u,batch,I)
   u = append!(u,batch)
   u,false
@@ -90,7 +93,6 @@ prob2 = EnsembleProblem(prob,prob_func=prob_func,output_func=output_func,reducti
 sim = solve(prob2,Tsit5(),trajectories=100,batch_size=20)
 @test sim.converged == false
 
-Random.seed!(100)
 reduction = function (u,batch,I)
   u+sum(batch),false
 end
@@ -105,4 +107,4 @@ output_func = function (sol,i)
 end
 prob2 = EnsembleProblem(prob,prob_func=prob_func,output_func=output_func)
 sim2 = solve(prob2,Tsit5(),trajectories=2)
-@test !sim2.converged && typeof(sim2.u) == Vector{SomeUserType}
+@test sim2.converged && typeof(sim2.u) == Vector{SomeUserType}