NonlinearSolve: how to pass parameters including a mutable workspace array

Specific Domains Optimization (Mathematical)
1

I would like to solve a 1D equation using NonlinearSolve.

Appended below is a toy MNWE followed by the error message generated.

The issue I’m encountering is that in the actual problem I need to pass a workspace array ws_array for the calculation. After reading the documentation and various threads on Julia Discourse, ComponentArrays seemed to be the way to proceed.

# test_NonlinearSolve.jl

using ComponentArrays
using NonlinearSolve

@inline function a_function(
			t::Float64,
			cp_v::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(q = 1, p = 2, g = 3, l = 4, m = 5, w = ViewAxis(6:21, Shaped1DAxis((16,))))}}}
		)::Float64

	@static_unpack q, p, g, l, m = cp_v
	cp_v.w[1] = sin(2.0*q)
	cp_v.w[2] = l^(-1)
	cp_v.w[3] = sin(q)
	cp_v.w[4] = m^(-1)
	cp_v.w[5] = cos(q)
	cp_v.w[6] = cp_v.w[5]^2
	cp_v.w[7] = cp_v.w[3]^2
	cp_v.w[6] =  - (1.0/2.0)*cp_v.w[6] + cp_v.w[7]
	cp_v.w[7] = g^2
	cp_v.w[6] = cp_v.w[4]*cp_v.w[6]*cp_v.w[7]*p^2
	cp_v.w[8] = (1.0/2.0)*cp_v.w[5]
	cp_v.w[9] = cp_v.w[2]^3
	cp_v.w[8] =  - g*p^4*cp_v.w[4]^3*cp_v.w[9]*cp_v.w[8]
	cp_v.w[6] = cp_v.w[6] + cp_v.w[8]
	cp_v.w[6] = cp_v.w[6]*cp_v.w[9]
	cp_v.w[8] = cp_v.w[3]*cp_v.w[5]
	cp_v.w[9] =  - (1.0/3.0)*cp_v.w[1] + (1.0/2.0)*cp_v.w[8]
	cp_v.w[9] = m*cp_v.w[3]*cp_v.w[9]*g^3
	cp_v.w[6] = cp_v.w[9] + (1.0/3.0)*cp_v.w[6]
	cp_v.w[6] = t*cp_v.w[2]*cp_v.w[6]
	cp_v.w[8] = (1.0/2.0)*cp_v.w[1] - cp_v.w[8]
	cp_v.w[7] = p*cp_v.w[8]*cp_v.w[7]*cp_v.w[2]^2
	cp_v.w[6] = (1.0/2.0)*cp_v.w[7] + cp_v.w[6]
	f_t = t^2*cp_v.w[6]
	return f_t
end

function main()
    # Workspace array
    ws_array::Array{Float64, 1} = zeros(Float64, 16)

    # Initial value of scalar t
    t_o::Float64 = 0.5

    cmpnt_vec = 
        ComponentVector(
            q = 0.7552444383914849, 
            p = -0.31175068141804535, 
            g = 6.67408, l = 1.0, m = 0.1, 
            w = ws_array
        )

    @show typeof(cmpnt_vec)
    println("")

    # Test a_function
    @show f_t = a_function(t_o, cmpnt_vec)
    println("")

    # Test NonlinearSolve
    nl_prblm = NonlinearProblem(a_function, t_o, cmpnt_vec)
    @show t_f = solve(nl_prblm, SimpleNewtonRaphson(), abstol = 1.0e-12)
    println("")
end

begin
    main()
end

However, this code generates the following error message:

julia> include("test_NonlinearSolve.jl")
.  .  .
ERROR: LoadError: MethodError: no method matching a_function(::ForwardDiff.Dual{ForwardDiff.Tag{…}, Float64, 1}, ::ComponentVector{Float64, Vector{…}, Tuple{…}})
The function `a_function` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  a_function(::Float64, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(q = 1, p = 2, g = 3, l = 4, m = 5, w = ViewAxis(6:21, Shaped1DAxis((16,))))}}})
   @ Main ~/projects/Hamiltonian/test/test_NonlinearSolve.jl:6
  a_function(::Float64, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(q = 1, p = 2, g = 3, l = 4, m = 5, w = ViewAxis(6:69, Shaped1DAxis((16,))))}}})
   @ Main ~/projects/Hamiltonian/test/test_NonlinearSolve.jl:6

