From 538ea1e8ba98f7a0e84d5ef299387bae68a24296 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Wed, 8 Jul 2026 05:41:34 -0400 Subject: [PATCH] Document public API coverage Co-Authored-By: Chris Rackauckas --- docs/pages.jl | 1 + docs/src/manual/api.md | 21 +++ src/ensemblegpuarray/kernels.jl | 14 +- src/ensemblegpukernel/gpukernel_algorithms.jl | 152 ++++++++++++++---- 4 files changed, 156 insertions(+), 32 deletions(-) create mode 100644 docs/src/manual/api.md diff --git a/docs/pages.jl b/docs/pages.jl index 70c39559..b727deb5 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -32,5 +32,6 @@ pages = [ "manual/backends.md", "manual/optimal_trajectories.md", "manual/choosing_ensembler.md", + "manual/api.md", ], ] diff --git a/docs/src/manual/api.md b/docs/src/manual/api.md new file mode 100644 index 00000000..d343fb6d --- /dev/null +++ b/docs/src/manual/api.md @@ -0,0 +1,21 @@ +# API + +## Reexported Ensemble API + +DiffEqGPU reexports these ensemble and callback helper APIs from SciMLBase.jl and +DiffEqBase.jl for compatibility with the SciML ensemble interface: + + - `EnsembleProblem` + - `EnsembleSolution` + - `EnsembleSerial` + - `EnsembleThreads` + - `EnsembleDistributed` + - `BrownFullBasicInit` + - `CheckInit` + - `terminate!` + +## Lower-Level Algorithms + +```@docs +LinSolveGPUSplitFactorize +``` diff --git a/src/ensemblegpuarray/kernels.jl b/src/ensemblegpuarray/kernels.jl index d1d4925d..29460535 100644 --- a/src/ensemblegpuarray/kernels.jl +++ b/src/ensemblegpuarray/kernels.jl @@ -368,7 +368,19 @@ end ### GPU Factorization """ -A parameter-parallel `SciMLLinearSolveAlgorithm`. + LinSolveGPUSplitFactorize() + LinSolveGPUSplitFactorize(len, nfacts) + +A parameter-parallel `SciMLLinearSolveAlgorithm` for applying pre-factorized +per-trajectory linear systems on a KernelAbstractions backend. + +# Arguments + + - `len`: the size of each factored linear system. + - `nfacts`: the number of factorizations stored in the batched factorization array. + +Most users do not need to construct this directly; `EnsembleGPUArray` installs it for +compatible stiff ensemble solves. """ struct LinSolveGPUSplitFactorize <: LinearSolve.SciMLLinearSolveAlgorithm len::Int diff --git a/src/ensemblegpukernel/gpukernel_algorithms.jl b/src/ensemblegpukernel/gpukernel_algorithms.jl index 5d5dbfda..47f241c9 100644 --- a/src/ensemblegpukernel/gpukernel_algorithms.jl +++ b/src/ensemblegpukernel/gpukernel_algorithms.jl @@ -1,65 +1,141 @@ """ -GPUTsit5() + GPUTsit5() -A specialized implementation of the 5th order `Tsit5` method specifically for kernel -generation with EnsembleGPUKernel. For a similar CPU implementation, see -SimpleATsit5 from SimpleDiffEq.jl. +Fifth-order Tsitouras Runge-Kutta method specialized for `EnsembleGPUKernel` ODE +solves. + +Use `GPUTsit5` as the ODE algorithm when solving an `EnsembleProblem` with +`EnsembleGPUKernel`: + +```julia +solve( + ensemble_prob, GPUTsit5(), EnsembleGPUKernel(backend); + trajectories = 10_000, adaptive = false, dt = 0.1f0 +) +``` + +`GPUTsit5` supports the `EnsembleGPUKernel` restrictions, including out-of-place ODE +functions over GPU-compatible static state containers. For a similar CPU implementation, +see `SimpleATsit5` from SimpleDiffEq.jl. """ struct GPUTsit5 <: GPUODEAlgorithm end """ -GPUVern7() + GPUVern7() -A specialized implementation of the 7th order `Vern7` method specifically for kernel -generation with EnsembleGPUKernel. +Seventh-order Verner Runge-Kutta method specialized for `EnsembleGPUKernel` ODE +solves. + +Use `GPUVern7` for non-stiff ODE ensemble solves that satisfy the `EnsembleGPUKernel` +kernel-generation restrictions: + +```julia +solve( + ensemble_prob, GPUVern7(), EnsembleGPUKernel(backend); + trajectories = 10_000, adaptive = false, dt = 0.1f0 +) +``` """ struct GPUVern7 <: GPUODEAlgorithm end """ -GPUVern9() + GPUVern9() + +Ninth-order Verner Runge-Kutta method specialized for `EnsembleGPUKernel` ODE solves. -A specialized implementation of the 9th order `Vern9` method specifically for kernel -generation with EnsembleGPUKernel. +Use `GPUVern9` for high-accuracy non-stiff ODE ensemble solves that satisfy the +`EnsembleGPUKernel` kernel-generation restrictions: + +```julia +solve( + ensemble_prob, GPUVern9(), EnsembleGPUKernel(backend); + trajectories = 10_000, adaptive = false, dt = 0.1f0 +) +``` """ struct GPUVern9 <: GPUODEAlgorithm end """ -GPURosenbrock23() + GPURosenbrock23(; autodiff = Val{true}()) + +Second/third-order Rosenbrock-W method specialized for stiff `EnsembleGPUKernel` ODE +solves. + +# Keyword Arguments -A specialized implementation of the W-method `Rosenbrock23` method specifically for kernel -generation with EnsembleGPUKernel. + - `autodiff`: whether automatic differentiation is used for derivative generation. + Pass `Val{false}()` when providing the required derivatives manually. + +```julia +solve( + ensemble_prob, GPURosenbrock23(), EnsembleGPUKernel(backend); + trajectories = 10_000 +) +``` """ struct GPURosenbrock23{AD} <: GPUODEImplicitAlgorithm{AD} end """ -GPURodas4() + GPURodas4(; autodiff = Val{true}()) + +Fourth-order Rosenbrock method specialized for stiff `EnsembleGPUKernel` ODE solves. + +# Keyword Arguments -A specialized implementation of the `Rodas4` method specifically for kernel -generation with EnsembleGPUKernel. + - `autodiff`: whether automatic differentiation is used for derivative generation. + Pass `Val{false}()` when providing the required derivatives manually. + +```julia +solve(ensemble_prob, GPURodas4(), EnsembleGPUKernel(backend); trajectories = 10_000) +``` """ struct GPURodas4{AD} <: GPUODEImplicitAlgorithm{AD} end """ -GPURodas5P() + GPURodas5P(; autodiff = Val{true}()) + +Fifth-order Rosenbrock method specialized for stiff `EnsembleGPUKernel` ODE solves. + +# Keyword Arguments -A specialized implementation of the `Rodas5P` method specifically for kernel -generation with EnsembleGPUKernel. + - `autodiff`: whether automatic differentiation is used for derivative generation. + Pass `Val{false}()` when providing the required derivatives manually. + +```julia +solve(ensemble_prob, GPURodas5P(), EnsembleGPUKernel(backend); trajectories = 10_000) +``` """ struct GPURodas5P{AD} <: GPUODEImplicitAlgorithm{AD} end """ -GPUKvaerno3() + GPUKvaerno3(; autodiff = Val{true}()) + +Third-order ESDIRK method specialized for stiff `EnsembleGPUKernel` ODE solves. + +# Keyword Arguments -A specialized implementation of the `Kvaerno3` method specifically for kernel -generation with EnsembleGPUKernel. + - `autodiff`: whether automatic differentiation is used for derivative generation. + Pass `Val{false}()` when providing the required derivatives manually. + +```julia +solve(ensemble_prob, GPUKvaerno3(), EnsembleGPUKernel(backend); trajectories = 10_000) +``` """ struct GPUKvaerno3{AD} <: GPUODEImplicitAlgorithm{AD} end """ -GPUKvaerno5() + GPUKvaerno5(; autodiff = Val{true}()) + +Fifth-order ESDIRK method specialized for stiff `EnsembleGPUKernel` ODE solves. + +# Keyword Arguments -A specialized implementation of the `Kvaerno5` method specifically for kernel -generation with EnsembleGPUKernel. + - `autodiff`: whether automatic differentiation is used for derivative generation. + Pass `Val{false}()` when providing the required derivatives manually. + +```julia +solve(ensemble_prob, GPUKvaerno5(), EnsembleGPUKernel(backend); trajectories = 10_000) +``` """ struct GPUKvaerno5{AD} <: GPUODEImplicitAlgorithm{AD} end @@ -72,17 +148,31 @@ for Alg in [:GPURosenbrock23, :GPURodas4, :GPURodas5P, :GPUKvaerno3, :GPUKvaerno end """ -GPUEM() + GPUEM() + +Euler-Maruyama method with weak order 1.0 specialized for `EnsembleGPUKernel` SDE +solves. -A specialized implementation of the Euler-Maruyama `GPUEM` method with weak order 1.0. Made specifically for kernel -generation with EnsembleGPUKernel. +```julia +solve( + ensemble_prob, GPUEM(), EnsembleGPUKernel(backend); + trajectories = 10_000, adaptive = false, dt = 0.1f0 +) +``` """ struct GPUEM <: GPUSDEAlgorithm end """ -GPUSIEA() + GPUSIEA() + +Weak order 2.0 SIEA method for Ito SDEs specialized for `EnsembleGPUKernel` SDE +solves. -A specialized implementation of the weak order 2.0 for Ito SDEs `GPUSIEA` method specifically for kernel -generation with EnsembleGPUKernel. +```julia +solve( + ensemble_prob, GPUSIEA(), EnsembleGPUKernel(backend); + trajectories = 10_000, adaptive = false, dt = 0.1f0 +) +``` """ struct GPUSIEA <: GPUSDEAlgorithm end