JuMP-ing with AlgebraicJulia II: A practical optimization model

optimization
visualization
Looking at using ACSets to structure and solve a fundamental network optimization model.
Author

Sean L. Wu

Published

March 26, 2025

This is the second blog post in a series about tools from AlgebraicJulia to ease mathematical programming using the JuMP framework. If you haven’t seen JuMP-ing with AlgebraicJulia I: The IJKLM model yet, I suggest you take a look at it first!

The multi-commodity network flow problem is a foundational model in many fields which rely on mathematical programming such as logistics and telecommunications routing. Models of this type are called network flow problems and optimal solutions tell us the best way to send things on edges between vertices such that demands are fulfilled, subject to various flow and supply constraints. There is a JuMP tutorial page which describes a dataset used to parameterize the multi-commodity network flow problem, and solves it with JuMP. Here we show how the same model can be expressed within a single ACSet, and how to use some of the graph visualization tools of Catlab (using GraphViz) to inspect your data and optimization model.

In this variant of the model, we consider a set of vertices i \in \mathcal{V}, each of which has a (potentially zero) supply capacity u^s_{i,p}, supply cost c^s_{i,p}, and demand d_{i,p} for each commodity p \in P. The vertices are connected by a set of edges (i, j) \in \mathcal{E}, which have a shipment cost c^x_{i,j,p} and a total flow capacity of u^x_{i,j}.

The decision variables are the purchased quantity of each commodity at each vertex s_{i,p}, as well as the optimal shipment of each commodity along each edge x_{i,j,p} that minimizes the total cost.

\begin{aligned} \min \;\; & \sum_{(i,j)\in\mathcal{E}, p \in P} c^x_{i,j,p} x_{i,j,p} + \sum_{i\in\mathcal{V}, p \in P} c^s_{i,p} s_{i,p} \\ s.t. \;\; & s_{i,p} + \sum_{(j, i) \in \mathcal{E}} x_{j,i,p} - \sum_{(i,j) \in \mathcal{E}} x_{i,j,p} = d_{i,p} & \forall i \in \mathcal{V}, p \in P \\ & x_{i,j,p} \ge 0 & \forall (i, j) \in \mathcal{E}, p \in P \\ & \sum_{p \in P} x_{i,j,p} \le u^x_{i,j} & \forall (i, j) \in \mathcal{E} \\ & 0 \le s_{i,p} \le u^s_{i,p} & \forall i \in \mathcal{V}, p \in P \end{aligned}

Data Import

First let’s load the data that is used to parameterize the model. Using the SQLite.jl library, we get each table from the relational database as a DataFrame, as in the JuMP tutorial which can be consulted for details.

Code
import SQLite, SQLite.DBInterface
using DataFrames

filename = joinpath(@__DIR__, "commodity_nz.db");
db = SQLite.DB(filename)

function get_table(db, table)
    query = DBInterface.execute(db, "SELECT * FROM $table")
    return DataFrames.DataFrame(query)
end

df_shipping = get_table(db, "shipping")
df_products = get_table(db, "products")
df_supply = get_table(db, "supply")
df_demand = get_table(db, "demand")

df_cost = DataFrames.leftjoin(df_shipping, df_products; on = [:product])
df_cost.flow_cost = df_cost.cost_per_km .* df_cost.distance_km

Unlike the tutorial, we make df_edges containing the (src, tgt) pairs, and df_productvertex which is the product of vertices and products (commodities). We then add the supply capacities, costs, and demands to each element of the product.

Code
places = unique([df_cost.origin; df_cost.destination])
products = unique(df_cost.product)

df_edges=unique(df_cost[:, [:origin, :destination]])
transform!(df_edges, :origin => ByRow(o -> begin
    findfirst(o .== places)
end) => :src)
transform!(df_edges, :destination => ByRow(o -> begin
    findfirst(o .== places)
end) => :tgt)

df_productvertex = crossjoin(DataFrame(place=places), DataFrame(product=products))
leftjoin!(df_productvertex, df_supply, on=[:place=>:origin, :product])
leftjoin!(df_productvertex, df_demand, on=[:place=>:destination, :product])
df_productvertex[ismissing.(df_productvertex.capacity), :capacity] .= 0.0
df_productvertex[ismissing.(df_productvertex.cost), :cost] .= 0.0
df_productvertex[ismissing.(df_productvertex.demand), :demand] .= 0.0

ACSet Structure

