diff --git a/Project.toml b/Project.toml index e319649..482f08d 100644 --- a/Project.toml +++ b/Project.toml @@ -25,12 +25,17 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" PGLib = "07a8691f-3d11-4330-951b-3c50f98338be" PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" [targets] -test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels"] +test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA", "MLUtils", "Optimisers", "Random"] diff --git a/test/power.jl b/test/power.jl index 2f43016..a3f6409 100644 --- a/test/power.jl +++ b/test/power.jl @@ -16,7 +16,7 @@ end _convert_array(data::N, backend) where {names,N<:NamedTuple{names}} = NamedTuple{names}(ExaModels.convert_array(d, backend) for d in data) -function _parse_ac_data_raw(filename) +function _parse_ac_data_raw(filename; T=Float64) ref = _build_power_ref(filename) # FIXME: only parse once arcdict = Dict(a => k for (k, a) in enumerate(ref[:arcs])) busdict = Dict(k => i for (i, (k, _)) in enumerate(ref[:bus])) @@ -27,24 +27,24 @@ function _parse_ac_data_raw(filename) begin loads = [ref[:load][l] for l in ref[:bus_loads][k]] shunts = [ref[:shunt][s] for s in ref[:bus_shunts][k]] - pd = sum(load["pd"] for load in loads; init = 0.0) - qd = sum(load["qd"] for load in loads; init = 0.0) - gs = sum(shunt["gs"] for shunt in shunts; init = 0.0) - bs = sum(shunt["bs"] for shunt in shunts; init = 0.0) + pd = T(sum(load["pd"] for load in loads; init = 0.0)) + qd = T(sum(load["qd"] for load in loads; init = 0.0)) + gs = T(sum(shunt["gs"] for shunt in shunts; init = 0.0)) + bs = T(sum(shunt["bs"] for shunt in shunts; init = 0.0)) (i = busdict[k], pd = pd, gs = gs, qd = qd, bs = bs) end for (k, _) in ref[:bus] ], gen = [ ( i = gendict[k], - cost1 = v["cost"][1], - cost2 = v["cost"][2], - cost3 = v["cost"][3], + cost1 = T(v["cost"][1]), + cost2 = T(v["cost"][2]), + cost3 = T(v["cost"][3]), bus = busdict[v["gen_bus"]], ) for (k, v) in ref[:gen] ], arc = [ - (i = k, rate_a = ref[:branch][l]["rate_a"], bus = busdict[i]) for + (i = k, rate_a = T(ref[:branch][l]["rate_a"]), bus = busdict[i]) for (k, (l, i, _)) in enumerate(ref[:arcs]) ], branch = [ @@ -67,34 +67,34 @@ function _parse_ac_data_raw(filename) t_idx = t_idx, f_bus = busdict[branch["f_bus"]], t_bus = busdict[branch["t_bus"]], - c1 = (-g * tr - b * ti) / ttm, - c2 = (-b * tr + g * ti) / ttm, - c3 = (-g * tr + b * ti) / ttm, - c4 = (-b * tr - g * ti) / ttm, - c5 = (g + g_fr) / ttm, - c6 = (b + b_fr) / ttm, - c7 = (g + g_to), - c8 = (b + b_to), - rate_a_sq = branch["rate_a"]^2, + c1 = T((-g * tr - b * ti) / ttm), + c2 = T((-b * tr + g * ti) / ttm), + c3 = T((-g * tr + b * ti) / ttm), + c4 = T((-b * tr - g * ti) / ttm), + c5 = T((g + g_fr) / ttm), + c6 = T((b + b_fr) / ttm), + c7 = T((g + g_to)), + c8 = T((b + b_to)), + rate_a_sq = T(branch["rate_a"]^2), ) end for (i, branch_raw) in ref[:branch] ], ref_buses = [busdict[i] for (i, _) in ref[:ref_buses]], - vmax = [v["vmax"] for (_, v) in ref[:bus]], - vmin = [v["vmin"] for (_, v) in ref[:bus]], - pmax = [v["pmax"] for (_, v) in ref[:gen]], - pmin = [v["pmin"] for (_, v) in ref[:gen]], - qmax = [v["qmax"] for (_, v) in ref[:gen]], - qmin = [v["qmin"] for (_, v) in ref[:gen]], - rate_a = [ref[:branch][l]["rate_a"] for (l, _, _) in ref[:arcs]], - angmax = [b["angmax"] for (_, b) in ref[:branch]], - angmin = [b["angmin"] for (_, b) in ref[:branch]], + vmax = [T(v["vmax"]) for (_, v) in ref[:bus]], + vmin = [T(v["vmin"]) for (_, v) in ref[:bus]], + pmax = [T(v["pmax"]) for (_, v) in ref[:gen]], + pmin = [T(v["pmin"]) for (_, v) in ref[:gen]], + qmax = [T(v["qmax"]) for (_, v) in ref[:gen]], + qmin = [T(v["qmin"]) for (_, v) in ref[:gen]], + rate_a = [T(ref[:branch][l]["rate_a"]) for (l, _, _) in ref[:arcs]], + angmax = [T(b["angmax"]) for (_, b) in ref[:branch]], + angmin = [T(b["angmin"]) for (_, b) in ref[:branch]], ) end _parse_ac_data(filename) = _parse_ac_data_raw(filename) -function _parse_ac_data(filename, backend) - _convert_array(_parse_ac_data_raw(filename), backend) +function _parse_ac_data(filename, backend; T=Float64) + _convert_array(_parse_ac_data_raw(filename, T=T), backend) end function create_ac_power_model( @@ -188,9 +188,101 @@ function create_ac_power_model( return ExaModel(c; prod = prod) end -# TODO: parametric version +# Parametric version -function create_power_models(backend = OpenCLBackend()) +function create_parametric_ac_power_model(filename::String = "pglib_opf_case14_ieee.m"; + prod::Bool = true, backend = OpenCLBackend(), T=Float64, kwargs...) + data = _parse_ac_data(filename, backend, T=T) + c = ExaCore(T; backend = backend) + + va = variable(c, length(data.bus);) + vm = variable( + c, + length(data.bus); + start = fill!(similar(data.bus, T), 1.0), + lvar = data.vmin, + uvar = data.vmax, + ) + + pg = variable(c, length(data.gen); lvar = data.pmin, uvar = data.pmax) + qg = variable(c, length(data.gen); lvar = data.qmin, uvar = data.qmax) + + @allowscalar pd = parameter(c, [b.pd for b in data.bus]) + @allowscalar qd = parameter(c, [b.qd for b in data.bus]) + + p = variable(c, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + q = variable(c, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + + o = objective(c, g.cost1 * pg[g.i]^2 + g.cost2 * pg[g.i] + g.cost3 for g in data.gen) + + # Reference bus angle ------------------------------------------------------ + c1 = constraint(c, va[i] for i in data.ref_buses) + + # Branch power-flow equations --------------------------------------------- + constraint( + c, + (p[b.f_idx] - b.c5 * vm[b.f_bus]^2 - + b.c3 * (vm[b.f_bus] * vm[b.t_bus] * cos(va[b.f_bus] - va[b.t_bus])) - + b.c4 * (vm[b.f_bus] * vm[b.t_bus] * sin(va[b.f_bus] - va[b.t_bus])) for + b in data.branch), + ) + + constraint( + c, + (q[b.f_idx] + b.c6 * vm[b.f_bus]^2 + + b.c4 * (vm[b.f_bus] * vm[b.t_bus] * cos(va[b.f_bus] - va[b.t_bus])) - + b.c3 * (vm[b.f_bus] * vm[b.t_bus] * sin(va[b.f_bus] - va[b.t_bus])) for + b in data.branch), + ) + + constraint( + c, + (p[b.t_idx] - b.c7 * vm[b.t_bus]^2 - + b.c1 * (vm[b.t_bus] * vm[b.f_bus] * cos(va[b.t_bus] - va[b.f_bus])) - + b.c2 * (vm[b.t_bus] * vm[b.f_bus] * sin(va[b.t_bus] - va[b.f_bus])) for + b in data.branch), + ) + + constraint( + c, + (q[b.t_idx] + b.c8 * vm[b.t_bus]^2 + + b.c2 * (vm[b.t_bus] * vm[b.f_bus] * cos(va[b.t_bus] - va[b.f_bus])) - + b.c1 * (vm[b.t_bus] * vm[b.f_bus] * sin(va[b.t_bus] - va[b.f_bus])) for + b in data.branch), + ) + + # Angle difference limits -------------------------------------------------- + constraint( + c, + va[b.f_bus] - va[b.t_bus] for b in data.branch; lcon = data.angmin, ucon = data.angmax, + ) + + # Apparent power thermal limits ------------------------------------------- + constraint( + c, + p[b.f_idx]^2 + q[b.f_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf), + ) + constraint( + c, + p[b.t_idx]^2 + q[b.t_idx]^2 - b.rate_a_sq for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf), + ) + + # Power balance at each bus ----------------------------------------------- + load_balance_p = constraint(c, pd[b.i] + b.gs * vm[b.i]^2 for b in data.bus) + load_balance_q = constraint(c, qd[b.i] - b.bs * vm[b.i]^2 for b in data.bus) + + # Map arc & generator variables into the bus balance equations + constraint!(c, load_balance_p, a.bus => p[a.i] for a in data.arc) + constraint!(c, load_balance_q, a.bus => q[a.i] for a in data.arc) + constraint!(c, load_balance_p, g.bus => -pg[g.i] for g in data.gen) + constraint!(c, load_balance_q, g.bus => -qg[g.i] for g in data.gen) + + return ExaModel(c; prod = prod) +end + +function create_power_models(backend = OpenCLBackend(), T=Float64) models = ExaModel[] push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend)) names = ["AC-OPF – IEEE-14"] diff --git a/test/runtests.jl b/test/runtests.jl index ae48720..1b1dfd1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,6 +14,16 @@ using PGLib using LinearAlgebra using OpenCL, pocl_jll, AcceleratedKernels + +using Lux +using LuxCUDA +using Lux.Training +using MLUtils +using Optimisers +using CUDA +using Random +import GPUArraysCore: @allowscalar + ExaModels.convert_array(x, ::OpenCLBackend) = CLArray(x) ExaModels.sort!(array::CLArray; lt = isless) = AcceleratedKernels.sort!(array; lt=lt) function Base.findall(f::F, bitarray::CLArray) where {F<:Function} @@ -24,10 +34,7 @@ function Base.findall(f::F, bitarray::CLArray) where {F<:Function} end Base.findall(bitarray::CLArray) = Base.findall(identity, bitarray) -import GPUArraysCore: @allowscalar - if haskey(ENV, "BNK_TEST_CUDA") - using CUDA @info "CUDA detected" end @@ -37,4 +44,5 @@ include("power.jl") include("test_viols.jl") include("test_diff.jl") include("api.jl") -include("config.jl") \ No newline at end of file +include("config.jl") +include("test_penalty.jl") \ No newline at end of file diff --git a/test/test_penalty.jl b/test/test_penalty.jl new file mode 100644 index 0000000..7723f61 --- /dev/null +++ b/test/test_penalty.jl @@ -0,0 +1,100 @@ +function feed_forward_builder( + num_p::Integer, + num_y::Integer, + hidden_layers::AbstractVector{<:Integer}; + activation = relu, +) + """ + Builds a Chain of Dense layers with Lux + """ + # Combine all layers: input size, hidden sizes, output size + layer_sizes = [num_p; hidden_layers; num_y] + + # Build up a list of Dense layers + dense_layers = Any[] + for i in 1:(length(layer_sizes)-1) + if i < length(layer_sizes) - 1 + # Hidden layers with activation + push!(dense_layers, Dense(layer_sizes[i], layer_sizes[i+1], activation)) + else + # Final layer with no activation + push!(dense_layers, Dense(layer_sizes[i], layer_sizes[i+1])) + end + end + + return Chain(dense_layers...) +end + +function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = gpu_device(), + backend=CPU(), + batch_size = 32, + dataset_size = 3200, + rng = Random.default_rng(), + T = Float64, +) + model = create_parametric_ac_power_model(filename; backend = backend, T=T) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) + bm_all = BNK.BatchModel(model, dataset_size, config=BNK.BatchModelConfig(:full)) + + function PenaltyLoss(model, ps, st, Θ) + X̂ , st_new = model(Θ, ps, st) + + obj = BNK.objective!(bm, X̂, Θ) + Vc, Vb = BNK.all_violations!(bm, X̂, Θ) + + return sum(obj) + 1000 * sum(Vc) + 1000 * sum(Vb), st_new, (;obj=sum(obj), Vc=sum(Vc), Vb=sum(Vb)) + end + + nvar = model.meta.nvar + ncon = model.meta.ncon + nθ = length(model.θ) + + Θ_train = randn(T, nθ, dataset_size) |> dev_gpu + + lux_model = feed_forward_builder(nθ, nvar, [320, 320]) + + ps_model, st_model = Lux.setup(rng, lux_model) + ps_model = ps_model |> dev_gpu + st_model = st_model |> dev_gpu + + X̂ , _ = lux_model(Θ_train, ps_model, st_model) + + y = BNK.objective!(bm_all, X̂, Θ_train) + + @test length(y) == dataset_size + Vc, Vb = BNK.all_violations!(bm_all, X̂, Θ_train) + @test size(Vc) == (ncon, dataset_size) + @test size(Vb) == (nvar, dataset_size) + + lagrangian_prev = sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb) + + train_state = Training.TrainState(lux_model, ps_model, st_model, Optimisers.Adam(1e-5)) + + data = DataLoader((Θ_train); batchsize=batch_size, shuffle=true) .|> dev_gpu + for (Θ) in data + _, loss_val, stats, train_state = Training.single_train_step!( + AutoZygote(), # AD backend + PenaltyLoss, + (Θ), # data + train_state + ) + end + + X̂ , st_model = lux_model(Θ_train, ps_model, st_model) + + y = BNK.objective!(bm_all, X̂, Θ_train) + Vc, Vb = BNK.all_violations!(bm_all, X̂, Θ_train) + lagrangian_new = sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb) + + @test lagrangian_new < lagrangian_prev +end + +@testset "Penalty Training" begin + backend, dev = if haskey(ENV, "BNK_TEST_CUDA") + CUDABackend(), gpu_device() + else + CPU(), cpu_device() + end + + test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = dev, backend=backend, T=Float32) +end \ No newline at end of file