I am writing a JuMP model, which has a quadratic term in the objective function. It looks something like the following:
@expression(model, term, sum( ( a[i] - b[i] * sum( x[(i,j2)] for j2 in dict_j[i] ) )* x[(i,j)] for (i,j) in term_idxs ) );
where x[(i,j)] is a sparse variable (with a pre-constructed subset of I x J called “term_idxs”, as also referred to in the performance considerations of Variables · JuMP), a/b are parameters, and dict_j[i] returns a subset of j to sum over in the inner sum.
The original model term is more complex and the real x has a few more dimensions, but the code above serves as an illustrative example.
Even though x is sparse, there are still quite a few feasible pre-constructed combinations. Generating the expression above takes more than 50% of the model generation time, which made me wonder about parallelizing the summation.
I know that parallelism in general seems to be difficult in the context of JuMP models, but is it also if I simply want to compute the sum in parallel and add it as an expression?
I messed around a little with different ideas, but except for “higher than single threaded CPU usage”, I did not achieve any improvements. For example, a smaller scale example using Folds.jl · Folds did produce an equivalent expression, but it was significantly slower than the original implementation. This is because computing a QuadExpr and passing it to the @expression macro seems to be significantly slower also in the non-parallel case.
Despite searching for quite a bit, I did not really find any approaches to parallelizing a sum in JuMP models.
Is there a recommended way of handling this and if yes, what would it look like?
Thanks a lot to everyone answering topics on this forum, it has proven really helpful in the past!
Btw, this is my first post here, so please give me a note in case I made any mistakes.
Welcome! It’s a little easier to provide feedback if you can produce a minimal working example: Please read: make it easier to help you. Without one, we have to make some guesses. What are a and b? What is term_idxs? Are there lots of (i, j) with the same i? How many terms are we talking about?
I did not really find any approaches to parallelizing a sum in JuMP models.
There is no way to use parallelism to make problem construction faster.
One suggestion I would try is to factor out the terms that do not depend on j. Assuming I didn’t make a mistake (I didn’t run the code because I don’t have the data):
I = unique(first.(term_idxs))
@expression(model, term_i[i=I], sum(x[(i, j)] for j in dict_j[i]))
@expression(model, sum((a[i] - b[i] * term_i[i]) * term_i[i] for i in I))
This should greatly reduce the number of loop iterations. Otherwise you’re looping over all term_idxs and for each term you’re looping over dict_j[I].
You should also read the Julia performance tips if you haven’t already: Performance Tips · The Julia Language. I assume you aren’t using global variables, etc?
Thanks a lot for the quick reply @odow.
Pre-calculating the sum is actually a great solution in my case, can’t believe I did not think of that before! In the real case, it did cut generation time by ~90%.
Just a quick note for future readers of this, there is a minor issue with the reformulation. It should instead look something like this:
@expression(model, term_i[i in I], sum( x[(i,j2)] for j2 in dict_j[i] ) )
@expression(model, term_fast, sum( ( a[i] - b[i] * term_i[i])* x[(i,j)] for (i,j) in term_idxs ) );
I have also put together a quick MWE:
using JuMP
I = 1:1000
J = 1:1000
term_idxs = Set( (i,j) for i in I, j in J)
dict_j = Dict(i => Set( rand(1:100) for n in 1:3 ) for i in I)
a = Dict(i => rand(1:100) for i in I)
b = Dict(i => rand(1:100) for i in I)
model = Model()
@variables model begin
x[term_idxs] >= 0
end;
@expression(model, term_orig, sum( ( a[i] - b[i] * sum( x[(i,j2)] for j2 in dict_j[i] ) )* x[(i,j)] for (i,j) in term_idxs ) );
@expression(model, term_i[i in I], sum( x[(i,j2)] for j2 in dict_j[i] ) );
@expression(model, term_fast, sum( ( a[i] - b[i] * term_i[i])* x[(i,j)] for (i,j) in term_idxs ) );
@assert term_orig == term_fast
Of course using Tuples of (i,j) does not really make sense in the MWE since x is not sparse…