Now we’ll generate a schema which is appropriate to represent the multi-commodity network flow problem. We use schema inheritance to inherit from SchGraph, the schema for simple graphs. Our new objects include Commodity which is the set of commodities considered, CommodityVertex which is the product of Commodity and V, and Shipping which is the product of Commodity and E. With these additional objects, we can conveniently store all the data and variables of the optimization model. We add some additional AttrTypes to the schema as well, DecisionVar which will store the JuMP decision variables, and NumVar which stores numeric data (used to parameterize the model).

Code
using Catlab

@present SchMultiCommodity <: SchGraph begin    
    (CommodityVertex,Commodity,Shipping)::Ob
    cv_c::Hom(CommodityVertex,Commodity)
    cv_v::Hom(CommodityVertex,V)
    s_p::Hom(Shipping,Commodity)
    s_e::Hom(Shipping,E)

    Label::AttrType
    vlabel::Attr(V,Label)
    plabel::Attr(Commodity,Label)

    NumVar::AttrType
    supplycap::Attr(CommodityVertex,NumVar) # supply capacity
    demand::Attr(CommodityVertex,NumVar) # demand
    purcost::Attr(CommodityVertex,NumVar) # purchase cost
    shipcost::Attr(Shipping,NumVar) # shipment cost
    flowcap::Attr(E,NumVar) # flow capacity

    DecisionVar::AttrType
    s::Attr(CommodityVertex,DecisionVar) # optimal supply
    x::Attr(Shipping,DecisionVar) # optimal shipment
end

to_graphviz(SchMultiCommodity, graph_attrs=Dict(:size=>"6",:ratio=>"fill"))

We next generate the ACSet types from the schema. Note that type inheritance for ACSets is handled seperately from the schema inheritance, so we inherit from HasGraph to take advantage of the Graphs module of Catlab.

Code
@abstract_acset_type AbstractMultiCommodity <: HasGraph
@acset_type MultiCommodity(SchMultiCommodity, index=[:src,:tgt,:cv_c,:cv_v,:s_p,:s_e]) <: AbstractMultiCommodity

Now we are ready to generate an instance of an ACSet of this type with the data we set up previously. Like the original tutorial, we assume that the capacity of each edge is 30.

Code
using JuMP

capacity = 30

multinet_acs = @acset MultiCommodity{String,Float64,JuMP.VariableRef} begin
    V=length(places)
    vlabel=places
    Commodity=length(products)
    plabel=products
    E=nrow(df_edges)
    src=df_edges.src
    tgt=df_edges.tgt
    flowcap=capacity
end

for r in eachrow(df_productvertex)
    v=only(incident(multinet_acs, r.place, :vlabel))
    p=only(incident(multinet_acs, r.product, :plabel))
    add_part!(
        multinet_acs, :CommodityVertex, 
        cv_v=v, cv_c=p,
        supplycap=r.capacity, demand=r.demand, purcost=r.cost
    )
end

for r in eachrow(df_cost)
    src=only(incident(multinet_acs, r.origin, :vlabel))
    tgt=only(incident(multinet_acs, r.destination, :vlabel))
    product=only(incident(multinet_acs, r.product, :plabel))
    add_part!(
        multinet_acs, :Shipping,
        s_p=product, s_e=only(edges(multinet_acs, src, tgt)),
        shipcost=r.flow_cost
    )
end

Next we’d like to visualize the data and the network we created. Catlab integrates with GraphViz for visualizations of graph-like objects, and because we want to customize how the vertices and edges are displayed, we write some code that generates the labels for vertices below. We rely on the ability of GraphViz to generate HTML-like labels for vertices. For more information on the syntax, please see the GraphViz documentation.

Below we write two versions of the method that will generate the label for a vertex. make_node_label_optim should be run on an ACSet of type AbstractMultiCommodity after it has been optimized, to print out a table indicating how much of each commodity was supplied by that vertex. The make_node_label_parameters prints out the data that parameterizes the model for each vertex/commodity combination: the maximum supply capacity, demand, and supply cost.

Code
import Catlab.Graphics.Graphviz.Html