Stacktrace:
  [1] (::NonlinearFunction{…})(::ForwardDiff.Dual{…}, ::Vararg{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/dld4W/src/scimlfunctions.jl:2604
  [2] pushforward(f::NonlinearFunction{…}, backend::AutoForwardDiff{…}, x::Float64, tx::Tuple{…}, contexts::DifferentiationInterface.Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/D3LUI/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:40
  [3] derivative(f::NonlinearFunction{…}, backend::AutoForwardDiff{…}, x::Float64, contexts::DifferentiationInterface.Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/D3LUI/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:206
  [4] compute_jacobian!!(::Nothing, prob::SciMLBase.ImmutableNonlinearProblem{…}, autodiff::AutoForwardDiff{…}, fx::Float64, x::Float64, ::SimpleNonlinearSolve.Utils.DINoPreparation)
    @ SimpleNonlinearSolve.Utils ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/utils.jl:110
  [5] #__solve#32
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/raphson.jl:55 [inlined]
  [6] __solve
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/raphson.jl:34 [inlined]
  [7] #simplenonlinearsolve_solve_up#40
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/SimpleNonlinearSolve.jl:123 [inlined]
  [8] simplenonlinearsolve_solve_up
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/SimpleNonlinearSolve.jl:118 [inlined]
  [9] #solve#39
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/SimpleNonlinearSolve.jl:109 [inlined]
 [10] solve
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/SimpleNonlinearSolve.jl:92 [inlined]
 [11] #solve#36
    @ ~/.julia/packages/SimpleNonlinearSolve/GiCJO/src/SimpleNonlinearSolve.jl:64 [inlined]

which, if I understand correctly, ComponentArrays is not compatible with the default ForwardDiff.

How to solve this issue? Should I use SciMLStructs and, if so, then how?

2

Hi :slight_smile: First of all, thanks for providing a minimal example and the full error message.

Regarding your issue: Try removing all type annotations from the function signature, i.e. turn this

@inline function a_function(
			t::Float64,
			cp_v::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(q = 1, p = 2, g = 3, l = 4, m = 5, w = ViewAxis(6:21, Shaped1DAxis((16,))))}}}
		)::Float64

into this

function a_function(t, cp_v)

What happens here is not that ComponentArrays is incompatible with ForwardDiff, but rather a_function is incompatible. ForwardDiff passes Dual numbers through the functions it wants to differentiate (see e.g. here for an explanation). If you restrict them to only take Float64 as input you will get the error message you posted (note the type of the first argument is not Float64, that’s the main issue

ERROR: LoadError: MethodError: no method matching a_function(::ForwardDiff.Dual{ForwardDiff.Tag{…}, Float64, 1}, ::ComponentVector{Float64, Vector{…}, Tuple{…}})
The function `a_function` exists, but no method is defined for this combination

The second issue that will appear after fixing the signature is that the ComponentVector you use as a “buffer” for storing intermediate values is also fixed to be of type ComponentVector{Float64, …} and ws_array is a Vector{Float64}. So as soon as you try to do something like cp_v.w[…] = t * … it won’t work, because t will be a Dual and not a Float64 (for ForwardDiff to work). But a Vector{Float64} can only accept Float64 values. The solution is to either not use the array at all (you don’t need it, as far as I can see) or to use a “diff cache” which contains/creates arrays of all the different types you might need.

[EDIT: Of course, if you really need the array for some reason the following doesn’t really apply. But perhaps then you could clarify why you need it – there might still be ways around using the array.] Why you don’t need a component array is because you just compute a single number from a fixed set of inputs and operations. In a_function you are basically maintaining your own “function stack” by re-using all the entries of the vector w. This is quite confusing and most likely slower than just using the real function stack (i.e. just make regular variables inside the function)… Just use a NamedTuple for all the parameters and get rid of the vector all together (you can search\replace all cp_v[i] by some variable w_i if you want, or, ideally, try to make it more transparent what happens. Something like this:

# test_NonlinearSolve.jl

using NonlinearSolve

function b_function(
			t,
			parameters
		)

	(; q, p, g, l, m) = parameters
	w_1 = sin(2.0*q)
	w_2 = l^(-1)
	w_3 = sin(q)
	w_4 = m^(-1)
	w_5 = cos(q)
	w_6 = w_5^2
	w_7 = w_3^2
	w_6 =  - (1.0/2.0)*w_6 + w_7
	w_7 = g^2
	w_6 = w_4*w_6*w_7*p^2
	w_8 = (1.0/2.0)*w_5
	w_9 = w_2^3
	w_8 =  - g*p^4*w_4^3*w_9*w_8
	w_6 = w_6 + w_8
	w_6 = w_6*w_9
	w_8 = w_3*w_5
	w_9 =  - (1.0/3.0)*w_1 + (1.0/2.0)*w_8
	w_9 = m*w_3*w_9*g^3
	w_6 = w_9 + (1.0/3.0)*w_6
	w_6 = t*w_2*w_6
	w_8 = (1.0/2.0)*w_1 - w_8
	w_7 = p*w_8*w_7*w_2^2
	w_6 = (1.0/2.0)*w_7 + w_6
	f_t = t^2*w_6
	return f_t
end

function main()
    t_o = 0.5

    parameters = (;
            q = 0.7552444383914849, 
            p = -0.31175068141804535, 
            g = 6.67408, l = 1.0, m = 0.1, 
        )

    # Test b_function
    @show f_t = b_function(t_o, parameters)
    println("")

    # Test NonlinearSolve
    nl_prblm = NonlinearProblem(b_function, t_o, parameters)
    @show t_f = solve(nl_prblm, SimpleNewtonRaphson(), abstol = 1.0e-12)
    println("")
end

begin
    main()
end

Note that removing the type annotations in the function signature has no impact on performance, but makes all the difference for writing generic code (working with regular floats and also duals). Unless you need dispatch (multiple methods which do different things based on the input type) type annotations of function arguments are usually unnecessary and often harmful when composing code (like here).

Also note that I removed @inline. No idea if it will have an impact here, but I doubt it (in my experience it was rarely necessary to speed things up, but you can test the speed of the solve with and without it if you are unsure).

1 Like
3

Hello. Thank for your detailed explanation and example solution.
Much appreciated.

I was wondering why the type annotations were not specified for functions in the NonlinearSolve documentation, now I understand that it’s due to the passing of Dual numbers as arguments.

Also, you’re quite right that I don’t need the workspace array ws_array. In the full project, the code is first computed using Symbolics.jl. It is then processed by a non-Julia symbolic application called FORM which reduces the number of floating point operations by about a factor of 8. The output is generated as Fortran code [which is converted back to Julia]. The Fortran compiler prefers arrays to scalars, but I can switch this to generate scalars instead as in your solution.

As an aside, I’ve yet to find a Julia package that can reduce floating point ops via symbolic code optimization as effectively as FORM. Are you aware of any such Julia package?

1 Like
4

Great!

Hmm, I don’t know any package like that, no. Symbolics.jl has the simplify function with various options and also allows to define custom term-rewriting rules. I don’t think the default simplifications are quite as powerful as what you are describing though. But there might be packages that implement something like Horner’s method (like here). There is also DynamicExpressions.jl which might be useful in implementing something like this directly in Julia. But if FOAM already exists for that purpose, just calling it like you do now is probably the easiest…

Just out of curiosity: How large of expressions are we talking about and does the rewriting really have a significant impact on the runtime in your final application? It might be best to open a new question with some more details on these details. Other people with more experience on this kind of optimization might also chime in :slight_smile: