refactor: move DFSane into NonlinearSolveSpectralMethods

Author: avik-pal

Project.toml

@@ -19,9 +19,9 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
 NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
+NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
-Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -99,7 +99,6 @@ PETSc = "0.2"
 Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
-Printf = "1.10"
 Random = "1.10"
 ReTestItems = "1.24"
 RecursiveArrayTools = "3.27"

common/nlls_problem_workloads.jl

@@ -0,0 +1,12 @@
+using SciMLBase: NonlinearProblem, NonlinearFunction
+
+nonlinear_functions = (
+    (NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1),
+    (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]),
+    (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])
+)
+
+nonlinear_problems = NonlinearProblem[]
+for (fn, u0) in nonlinear_functions
+    push!(nonlinear_problems, NonlinearProblem(fn, u0, 2.0))
+end

common/nonlinear_problem_workloads.jl

@@ -24,8 +24,10 @@ using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinear
 using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator
 using SciMLOperators: AbstractSciMLOperator, IdentityOperator
 using StaticArraysCore: StaticArray, SMatrix, SArray, MArray
+using SymbolicIndexingInterface: SymbolicIndexingInterface
 
 const DI = DifferentiationInterface
+const SII = SymbolicIndexingInterface
 
 include("public.jl")
 include("utils.jl")
@@ -49,6 +51,8 @@ include("descent/damped_newton.jl")
 include("descent/dogleg.jl")
 include("descent/geodesic_acceleration.jl")
 
+include("solve.jl")
+
 # Unexported Public API
 @compat(public, (L2_NORM, Linf_NORM, NAN_CHECK, UNITLESS_ABS2, get_tolerance))
 @compat(public, (nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution))

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -218,10 +218,111 @@ concrete_jac(v::Bool) = v
 concrete_jac(::Val{false}) = false
 concrete_jac(::Val{true}) = true
 
+"""
+    AbstractNonlinearSolveCache
+
+Abstract Type for all NonlinearSolveBase Caches.
+
+### Interface Functions
+
+  - `get_fu(cache)`: get the residual.
+  - `get_u(cache)`: get the current state.
+  - `set_fu!(cache, fu)`: set the residual.
+  - `has_time_limit(cache)`: whether or not the solver has a maximum time limit.
+  - `not_terminated(cache)`: whether or not the solver has terminated.
+
+  - `SciMLBase.set_u!(cache, u)`: set the current state.
+  - `SciMLBase.reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state
+    `u0` and any additional keyword arguments.
+  - `SciMLBase.step!(cache; kwargs...)`: See [`SciMLBase.step!`](@ref) for more details.
+  - `SciMLBase.isinplace(cache)`: whether or not the solver is inplace.
+
+Additionally implements `SymbolicIndexingInterface` interface Functions.
+
+#### Expected Fields in Sub-Types
+
+For the default interface implementations we expect the following fields to be present in
+the cache:
+
+  - `fu`: the residual.
+  - `u`: the current state.
+  - `maxiters`: the maximum number of iterations.
+  - `nsteps`: the number of steps taken.
+  - `force_stop`: whether or not the solver has been forced to stop.
+  - `retcode`: the return code.
+  - `stats`: `NLStats` object.
+  - `alg`: the algorithm.
+  - `maxtime`: the maximum time limit for the solver. (Optional)
+  - `timer`: the timer for the solver. (Optional)
+  - `total_time`: the total time taken by the solver. (Optional)
+"""
 abstract type AbstractNonlinearSolveCache <: AbstractNonlinearSolveBaseAPI end
 
