Hi Everyone,
I’m finding that using DiscreteProblem on my toy example is about 4x slower than a vanilla loop with pre-allocated output.
DiscreteProblem version
Base Julia version
My model is as follows:
function sir_map!(du,u,p,t)
(S,I,R) = u
(β,c,γ,δt) = p
N = S+I+R
infection = rate_to_proportion(β*c*I/N,δt)*S
recovery = rate_to_proportion(γ,δt)*I
@inbounds begin
du[1] = S-infection
du[2] = I+infection-recovery
du[3] = R+recovery
end
nothing
end;
And my base code loop for solving it is as follows:
function solve_map(f, u0, nsteps, p)
# Pre-allocate array with correct type
sol = similar(u0, length(u0), nsteps + 1)
# Initialize the first column with the initial state
sol[:, 1] = u0
# Iterate over the time steps
@inbounds for t in 2:nsteps+1
u = @view sol[:, t-1] # Get the current state
du = @view sol[:, t] # Prepare the next state
f(du, u, p, t) # Call the function to update du
end
return sol
end;
I solve the above using DiscreteProblem as follows.
prob_map = DiscreteProblem(sir_map!,u0,tspan,p)
sol_map = solve(prob_map,FunctionMap());
I suspect that the time difference between the implementations is due to not pre-allocating the output for the solution of DiscreteProblem; any suggestions on how to speed it up?