diff --git a/Project.toml b/Project.toml index f3f8715ee..611681111 100644 --- a/Project.toml +++ b/Project.toml @@ -158,7 +158,7 @@ PartitionedSolvers = "0.3" Pkg = "1.10" PrecompileTools = "1.2" Preferences = "1.4" -PureKLU = "1.0.1" +PureKLU = "1.1" PureUMFPACK = "0.1" Random = "1.10" RecursiveArrayTools = "3.37, 4" diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 4acc7e9ad..3d902d4fc 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -284,6 +284,18 @@ function SciMLBase.init(prob::DualAbstractLinearProblem, alg::SparspakFactorizat return __init(prob, alg, args...; kwargs...) end +# Duals only in b (A is primal): route PureKLU to a plain LinearCache and solve natively. +# PureKLU's mixed-type `ldiv!` (primal KLU factor against a Dual RHS, from PureKLUForwardDiffExt) +# keeps the factorization in Float64 and pushes the duals through the back-substitution, so +# there is no reason to build the split DualLinearCache here. This is type-stable: dispatch is +# purely on the problem +# subtype (b-dual / A-plain) and the alg type, so `init` always returns a `LinearCache` for +# this method. It also gets correct factorization reuse across b-only `reinit!`s for free, +# which the split path's `reinit!` does not (it always marks the inner cache fresh). +function SciMLBase.init(prob::DualBLinearProblem, alg::PureKLUFactorization, args...; kwargs...) + return __init(prob, alg, args...; kwargs...) +end + # NOTE: Removed the runtime conditional for DefaultLinearSolver that checked for # GenericLUFactorization. Now always use __dual_init for type stability. function SciMLBase.init(prob::DualAbstractLinearProblem, alg::DefaultLinearSolver, args...; kwargs...) diff --git a/test/Core/forwarddiff_overloads.jl b/test/Core/forwarddiff_overloads.jl index 978a4faa3..24dd6f5aa 100644 --- a/test/Core/forwarddiff_overloads.jl +++ b/test/Core/forwarddiff_overloads.jl @@ -280,6 +280,44 @@ plain_A = ForwardDiff.value.(A) prob = LinearProblem(sparse(plain_A), b) @test ≈(solve(prob, PureKLUFactorization()), plain_A \ b, rtol = 1.0e-9) +# Mixed-type ldiv!: a primal (Float64) KLU factorization backsolving a Dual RHS +# without promoting A. The factor stays Float64; duals ride through the +# back-substitution (value + each partial column solved in one multi-RHS solve). +@testset "PureKLU primal factor \\ Dual RHS (mixed ldiv!)" begin + Asp = sparse(2.0I, 5, 5) + sparse(plain_A[1, 1] * 0.0I, 5, 5) + for i in 1:4 + Asp[i, i + 1] = 0.3 + Asp[i + 1, i] = 0.2 + end + for nchunk in (1, 2, 3) + bd = [ + ForwardDiff.Dual{Nothing, Float64, nchunk}( + Float64(i), ForwardDiff.Partials(ntuple(k -> sin(i + k), nchunk)) + ) for i in 1:5 + ] + cache = LinearSolve.__init(LinearProblem(Asp, bd), PureKLUFactorization()) + @test eltype(cache.A) == Float64 # A not promoted + u = solve!(cache).u + uref = Matrix{eltype(bd)}(Asp) \ bd + @test isapprox(ForwardDiff.value.(u), ForwardDiff.value.(uref); rtol = 1.0e-10) + @test all( + isapprox( + ForwardDiff.partials(u[i], j), ForwardDiff.partials(uref[i], j); + rtol = 1.0e-8, atol = 1.0e-12 + ) for i in 1:5, j in 1:nchunk + ) + end + + # Duals-only-in-b is routed to a plain LinearCache (native solve), not the split + # DualLinearCache, and that routing is type-stable. + Adual, bdual = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)]) + plain_Asp = sparse(ForwardDiff.value.(Adual)) + bprob = LinearProblem(plain_Asp, bdual) + @test init(bprob, PureKLUFactorization()) isa LinearSolve.LinearCache + @test (@inferred init(bprob, PureKLUFactorization())) isa LinearSolve.LinearCache + @test ≈(solve(bprob, PureKLUFactorization()), Adual \ bdual, rtol = 1.0e-9) +end + A, b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)]) prob = LinearProblem(sparse(A), sparse(b))