Efficient way to implement multi-index nonlinear constraints in ExaModels? (PCE-expanded AC power flow)

Specific Domains Optimization (Mathematical)
, ,
1

Hi all,

I’m implementing a PCE-expanded AC power flow model. In JuMP, one of my nonlinear constraints looks like:

@constraint(opf, pfP[i=1:n_bus, k=1:K],
    p[i,k] * T2_diag[k] ==
        sum(
            (G[i,j]*(E[i,l1]*E[j,l2] + F[i,l1]*F[j,l2])
             + B[i,j]*(F[i,l1]*E[j,l2] - E[i,l1]*F[j,l2])) * val
            for (l1, l2, val) in T3_zip[k], j in 1:n_bus
        )
)

This works well in JuMP and converts to ExaModels quickly.

However, when I try to implement the same constraint directly in ExaModels (using generators, flattened tensor data, or custom kernel functions), the model construction becomes much slower and the resulting nnzj/nnzh are 3–7× larger than the JuMP-generated ExaModel.

So my questions are:

  • Is there a recommended pattern for writing large multi-index nonlinear sums in ExaModels?
  • Why does JuMP → ExaModels generate a much smaller expression graph than my hand-written version?
  • Are there examples or documentation on expressing tensor-structured nonlinear constraints efficiently in ExaModels?

Any advice or example code would be greatly appreciated.

Thanks!

2

With ExaModels, model construction will generally be faster with fewer constraint and objective calls. If you have something like

for i in 1:ncons
    constraint(m, x[i]*A[i])
end

model construction (and first jacobian & hessian evaluation) will be faster if it is rewritten with a single constraint call. There is some discussion here.

The nnzh reported by ExaModels is misleading as the Hessian is stored in a non-compact format. The Hessian is stored as a vector of values and two vectors of indices with

H[i,j] = sum(hvals[k] for k in eachindex(hvals) if hrow[k]==i && hcol[k]==j)

hrow and hcol can and often do contain repeats. The reported nnzh is simply the length of hvals and not the true number of non-zeros in the Hessian. I think nnzj works similarly. The Jacobian is definitely stored in the same format, but I haven’t checked that nnzj is simply the length of jvals.

I’m not sure how best to write large multi-index sums. I started a related discussion, and there are two related issues, 102 and 167.

You might get some more concrete advice if you share your code or, better yet, a small working example.