SciML / DiffEqGPU.jl

GPU-acceleration routines for DifferentialEquations.jl and the broader SciML scientific machine learning ecosystem
https://docs.sciml.ai/DiffEqGPU/stable/
MIT License
272 stars 27 forks source link

Add Rosenbrock23 with EnsembleGPUKernel #252

Closed utkarsh530 closed 1 year ago

utkarsh530 commented 1 year ago

Example:

using DiffEqGPU, StaticArrays, LinearAlgebra, DiffEqBase

using CUDA

backend = CUDABackend()

function lorenz(u, p, t)
    σ = p[1]
    ρ = p[2]
    β = p[3]
    du1 = σ * (u[2] - u[1])
    du2 = u[1] * (ρ - u[3]) - u[2]
    du3 = u[1] * u[2] - β * u[3]
    return SVector{3}(du1, du2, du3)
end

function lorenz_jac(u, p, t)
    σ = p[1]
    ρ = p[2]
    β = p[3]
    x = u[1]
    y = u[2]
    z = u[3]
    J11 = -σ
    J21 = ρ - z
    J31 = y
    J12 = σ
    J22 = -1
    J32 = x
    J13 = 0
    J23 = -x
    J33 = -β
    return SMatrix{3,3}(J11,J21,J31,J12,J22,J32,J13,J23,J33)
end

function lorenz_tgrad(u, p, t)
    return SVector{3, eltype(u)}(0.0, 0.0, 0.0)
end

u0 = @SVector [1.0f0; 0.0f0; 0.0f0]
tspan = (0.0f0, 10.0f0)
p = @SVector [10.0f0, 28.0f0, 8 / 3.0f0]

func = ODEFunction(lorenz, jac = lorenz_jac, tgrad = lorenz_tgrad)
prob = ODEProblem{false}(func, u0, tspan, p)

prob_func = (prob, i, repeat) -> remake(prob, p = p)
monteprob = EnsembleProblem(prob, prob_func = prob_func, safetycopy = false)

sol = solve(monteprob, GPURosenbrock23(), EnsembleGPUKernel(backend, 0.0), trajectories = 10,
                      adaptive = false, dt = 0.01f0)
utkarsh530 commented 1 year ago

Requires #241