"""
Make a node label for a model that has been optimized.
"""
function make_node_label_optim(v, acs, label_dict)
    commodity_v = incident(acs, v, :cv_v)
    label = Vector{String}()
    push!(
        label, """
            <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
            <TR><TD COLSPAN="2">$(acs[v, :vlabel])</TD></TR>
            <TR><TD>Product</TD><TD>Purchased</TD></TR>
        """
    )
    for cv in commodity_v
        if JuMP.value(acs[cv, :s]) > 0
            emoji = label_dict[acs[cv, (:cv_c, :plabel)]]
            push!(
                label, """
                    <TR><TD>$(emoji)</TD><TD>$(JuMP.value(acs[cv, :s]))</TD></TR>
                """
            )
        end
    end
    push!(
        label, """
            </TABLE>
        """
    )
end

"""
Make a node label for the model parameters.
"""
function make_node_label_parameters(v, acs, label_dict)
    commodity_v = incident(acs, v, :cv_v)
    label = Vector{String}()
    push!(
        label, """
            <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
            <TR><TD COLSPAN="4">$(acs[v, :vlabel])</TD></TR>
            <TR><TD>Product</TD><TD>Supply Capacity</TD><TD>Demand</TD><TD>Purchase Cost</TD></TR>
        """
    )
    for cv in commodity_v
        emoji = label_dict[acs[cv, (:cv_c, :plabel)]]
        push!(
            label, """
                <TR><TD>$(emoji)</TD><TD>$(acs[cv, :supplycap])</TD><TD>$(acs[cv, :demand])</TD><TD>$(acs[cv, :purcost])</TD></TR>
            """
        )
    end
    push!(
        label, """
            </TABLE>
        """
    )
end

make_node_label_dict = Dict(:optim=>make_node_label_optim, :parameters=>make_node_label_parameters)

Likewise we have two similar methods for edges. The first, to be run post-optimization, prints the amount of each commodity shipped on each edge. The second prints the shipping cost of each commodity, producing a visualization of the data which parameterizes the mathematical model. Finally we define make_property_graph which generates a property graph using these methods. Because our ACSet type inherits from HasGraph, we can use the to_graphviz_property_graph method defined on any graph-like type to do the hard work, and then simply modify edge and vertex labels.

Code
"""
Make an edge label for an optmized model.
"""
function make_edge_label_optim(e, acs, label_dict)
    ship = incident(acs, e, :s_e)
    label = Vector{String}()
    edge_label = acs[e, (:s_e, :src, :vlabel)] * " → " * acs[e, (:s_e, :tgt, :vlabel)]
    push!(
        label, """
            <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
            <TR><TD COLSPAN="2">$(edge_label)</TD></TR>
            <TR><TD>Product</TD><TD>Shipped</TD></TR>
        """
    )
    for s in ship
        if JuMP.value(acs[s, :x]) > 0.0
            emoji = label_dict[acs[s, (:s_p, :plabel)]]
            push!(
                label, """
                    <TR><TD>$(emoji)</TD><TD>$(JuMP.value(acs[s, :x]))</TD></TR>
                """
            )
        end
    end
    push!(
        label, """
            </TABLE>
        """
    )
end

"""
Make an edge label for the parameters of the model.
"""
function make_edge_label_parameters(e, acs, label_dict)
    ship = incident(acs, e, :s_e)
    label = Vector{String}()
    edge_label = acs[e, (:s_e, :src, :vlabel)] * " → " * acs[e, (:s_e, :tgt, :vlabel)]
    push!(
        label, """
            <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0">
            <TR><TD COLSPAN="2">$(edge_label)</TD></TR>
            <TR><TD>Product</TD><TD>Shipping Cost</TD></TR>
        """
    )
    for s in ship
        emoji = label_dict[acs[s, (:s_p, :plabel)]]
        push!(
            label, """
                <TR><TD>$(emoji)</TD><TD>$(acs[s, :shipcost])</TD></TR>
            """
        )
    end
    push!(
        label, """
            </TABLE>
        """
    )
end

make_edge_label_dict = Dict(:optim=>make_edge_label_optim,:parameters=>make_edge_label_parameters)

function make_property_graph(acs::T; labels, kwargs...) where {T<:AbstractMultiCommodity}
    pg = to_graphviz_property_graph(acs; kwargs...)   
    label_dict = Dict("milk"=>"🥛", "kiwifruit"=>"🥝")
    for v in parts(acs, :V)
        label = make_node_label_dict[labels](v, acs, label_dict)
        set_vprops!(pg, v, label=Html(join(label)), shape="plain")
    end
    for e in parts(acs, :E)
        label = make_edge_label_dict[labels](e, acs, label_dict)
        set_eprops!(pg, e, label=Html(join(label)), shape="plain")
    end
    return pg
end

