Why does EnsembleGPUArray not save at the given time points the solution?

Specific Domains GPU
1

I am trying to solve the same ODE for multiple parameters on the GPU, however as you can see in the provided code, the size of the solution trajectories is not the expected one. I get solution that are of the size 7*6, but i want the solutions to be saved at all time points in time_points. Does anybody know where the mistake is? I already varied the batchsize because of the synchronized time stepping, but that didn’t help.

using DiffEqGPU, OrdinaryDiffEq
function Glucose_Yeast(du,u,p,t)
    du[1] = p[1] - p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10]))
    du[2] = 2 * p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10])) - p[3]*u[2]*(p[12]-u[5]) - p[7]*u[2]*u[5]
    du[3] = p[3]*u[2]*(p[12]-u[5]) - p[4]*u[3]*(p[13] - u[6])
    du[4] = p[4]*u[3]*(p[13] - u[6]) - p[5]*u[4]*u[5] - p[9]*(u[4]-u[7])
    du[5] = p[3]*u[2]*(p[12]-u[5]) - p[5]*u[4]*u[5] - p[7] * u[2] * u[5]
    du[6] = -2 * p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10])) + 2 * p[4]*u[3]*(p[13] - u[6]) - p[6] * u[6]
    du[7] = p[14]*p[9] * (u[4] - u[7]) - p[9] * u[7]
  end

time_points = [5, 15, 23,30, 35, 50, 65, 75, 99, 120, 140, 150, 170, 175, 198, 200]

parameters = [2.5, 100, 6, 16, 100,1.28, 12, 1.8, 13, 4, 0.52, 1, 4, 0.1]

no_samples = 1000
tr = (0,200)
u0 = [0.8328305840762804, 0.07874494552628386, 0.29149315284561084, 0.7805647100582624, 0.7934642151447182, 0.14760486012152874, 0.4416625406969794]


num_pars = length(parameters)
lbs = zeros(num_pars)
ubs = zeros(num_pars)
for j in 1:num_pars
    lbs[j] = 1e-8
    ubs[j] = 2*parameters[j]
end
s = SobolSeq(lbs, ubs)
s = Sobol.skip(s, no_samples)

Pre_Sobol_Matrix = zeros((num_pars, no_samples))
for j in 1:no_samples
    Pre_Sobol_Matrix[:,j] = Sobol.next!(s)
end

function GPU_solver(ODE, parameters, u0, tr, Pre_Sobol_Matrix, no_samples, time_points)
    prob = ODEProblem(ODE,u0,tr, parameters[1])
    function prob_func(prob,l,repeat)

        #Create a new parameter set
        new_p =  collect(Pre_Sobol_Matrix[:,l])
    
        #Launch a new trajectory
        remake(prob, p = new_p)
    end
    monteprob = EnsembleProblem(prob, prob_func = prob_func)

    println("Start solving")

    @time sol = solve(
            monteprob,
            Tsit5(),
            EnsembleGPUArray(), 
            trajectories = no_samples,
            batchsize = 100,
            rtol=1e-6,
            atol=1e-6, 
            saveat = time_points 
        )
    return sol
end

sol =  GPU_solver(Glucose_Yeast, parameters, u0, tr, Pre_Sobol_Matrix, no_samples, time_points)

for i in 1:no_samples
    println(size(sol[:,:,i]))
end
2

The solver throws a warning:

┌ Warning: Interrupted. Larger maxiters is needed. If you are using an integrator for non-stiff ODEs or an automatic switching algorithm (the default), you may want to consider using a method for stiff equations. See the solver pages for more details (e.g. https://diffeq.sciml.ai/stable/solvers/ode_solve/#Stiff-Problems).
└ @ SciMLBase C:\Users\accou\.julia\dev\SciMLBase\src\integrator_interface.jl:352

This means it exited early and Tsit5 is likely not the right solver for this. In general, chemical reaction networks are stiff so Tsit5 was a bad choice here. Try Rodas5 instead.

BTW, more correct code:

using DiffEqGPU, OrdinaryDiffEq, Sobol
function Glucose_Yeast(du,u,p,t)
    du[1] = p[1] - p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10]))
    du[2] = 2 * p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10])) - p[3]*u[2]*(p[12]-u[5]) - p[7]*u[2]*u[5]
    du[3] = p[3]*u[2]*(p[12]-u[5]) - p[4]*u[3]*(p[13] - u[6])
    du[4] = p[4]*u[3]*(p[13] - u[6]) - p[5]*u[4]*u[5] - p[9]*(u[4]-u[7])
    du[5] = p[3]*u[2]*(p[12]-u[5]) - p[5]*u[4]*u[5] - p[7] * u[2] * u[5]
    du[6] = -2 * p[2] * u[1] * u[6] / (1 + u[6] / (p[11]^p[10])) + 2 * p[4]*u[3]*(p[13] - u[6]) - p[6] * u[6]
    du[7] = p[14]*p[9] * (u[4] - u[7]) - p[9] * u[7]
  end

time_points = [5, 15, 23,30, 35, 50, 65, 75, 99, 120, 140, 150, 170, 175, 198, 200]
parameters = [2.5, 100, 6, 16, 100,1.28, 12, 1.8, 13, 4, 0.52, 1, 4, 0.1]
no_samples = 1000
tr = (0,200)
u0 = [0.8328305840762804, 0.07874494552628386, 0.29149315284561084, 0.7805647100582624, 0.7934642151447182, 0.14760486012152874, 0.4416625406969794]


num_pars = length(parameters)
lbs = zeros(num_pars)
ubs = zeros(num_pars)
for j in 1:num_pars
    lbs[j] = 1e-8
    ubs[j] = 2*parameters[j]
end
s = SobolSeq(lbs, ubs)
s = Sobol.skip(s, no_samples)

Pre_Sobol_Matrix = zeros((num_pars, no_samples))
for j in 1:no_samples
    Pre_Sobol_Matrix[:,j] = Sobol.next!(s)
end

function GPU_solver(ODE, parameters, u0, tr, Pre_Sobol_Matrix, no_samples, time_points)
    prob = ODEProblem(ODE,u0,tr, parameters[1])
    function prob_func(prob,l,repeat)

        #Create a new parameter set
        new_p =  collect(Pre_Sobol_Matrix[:,l])
    
        #Launch a new trajectory
        remake(prob, p = new_p)
    end
    monteprob = EnsembleProblem(prob, prob_func = prob_func)

    println("Start solving")

    @time sol = solve(
            monteprob,
            Tsit5(),
            EnsembleGPUArray(), 
            trajectories = no_samples,
            batch_size = 100,
            reltol=1e-6,
            abstol=1e-6, 
            saveat = time_points 
        )
    return sol
end

sol =  GPU_solver(Glucose_Yeast, parameters, u0, tr, Pre_Sobol_Matrix, no_samples, time_points)

for i in 1:no_samples
    println(size(sol[:,:,i]))
end