BDQ v5 w help from Chris

BDQ_v5.jl

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]