Проблемы производительности с формулировками «суммировать, если»

Цель данного руководства — описание распространенной проблемы производительности, которая может возникнуть при суммировании вида sum(x[a] for a in list if condition(a)). Эта проблема особенно часто возникает в моделях с графовой или сетевой структурой.

Это более сложное руководство, чем другие руководства по началу работы. Оно помещено в раздел «Начало работы», так как это одна из наиболее распространенных причин проблем с производительностью, с которыми сталкиваются пользователи, когда начинают использовать JuMP для написания крупных программ. Если вы только приступаете к работе с JuMP, то можете бегло просмотреть это руководство и вернуться к нему, когда напишете несколько моделей JuMP.

Требуемые пакеты

В этом руководстве используются следующие пакеты.

using JuMP
import Plots

Данные

В качестве наглядного примера рассмотрим задачу о сетевом потоке, похожую на задачи в разделах Network flow problems и The network multi-commodity flow problem.

Вот функция, которая строит случайный граф. Конкретные детали не имеют значения.

function build_random_graph(num_nodes::Int, num_edges::Int)
    nodes = 1:num_nodes
    edges = Pair{Int,Int}[i - 1 => i for i in 2:num_nodes]
    while length(edges) < num_edges
        edge = rand(nodes) => rand(nodes)
        if !(edge in edges)
            push!(edges, edge)
        end
    end
    function demand(n)
        if n == 1
            return -1
        elseif n == num_nodes
            return 1
        else
            return 0
        end
    end
    return nodes, edges, demand
end

nodes, edges, demand = build_random_graph(4, 8)

(1:4, [1 => 2, 2 => 3, 3 => 4, 2 => 2, 3 => 3, 1 => 1, 3 => 2, 1 => 4], Main.demand)

Цель — определить поток товара через каждое ребро в edges для удовлетворения спроса demand(n) каждого узла n в nodes.

Математическая формулировка:

Примитивная модель

Вначале у вас может получиться такая модель:

model = Model()
@variable(model, flows[e in edges] >= 0)
@constraint(
    model,
    [n in nodes],
    sum(flows[(i, j)] for (i, j) in edges if j == n) -
    sum(flows[(i, j)] for (i, j) in edges if i == n) == demand(n)
);

Преимущество такой формулировки в том, что она очень похожа на математическую формулировку задачи о сетевом потоке.

Недостаток же ее не так очевиден. На внутреннем уровне макрос JuMP @constraint расширяется примерно так:

model = Model()
@variable(model, flows[e in edges] >= 0)
for n in nodes
    flow_in = AffExpr(0.0)
    for (i, j) in edges
        if j == n
            add_to_expression!(flow_in, flows[(i, j)])
        end
    end
    flow_out = AffExpr(0.0)
    for (i, j) in edges
        if i == n
            add_to_expression!(flow_out, flows[(i, j)])
        end
    end
    @constraint(model, flow_in - flow_out == demand(n))
end

Эта формулировка включает два цикла for, с выполнением цикла по каждому ребру (дважды) для каждого узла. В нотации «O» большое время выполнения будет составлять . При большом количестве узлов и ребер выполнение этого цикла может занимать много времени.

Создадим функцию для тестирования производительности нашей формулировки:

function build_naive_model(nodes, edges, demand)
    model = Model()
    @variable(model, flows[e in edges] >= 0)
    @constraint(
        model,
        [n in nodes],
        sum(flows[(i, j)] for (i, j) in edges if j == n) -
        sum(flows[(i, j)] for (i, j) in edges if i == n) == demand(n)
    )
    return model
end

nodes, edges, demand = build_random_graph(1_000, 2_000)
@elapsed build_naive_model(nodes, edges, demand)

0.164541712

Хорошим способом тестирования производительности является измерение времени выполнения в широком диапазоне размеров входных данных. Из анализа «О» большого можно предположить, что удвоение количества узлов и ребер приведет к четырехкратному увеличению времени выполнения.