Finally, we may visualize the data which parameterizes the optimization model.

Code
to_graphviz(make_property_graph(multinet_acs, labels=:parameters, graph_attrs=Dict(:rankdir=>"TD")))

Optimization model in JuMP

Now we are ready to write the optimization model itself. We take a different approach and write methods that will add the appropriate decision variables, constraints, and objective to the ACSet. Each of them takes as argument the ACSet object and the JuMP model.

Code
"""
Adds `s` and `x`, decision variables for how much supply to procure of a product at each vertex,
and how much of each product to ship along edges.
"""
function add_variables!(acs, model)
    acs[:, :s] = @variable(model, 0  s[i=parts(acs, :CommodityVertex)]  acs[i, :supplycap])
    acs[:, :x] = @variable(model, 0  x[i=parts(acs, :Shipping)])
end

"""
Add constraints to respect total flow capacity of all products along each edge.
"""
function add_flow_constraint!(acs, model)
    for i in parts(acs, :E)
        @constraint(
            model,
            sum(acs[incident(acs, i, :s_e), :x])  acs[i, :flowcap]
        )
    end
end

"""
Add constraint that, for each product, vertex pair, the supply plus inbound
shipping, minus outbound shipping, must equal the demand.
"""
function add_conservation_constraint!(acs, model)
    for i in parts(acs, :CommodityVertex)
        v = acs[i, :cv_v]
        p = acs[i, :cv_c]

        # inbound shipping
        inbound = intersect(incident(acs, v, (:s_e, :tgt)), incident(acs, p, :s_p))

        # outbound shipping
        outbound = intersect(incident(acs, v, (:s_e, :src)), incident(acs, p, :s_p))

        @constraint(
            model,
            acs[i, :s] + 
                sum(acs[inbound, :x], init=zero(JuMP.VariableRef)) - 
                sum(acs[outbound, :x], init=zero(JuMP.VariableRef)) == acs[i, :demand]
        )
    end
end

"""
Minimize costs of shipping and costs of purchase
"""
function add_objective!(acs, model)
    @objective(
        model, Min,
        sum(acs[:, :shipcost] .* acs[:, :x]) + 
            sum(acs[:, :purcost] .* acs[:, :s])
    )
end

Finally we are ready to generate and solve the model. We import HiGHS which contains the solvers we will use. We can then add the decision variables, constraints, and objective to the model, and solve it. We then copy the optimized values of the decision variables into the acset, and visualize the solution.

Code
# make the JuMP model
using HiGHS

model = JuMP.Model(HiGHS.Optimizer)
set_silent(model)

add_variables!(multinet_acs, model)
add_flow_constraint!(multinet_acs, model)
add_conservation_constraint!(multinet_acs, model)
add_objective!(multinet_acs, model)

optimize!(model)

to_graphviz(make_property_graph(multinet_acs, labels=:optim, graph_attrs=Dict(:rankdir=>"TD")))

We may format the solution in the same way as the original tutorial to confirm that they are the same.

Code
"""
Format optimized decision variables for comparison. 
"""
function format_output(acs)
    df_solution_flow = DataFrame(tables(acs).Shipping)
    df_solution_flow.x = JuMP.value.(df_solution_flow.x)
    df_solution_flow = df_solution_flow[df_solution_flow.x .> 0.0, :]
    df_solution_flow.origin = acs[df_solution_flow.s_e, (:src, :vlabel)]
    df_solution_flow.destination = acs[df_solution_flow.s_e, (:tgt, :vlabel)]
    df_solution_flow.product = acs[df_solution_flow.s_p, :plabel]
    select!(df_solution_flow, Not([:s_p,:s_e,:shipcost]))
    rename!(df_solution_flow, Dict(:x => :x_flow))
    return df_solution_flow
end

format_output(multinet_acs)
9×4 DataFrame
Row x_flow origin destination product
Float64 String String String
1 10.0 waikato auckland milk
2 2.0 waikato wellington milk
3 2.0 tauranga auckland milk
4 2.0 tauranga waikato milk
5 4.0 christchurch auckland milk
6 4.0 auckland christchurch kiwifruit
7 20.0 waikato auckland kiwifruit
8 2.0 waikato wellington kiwifruit
9 22.0 tauranga waikato kiwifruit

I hope that this post has shown you how ACSets can helpfully structure optimization problems, joining data with the mathematical model in a single data structure, especially those models which can be represented as elaborations of graphs (of which there are many)!