Best practices for use of structs with NonlinearSolve.jl

General Usage
1

I’m working on an economic model where I need to solve a system of equations.

Most of those equations are defined implicitly from a sort of mean-field game. Basically I construct several multidimensional arrays, construct an interpolated function from them, and then integrate that function over a discretized distribution. The solution of the system is the difference between those integrals and a vector of constants.

My code can be a lot neater and often faster (due to inplace mutation and stack allocation) when I organize all those arrays within a struct.

However, I’ve been struggling to get this to work well with automatic differentiation and NonlinearSolve.jl

The gist of the issue seems to be that NonlinearSolve sometimes passes a ForwardDiff dual type through only some of the variables to be solved for in the system.

As a result, I often need to do something like:

@with_kw struct mystruct{T1,T2,T3}
    A1::Array{T1,2} = fill(0.0, 2, 100)
    A2::Array{T2,2} = fill(0.0, 2, 100)
    A3::Array{T3,3} = fill(0.0, 2, 100)
end

And then within the function that sets up the system I have to do something like

function mysystem(x)
    S = mystruct{eltype(x[1]),eltype(x[2]),eltype(x[3])}()
    ...
end

This works okay but it’s cumbersome because I usually have many more than three arrays and because I don’t always know which element of x I should use to type each element of the struct.

Is it possible to somehow force the x vector to be all the same type? Alternatively, can I somehow promote non-dual type elements to dual and preserve type stability?

2

You could also create your struct with a single type parameter and promote the input types inside that struct’s (inner) constructor?