run_times = Float64[]
factors = 1:10
for factor in factors
    graph = build_random_graph(1_000 * factor, 5_000 * factor)
    push!(run_times, @elapsed build_naive_model(graph...))
end
Plots.plot(; xlabel = "Factor", ylabel = "Runtime [s]")
Plots.scatter!(factors, run_times; label = "Actual")
a, b = hcat(ones(10), factors .^ 2) \ run_times
Plots.plot!(factors, a .+ b * factors .^ 2; label = "Quadratic fit")

Как и ожидалось, время выполнения демонстрирует квадратичное масштабирование: если удвоить количество узлов и ребер, время выполнения увеличится в четыре раза.

Кэширование

Чтобы улучшить формулировку, можно кэшировать список входящих и исходящих узлов для каждого узла n:

out_nodes = Dict(n => Int[] for n in nodes)
in_nodes = Dict(n => Int[] for n in nodes)
for (i, j) in edges
    push!(out_nodes[i], j)
    push!(in_nodes[j], i)
end

Внесем соответствующее изменение в модель:

model = Model()
@variable(model, flows[e in edges] >= 0)
@constraint(
    model,
    [n in nodes],
    sum(flows[(i, n)] for i in in_nodes[n]) -
    sum(flows[(n, j)] for j in out_nodes[n]) == demand(n)
);

Преимущество этой формулировки в том, что теперь перебирается out_nodes[n], а не edges для каждого узла n, поэтому время выполнения составляет .

Создадим новую функцию для тестирования производительности нашей формулировки:

function build_cached_model(nodes, edges, demand)
    out_nodes = Dict(n => Int[] for n in nodes)
    in_nodes = Dict(n => Int[] for n in nodes)
    for (i, j) in edges
        push!(out_nodes[i], j)
        push!(in_nodes[j], i)
    end
    model = Model()
    @variable(model, flows[e in edges] >= 0)
    @constraint(
        model,
        [n in nodes],
        sum(flows[(i, n)] for i in in_nodes[n]) -
        sum(flows[(n, j)] for j in out_nodes[n]) == demand(n)
    )
    return model
end

nodes, edges, demand = build_random_graph(1_000, 2_000)
@elapsed build_cached_model(nodes, edges, demand)

0.197684169

Анализ

Теперь можно проанализировать разницу во времени выполнения двух формулировок:

run_times_naive = Float64[]
run_times_cached = Float64[]
factors = 1:10
for factor in factors
    graph = build_random_graph(1_000 * factor, 5_000 * factor)
    push!(run_times_naive, @elapsed build_naive_model(graph...))
    push!(run_times_cached, @elapsed build_cached_model(graph...))
end
Plots.plot(; xlabel = "Factor", ylabel = "Runtime [s]")
Plots.scatter!(factors, run_times_naive; label = "Actual")
a, b = hcat(ones(10), factors .^ 2) \ run_times_naive
Plots.plot!(factors, a .+ b * factors .^ 2; label = "Quadratic fit")
Plots.scatter!(factors, run_times_cached; label = "Cached")
a, b = hcat(ones(10), factors) \ run_times_cached
Plots.plot!(factors, a .+ b * factors; label = "Linear fit")

Хотя кэшированной модели приходится создавать in_nodes и out_nodes, она асимптотически быстрее примитивной модели и масштабируется с factor линейно, а не квадратично.

Вывод

При написании кода с условиями типа sum-if, например @constraint(model, [a in set], sum(x[b] for b in list if condition(a, b)), можно улучшить производительность путем кэширования элементов, для которых condition(a, b) имеет значение true.

Наконец, следует понимать, что такое поведение характерно не только для JuMP, но и в целом для любых компьютерных программ. (Программы на Python, использующие Pyomo или gurobipy, также выиграют от такого кэширования.)

Понимание нотации «О» большого и алгоритмической сложности — полезный при отладке навык независимо от типа разрабатываемой программы.