feat: nicer generalized printing of structs/results

Author: avik-pal

lib/NonlinearSolveBase/src/abstract_types.jl

@@ -10,14 +10,7 @@ end
 abstract type AbstractNonlinearSolveBaseAPI end # Mostly used for pretty-printing
 
 function Base.show(io::IO, ::MIME"text/plain", alg::AbstractNonlinearSolveBaseAPI)
-    main_name = nameof(typeof(alg))
-    modifiers = String[]
-    for field in fieldnames(typeof(alg))
-        val = getfield(alg, field)
-        Utils.is_default_value(val, field, getfield(alg, field)) && continue
-        push!(modifiers, "$(field) = $(val)")
-    end
-    print(io, "$(main_name)($(join(modifiers, ", ")))")
+    print(io, Utils.clean_sprint_struct(alg))
     return
 end
 
@@ -198,12 +191,9 @@ Abstract Type for all NonlinearSolveBase Algorithms.
 
   - `concrete_jac(alg)`: whether or not the algorithm uses a concrete Jacobian. Defaults
     to `nothing`.
-  - `get_name(alg)`: get the name of the algorithm.
 """
 abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end
 
-get_name(alg::AbstractNonlinearSolveAlgorithm) = Utils.safe_getproperty(alg, Val(:name))
-
 """
     concrete_jac(alg::AbstractNonlinearSolveAlgorithm)::Bool
 
@@ -218,6 +208,18 @@ concrete_jac(v::Bool) = v
 concrete_jac(::Val{false}) = false
 concrete_jac(::Val{true}) = true
 
