# Modify custom solver to deal with multidimensional arrays

Some time ago @longemen3000 shared with me his custom nl equation solver written with `NLSolvers` here. It works well when the system is provided as a vector. I tried to modify it to accept systems provided with multidimensional arrays, but I didn’t succeed. Where did I go wrong?

Here is an illustration of the problem:

``````using NLSolvers
using ForwardDiff

function customnlsolve(f!, x0, method=NLSolvers.TrustRegion(NLSolvers.Newton(), NLSolvers.Dogleg()), options=NLSolvers.NEqOptions())
siz = size(x0)
#xcache = zeros(eltype(x0), len)
xcache = similar(x0)
#Fcache = zeros(eltype(x0), len)
Fcache = zeros(eltype(x0), siz...)
Jcache = zeros(eltype(x0), *(siz...), *(siz...))
jconfig = ForwardDiff.JacobianConfig(f!, x0, x0)

function j!(J, x)
#@show J
ForwardDiff.jacobian!(J, f!, Fcache, x, jconfig)
end

function fj!(F, J, x)
#@show J, F
ForwardDiff.jacobian!(J, f!, F, x, jconfig)
F, J
end

function jv!(x)
function JacV(Fv, v)
ForwardDiff.jacobian!(Jcache, f!, Fcache, v, jconfig)
Fv .= Jcache * v
end
return LinearMap(JacV, length(x))
end

vectorobj = NLSolvers.VectorObjective(f!, j!, fj!, jv!)
vectorprob = NEqProblem(vectorobj)
return NLSolvers.solve(vectorprob, x0, method, options)
end
``````

system supplied with vectors:

``````function sys2!(sto2, x)
sto2[1] = 1 - x[1]
sto2[2] = 10(x[2]-x[1]^2)
end

sto2 = zeros(2)
x02 = rand(2)
``````

it works:

``````julia> customnlsolve(sys2!, x02)
Results of solving non-linear equations

* Algorithm:
Newton's method with default linsolve with Dogleg{Nothing}

* Candidate solution:
Final residual 2-norm:      2.89e-15
Final residual Inf-norm:    2.89e-15

Initial residual 2-norm:    7.09e+00
Initial residual Inf-norm:  7.09e+00

* Convergence measures
|F(x')|               = 2.89e-15 <= 0.00e+00 (false)

* Work counters
Seconds run:   6.60e-05
Iterations:    2
``````

with a system supplied with arrays, it throws an error:

``````function sys4!(sto, xx)
x = xx[1, 1]
y = xx[1, 2]
z = xx[2, 1]
t = xx[2, 2]
sto[1, 1] = 4.5e-4*x - x*y
sto[1, 2] = z - 1e-14/t
sto[2, 1] = y - 0.02 - x
sto[2, 2] = z - t - y
return sto
end

sto4 = zeros(2, 2)
x04 = rand(2, 2)
``````
``````julia> customnlsolve(sys4!, x04)
ERROR: DimensionMismatch("new dimensions (4, 4) must be consistent with array size 4")
Stacktrace:
[1] (::Base.var"#throw_dmrsa#272")(dims::Tuple{Int64, Int64}, len::Int64)
@ Base ./reshapedarray.jl:41
[2] reshape
@ ./reshapedarray.jl:45 [inlined]
[3] reshape
@ ./reshapedarray.jl:116 [inlined]
[4] extract_jacobian!(#unused#::Type{ForwardDiff.Tag{typeof(sys4!), Float64}}, result::Matrix{Float64}, ydual::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, n::Int64)
@ ForwardDiff ~/.julia/packages/ForwardDiff/CkdHU/src/jacobian.jl:115
[5] vector_mode_jacobian!(result::Matrix{Float64}, f!::typeof(sys4!), y::Matrix{Float64}, x::Matrix{Float64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}})
@ ForwardDiff ~/.julia/packages/ForwardDiff/CkdHU/src/jacobian.jl:178
[6] jacobian!
@ ~/.julia/packages/ForwardDiff/CkdHU/src/jacobian.jl:78 [inlined]
[7] jacobian!
@ ~/.julia/packages/ForwardDiff/CkdHU/src/jacobian.jl:76 [inlined]
[8] fj!
@ ./REPL[3]:17 [inlined]
[9] upto_hessian(nr::NLSolvers.NormedResiduals{Matrix{Float64}, Matrix{Float64}, NLSolvers.VectorObjective{typeof(sys4!), var"#j!#1"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}}, var"#fj!#2"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}}, var"#jv!#3"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}, Matrix{Float64}}}}, Fx::Matrix{Float64}, Jx::Matrix{Float64}, x::Matrix{Float64})
@ NLSolvers ~/.julia/packages/NLSolvers/nsE1G/src/nlsolve/trustregions/newton.jl:35
[10] upto_hessian
@ ~/.julia/packages/NLSolvers/nsE1G/src/optimize/problem_types.jl:26 [inlined]
[11] prepare_variables(prob::OptimizationProblem{NLSolvers.NormedResiduals{Matrix{Float64}, Matrix{Float64}, NLSolvers.VectorObjective{typeof(sys4!), var"#j!#1"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}}, var"#fj!#2"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}}, var"#jv!#3"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}, Matrix{Float64}}}}, Nothing, NLSolvers.Euclidean{Tuple{0}}, Nothing, NLSolvers.InPlace, Nothing}, approach::TrustRegion{Newton{Direct, typeof(NLSolvers.DefaultNewtonLinsolve), Nothing, Nothing}, Dogleg{Nothing}, NLSolvers.BTR{Nothing}}, x0::Matrix{Float64}, ∇fz::Matrix{Float64}, B::Nothing)
@ NLSolvers ~/.julia/packages/NLSolvers/nsE1G/src/optimize/problem_types.jl:204
[12] solve(problem::OptimizationProblem{NLSolvers.NormedResiduals{Matrix{Float64}, Matrix{Float64}, NLSolvers.VectorObjective{typeof(sys4!), var"#j!#1"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}}, var"#fj!#2"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}}, var"#jv!#3"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}, Matrix{Float64}}}}, Nothing, NLSolvers.Euclidean{Tuple{0}}, Nothing, NLSolvers.InPlace, Nothing}, s0::Tuple{Matrix{Float64}, Nothing}, approach::TrustRegion{Newton{Direct, typeof(NLSolvers.DefaultNewtonLinsolve), Nothing, Nothing}, Dogleg{Nothing}, NLSolvers.BTR{Nothing}}, options::OptimizationOptions{Float64, Float64, Float64, Int64, NLSolvers.var"#10#14", NLSolvers.var"#11#15"})
@ NLSolvers ~/.julia/packages/NLSolvers/nsE1G/src/optimize/trustregions/optimize/inplace_loop.jl:10
[13] solve
@ ~/.julia/packages/NLSolvers/nsE1G/src/optimize/trustregions/trustregions.jl:37 [inlined]
[14] solve(prob::NEqProblem{NLSolvers.VectorObjective{typeof(sys4!), var"#j!#1"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}}, var"#fj!#2"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}}, var"#jv!#3"{typeof(sys4!), ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4, Tuple{Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sys4!), Float64}, Float64, 4}}}}, Matrix{Float64}, Matrix{Float64}}}, Nothing, NLSolvers.Euclidean{Tuple{0}}, NLSolvers.InPlace}, x::Matrix{Float64}, approach::TrustRegion{Newton{Direct, typeof(NLSolvers.DefaultNewtonLinsolve), Nothing, Nothing}, Dogleg{Nothing}, NLSolvers.BTR{Nothing}}, options::NEqOptions{Float64, Int64, Nothing})
@ NLSolvers ~/.julia/packages/NLSolvers/nsE1G/src/nlsolve/trustregions/newton.jl:70
[15] customnlsolve(f!::Function, x0::Matrix{Float64}, method::TrustRegion{Newton{Direct, typeof(NLSolvers.DefaultNewtonLinsolve), Nothing, Nothing}, Dogleg{Nothing}, NLSolvers.BTR{Nothing}}, options::NEqOptions{Float64, Int64, Nothing})
@ Main ./REPL[3]:31
[16] customnlsolve(f!::Function, x0::Matrix{Float64})
@ Main ./REPL[3]:2
[17] top-level scope
@ REPL[11]:1
``````

