Efficient way to implement multi-index nonlinear constraints in ExaModels?

[Yanlingogo]

Hi all,

I’m implementing a PCE-expanded AC power flow model. In JuMP, one of my nonlinear constraints in JuMP form 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
        )
)

And the code in ExaModel is below, the main difference is that I flatten the elements in T3_zip:

constraint(
    core,
    (
        let i = i, k = k,
            l1_arr = T3_zip.l1,
            l2_arr = T3_zip.l2,
            val_arr = T3_zip.val,
            p_start = p_start,
            p_end   = p_end

            # active injection at bus i and chaos index k
            p_ik = sum(Cgpar[i, g] * Pgen[g, k] for g in 1:n_gen) - Pdpar[i, k]

            rhs = sum(
                (
                    Gpar[i, j] * (E[i, l1_arr[p]] * E[j, l2_arr[p]] +
                                  F[i, l1_arr[p]] * F[j, l2_arr[p]]) +
                    Bpar[i, j] * (F[i, l1_arr[p]] * E[j, l2_arr[p]] -
                                  E[i, l1_arr[p]] * F[j, l2_arr[p]])
                ) * val_arr[p]
                for p in p_start:p_end, j in 1:n_bus;
                init = 0.0
            )

            p_ik * T2par[k] - rhs
        end
        for i in 1:n_bus
    );
    lcon = 0.0,
    ucon = 0.0
)

I found the function ExaModel(JuMPmodel) works well. I think if I directly code the model for ExaModel, the execution will be more efficient, but results turned out to be different. the model construction with second part code 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!

[sshin23]

The best way to handle summations within constraint is to use constraint!. A tutorial can be found here:

https://madsuite.org/powertech2025/dev/1-powerflow/#Implementing-the-power-flow-equations-with-ExaModels