Skip to content

Instantly share code, notes, and snippets.

@nmatzke

nmatzke/BDQ_v5.jl

Created Apr 22, 2019
Embed
What would you like to do?
BDQ v5 w help from Chris
using DifferentialEquations
using BenchmarkTools # for @benchmark
using LSODA # for lsoda()
using Sundials # for CVODE_BDF()
# State-dependent birth-death process:
# Problem setup
n = 500 # Number of states
two = 2.0 # Constant (but missing in some published equations)
# Define Q
# Here, allow transitions only between adjacent states
# (may be different in other problems)
Qi = collect(1:(n-1))
Qi = vcat(Qi, collect(2:n))
Qj = collect(2:n)
Qj = vcat(Qj, collect(1:(n-1)))
Qr = vcat(repeat([0.1],(n-1)), repeat([0.01],(n-1)))
# Define lambdas (L) (birthrates) (for now, just boring i->i,i births)
Li = collect(1:n)
Lj = collect(1:n)
Lk = collect(1:n)
Lr = repeat([0.5], n)
# Define mus (death rates)
mu = repeat([0.5], n)
# Likelihood Equation: for each state i there are 2 equations:
# * Et[i] is the probability of eventual death of an individual
# in state i at time t, by the time of observation (t0=0.0)
# * Dt[i] is the Likelihood of the observed data at t0=0.0 given
# the individual is in state i at time t
#
# The calculation of dEi and dDi for state i involves many
# ==i and !=i operations across Q and L. These only need to be
# done once per problem (may or may not save time to
# pre-calculate).
Li_eq_i = Any[]
Qi_not_i = Any[]
Lj_sub_i = Any[]
Lk_sub_i = Any[]
Qj_not_i = Any[]
for i in 1:n
push!(Li_eq_i, Li .== i)
push!(Qi_not_i, Qi .!= i)
push!(Lj_sub_i, Lj[Li .== i])
push!(Lk_sub_i, Lk[Li .== i])
push!(Qj_not_i, Qj[Qi .!= i])
end
# Parameters
p = (n=n, two=two, Qi=Qi, Qj=Qj, Qr=Qr, Li=Li, Lj=Lj, Lk=Lk, Lr=Lr, mu=mu, Li_eq_i=Li_eq_i, Qi_not_i=Qi_not_i, Lj_sub_i=Lj_sub_i, Lk_sub_i=Lk_sub_i, Qj_not_i=Qj_not_i)
# Likelihood function
function BDQ_readable_v5!(du,u,p,t)
n, two, Qi, Qj, Qr, Li, Lj, Lk, Lr, mu = p
Et = @view u[1:n]
Dt = @view u[(n+1):(2*n)]
dEt = @view du[1:n]
dDt = @view du[(n+1):(2*n)]
@inbounds for i in 1:n
#Lj_sub_i = @view Lj[p.Li_eq_i[i]]
#Lk_sub_i = @view Lk[p.Li_eq_i[i]]
Lr_sub_i = @view Lr[p.Li_eq_i[i]]
Qr_not_i = @view Qr[p.Qi_not_i[i]]
#Qj_not_i = @view Qj[p.Qi_not_i[i]]
Et_Qj_not_i = @view Et[p.Qj_not_i[i]]
Et_Lj_sub_i = @view Et[p.Lj_sub_i[i]]
Et_Lk_sub_i = @view Et[p.Lk_sub_i[i]]
Dt_Lk_sub_i = @view Dt[p.Lk_sub_i[i]]
Dt_Lj_sub_i = @view Dt[p.Lj_sub_i[i]]
Dt_Qj_not_i = @view Dt[p.Qj_not_i[i]]
dEt[i] = mu[i] +
-(sum(Lr_sub_i) + sum(Qr_not_i) + mu[i])*Et[i] +
(sum(Qr_not_i .* Et_Qj_not_i)) +
(two * sum(Lr_sub_i .* Et_Lj_sub_i .* Et_Lk_sub_i))
dDt[i] = -(sum(Lr_sub_i) + sum(Qr_not_i) + mu[i])*Dt[i] +
(sum(Qr_not_i .* Dt_Qj_not_i)) +
(sum(Lr_sub_i .*
(Dt_Lk_sub_i.*Et_Lj_sub_i
.+ Dt_Lj_sub_i.*Et_Lk_sub_i) ))
end
end
u = repeat([0.0], 2*n)
du = repeat([0.0], 2*n)
u0 = repeat([0.0], 2*n)
u0[n+1] = 1.0
tspan = (0.0, 1.0)
prob = ODEProblem(BDQ_readable_v5!, u0, tspan, p)
@benchmark sol_CVODE = solve(prob, CVODE_BDF(linear_solver=:GMRES), dense=false, save_everystep=false, save_end=true, save_start=true)
# memory estimate: 2.72 GiB
# allocs estimate: 5711652
# --------------
# minimum time: 1.638 s (16.58% GC)
# median time: 1.707 s (18.34% GC)
# mean time: 1.690 s (17.78% GC)
# maximum time: 1.726 s (18.35% GC)
# --------------
# samples: 3
# evals/sample: 1
@benchmark sol_Tsit5 = solve(prob, Tsit5(), dense=false, save_everystep=false, save_end=true, save_start=false)
# memory estimate: 3.29 GiB
# allocs estimate: 6928602
# --------------
# minimum time: 2.158 s (18.41% GC)
# median time: 2.166 s (18.40% GC)
# mean time: 2.166 s (18.41% GC)
# maximum time: 2.174 s (18.33% GC)
# --------------
# samples: 3
# evals/sample: 1
@benchmark sol_Vern9 = solve(prob, Vern9(), dense=false, save_everystep=false, save_end=true, save_start=false)
# memory estimate: 6.43 GiB
# allocs estimate: 13521198
# --------------
# minimum time: 4.155 s (18.08% GC)
# median time: 4.173 s (18.19% GC)
# mean time: 4.173 s (18.19% GC)
# maximum time: 4.190 s (18.31% GC)
# --------------
# samples: 2
# evals/sample: 1
sol_CVODE = solve(prob, CVODE_BDF(linear_solver=:GMRES), dense=false, save_everystep=false, save_end=true, save_start=true)
# [0.353391, 0.353391, 0.353391, 0.353391, 0.353391, 0.353391, 0.353391, 0.353391, 0.353391, 0.353391 … 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5, 8.13777e-5]
sol_Tsit5 = solve(prob, Tsit5(), dense=false, save_everystep=false, save_end=true, save_start=false)
# [0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296 … 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5, 8.14129e-5]
sol_Vern9 = solve(prob, Vern9(), dense=false, save_everystep=false, save_end=true, save_start=false)
# [0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296, 0.353296 … 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5, 8.14135e-5]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.