Create functionality to approximate NLP functions present in a optimization model.

[andrewrosemberg]

Create functionality to approximate NLP functions present in a optimization model by ML proxies (e.g. Neural Network).

For example, if we have an optimization problem dependent on non-linear (NL) expressions:

function build_test_nlp_model()
    model = Model()

    @variable(model, x[i = 1:2]);

    @variable(model, y[i = 1:2] >= 0.0);

    ex1 = sin(x[1])
    ex2 = cos(x[2])

    cons = @NLconstraint(model, ex1 == ex2)

    @objective(model, Min, sum(x) + sum(y))

    return model, cons
end

We might be able to extract and evaluate NLP expressions:

function constraints_nlp_evaluator(model, x)
    d = NLPEvaluator(model)
    MOI.initialize(d, [:ExprGraph])

    f = zeros(length(model.nlp_model.constraints))

    MOI.eval_constraint(d, f, x)

    return f
end

model, cons = build_test_nlp_model()
input = [[i] for i in rand(1000)]
output = constraints_nlp_evaluator.(model, input)

After fitting a proxy to these expressions/functions, we can add them back to the optimization problem:

flux_model = Chain(Dense(3, 3, relu), Dense(3, 1))

ex = flux_model(x)[1]

cons = @constraint(model, ex == 1)

[andrewrosemberg]

Example PowerModels:

If we want to approximate the following functions from PowerModels (i.e. the Power Flow equations):

import PowerModels: constraint_ohms_yt_from, constraint_ohms_yt_to, constraint_ohms_y_from, constraint_ohms_y_to

function _function_ohms_yt_from(branch::Dict)
    g, b = calc_branch_y(branch)
    tr, ti = calc_branch_t(branch)
    g_fr = branch["g_fr"]
    b_fr = branch["b_fr"]
    tm = branch["tap"]

    return  (vm_fr, vm_to, va_fr, va_to) ->  ((g+g_fr)/tm^2*vm_fr^2 + (-g*tr+b*ti)/tm^2*(vm_fr*vm_to*cos(va_fr-va_to)) + (-b*tr-g*ti)/tm^2*(vm_fr*vm_to*sin(va_fr-va_to)), # from
            -(b+b_fr)/tm^2*vm_fr^2 - (-b*tr-g*ti)/tm^2*(vm_fr*vm_to*cos(va_fr-va_to)) + (-g*tr+b*ti)/tm^2*(vm_fr*vm_to*sin(va_fr-va_to))) # to

end

function function_ohms_yt_from(branch::Dict)
    _function_ohms_yt_from(branch)
end

function PowerModels.constraint_ohms_yt_from(pm::AbstractACPModel, i::Int; nw::Int=nw_id_default)
    branch = ref(pm, nw, :branch, i)
    f_bus = branch["f_bus"]
    t_bus = branch["t_bus"]
    f_idx = (i, f_bus, t_bus)

    p_fr  = var(pm, nw,  :p, f_idx)
    q_fr  = var(pm, nw,  :q, f_idx)
    vm_fr = var(pm, nw, :vm, f_bus)
    vm_to = var(pm, nw, :vm, t_bus)
    va_fr = var(pm, nw, :va, f_bus)
    va_to = var(pm, nw, :va, t_bus)

    f_owms = function_ohms_yt_from(branch)
    f_owms_p, f_owms_q = f_owms(vm_fr, vm_to, va_fr, va_to)

    #constraint_ohms_yt_from(pm, nw, f_bus, t_bus, f_idx, t_idx, g, b, g_fr, b_fr, tr, ti, tm)
    JuMP.@constraint(pm.model, p_fr == f_owms_p) # @NL
    JuMP.@constraint(pm.model, q_fr == f_owms_q) # @NL
end