+function Base.show(io::IO, ::MIME"text/plain", alg::AbstractNonlinearSolveAlgorithm)
+    print(io, Utils.clean_sprint_struct(alg, 0))
+    return
+end
+
+function show_nonlinearsolve_algorithm(
+        io::IO, alg::AbstractNonlinearSolveAlgorithm, name, indent::Int = 0
+)
+    print(io, name)
+    print(io, Utils.clean_sprint_struct(alg, indent))
+end
+
 """
     AbstractNonlinearSolveCache
 
@@ -299,30 +301,34 @@ 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),")
+function Base.show(io::IO, ::MIME"text/plain", cache::AbstractNonlinearSolveCache)
+    return show_nonlinearsolve_cache(io, cache)
+end
 
-#     println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",")
-#     println(io, " "^(indent + 4) * "retcode = ", cache.retcode)
-#     print(io, " "^(indent) * ")")
-# end
+function show_nonlinearsolve_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0)
+    println(io, "$(nameof(typeof(cache)))(")
+    show_nonlinearsolve_algorithm(
+        io,
+        cache.alg,
+        (" "^(indent + 4)) * "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/descent/geodesic_acceleration.jl

@@ -49,13 +49,12 @@ get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent)
     last_step_accepted::Bool
 end
 
-# XXX: Implement
-# function __reinit_internal!(
-#         cache::GeodesicAccelerationCache, args...; p = cache.p, kwargs...)
-#     cache.p = p
-#     cache.last_step_accepted = false
-# end
+function InternalAPI.reinit!(cache::GeodesicAccelerationCache; p = cache.p, kwargs...)
+    cache.p = p
+    cache.last_step_accepted = false
+end
 
+# XXX: Implement
 # @internal_caches GeodesicAccelerationCache :descent_cache
 
 function get_velocity(cache::GeodesicAccelerationCache)

lib/NonlinearSolveBase/src/utils.jl

@@ -115,6 +115,7 @@ unwrap_val(::Val{x}) where {x} = unwrap_val(x)
 
 is_default_value(::Any, ::Symbol, ::Nothing) = true
 is_default_value(::Any, ::Symbol, ::Missing) = true
+is_default_value(::Any, ::Symbol, val::Int) = val == typemax(typeof(val))
 is_default_value(::Any, ::Symbol, ::Any) = false
 
 maybe_symmetric(x) = Symmetric(x)
@@ -178,7 +179,6 @@ function condition_number(J::AbstractVector)
 end
 condition_number(::Any) = -1
 
-# XXX: Move to NonlinearSolveQuasiNewton
 # compute `pinv` if `inv` won't work
 maybe_pinv!!_workspace(A) = nothing, A
 
@@ -239,4 +239,47 @@ function make_identity!!(A::AbstractMatrix{T}, α) where {T}
     return A
 end
 
+function clean_sprint_struct(x)
+    x isa Symbol && return "$(Meta.quot(x))"
+    x isa Number && return string(x)
+    (!Base.isstructtype(typeof(x)) || x isa Val) && return string(x)
+
+    modifiers = String[]
+    name = nameof(typeof(x))
+    for field in fieldnames(typeof(x))
+        val = getfield(x, field)
+        if field === :name
+            name = val
+            continue
+        end
+        is_default_value(x, field, val) && continue
+        push!(modifiers, "$(field) = $(clean_sprint_struct(val))")
+    end
+
+    return "$(nameof(typeof(x)))($(join(modifiers, ", ")))"
+end
+
+function clean_sprint_struct(x, indent::Int)
+    x isa Symbol && return "$(Meta.quot(x))"
+    x isa Number && return string(x)
+    (!Base.isstructtype(typeof(x)) || x isa Val) && return string(x)
+
+    modifiers = String[]
+    name = nameof(typeof(x))
+    for field in fieldnames(typeof(x))
+        val = getfield(x, field)
+        if field === :name
+            name = val
+            continue
+        end
+        is_default_value(x, field, val) && continue
+        push!(modifiers, "$(field) = $(clean_sprint_struct(val, indent + 4))")
+    end
+    spacing = " "^indent * "    "
+    spacing_last = " "^indent
+
+    length(modifiers) == 0 && return "$(nameof(typeof(x)))()"
+    return "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))"
+end
+
 end

lib/NonlinearSolveQuasiNewton/src/initialization.jl

@@ -151,12 +151,13 @@ NonlinearSolveBase.jacobian_initialized_preinverted(::BroydenLowRankInitializati
 
 function InternalAPI.init(
         prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization,
-        solver, f::F, fu, u, p; internalnorm::IN = L2_NORM, kwargs...
+        solver, f::F, fu, u, p;
+        internalnorm::IN = L2_NORM, maxiters = 1000, kwargs...
 ) where {F, IN}
     if u isa Number # Use the standard broyden
         return InternalAPI.init(
             prob, IdentityInitialization(true, FullStructure()),
-            solver, f, fu, u, p; internalnorm, kwargs...
+            solver, f, fu, u, p; internalnorm, maxiters, kwargs...
         )
     end
     # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays
@@ -212,7 +213,7 @@ Base.adjoint(op::BroydenLowRankJacobian{<:Real}) = transpose(op)
 
 # Storing the transpose to ensure contiguous memory on splicing
 function BroydenLowRankJacobian(
-    fu::StaticArray, u::StaticArray; alpha = true, threshold::Val = Val(10)
+        fu::StaticArray, u::StaticArray; alpha = true, threshold::Val = Val(10)
 )
     T = promote_type(eltype(u), eltype(fu))
     U = MArray{Tuple{prod(Size(fu)), Utils.unwrap_val(threshold)}, T}(undef)
@@ -265,8 +266,8 @@ function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowR
 end
 
 function LinearAlgebra.mul!(
-    J::BroydenLowRankJacobian, u::AbstractArray, vᵀ::LinearAlgebra.AdjOrTransAbsVec,
-    α::Bool, β::Bool
+        J::BroydenLowRankJacobian, u::AbstractArray, vᵀ::LinearAlgebra.AdjOrTransAbsVec,
+        α::Bool, β::Bool
 )
     @assert α & β
     idx_update = mod1(J.idx + 1, size(J.U, 2))

lib/NonlinearSolveQuasiNewton/src/reset_conditions.jl

@@ -103,7 +103,7 @@ function InternalAPI.init(
 end
 
 @concrete struct IllConditionedJacobianResetCache <: AbstractResetConditionCache
-    condition_number_threshold <: Number
+    condition_number_threshold
 end
 
 # NOTE: we don't need a reinit! since we establish the threshold based on the eltype

lib/NonlinearSolveQuasiNewton/src/solve.jl

@@ -37,22 +37,6 @@ examples include [`Broyden`](@ref)'s Method.
     name::Symbol
 end
 
-# XXX: Implement
-# function __show_algorithm(io::IO, alg::QuasiNewtonAlgorithm, name, indent)
-#     modifiers = String[]
-#     __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)")
-#     __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)")
-#     push!(modifiers, "descent = $(alg.descent)")
-#     push!(modifiers, "update_rule = $(alg.update_rule)")
-#     push!(modifiers, "reinit_rule = $(alg.reinit_rule)")
-#     push!(modifiers, "max_resets = $(alg.max_resets)")
-#     push!(modifiers, "initialization = $(alg.initialization)")
-#     store_inverse_jacobian(alg.update_rule) && push!(modifiers, "inverse_jacobian = true")
-#     spacing = " "^indent * "    "
-#     spacing_last = " "^indent
-#     print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))")
-# end
-
 function QuasiNewtonAlgorithm(;
         linesearch = missing, trustregion = missing, descent, update_rule, reinit_rule,
         initialization, max_resets::Int = typemax(Int), name::Symbol = :unknown,

lib/NonlinearSolveSpectralMethods/src/solve.jl

@@ -32,18 +32,6 @@ function GeneralizedDFSane(;
     return GeneralizedDFSane(linesearch, sigma_min, sigma_max, sigma_1, name)
 end
 
-# XXX: Add
-# function __show_algorithm(io::IO, alg::GeneralizedDFSane, name, indent)
-#     modifiers = String[]
-#     __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)")
-#     push!(modifiers, "σ_min = $(alg.σ_min)")
-#     push!(modifiers, "σ_max = $(alg.σ_max)")
-#     push!(modifiers, "σ_1 = $(alg.σ_1)")
-#     spacing = " "^indent * "    "
-#     spacing_last = " "^indent
-#     print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))")
-# end
-
 @concrete mutable struct GeneralizedDFSaneCache <: AbstractNonlinearSolveCache
     # Basic Requirements
     fu

src/NonlinearSolve.jl

@@ -32,6 +32,8 @@ using NonlinearSolveBase: NonlinearSolveBase,
                           reset_timer!, @static_timeit,
                           init_nonlinearsolve_trace, update_trace!, reset!
 
+using NonlinearSolveQuasiNewton: Broyden, Klement
+
 # XXX: Remove
 import NonlinearSolveBase: InternalAPI, concrete_jac, supports_line_search,
                            supports_trust_region, last_step_accepted, get_linear_solver,
@@ -149,12 +151,6 @@ include("internal/forward_diff.jl") # we need to define after the algorithms
 
     @compile_workload begin
         @sync begin
-            for T in (Float32, Float64), (fn, u0) in nlfuncs
-                Threads.@spawn NonlinearProblem(fn, T.(u0), T(2))
-            end
-            for (fn, u0) in nlfuncs
-                Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0)
-            end
             for prob in probs_nls, alg in nls_algs
                 Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false)
             end