NonlinearSolve with in-place mutation of intermediate variables

General Usage Performance
1

I’m working on an economic model involving a dynamic programming problem. Solving this problem involves mutating several large’ish (e.g., 100 x 3 x 15 x 200) arrays at each iteration. To make the code efficient, I set up immutable structures containing these arrays and then mutate the elements in the arrays on each iteration to avoid allocations. That part of the code works well and runs really fast.

Now I want to internally calibrate some parameters so that the solution to the model matches some empirical moments. I have been using NonlinearSolve.jl for this purpose, setting up an equation system where the equations are the deviations of model-predicted moments from their empirical counterparts.

Below is a minimal working example:

using Parameters
using Accessors
using NonlinearSolve

@with_kw struct ParExo{T} #Parameters to be calibrated externally to nonlinear system problem
    shiftY::Vector{T} = [0.0; 0.0]
end

@with_kw struct ParEndo{T}  #Parameters to be calibrated internal to nonlinear system problem
    shiftX::Vector{T} = [0.0; 0.0]
end

@with_kw struct ObjectiveValues{T} #Container for objective values (to avoid allocations)
    fv::Vector{T} = [0.0; 0.0]
end

function ModMoments!(res,x,p,pEndo,pExo) 

    @reset pEndo.shiftX[:] .= x  #Reset the internally calibrated parameters

    T = typeof(x[1])

    OV = ObjectiveValues{T}()

    OV.fv[:] .= ([0.0; 0.0] .+ pEndo.shiftX[:]).^3 .+ pExo.shiftY[:]

    res = OV.fv[:] .- [0.0; 0.0]

    return nothing

end


function Calib(PEndo)

    #Do some external parameters
    pExo = ParExo()
    pExo.shiftY[:] .= [3.0; 3.0]

    #x0 = [1.2; 1.2]
    x0 = [0.1; 0.1]

    p = 1.0

    prob = NonlinearProblem{true}((res, x, p) -> ModMoments!(res, x,p,PEndo,pExo),x0,p)  #Need to set the "inplacedness manually to ensure type stability'

    sol = solve(prob, NewtonRaphson())

    return sol

end

Calib(ParEndo())

The ObjectiveValues structure here is a stand-in for the structure containing the large number of arrays that need to be mutated as part of the dynamic programming problem solution.

In order for the automatic differentiation to work, I had to create a new instance within the ModMoments! function with the same type as x so that it would accept a ForwardDiff.Dual type. But I think this means that the structure is being recreated each time the function is called.

Is there a better way of doing this?

2

Looks like a problem for Home · PreallocationTools.jl

3

Thank you! That does indeed look like the solution.

I’ll have to do some work to figure out how to make the syntax work, but it seems like what I need.