diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index eac8dc1e9..4baef6c55 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -75,7 +75,7 @@ include("forward_diff.jl") # public for NonlinearSolve.jl and subpackages to use @compat(public, (InternalAPI, supports_line_search, supports_trust_region, set_du!)) @compat(public, (construct_linear_solver, needs_square_A, needs_concrete_A)) -@compat(public, (construct_jacobian_cache,)) +@compat(public, (construct_jacobian_cache, reused_jacobian)) @compat(public, (assert_extension_supported_termination_condition, construct_extension_function_wrapper, construct_extension_jac)) diff --git a/lib/NonlinearSolveBase/src/jacobian.jl b/lib/NonlinearSolveBase/src/jacobian.jl index 04a61789f..46c2f0c4a 100644 --- a/lib/NonlinearSolveBase/src/jacobian.jl +++ b/lib/NonlinearSolveBase/src/jacobian.jl @@ -87,7 +87,7 @@ function construct_jacobian_cache( end end - return JacobianCache(J, f, fu, u, p, stats, autodiff, di_extras) + return JacobianCache(J, f, fu, p, stats, autodiff, di_extras) end function construct_jacobian_cache( @@ -96,7 +96,7 @@ function construct_jacobian_cache( linsolve = missing ) if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f) - return JacobianCache(u, f, fu, u, p, stats, autodiff, nothing) + return JacobianCache(fu, f, fu, p, stats, autodiff, nothing) end if autodiff === nothing throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \ @@ -107,74 +107,76 @@ function construct_jacobian_cache( @assert !(autodiff isa AutoSparse) "`autodiff` cannot be `AutoSparse` for scalar \ nonlinear problems." di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p)) - return JacobianCache(u, f, fu, u, p, stats, autodiff, di_extras) + return JacobianCache(u, f, fu, p, stats, autodiff, di_extras) end @concrete mutable struct JacobianCache <: AbstractJacobianCache J f <: NonlinearFunction fu - u p stats::NLStats autodiff di_extras end -function InternalAPI.reinit!(cache::JacobianCache; p = cache.p, u0 = cache.u, kwargs...) - cache.u = u0 +function InternalAPI.reinit!(cache::JacobianCache; p = cache.p, kwargs...) cache.p = p end -# Core Computation -(cache::JacobianCache)(u) = cache(cache.J, u, cache.p) -function (cache::JacobianCache{<:JacobianOperator})(::Nothing) - return StatefulJacobianOperator(cache.J, cache.u, cache.p) -end -(cache::JacobianCache)(::Nothing) = cache.J +# Deprecations +(cache::JacobianCache{<:Number})(::Nothing) = error("Please report a bug to NonlinearSolve.jl") +(cache::JacobianCache{<:JacobianOperator})(::Nothing) = error("Please report a bug to NonlinearSolve.jl") +(cache::JacobianCache)(::Nothing) = error("Please report a bug to NonlinearSolve.jl") -## Operator -function (cache::JacobianCache{<:JacobianOperator})(J::JacobianOperator, u, p = cache.p) - return StatefulJacobianOperator(J, u, p) -end +reused_jacobian(cache::JacobianCache, u) = cache.J +reused_jacobian(cache::JacobianCache{<:JacobianOperator}, u) = StatefulJacobianOperator(cache.J, u, cache.p) +# Core Computation ## Numbers -function (cache::JacobianCache{<:Number})(::Number, u, p = cache.p) +function (cache::JacobianCache{<:Number})(u) cache.stats.njacs += 1 - cache.J = if SciMLBase.has_jac(cache.f) - cache.f.jac(u, p) - elseif SciMLBase.has_vjp(cache.f) - cache.f.vjp(one(u), u, p) - elseif SciMLBase.has_jvp(cache.f) - cache.f.jvp(one(u), u, p) + + (; f, J, p) = cache + cache.J = if SciMLBase.has_jac(f) + f.jac(u, p) + elseif SciMLBase.has_vjp(f) + f.vjp(one(u), u, p) + elseif SciMLBase.has_jvp(f) + f.jvp(one(u), u, p) else - DI.derivative(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) + DI.derivative(f, cache.di_extras, cache.autodiff, u, Constant(p)) end return cache.J end ## Actually Compute the Jacobian -function (cache::JacobianCache)(J::Union{AbstractMatrix, Nothing}, u, p = cache.p) +function (cache::JacobianCache)(u) cache.stats.njacs += 1 - if SciMLBase.isinplace(cache.f) - if SciMLBase.has_jac(cache.f) - cache.f.jac(J, u, p) + (; f, J, p) = cache + if SciMLBase.isinplace(f) + if SciMLBase.has_jac(f) + f.jac(J, u, p) else DI.jacobian!( - cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p) + f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p) ) end return J else if SciMLBase.has_jac(cache.f) - cache.J = cache.f.jac(u, p) + cache.J = f.jac(u, p) else - cache.J = DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) + cache.J = DI.jacobian(f, cache.di_extras, cache.autodiff, u, Constant(p)) end return cache.J end end +function (cache::JacobianCache{<:JacobianOperator})(u) + return StatefulJacobianOperator(cache.J, u, cache.p) +end + # Sparse Automatic Differentiation function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) if f.sparsity === nothing diff --git a/lib/NonlinearSolveBase/src/tracing.jl b/lib/NonlinearSolveBase/src/tracing.jl index c7ae3f542..6caba9f4c 100644 --- a/lib/NonlinearSolveBase/src/tracing.jl +++ b/lib/NonlinearSolveBase/src/tracing.jl @@ -222,12 +222,13 @@ function update_trace!(cache, α = true; uses_jac_inverse = Val(false)) trace === missing && return nothing J = Utils.safe_getproperty(cache, Val(:J)) + du = SciMLBase.get_du(cache) if J === missing update_trace!( - trace, cache.nsteps + 1, get_u(cache), get_fu(cache), nothing, cache.du, α + trace, cache.nsteps + 1, get_u(cache), get_fu(cache), nothing, du, α ) else J = uses_jac_inverse isa Val{true} ? Utils.Pinv(cache.J) : cache.J - update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), J, cache.du, α) + update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), J, du, α) end end diff --git a/lib/NonlinearSolveBase/src/wrappers.jl b/lib/NonlinearSolveBase/src/wrappers.jl index 0b96836ad..f24508835 100644 --- a/lib/NonlinearSolveBase/src/wrappers.jl +++ b/lib/NonlinearSolveBase/src/wrappers.jl @@ -101,13 +101,10 @@ function construct_extension_jac( J_no_scalar = can_handle_scalar isa Val{false} && prob.u0 isa Number ? @closure(u->[Jₚ(u[1])]) : Jₚ - J_final = if can_handle_oop isa Val{false} && !SciMLBase.isinplace(prob) - @closure((J, u)->copyto!(J, J_no_scalar(u))) - else - J_no_scalar - end + J_final(J, u) = copyto!(J, J_no_scalar(u)) + J_final(u) = J_no_scalar(u) initial_jacobian isa Val{false} && return J_final - return J_final, Jₚ(nothing) + return J_final, reused_jacobian(Jₚ, u0) end diff --git a/lib/NonlinearSolveFirstOrder/Project.toml b/lib/NonlinearSolveFirstOrder/Project.toml index b18cfa151..bda7c9f55 100644 --- a/lib/NonlinearSolveFirstOrder/Project.toml +++ b/lib/NonlinearSolveFirstOrder/Project.toml @@ -57,7 +57,7 @@ SciMLBase = "2.116" SciMLJacobianOperators = "0.1.0" Setfield = "1.1.1" SparseArrays = "1.10" -SparseConnectivityTracer = "1, 1" +SparseConnectivityTracer = "1" SparseMatrixColorings = "0.4.5" StableRNGs = "1" StaticArrays = "1.9.8" diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl index 0a6c009c6..dba01edfe 100644 --- a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -21,7 +21,7 @@ using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, Utils, InternalAPI, get_timer_output, @static_timeit, update_trace!, L2_NORM, NonlinearSolvePolyAlgorithm, NewtonDescent, DampedNewtonDescent, GeodesicAcceleration, - Dogleg, NonlinearSolveForwardDiffCache + Dogleg, NonlinearSolveForwardDiffCache, reused_jacobian using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode, NonlinearFunction, NonlinearLeastSquaresProblem, NonlinearProblem, NoSpecialize diff --git a/lib/NonlinearSolveFirstOrder/src/solve.jl b/lib/NonlinearSolveFirstOrder/src/solve.jl index a23e57ac2..874f5dbc4 100644 --- a/lib/NonlinearSolveFirstOrder/src/solve.jl +++ b/lib/NonlinearSolveFirstOrder/src/solve.jl @@ -55,8 +55,6 @@ end u u_cache p - du # Aliased to `get_du(descent_cache)` - J # Aliased to `jac_cache.J` alg <: GeneralizedFirstOrderAlgorithm prob <: AbstractNonlinearProblem globalization <: Union{Val{:LineSearch}, Val{:TrustRegion}, Val{:None}} @@ -91,6 +89,13 @@ end initializealg end +function SciMLBase.get_du(cache::GeneralizedFirstOrderAlgorithmCache) + SciMLBase.get_du(cache.descent_cache) +end +function NonlinearSolveBase.set_du!(cache::GeneralizedFirstOrderAlgorithmCache, δu) + NonlinearSolveBase.set_du!(cache.descent_cache, δu) +end + function InternalAPI.reinit_self!( cache::GeneralizedFirstOrderAlgorithmCache, args...; p = cache.p, u0 = cache.u, alias_u0::Bool = hasproperty(cache, :alias_u0) ? cache.alias_u0 : false, @@ -165,7 +170,7 @@ function SciMLBase.__init( prob, alg, prob.f, fu, u, prob.p; stats, alg.autodiff, linsolve, alg.jvp_autodiff, alg.vjp_autodiff ) - J = jac_cache(nothing) + J = reused_jacobian(jac_cache, u) descent_cache = InternalAPI.init( prob, alg.descent, J, fu, u; stats, abstol, reltol, internalnorm, @@ -212,7 +217,7 @@ function SciMLBase.__init( ) cache = GeneralizedFirstOrderAlgorithmCache( - fu, u, u_cache, prob.p, du, J, alg, prob, globalization, + fu, u, u_cache, prob.p, alg, prob, globalization, jac_cache, descent_cache, linesearch_cache, trustregion_cache, stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, 0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs, @@ -233,7 +238,7 @@ function InternalAPI.step!( J = cache.jac_cache(cache.u) new_jacobian = true else - J = cache.jac_cache(nothing) + J = reused_jacobian(cache.jac_cache, cache.u) new_jacobian = false end end diff --git a/lib/NonlinearSolveHomotopyContinuation/Project.toml b/lib/NonlinearSolveHomotopyContinuation/Project.toml index 21a839d47..acb136a57 100644 --- a/lib/NonlinearSolveHomotopyContinuation/Project.toml +++ b/lib/NonlinearSolveHomotopyContinuation/Project.toml @@ -19,6 +19,9 @@ TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c" [sources.NonlinearSolveBase] path = "../NonlinearSolveBase" +[sources.NonlinearSolve] +path = "../.." + [compat] ADTypes = "1.11.0" Aqua = "0.8" diff --git a/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl b/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl index a7f3d93c3..6bc55c48e 100644 --- a/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl +++ b/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl @@ -18,7 +18,7 @@ using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, AbstractApproximateJacobianUpdateRule, AbstractDescentDirection, AbstractApproximateJacobianUpdateRuleCache, Utils, InternalAPI, get_timer_output, @static_timeit, - update_trace!, L2_NORM, NewtonDescent + update_trace!, L2_NORM, NewtonDescent, reused_jacobian using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode, NonlinearProblem, NonlinearFunction, NoSpecialize using SciMLOperators: AbstractSciMLOperator diff --git a/lib/NonlinearSolveQuasiNewton/src/initialization.jl b/lib/NonlinearSolveQuasiNewton/src/initialization.jl index f730b541b..413cd0896 100644 --- a/lib/NonlinearSolveQuasiNewton/src/initialization.jl +++ b/lib/NonlinearSolveQuasiNewton/src/initialization.jl @@ -125,7 +125,7 @@ function InternalAPI.init( jac_cache = NonlinearSolveBase.construct_jacobian_cache( prob, solver, prob.f, fu, u, p; stats, autodiff, linsolve ) - J = alg.structure(jac_cache(nothing)) + J = alg.structure(reused_jacobian(jac_cache, u)) return InitializedApproximateJacobianCache( J, alg.structure, alg, jac_cache, false, internalnorm ) diff --git a/lib/NonlinearSolveQuasiNewton/src/solve.jl b/lib/NonlinearSolveQuasiNewton/src/solve.jl index f49c86df8..e24c6f45e 100644 --- a/lib/NonlinearSolveQuasiNewton/src/solve.jl +++ b/lib/NonlinearSolveQuasiNewton/src/solve.jl @@ -56,7 +56,6 @@ end u u_cache p - du # Aliased to `get_du(descent_cache)` J # Aliased to `initialization_cache.J` if !inverted_jac alg <: QuasiNewtonAlgorithm prob <: AbstractNonlinearProblem @@ -98,6 +97,13 @@ end initializealg end +function SciMLBase.get_du(cache::QuasiNewtonCache) + SciMLBase.get_du(cache.descent_cache) +end +function NonlinearSolveBase.set_du!(cache::QuasiNewtonCache, δu) + NonlinearSolveBase.set_du!(cache.descent_cache, δu) +end + function NonlinearSolveBase.get_abstol(cache::QuasiNewtonCache) NonlinearSolveBase.get_abstol(cache.termination_cache) end @@ -220,7 +226,7 @@ function SciMLBase.__init( ) cache = QuasiNewtonCache( - fu, u, u_cache, prob.p, du, J, alg, prob, globalization, + fu, u, u_cache, prob.p, J, alg, prob, globalization, initialization_cache, descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache, inv_workspace, stats, 0, 0, alg.max_resets, maxiters, maxtime, @@ -269,7 +275,7 @@ function InternalAPI.step!( elseif recompute_jacobian === nothing # Standard Step reinit = InternalAPI.solve!( - cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du + cache.reinit_rule_cache, cache.J, cache.fu, cache.u, SciMLBase.get_du(cache) ) reinit && (countable_reinit = true) elseif recompute_jacobian diff --git a/lib/NonlinearSolveSpectralMethods/src/solve.jl b/lib/NonlinearSolveSpectralMethods/src/solve.jl index b4bb23d43..9bfd5709c 100644 --- a/lib/NonlinearSolveSpectralMethods/src/solve.jl +++ b/lib/NonlinearSolveSpectralMethods/src/solve.jl @@ -72,6 +72,13 @@ end initializealg end +function SciMLBase.get_du(cache::GeneralizedDFSaneCache) + cache.du +end +function NonlinearSolveBase.set_du!(cache::GeneralizedDFSaneCache, δu) + cache.du = δu +end + function InternalAPI.reinit_self!( cache::GeneralizedDFSaneCache, args...; p = cache.p, u0 = cache.u, alias_u0::Bool = hasproperty(cache, :alias_u0) ? cache.alias_u0 : false, diff --git a/lib/SCCNonlinearSolve/Project.toml b/lib/SCCNonlinearSolve/Project.toml index 3d51de313..2ecd0e358 100644 --- a/lib/SCCNonlinearSolve/Project.toml +++ b/lib/SCCNonlinearSolve/Project.toml @@ -36,6 +36,7 @@ julia = "1.10" [sources] NonlinearSolveBase = {path = "../NonlinearSolveBase"} NonlinearSolveFirstOrder = {path = "../NonlinearSolveFirstOrder"} +NonlinearSolve = {path = "../.."} [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"