diff --git a/src/gaussnewton.jl b/src/gaussnewton.jl index f8c0c7fbe..42521189e 100644 --- a/src/gaussnewton.jl +++ b/src/gaussnewton.jl @@ -1,6 +1,6 @@ """ - GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - adkwargs...) + GaussNewton(; concrete_jac = nothing, linsolve = nothing, + precs = DEFAULT_PRECS, adkwargs...) An advanced GaussNewton implementation with support for efficient handling of sparse matrices via colored automatic differentiation and preconditioned linear solvers. Designed @@ -41,8 +41,8 @@ for large-scale and numerically-difficult nonlinear least squares problems. precs end -function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - adkwargs...) +function GaussNewton(; concrete_jac = nothing, linsolve = CholeskyFactorization(), + precs = DEFAULT_PRECS, adkwargs...) ad = default_adargs_to_adtype(; adkwargs...) return GaussNewton{_unwrap_val(concrete_jac)}(ad, linsolve, precs) end @@ -97,8 +97,8 @@ function perform_step!(cache::GaussNewtonCache{true}) __matmul!(Jᵀf, J', fu1) # u = u - J \ fu - linres = dolinsolve(alg.precs, linsolve; A = JᵀJ, b = _vec(Jᵀf), linu = _vec(du), - p, reltol = cache.abstol) + linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(JᵀJ), b = _vec(Jᵀf), + linu = _vec(du), p, reltol = cache.abstol) cache.linsolve = linres.cache @. u = u - du f(cache.fu_new, u, p) @@ -125,8 +125,8 @@ function perform_step!(cache::GaussNewtonCache{false}) if linsolve === nothing cache.du = fu1 / cache.J else - linres = dolinsolve(alg.precs, linsolve; A = cache.JᵀJ, b = _vec(cache.Jᵀf), - linu = _vec(cache.du), p, reltol = cache.abstol) + linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.JᵀJ), + b = _vec(cache.Jᵀf), linu = _vec(cache.du), p, reltol = cache.abstol) cache.linsolve = linres.cache end cache.u = @. u - cache.du # `u` might not support mutation diff --git a/src/jacobian.jl b/src/jacobian.jl index 10ac69ec6..eddfdaef4 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -95,14 +95,14 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii Jᵀfu = J' * fu end - linprob = LinearProblem(needsJᵀJ ? JᵀJ : J, needsJᵀJ ? _vec(Jᵀfu) : _vec(fu); - u0 = _vec(du)) + linprob = LinearProblem(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J, + needsJᵀJ ? _vec(Jᵀfu) : _vec(fu); u0 = _vec(du)) weight = similar(u) recursivefill!(weight, true) - Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing, - nothing)..., weight) + Pl, Pr = wrapprecs(alg.precs(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J, nothing, u, p, + nothing, nothing, nothing, nothing, nothing)..., weight) linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr, linsolve_kwargs...) @@ -119,6 +119,12 @@ __init_JᵀJ(J::Number) = zero(J) __init_JᵀJ(J::AbstractArray) = J' * J __init_JᵀJ(J::StaticArray) = MArray{Tuple{size(J, 2), size(J, 2)}, eltype(J)}(undef) +__maybe_symmetric(x) = Symmetric(x) +__maybe_symmetric(x::Number) = x +# LinearSolve with `nothing` doesn't dispatch correctly here +__maybe_symmetric(x::StaticArray) = x +__maybe_symmetric(x::SparseArrays.AbstractSparseMatrix) = x + ## Special Handling for Scalars function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p, ::Val{false}; linsolve_with_JᵀJ::Val{needsJᵀJ} = Val(false), diff --git a/src/levenberg.jl b/src/levenberg.jl index 8b0e5728e..07a8b8251 100644 --- a/src/levenberg.jl +++ b/src/levenberg.jl @@ -13,8 +13,8 @@ numerically-difficult nonlinear systems. ### Keyword Arguments - `autodiff`: determines the backend used for the Jacobian. Note that this argument is - ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to - `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. + ignored if an analytical Jacobian is passed, as that will be used instead. Defaults to + `AutoForwardDiff()`. Valid choices are types from ADTypes.jl. - `concrete_jac`: whether to build a concrete Jacobian. If a Krylov-subspace method is used, then the Jacobian will not be constructed and instead direct Jacobian-vector products `J*v` are computed using forward-mode automatic differentiation or finite differencing @@ -203,8 +203,8 @@ function perform_step!(cache::LevenbergMarquardtCache{true}) # The following lines do: cache.v = -cache.mat_tmp \ cache.u_tmp mul!(cache.u_tmp, J', fu1) @. cache.mat_tmp = JᵀJ + λ * DᵀD - linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp, b = _vec(cache.u_tmp), - linu = _vec(cache.du), p = p, reltol = cache.abstol) + linres = dolinsolve(alg.precs, linsolve; A = __maybe_symmetric(cache.mat_tmp), + b = _vec(cache.u_tmp), linu = _vec(cache.du), p = p, reltol = cache.abstol) cache.linsolve = linres.cache @. cache.v = -cache.du @@ -280,8 +280,8 @@ function perform_step!(cache::LevenbergMarquardtCache{false}) if linsolve === nothing cache.v = -cache.mat_tmp \ (J' * fu1) else - linres = dolinsolve(alg.precs, linsolve; A = -cache.mat_tmp, b = _vec(J' * fu1), - linu = _vec(cache.v), p, reltol = cache.abstol) + linres = dolinsolve(alg.precs, linsolve; A = -__maybe_symmetric(cache.mat_tmp), + b = _vec(J' * fu1), linu = _vec(cache.v), p, reltol = cache.abstol) cache.linsolve = linres.cache end @@ -291,7 +291,7 @@ function perform_step!(cache::LevenbergMarquardtCache{false}) cache.a = -cache.mat_tmp \ _vec(J' * ((2 / h) .* ((f(u .+ h .* v, p) .- fu1) ./ h .- J * v))) else - linres = dolinsolve(alg.precs, linsolve; A = -cache.mat_tmp, + linres = dolinsolve(alg.precs, linsolve; b = _mutable(_vec(J' * ((2 / h) .* ((f(u .+ h .* v, p) .- fu1) ./ h .- J * v)))), linu = _vec(cache.a), p, reltol = cache.abstol) diff --git a/test/nonlinear_least_squares.jl b/test/nonlinear_least_squares.jl index c9fda61ba..9ea4b0eca 100644 --- a/test/nonlinear_least_squares.jl +++ b/test/nonlinear_least_squares.jl @@ -27,10 +27,15 @@ prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function; resid_prototype = zero(y_target)), θ_init, x) nlls_problems = [prob_oop, prob_iip] -solvers = [GaussNewton(), LevenbergMarquardt(), LSOptimSolver(:lm), LSOptimSolver(:dogleg)] +solvers = [ + GaussNewton(), + LevenbergMarquardt(), + LSOptimSolver(:lm), + LSOptimSolver(:dogleg), +] for prob in nlls_problems, solver in solvers - @time sol = solve(prob, solver; maxiters = 1000, abstol = 1e-8) + @time sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test SciMLBase.successful_retcode(sol) @test norm(sol.resid) < 1e-6 end