User defined functions with vector constraints

I’ve been thrown in the deep end with a problem that looks to be possible with current Julia tooling, but being new to the language and this particular subset of the mathematics I’m having a hard time.

In a MWE format: I have a set of equations

\begin{align} \frac{\rm{d}T}{\rm{d}t} &= f_T(A(t), T(t), D(t), O(t)) \\ \frac{\rm{d}M}{\rm{d}t} &= f_M(A(t), M(t), D(t)) \\ \frac{\rm{d}S}{\rm{d}t} &= f_S(E(t), M(t), D(t)) \\ \frac{\rm{d}D}{\rm{d}t} &= f_D(A(t), D(t)) \\ S &= A + T +M\\ \end{align}

Where E(t) and O(t) are known. Solving this set for T, M, S, D, A can be achieved by using the DAEProblem method of DifferentialEquations.jl since this is clearly a Differential Algebraic Equation form, for a given time span (this is in years, so tspan = (1750, 2100) is a good enough example.

This part is implemented, working & fine.

Additionally, I have a nonlinear optimisation problem which I have running through JuMP.jl using the Ipopt.jl sover. There are 21 variables (each of which are vectors representing time) and 18 constraints (excluding simple start / end conditions).

I wont dump the whole lot of this here, but two examples of variables and constraints I have are:

@variable(model, Λ[1:N]);
@NLconstraint(model, [i=1:N], Λ[i] == YGROSS[i] * θ₁[i] * μ[i]^θ₂);
@variable(model, K[1:N] >= 100.0);
@constraint(model, [i=1:N-1], K[i+1] <= (1-δk) * K[i] + I[i]);

N = 100 and the rest of the values are defined elsewhere.

This model is also functioning well, and has been confirmed to be giving me valid output.

Now, I wish to merge the two systems and am hitting a roadblock in my understanding of some of the JuMP documentation.

The DAE system uses pre-calculated values of E(t) and O(t), which I’d ultimately like to generate from my NLOpt system. The ways I’ve considered to do the merge so far:

  • Using callbacks from JuMP, sending values to DAESolve each step, then calling back to JuMP. Doesn’t make too much sense because the callback system isn’t implemented for Ipopt. Perhaps on a lower level? But that seems a little awkward in general.
  • Discretize the DAE system and incorporate it into the NLOpt structure
  • Refactor and make the NLOpt problem continuous

The third option may be on the cards in the future, but for the moment I’d be happy for a solution. So point 2 here makes most sense.

I can easily add the algebraic constraint to the JuMP model to capture A.

@constraint(model, [i=1:N], A[i] == S[i] - T[i] - M[i]);

For the rest I tried an Explicit Euler approach which is frankly too noisy. I know I should be using Sundials.jl or NLsolve.jl (in the future) if I want to use an Implicit Euler method, but then I have issue following how I could arrange my DAE so it fits into JuMP syntax whilst invoking Sundials for the root finding.

Which brings me to my question: is the ForwardDiff.jl AD something I can leverage here?

For example:

ΔS(E, M, D) = ...
JuMP.register(model, :ΔS, 3, ΔS, autodiff=true);

From what I read in the docs concerning nonlinear user-defined functions, I can register scalar functions, not vector ones. That should be fine, but what I don’t get is what value the result should be assigned to if I indeed have vectors.

@NLconstraint(model, [i=1:N], S[i] == ΔS(E[i], M[i], D[i]));
@NLconstraint(model, [i=1:N-1], S[i+1] == ΔS(E[i], M[i], D[i]));
@NLconstraint(model, [i=1:N-1], S[i+1] == S[i] + step*(ΔS(E[i], M[i], D[i])));

Realistically, I could be asking the wrong question here too, so any nudges in the right direction would be greatly appreciated.