How to create an user-defined function with multiple vector as inputs

Hi @PatrickeTownsend, welcome to the forum.

I think you’re slightly off in how you’re you’re approaching JuMP’s user-defined functions.

With a slightly simpler example, your code looks something like this:

using JuMP
function state_f(x, y...)
    x[1] = sum(sin(y[i]) for i in 1:3)
    x[2] = sum(cos(y[i]) for i in 1:3)
end
model = Model()
register(model, :state_f, 4, state_f; autodiff = true)
@variable(model, x[1:2])
@variable(model, y[1:3])
state_f(x, y...)
@optimize(model, Min, sum(x))

This has a few things wrong with it:

  • User-defined functions must take scalars as input
  • User-defined functions must return a scalar
  • User-defined functions must be used in @NL macros

This means that you cannot write a function like StateFunc which takes a vector as the first argument and modifies in-place, and it also means that calling StateFunc and Collocation_state outside the macros is not adding a constraint to the model. It’s just evaluating the function.

Relatedly:

It is possible to create an user defined function with multiple arrays as inputs?

is “no” because you cannot have a user-defined function which takes even one array as input :smile:.

To change my example somewhat, here’s something that you could write in JuMP:

using JuMP
state_f1(y...) = sum(sin(y[i]) for i in 1:3)
state_f2(y...) = sum(cos(y[i]) for i in 1:3)
model = Model()
register(model, :state_f1, 3, state_f; autodiff = true)
register(model, :state_f2, 3, state_f; autodiff = true)
@variable(model, x[1:2])
@variable(model, y[1:3])
@NLconstraint(model, x[1] == state_f1(y...))
@NLconstraint(model, x[2] == state_f2(y...))
@optimize(model, Min, sum(x))

Now writing out state_f1, state_f2 etc can be painful. So you can write functions which return a vector as output, but then you need to follow this tutorial: Tips and tricks · JuMP

1 Like