Fix GPUKvaerno3/5 nonsingular mass-matrix explicit stages

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
 

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

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)
 

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
 

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
 

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)

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

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
 
 # ============================================================================