[WIP] Multi-Step Methods

Project.toml

@@ -1,10 +1,11 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.5.4"
+version = "3.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
@@ -55,12 +56,13 @@ NonlinearSolveZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.6"
+Accessors = "0.1.32"
 Aqua = "0.8"
 ArrayInterface = "7.7"
 BandedMatrices = "1.4"
 BenchmarkTools = "1.4"
-ConcreteStructs = "0.2.3"
 CUDA = "5.1"
+ConcreteStructs = "0.2.3"
 DiffEqBase = "6.146.0"
 Enzyme = "0.11.11"
 FastBroadcast = "0.2.8"

docs/src/devdocs/internal_interfaces.md

@@ -13,6 +13,7 @@ NonlinearSolve.AbstractNonlinearSolveCache
 ```@docs
 NonlinearSolve.AbstractDescentAlgorithm
 NonlinearSolve.AbstractDescentCache
+NonlinearSolve.DescentResult
 ```
 
 ## Approximate Jacobian

docs/src/tutorials/large_systems.md

@@ -2,8 +2,8 @@
 
 This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving
 ill-conditioned nonlinear systems requires specializing the linear solver on properties of
-the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` linear solve and the
-``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of
+the Jacobian in order to cut down on the `\mathcal{O}(n^3)` linear solve and the
+`\mathcal{O}(n^2)` back-solves. This tutorial is designed to explain the advanced usage of
 NonlinearSolve.jl by solving the steady state stiff Brusselator partial differential
 equation (BRUSS) using NonlinearSolve.jl.
 

src/NonlinearSolve.jl

@@ -8,9 +8,9 @@ import Reexport: @reexport
 import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload
 
 @recompile_invalidations begin
-    using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays,
-          LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf,
-          SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools
+    using Accessors, ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures,
+          LazyArrays, LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences,
+          Printf, SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools
 
     import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing
     import DiffEqBase: AbstractNonlinearTerminationMode,
@@ -45,11 +45,13 @@ include("adtypes.jl")
 include("timer_outputs.jl")
 include("internal/helpers.jl")
 
+include("descent/common.jl")
 include("descent/newton.jl")
 include("descent/steepest.jl")
 include("descent/dogleg.jl")
 include("descent/damped_newton.jl")
 include("descent/geodesic_acceleration.jl")
+include("descent/multistep.jl")
 
 include("internal/operators.jl")
 include("internal/jacobian.jl")
@@ -69,6 +71,7 @@ include("core/spectral_methods.jl")
 
 include("algorithms/raphson.jl")
 include("algorithms/pseudo_transient.jl")
+include("algorithms/multistep.jl")
 include("algorithms/broyden.jl")
 include("algorithms/klement.jl")
 include("algorithms/lbroyden.jl")
@@ -138,7 +141,8 @@ include("default.jl")
 end
 
 # Core Algorithms
-export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane
+export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane,
+       MultiStepNonlinearSolver
 export GaussNewton, LevenbergMarquardt, TrustRegion
 export NonlinearSolvePolyAlgorithm,
        RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
@@ -152,7 +156,9 @@ export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, Genera
 
 # Descent Algorithms
 export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent,
-       GeodesicAcceleration
+       GeodesicAcceleration, GenericMultiStepDescent
+## Multistep Algorithms
+export MultiStepSchemes
 
 # Globalization
 ## Line Search Algorithms

src/abstract_types.jl

@@ -66,8 +66,8 @@
 ### `__internal_solve!` specification
 
 ```julia
-δu, success, intermediates = __internal_solve!(cache::AbstractDescentCache, J, fu, u,
-    idx::Val; skip_solve::Bool = false, kwargs...)
+descent_result = __internal_solve!(cache::AbstractDescentCache, J, fu, u, idx::Val;
+    skip_solve::Bool = false, kwargs...)
 ```
 
   - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`).
@@ -79,21 +79,19 @@
     direction was rejected and we want to try with a modified trust region.
   - `kwargs`: keyword arguments to pass to the linear solver if there is one.
 
