How to speed up AffineExpr generation?

Specific Domains Optimization (Mathematical)
1

I have a problem with 3000×3000 matrix of JuMP variables, and several millions of constraints in those. In the naive formulation the size is what obstructs solving the problem in reasonable time;

I decided to play smart and broken-up the problem to several smaller matrices (P[i]'s) and transformations (U[i]'s) reducing the number of JuMP variables to ~10000. However now I have a few dozens of 3000×3000 matrices of JuMP’s GenericAffExpr{Float64,JuMP.Variable} to generate and sum. Doing this directly takes ridiculous amounts of memory (I’ve run out of at 30G) and time;

I suspect I do something wrong with problem generation; here is the code:

m = JuMP.Model();
l = size(U,1)
P = Vector{Any}(l)
for i in 1:len
    u = size(U[i],2)
    JuMP.@variable(m, P[i][1:u,1:u], SDP)
    JuMP.@SDconstraint(m, P[i] >= 0)
end

this generates len SDP matrices (P) of the appropriate sizes which I transform using matrices U and assemble to the original one:

PP = JuMP.@expression(m, sum(U[k]*P[k]*(U[k]') for k in 1:len))

I know this is not the efficient way (? because of matrix multiplication which involves a lot of + between AffExpr, so should I write my own matrix multiplication routines using AffExpr and/or @expression, ala

*{T}(A::Array{T, 2}, B::Array{JuMP.Variable,2}) = ...
*{T}(A::Array{JuMP.GenericAffExpr{T,JuMP.Variable}, 2}, B::Array{T,2}) = ...

or there is a clever way around this?

2

You didn’t post the entire model. Depending on what other constraints are present and how your objective function is formulated, you might get better behavior if you change variables so that Z[i] defined to be U[i]*P[i]*U[i]', where Z[i] is the new decision variable. Note that P>=0 implies Z>=0, and the converse is true if the columns of U[i] are independent.

3

The rest of the model is not involved, but here is the whole problem:

m = JuMP.Model();
l = size(U,1)
P = Vector{Any}(l)
for i in 1:len
    u = size(U[i],2)
    JuMP.@variable(m, P[i][1:u,1:u], SDP)
    JuMP.@SDconstraint(m, P[i] >= 0)
end

PP = JuMP.@expression(m, sum(U[k]*P[k]*(U[k]') for k in 1:len))

JuMP.@variable(m, x >= 0)
for (A, b, c) in zip(Ai's, B, C)
    JuMP.@constraint(m, sum(PP[i,j] for (i,j) in A) == b - x*c)
end
JuMP.@objective(m, Max, x)

So this is standard formulation as ∑ᵢⱼ AᵢⱼₖPPᵢⱼ + Cₖx = Bₖ, or if You prefer ⟨Aₖ,PP⟩ + Cₖx = Bₖ;

I didn’t get the Z[i] thing. If I set the Z[i] to be the new this is exactly what was happening before, because the number of variables (size of Z[i]) is huge. The whole point of doing this is to exploit structure of PP and “group variables known to be the same under the same name”.

U[i]'s are orthogonal.

PS. I tried to replace matrix multiplication with this:

function getcoefficient(k::Int, i::Int, j::Int)
    b = Array{JuMP.GenericAffExpr{Float64,JuMP.Variable}, 1}(size(U[k],2));
    for l in 1:size(P[k],2)
        b[l] = JuMP.AffExpr(P[k][:,l], full(U[k][i,:]), 0.0)
    end
    return dot(b, U[k][j,:])
end

getcoefficient(model::JuMP.Model, i::Int, j::Int) = JuMP.@expression(model, sum(getcoefficient(k,i,j) for k in 1:length(U)))

however getcoefficient still takes ~ 0.5 ms

4

The standard Julia performance tips apply to generating JuMP models:

  1. Don’t run things in global scope
  2. Use containers with concrete types (i.e., not Vector{Any}(l))
  3. Avoid temporary objects and vectorized operations. This applies particularly to JuMP because operations involving AffExpr are expensive and their cost grows with the size of the expression.
  4. Use the memory and time profiling tools to identify bottlenecks

Side note,

    JuMP.@SDconstraint(m, P[i] >= 0)

is redundant here because the SDP tag tells JuMP to enforce PSDness of the matrix of variables.

1 Like
5

Thanks for the tips;
All the AffExpr generation happens inside @expression macro;
model generation is actually wrapped inside a function;
generation of m as in the first post takes 0.000423s, it’s not the issue of P being Vector{Any}
(in fact, while computing U[k]*P[k]*(U[k]') julia correctly infers all the types);

Question:
Can JuMP perform any type of arithmetic on variables to simplify the expressions (i.e. reduce length of .vars and .coeffs)? The problem as I perceive right now is that the expressions I generate are of length exceeding millions (with only ~10000 actual variables).

6

If your expressions are completely dense, you may consider collecting all the coefficients manually into a dense vector and then generating each expression via dot(coefs,vars) at the end. JuMP’s data structures are designed for sparse problems and we don’t collect duplicate coefficients until the end when we generate the problem.

7

If I understand correctly, the issue you face is as follows: you have linear constaints that are written with matrix-multiplications UPU’, P is the variable, U is the coefficient, on both sides. If you wrote the linear operator explicitly as a matrix, it would be a huge Kronecker-product matrix kron(U,U). So you would prefer to keep the solver from ever explicitly forming kron(U,U) and instead work implicitly with the linear constraint.

There is certainly academic work on this issue: One idea is to use iterative rather than direct methods to solve the linear system that arises in interior point methods, in which case the implicit representation of the constraint should be fine. Another idea is to abandon interior-point methods completely and switch to a first-order method. However, I don’t know if JuMP supports any of this, and I also don’t know which interior-point solvers support iterative methods for linear systems for an implicitly-specified constraint. Possibly Miles Lubin could be of more assistance.

8

@miles.lubin, thanks, indeed the growing expressions proved to be the problem on the memory side; I wrote something like this to keep expressions from growing:

function gather{T}(vars::Vector{JuMP.Variable}, coeffs::Vector{T})
    uvars = unique(vars)
    new_coeffs = Vector{T}(length(uvars))
    for i in 1:length(uvars) 
        new_coeffs[i] = sum(coeffs[vars .== uvars[i]])
    end
    return uvars, new_coeffs
end

function gather!(c::JuMP.AffExpr)
    c.vars, c.coeffs = gather(c.vars, c.coeffs)  
    return c
end

gather(c::JuMP.AffExpr) = JuMP.dot(gather(c.vars, c.coeffs)...)

I hoped to reuse the old code, but I should probably avoid forming the expression completely and directly compute the gathered version.

9

In the context of SDP, JuMP’s functionality can be boiled down to generating the explicit A matrix, b and c vectors and supporting data structures of the MathProgBase conic format which a number of solvers recognize. It’s unlikely this will change substantially any time soon, and I leave it up to users to decide if they need a tool that performs this job. What you’re mentioning is of course interesting and sounds similar to what Steven Diamond has worked on (paper).

10

@Stephen_Vavasis, Yes I tried to not form the product kron(U,U). But the right solution (maybe only for my, particular problem) turned out to be easier than suggested “Matrix-free optimisation”.

Instead of transforming matrices of JuMP.Variables, I can actually transform the constraints in the opposite direction. This requires that all As, B and C are amenable to the same symmetry that guided the decomposition of PP into P[i]'.

(the first solution with gather shortening the arising AffExprs works but is slower)