Maybe you could wrap the output of the input function with `vec`?

wouldn’t I have to change that in-place? is it possible to reshape an array in place?
for example, this doesn’t work:

``````function sys4!(sto, xx)
x = xx[1, 1]
y = xx[1, 2]
z = xx[2, 1]
t = xx[2, 2]
sto[1, 1] = 4.5e-4*x - x*y
sto[1, 2] = z - 1e-14/t
sto[2, 1] = y - 0.02 - x
sto[2, 2] = z - t - y
sto = vec(sto)
return sto
end

customnlsolve(sys4!, vec(x04))
``````

you also need to reshape at input. basically you need to:

``````vector input -> reshape to matrix -> process -> reshape to vector -> output
``````

in your case, you would need to do:

``````function sys4!(sto, xx)
sto = reshape(sto,2,2)
xx = reshape(xx,2,2)
x = xx[1, 1]
y = xx[1, 2]
z = xx[2, 1]
t = xx[2, 2]
sto[1, 1] = 4.5e-4*x - x*y
sto[1, 2] = z - 1e-14/t
sto[2, 1] = y - 0.02 - x
sto[2, 2] = z - t - y
sto = vec(sto)
return sto
end
``````
1 Like

Thank you again mate!
Was there anything I could do to modify directly the solver? I thought playing with the shapes of `Fcache` and `Jcache` would do it.

mmm, i don’t really know. basically all the routines want a `AbstractVector` as input. because the shape of the jacobian is different than the shape of the input. when using gradient-based methods (or in Non Linear solving, jacobian-free methods) the algorithms just iterate over `eachindex`, skipping the shape directly.
On a non-related note. two things (discovered by playing with that nlsolve implementation):

• , specify a chunk size on `jconfig `, there is a type instability there:
``````    jconfig = ForwardDiff.JacobianConfig(f!, x0, x0,ForwardDiff.Chunk{2}())
``````
• You could use `MVector` from StaticArrays.jl, for a good speedup (around 50% faster). the library also has a full out-of-place mode for `SVector` (that i haven’t tested yet).
• the code is being used here.
1 Like

This is interesting, because `NLsolve` which `NLSolvers` will eventually replace, accepts multidimensional arrays as inputs.

The issue here is that I have many equations, around 100, and `StaticArrays` doesn’t scale to that number.

Thank you!

1 Like

yeah, in the NLSolve.jl code, the same `vec` trick is present. for example:

I will open an issue and ask.

but this works:

``````using NLsolve

function sys4!(sto, xx)
#    sto = reshape(sto,2,2)
#    xx = reshape(xx,2,2)
x = xx[1, 1]
y = xx[1, 2]
z = xx[2, 1]
t = xx[2, 2]
sto[1, 1] = 4.5e-4*x - x*y
sto[1, 2] = z - 1e-14/t
sto[2, 1] = y - 0.02 - x
sto[2, 2] = z - t - y
#     sto = vec(sto)
return sto
end

sto4 = zeros(2, 2)
x04 = rand(2, 2)

nlsolve(sys4!, x04)
``````

yes, it should support it in the future, vec is the solution I used in NLsolve

1 Like