acset vs dicts for jump model generation

discourse_jump.jl

# from question: https://discourse.julialang.org/t/speeding-up-jump-model-creation-with-sets-that-depend-on-other-indexes/107333

using Catlab, JuMP
using BenchmarkTools

@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 = [(i,j) for i in 1:num_vertices for j in i+1:num_vertices]
add_parts!(
    NetData, :E, length(E),
    src = first.(E), tgt = last.(E),
    te = [T_options[(a+b)%3+1] for (a,b) in E]
)

# ------------------------------------------------------------
# build model from acset

function create_jump_model_acset(acs)
    model = Model()

    V = parts(acs,:V)
    E = parts(acs,:E)

    @variable(
        model, 
        flow[e=E, acs[e,:te]] ≥ 0
    )

    @objective(
        model, 
        Min,
        sum(flow[e,t] for e=E, t=acs[e,:te])
    )

    @expression(
        model,
        incoming[v=V,Tᵥ=acs[v,:tv]],
        sum(
            length(Tₑ ∩ Tᵥ) * flow[e,Tₑ]
            for e in incident(acs,v,:tgt), Tₑ in acs[e,:te]
        )
    )

    @expression(
        model,
        outgoing[v=V,Tᵥ=acs[v,:tv]],
        sum(
            length(Tₑ ∩ Tᵥ) * flow[e,Tₑ]
            for e in incident(acs,v,:src), Tₑ in acs[e,:te]
        )
    )

    @constraint(
        model,
        balance[v=V,Tᵥ=acs[v,:tv]],
        incoming[v,Tᵥ] == outgoing[v,Tᵥ]
    )

    return model
end


# ------------------------------------------------------------
# existing impl

TV = Dict(
    v => T_options[v%3+1]
    for v in V
)
TE = Dict(
    (a, b) => T_options[(a+b)%3+1]
    for (a, b) in E
)

function create_jump_model_original(V,E,TV,TE)
    model = Model()

    @variable(model, flow[e ∈ E, TE[e]] ≥ 0)

    @objective(model, Min,
        sum(flow[e, T] for e in E, T ∈ TE[e])
    )

    @expression(model,
        incoming[v ∈ V, Tᵥ ∈ TV[v]],
        sum(
            length(Tₑ ∩ Tᵥ) * flow[(src, v), Tₑ]
            for src in [src for (src, dst) in E if dst == v], Tₑ ∈ TE[(src, v)]
        )
    )
    @expression(model,
        outgoing[v ∈ V, Tᵥ ∈ TV[v]],
        sum(
            length(Tₑ ∩ Tᵥ) * flow[(v, dst), Tₑ]
            for dst in [dst for (src, dst) in E if src == v], Tₑ ∈ TE[(v, dst)]
        )
    )
    @constraint(model,
        balance[v ∈ V, Tᵥ ∈ TV[v]],
        incoming[v, Tᵥ] == outgoing[v, Tᵥ]
    )

    return model
end


# ------------------------------------------------------------
# benchmark

@benchmark create_jump_model_acset(NetData) setup=(create_jump_model_acset(NetData))
@benchmark create_jump_model_original(V,E,TV,TE) setup=(create_jump_model_original(V,E,TV,TE))

@time create_jump_model_acset(NetData)
@time create_jump_model_original(V,E,TV,TE)

mod_acset = create_jump_model_acset(NetData)
mod_orig = create_jump_model_original(V,E,TV,TE) 

@assert length(mod_acset[:flow]) == length(mod_orig[:flow])
@assert length(mod_acset[:incoming]) == length(mod_orig[:incoming])
@assert length(mod_acset[:outgoing]) == length(mod_orig[:outgoing])
@assert length(mod_acset[:balance]) == length(mod_orig[:balance])