diff --git a/src/ensemblegpukernel/gpukernel_algorithms.jl b/src/ensemblegpukernel/gpukernel_algorithms.jl index 6f177384..7b013bc7 100644 --- a/src/ensemblegpukernel/gpukernel_algorithms.jl +++ b/src/ensemblegpukernel/gpukernel_algorithms.jl @@ -185,6 +185,9 @@ struct GPURodas5P{AD} <: GPUODEImplicitAlgorithm{AD} end Third-order ESDIRK method specialized for stiff `EnsembleGPUKernel` ODE solves. +Supports nonsingular mass matrices (`M * u' = f(u,p,t)`). Singular mass matrices +(DAEs) are rejected; use `GPURosenbrock23`, `GPURodas4`, or `GPURodas5P` instead. + # Keyword Arguments - `autodiff`: whether automatic differentiation is used for derivative generation. @@ -201,7 +204,8 @@ A `GPUKvaerno3` algorithm selector. # Throws -`solve` throws when the problem does not provide GPU-compatible derivative information. +`solve` throws when the problem does not provide GPU-compatible derivative information, +or when the mass matrix is singular. """ struct GPUKvaerno3{AD} <: GPUODEImplicitAlgorithm{AD} end @@ -210,6 +214,9 @@ struct GPUKvaerno3{AD} <: GPUODEImplicitAlgorithm{AD} end Fifth-order ESDIRK method specialized for stiff `EnsembleGPUKernel` ODE solves. +Supports nonsingular mass matrices (`M * u' = f(u,p,t)`). Singular mass matrices +(DAEs) are rejected; use `GPURosenbrock23`, `GPURodas4`, or `GPURodas5P` instead. + # Keyword Arguments - `autodiff`: whether automatic differentiation is used for derivative generation. @@ -226,7 +233,8 @@ A `GPUKvaerno5` algorithm selector. # Throws -`solve` throws when the problem does not provide GPU-compatible derivative information. +`solve` throws when the problem does not provide GPU-compatible derivative information, +or when the mass matrix is singular. """ struct GPUKvaerno5{AD} <: GPUODEImplicitAlgorithm{AD} end diff --git a/src/ensemblegpukernel/linalg/linsolve.jl b/src/ensemblegpukernel/linalg/linsolve.jl index 89b0bd64..d8f934db 100644 --- a/src/ensemblegpukernel/linalg/linsolve.jl +++ b/src/ensemblegpukernel/linalg/linsolve.jl @@ -106,3 +106,29 @@ end LUp.U \ (LUp.L \ $(length(Sb) > 1 ? :(b[LUp.p, :]) : :(b[LUp.p]))) end end + +const _KVAERNO_SINGULAR_MASS_MSG = "GPUKvaerno3 and GPUKvaerno5 do not support DAEs (singular mass matrices). Use GPURosenbrock23, GPURodas4, or GPURodas5P instead." + +@inline function _check_kvaerno_mass_matrix(alg, mass_matrix) + if alg isa Union{GPUKvaerno3, GPUKvaerno5} && _is_singular_mass_matrix(mass_matrix) + throw(ArgumentError(_KVAERNO_SINGULAR_MASS_MSG)) + end + return nothing +end + +# Includes UniformScaling{Bool}: `I` (λ=true) is nonsingular; `false*I` (λ=false) is singular. +@inline _is_singular_mass_matrix(M::LinearAlgebra.UniformScaling) = iszero(M.λ) +@inline _is_singular_mass_matrix(M) = iszero(det(M)) + +# Explicit first-stage / FSAL derivative for ESDIRK: solve M * du = f(u,p,t). +# Singular M is rejected on the host by `_check_kvaerno_mass_matrix` before launch. +@inline function _kvaerno_explicit_du(f, u, p, t) + return _mass_matrix_solve_du(f.mass_matrix, f(u, p, t)) +end + +@inline function _mass_matrix_solve_du(M::LinearAlgebra.UniformScaling, rhs) + return rhs / M.λ +end +@inline function _mass_matrix_solve_du(M::StaticMatrix, rhs) + return linear_solve(M, rhs) +end diff --git a/src/ensemblegpukernel/lowerlevel_solve.jl b/src/ensemblegpukernel/lowerlevel_solve.jl index b2812956..ad1e85db 100644 --- a/src/ensemblegpukernel/lowerlevel_solve.jl +++ b/src/ensemblegpukernel/lowerlevel_solve.jl @@ -88,6 +88,8 @@ function vectorized_solve( backend = get_backend(probs) backend = maybe_prefer_blocks(backend) + _check_kvaerno_mass_matrix(alg, prob.f.mass_matrix) + prob = convert(ImmutableODEProblem, prob) dt = convert(eltype(prob.tspan), dt) saveat_converted = nothing @@ -277,6 +279,8 @@ function vectorized_asolve( backend = get_backend(probs) backend = maybe_prefer_blocks(backend) + _check_kvaerno_mass_matrix(alg, prob.f.mass_matrix) + # Get the time type from the problem Tt = eltype(prob.tspan) diff --git a/src/ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl b/src/ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl index ba4eca59..c7a2dd64 100644 --- a/src/ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl +++ b/src/ensemblegpukernel/perform_step/gpu_kvaerno3_perform_step.jl @@ -25,7 +25,7 @@ end if integ.u_modified - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) integ.u_modified = false else @inbounds k1 = integ.k1 @@ -79,7 +79,7 @@ k2 = z₄ ./ dt @inbounds begin # Necessary for interpolation - integ.k1 = f(integ.u, p, t) + integ.k1 = _kvaerno_explicit_du(f, integ.u, p, t) integ.k2 = k2 end @@ -113,7 +113,7 @@ end α32 = integ.tab if integ.u_modified - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) integ.u_modified = false else @inbounds k1 = integ.k1 @@ -134,7 +134,7 @@ end # FSAL Step 1 - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) z₁ = dt * k1 diff --git a/src/ensemblegpukernel/perform_step/gpu_kvaerno5_perform_step.jl b/src/ensemblegpukernel/perform_step/gpu_kvaerno5_perform_step.jl index 56de80b6..9c885839 100644 --- a/src/ensemblegpukernel/perform_step/gpu_kvaerno5_perform_step.jl +++ b/src/ensemblegpukernel/perform_step/gpu_kvaerno5_perform_step.jl @@ -27,7 +27,7 @@ end if integ.u_modified - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) integ.u_modified = false else @inbounds k1 = integ.k1 @@ -110,7 +110,7 @@ k2 = z₇ ./ dt @inbounds begin # Necessary for interpolation - integ.k1 = f(integ.u, p, t) + integ.k1 = _kvaerno_explicit_du(f, integ.u, p, t) integ.k2 = k2 end @@ -146,7 +146,7 @@ end @unpack α31, α32, α41, α42, α43, α51, α52, α53, α61, α62, α63 = integ.tab if integ.u_modified - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) integ.u_modified = false else @inbounds k1 = integ.k1 @@ -167,7 +167,7 @@ end # FSAL Step 1 - k1 = f(uprev, p, t) + k1 = _kvaerno_explicit_du(f, uprev, p, t) z₁ = dt * k1 diff --git a/src/solve.jl b/src/solve.jl index a700d788..9f15063e 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -406,6 +406,8 @@ function batch_solve_up_kernel( )... ) + _check_kvaerno_mass_matrix(alg, ensembleprob.prob.f.mass_matrix) + dev = ensemblealg.dev probs = _kernel_transfer(dev, adapted_probs) @@ -422,6 +424,7 @@ function batch_solve_up_kernel( kwargs..., callback = _callback ) end + solus = _kernel_transfer(CPU(), us) solts = _kernel_transfer(CPU(), ts) return (solts, solus) diff --git a/test/gpu_kernel_de/stiff_ode/gpu_ode_mass_matrix.jl b/test/gpu_kernel_de/stiff_ode/gpu_ode_mass_matrix.jl index 3e089675..c5ccdc99 100644 --- a/test/gpu_kernel_de/stiff_ode/gpu_ode_mass_matrix.jl +++ b/test/gpu_kernel_de/stiff_ode/gpu_ode_mass_matrix.jl @@ -67,3 +67,35 @@ sol = solve( expected = iszero(initial) ? initial : exp(-1.0f0 / mass.λ) @test all(s -> isapprox(s.u[end][1], expected; atol = 2.0f-6), sol.u) end + +# Analytic check for M * u' = -u with nonsingular M = 2I (issue #530). +# Expected endpoint is u0 * exp(-t/2). Float32 so Metal/oneAPI backends can run it. +@testset "Kvaerno nonsingular mass matrix" begin + f(u, p, t) = -u + jac(u, p, t) = -one(SMatrix{2, 2, Float32}) + expected = exp(-0.5f0) * SVector(1.0f0, 2.0f0) + for alg in (GPUKvaerno3(), GPUKvaerno5()), adaptive in (false, true) + prob = ODEProblem( + ODEFunction{false}(f; jac, mass_matrix = 2.0f0 * I), + SVector(1.0f0, 2.0f0), (0.0f0, 1.0f0) + ) + sol = solve( + EnsembleProblem(prob; safetycopy = false), alg, + EnsembleGPUKernel(backend, 0.0); trajectories = 2, adaptive, + dt = 0.01f0, abstol = 1.0f-8, reltol = 1.0f-8, save_everystep = false + ) + @test sol.u[1].u[end] ≈ expected rtol = 1.0f-5 atol = 1.0f-5 + end +end + +@testset "Kvaerno rejects singular mass matrix" begin + f = ODEFunction{false}((u, p, t) -> -u; mass_matrix = 0.0f0 * I) + prob = ODEProblem(f, SVector(0.0f0), (0.0f0, 1.0f0)) + for alg in (GPUKvaerno3(), GPUKvaerno5()) + @test_throws ArgumentError solve( + EnsembleProblem(prob; safetycopy = false), alg, + EnsembleGPUKernel(backend, 0.0); trajectories = 2, adaptive = false, + dt = 0.01f0, save_everystep = false + ) + end +end diff --git a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl index 208df41e..09f1305c 100644 --- a/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl +++ b/test/gpu_kernel_de/stiff_ode/gpu_ode_modelingtoolkit_dae.jl @@ -92,27 +92,21 @@ end end @testset "GPUKvaerno3 DAE" begin - sol = solve( + @test_throws ArgumentError solve( monteprob, GPUKvaerno3(), EnsembleGPUKernel(backend), trajectories = 2, dt = 0.001f0, adaptive = false ) - @test length(sol.u) == 2 - @test !any(isnan, sol.u[1].u[end]) - @test abs(sol.u[1].u[end][1] + sol.u[1].u[end][2] - 1.0f0) < 0.01f0 end @testset "GPUKvaerno5 DAE" begin - sol = solve( + @test_throws ArgumentError solve( monteprob, GPUKvaerno5(), EnsembleGPUKernel(backend), trajectories = 2, dt = 0.001f0, adaptive = false ) - @test length(sol.u) == 2 - @test !any(isnan, sol.u[1].u[end]) - @test abs(sol.u[1].u[end][1] + sol.u[1].u[end][2] - 1.0f0) < 0.01f0 end # ============================================================================