-#### Returned values
-
-  - `δu`: the descent direction.
-  - `success`: Certain Descent Algorithms can reject a descent direction for example
-    `GeodesicAcceleration`.
-  - `intermediates`: A named tuple containing intermediates computed during the solve.
-    For example, `GeodesicAcceleration` returns `NamedTuple{(:v, :a)}` containing the
-    "velocity" and "acceleration" terms.
+Returns a result of type [`DescentResult`](@ref).
 
 ### Interface Functions
 
   - `get_du(cache)`: get the descent direction.
   - `get_du(cache, ::Val{N})`: get the `N`th descent direction.
   - `set_du!(cache, δu)`: set the descent direction.
   - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction.
+  - `get_internal_cache(cache, ::Val{field})`: get the internal cache field.
+  - `get_internal_cache(cache, field::Val, ::Val{N})`: get the `N`th internal cache field.
+  - `set_internal_cache!(cache, value, ::Val{field})`: set the internal cache field.
+  - `set_internal_cache!(cache, value, field::Val, ::Val{N})`: set the `N`th internal cache
+    field.
   - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the
     cache has a `last_step_accepted` field and returns it if it does, else returns `true`.
 """
@@ -105,6 +103,29 @@
 set_du!(cache::AbstractDescentCache, δu) = (cache.δu = δu)
 set_du!(cache::AbstractDescentCache, δu, ::Val{1}) = set_du!(cache, δu)
 set_du!(cache::AbstractDescentCache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu)
+function get_internal_cache(cache::AbstractDescentCache, ::Val{field}) where {field}
+    return getproperty(cache, field)
+end
+function get_internal_cache(cache::AbstractDescentCache, field::Val, ::Val{1})
+    return get_internal_cache(cache, field)
+end
+function get_internal_cache(
+        cache::AbstractDescentCache, ::Val{field}, ::Val{N}) where {field, N}
+    true_field = Symbol(string(field), "s")  # Julia 1.10 compiles this away
+    return getproperty(cache, true_field)[N]
+end
+function set_internal_cache!(cache::AbstractDescentCache, value, ::Val{field}) where {field}
+    return setproperty!(cache, field, value)
+end
+function set_internal_cache!(
+        cache::AbstractDescentCache, value, field::Val, ::Val{1})
+    return set_internal_cache!(cache, value, field)
+end
+function set_internal_cache!(
+        cache::AbstractDescentCache, value, ::Val{field}, ::Val{N}) where {field, N}
+    true_field = Symbol(string(field), "s")  # Julia 1.10 compiles this away
+    return setproperty!(cache, true_field, value, N)
+end
 
 function last_step_accepted(cache::AbstractDescentCache)
     hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted

src/algorithms/multistep.jl

@@ -0,0 +1,10 @@
+function MultiStepNonlinearSolver(; concrete_jac = nothing, linsolve = nothing,
+        scheme = MSS.PotraPtak3, precs = DEFAULT_PRECS, autodiff = nothing,
+        vjp_autodiff = nothing, linesearch = NoLineSearch())
+    forward_ad = ifelse(autodiff isa ADTypes.AbstractForwardMode, autodiff, nothing)
+    scheme_concrete = apply_patch(
+        scheme, (; autodiff, vjp_autodiff, jvp_autodiff = forward_ad))
+    descent = GenericMultiStepDescent(; scheme = scheme_concrete, linsolve, precs)
+    return GeneralizedFirstOrderAlgorithm(; concrete_jac, name = MSS.display_name(scheme),
+        descent, jacobian_ad = autodiff, linesearch, reverse_ad = vjp_autodiff, forward_ad)
+end

src/core/approximate_jacobian.jl

@@ -163,8 +163,7 @@
 
         linsolve = get_linear_solver(alg.descent)
         initialization_cache = __internal_init(prob, alg.initialization, alg, f, fu, u, p;
-            linsolve,
-            maxiters, internalnorm)
+            linsolve, maxiters, internalnorm)
 
         abstol, reltol, termination_cache = init_termination_cache(abstol, reltol, fu, u,
             termination_condition)
@@ -222,9 +221,7 @@
     new_jacobian = true
     @static_timeit cache.timer "jacobian init/reinit" begin
         if get_nsteps(cache) == 0  # First Step is special ignore kwargs
-            J_init = __internal_solve!(cache.initialization_cache,
-                cache.fu,
-                cache.u,
+            J_init = __internal_solve!(cache.initialization_cache, cache.fu, cache.u,
                 Val(false))
             if INV
                 if jacobian_initialized_preinverted(cache.initialization_cache.alg)
@@ -283,54 +280,65 @@
     @static_timeit cache.timer "descent" begin
         if cache.trustregion_cache !== nothing &&
            hasfield(typeof(cache.trustregion_cache), :trust_region)
-            δu, descent_success, descent_intermediates = __internal_solve!(
-                cache.descent_cache,
-                J, cache.fu, cache.u; new_jacobian,
-                trust_region = cache.trustregion_cache.trust_region)
+            descent_result = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u;
+                new_jacobian, trust_region = cache.trustregion_cache.trust_region)
         else
-            δu, descent_success, descent_intermediates = __internal_solve!(
-                cache.descent_cache,
-                J, cache.fu, cache.u; new_jacobian)
+            descent_result = __internal_solve!(cache.descent_cache, J, cache.fu, cache.u;
+                new_jacobian)
         end
     end
 
-    if descent_success
-        if GB === :LineSearch
-            @static_timeit cache.timer "linesearch" begin
-                needs_reset, α = __internal_solve!(cache.linesearch_cache, cache.u, δu)
-            end
-            if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period
-                cache.force_reinit = true
-            else
-                @static_timeit cache.timer "step" begin
-                    @bb axpy!(α, δu, cache.u)
-                    evaluate_f!(cache, cache.u, cache.p)
-                end
-            end
-        elseif GB === :TrustRegion
-            @static_timeit cache.timer "trustregion" begin
-                tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache, J,
-                    cache.fu, cache.u, δu, descent_intermediates)
-                if tr_accepted
-                    @bb copyto!(cache.u, u_new)
-                    @bb copyto!(cache.fu, fu_new)
-                end
-                if hasfield(typeof(cache.trustregion_cache), :shrink_counter) &&
-                   cache.trustregion_cache.shrink_counter > cache.max_shrink_times
-                    cache.retcode = ReturnCode.ShrinkThresholdExceeded
-                    cache.force_stop = true
-                end
-            end
-            α = true
-        elseif GB === :None
+    if descent_result.success
+        if GB === :None
             @static_timeit cache.timer "step" begin
-                @bb axpy!(1, δu, cache.u)
+                if descent_result.u !== missing
+                    @bb copyto!(cache.u, descent_result.u)
+                elseif descent_result.δu !== missing
+                    @bb axpy!(1, descent_result.δu, cache.u)
+                else
+                    error("This shouldn't occur. `$(cache.alg.descent)` is incorrectly \
+                           specified.")
+                end
                 evaluate_f!(cache, cache.u, cache.p)
             end
             α = true
         else
-            error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \
-                :TrustRegion, :None)")
+            δu = descent_result.δu
+            @assert δu!==missing "Descent Supporting LineSearch or TrustRegion must return a `δu`."
+
+            if GB === :LineSearch
+                @static_timeit cache.timer "linesearch" begin
+                    needs_reset, α = __internal_solve!(cache.linesearch_cache, cache.u, δu)
+                end
+                if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period
+                    cache.force_reinit = true
+                else
+                    @static_timeit cache.timer "step" begin
+                        @bb axpy!(α, δu, cache.u)
+                        evaluate_f!(cache, cache.u, cache.p)
+                    end
+                end
+            elseif GB === :TrustRegion
+                @static_timeit cache.timer "trustregion" begin
+                    tr_accepted, u_new, fu_new = __internal_solve!(cache.trustregion_cache,
+                        J, cache.fu, cache.u, δu, descent_result.extras)
+                    if tr_accepted
+                        @bb copyto!(cache.u, u_new)
+                        @bb copyto!(cache.fu, fu_new)
+                        α = true
+                    else
+                        α = false
+                    end
+                    if hasfield(typeof(cache.trustregion_cache), :shrink_counter) &&
+                       cache.trustregion_cache.shrink_counter > cache.max_shrink_times
+                        cache.retcode = ReturnCode.ShrinkThresholdExceeded
+                        cache.force_stop = true
+                    end
+                end
+            else
+                error("Unknown Globalization Strategy: $(GB). Allowed values are \
+                       (:LineSearch, :TrustRegion, :None)")
+            end
         end
         check_and_update!(cache, cache.fu, cache.u, cache.u_cache)
     else