The actual numerics functions
begin
abstract type AbstractIntegrationRule end
struct TrapezoidalIntegrationRule <: AbstractIntegrationRule end
struct JumpConfig{T<:AbstractIntegrationRule}
n::Int # mesh points
Δt::Float64 # time step
integration_rule::T
end
function integration_JuMP!(model, u, du, jc, ::TrapezoidalIntegrationRule)
@variable(model, δt[1:jc.n] == jc.Δt)
for t in 2:jc.n
s = t - 1
for k in 1:length(u)
@constraint(model, u[k][t] == u[k][s] + 0.5 * δt[t] * (du[k][s] + du[k][t]))
end
end
end
function fast_sim(
u0, jc::JumpConfig
)
n = jc.n
model = Model(Ipopt.Optimizer)
# dynamical variables
u = @JuMP.variables(model,begin
0 <= φ_R[1:n] <= 0.45
0 <= c_pc[1:n]
0 <= M[1:n]
end);
# optimization knob
@variable(model, 0 <= α_R[1:n] <= 0.45 )
# helper functions
@expression(model, γ[j = 1:n], 10 * c_pc[j]/(c_pc[j] + 0.03))
@expression(model, λ[j = 1:n], γ[j] * φ_R[j])
du = Vector{Vector{NonlinearExpr}}(undef, length(u))
# growth dynamics
du[1] = @expression(model, δφ_R[j = 1:n], (α_R[j] - φ_R[j]) * λ[j])
du[2] = @expression(model, δc_pc[j =1:n], 5*(0.45 - φ_R[j]) - (1 + c_pc[j]) * λ[j] )
du[3] = @expression(model, δM[j = 1:n], M[j] * λ[j])
# integration
integration_JuMP!(model, u, du, jc, TrapezoidalIntegrationRule())
@objective(model, Max, M[n])
set_silent(model)
# one set of initial conditions
for i in 1:length(u)
fix(u[i][1], u0[i]; force = true)
end
optimize!(model)
Ls = value.(λ)
Rs = value.(φ_R)
aRs = value.(α_R)
cpcs = value.(c_pc)
return Rs, cpcs, aRs, Ls
end
function slow_sim(
u0, jc::JumpConfig; Kd = 1e-6, k_on = 4e6, C_max = 1e-6
)
n = jc.n
model = Model(Ipopt.Optimizer)
# dynamical variables
u = @JuMP.variables(model,begin
0 <= φ_R[1:n] <= 0.45
0 <= c_pc[1:n]
0 <= C_bnd[1:n]
0 <= C_in[1:n]
0 <= M[1:n]
end);
# optimization knob
@variable(model, 0 <= α_R[1:n] <= 0.45 )
# helper functions
@expression(model, C_out[j = 1:n], C_max)
@expression(model, C_Rf[j = 1:n], φ_R[j] * 6.4e-5 - C_bnd[j])
@expression(model, φ_Rf[j = 1:n], φ_R[j] - C_bnd[j] / 6.4e-5)
@constraint(model, φ_Rf .>= 0)
@expression(model, φ_M[j = 1:n], 0.45 - φ_R[j])
@expression(model, γ[j = 1:n], 10 * c_pc[j]/(c_pc[j] + 0.03))
@expression(model, λ[j = 1:n], γ[j] * φ_Rf[j])
du = Vector{Vector{NonlinearExpr}}(undef, length(u))
# growth dynamics
k_off = k_on * Kd # 0.4 # 40
du[1] = @expression(model, δφ_R[j = 1:n], (α_R[j] - φ_R[j]) * λ[j])
# 1/s in mass fractions -> multiply by ρ to get conc.
du[2] = @expression(model, δc_pc[j =1:n], 5*(0.45 - φ_R[j]) - (1 + c_pc[j]) * λ[j] )
du[3] = @expression(model, δC_bnd[j = 1:n], k_on*C_in[j]*C_Rf[j] - k_off*C_bnd[j] - λ[j]*C_bnd[j])
du[4] = @expression(model, δC_in[j = 1:n],
1.5*(C_out[j] - C_in[j]) + k_off*C_bnd[j] - k_on*C_Rf[j]*C_in[j] - λ[j]*C_in[j]
)
du[5] = @expression(model, δM[j = 1:n], M[j] * λ[j])
# integration
integration_JuMP!(model, u, du, jc, TrapezoidalIntegrationRule())
@objective(model, Max, M[n])
set_silent(model)
for i in 1:length(u)
fix(u[i][1], u0[i]; force = true)
end
optimize!(model)
Ls = value.(λ)
Rs = value.(φ_R)
a_Rs = value.(α_R)
cpcs = value.(c_pc)
C_bnds = value.(C_bnd)
C_ins = value.(C_in)
return Rs, a_Rs, cpcs, C_bnds, C_ins, Ls
end
end
Example code to run it:
begin
u0_fast = [0.1, 0.05, 0.001]
steps = 1000
δt = 1/250
jc = JumpConfig(steps, δt, TrapezoidalIntegrationRule())
fast_soln = fast_sim(u0_fast, jc)
end
It takes ~2s to run the above in a Pluto NB ^
u0_slow = [0.1, 0.05, 0.0, 0.0, 0.001]
slow_soln = slow_sim(u0_slow, jc)
and ~30s to run ^ the slow sim.
I was able to speed it up by adding a constraint, \lambda \geq 0, but still looking for it to be faster.
The speed is significantly affected by the magnitudes of k_on
and C_max
; I have them set lower than I want them to get the simulation to run. Something about the significant difference in the parameter magnitudes is significantly slowing down the optimization.
Here are also some plotting utilities for plotting the slow solution, in case it is useful. The results from the above optimizations can be plotted like plot_soln(slow_soln, jc)
.
My goal is to figure out how to get the slow_sim
faster, fast_sim
was just an intermediate step that I am sharing for background.
const green = colorant"#04C13C"
const purp = colorant"#BA79BE"
const pink = colorant"#FA007B"
const blue = colorant"#05A8C9"
const brown = colorant"#8B4513"
const orange = colorant"#FFA500"
const thick_line = 3.5
const thin_line = 2
const extra_thin_line = 1
const opt_color = :gray
function plot_soln(soln, jc::JumpConfig)
Rs = soln[1]
aRs = soln[2]
cpcs = soln[3]
normed_cpcs = cpcs ./ 0.03
C_bnds = soln[4]
C_ins = soln[5]
Ls = soln[6]
ts = cumsum([0; repeat([jc.Δt], jc.n)])[1:end-1]
plt_growth = λ_plot(ts, Ls)
plt_pfractions = φR_plot(ts, Rs, aRs)
plt_pcursors = precursor_plot(ts, normed_cpcs, xlab=true)
plt_ABX = ABX_plot(ts, C_ins, C_bnds, xlab=true)
p = plot(
plt_pfractions, plt_growth, plt_pcursors, plt_ABX;
layout = grid(2,2), size = (700,700), legend=:right
)
return p
end
function precursor_plot(ts, normed_c_pcs; subset_idxs=nothing, xlab=false)
if xlab
xlab = "Time (hrs)"
else
xlab = ""
end
plt = plot_series(
ts, normed_c_pcs, blue;
title = "Precursors",
xlab = xlab,
ylab = "Concentration (Km)",
subset_idxs = subset_idxs,
)
return plt
end
function φR_plot(ts, φ_Rs, α_Rs; subset_idxs=nothing, xlab=false)
if xlab
xlab = "Time (hrs)"
else
xlab = ""
end
plt = plot_series(
ts, φ_Rs, purp;
title = "Protein Fractions",
xlab = xlab,
ylab = "Mass Fraction",
subset_idxs = subset_idxs,
)
plot_series!(
plt, ts, α_Rs, pink;
subset_idxs = subset_idxs, lw = extra_thin_line, ls = :dashdot,
)
return plt
end
function λ_plot(ts, λs; subset_idxs=nothing, xlab=false)
if xlab
xlab = "Time (hrs)"
else
xlab = ""
end
plt = plot_series(
ts, λs, green;
title = "Growth Rate",
xlab = xlab,
ylab = "Rate (1/hr)",
subset_idxs = subset_idxs,
)
return plt
end
function ABX_plot(ts, C_ins, C_bnds; subset_idxs=nothing, xlab=false)
if xlab
xlab = "Time (hrs)"
else
xlab = ""
end
plt = plot_series(
ts, C_bnds, orange;
title = "ABX",
xlab = xlab,
ylab = "Concentration ([M])",
subset_idxs = subset_idxs,
label = "Bound to Ribosome"
)
plot_series!(
plt, ts, C_ins, brown;
subset_idxs = subset_idxs,
label = "Intracellular",
)
return plt
end
function plot_series(
ts, xs, color::RGB;
title, xlab, ylab,
subset_idxs=nothing, lw = thick_line,
ls = :solid, label = nothing
)
xs, n_plots = subset_series(xs, subset_idxs)
colors = set_colors(n_plots, color)
plt = plot(title = title, xlabel = xlab, ylabel = ylab)
for i in 1:n_plots
if i == n_plots
lab = label
else
lab = nothing
end
plot!(
ts,
xs[:,i];
label = lab,
linecolor = colors[i],
linewidth = lw, linestyle = ls,
)
end
return plt
end
function plot_series!(
plt, ts, xs, color::RGB;
subset_idxs=nothing, lw = thick_line,
ls = :solid, label = nothing
)
xs, n_plots = subset_series(xs, subset_idxs)
colors = set_colors(n_plots, color)
for i in 1:n_plots
if i == n_plots
lab = label
else
lab = nothing
end
plot!(
plt,
ts,
xs[:,i];
label = lab,
linecolor = colors[i],
linewidth = lw, linestyle = ls,
)
end
return plt
end
function set_colors(n_plots, base)
if n_plots == 1
return [base]
else
return gen_color_range(n_plots, base)
end
end
function subset_series(series, subset_idxs)
if series isa Matrix
if subset_idxs != nothing
series = series[:, subset_idxs]
end
n_plots = size(series)[2]
elseif series isa Vector
n_plots = 1
else
error("series must be a matrix or vector")
end
return series, n_plots
end
function gen_color_range(n, end_color, start_color=colorant"cornsilk3")
if n == 1
return [end_color]
end
# rescale
xs = (0:n-1)/(5/4*(n-1)) .+ 0.2
cs = ColorScheme(range(start_color, end_color))
return get(cs, xs)
end