Lux integration test

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"]

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"]

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")
+include("config.jl")
+include("test_penalty.jl")

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