-function get_u end
-function get_fu end
+get_u(cache::AbstractNonlinearSolveCache) = cache.u
+get_fu(cache::AbstractNonlinearSolveCache) = cache.fu
+set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu)
+SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u)
+
+function has_time_limit(cache::AbstractNonlinearSolveCache)
+    maxtime = Utils.safe_getproperty(cache, Val(:maxtime))
+    return maxtime !== missing && maxtime !== nothing
+end
+
+function not_terminated(cache::AbstractNonlinearSolveCache)
+    return !cache.force_stop && cache.nsteps < cache.maxiters
+end
+
+function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache; kwargs...)
+    return InternalAPI.reinit!(cache; kwargs...)
+end
+function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache, u0; kwargs...)
+    return InternalAPI.reinit!(cache; u0, kwargs...)
+end
+
+SciMLBase.isinplace(cache::AbstractNonlinearSolveCache) = SciMLBase.isinplace(cache.prob)
+
+## SII Interface
+SII.symbolic_container(cache::AbstractNonlinearSolveCache) = cache.prob
+SII.parameter_values(cache::AbstractNonlinearSolveCache) = SII.parameter_values(cache.prob)
+SII.state_values(cache::AbstractNonlinearSolveCache) = SII.state_values(cache.prob)
+
+function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol)
+    if sym === :ps
+        !hasfield(typeof(cache), :ps) && return SII.ParameterIndexingProxy(cache)
+        return getfield(cache, :ps)
+    end
+    return getfield(cache, sym)
+end
+
+Base.getindex(cache::AbstractNonlinearSolveCache, sym) = SII.getu(cache, sym)(cache)
+function Base.setindex!(cache::AbstractNonlinearSolveCache, val, sym)
+    return SII.setu(cache, sym)(cache, val)
+end
+
+# XXX: Implement this
+# function Base.show(io::IO, cache::AbstractNonlinearSolveCache)
+#     __show_cache(io, cache, 0)
+# end
+
+# function __show_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0)
+#     println(io, "$(nameof(typeof(cache)))(")
+#     __show_algorithm(io, cache.alg,
+#         (" "^(indent + 4)) * "alg = " * string(get_name(cache.alg)), indent + 4)
+
+#     ustr = sprint(show, get_u(cache); context = (:compact => true, :limit => true))
+#     println(io, ",\n" * (" "^(indent + 4)) * "u = $(ustr),")
+
+#     residstr = sprint(show, get_fu(cache); context = (:compact => true, :limit => true))
+#     println(io, (" "^(indent + 4)) * "residual = $(residstr),")
+
+#     normstr = sprint(
+#         show, norm(get_fu(cache), Inf); context = (:compact => true, :limit => true))
+#     println(io, (" "^(indent + 4)) * "inf-norm(residual) = $(normstr),")
+
+#     println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",")
+#     println(io, " "^(indent + 4) * "retcode = ", cache.retcode)
+#     print(io, " "^(indent) * ")")
+# end
 
 """
     AbstractLinearSolverCache

lib/NonlinearSolveBase/src/abstract_types.jl

@@ -0,0 +1,107 @@
+function SciMLBase.__solve(
+        prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveAlgorithm, args...;
+        kwargs...
+)
+    cache = SciMLBase.init(prob, alg, args...; kwargs...)
+    return SciMLBase.solve!(cache)
+end
+
+function CommonSolve.solve!(cache::AbstractNonlinearSolveCache)
+    while not_terminated(cache)
+        SciMLBase.step!(cache)
+    end
+
+    # The solver might have set a different `retcode`
+    if cache.retcode == ReturnCode.Default
+        cache.retcode = ifelse(
+            cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success
+        )
+    end
+
+    # XXX: Implement this
+    # update_from_termination_cache!(cache.termination_cache, cache)
+
+    update_trace!(
+        cache.trace, cache.nsteps, get_u(cache), get_fu(cache), nothing, nothing, nothing;
+        last = Val(true)
+    )
+
+    return SciMLBase.build_solution(
+        cache.prob, cache.alg, get_u(cache), get_fu(cache);
+        cache.retcode, cache.stats, cache.trace
+    )
+end
+
+"""
+    step!(cache::AbstractNonlinearSolveCache;
+        recompute_jacobian::Union{Nothing, Bool} = nothing)
+
+Performs one step of the nonlinear solver.
+
+### Keyword Arguments
+
+  - `recompute_jacobian`: allows controlling whether the jacobian is recomputed at the
+    current step. If `nothing`, then the algorithm determines whether to recompute the
+    jacobian. If `true` or `false`, then the jacobian is recomputed or not recomputed,
+    respectively. For algorithms that don't use jacobian information, this keyword is
+    ignored with a one-time warning.
+"""
+function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...)
+    not_terminated(cache) || return
+
+    has_time_limit(cache) && (time_start = time())
+
+    res = @static_timeit cache.timer "solve" begin
+        InternalAPI.step!(cache, args...; kwargs...)
+    end
+
+    cache.stats.nsteps += 1
+    cache.nsteps += 1
+
+    if has_time_limit(cache)
+        cache.total_time += time() - time_start
+
+        if !cache.force_stop && cache.retcode == ReturnCode.Default &&
+           cache.total_time ≥ cache.maxtime
+            cache.retcode = ReturnCode.MaxTime
+            cache.force_stop = true
+        end
+    end
+
+    return res
+end
+
+# Some algorithms don't support creating a cache and doing `solve!`, this unfortunately
+# makes it difficult to write generic code that supports caching. For the algorithms that
+# don't have a `__init` function defined, we create a "Fake Cache", which just calls
+# `__solve` from `solve!`
+# Warning: This doesn't implement all the necessary interface functions
+@concrete mutable struct NonlinearSolveNoInitCache <: AbstractNonlinearSolveCache
+    prob
+    alg
+    args
+    kwargs::Any
+end
+
+function SciMLBase.reinit!(
+        cache::NonlinearSolveNoInitCache, u0 = cache.prob.u0; p = cache.prob.p, kwargs...
+)
+    cache.prob = SciMLBase.remake(cache.prob; u0, p)
+    cache.kwargs = merge(cache.kwargs, kwargs)
+    return cache
+end
+
+function Base.show(io::IO, ::MIME"text/plain", cache::NonlinearSolveNoInitCache)
+    print(io, "NonlinearSolveNoInitCache(alg = $(cache.alg))")
+end
+
+function SciMLBase.__init(
+        prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveAlgorithm, args...;
+        kwargs...
+)
+    return NonlinearSolveNoInitCache(prob, alg, args, kwargs)
+end
+
+function CommonSolve.solve!(cache::NonlinearSolveNoInitCache)
+    return CommonSolve.solve(cache.prob, cache.alg, cache.args...; cache.kwargs...)
+end

lib/NonlinearSolveBase/src/solve.jl

@@ -9,7 +9,7 @@ using SciMLOperators: AbstractSciMLOperator
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, NonlinearFunction
 using StaticArraysCore: StaticArray, SArray
 
-using ..NonlinearSolveBase: L2_NORM, Linf_NORM
+using ..NonlinearSolveBase: NonlinearSolveBase, L2_NORM, Linf_NORM
 
 is_extension_loaded(::Val) = false
 
@@ -145,6 +145,26 @@ function evaluate_f!!(f::NonlinearFunction, fu, u, p)
     return f(u, p)
 end
 
+function evaluate_f(prob::AbstractNonlinearProblem, u)
+    if SciMLBase.isinplace(prob)
+        fu = prob.f.resid_prototype === nothing ? similar(u) :
+             similar(prob.f.resid_prototype)
+        prob.f(fu, u, prob.p)
+    else
+        fu = prob.f(u, prob.p)
+    end
+    return fu
+end
+
+function evaluate_f!(cache, u, p)
+    cache.stats.nf += 1
+    if SciMLBase.isinplace(cache)
+        cache.prob.f(NonlinearSolveBase.get_fu(cache), u, p)
+    else
+        NonlinearSolveBase.set_fu!(cache, cache.prob.f(u, p))
+    end
+end
+
 function make_sparse end
 
 condition_number(J::AbstractMatrix) = cond(J)

lib/NonlinearSolveBase/src/utils.jl

@@ -3,14 +3,36 @@ uuid = "26075421-4e9a-44e1-8bd1-420ed7ad02b2"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 version = "1.0.0"
 
+[deps]
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
+NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+
 [compat]
 Aqua = "0.8"
+CommonSolve = "0.2.4"
+ConcreteStructs = "0.2.3"
+DiffEqBase = "6.155.3"
 ExplicitImports = "1.5"
 Hwloc = "3"
 InteractiveUtils = "<0.0.1, 1"
+LineSearch = "0.1.4"
+LinearAlgebra = "1.11.0"
+MaybeInplace = "0.1.4"
 NonlinearProblemLibrary = "0.1.2"
+NonlinearSolveBase = "1.1"
 Pkg = "1.10"
+PrecompileTools = "1.2"
 ReTestItems = "1.24"
+Reexport = "1"
+SciMLBase = "2.54"
 StableRNGs = "1"
 Test = "1.10"
 julia = "1.10"

lib/NonlinearSolveSpectralMethods/Project.toml

@@ -1,3 +1,41 @@
 module NonlinearSolveSpectralMethods
 
+using Reexport: @reexport
+using PrecompileTools: @compile_workload, @setup_workload
+
+using CommonSolve: CommonSolve
+using ConcreteStructs: @concrete
+using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches
+using LinearAlgebra: dot
+using LineSearch: RobustNonMonotoneLineSearch
+using MaybeInplace: @bb
+using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm,
+                          AbstractNonlinearSolveCache, Utils, InternalAPI, get_timer_output,
+                          @static_timeit, update_trace!
+using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode
+
+include("dfsane.jl")
+
+include("solve.jl")
+
+@setup_workload begin
+    include(joinpath(
+        @__DIR__, "..", "..", "..", "common", "nonlinear_problem_workloads.jl"
+    ))
+
+    algs = [DFSane()]
+
+    @compile_workload begin
+        @sync begin
+            for prob in nonlinear_problems, alg in algs
+                Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false)
+            end
+        end
+    end
+end
+
+@reexport using SciMLBase, NonlinearSolveBase
+
+export GeneralizedDFSane, DFSane
+
 end