Skip to content

Don't setup a JacVec if concrete_jac==true, and use non SciMLOperator W_prototype #2491

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Don't setup a JacVec if concrete_jac==true, and use non SciMLOperator W_prototype #2491

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
426483c
don't use jacvec with concrete_jac
Oct 10, 2024
8dba9c5
update comment
Oct 10, 2024
83cc7ce
use non SciMLOperator W_prototype
Oct 10, 2024
851e3aa
simplify logic slightly
Oct 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 8 additions & 10 deletions lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl
View file
Open in desktop
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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how would this not be an issue to drop?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if iter<=1 we already returned on the first line (I changed that a few months ago)

(!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,7 +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)) &&
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
Expand All @@ -716,10 +716,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)
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
# 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)
Expand All @@ -734,9 +734,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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how would this work?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

by using the Jacobian. if you scroll up to the w_operator ops, they all check for whether ja cacvec exists

end
else
J = if !IIP && DiffEqBase.has_jac(f)
Expand Down
Loading