function _function_ohms_yt_to(branch::Dict)
    g, b = calc_branch_y(branch)
    tr, ti = calc_branch_t(branch)
    g_to = branch["g_to"]
    b_to = branch["b_to"]
    tm = branch["tap"]

    return  (vm_fr, vm_to, va_fr, va_to) ->  ((g+g_to)*vm_to^2 + (-g*tr-b*ti)/tm^2*(vm_to*vm_fr*cos(va_to-va_fr)) + (-b*tr+g*ti)/tm^2*(vm_to*vm_fr*sin(va_to-va_fr)), # from
            -(b+b_to)*vm_to^2 - (-b*tr+g*ti)/tm^2*(vm_to*vm_fr*cos(va_to-va_fr)) + (-g*tr-b*ti)/tm^2*(vm_to*vm_fr*sin(va_to-va_fr))) # to

end

function function_ohms_yt_to(branch::Dict)
    _function_ohms_yt_to(branch)
end

function PowerModels.constraint_ohms_yt_to(pm::AbstractACPModel, i::Int; nw::Int=nw_id_default)
    branch = ref(pm, nw, :branch, i)
    t_bus = branch["t_bus"]
    f_bus = branch["f_bus"]
    t_idx = (i, t_bus, f_bus)

    p_to  = var(pm, nw,  :p, t_idx)
    q_to  = var(pm, nw,  :q, t_idx)
    vm_fr = var(pm, nw, :vm, f_bus)
    vm_to = var(pm, nw, :vm, t_bus)
    va_fr = var(pm, nw, :va, f_bus)
    va_to = var(pm, nw, :va, t_bus)

    f_owms = function_ohms_yt_to(branch)
    f_owms_p, f_owms_q = f_owms(vm_fr, vm_to, va_fr, va_to)

    # constraint_ohms_yt_to(pm, nw, f_bus, t_bus, f_idx, t_idx, g, b, g_to, b_to, tr, ti, tm)
    JuMP.@constraint(pm.model, p_to == f_owms_p) # @NL
    JuMP.@constraint(pm.model, q_to == f_owms_q) # @NL
end

We can use the following pipeline:

using PowerModels
using PGLib
using MLJFlux
using Gurobi
using Flux
using MLJ
using L2O
using JuMP
using CUDA
using Logging

matpower_case_name = "pglib_opf_case5_pjm"

network_data = make_basic_network(pglib(matpower_case_name))

branch = network_data["branch"]["2"]
f_bus_index = branch["f_bus"]
f_bus = network_data["bus"]["$f_bus_index"]
t_bus_index = branch["t_bus"]
t_bus = network_data["bus"]["$t_bus_index"]

f_owms = function_ohms_yt_from(branch)

num_samples = 1000

# Define Model
model = MultitargetNeuralNetworkRegressor(;
    builder=FullyConnectedBuilder([64, 32, 8]),
    rng=123,
    epochs=100,
    optimiser=optimiser,
    acceleration=CUDALibs(),
    batch_size=32,
)

# Define the machine
_vm_fr, _vm_to = rand(f_bus["vmin"]:0.0001:f_bus["vmax"], num_samples), rand(t_bus["vmin"]:0.0001:t_bus["vmax"], num_samples)
_va_fr, _va_to = rand(branch["angmin"]:0.0001:branch["angmax"], num_samples), rand(branch["angmin"]:0.0001:branch["angmax"], num_samples)
X = [_vm_fr _vm_to _va_fr _va_to]
y = [i[1] for i in f_owms_val][:,:]

mach = machine(model, X, y)
fit!(mach; verbosity=2)

# Make predictions
predictions = predict(mach, [fill(vm_fr, num_samples) fill(vm_to, num_samples) va_fr va_to])

# Have PowerModels now call the function that uses the Neural Network

function function_ohms_yt_from(::Dict)
    return (vm_fr, vm_to, va_fr, va_to) -> model([vm_fr, vm_to, va_fr, va_to])
end

function function_ohms_yt_to(branch::Dict)
    return (vm_fr, vm_to, va_fr, va_to) -> model([vm_fr, vm_to, va_fr, va_to])
end

pm = instantiate_model(
    network_data,
    ACPPowerModel,
    PowerModels.build_opf;
    setting=Dict("output" => Dict("branch_flows" => true, "duals" => true)),
)

# solve
result = optimize_model!(pm, optimizer=Gurobi.Optimizer)