Tutorial utilising MTK to highlight the utility of DiffEqGPU.jl

[H-Sax]

Is your feature request related to a problem? Please describe.

It would be nice to see a tutorial where MTK examples show how to create a GPU compatible form to utilise the DiffEqGPU solvers

[ChrisRackauckas]

The simple thing:

using OrdinaryDiffEq, ModelingToolkit
using SymbolicIndexingInterface
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters σ ρ β
@variables x(t) y(t) z(t)

eqs = [D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@mtkbuild sys = ODESystem(eqs, t)

u0 = [D(x) => 2f0,
    x => 1f0,
    y => 0f0,
    z => 0f0]

p = [σ => 28f0,
    ρ => 10f0,
    β => 8f0 / 3f0]

tspan = (0f0, 100f0)
prob = ODEProblem(sys, u0, tspan, p, jac = true)

prob_func = function (prob, i, repeat)
	prob = deepcopy(prob)
	prob.p.tunable .= rand(Float32, 3) .* prob.p.tunable
	prob
end

monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
    saveat = 1.0f0);

Should just work. But you probably don't want to remake using such low level stuff and instead use the SymbolicIndexingInterface. For that you probably want the following:

probchanger = SymbolicIndexingInterface.setp_oop(prob, [σ, ρ, β]);
prob_func = (prob, i, repeat) -> (probchanger(prob, rand(Float32, 3) .* [28f0, 10f0, 8f0/3f0]); prob)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)
sol = solve(monteprob, Tsit5(), EnsembleThreads(), trajectories = 10_000, saveat = 1.0f0);

using DiffEqGPU, CUDA
sol = solve(monteprob, Tsit5(), EnsembleGPUArray(CUDA.CUDABackend()), trajectories = 10_000,
    saveat = 1.0f0);

@AayushSabharwal is going to improve setsym_oop which allows for mixing initial conditions and parameters in what's set to make it even easier. I still need to test this works in a GPU kernel too, you might need to flat eval of the RGF so it's cache-free.

[utkarsh530]

https://github.com/utkarsh530/GPUODEBenchmarks/blob/1a90256220a8a557b238e37cb3101ba7fd5319fc/GPU_ODE_Julia/src/GPU_ODE_Julia.jl#L5

I had this written up to compile with just SArrays