Modify custom solver to deal with multidimensional arrays

General Usage
,
1

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
2

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

3

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))
4

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
5

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.

6

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
7

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
8

yeah, in the NLSolve.jl code, the same vec trick is present. for example:
https://github.com/JuliaNLSolvers/NLsolve.jl/blob/219672802ed79b3d8ae13fdc83febce3bd0acd37/src/solvers/trust_region.jl#L51

9

I will open an issue and ask.

EDIT: here https://github.com/JuliaNLSolvers/NLSolvers.jl/issues/25

10

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)
11

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

1 Like