Example of nodal problem formulations with JuMP

nodal-jump.jl

using JuMP

function localVar(m, name = "", expose = false) 
    x = JuMP.Variable(m, -Inf, Inf, :Cont, name, NaN)
    if expose
        JuMP.registervar(m, name, x)
    end
    x
end


############################################################
##
## The following is a way to collect nodal variables and equations. `refbranch`
## registers that a branch is connected to a node and gives the branch flow into
## the node. Then when the model has been specified, `addnodalequations` adds
## constraints based on these: at each registered node, the sum of branch flows
## into the node is set to zero.. 
##
############################################################
const nodedict = Dict()

function refbranch(m, n, x)
    if !haskey(m.ext, :NodeDict)
        m.ext[:NodeDict] = Dict{Any,Any}(n => x)  
    else    
        mdict = m.ext[:NodeDict]
        if haskey(mdict, n)
            mdict[n] += x
        else 
            mdict[n] = x
        end
    end
end

function addnodalequations(m)
    for (k,v) in m.ext[:NodeDict]
        if isa(k, Variable)
            @addConstraint(m, v == 0.0)
        end
    end
end

############################################################
##
## Component definitions
##
############################################################

function branch(m, n1, n2, i)
    refbranch(m, n1, i)
    refbranch(m, n2, -i)
end
    

function resistor(m::Model, n1, n2, R)
    # @show i = localVar(m, "resistor_I")
    i = localVar(m, string("resistor_I_", n1, "-", n2))
    branch(m, n1, n2, i)
    @addConstraint(m, n1 - n2 == i * R)
end

function currentsource(m::Model, n1, n2, i)
    branch(m, n1, n2, -i)
end

function voltagesource(m::Model, n1, n2, v)
    i = localVar(m, string("source_I_", n1, "-", n2))
    branch(m, n1, n2, i)
    @addConstraint(m, n1 - n2 == v)
end

function powermeter(m::Model, P, mdl, n1, n2, args...)
    ## This meter wraps another model `mdl`
    n3 = localVar(m, string("powermeter_internal_V_", n1, "-", n2))
    i = localVar(m, string("powermeter_I_", n1, "-", n2))
    branch(m, n1, n3, i)
    mdl(m, n3, n2, args...)
    @addConstraint(m, n1 == n3)
    @addConstraint(m, i * (n1 - n2) == P)
end

function currentmeter(m::Model, I, mdl, n1, n2, args...)
    n3 = localVar(m, string("powermeter_internal_V_", n1, "-", n2))
    branch(m, n1, n3, I)
    @addConstraint(m, n1 == n3)
    mdl(m, n3, n2, args...)
end

############################################################
##
## Example models to solve
##
############################################################

function submodel(m, v1)
    @defVar(m, v2)
    @defVar(m, v3)
    resistor(m, v1, v2, 3)
    resistor(m, v2,  0, 6)
    resistor(m, v2, v3, 3)
    resistor(m, v3,  0, 3)
end

function model1()   # use a current source
    m = Model()
    @defVar(m, v1)  # define voltage nodes
    resistor(m, v1,  0, 6)   # connection to ground
    submodel(m, v1)
    currentsource(m, v1, 0.0, 10.0)    
    addnodalequations(m)   # fix up nodal equations
    m
end

function model2()     # use a voltage source using a constant
    m = Model()
    v1 = 10.0       # a known voltage, a voltage source
    submodel(m, v1)
    addnodalequations(m)   # fix up nodal equations
    m
end

function model3()     # use a voltage source
    m = Model()
    @defVar(m, v1)  # define voltage nodes
    submodel(m, v1)
    voltagesource(m, v1, 0.0, 10.0)    
    addnodalequations(m)   # fix up nodal equations
    m
end

m1 = model1()
m2 = model2()
m3 = model3()
solve(m1)
@show m1.colNames
@show m1.colVal
solve(m2)
@show m2.colNames
@show m2.colVal
solve(m3)
@show m3.colNames
@show m3.colVal


############################################################
##
## Optimization problem
##
############################################################

function opt1()
    m = Model()
    v1 = 10.0       # voltage source
    @defVar(m, Rout)  # Rout is the variable that can be twiddled 
    @defVar(m, P)     # output power to maximize
    @defVar(m, v2)
    @defVar(m, v3)
    resistor(m, v1, v2, 3)
    resistor(m, v2,  0, 6)
    resistor(m, v2, v3, 3)
    powermeter(m, P, resistor, v3, 0, Rout)
    addnodalequations(m)
    setObjective(m, :Max, P)   # maximize power output
    (m, Rout, P)
end

(o1, Rout, P) = opt1()
s = solve(o1)
println("Max power = $(getValue(P)) with Rout = $(getValue(Rout))")