Skip to content

Commit 39667c5

Browse files
Wrap Jacobian with both dense and sparse signatures when sparsity pattern present
When an ODEFunction has a sparse sparsity pattern but no jac_prototype, the solver's build_J_W may create a SparseMatrixCSC J from the sparsity pattern (when using AutoSparse). The previous code only wrapped the Jacobian with a dense Matrix{Float64} signature, causing "No matching function wrapper was found!" at solve time. Now detects non-dense sparsity patterns and creates a FunctionWrappersWrapper with both dense and sparse matrix signatures so the Jacobian call works regardless of which matrix type the solver allocates. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
1 parent 92bfd6b commit 39667c5

File tree

1 file changed

+23
-4
lines changed

1 file changed

+23
-4
lines changed

src/solve.jl

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -810,12 +810,31 @@ function promote_f(
810810
if f.tgrad !== nothing && !(f.tgrad isa FunctionWrappersWrappers.FunctionWrappersWrapper)
811811
f = @set f.tgrad = wrapfun_jac_iip(f.tgrad, (u0, u0, p, t))
812812
end
813-
# Wrap the Jacobian if present, so its type is also erased
813+
# Wrap the Jacobian if present, so its type is also erased.
814+
# Include both dense and sparse matrix signatures when the function
815+
# has a sparsity pattern, since the solver may use either depending on
816+
# the autodiff configuration (AutoSparse creates sparse J from sparsity).
814817
if f.jac !== nothing && !(f.jac isa FunctionWrappersWrappers.FunctionWrappersWrapper)
815818
n = length(u0)
816-
J_proto = f.jac_prototype !== nothing ? similar(f.jac_prototype, uElType) :
817-
zeros(uElType, n, n)
818-
f = @set f.jac = wrapfun_jac_iip(f.jac, (J_proto, u0, p, t))
819+
if f.jac_prototype !== nothing
820+
J_proto = similar(f.jac_prototype, uElType)
821+
f = @set f.jac = wrapfun_jac_iip(f.jac, (J_proto, u0, p, t))
822+
elseif isdefined(f, :sparsity) && f.sparsity isa AbstractMatrix &&
823+
typeof(similar(f.sparsity, uElType)) !== Matrix{uElType}
824+
# The sparsity pattern is a non-dense matrix (e.g. SparseMatrixCSC).
825+
# The solver may call the Jacobian with either a dense or sparse matrix
826+
# depending on the autodiff config, so wrap for both signatures.
827+
J_dense = zeros(uElType, n, n)
828+
J_sparse = similar(f.sparsity, uElType)
829+
f = @set f.jac = FunctionWrappersWrappers.FunctionWrappersWrapper(
830+
Void(f.jac),
831+
(typeof((J_dense, u0, p, t)), typeof((J_sparse, u0, p, t))),
832+
(Nothing, Nothing)
833+
)
834+
else
835+
J_proto = zeros(uElType, n, n)
836+
f = @set f.jac = wrapfun_jac_iip(f.jac, (J_proto, u0, p, t))
837+
end
819838
end
820839
return unwrapped_f(f, wrapfun_iip(f.f, (u0, u0, p, t), Val(CS)))
821840
else

0 commit comments

Comments
 (0)