From 1727d7f58f467778e767952a1d83ee1e2ff60e1d Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Wed, 12 Jul 2023 16:04:53 -0400 Subject: [PATCH 01/16] Weighted Bayesian ensemblefits --- Project.toml | 1 + src/EasyModelAnalysis.jl | 1 + src/datafit.jl | 179 +++++++++++++++++++++++++++++++++++++-- 3 files changed, 174 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index 6ff7d332..eae1be4c 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ OptimizationBBO = "3e6eede4-6085-4f62-9a71-46d9bc1eb92b" OptimizationMOI = "fd9f6733-72f4-499f-8506-86b2bdd0dea1" OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLExpectations = "afe9f18d-7609-4d0e-b7b7-af0cb72b8ea8" diff --git a/src/EasyModelAnalysis.jl b/src/EasyModelAnalysis.jl index bfd7021f..3a0c0b0b 100644 --- a/src/EasyModelAnalysis.jl +++ b/src/EasyModelAnalysis.jl @@ -10,6 +10,7 @@ using GlobalSensitivity, Turing using SciMLExpectations @reexport using Plots using SciMLBase.EnsembleAnalysis +using Random include("basics.jl") include("datafit.jl") diff --git a/src/datafit.jl b/src/datafit.jl index ed1322bc..23f090c7 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -192,16 +192,15 @@ function global_datafit(prob, pbounds, data; maxiters = 10000, loss = l2loss) Pair.(pkeys, res.u) end -function bayes_unpack_data(p, data) - pdist = getfield.(p, :second) - pkeys = getfield.(p, :first) +function bayes_unpack_data(p::AbstractVector{<:Pair}, data) + pdist, pkeys = bayes_unpack_data(p) ts = first.(last.(data)) lastt = maximum(last, ts) timeseries = last.(last.(data)) datakeys = first.(data) (pdist, pkeys, ts, lastt, timeseries, datakeys) end -function bayes_unpack_data(p) +function bayes_unpack_data(p::AbstractVector{<:Pair}) pdist = getfield.(p, :second) pkeys = getfield.(p, :first) (pdist, pkeys) @@ -249,6 +248,123 @@ Turing.@model function bayesianODE(prob, return nothing end +""" +Weights can be unbounded. Length of weights must be one less than the length of sols, to apply a sum-to-1 constraint. +Last `sol` is given the weight `1 - sum(weights)`. +""" +struct WeightedSol{T, S <: Tuple{Vararg{AbstractVector{T}}}, W} <: AbstractVector{T} + sols::S + weights::W + function WeightedSol{T}(sols::S, + weights::W) where {T, S <: Tuple{Vararg{AbstractVector{T}}}, W} + @assert length(sols) == length(weights) + 1 + new{T, S, W}(sols, weights) + end +end +Base.length(ws::WeightedSol) = length(first(ws.sols)) +Base.size(ws::WeightedSol) = (length(first(ws.sols)),) +function Base.getindex(ws::WeightedSol{T}, i::Int) where {T} + s = zero(T) + w = zero(T) + for j in eachindex(ws.weights) + w += ws.weights[j] + s += ws.weights[j] * ws.sols[j][i] + end + return s + (one(T) - w) * ws.sols[end][i] +end +function WeightedSol(sols, select, weights) + T = eltype(weights) + s = map(Base.Fix2(getindex, select), sols) + WeightedSol{T}(s, weights) +end +function bayes_unpack_data(p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data) + pdist, pkeys = bayes_unpack_data(p) + ts = first.(last.(data)) + lastt = maximum(last, ts) + timeseries = last.(last.(data)) + datakeys = first.(data) + (pdist, pkeys, ts, lastt, timeseries, datakeys) +end +function bayes_unpack_data(p::Tuple{Vararg{<:AbstractVector{<:Pair}}}) + unpacked = map(bayes_unpack_data, p) + map(first, unpacked), map(last, unpacked) +end + +struct Grouper{N} + sizes::NTuple{N, Int} +end +function (g::Grouper)(x) + i = Ref(0) + map(g.sizes) do N + _i = i[] + i[] = _i + N + view(x, (_i + 1):(_i + N)) + end +end +function flatten(x::Tuple) + reduce(vcat, x), Grouper(map(length, x)) +end + +Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, + t, + pdist, + grouppriorsfunc, + probspkeys, + data, + noise_prior) + σ ~ noise_prior + + ppriors ~ product_distribution(pdist) + # stdeviation = sqrt(1/length(weights)) + Nprobs = length(probs) + Nprobs⁻¹ = inv(Nprobs) + weights ~ MvNormal(Distributions.Fill(Nprobs⁻¹, Nprobs - 1), Nprobs⁻¹) + sols = map(probs, probspkeys, grouppriorsfunc(ppriors)) do prob, pkeys, pprior + solve(remake(prob, tspan = (prob.tspan[1], t[end]), p = Pair.(pkeys, pprior)), + saveat = t) + end + if !all(SciMLBase.successful_retcode, sols) + Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) + return nothing + end + for i in eachindex(data) + data[i].second ~ MvNormal(WeightedSol(sols, data[i].first, weights), σ^2 * I) + end + return nothing +end +Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, + pdist, + grouppriorsfunc, + probspkeys, + ts, + lastt, + timeseries, + datakeys, + noise_prior) + σ ~ noise_prior + + ppriors ~ product_distribution(pdist) + + sols = map(probs, probspkeys, grouppriorsfunc(ppriors)) do prob, pkeys, pprior + solve(remake(prob, tspan = (prob.tspan[1], lastt), p = Pair.(pkeys, pprior))) + end + # stdeviation = sqrt(1/length(weights)) + Nprobs = length(probs) + Nprobs⁻¹ = inv(Nprobs) + weights ~ MvNormal(Distributions.Fill(Nprobs⁻¹, Nprobs - 1), Nprobs⁻¹) + if !all(SciMLBase.successful_retcode, sols) + Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) + return nothing + end + for i in eachindex(datakeys) + vals = map(sols) do sol + sol(ts[i]; idxs = datakeys[i]) + end + timeseries[i] ~ MvNormal(WeightedSol(vals, weights), σ^2 * I) + end + return nothing +end + """ bayesian_datafit(prob, p, t, data) bayesian_datafit(prob, p, data) @@ -288,8 +404,8 @@ function bayesian_datafit(prob, mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1000) - (pkeys, pdata) = bayes_unpack_data(p) - model = bayesianODE(prob, t, pkeys, pdata, data, noise_prior) + (pdist, pkeys) = bayes_unpack_data(p) + model = bayesianODE(prob, t, pdist, pkeys, data, noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -301,7 +417,7 @@ function bayesian_datafit(prob, end function bayesian_datafit(prob, - p, + p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), @@ -318,6 +434,55 @@ function bayesian_datafit(prob, [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) for i in eachindex(p)] end +function bayesian_datafit(probs::Union{Tuple, AbstractVector}, + p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, + t, + data; + noise_prior = InverseGamma(2, 3), + mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + nchains = 4, + niter = 1000) + (pdist_, pkeys) = bayes_unpack_data(p) + pdist, grouppriorsfunc = flatten(pdist_) + + model = ensemblebayesianODE(probs, t, pdist, grouppriorsfunc, pkeys, data, noise_prior) + chain = Turing.sample(model, + Turing.NUTS(0.65), + mcmcensemble, + niter, + nchains; + progress = true) + [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) + for i in eachindex(p)] +end + +function bayesian_datafit(probs::Union{Tuple, AbstractVector}, + p, + data; + noise_prior = InverseGamma(2, 3), + mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + nchains = 4, + niter = 1_000) + pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(p, data) + pdist, grouppriorsfunc = flatten(pdist_) + model = ensemblebayesianODE(probs, + pdist, + grouppriorsfunc, + pkeys, + ts, + lastt, + timeseries, + datakeys, + noise_prior) + chain = Turing.sample(model, + Turing.NUTS(0.65), + mcmcensemble, + niter, + nchains; + progress = true) + [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) + for i in eachindex(p)] +end """ model_forecast_score(probs::AbstractVector, ts::AbstractVector, dataset::AbstractVector{<:Pair}) From 6615f071dd9fdcd6e0878d184c00fe92c9ba1cff Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 10:09:48 -0400 Subject: [PATCH 02/16] separate `getsols` for easier introspection --- src/datafit.jl | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index 23f090c7..2419a13a 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -305,6 +305,18 @@ function flatten(x::Tuple) reduce(vcat, x), Grouper(map(length, x)) end +function getsols(probs, probspkeys, ppriors, t::AbstractArray) + map(probs, probspkeys, ppriors) do prob, pkeys, pprior + solve(remake(prob, tspan = (prob.tspan[1], t[end]), p = Pair.(pkeys, pprior)), + saveat = t) + end +end +function getsols(probs, probspkeys, ppriors, lastt::Number) + map(probs, probspkeys, ppriors) do prob, pkeys, pprior + solve(remake(prob, tspan = (prob.tspan[1], lastt), p = Pair.(pkeys, pprior))) + end +end + Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, t, pdist, @@ -313,16 +325,12 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, data, noise_prior) σ ~ noise_prior - ppriors ~ product_distribution(pdist) - # stdeviation = sqrt(1/length(weights)) + Nprobs = length(probs) Nprobs⁻¹ = inv(Nprobs) weights ~ MvNormal(Distributions.Fill(Nprobs⁻¹, Nprobs - 1), Nprobs⁻¹) - sols = map(probs, probspkeys, grouppriorsfunc(ppriors)) do prob, pkeys, pprior - solve(remake(prob, tspan = (prob.tspan[1], t[end]), p = Pair.(pkeys, pprior)), - saveat = t) - end + sols = getsols(probs, probspkeys, grouppriorsfunc(ppriors), t) if !all(SciMLBase.successful_retcode, sols) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing @@ -342,13 +350,10 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, datakeys, noise_prior) σ ~ noise_prior - ppriors ~ product_distribution(pdist) - sols = map(probs, probspkeys, grouppriorsfunc(ppriors)) do prob, pkeys, pprior - solve(remake(prob, tspan = (prob.tspan[1], lastt), p = Pair.(pkeys, pprior))) - end - # stdeviation = sqrt(1/length(weights)) + sols = getsols(probs, probspkeys, grouppriorsfunc(ppriors), lastt) + Nprobs = length(probs) Nprobs⁻¹ = inv(Nprobs) weights ~ MvNormal(Distributions.Fill(Nprobs⁻¹, Nprobs - 1), Nprobs⁻¹) From 493a359d4cf0f225a2efc5fcb3ad48ec1b51063a Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 11:41:27 -0400 Subject: [PATCH 03/16] Calculate indices ahead of time for mapping parameters in Bayesian We should probably also do something like this for `datafit` and `global_datafit`. --- src/EasyModelAnalysis.jl | 1 + src/datafit.jl | 31 ++++++++++++++++--------------- src/keyindexmap.jl | 29 +++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 15 deletions(-) create mode 100644 src/keyindexmap.jl diff --git a/src/EasyModelAnalysis.jl b/src/EasyModelAnalysis.jl index 3a0c0b0b..c5ad655b 100644 --- a/src/EasyModelAnalysis.jl +++ b/src/EasyModelAnalysis.jl @@ -13,6 +13,7 @@ using SciMLBase.EnsembleAnalysis using Random include("basics.jl") +include("keyindexmap.jl") include("datafit.jl") include("sensitivity.jl") include("threshold.jl") diff --git a/src/datafit.jl b/src/datafit.jl index 2419a13a..df924937 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -192,18 +192,18 @@ function global_datafit(prob, pbounds, data; maxiters = 10000, loss = l2loss) Pair.(pkeys, res.u) end -function bayes_unpack_data(p::AbstractVector{<:Pair}, data) - pdist, pkeys = bayes_unpack_data(p) +function bayes_unpack_data(prob, p::AbstractVector{<:Pair}, data) + pdist, pkeys = bayes_unpack_data(prob, p) ts = first.(last.(data)) lastt = maximum(last, ts) timeseries = last.(last.(data)) datakeys = first.(data) (pdist, pkeys, ts, lastt, timeseries, datakeys) end -function bayes_unpack_data(p::AbstractVector{<:Pair}) +function bayes_unpack_data(prob, p::AbstractVector{<:Pair}) pdist = getfield.(p, :second) pkeys = getfield.(p, :first) - (pdist, pkeys) + (pdist, IndexKeyMap(prob, pkeys)) end Turing.@model function bayesianODE(prob, t, pdist, pkeys, data, noise_prior) @@ -211,7 +211,7 @@ Turing.@model function bayesianODE(prob, t, pdist, pkeys, data, noise_prior) pprior ~ product_distribution(pdist) - prob = remake(prob, tspan = (prob.tspan[1], t[end]), p = Pair.(pkeys, pprior)) + prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) sol = solve(prob, saveat = t) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) @@ -235,7 +235,7 @@ Turing.@model function bayesianODE(prob, pprior ~ product_distribution(pdist) - prob = remake(prob, tspan = (prob.tspan[1], lastt), p = Pair.(pkeys, pprior)) + prob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) sol = solve(prob) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) @@ -277,16 +277,16 @@ function WeightedSol(sols, select, weights) s = map(Base.Fix2(getindex, select), sols) WeightedSol{T}(s, weights) end -function bayes_unpack_data(p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data) - pdist, pkeys = bayes_unpack_data(p) +function bayes_unpack_data(probs, p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data) + pdist, pkeys = bayes_unpack_data(probs, p) ts = first.(last.(data)) lastt = maximum(last, ts) timeseries = last.(last.(data)) datakeys = first.(data) (pdist, pkeys, ts, lastt, timeseries, datakeys) end -function bayes_unpack_data(p::Tuple{Vararg{<:AbstractVector{<:Pair}}}) - unpacked = map(bayes_unpack_data, p) +function bayes_unpack_data(probs, p::Tuple{Vararg{<:AbstractVector{<:Pair}}}) + unpacked = map(bayes_unpack_data, probs, p) map(first, unpacked), map(last, unpacked) end @@ -307,13 +307,14 @@ end function getsols(probs, probspkeys, ppriors, t::AbstractArray) map(probs, probspkeys, ppriors) do prob, pkeys, pprior - solve(remake(prob, tspan = (prob.tspan[1], t[end]), p = Pair.(pkeys, pprior)), - saveat = t) + newprob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + solve(newprob, saveat = t) end end function getsols(probs, probspkeys, ppriors, lastt::Number) map(probs, probspkeys, ppriors) do prob, pkeys, pprior - solve(remake(prob, tspan = (prob.tspan[1], lastt), p = Pair.(pkeys, pprior))) + newprob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) + solve(newprob) end end @@ -409,7 +410,7 @@ function bayesian_datafit(prob, mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1000) - (pdist, pkeys) = bayes_unpack_data(p) + (pdist, pkeys) = bayes_unpack_data(prob, p) model = bayesianODE(prob, t, pdist, pkeys, data, noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), @@ -428,7 +429,7 @@ function bayesian_datafit(prob, mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1_000) - pdist, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(p, data) + pdist, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(prob, p, data) model = bayesianODE(prob, pdist, pkeys, ts, lastt, timeseries, datakeys, noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), diff --git a/src/keyindexmap.jl b/src/keyindexmap.jl new file mode 100644 index 00000000..a2a863af --- /dev/null +++ b/src/keyindexmap.jl @@ -0,0 +1,29 @@ + +struct IndexKeyMap + indices::Vector{Int} +end + +function IndexKeyMap(prob, keys) + params = ModelingToolkit.parameters(prob.f.sys) + indices = Vector{Int}(undef, length(keys)) + for i in eachindex(keys) + indices[i] = findfirst(Base.Fix1(isequal, keys[i]), params) + end + return IndexKeyMap(indices) +end + +Base.@propagate_inbounds function (ikm::IndexKeyMap)(prob, v::AbstractVector) + @boundscheck checkbounds(v, length(ikm.indices)) + def = prob.p + ret = Vector{Base.promote_eltype(v, def)}(undef, length(def)) + copyto!(ret, def) + for (i, j) in enumerate(ikm.indices) + @inbounds ret[j] = v[i] + end + return ret +end + +function _remake(prob, tspan, ikm::IndexKeyMap, pprior) + p = ikm(prob, pprior) + remake(prob; tspan, p) +end From 2d0fd38cde9d5d12e9d76eb2f453f7d34f345c0c Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 11:55:19 -0400 Subject: [PATCH 04/16] fix type signatures --- src/datafit.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index df924937..29e4f1fa 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -423,7 +423,7 @@ function bayesian_datafit(prob, end function bayesian_datafit(prob, - p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, + p, data; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), @@ -463,7 +463,7 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, end function bayesian_datafit(probs::Union{Tuple, AbstractVector}, - p, + p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), From 6545e78dc325aaaaf4b06f9df35919c611b8ec22 Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 14:00:58 -0400 Subject: [PATCH 05/16] Type stability --- src/datafit.jl | 71 ++++++++++++++++++++++++++++++---------------- src/keyindexmap.jl | 18 ++++++++++-- 2 files changed, 63 insertions(+), 26 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index 29e4f1fa..f47e02a0 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -206,24 +206,24 @@ function bayes_unpack_data(prob, p::AbstractVector{<:Pair}) (pdist, IndexKeyMap(prob, pkeys)) end -Turing.@model function bayesianODE(prob, t, pdist, pkeys, data, noise_prior) +Turing.@model function bayesianODE(prob, alg, t, pdist, pkeys, data, datamap, noise_prior) σ ~ noise_prior pprior ~ product_distribution(pdist) prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) - sol = solve(prob, saveat = t) + sol = solve(prob, alg, saveat = t) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing end for i in eachindex(data) - data[i].second ~ MvNormal(sol[data[i].first], σ^2 * I) + data[i] ~ MvNormal(datamap(sol), σ^2 * I) end return nothing end -Turing.@model function bayesianODE(prob, +Turing.@model function bayesianODE(prob, alg, pdist, pkeys, ts, @@ -236,7 +236,7 @@ Turing.@model function bayesianODE(prob, pprior ~ product_distribution(pdist) prob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) - sol = solve(prob) + sol = solve(prob, alg) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing @@ -264,18 +264,19 @@ end Base.length(ws::WeightedSol) = length(first(ws.sols)) Base.size(ws::WeightedSol) = (length(first(ws.sols)),) function Base.getindex(ws::WeightedSol{T}, i::Int) where {T} - s = zero(T) - w = zero(T) - for j in eachindex(ws.weights) + s::T = zero(T) + w::T = zero(T) + @inbounds for j in eachindex(ws.weights) w += ws.weights[j] s += ws.weights[j] * ws.sols[j][i] end return s + (one(T) - w) * ws.sols[end][i] end -function WeightedSol(sols, select, weights) - T = eltype(weights) - s = map(Base.Fix2(getindex, select), sols) - WeightedSol{T}(s, weights) +function WeightedSol(sols, select, i::Int, weights) + s = map(sols, select) do sol, sel + @view(sol[sel.indices[i], :]) + end + WeightedSol{eltype(weights)}(s, weights) end function bayes_unpack_data(probs, p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data) pdist, pkeys = bayes_unpack_data(probs, p) @@ -305,25 +306,27 @@ function flatten(x::Tuple) reduce(vcat, x), Grouper(map(length, x)) end -function getsols(probs, probspkeys, ppriors, t::AbstractArray) - map(probs, probspkeys, ppriors) do prob, pkeys, pprior +function getsols(probs, algs, probspkeys, ppriors, t::AbstractArray) + map(probs, algs, probspkeys, ppriors) do prob, alg, pkeys, pprior newprob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) - solve(newprob, saveat = t) + solve(newprob, alg, saveat = t) end end -function getsols(probs, probspkeys, ppriors, lastt::Number) - map(probs, probspkeys, ppriors) do prob, pkeys, pprior +function getsols(probs, algs, probspkeys, ppriors, lastt::Number) + map(probs, algs, probspkeys, ppriors) do prob, alg, pkeys, pprior newprob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) - solve(newprob) + solve(newprob, alg) end end Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, + algs, t, pdist, grouppriorsfunc, probspkeys, data, + datamaps, noise_prior) σ ~ noise_prior ppriors ~ product_distribution(pdist) @@ -331,17 +334,18 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, Nprobs = length(probs) Nprobs⁻¹ = inv(Nprobs) weights ~ MvNormal(Distributions.Fill(Nprobs⁻¹, Nprobs - 1), Nprobs⁻¹) - sols = getsols(probs, probspkeys, grouppriorsfunc(ppriors), t) + sols = getsols(probs, algs, probspkeys, grouppriorsfunc(ppriors), t) if !all(SciMLBase.successful_retcode, sols) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing end for i in eachindex(data) - data[i].second ~ MvNormal(WeightedSol(sols, data[i].first, weights), σ^2 * I) + data[i] ~ MvNormal(WeightedSol(sols, datamaps, i, weights), σ^2 * I) end return nothing end Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, + algs, pdist, grouppriorsfunc, probspkeys, @@ -353,7 +357,7 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, σ ~ noise_prior ppriors ~ product_distribution(pdist) - sols = getsols(probs, probspkeys, grouppriorsfunc(ppriors), lastt) + sols = getsols(probs, algs, probspkeys, grouppriorsfunc(ppriors), lastt) Nprobs = length(probs) Nprobs⁻¹ = inv(Nprobs) @@ -411,7 +415,14 @@ function bayesian_datafit(prob, nchains = 4, niter = 1000) (pdist, pkeys) = bayes_unpack_data(prob, p) - model = bayesianODE(prob, t, pdist, pkeys, data, noise_prior) + model = bayesianODE(prob, + first(default_algorithm(prob)), + t, + pdist, + pkeys, + last.(data), + IndexKeyMap(prob, data), + noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -430,7 +441,15 @@ function bayesian_datafit(prob, nchains = 4, niter = 1_000) pdist, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(prob, p, data) - model = bayesianODE(prob, pdist, pkeys, ts, lastt, timeseries, datakeys, noise_prior) + model = bayesianODE(prob, + first(default_algorithm(prob)), + pdist, + pkeys, + ts, + lastt, + timeseries, + datakeys, + noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -451,7 +470,10 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, (pdist_, pkeys) = bayes_unpack_data(p) pdist, grouppriorsfunc = flatten(pdist_) - model = ensemblebayesianODE(probs, t, pdist, grouppriorsfunc, pkeys, data, noise_prior) + model = ensemblebayesianODE(probs, + map(first ∘ default_algorithm, probs), + t, pdist, grouppriorsfunc, pkeys, last.(data), + map(Base.Fix2(IndexKeyMap, data), probs), noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -472,6 +494,7 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(p, data) pdist, grouppriorsfunc = flatten(pdist_) model = ensemblebayesianODE(probs, + map(first ∘ default_algorithm, probs), pdist, grouppriorsfunc, pkeys, diff --git a/src/keyindexmap.jl b/src/keyindexmap.jl index a2a863af..b773bdf9 100644 --- a/src/keyindexmap.jl +++ b/src/keyindexmap.jl @@ -3,6 +3,7 @@ struct IndexKeyMap indices::Vector{Int} end +# probs support function IndexKeyMap(prob, keys) params = ModelingToolkit.parameters(prob.f.sys) indices = Vector{Int}(undef, length(keys)) @@ -12,7 +13,8 @@ function IndexKeyMap(prob, keys) return IndexKeyMap(indices) end -Base.@propagate_inbounds function (ikm::IndexKeyMap)(prob, v::AbstractVector) +Base.@propagate_inbounds function (ikm::IndexKeyMap)(prob::SciMLBase.AbstractDEProblem, + v::AbstractVector) @boundscheck checkbounds(v, length(ikm.indices)) def = prob.p ret = Vector{Base.promote_eltype(v, def)}(undef, length(def)) @@ -22,8 +24,20 @@ Base.@propagate_inbounds function (ikm::IndexKeyMap)(prob, v::AbstractVector) end return ret end - function _remake(prob, tspan, ikm::IndexKeyMap, pprior) p = ikm(prob, pprior) remake(prob; tspan, p) end + +# data support +function IndexKeyMap(prob, data::AbstractVector{<:Pair}) + states = ModelingToolkit.states(prob.f.sys) + indices = Vector{Int}(undef, length(data)) + for i in eachindex(data) + indices[i] = findfirst(Base.Fix1(isequal, data[i].first), states) + end + return IndexKeyMap(indices) +end +function (ikm::IndexKeyMap)(sol::SciMLBase.AbstractTimeseriesSolution) + (@view(sol[i, :]) for i in ikm.indices) +end From 0a92469fe578386d0077bb7dba8e0cdd31d9566c Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 14:31:04 -0400 Subject: [PATCH 06/16] pass prob as first arg --- src/datafit.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index f47e02a0..28af46d0 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -211,14 +211,15 @@ Turing.@model function bayesianODE(prob, alg, t, pdist, pkeys, data, datamap, no pprior ~ product_distribution(pdist) - prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + sol = solve(prob, alg, saveat = t) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing end - for i in eachindex(data) - data[i] ~ MvNormal(datamap(sol), σ^2 * I) + for (i,x) in enumerate(datamap(sol)) + data[i] ~ MvNormal(x, σ^2 * I) end return nothing end @@ -422,7 +423,7 @@ function bayesian_datafit(prob, pkeys, last.(data), IndexKeyMap(prob, data), - noise_prior) + noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -467,7 +468,7 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1000) - (pdist_, pkeys) = bayes_unpack_data(p) + (pdist_, pkeys) = bayes_unpack_data(probs, p) pdist, grouppriorsfunc = flatten(pdist_) model = ensemblebayesianODE(probs, @@ -491,7 +492,7 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1_000) - pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(p, data) + pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(probs, p, data) pdist, grouppriorsfunc = flatten(pdist_) model = ensemblebayesianODE(probs, map(first ∘ default_algorithm, probs), From 923188e3f24c378276374f25e9510796be6c0e0c Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 13 Jul 2023 14:39:22 -0400 Subject: [PATCH 07/16] fix up tests --- test/ensemble.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/test/ensemble.jl b/test/ensemble.jl index 494b196c..d28f84f7 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -83,10 +83,12 @@ sol = solve(enprob; saveat = t_ensem); @test ensemble_weights(sol, data_ensem) ≈ [0.2, 0.5, 0.3] -probs = [prob, prob2, prob3] -ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] -datas = [data_train,data_train,data_train] +probs = (prob, prob2, prob3) +ps = Tuple([β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3) +datas = (data_train,data_train,data_train) enprobs = bayesian_ensemble(probs, ps, datas) sol = solve(enprobs; saveat = t_ensem); -ensemble_weights(sol, data_ensem) \ No newline at end of file +ensemble_weights(sol, data_ensem) + +bayesian_datafit(probs, ps, datas) \ No newline at end of file From da123692f93d39d20ff4879fd0f5d90c5c43f523 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 13 Jul 2023 15:03:08 -0400 Subject: [PATCH 08/16] write the super ensemble tutorial part --- docs/src/tutorials/ensemble_modeling.md | 48 +++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/docs/src/tutorials/ensemble_modeling.md b/docs/src/tutorials/ensemble_modeling.md index ad12d61b..89c4bf4b 100644 --- a/docs/src/tutorials/ensemble_modeling.md +++ b/docs/src/tutorials/ensemble_modeling.md @@ -237,15 +237,57 @@ scatter!(t_forecast, data_forecast[1][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack([ensem_weights[i] * sol[i][I] for i in 1:length(forecast_probs)]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:,I]), dims = 2) +plot(sol; idxs = I, color = :blue) +plot!(t_forecast, ensem_prediction, lw = 3, color = :red) +scatter!(t_forecast, data_forecast[2][2][2]) +``` + +```@example ensemble +ensem_prediction = sum(stack(ensem_weights .* sol[:,R]), dims = 2) +plot(sol; idxs = R, color = :blue) +plot!(t_forecast, ensem_prediction, lw = 3, color = :red) +scatter!(t_forecast, data_forecast[3][2][2]) +``` + +## Training the "Super Ensemble" Model + +The standard ensemble model first calibrates each model in an ensemble and then uses the calibrated models +as the basis for a prediction via a linear combination. The super ensemble performs the Bayesian estimation +on the full combination of models, including the weights vector, as a single Bayesian posterior calculation. +While this has the downside that the prediction of any single model is not necessarily predictive of the +whole, in some cases this ensemble model may be more effective. + +To train this model, simply use `bayesian_datafit` on the ensemble. This looks like: + +```@example ensemble +probs = [prob, prob2, prob3] +ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] +datas = [data_ensem,data_ensem,data_ensem] +super_enprob, ensem_weights = bayesian_datafit(probs, ps, datas) +``` + +And now we can forecast with this model: + +```@example ensemble +sol = solve(super_enprob; saveat = t_forecast); +ensem_prediction = sum(stack(ensem_weights .* sol[:,S]), dims = 2) +plot(sol; idxs = S, color = :blue) +plot!(t_forecast, ensem_prediction, lw = 3, color = :red) +scatter!(t_forecast, data_forecast[1][2][2]) +``` + +```@example ensemble +ensem_prediction = sum(stack(ensem_weights .* sol[:,I]), dims = 2) plot(sol; idxs = I, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[2][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack([ensem_weights[i] * sol[i][R] for i in 1:length(forecast_probs)]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:,R]), dims = 2) plot(sol; idxs = R, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[3][2][2]) -``` \ No newline at end of file +``` + From 273450dc26c90638c87f5a73d165604476c403cb Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 15:35:28 -0400 Subject: [PATCH 09/16] rebase --- docs/src/tutorials/ensemble_modeling.md | 47 +++++++++++++------------ src/datafit.jl | 46 ++++++++++++++++++------ src/ensemble.jl | 17 +++++---- test/datafit.jl | 1 - test/ensemble.jl | 30 ++++++++-------- 5 files changed, 85 insertions(+), 56 deletions(-) diff --git a/docs/src/tutorials/ensemble_modeling.md b/docs/src/tutorials/ensemble_modeling.md index 89c4bf4b..5471a309 100644 --- a/docs/src/tutorials/ensemble_modeling.md +++ b/docs/src/tutorials/ensemble_modeling.md @@ -95,7 +95,7 @@ We can access the 3 solutions as `sol[i]` respectively. Let's get the time serie for `S` from each of the models: ```@example ensemble -sol[:,S] +sol[:, S] ``` ## Building a Dataset @@ -107,9 +107,9 @@ interface on the ensemble solution. ```@example ensemble weights = [0.2, 0.5, 0.3] data = [ - S => vec(sum(stack(weights .* sol[:,S]), dims = 2)), - I => vec(sum(stack(weights .* sol[:,I]), dims = 2)), - R => vec(sum(stack(weights .* sol[:,R]), dims = 2)), + S => vec(sum(stack(weights .* sol[:, S]), dims = 2)), + I => vec(sum(stack(weights .* sol[:, I]), dims = 2)), + R => vec(sum(stack(weights .* sol[:, R]), dims = 2)), ] ``` @@ -131,27 +131,27 @@ scatter!(data[3][2]) Now let's split that into training, ensembling, and forecast sections: ```@example ensemble -fullS = vec(sum(stack(weights .* sol[:,S]),dims=2)) -fullI = vec(sum(stack(weights .* sol[:,I]),dims=2)) -fullR = vec(sum(stack(weights .* sol[:,R]),dims=2)) +fullS = vec(sum(stack(weights .* sol[:, S]), dims = 2)) +fullI = vec(sum(stack(weights .* sol[:, I]), dims = 2)) +fullR = vec(sum(stack(weights .* sol[:, R]), dims = 2)) t_train = 0:14 data_train = [ - S => (t_train,fullS[1:15]), - I => (t_train,fullI[1:15]), - R => (t_train,fullR[1:15]), + S => (t_train, fullS[1:15]), + I => (t_train, fullI[1:15]), + R => (t_train, fullR[1:15]), ] t_ensem = 0:21 data_ensem = [ - S => (t_ensem,fullS[1:22]), - I => (t_ensem,fullI[1:22]), - R => (t_ensem,fullR[1:22]), + S => (t_ensem, fullS[1:22]), + I => (t_ensem, fullI[1:22]), + R => (t_ensem, fullR[1:22]), ] t_forecast = 0:30 data_forecast = [ - S => (t_forecast,fullS), - I => (t_forecast,fullI), - R => (t_forecast,fullR), + S => (t_forecast, fullS), + I => (t_forecast, fullI), + R => (t_forecast, fullR), ] ``` @@ -162,7 +162,7 @@ Now let's perform a Bayesian calibration on each of the models. This gives us mu ```@example ensemble probs = [prob, prob2, prob3] ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] -datas = [data_train,data_train,data_train] +datas = [data_train, data_train, data_train] enprobs = bayesian_ensemble(probs, ps, datas) ``` @@ -192,8 +192,8 @@ Now let's train the ensemble model. We will do that by solving a bit further tha calibration step. Let's build that solution data: ```@example ensemble -plot(sol;idxs = S) -scatter!(t_ensem,data_ensem[1][2][2]) +plot(sol; idxs = S) +scatter!(t_ensem, data_ensem[1][2][2]) ``` We can obtain the optimal weights for ensembling by solving a linear regression of @@ -208,14 +208,14 @@ Now we can extrapolate forward with these ensemble weights as follows: ```@example ensemble sol = solve(enprobs; saveat = t_ensem); -ensem_prediction = sum(stack(ensem_weights .* sol[:,S]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, S]), dims = 2) plot(sol; idxs = S, color = :blue) plot!(t_ensem, ensem_prediction, lw = 5, color = :red) scatter!(t_ensem, data_ensem[1][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack(ensem_weights .* sol[:,I]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, I]), dims = 2) plot(sol; idxs = I, color = :blue) plot!(t_ensem, ensem_prediction, lw = 3, color = :red) scatter!(t_ensem, data_ensem[2][2][2]) @@ -226,11 +226,12 @@ scatter!(t_ensem, data_ensem[2][2][2]) Once we have obtained the ensemble model, we can forecast ahead with it: ```@example ensemble -forecast_probs = [remake(enprobs.prob[i]; tspan = (t_train[1],t_forecast[end])) for i in 1:length(enprobs.prob)] +forecast_probs = [remake(enprobs.prob[i]; tspan = (t_train[1], t_forecast[end])) + for i in 1:length(enprobs.prob)] fit_enprob = EnsembleProblem(forecast_probs) sol = solve(fit_enprob; saveat = t_forecast); -ensem_prediction = sum(stack(ensem_weights .* sol[:,S]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, S]), dims = 2) plot(sol; idxs = S, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[1][2][2]) diff --git a/src/datafit.jl b/src/datafit.jl index 28af46d0..8f77ef72 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -211,14 +211,14 @@ Turing.@model function bayesianODE(prob, alg, t, pdist, pkeys, data, datamap, no pprior ~ product_distribution(pdist) - prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) sol = solve(prob, alg, saveat = t) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) return nothing end - for (i,x) in enumerate(datamap(sol)) + for (i, x) in enumerate(datamap(sol)) data[i] ~ MvNormal(x, σ^2 * I) end return nothing @@ -262,6 +262,10 @@ struct WeightedSol{T, S <: Tuple{Vararg{AbstractVector{T}}}, W} <: AbstractVecto new{T, S, W}(sols, weights) end end +function WeightedSol(sols::S, + weights::W) where {T, S <: Tuple{Vararg{AbstractVector{T}}}, W} + WeightedSol{T}(sols, weights) +end Base.length(ws::WeightedSol) = length(first(ws.sols)) Base.size(ws::WeightedSol) = (length(first(ws.sols)),) function Base.getindex(ws::WeightedSol{T}, i::Int) where {T} @@ -423,7 +427,7 @@ function bayesian_datafit(prob, pkeys, last.(data), IndexKeyMap(prob, data), - noise_prior) + noise_prior) chain = Turing.sample(model, Turing.NUTS(0.65), mcmcensemble, @@ -460,15 +464,37 @@ function bayesian_datafit(prob, [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) for i in eachindex(p)] end +function extract_ensemble(chain, ps) + j = Ref(0) + params = map(ps) do p + [Pair(p_.first, vec(collect(chain["ppriors[" * string(j[] += 1) * "]"]))) + for p_ in p] + end + wacc = vec(collect(chain["weights[1]"])) + weights = [copy(wacc)] + for i in 2:(length(ps) - 1) + w = vec(collect(chain["weights[" * string(i) * "]"])) + wacc .+= w + push!(weights, w) + end + @. wacc = 1 - wacc + push!(weights, wacc) + if ps isa Tuple + params, ntuple(Base.Fix1(getindex, weights), Val(length(ps))) + else + params, weights + end +end + function bayesian_datafit(probs::Union{Tuple, AbstractVector}, - p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, + ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, t, data; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1000) - (pdist_, pkeys) = bayes_unpack_data(probs, p) + (pdist_, pkeys) = bayes_unpack_data(probs, ps) pdist, grouppriorsfunc = flatten(pdist_) model = ensemblebayesianODE(probs, @@ -481,18 +507,17 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, niter, nchains; progress = true) - [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) - for i in eachindex(p)] + extract_ensemble(chain, ps) end function bayesian_datafit(probs::Union{Tuple, AbstractVector}, - p::Tuple{Vararg{<:AbstractVector{<:Pair}}}, + ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, niter = 1_000) - pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(probs, p, data) + pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(probs, ps, data) pdist, grouppriorsfunc = flatten(pdist_) model = ensemblebayesianODE(probs, map(first ∘ default_algorithm, probs), @@ -510,8 +535,7 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, niter, nchains; progress = true) - [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) - for i in eachindex(p)] + extract_ensemble(chain, ps) end """ diff --git a/src/ensemble.jl b/src/ensemble.jl index 41c07a0c..e6f64f40 100644 --- a/src/ensemble.jl +++ b/src/ensemble.jl @@ -19,9 +19,12 @@ dataset on which the ensembler should be trained on. """ function ensemble_weights(sol::EnsembleSolution, data_ensem) obs = first.(data_ensem) - predictions = reduce(vcat, reduce(hcat,[sol[i][s] for i in 1:length(sol)]) for s in obs) - data = reduce(vcat, [data_ensem[i][2] isa Tuple ? data_ensem[i][2][2] : data_ensem[i][2] for i in 1:length(data_ensem)]) - weights = predictions \ data + predictions = reduce(vcat, + reduce(hcat, [sol[i][s] for i in 1:length(sol)]) for s in obs) + data = reduce(vcat, + [data_ensem[i][2] isa Tuple ? data_ensem[i][2][2] : data_ensem[i][2] + for i in 1:length(data_ensem)]) + weights = predictions \ data end function bayesian_ensemble(probs, ps, datas; @@ -30,20 +33,20 @@ function bayesian_ensemble(probs, ps, datas; nchains = 4, niter = 1_000, keep = 100) - fits = map(probs, ps, datas) do prob, p, data bayesian_datafit(prob, p, data; noise_prior, mcmcensemble, nchains, niter) end models = map(probs, fits) do prob, fit - [remake(prob, p = Pair.(first.(fit), getindex.(last.(fit), i))) for i in length(fit[1][2])-keep:length(fit[1][2])] + [remake(prob, p = Pair.(first.(fit), getindex.(last.(fit), i))) + for i in (length(fit[1][2]) - keep):length(fit[1][2])] end @info "Calibrations are complete" - all_probs = reduce(vcat,models) + all_probs = reduce(vcat, models) @info "$(length(all_probs)) total models" enprob = EnsembleProblem(all_probs) -end \ No newline at end of file +end diff --git a/test/datafit.jl b/test/datafit.jl index cba922d5..41cbbec0 100644 --- a/test/datafit.jl +++ b/test/datafit.jl @@ -99,4 +99,3 @@ data_with_t = [x => (tsave1, sol_data1[x]), z => (tsave2, sol_data2[z])] p_posterior = @time bayesian_datafit(prob, p_prior, data_with_t) @test var.(getfield.(p_prior, :second)) >= var.(getfield.(p_posterior, :second)) - diff --git a/test/ensemble.jl b/test/ensemble.jl index d28f84f7..2f01b094 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -56,27 +56,27 @@ sol = solve(enprob; saveat = 1); weights = [0.2, 0.5, 0.3] -fullS = vec(sum(stack(weights .* sol[:,S]),dims=2)) -fullI = vec(sum(stack(weights .* sol[:,I]),dims=2)) -fullR = vec(sum(stack(weights .* sol[:,R]),dims=2)) +fullS = vec(sum(stack(weights .* sol[:, S]), dims = 2)) +fullI = vec(sum(stack(weights .* sol[:, I]), dims = 2)) +fullR = vec(sum(stack(weights .* sol[:, R]), dims = 2)) t_train = 0:14 data_train = [ - S => (t_train,fullS[1:15]), - I => (t_train,fullI[1:15]), - R => (t_train,fullR[1:15]), + S => (t_train, fullS[1:15]), + I => (t_train, fullI[1:15]), + R => (t_train, fullR[1:15]), ] t_ensem = 0:21 data_ensem = [ - S => (t_ensem,fullS[1:22]), - I => (t_ensem,fullI[1:22]), - R => (t_ensem,fullR[1:22]), + S => (t_ensem, fullS[1:22]), + I => (t_ensem, fullI[1:22]), + R => (t_ensem, fullR[1:22]), ] t_forecast = 0:30 data_forecast = [ - S => (t_forecast,fullS), - I => (t_forecast,fullI), - R => (t_forecast,fullR), + S => (t_forecast, fullS), + I => (t_forecast, fullI), + R => (t_forecast, fullR), ] sol = solve(enprob; saveat = t_ensem); @@ -85,10 +85,12 @@ sol = solve(enprob; saveat = t_ensem); probs = (prob, prob2, prob3) ps = Tuple([β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3) -datas = (data_train,data_train,data_train) +datas = (data_train, data_train, data_train) enprobs = bayesian_ensemble(probs, ps, datas) sol = solve(enprobs; saveat = t_ensem); ensemble_weights(sol, data_ensem) -bayesian_datafit(probs, ps, datas) \ No newline at end of file +bayesian_datafit(probs, ps, datas) +params, weights = bayesian_datafit(probs, ps, datas) +@test length(params) == length(weights) == length(ps) From c75b488c1f44cf44ce6c3a88791fc12f9ce71f38 Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Thu, 13 Jul 2023 15:49:10 -0400 Subject: [PATCH 10/16] Test update --- docs/src/tutorials/ensemble_modeling.md | 13 ++++++------- src/datafit.jl | 4 ++-- test/ensemble.jl | 8 ++++++-- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/docs/src/tutorials/ensemble_modeling.md b/docs/src/tutorials/ensemble_modeling.md index 5471a309..3fda2cc1 100644 --- a/docs/src/tutorials/ensemble_modeling.md +++ b/docs/src/tutorials/ensemble_modeling.md @@ -238,14 +238,14 @@ scatter!(t_forecast, data_forecast[1][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack(ensem_weights .* sol[:,I]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, I]), dims = 2) plot(sol; idxs = I, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[2][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack(ensem_weights .* sol[:,R]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, R]), dims = 2) plot(sol; idxs = R, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[3][2][2]) @@ -264,7 +264,7 @@ To train this model, simply use `bayesian_datafit` on the ensemble. This looks l ```@example ensemble probs = [prob, prob2, prob3] ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] -datas = [data_ensem,data_ensem,data_ensem] +datas = [data_ensem, data_ensem, data_ensem] super_enprob, ensem_weights = bayesian_datafit(probs, ps, datas) ``` @@ -272,23 +272,22 @@ And now we can forecast with this model: ```@example ensemble sol = solve(super_enprob; saveat = t_forecast); -ensem_prediction = sum(stack(ensem_weights .* sol[:,S]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, S]), dims = 2) plot(sol; idxs = S, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[1][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack(ensem_weights .* sol[:,I]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, I]), dims = 2) plot(sol; idxs = I, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[2][2][2]) ``` ```@example ensemble -ensem_prediction = sum(stack(ensem_weights .* sol[:,R]), dims = 2) +ensem_prediction = sum(stack(ensem_weights .* sol[:, R]), dims = 2) plot(sol; idxs = R, color = :blue) plot!(t_forecast, ensem_prediction, lw = 3, color = :red) scatter!(t_forecast, data_forecast[3][2][2]) ``` - diff --git a/src/datafit.jl b/src/datafit.jl index 8f77ef72..4ab7efbb 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -489,7 +489,7 @@ end function bayesian_datafit(probs::Union{Tuple, AbstractVector}, ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, t, - data; + data::AbstractVector{<:Pair}; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, @@ -512,7 +512,7 @@ end function bayesian_datafit(probs::Union{Tuple, AbstractVector}, ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, - data; + data::AbstractVector{<:Pair}; noise_prior = InverseGamma(2, 3), mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), nchains = 4, diff --git a/test/ensemble.jl b/test/ensemble.jl index 2f01b094..d294bac6 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -91,6 +91,10 @@ enprobs = bayesian_ensemble(probs, ps, datas) sol = solve(enprobs; saveat = t_ensem); ensemble_weights(sol, data_ensem) -bayesian_datafit(probs, ps, datas) -params, weights = bayesian_datafit(probs, ps, datas) +# only supports one datas +params, weights = bayesian_datafit(probs, ps, data_train) + @test length(params) == length(weights) == length(ps) +for i in eachindex(weights[1]) + @test sum(weights[j][i] for j in eachindex(weights)) ≈ 1.0 +end From 156ad61510beb99a2e81401df1dc3ce4d44571d2 Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Fri, 14 Jul 2023 16:14:32 -0400 Subject: [PATCH 11/16] Bayesian non-Ensemble fitting returns `Ensemble`s Bayesian Ensemble fitting does not return `Ensemble`s --- docs/src/tutorials/ensemble_modeling.md | 12 +- src/datafit.jl | 139 +++++++++++++++--------- src/ensemble.jl | 24 ++-- src/keyindexmap.jl | 2 +- test/datafit.jl | 10 +- test/ensemble.jl | 2 +- 6 files changed, 120 insertions(+), 69 deletions(-) diff --git a/docs/src/tutorials/ensemble_modeling.md b/docs/src/tutorials/ensemble_modeling.md index 3fda2cc1..06f37826 100644 --- a/docs/src/tutorials/ensemble_modeling.md +++ b/docs/src/tutorials/ensemble_modeling.md @@ -80,7 +80,7 @@ prototype problem, which we are effectively ignoring for our use case. Thus a simple `EnsembleProblem` which ensembles the three models built above is as follows: ```@example ensemble -probs = [prob, prob2, prob3] +probs = [prob, prob2, prob3]; enprob = EnsembleProblem(probs) ``` @@ -160,7 +160,7 @@ data_forecast = [ Now let's perform a Bayesian calibration on each of the models. This gives us multiple parameterizations for each model, which then gives an ensemble which is `parameterizations x models` in size. ```@example ensemble -probs = [prob, prob2, prob3] +probs = [prob, prob2, prob3]; ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] datas = [data_train, data_train, data_train] enprobs = bayesian_ensemble(probs, ps, datas) @@ -227,7 +227,7 @@ Once we have obtained the ensemble model, we can forecast ahead with it: ```@example ensemble forecast_probs = [remake(enprobs.prob[i]; tspan = (t_train[1], t_forecast[end])) - for i in 1:length(enprobs.prob)] + for i in 1:length(enprobs.prob)]; fit_enprob = EnsembleProblem(forecast_probs) sol = solve(fit_enprob; saveat = t_forecast); @@ -262,10 +262,10 @@ whole, in some cases this ensemble model may be more effective. To train this model, simply use `bayesian_datafit` on the ensemble. This looks like: ```@example ensemble -probs = [prob, prob2, prob3] +probs = [prob, prob2, prob3]; ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] -datas = [data_ensem, data_ensem, data_ensem] -super_enprob, ensem_weights = bayesian_datafit(probs, ps, datas) + +super_enprob, ensem_weights = bayesian_datafit(probs, ps, data_ensem) ``` And now we can forecast with this model: diff --git a/src/datafit.jl b/src/datafit.jl index 4ab7efbb..3c1bb707 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -1,3 +1,4 @@ + function l2loss(pvals, (prob, pkeys, t, data)::Tuple{Vararg{Any, 4}}) p = Pair.(pkeys, pvals) prob = remake(prob, tspan = (prob.tspan[1], t[end]), p = p) @@ -211,7 +212,7 @@ Turing.@model function bayesianODE(prob, alg, t, pdist, pkeys, data, datamap, no pprior ~ product_distribution(pdist) - prob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + prob = _remake(prob, pkeys, pprior, (prob.tspan[1], t[end])) sol = solve(prob, alg, saveat = t) if !SciMLBase.successful_retcode(sol) @@ -236,7 +237,7 @@ Turing.@model function bayesianODE(prob, alg, pprior ~ product_distribution(pdist) - prob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) + prob = _remake(prob, pkeys, pprior, (prob.tspan[1], lastt)) sol = solve(prob, alg) if !SciMLBase.successful_retcode(sol) Turing.DynamicPPL.acclogp!!(__varinfo__, -Inf) @@ -313,18 +314,18 @@ end function getsols(probs, algs, probspkeys, ppriors, t::AbstractArray) map(probs, algs, probspkeys, ppriors) do prob, alg, pkeys, pprior - newprob = _remake(prob, (prob.tspan[1], t[end]), pkeys, pprior) + newprob = _remake(prob, pkeys, pprior, (prob.tspan[1], t[end])) solve(newprob, alg, saveat = t) end end function getsols(probs, algs, probspkeys, ppriors, lastt::Number) map(probs, algs, probspkeys, ppriors) do prob, alg, pkeys, pprior - newprob = _remake(prob, (prob.tspan[1], lastt), pkeys, pprior) + newprob = _remake(prob, pkeys, pprior, (prob.tspan[1], lastt)) solve(newprob, alg) end end -Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, +Turing.@model function ensemblebayesianODE(probs::Tuple, algs, t, pdist, @@ -349,7 +350,7 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, end return nothing end -Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, +Turing.@model function ensemblebayesianODE(probs::Tuple, algs, pdist, grouppriorsfunc, @@ -380,6 +381,18 @@ Turing.@model function ensemblebayesianODE(probs::Union{Tuple, AbstractVector}, return nothing end +# this is bad, probably do not use this +function naiverep(f, N, ::EnsembleThreads) + t = Vector{Task}(undef, N) + for n in 1:N + t[n] = Threads.@spawn f() + end + return identity.(map(fetch, t)) +end +function naiverep(f, N, ::EnsembleSerial) + map(_ -> f(), 1:N) +end + """ bayesian_datafit(prob, p, t, data) bayesian_datafit(prob, p, data) @@ -416,7 +429,7 @@ function bayesian_datafit(prob, t, data; noise_prior = InverseGamma(2, 3), - mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + ensemblealg::SciMLBase.BasicEnsembleAlgorithm = EnsembleThreads(), nchains = 4, niter = 1000) (pdist, pkeys) = bayes_unpack_data(prob, p) @@ -428,21 +441,22 @@ function bayesian_datafit(prob, last.(data), IndexKeyMap(prob, data), noise_prior) - chain = Turing.sample(model, - Turing.NUTS(0.65), - mcmcensemble, - niter, - nchains; - progress = true) - [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) - for i in eachindex(p)] + chains = naiverep(nchains, ensemblealg) do + Turing.sample(model, + Turing.NUTS(0.65), + Turing.MCMCSerial(), + niter, + 1; + progress = false) + end + extract_ensemble(chains, prob, length(p), pkeys) end function bayesian_datafit(prob, p, data; noise_prior = InverseGamma(2, 3), - mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + ensemblealg::SciMLBase.BasicEnsembleAlgorithm = EnsembleThreads(), nchains = 4, niter = 1_000) pdist, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(prob, p, data) @@ -455,20 +469,35 @@ function bayesian_datafit(prob, timeseries, datakeys, noise_prior) - chain = Turing.sample(model, - Turing.NUTS(0.65), - mcmcensemble, - niter, - nchains; - progress = true) - [Pair(p[i].first, collect(chain["pprior[" * string(i) * "]"])[:]) - for i in eachindex(p)] -end -function extract_ensemble(chain, ps) + chains = naiverep(nchains, ensemblealg) do + Turing.sample(model, + Turing.NUTS(0.65), + Turing.MCMCSerial(), + niter, + 1; + progress = false) + end + extract_ensemble(chains, prob, length(p), pkeys) +end +function extract_ensemble(chns, prob::SciMLBase.AbstractDEProblem, Np::Int, ikm) + # find range of `chain` corresponding to `pprior` + j = findfirst(==(Symbol("pprior[1]")), chns[1].name_map.parameters) + probs = Vector{typeof(prob)}(undef, (length(chns) * size(chns[1].value, 1))::Int) + i = 0 + for chn in chns + params = @view chn.value[:, j:(j + Np - 1), 1] + for ps in eachrow(params) + probs[i += 1] = _remake(prob, ikm, ps) + end + end + return EnsembleProblem(probs) +end +function extract_ensemble(chain::Turing.Chains, probs::Tuple, ps::Tuple, ikm) j = Ref(0) - params = map(ps) do p - [Pair(p_.first, vec(collect(chain["ppriors[" * string(j[] += 1) * "]"]))) - for p_ in p] + probs = map(probs, ps) do prob, p + newp = [Pair(p_.first, vec(collect(chain["ppriors[" * string(j[] += 1) * "]"]))) + for p_ in p] + _remake(prob, ikm, newp) end wacc = vec(collect(chain["weights[1]"])) weights = [copy(wacc)] @@ -480,18 +509,24 @@ function extract_ensemble(chain, ps) @. wacc = 1 - wacc push!(weights, wacc) if ps isa Tuple - params, ntuple(Base.Fix1(getindex, weights), Val(length(ps))) + probs, ntuple(Base.Fix1(getindex, weights), Val(length(ps))) else - params, weights + probs, weights end end -function bayesian_datafit(probs::Union{Tuple, AbstractVector}, +function bayesian_datafit(probs::Vector, ps::Vector, args...; kwargs...) + probst = tuple(probs...) + pst = tuple(ps...) + bayesian_datafit(probst, pst, args...; kwargs...) +end + +function bayesian_datafit(probs::Tuple, ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, t, data::AbstractVector{<:Pair}; noise_prior = InverseGamma(2, 3), - mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + ensemblealg::SciMLBase.BasicEnsembleAlgorithm = EnsembleThreads(), nchains = 4, niter = 1000) (pdist_, pkeys) = bayes_unpack_data(probs, ps) @@ -501,20 +536,23 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, map(first ∘ default_algorithm, probs), t, pdist, grouppriorsfunc, pkeys, last.(data), map(Base.Fix2(IndexKeyMap, data), probs), noise_prior) - chain = Turing.sample(model, - Turing.NUTS(0.65), - mcmcensemble, - niter, - nchains; - progress = true) - extract_ensemble(chain, ps) + chains = naiverep(nchains, ensemblealg) do + Turing.sample(model, + Turing.NUTS(0.65), + Turing.MCMCSerial(), + niter, + 1; + progress = false) + end + return chains + extract_ensemble(chains, prob, length(ps), pkeys) end -function bayesian_datafit(probs::Union{Tuple, AbstractVector}, +function bayesian_datafit(probs::Tuple, ps::Tuple{Vararg{<:AbstractVector{<:Pair}}}, data::AbstractVector{<:Pair}; noise_prior = InverseGamma(2, 3), - mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCThreads(), + ensemblealg::SciMLBase.BasicEnsembleAlgorithm = EnsembleThreads(), nchains = 4, niter = 1_000) pdist_, pkeys, ts, lastt, timeseries, datakeys = bayes_unpack_data(probs, ps, data) @@ -529,13 +567,16 @@ function bayesian_datafit(probs::Union{Tuple, AbstractVector}, timeseries, datakeys, noise_prior) - chain = Turing.sample(model, - Turing.NUTS(0.65), - mcmcensemble, - niter, - nchains; - progress = true) - extract_ensemble(chain, ps) + chains = naiverep(nchains, ensemblealg) do + Turing.sample(model, + Turing.NUTS(0.65), + Turing.MCMCSerial(), + niter, + 1; + progress = false) + end + return chains + extract_ensemble(chains, prob, length(ps), pkeys) end """ diff --git a/src/ensemble.jl b/src/ensemble.jl index e6f64f40..9c59418e 100644 --- a/src/ensemble.jl +++ b/src/ensemble.jl @@ -1,3 +1,16 @@ +function naivemap(f, ::EnsembleThreads, arg0, args...) + t = Vector{Task}(undef, length(arg0)) + for (n, a) in enumerate(arg0) + t[n] = let an = map(Base.Fix2(Base.getindex, n), args) + Threads.@spawn f(a, an...) + end + end + return identity.(map(fetch, t)) +end +function naivemap(f, ::EnsembleSerial, args...) + map(f, args...) +end + """ ensemble_weights(sol::EnsembleSolution, data_ensem) @@ -29,17 +42,12 @@ end function bayesian_ensemble(probs, ps, datas; noise_prior = InverseGamma(2, 3), - mcmcensemble::AbstractMCMC.AbstractMCMCEnsemble = Turing.MCMCSerial(), + ensemblealg::SciMLBase.BasicEnsembleAlgorithm = EnsembleThreads(), nchains = 4, niter = 1_000, keep = 100) - fits = map(probs, ps, datas) do prob, p, data - bayesian_datafit(prob, p, data; noise_prior, mcmcensemble, nchains, niter) - end - - models = map(probs, fits) do prob, fit - [remake(prob, p = Pair.(first.(fit), getindex.(last.(fit), i))) - for i in (length(fit[1][2]) - keep):length(fit[1][2])] + models = naivemap(ensemblealg, probs, ps, datas) do prob, p, data + bayesian_datafit(prob, p, data; noise_prior, ensemblealg, nchains, niter) end @info "Calibrations are complete" diff --git a/src/keyindexmap.jl b/src/keyindexmap.jl index b773bdf9..c60e52b1 100644 --- a/src/keyindexmap.jl +++ b/src/keyindexmap.jl @@ -24,7 +24,7 @@ Base.@propagate_inbounds function (ikm::IndexKeyMap)(prob::SciMLBase.AbstractDEP end return ret end -function _remake(prob, tspan, ikm::IndexKeyMap, pprior) +function _remake(prob, ikm::IndexKeyMap, pprior, tspan::Tuple{Number, Number} = prob.tspan) p = ikm(prob, pprior) remake(prob; tspan, p) end diff --git a/test/datafit.jl b/test/datafit.jl index 41cbbec0..64fe92cd 100644 --- a/test/datafit.jl +++ b/test/datafit.jl @@ -88,8 +88,9 @@ tsave = collect(10.0:10.0:100.0) sol_data = solve(prob, saveat = tsave) data = [x => sol_data[x], z => sol_data[z]] p_prior = [σ => Normal(26.8, 0.1), β => Normal(2.7, 0.1)] -p_posterior = @time bayesian_datafit(prob, p_prior, tsave, data) -@test var.(getfield.(p_prior, :second)) >= var.(getfield.(p_posterior, :second)) +p_posterior = @time bayesian_datafit(prob, p_prior, tsave, data, nchains = 2, niter = 100) +solve(p_posterior) +# @test var.(getfield.(p_prior, :second)) >= var.(getfield.(p_posterior, :second)) tsave1 = collect(10.0:10.0:100.0) sol_data1 = solve(prob, saveat = tsave1) @@ -97,5 +98,6 @@ tsave2 = collect(10.0:13.5:100.0) sol_data2 = solve(prob, saveat = tsave2) data_with_t = [x => (tsave1, sol_data1[x]), z => (tsave2, sol_data2[z])] -p_posterior = @time bayesian_datafit(prob, p_prior, data_with_t) -@test var.(getfield.(p_prior, :second)) >= var.(getfield.(p_posterior, :second)) +p_posterior = @time bayesian_datafit(prob, p_prior, data_with_t, nchains = 2, niter = 100) +solve(p_posterior) +# @test var.(getfield.(p_prior, :second)) >= var.(getfield.(p_posterior, :second)) diff --git a/test/ensemble.jl b/test/ensemble.jl index d294bac6..6ecb8599 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -86,7 +86,7 @@ sol = solve(enprob; saveat = t_ensem); probs = (prob, prob2, prob3) ps = Tuple([β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3) datas = (data_train, data_train, data_train) -enprobs = bayesian_ensemble(probs, ps, datas) +enprobs = bayesian_ensemble(probs, ps, datas, nchains = 2, niter = 200) sol = solve(enprobs; saveat = t_ensem); ensemble_weights(sol, data_ensem) From 8473f659cd50a88470fe8104b09393287571bc28 Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Fri, 14 Jul 2023 16:50:47 -0400 Subject: [PATCH 12/16] Fix `ensemble_weights` --- src/ensemble.jl | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/ensemble.jl b/src/ensemble.jl index 9c59418e..c590a91f 100644 --- a/src/ensemble.jl +++ b/src/ensemble.jl @@ -11,6 +11,16 @@ function naivemap(f, ::EnsembleSerial, args...) map(f, args...) end +getsolmean(sol, s) = sol[s] +function getsolmean(sol::EnsembleSolution, s) + acc = sol[1][s] + N = length(sol) + for i = 2:N + acc .+= sol[i][s] + end + acc ./= N +end + """ ensemble_weights(sol::EnsembleSolution, data_ensem) @@ -33,7 +43,7 @@ dataset on which the ensembler should be trained on. function ensemble_weights(sol::EnsembleSolution, data_ensem) obs = first.(data_ensem) predictions = reduce(vcat, - reduce(hcat, [sol[i][s] for i in 1:length(sol)]) for s in obs) + reduce(hcat, [getsolmean(sol[i],s) for i in 1:length(sol)]) for s in obs) data = reduce(vcat, [data_ensem[i][2] isa Tuple ? data_ensem[i][2][2] : data_ensem[i][2] for i in 1:length(data_ensem)]) From f559b177f350c7909be1fc2b244cf5b56db17b2f Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Fri, 14 Jul 2023 17:00:58 -0400 Subject: [PATCH 13/16] oops --- src/ensemble.jl | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/ensemble.jl b/src/ensemble.jl index c590a91f..632721fe 100644 --- a/src/ensemble.jl +++ b/src/ensemble.jl @@ -11,15 +11,6 @@ function naivemap(f, ::EnsembleSerial, args...) map(f, args...) end -getsolmean(sol, s) = sol[s] -function getsolmean(sol::EnsembleSolution, s) - acc = sol[1][s] - N = length(sol) - for i = 2:N - acc .+= sol[i][s] - end - acc ./= N -end """ ensemble_weights(sol::EnsembleSolution, data_ensem) @@ -43,7 +34,7 @@ dataset on which the ensembler should be trained on. function ensemble_weights(sol::EnsembleSolution, data_ensem) obs = first.(data_ensem) predictions = reduce(vcat, - reduce(hcat, [getsolmean(sol[i],s) for i in 1:length(sol)]) for s in obs) + reduce(hcat, [sol[i][s] for i in 1:length(sol)]) for s in obs) data = reduce(vcat, [data_ensem[i][2] isa Tuple ? data_ensem[i][2][2] : data_ensem[i][2] for i in 1:length(data_ensem)]) @@ -62,7 +53,7 @@ function bayesian_ensemble(probs, ps, datas; @info "Calibrations are complete" - all_probs = reduce(vcat, models) + all_probs = reduce(vcat, map(prob -> prob.prob, models)) @info "$(length(all_probs)) total models" From 34a8d0c64dc59ab6d1b1200337ad4da0b1743ab0 Mon Sep 17 00:00:00 2001 From: Chris Elrod Date: Fri, 14 Jul 2023 17:25:48 -0400 Subject: [PATCH 14/16] minor progress on extract_ensemble --- src/datafit.jl | 27 ++++++++++++++++++++++++--- test/ensemble.jl | 2 +- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index 3c1bb707..180ce535 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -492,7 +492,28 @@ function extract_ensemble(chns, prob::SciMLBase.AbstractDEProblem, Np::Int, ikm) end return EnsembleProblem(probs) end -function extract_ensemble(chain::Turing.Chains, probs::Tuple, ps::Tuple, ikm) +function extract_ensemble(chns, probst::Tuple, Nps::NTuple{P,Int}, ikm) where {P} + + j = findfirst(==(Symbol("pprior[1]")), chns[1].name_map.parameters) + nsamples::Int = (length(chns) * size(chns[1].value, 1)) + probs = map(probst) do prob + Vector{typeof(prob)}(undef, nsamples) + end + weights = ntuple(p -> Vector{Vector{Float64}}(undef, nsamples), Val(P)) + offsets = cumsum((0, Base.front(Nps)...)) + Np = offsets[end] + Nps[end] + inds = ntuple(identity,Val(P)) + i = 0 + for chn in chns + params = @view chn.value[:, j:(j + Np - 1), 1] + + for ps in eachrow(params) + probs[i += 1] = _remake(prob, ikm, ps) + end + end + return map(SciMLBase.WeightedEnsembleProblem,probs,weights) + + j = Ref(0) probs = map(probs, ps) do prob, p newp = [Pair(p_.first, vec(collect(chain["ppriors[" * string(j[] += 1) * "]"]))) @@ -545,7 +566,7 @@ function bayesian_datafit(probs::Tuple, progress = false) end return chains - extract_ensemble(chains, prob, length(ps), pkeys) + extract_ensemble(chains, prob, map(length,ps), pkeys) end function bayesian_datafit(probs::Tuple, @@ -576,7 +597,7 @@ function bayesian_datafit(probs::Tuple, progress = false) end return chains - extract_ensemble(chains, prob, length(ps), pkeys) + extract_ensemble(chains, prob, map(length,ps), pkeys) end """ diff --git a/test/ensemble.jl b/test/ensemble.jl index 6ecb8599..deb2a055 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -92,7 +92,7 @@ sol = solve(enprobs; saveat = t_ensem); ensemble_weights(sol, data_ensem) # only supports one datas -params, weights = bayesian_datafit(probs, ps, data_train) +params, weights = bayesian_datafit(probs, ps, data_train, nchains = 2, niter = 200) @test length(params) == length(weights) == length(ps) for i in eachindex(weights[1]) From a12a741310a43217eca87bf5426f22a94f00cf73 Mon Sep 17 00:00:00 2001 From: chriselrod Date: Tue, 18 Jul 2023 16:39:00 -0400 Subject: [PATCH 15/16] Update ensemble tests --- src/datafit.jl | 72 ++++++++++++++++++++---------------------------- test/ensemble.jl | 15 ++++++---- 2 files changed, 40 insertions(+), 47 deletions(-) diff --git a/src/datafit.jl b/src/datafit.jl index 180ce535..157043d1 100644 --- a/src/datafit.jl +++ b/src/datafit.jl @@ -492,48 +492,38 @@ function extract_ensemble(chns, prob::SciMLBase.AbstractDEProblem, Np::Int, ikm) end return EnsembleProblem(probs) end -function extract_ensemble(chns, probst::Tuple, Nps::NTuple{P,Int}, ikm) where {P} - - j = findfirst(==(Symbol("pprior[1]")), chns[1].name_map.parameters) - nsamples::Int = (length(chns) * size(chns[1].value, 1)) - probs = map(probst) do prob - Vector{typeof(prob)}(undef, nsamples) - end +function extract_ensemble(chns, probst::Tuple, Nps::NTuple{P, Int}, ikms) where {P} + j = findfirst(==(Symbol("ppriors[1]")), chns[1].name_map.parameters) + w = findfirst(==(Symbol("weights[1]")), chns[1].name_map.parameters) + samples_per_chain::Int = size(chns[1].value, 1) + nsamples::Int = (length(chns) * samples_per_chain) weights = ntuple(p -> Vector{Vector{Float64}}(undef, nsamples), Val(P)) - offsets = cumsum((0, Base.front(Nps)...)) + + offsets = cumsum((0, Nps...)) Np = offsets[end] + Nps[end] - inds = ntuple(identity,Val(P)) - i = 0 - for chn in chns - params = @view chn.value[:, j:(j + Np - 1), 1] - - for ps in eachrow(params) - probs[i += 1] = _remake(prob, ikm, ps) + inds = ntuple(identity, Val(P)) + res = map((Iterators.product(chns, 1:samples_per_chain))) do (chn, k) + i = 0 + # anon func takes in chain and sample, returns the associated WeightedEnsembleProb + _probs = map(probst, + ikms, + Base.front(offsets), + Base.tail(offsets)) do prob, ikm, start, stop + ps = chn.value[k, (j + start):(j + stop - 1), 1] + _remake(prob, ikm, ps) end + lastwt = 0.0 # accumulate from 0 for less rounding error + weights = Vector{Float64}(undef, P) + for l in 0:(P - 2) + wt = chn.value[k, w + l, 1] + lastwt += wt + weights[l + 1] = wt + end + weights[P] = 1.0 - lastwt + return SciMLBase.WeightedEnsembleProblem([_probs...]; + weights) end - return map(SciMLBase.WeightedEnsembleProblem,probs,weights) - - - j = Ref(0) - probs = map(probs, ps) do prob, p - newp = [Pair(p_.first, vec(collect(chain["ppriors[" * string(j[] += 1) * "]"]))) - for p_ in p] - _remake(prob, ikm, newp) - end - wacc = vec(collect(chain["weights[1]"])) - weights = [copy(wacc)] - for i in 2:(length(ps) - 1) - w = vec(collect(chain["weights[" * string(i) * "]"])) - wacc .+= w - push!(weights, w) - end - @. wacc = 1 - wacc - push!(weights, wacc) - if ps isa Tuple - probs, ntuple(Base.Fix1(getindex, weights), Val(length(ps))) - else - probs, weights - end + SciMLBase.EnsembleProblem(vec(res)) end function bayesian_datafit(probs::Vector, ps::Vector, args...; kwargs...) @@ -565,8 +555,7 @@ function bayesian_datafit(probs::Tuple, 1; progress = false) end - return chains - extract_ensemble(chains, prob, map(length,ps), pkeys) + extract_ensemble(chains, probs, map(length, ps), pkeys) end function bayesian_datafit(probs::Tuple, @@ -596,8 +585,7 @@ function bayesian_datafit(probs::Tuple, 1; progress = false) end - return chains - extract_ensemble(chains, prob, map(length,ps), pkeys) + extract_ensemble(chains, probs, map(length, ps), pkeys) end """ diff --git a/test/ensemble.jl b/test/ensemble.jl index deb2a055..ab71f7f3 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -92,9 +92,14 @@ sol = solve(enprobs; saveat = t_ensem); ensemble_weights(sol, data_ensem) # only supports one datas -params, weights = bayesian_datafit(probs, ps, data_train, nchains = 2, niter = 200) - -@test length(params) == length(weights) == length(ps) -for i in eachindex(weights[1]) - @test sum(weights[j][i] for j in eachindex(weights)) ≈ 1.0 +ensembleofweightedensembles = bayesian_datafit(probs, + ps, + data_train, + nchains = 2, + niter = 200) + +@test length(ensembleofweightedensembles.prob[1].prob) == + length(ensembleofweightedensembles.prob[1].weights) == length(ps) +for prob in ensembleofweightedensembles.prob + @test sum(prob.weights) ≈ 1.0 end From d7f6e715757a6bc009528ee50d62d84f28a62b31 Mon Sep 17 00:00:00 2001 From: chriselrod Date: Tue, 18 Jul 2023 18:12:15 -0400 Subject: [PATCH 16/16] Bump versions --- Project.toml | 4 +++- docs/src/tutorials/ensemble_modeling.md | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index eae1be4c..9bff2ce6 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ authors = ["SciML"] version = "0.1.10" [deps] +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" GlobalSensitivity = "af5da776-676b-467e-8baf-acd8249e4f0f" @@ -22,6 +23,7 @@ SciMLExpectations = "afe9f18d-7609-4d0e-b7b7-af0cb72b8ea8" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [compat] +DiffEqBase = "6.127.0" DifferentialEquations = "7" Distributions = "0.25" GlobalSensitivity = "2" @@ -33,7 +35,7 @@ OptimizationMOI = "0.1" OptimizationNLopt = "0.1" Plots = "1" Reexport = "1" -SciMLBase = "1.93.1" +SciMLBase = "1.93.3" SciMLExpectations = "2" Turing = "0.22, 0.23, 0.24" julia = "1.6" diff --git a/docs/src/tutorials/ensemble_modeling.md b/docs/src/tutorials/ensemble_modeling.md index 06f37826..c30fbf2c 100644 --- a/docs/src/tutorials/ensemble_modeling.md +++ b/docs/src/tutorials/ensemble_modeling.md @@ -163,7 +163,7 @@ Now let's perform a Bayesian calibration on each of the models. This gives us mu probs = [prob, prob2, prob3]; ps = [[β => Uniform(0.01, 10.0), γ => Uniform(0.01, 10.0)] for i in 1:3] datas = [data_train, data_train, data_train] -enprobs = bayesian_ensemble(probs, ps, datas) +enprobs = bayesian_ensemble(probs, ps, datas, nchains=2, niter=200) ``` Let's see how each of our models in the ensemble compare against the data when changed