I share a toy model that I use for this purpose. The first is the function to define the ODEs, the second is the main code.
function pk_model!(du, u, p, t)
# PROBLEM
# Model: `pk1`
# - a simple one compartment pk mode
# PARAMETERS
CL = p[1]; # Clearance (volume/time)
V = p[2]; # Central volume (volume)
KA = p[3]; # Absorption rate constant 1 (1/time)
# CMT
EV = u[1]; # extravascular compartment (mass)
CENT = u[2]; # Central compartment (mass)
#define Central compartment concentration
CP = CENT/V;
# ODEs
dxdt_EV = -KA *EV;
dxdt_CENT = KA *EV - CL*CP;
# Function outputs
du[1] = dxdt_EV;
du[2] = dxdt_CENT
# Auxiliary output to calculate AUC, for sensitivity analysis
dxdt_AUC = CENT;
du[3] = dxdt_AUC;
end
using DifferentialEquations
using GlobalSensitivity, QuasiMonteCarlo, Statistics, Random
using DiffEqGPU, CUDA
using JLD2
using ComponentArrays
using OrderedCollections
include("pk_model_sens.jl")
Dose_mg = 100.0f0;
u0 = [Dose_mg, 0.0f0, 0.0f0];
param = ComponentArray(CL = 1.0f0, V = 20.0f0, KA = 1.0f0);
#Define the simulation timespan
tspan = (0.f0, 20.f0); #hr
t = collect(range(0.0f0, stop=20.0f0, length=20_0));
# Define an ode problem
# note; the pk_model! is a function defined in 'pk_model.jl'
prob = ODEProblem(pk_model!, u0, tspan, collect(param));
#define parameters to be explored and their ranges
params = OrderedDict(:CL=>[0.1f0, 10.f0], :KA=>[0.01f0,10.0f0]);
param_names = collect(keys(params));
param_values = collect(values(params));
Random.seed!(1235)
# number of samples
N = 10_000;
# lower and upper bound for model parameters
lb = reduce(hcat,param_values)[1,:];
ub = reduce(hcat,param_values)[2,:];
# Generate design matrices based on QMC
sampler = SobolSample();
A,B = QuasiMonteCarlo.generate_design_matrices(N,lb,ub,sampler);
### GSA using chunks
ensemble_chunked = function (p)
# global parameterset = p
Np::Int64 = size(p,2)
chunks::Int64 = 10_0;
# println("Running ", Np, " cases.")
chunksize::Int64 = Np/chunks
# println("Using ", chunks, " chunk(s) of size ", chunksize, ".")
out = zeros(2,Np)
for k in 1:chunks
offset = Int((k - 1) * Np/chunks)
startindx = Int(offset + 1)
endindx = Int(k * Np/chunks)
# println("Starting chunk ", k, ", from ", startindx, " to ", endindx, " with offset ", offset, ".")
pchunk = p[:,startindx:endindx]
# function prob_func(prob,i,repeat)
# remake(prob; p=pchunk[:,i])
# end
function prob_func(prob,i,repeat)
p_pop = deepcopy(param);
p_pop[param_names] =pchunk[:,i];
remake(prob;p=collect(p_pop))
end
ensemble_prob = EnsembleProblem(prob,prob_func=prob_func)
ensemble_sol = solve(ensemble_prob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend());saveat = t,trajectories=chunksize)
Threads.@threads for i in 1:chunksize
out[1,i+offset] = ensemble_sol[i][3,end]
out[2,i+offset] = maximum(ensemble_sol[i][2,:])
end
end
out
end
sobol_chunks = gsa(ensemble_chunked,Sobol(),A,B,batch=true);
save_object("chunks_GPU.jld2", (sobol_chunks=sobol_chunks, param_names=param_names))
sobol_chunks, param_names = load_object("chunks_GPU.jld2");