Reducing number of allocations

Specific Domains Optimization (Mathematical)
1

I am playing around with JuMP to compare its performance to pyomo. This is the model I am using:

using JuMP
using GLPK
using LinearAlgebra
using BenchmarkTools

function run_market_opt(max_production::AbstractVector{T}, bids::AbstractVector{T}, demand::T) where T <: Real
    n_plants = length(max_production)
    model = Model(GLPK.Optimizer, add_bridges=false)
    set_string_names_on_creation(model, false)
    set_silent(model)

    @variable(model, production[1:n_plants] >= zero(T))
    @constraint(model, production .<= max_production)
    @constraint(model, sum(production) == demand)
    @objective(model, Min, dot(production, bids))

    optimize!(model)
    @assert is_solved_and_feasible(model)

    return value.(production)
end

##
n_plants = 10000
max_production = 100 .* rand(Float64, n_plants)
bids = rand(Float64, n_plants) 
demand = 30.0 * n_plants

@btime production_jump = run_market_opt(max_production, bids, demand);

with this I get:

1.018 s (350848 allocations: 25.42 MiB)

I was expecting JuMP to be faster than pyomo. This is the equivalent pyomo model:


import numpy as np
import pyomo.environ as pyo
from time import time

def run_market_opt(max_production, bids, demand):
    n_plants = len(max_production)

    model = pyo.ConcreteModel()
    model.I = pyo.RangeSet(0, n_plants - 1)

    model.production = pyo.Var(model.I, domain=pyo.NonNegativeReals)

    def max_prod_rule(model, i):
        return model.production[i] <= max_production[i]

    model.max_production_constraint = pyo.Constraint(model.I, rule=max_prod_rule)

    def demand_rule(model):
        return sum(model.production[i] for i in model.I) == demand

    model.demand_constraint = pyo.Constraint(rule=demand_rule)

    def obj_rule(model):
        return sum(model.production[i] * bids[i] for i in model.I)

    model.objective = pyo.Objective(rule=obj_rule, sense=pyo.minimize)

    solver = pyo.SolverFactory("glpk")
    results = solver.solve(model, tee=False)

    if (results.solver.status == pyo.SolverStatus.ok) and (
        results.solver.termination_condition == pyo.TerminationCondition.optimal
    ):
        return np.array([pyo.value(model.production[i]) for i in model.I])
    else:
        raise RuntimeError("Optimization failed to find a solution")


# %%
n_plants = 10000
max_production = 100 * np.random.rand(n_plants)
bids = np.random.rand(n_plants)
demand = 30 * n_plants

# %%
t1 = time()
production_pyomo = run_market_opt(max_production, bids, demand)
t2 = time()
print(f"Pyomo took {t2 - t1} seconds")

which gives

Pyomo took 0.874971866607666 seconds

Is there something I can do to speed-up the JuMP version? Could I do something to help reduce the number of allocations?

Thanks!

2

This should allocate a vector each time, maybe try

all(zip(production, max_production)) do (p, m)
    p <= m
end

For a functional version that should not allocate and return early

3

Thanks for the suggestion. Is this what you had in mind?

    for (p, m) in zip(production, max_production)
        @constraint(model, p <= m)
    end

I still get the same number of allocations.

4

I would write your code like:

julia> using JuMP

julia> using GLPK

julia> function run_market_opt(max_production, bids, demand)
           n_plants = length(max_production)
           model = Model(GLPK.Optimizer)
           set_silent(model)
           @variable(model, 0 <= production[i in 1:n_plants] <= max_production[i])
           @constraint(model, sum(production) == demand)
           @objective(model, Min, bids' * production)
           optimize!(model)
           @assert is_solved_and_feasible(model)
           @show solve_time(model)
           return value.(production)
       end
run_market_opt (generic function with 1 method)

julia> begin
           n_plants = 10000
           max_production = 100 .* rand(Float64, n_plants)
           bids = rand(Float64, n_plants) 
           demand = 30.0 * n_plants
       end;

julia> @time production_jump = run_market_opt(max_production, bids, demand);
solve_time(model) = 0.7042160034179688
  0.793588 seconds (342.42 k allocations: 24.595 MiB, 7.68% compilation time)

julia> @time production_jump = run_market_opt(max_production, bids, demand);
solve_time(model) = 0.7236940860748291
  0.756984 seconds (282.69 k allocations: 20.359 MiB)

julia> @time production_jump = run_market_opt(max_production, bids, demand);
solve_time(model) = 0.7511911392211914
  0.804027 seconds (282.69 k allocations: 20.359 MiB, 2.72% gc time)

Note that almost all the time is spent inside the solver. JuMP will be about the same speed as Pyomo for this trivial problem.

Also note that you can do much better than GLPK:

julia> using JuMP

julia> using HiGHS

julia> function run_market_opt_highs(max_production, bids, demand)
           n_plants = length(max_production)
           model = Model(HiGHS.Optimizer)
           set_silent(model)
           set_attribute(model, "presolve", "off")  # HiGHS presolve is actually slow
           @variable(model, 0 <= production[i in 1:n_plants] <= max_production[i])
           @constraint(model, sum(production) == demand)
           @objective(model, Min, bids' * production)
           optimize!(model)
           @assert is_solved_and_feasible(model)
           @show solve_time(model)
           return value.(production)
       end
run_market_opt_highs (generic function with 1 method)

julia> @time production_jump = run_market_opt_highs(max_production, bids, demand);
solve_time(model) = 0.06296286405995488
  0.169285 seconds (342.61 k allocations: 25.097 MiB, 34.77% compilation time)

julia> @time production_jump = run_market_opt_highs(max_production, bids, demand);
solve_time(model) = 0.057109951972961426
  0.104776 seconds (282.72 k allocations: 20.856 MiB)

julia> @time production_jump = run_market_opt_highs(max_production, bids, demand);
solve_time(model) = 0.05831647792365402
  0.103382 seconds (282.72 k allocations: 20.856 MiB)

For some reason, HiGHS presolve is very slow for this model, so I turned it off. (Because the model is trivial, it actually helps to just go straight to a simplex solve, which is what GLPK does.) In larger examples, you’ll want to leave presolve on.

1 Like
5

Thank you, this is very helpful

1 Like