Skip to content

Commit

Permalink
Adding Oscar Smith's bugfix for derivative-utils, to allow for use of…
Browse files Browse the repository at this point in the history
… Krylov_GMRES

with a concrete jacobian.
  • Loading branch information
LeeStrobel committed Oct 11, 2024
1 parent df11489 commit ac2f970
Showing 1 changed file with 9 additions and 13 deletions.
22 changes: 9 additions & 13 deletions lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -391,7 +391,7 @@ function do_newJW(integrator, alg, nlsolver, repeat_step)::NTuple{2, Bool}
# TODO: add `isJcurrent` support for Rosenbrock solvers
if !isnewton(nlsolver)
isfreshJ = !(integrator.alg isa CompositeAlgorithm) &&
(integrator.iter > 1 && errorfail && !integrator.u_modified)
(!integrator.u_modified)
return !isfreshJ, true
end
isfirstcall(nlsolver) && return true, true
Expand Down Expand Up @@ -684,9 +684,9 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
# TODO - if jvp given, make it SciMLOperators.FunctionOperator
# TODO - make mass matrix a SciMLOperator so it can be updated with time. Default to IdentityOperator
islin, isode = islinearfunction(f, alg)
if isdefined(f, :W_prototype) && (f.W_prototype isa AbstractSciMLOperator)
# We use W_prototype when it is provided as a SciMLOperator, and in this case we require jac_prototype to be a SciMLOperator too.
if !(f.jac_prototype isa AbstractSciMLOperator)
if isdefined(f, :W_prototype) && !isnothing(f.W_prototype)
# If W_prototype is a SciMLOperator, we require jac_prototype to be a SciMLOperator too.
if f.W_prototype isa AbstractSciMLOperator && !(f.jac_prototype isa AbstractSciMLOperator)
error("SciMLOperator for W_prototype only supported when jac_prototype is a SciMLOperator, but got $(typeof(f.jac_prototype))")
end
W = f.W_prototype
Expand All @@ -703,8 +703,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
# If factorization, then just use the jac_prototype
J = similar(f.jac_prototype)
W = similar(J)
elseif (IIP && (concrete_jac(alg) === nothing || !concrete_jac(alg)) &&
alg.linsolve !== nothing &&
elseif (IIP && (concrete_jac(alg) != true) && alg.linsolve !== nothing &&
!LinearSolve.needs_concrete_A(alg.linsolve))
# If the user has chosen GMRES but no sparse Jacobian, assume that the dense
# Jacobian is a bad idea and create a fully matrix-free solver. This can
Expand All @@ -716,11 +715,10 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
J = jacvec
W = WOperator{IIP}(f.mass_matrix, dt, J, u, jacvec)
elseif alg.linsolve !== nothing && !LinearSolve.needs_concrete_A(alg.linsolve) ||
concrete_jac(alg) !== nothing && concrete_jac(alg)
# The linear solver does not need a concrete Jacobian, but the user has
concrete_jac(alg) == true
# The linear solver does not need a concrete Jacobian, but the user has
# asked for one. This will happen when the Jacobian is used in the preconditioner
# Thus setup JacVec and a concrete J, using sparsity when possible
_f = islin ? (isode ? f.f : f.f1.f) : f
# or when jvp computation is expensive. Therefore, use concrete J, and sparsity when possible _f = islin ? (isode ? f.f : f.f1.f) : f
J = if f.jac_prototype === nothing
ArrayInterface.undefmatrix(u)
else
Expand All @@ -734,9 +732,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
else
(u, p, t) -> _f(u, p, t)
end
jacvec = JacVec(__f, copy(u), p, t;
autodiff = alg_autodiff(alg), tag = OrdinaryDiffEqTag())
WOperator{IIP}(f.mass_matrix, dt, J, u, jacvec)
WOperator{IIP}(f.mass_matrix, dt, J, u, nothing)
end
else
J = if !IIP && DiffEqBase.has_jac(f)
Expand Down

0 comments on commit ac2f970

Please sign in to comment.