Speeding up JuMP model creation with sets that depend on other indexes

@abelsiqueira I hope you don’t mind me offering yet another “radically different data structure” :slightly_smiling_face:

You can consider using “ACSets” exported by the Catlab.jl package. They are a generalization of C-Sets, a structure from category theory, but basically end up being a lot like an in-memory relational database with (optional) equations asserting equivalence (aka commutativity) between paths.

While there are multiple ways to build a data base schema for your data, one is as a “decorated graph”. In the schema diagram below, E and V will be sent to Sets (in the mathematical sense), and TimeType will be sent to a Julia type. The arrows will be functions. The labeled ovals in the diagram are combinatorial data, and the dot (TimeType) is atomic data.

With that, you can build your model in the following manner. I personally think it is very readable this way, but your opinion may differ. Code like NetData[e,:src] walks forwards along arrows (foreign keys/functions), and code like incident(NetData, 1, :src) finds preimages, so for the example, it returns the set of edges whose source is vertex 1.

using Catlab, JuMP

@present NetSch(FreeSchema) begin
    (V,E)::Ob
    src::Hom(E,V)
    tgt::Hom(E,V)
    TimeType::AttrType
    te::Attr(E,TimeType)
    tv::Attr(V,TimeType)
end

to_graphviz(NetSch, graph_attrs=Dict(:dpi=>"72",:size=>"3.5",:ratio=>"expand"))

@acset_type NetType(NetSch, index=[:src,:tgt])

num_vertices = 100

T1 = [1:3, 4:6, 7:9, 10:12]
T2 = [1:4, 5:8, 9:12]
T3 = [1:6, 7:12]
T_options = [T1, T2, T3]

# make the acset
NetData = NetType{typeof(T1)}()

# add the verticies
V = 1:num_vertices
add_parts!(
    NetData, :V, length(V),
    tv = [T_options[v%3+1] for v in V]
)

# add the edges
E_list = [(i,j) for i in 1:num_vertices for j in i+1:num_vertices]
add_parts!(
    NetData, :E, length(E_list),
    src = first.(E_list), tgt = last.(E_list),
    te = [T_options[(a+b)%3+1] for (a,b) in E_list]
)

E = parts(NetData,:E)

# build the JuMP model 
model = JuMP.Model()
@variable(
    model, 
    flow[e=E, NetData[e,:te]] ≥ 0
)
@objective(
    model, 
    Min,
    sum(flow[e,t] for e=E, t=NetData[e,:te])
)
@expression(
    model,
    incoming[v=V,Tᵥ=NetData[v,:tv]],
    sum(
        length(Tₑ ∩ Tᵥ) * flow[e,Tₑ]
        for e in incident(NetData,v,:tgt), Tₑ in NetData[e,:te]
    )
)
@expression(
    model,
    outgoing[v=V,Tᵥ=NetData[v,:tv]],
    sum(
        length(Tₑ ∩ Tᵥ) * flow[e,Tₑ]
        for e in incident(NetData,v,:src), Tₑ in NetData[e,:te]
    )
)
@constraint(
    model,
    balance[v=V,Tᵥ=NetData[v,:tv]],
    incoming[v,Tᵥ] == outgoing[v,Tᵥ]
)

I’ll also state that this was just the most direct translation of your original code into an ACSet formulation. It may not be the most efficient, especially if all the times could reasonably be represented as a finite set rather than an arbitrary Julia type. Also, for larger models, I second @odow’s suggestion to do subsetting and filtering as much as possible outside of the JuMP code, for which various tools exist in Catlab (e.g. conjunctive queries).

1 Like