Here’s the code in a single file.

It might take me a while to actually understand what’s going on, but there are a few red flags that stand out. You should essentially never need to use `eval`

or `invoke`

latest, and the mix of `DifferentialEquations`

and `Symbolics`

with JuMP can cause issues. This code is also very far from what the C++ is doing, so it’s no wonder that there is a performance difference.

How long does their Ipopt code take to run? And is the RAM issue here because of the solve or during the setup? Asked a different way, if you remove `optimize!`

, is there still an issue?

```
using DifferentialEquations
using Symbolics
using Plots
using Parameters
using JuMP
using Ipopt
using StaticArrays
using Interpolations
using IOCapture
using JLD2
struct VarAssimMeta
vf::Function
vf_strs::Vector{String}
D::Int
P::Int
observed_dims::Vector{Int}
driven_dims::Vector{Int}
L::Int
I::Int
lb_states::Vector
ub_states::Vector
lb_params::Vector
ub_params::Vector
end
function VarAssimMeta(
vf::Function,
D::Int,
P::Int,
observed_dims::Vector{Int},
driven_dims::Vector{Int},
lb_states::Vector,
ub_states::Vector,
lb_params::Vector,
ub_params::Vector,
)
L = length(observed_dims)
I = length(driven_dims)
Symbolics.@variables u[1:D], p[1:P], f[1:I]
vf_syms = vf(u, p, f)
vf_strs = repr.(vf_syms)
VarAssimMeta(vf, vf_strs, D, P, observed_dims, driven_dims, L, I, lb_states, ub_states, lb_params, ub_params)
end
function jump_nl_substitute!(model, eqstr, prec, urec, frec)
npts = size(frec)[2]
@assert size(urec)[2] == npts
find_pats = [Regex("$x" * raw"\[(\d+)\]") for x in ["p", "u", "f"]]
repl_pats = [SubstitutionString("opt_p[\\1]"), SubstitutionString("opt_u[\\1,j]"), SubstitutionString("drive[\\1,j]")]
neweqstr = replace(eqstr, [find_pats[i] => repl_pats[i] for i in eachindex(find_pats)]...)
model_lambda_str = """(model, opt_p, opt_u, drive) -> @NLexpression(
model,
[j=1:$npts],
$neweqstr);
"""
model_lambda = eval(Meta.parse(model_lambda_str))
reshape(Base.invokelatest(model_lambda, model, prec, urec, frec), (1, npts))
end
abstract type AbstractVADiscretization end
struct SimpsonHermite <: AbstractVADiscretization end
function gen_dspe_jump_model(vam::VarAssimMeta, disc::SimpsonHermite, timepoints, data, drive_frozen)
N = length(timepoints) - 1
dt = diff(timepoints)
@unpack vf_strs, D, P, observed_dims, lb_states, ub_states, lb_params, ub_params = vam
L = length(observed_dims)
model = Model()
JuMP.@variables(model, begin
opt_x[1:D, 1:(N+1)]
opt_xm[1:D, 1:N]
0.0 <= opt_k[1:(2*N+1)] <= 1.0
opt_p[1:P]
end)
drive = vcat(drive_frozen, opt_k')
data_x = data[:, 1:2:(2*N+1)]
data_xm = data[:, 2:2:(2*N)]
drive_x = drive[:, 1:2:(2*N+1)]
drive_xm = drive[:, 2:2:(2*N)]
println("Generating timepoint and midpoint vector field subexpressions")
F = vcat([jump_nl_substitute!(model, eqstr, opt_p, opt_x, drive_x) for eqstr in vf_strs]...)
Fm = vcat([jump_nl_substitute!(model, eqstr, opt_p, opt_xm, drive_xm) for eqstr in vf_strs]...)
println("Generating Simpson residual subexpressions.")
S = @NLexpression(
model,
[i = 1:D, j = 1:N],
(opt_x[i, j] + (1 / 6) * dt[j] * F[i, j]) + (2 / 3) * dt[j] * Fm[i, j] + (-opt_x[i, j+1] + (1 / 6) * dt[j] * F[i, j+1])
)
println("Generating Hermite residual subexpressions.")
H = @NLexpression(
model,
[i = 1:D, j = 1:N],
((1 / 2) * opt_x[i, j] + (dt[j] / 8) * F[i, j]) + ((1 / 2) * opt_x[i, j+1] - (dt[j] / 8) * F[i, j+1]) - opt_xm[i, j],
)
@NLconstraint(model, [i = 1:D, j = 1:N], S[i, j] == 0.0)
@NLconstraint(model, [i = 1:D, j = 1:N], H[i, j] == 0.0)
for j = 1:D
set_lower_bound.(opt_x[j, :], lb_states[j])
set_lower_bound.(opt_xm[j, :], lb_states[j])
set_upper_bound.(opt_x[j, :], ub_states[j])
set_upper_bound.(opt_xm[j, :], ub_states[j])
end
for i = 1:P
set_lower_bound(opt_p[i], lb_params[i])
set_upper_bound(opt_p[i], ub_params[i])
end
@NLobjective(
model,
Min,
(
sum((opt_x[i, j] - data_x[i, j])^2 for i in observed_dims, j = 1:(N+1)) +
sum((opt_xm[i, j] - data_xm[i, j])^2 for i in observed_dims, j = 1:N) +
sum(opt_k[j]^2 for j = 1:(2*N+1))
) / (2 * N + 1)
)
println("Model generation complete!")
model
end
gtp = [1.0, -54.4, -77.0, 50.0, 0.3, 20.0, 120.0, -40.0, 15.0, 0.1, 0.4, 0.0, -60.0, -15.0, 1.0, 7.0, 0.0, -55.0, 30.0, 1.0, 5.0, 0.0];
function lorenz(u, p, t)
dx = 10.0 * (u[2] - u[1])
dy = u[1] * (28.0 - u[3]) - u[2]
dz = u[1] * u[2] - (8 / 3) * u[3]
SA[dx, dy, dz]
end
function scale_timeseries(ts, minv, maxv)
tsmax = maximum(ts)
tsmin = minimum(ts)
minv .+ (maxv - minv) * (ts .- tsmin) / (tsmax - tsmin)
end
lor_tmax = 80.0
lor_dt = 0.005
lor_u0 = SA[-1.31; 0.8; 19.77]
lor_tspan = (0.0, lor_tmax)
lor_prob = ODEProblem(lorenz, lor_u0, lor_tspan)
lor_sol = solve(lor_prob, RK4(), adaptive = true, dt = lor_dt);
lorfunc = linear_interpolation(scale_timeseries(lor_sol.t, 0.0, 300.0), scale_timeseries(lor_sol[1, :], -20.0, 20.0), extrapolation_bc = 0.0);
const tanhgain = 1.0;
@inline nlss(V, θ, σ) = 0.5 * (1 + tanh(tanhgain * (V - θ) / σ))
@inline nltc(V, θ, σ, τ0, τ1, τ2) = τ0 + τ1 * (1 - tanh(tanhgain * (V - θ) / σ)^2) + τ2 * (1 + tanh(tanhgain * (V - θ) / σ))
function hhvf_oop(u, p, t)
pvec, Iapp = p
Cm, EL, EK, ENa, gL, gK, gNa, θm, σm, τm0, τm1, τm2, θh, σh, τh0, τh1, τh2, θn, σn, τn0, τn1, τn2 = pvec
V, m, h, n = u
dV = (1 / Cm) * (gNa * m * m * m * h * (ENa - V) + gK * n * n * n * n * (EK - V) + gL * (EL - V) + Iapp(t))
dm = (nlss(V, θm, σm) - m) / nltc(V, θm, σm, τm0, τm1, τm2)
dh = (nlss(V, θh, σh) - h) / nltc(V, θh, σh, τh0, τh1, τh2)
dn = (nlss(V, θn, σn) - n) / nltc(V, θn, σn, τn0, τn1, τn2)
[dV, dm, dh, dn]
end
hh_tmax = 300.0
hh_saveat = 0.04
hh_u0 = [
-68.24221681836171
0.056029230048653705
0.7700232861002139
0.3402655968929933
];
hh_prob = ODEProblem(hhvf_oop, hh_u0, (0.0, hh_tmax), (gtp, lorfunc))
hh_sol = solve(hh_prob, RK4(), adaptive = true);
function hhvf_assim_controlled(u, p, f)
Cm, Isa, EL, EK, ENa, gL, gK, gNa, θm, σm, τm0, τm1, τm2, θh, σh, τh0, τh1, τh2, θn, σn, τn0, τn1, τn2 = p
I, data, k = f
V, m, h, n = u
dV = (1 / Cm) * (gNa * m * m * m * h * (ENa - V) + gK * n * n * n * n * (EK - V) + gL * (EL - V) + I / Isa) + k * (data - V)
dm = (nlss(V, θm, σm) - m) / nltc(V, θm, σm, τm0, τm1, τm2)
dh = (nlss(V, θh, σh) - h) / nltc(V, θh, σh, τh0, τh1, τh2)
dn = (nlss(V, θn, σn) - n) / nltc(V, θn, σn, τn0, τn1, τn2)
[dV, dm, dh, dn]
end
plower = [1.0, 0.0005, -100.0, -100.0, 0.0, 0.0, 0.0, 0.0, -100.0, 0.01, 0.01, 0.01, 0.0, -100.0, -100.0, 0.01, 0.01, 0.0, -100.0, 0.01, 0.01, 0.01, 0.0]
pupper = [1.0, 1.0, 0.0, 0.0, 100.0, 100.0, 1000.0, 1500.0, -0.01, 100.0, 10.0, 10.0, 0.0, -0.01, -0.01, 10.0, 10.0, 0.0, -0.01, 100.0, 10.0, 10.0, 0.0]
slower = [-150.0, 0.0, 0.0, 0.0]
supper = [150.0, 1.0, 1.0, 1.0];
vam = VarAssimMeta(hhvf_assim_controlled, 4, 23, [1], [1, 2, 3], slower, supper, plower, pupper)
timepoints = hh_sol.t
N = length(timepoints) - 1
midt = (hh_sol.t[1:end-1] + hh_sol.t[2:end]) / 2.0;
allt = vcat(vec(hcat(hh_sol.t[1:end-1], midt)'), hh_sol.t[end]);
middatamat = hcat(hh_sol.(midt)...);
alldatamat = hcat(hh_sol.(allt)...);
obsdatamat = alldatamat[1:1, :];
drive = lorfunc.(allt)';
noise = 0.1 * randn(2N + 1);
data = obsdatamat + noise';
drive_frozen = vcat(drive, data);
model = gen_dspe_jump_model(vam, SimpsonHermite(), timepoints, data, drive_frozen)
set_optimizer(model, Ipopt.Optimizer)
optimize!(model)
```