Handle non-vector inputs

Project.toml

@@ -50,7 +50,7 @@ PrecompileTools = "1"
 RecursiveArrayTools = "2"
 Reexport = "0.2, 1"
 SciMLBase = "2.4"
-SimpleNonlinearSolve = "0.1.22"
+SimpleNonlinearSolve = "0.1.23"
 SparseDiffTools = "2.6"
 StaticArraysCore = "1.4"
 UnPack = "1.0"

src/NonlinearSolve.jl

@@ -5,6 +5,7 @@ if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_m
 end
 
 using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
+import ArrayInterface: restructure
 import ForwardDiff
 
 import ADTypes: AbstractFiniteDifferencesMode

src/jacobian.jl

@@ -92,7 +92,7 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
     if needsJᵀJ
         JᵀJ = __init_JᵀJ(J)
         # FIXME: This needs to be handled better for JacVec Operator
-        Jᵀfu = J' * fu
+        Jᵀfu = J' * _vec(fu)
     end
 
     linprob = LinearProblem(needsJᵀJ ? __maybe_symmetric(JᵀJ) : J,

src/levenberg.jl

@@ -172,7 +172,7 @@ function SciMLBase.__init(prob::Union{NonlinearProblem{uType, iip},
     else
         d = similar(u)
         d .= min_damping_D
-        DᵀD = Diagonal(d)
+        DᵀD = Diagonal(_vec(d))
     end
 
     loss = internalnorm(fu1)
@@ -289,7 +289,7 @@ function perform_step!(cache::LevenbergMarquardtCache{false})
         cache.v = -cache.mat_tmp \ (J' * fu1)
     else
         linres = dolinsolve(alg.precs, linsolve; A = -__maybe_symmetric(cache.mat_tmp),
-            b = _vec(J' * fu1), linu = _vec(cache.v), p, reltol = cache.abstol)
+            b = _vec(J' * _vec(fu1)), linu = _vec(cache.v), p, reltol = cache.abstol)
         cache.linsolve = linres.cache
     end
 
@@ -301,7 +301,7 @@ function perform_step!(cache::LevenbergMarquardtCache{false})
     else
         linres = dolinsolve(alg.precs, linsolve;
             b = _mutable(_vec(J' *
-                              ((2 / h) .* ((f(u .+ h .* v, p) .- fu1) ./ h .- J * v)))),
+                         _vec(((2 / h) .* (_vec(f(u .+ h .* _restructure(u,v), p)) .- _vec(fu1) ./ h .- J * _vec(v)))))),
             linu = _vec(cache.a), p, reltol = cache.abstol)
         cache.linsolve = linres.cache
     end

src/trustRegion.jl

@@ -347,7 +347,7 @@ function perform_step!(cache::TrustRegionCache{true})
     if cache.make_new_J
         jacobian!!(J, cache)
         mul!(cache.H, J', J)
-        mul!(cache.g, J', fu)
+        mul!(_vec(cache.g), J', _vec(fu))
         cache.stats.njacs += 1
 
         # do not use A = cache.H, b = _vec(cache.g) since it is equivalent
@@ -378,9 +378,9 @@ function perform_step!(cache::TrustRegionCache{false})
     if make_new_J
         J = jacobian!!(cache.J, cache)
         cache.H = J' * J
-        cache.g = J' * fu
+        cache.g = _restructure(fu, J' * _vec(fu))
         cache.stats.njacs += 1
-        cache.u_gauss_newton = -1 .* (cache.H \ cache.g)
+        cache.u_gauss_newton = -1 .* _restructure(cache.g, cache.H \ _vec(cache.g))
     end
 
     # Compute the Newton step.
@@ -419,7 +419,7 @@ function trust_region_step!(cache::TrustRegionCache)
     cache.loss_new = get_loss(fu_new)
 
     # Compute the ratio of the actual reduction to the predicted reduction.
-    cache.r = -(loss - cache.loss_new) / (dot(du, g) + dot(du, H, du) / 2)
+    cache.r = -(loss - cache.loss_new) / (dot(_vec(du), _vec(g)) + dot(_vec(du), H, _vec(du)) / 2)
     @unpack r = cache
 
     if radius_update_scheme === RadiusUpdateSchemes.Simple
@@ -597,7 +597,7 @@ function dogleg!(cache::TrustRegionCache{true})
 
     # Take intersection of steepest descent direction and trust region if Cauchy point lies outside of trust region
     l_grad = norm(cache.g) # length of the gradient
-    d_cauchy = l_grad^3 / dot(cache.g, cache.H, cache.g) # distance of the cauchy point from the current iterate
+    d_cauchy = l_grad^3 / dot(_vec(cache.g), cache.H, _vec(cache.g)) # distance of the cauchy point from the current iterate
     if d_cauchy >= trust_r
         @. cache.du = -(trust_r / l_grad) * cache.g # step to the end of the trust region
         return
@@ -627,7 +627,7 @@ function dogleg!(cache::TrustRegionCache{false})
 
     ## Take intersection of steepest descent direction and trust region if Cauchy point lies outside of trust region
     l_grad = norm(cache.g)
-    d_cauchy = l_grad^3 / dot(cache.g, cache.H, cache.g) # distance of the cauchy point from the current iterate
+    d_cauchy = l_grad^3 / dot(_vec(cache.g), cache.H, _vec(cache.g)) # distance of the cauchy point from the current iterate
     if d_cauchy > trust_r # cauchy point lies outside of trust region
         cache.du = -(trust_r / l_grad) * cache.g # step to the end of the trust region
         return

src/utils.jl

@@ -74,6 +74,9 @@ end
 @inline _vec(v::Number) = v
 @inline _vec(v::AbstractVector) = v
 
+@inline _restructure(y,x) = restructure(y,x)
+@inline _restructure(y::Number,x::Number) = x
+
 DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing
 
 function dolinsolve(precs::P, linsolve; A = nothing, linu = nothing, b = nothing,

test/matrix_resizing.jl

@@ -0,0 +1,21 @@
+using NonlinearSolve, Test
+
+ff(u, p) = u .* u .- p
+u0 = rand(2,2)
+p = 2.0
+vecprob = NonlinearProblem(ff, vec(u0), p)
+prob = NonlinearProblem(ff, u0, p)
+
+for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg())
+    @test vec(solve(prob, alg).u) == solve(vecprob, alg).u
+end
+
+fiip(du, u, p) = (du .= u .* u .- p)
+u0 = rand(2,2)
+p = 2.0
+vecprob = NonlinearProblem(fiip, vec(u0), p)
+prob = NonlinearProblem(fiip, u0, p)
+
+for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg())
+    @test vec(solve(prob, alg).u) == solve(vecprob, alg).u
+end

test/runtests.jl

@@ -16,6 +16,7 @@ end
         @time @safetestset "Basic Tests + Some AD" include("basictests.jl")
         @time @safetestset "Sparsity Tests" include("sparse.jl")
         @time @safetestset "Polyalgs" include("polyalgs.jl")
+        @time @safetestset "Matrix Resizing" include("matrix_resizing.jl")
         @time @safetestset "Nonlinear Least Squares" include("nonlinear_least_squares.jl")
     end