Merge pull request #148 from Vaibhavdixit02/master

Make sample_u0 and save_idxs be compatible

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqBayes"
 uuid = "ebbdde9d-f333-5424-9be2-dbf1e9acfb5e"
 authors = ["Vaibhavdixit02 <vaibhavyashdixit@gmail.com>"]
-version = "2.9.1"
+version = "2.10.0"
 
 [deps]
 ApproxBayes = "f5f396d3-230c-5e07-80e6-9fadf06146cc"

src/abc_inference.jl

@@ -1,6 +1,17 @@
-function createabcfunction(prob, t, distancefunction, alg; save_idxs = nothing, kwargs...)
+function createabcfunction(prob, t, distancefunction, alg; save_idxs = nothing, sample_u0 = false, kwargs...)
     function simfunc(params, constants, data)
-        sol = concrete_solve(STANDARD_PROB_GENERATOR(prob, params), alg; saveat = t, save_idxs = save_idxs, kwargs...)
+        local u0
+        if sample_u0
+            u0 = save_idxs === nothing ? params[1:length(prob.u0)] : params[1:length(save_idxs)]
+            if length(u0) < length(prob.u0)
+                for i in length(u0):length(prob.u0)
+                    push!(u0,prob.u0[i])
+                end
+            end
+        else
+            u0 = prob.u0
+        end
+        sol = concrete_solve(STANDARD_PROB_GENERATOR(prob, params), alg, u0; saveat = t, save_idxs = save_idxs, kwargs...)
         if size(sol, 2) < length(t)
             return Inf,nothing
         else
@@ -12,9 +23,9 @@ end
 
 function abc_inference(prob::DiffEqBase.DEProblem, alg, t, data, priors; ϵ=0.001,
                        distancefunction = euclidean, ABCalgorithm = ABCSMC, progress = false,
-                       num_samples = 500, maxiterations = 10^5, save_idxs = nothing, kwargs...)
+                       num_samples = 500, maxiterations = 10^5, save_idxs = nothing, sample_u0 = false, kwargs...)
 
