add input preparation to ensemble problems

SciML · Oct 18, 2023 · 681dad8 · 681dad8
1 parent 578352c
commit 681dad8
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/src/ensemble/ensemble_problems.jl b/src/ensemble/ensemble_problems.jl
@@ -23,21 +23,25 @@ function EnsembleProblem(prob::AbstractVector{<:AbstractSciMLProblem}; kwargs...
         kwargs...)
 end
 function EnsembleProblem(prob;
-    output_func = DEFAULT_OUTPUT_FUNC,
     prob_func = DEFAULT_PROB_FUNC,
+    output_func = DEFAULT_OUTPUT_FUNC,
     reduction = DEFAULT_REDUCTION,
     u_init = nothing,
     safetycopy = prob_func !== DEFAULT_PROB_FUNC)
-    EnsembleProblem(prob, prob_func, output_func, reduction, u_init, safetycopy)
+    _prob_func = prepare_function(prob_func)
+    _output_func = prepare_function(output_func)
+    _reduction = prepare_function(reduction)
+    _u_init = prepare_initial_state(u_init)
+    EnsembleProblem(prob, _prob_func, _output_func, _reduction, _u_init, safetycopy)
 end
 
 function EnsembleProblem(; prob,
-    output_func = DEFAULT_OUTPUT_FUNC,
     prob_func = DEFAULT_PROB_FUNC,
+    output_func = DEFAULT_OUTPUT_FUNC,
     reduction = DEFAULT_REDUCTION,
     u_init = nothing, p = nothing,
     safetycopy = prob_func !== DEFAULT_PROB_FUNC)
-    EnsembleProblem(prob, prob_func, output_func, reduction, u_init, safetycopy)
+    EnsembleProblem(prob; prob_func, output_func, reduction, u_init, safetycopy)
 end
 
 struct WeightedEnsembleProblem{T1 <: AbstractEnsembleProblem, T2 <: AbstractVector} <: