Merge pull request #470 from chriselrod/addweightedensembleproblem

Add  `WeightedEnsembleProblem`

Author: ChrisRackauckas

src/ensemble/ensemble_problems.jl

@@ -35,3 +35,22 @@ function EnsembleProblem(; prob,
     safetycopy = prob_func !== DEFAULT_PROB_FUNC)
     EnsembleProblem(prob, prob_func, output_func, reduction, u_init, safetycopy)
 end
+
+struct WeightedEnsembleProblem{T1<:AbstractEnsembleProblem, T2<:AbstractVector} <: AbstractEnsembleProblem
+  ensembleprob::T1
+  weights::T2
+end
+Base.propertynames(e::WeightedEnsembleProblem) = (Base.propertynames(getfield(e, :ensembleprob))..., :weights)
+function Base.getproperty(e::WeightedEnsembleProblem, f::Symbol)
+  f === :weights && return getfield(e, :weights)
+  f === :ensembleprob && return getfield(e, :ensembleprob)
+  return getproperty(getfield(e, :ensembleprob), f)
+end
+function WeightedEnsembleProblem(args...; weights, kwargs...)
+  # TODO: allow skipping checks?
+  @assert sum(weights) ≈ 1
+  ep = EnsembleProblem(args...; kwargs...)
+  @assert length(ep.prob) == length(weights)
+  WeightedEnsembleProblem(ep, weights)
+end
+

src/ensemble/ensemble_solutions.jl

@@ -46,6 +46,15 @@ function EnsembleSolution(sim::T, elapsedTime,
         converged)
 end
 
+struct WeightedEnsembleSolution{T1<:AbstractEnsembleSolution, T2<:Number}
+    ensol::T1
+    weights::Vector{T2}
+    function WeightedEnsembleSolution(ensol, weights)
+        @assert length(weights) == length(ensol)
+        new{typeof(ensol), eltype(weights)}(ensol, weights)
+    end
+end
+
 function Base.reverse(sim::EnsembleSolution)
     EnsembleSolution(reverse(sim.u), sim.elapsedTime, sim.converged)
 end
@@ -213,3 +222,7 @@ end
 function (sol::AbstractEnsembleSolution)(args...; kwargs...)
     [s(args...; kwargs...) for s in sol]
 end
+
+Base.@propagate_inbounds function Base.getindex(sol::WeightedEnsembleSolution, S)
+    return [sum(stack(sol.weights .* sol.ensol[:, S]), dims = 2)]
+end