-    abcsetup = ABCalgorithm(createabcfunction(prob, t, distancefunction, alg; save_idxs = save_idxs, kwargs...),
+    abcsetup = ABCalgorithm(createabcfunction(prob, t, distancefunction, alg; save_idxs = save_idxs, sample_u0 = sample_u0, kwargs...),
                             length(priors),
                             ϵ,
                             ApproxBayes.Prior(priors);

src/dynamichmc_inference.jl

@@ -36,9 +36,15 @@ function (P::DynamicHMCPosterior)(θ)
     @unpack algorithm, problem, data, t, parameter_priors = P
     @unpack σ_priors, solve_kwargs, sample_u0, save_idxs = P
     T = eltype(parameters)
-    nu = length(problem.u0)
+    nu = save_idxs == nothing ? length(problem.u0) : length(save_idxs)
     u0 = convert.(T, sample_u0 ? parameters[1:nu] : problem.u0)
     p = convert.(T, sample_u0 ? parameters[(nu + 1):end] : parameters)
+    if length(u0) < length(problem.u0)
+      # assumes u is ordered such that the observed variables are in the begining, consistent with ordered theta 
+      for i in length(u0):length(problem.u0)
+          push!(u0, convert(T,problem.u0[i]))
+      end
+    end
     _saveat = t === nothing ? Float64[] : t
     sol = concrete_solve(problem, algorithm, u0, p; saveat = _saveat, save_idxs = save_idxs, solve_kwargs...)
     failure = size(sol, 2) < length(_saveat)

src/stan_inference.jl

@@ -57,9 +57,9 @@ function stan_inference(prob::DiffEqBase.DEProblem,t,data,priors = nothing;alg=:
   length_of_params = length(vars)
   if isnothing(diffeq_string)
     sys = first(ModelingToolkit.modelingtoolkitize(prob))
-    length_of_parameter = length(sys.ps)
+    length_of_parameter = length(sys.ps) + sample_u0 * length(save_idxs)
   else
-    length_of_parameter = length(prob.p) + sample_u0 * length(prob.u0)
+    length_of_parameter = length(prob.p) + sample_u0 * length(save_idxs)
   end
   if alg ==:rk45
     algorithm = "integrate_ode_rk45"
@@ -90,8 +90,17 @@ function stan_inference(prob::DiffEqBase.DEProblem,t,data,priors = nothing;alg=:
   priors_string = string(generate_priors(length_of_parameter,priors))
   stan_likelihood = stan_string(likelihood)
   if sample_u0
-    nu = length(prob.u0)
-    integral_string = "u_hat = $algorithm(sho, theta[1:$nu], t0, ts, theta[$(nu+1):$length_of_parameter], x_r, x_i, $reltol, $abstol, $maxiter);"
+    nu = length(save_idxs)
+    if nu < length(prob.u0)
+      u0 = "{"
+      for u_ in prob.u0[nu+1:length(prob.u0)]
+        u0 = u0*string(u_)
+      end
+      u0 = u0*"}"
+      integral_string = "u_hat = $algorithm(sho, append_array(theta[1:$nu],$u0), t0, ts, theta[$(nu+1):$length_of_parameter], x_r, x_i, $reltol, $abstol, $maxiter);"
+    else 
+      integral_string = "u_hat = $algorithm(sho, theta[1:$nu], t0, ts, theta[$(nu+1):$length_of_parameter], x_r, x_i, $reltol, $abstol, $maxiter);"
+    end
   else
     integral_string = "u_hat = $algorithm(sho, u0, t0, ts, theta, x_r, x_i, $reltol, $abstol, $maxiter);"
   end

src/turing_inference.jl

@@ -30,9 +30,15 @@ function turing_inference(
         for i in 1:length(likelihood_dist_priors)
             σ[i] ~ likelihood_dist_priors[i]
         end
-        nu = length(prob.u0)
+        nu = save_idxs === nothing ? length(prob.u0) : length(save_idxs)
         u0 = convert.(T, sample_u0 ? theta[1:nu] : prob.u0)
         p = convert.(T, sample_u0 ? theta[(nu + 1):end] : theta)
+        if length(u0) < length(prob.u0)
+            # assumes u is ordered such that the observed variables are in the begining, consistent with ordered theta 
+            for i in length(u0):length(prob.u0)
+                push!(u0, convert(T,prob.u0[i]))
+            end
+        end
         _saveat = isnothing(t) ? Float64[] : t
         sol = concrete_solve(prob, alg, u0, p; saveat = _saveat, progress = progress, save_idxs = save_idxs, kwargs...)
         failure = size(sol, 2) < length(_saveat)
@@ -55,4 +61,4 @@ function turing_inference(
     model = mf(data)
     chn = sample(model, sampler, num_samples; progress = progress)
     return chn
-end
+end

test/abc.jl

@@ -1,5 +1,5 @@
 using DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions, Distances, StatsBase, RecursiveArrayTools
-using Test
+using Test, Distributions
 
 # One parameter case
 f1 = @ode_def begin
@@ -20,15 +20,33 @@ bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
 
 @test mean(bayesian_result.parameters, weights(bayesian_result.weights)) ≈ 1.5 atol=0.1
 
+priors = [Normal(1.,0.01),Normal(1.,0.01),Normal(1.5,0.01)]
+bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
+                                num_samples=500,ϵ = 0.001,sample_u0=true)
+
+meanvals = mean(bayesian_result.parameters, weights(bayesian_result.weights), 1)
+@test meanvals[1] ≈ 1. atol=0.1
+@test meanvals[2] ≈ 1. atol=0.1
+@test meanvals[3] ≈ 1.5 atol=0.1
+
 sol = solve(prob1,Tsit5(),save_idxs=[1])
 randomized = VectorOfArray([(sol(t[i]) + .01randn(1)) for i in 1:length(t)])
 data = convert(Array,randomized)
-
+priors = [Normal(1.5,0.01)]
 bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
                                 num_samples=500,ϵ = 0.001,save_idxs=[1])
 
 @test mean(bayesian_result.parameters, weights(bayesian_result.weights)) ≈ 1.5 atol=0.1
 
+priors = [Normal(1.,0.01),Normal(1.5,0.01)]
+bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
+                                num_samples=500,ϵ = 0.001,sample_u0=true,save_idxs=[1])
+
+meanvals = mean(bayesian_result.parameters, weights(bayesian_result.weights), 1)
+@test meanvals[1] ≈ 1. atol=0.1
+@test meanvals[2] ≈ 1.5 atol=0.1
+
+
 # custom distance-function
 weights_ = ones(size(data)) # weighted data
 for i = 1:3:length(data)
@@ -42,6 +60,7 @@ distfn = function (d1, d2)
     end
     return sqrt(d)
 end
+priors = [Normal(1.5,0.01)]
 bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
                                 num_samples=500, ϵ = 0.001,
                                 distancefunction = distfn)

test/dynamicHMC.jl

@@ -42,6 +42,14 @@ bayesian_result = dynamichmc_inference(prob1, Tsit5(), t, data, (Normal(1.5, 1),
                                            as(Vector, asℝ₊, 1),mcmc_kwargs=mcmc_kwargs, save_idxs = [1])
 @test mean(p.parameters[1] for p in bayesian_result.posterior) ≈ p[1] atol = 0.1
 
+priors = [Normal(1.,0.001),Normal(1.5,0.001)]
+mcmc_kwargs = (initialization = (q = zeros(2 + 1),), reporter = reporter)
+bayesian_result = dynamichmc_inference(prob1, Tsit5(), t, data, priors,
+                                           as(Vector, asℝ₊, 2),mcmc_kwargs=mcmc_kwargs,save_idxs=[1],sample_u0=true)
+
+@test mean(p.parameters[1] for p in bayesian_result.posterior) ≈ 1. atol = 0.1
+@test mean(p.parameters[2] for p in bayesian_result.posterior) ≈ 1.5 atol = 0.1
+
 # With hand-code likelihood function
 weights_ = ones(size(data)) # weighted data
 for i = 1:3:length(data)

test/stan.jl

@@ -35,13 +35,22 @@ sdf  = CmdStan.read_summary(bayesian_result.model)
 sol = solve(prob1,Tsit5(),save_idxs=[1])
 randomized = VectorOfArray([(sol(t[i]) + .01 * randn(1)) for i in 1:length(t)])
 data = convert(Array,randomized)
-
+priors = [Truncated(Normal(1.5,0.1),0,2)]
 bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
                                  num_warmup=500,likelihood=Normal,save_idxs=[1])
 
 sdf  = CmdStan.read_summary(bayesian_result.model)
 @test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1.5 atol=3e-1
 
+
+priors = [Normal(1.,0.01),Normal(1.5,0.01)]
+bayesian_result = stan_inference(prob1,t,data,priors;num_samples=300,
+                                 num_warmup=500,likelihood=Normal,save_idxs=[1],sample_u0=true)
+
+sdf  = CmdStan.read_summary(bayesian_result.model)
+@test sdf[sdf.parameters .== :theta1, :mean][1] ≈ 1. atol=3e-1
+@test sdf[sdf.parameters .== :theta2, :mean][1] ≈ 1.5 atol=3e-1
+
 println("Four parameter case")
 f1 = @ode_def begin
   dx = a*x - b*x*y

test/turing.jl

@@ -32,12 +32,19 @@ bayesian_result = turing_inference(prob1,Tsit5(),t,data,priors;num_samples=500,s
 sol = solve(prob1,Tsit5(),save_idxs=[1])
 randomized = VectorOfArray([(sol(t[i]) + .01 * randn(1)) for i in 1:length(t)])
 data = convert(Array,randomized)
-
+priors = [Normal(1.5,0.01)]
 bayesian_result = turing_inference(prob1,Tsit5(),t,data,priors;num_samples=500,
                                    syms=[:a],save_idxs=[1])
 
 @test mean(get(bayesian_result,:a)[1]) ≈ 1.5 atol=3e-1
 
+priors = [Normal(1.,0.01),Normal(1.5,0.01)]
+bayesian_result = turing_inference(prob1,Tsit5(),t,data,priors;num_samples=500,sample_u0 =true,
+                                   syms=[:u1,:a],save_idxs=[1])
+
+@test mean(get(bayesian_result,:a)[1]) ≈ 1.5 atol=3e-1
+@test mean(get(bayesian_result,:u1)[1]) ≈ 1.0 atol=3e-1
+
 println("Four parameter case")
 f2 = @ode_def begin
     dx = a*x - b*x*y