From 2b84f4228a40a4bfb9d3825d9331294899684cdc Mon Sep 17 00:00:00 2001 From: Andrew Rosemberg Date: Wed, 9 Jul 2025 16:17:41 -0400 Subject: [PATCH 01/13] draft lux + bnk --- Project.toml | 6 ++- test/power.jl | 116 ++++++++++++++++++++++++++++++++++++++++++- test/test_penalty.jl | 78 +++++++++++++++++++++++++++++ 3 files changed, 198 insertions(+), 2 deletions(-) create mode 100644 test/test_penalty.jl diff --git a/Project.toml b/Project.toml index e319649..b6106bb 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,8 @@ BNKJuMP = "JuMP" [compat] ExaModels = "0.8.3" +Lux = "1.15.0" +LuxCUDA = "0.3.3" [extras] AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" @@ -31,6 +33,8 @@ PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" [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"] diff --git a/test/power.jl b/test/power.jl index 2f43016..43ed3dd 100644 --- a/test/power.jl +++ b/test/power.jl @@ -188,7 +188,121 @@ function create_ac_power_model( return ExaModel(c; prod = prod) end -# TODO: parametric version +# Parametric version + +function c_active_power_balance_demand_polar_nodemand(b, vm) + return + b.gs * vm^2 +end + +function c_reactive_power_balance_demand_polar_nodemand(b, vm) + return - b.bs * vm^2 +end + +function build_polar_opf_demand_parametric(data; backend = nothing, T=Float64, kwargs...) + core = ExaCore(T; backend = backend) + + va = variable(core, length(data.bus);) + vm = variable( + core, + length(data.bus); + start = fill!(similar(data.bus, Float64), 1.0), + lvar = data.vmin, + uvar = data.vmax, + ) + + pg = variable(core, length(data.gen); lvar = data.pmin, uvar = data.pmax) + qg = variable(core, length(data.gen); lvar = data.qmin, uvar = data.qmax) + + pd = parameter(core, [b.pd for b in data.bus]) + qd = parameter(core, [b.qd for b in data.bus]) + + p = variable(core, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + q = variable(core, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) + + o = objective( + core, gen_cost(g, pg[g.i]) for g in data.gen) + + c_ref_angle = constraint(core, c_ref_angle_polar(va[i]) for i in data.ref_buses) + + c_to_active_power_flow = constraint(core, c_to_active_power_flow_polar(b, p[b.f_idx], + vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + + c_to_reactive_power_flow = constraint(core, c_to_reactive_power_flow_polar(b, q[b.f_idx], + vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + + c_from_active_power_flow = constraint(core, c_from_active_power_flow_polar(b, p[b.t_idx], + vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + + c_from_reactive_power_flow = constraint(core, c_from_reactive_power_flow_polar(b, q[b.t_idx], + vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + + c_phase_angle_diff = constraint( + core, + c_phase_angle_diff_polar(b,va[b.f_bus],va[b.t_bus]) for b in data.branch; + lcon = data.angmin, + ucon = data.angmax, + ) + + c_active_power_balance = constraint(core, pd[b.i] + c_active_power_balance_demand_polar_nodemand(b, vm[b.i]) for b in data.bus) + + c_reactive_power_balance = constraint(core, qd[b.i] + c_reactive_power_balance_demand_polar_nodemand(b, vm[b.i]) for b in data.bus) + + c_active_power_balance_arcs = constraint!(core, c_active_power_balance, a.bus => p[a.i] for a in data.arc) + c_reactive_power_balance_arcs = constraint!(core, c_reactive_power_balance, a.bus => q[a.i] for a in data.arc) + + c_active_power_balance_gen = constraint!(core, c_active_power_balance, g.bus => -pg[g.i] for g in data.gen) + c_reactive_power_balance_gen = constraint!(core, c_reactive_power_balance, g.bus => -qg[g.i] for g in data.gen) + + c_from_thermal_limit = constraint( + core, c_thermal_limit(b,p[b.f_idx],q[b.f_idx]) for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf), + ) + + c_to_thermal_limit = constraint( + core, c_thermal_limit(b,p[b.t_idx],q[b.t_idx]) + for b in data.branch; + lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf), + ) + + model =ExaModel(core; kwargs...) + + vars = ( + va = va, + vm = vm, + pg = pg, + qg = qg, + p = p, + q = q + ) + + cons = ( + c_ref_angle = c_ref_angle, + c_to_active_power_flow = c_to_active_power_flow, + c_to_reactive_power_flow = c_to_reactive_power_flow, + c_from_active_power_flow = c_from_active_power_flow, + c_from_reactive_power_flow = c_from_reactive_power_flow, + c_phase_angle_diff = c_phase_angle_diff, + c_active_power_balance = c_active_power_balance, + c_reactive_power_balance = c_reactive_power_balance, + c_from_thermal_limit = c_from_thermal_limit, + c_to_thermal_limit = c_to_thermal_limit + ) + + return model, vars, cons, (pd=pd, qd=qd) +end + +function create_parametric_ac_power_model( + filename; + backend = OpenCLBackend(), + T = Float64, + kwargs..., +) + + data, _ = parse_ac_power_data(filename) + data = convert_data(data, backend) + + return build_polar_opf_demand_parametric(data, backend = backend, T=T, kwargs...) +end function create_power_models(backend = OpenCLBackend()) models = ExaModel[] diff --git a/test/test_penalty.jl b/test/test_penalty.jl new file mode 100644 index 0000000..a6e3f5c --- /dev/null +++ b/test/test_penalty.jl @@ -0,0 +1,78 @@ +using Lux +using LuxCUDA +using Lux.Training +using MLUtils +using Optimisers +using CUDA + +dev_gpu = gpu_device() + +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 + +model = create_parametric_ac_power_model("pglib_opf_case14_ieee.m"; backend = CUDABackend()) +bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) + +function PenaltyLoss(model, ps, st, Θ) + X̂ , _ = model(Θ, ps, st) + + y = BNK.objective!(bm, X̂, Θ) + Vc, Vb = BNK.all_violations!(bm, X̂, Θ) + + return sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb) +end + +MOD = CUDA +batch_size = 32 +dataset_size = 3200 + +nvar = model.meta.nvar +ncon = model.meta.ncon +nθ = length(model.θ) + +# X̂ = MOD.randn(nvar, dataset_size) +Θ_train = MOD.randn(nθ, dataset_size) + +lux_model = feed_forward_builder(nθ, nvar, [320, 320]) + +ps_model, st_model = Lux.setup(rng, lux_model) +X̂ , _ = model(Xtrain, ps_model, st_model) + +y = BNK.objective!(bm, X̂, Θ_train) +Vc, Vb = BNK.all_violations!(bm, X̂, Θ_train) + +train_state = Training.TrainState(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 \ No newline at end of file From 101f4085a90b11bc168f79caf2c3ba3b7c92ec82 Mon Sep 17 00:00:00 2001 From: Andrew Rosemberg Date: Wed, 9 Jul 2025 16:28:38 -0400 Subject: [PATCH 02/13] updat ignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 230c5ed..e659031 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /Manifest*.toml +*.out From 1ff340131ccd1216fa574a8645d12a490a990828 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 9 Jul 2025 17:13:17 -0400 Subject: [PATCH 03/13] running cpu --- Project.toml | 8 ++- test/power.jl | 152 ++++++++++++++++++------------------------- test/test_penalty.jl | 25 +++---- 3 files changed, 84 insertions(+), 101 deletions(-) diff --git a/Project.toml b/Project.toml index b6106bb..9d6239f 100644 --- a/Project.toml +++ b/Project.toml @@ -27,14 +27,16 @@ 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" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" PGLib = "07a8691f-3d11-4330-951b-3c50f98338be" PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" -Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" [targets] -test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA"] +test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA", "MLUtils", "Optimisers"] diff --git a/test/power.jl b/test/power.jl index 43ed3dd..21368e3 100644 --- a/test/power.jl +++ b/test/power.jl @@ -190,118 +190,96 @@ end # Parametric version -function c_active_power_balance_demand_polar_nodemand(b, vm) - return + b.gs * vm^2 -end - -function c_reactive_power_balance_demand_polar_nodemand(b, vm) - return - b.bs * vm^2 -end - -function build_polar_opf_demand_parametric(data; backend = nothing, T=Float64, kwargs...) - core = ExaCore(T; backend = backend) +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) + c = ExaCore(T; backend = backend) - va = variable(core, length(data.bus);) + va = variable(c, length(data.bus);) vm = variable( - core, + c, length(data.bus); start = fill!(similar(data.bus, Float64), 1.0), lvar = data.vmin, uvar = data.vmax, ) - pg = variable(core, length(data.gen); lvar = data.pmin, uvar = data.pmax) - qg = variable(core, length(data.gen); lvar = data.qmin, uvar = data.qmax) - - pd = parameter(core, [b.pd for b in data.bus]) - qd = parameter(core, [b.qd for b in data.bus]) - - p = variable(core, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) - q = variable(core, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a) - - o = objective( - core, gen_cost(g, pg[g.i]) for g in data.gen) - - c_ref_angle = constraint(core, c_ref_angle_polar(va[i]) for i in data.ref_buses) + pg = variable(c, length(data.gen); lvar = data.pmin, uvar = data.pmax) + qg = variable(c, length(data.gen); lvar = data.qmin, uvar = data.qmax) - c_to_active_power_flow = constraint(core, c_to_active_power_flow_polar(b, p[b.f_idx], - vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + pd = parameter(c, [b.pd for b in data.bus]) + qd = parameter(c, [b.qd for b in data.bus]) - c_to_reactive_power_flow = constraint(core, c_to_reactive_power_flow_polar(b, q[b.f_idx], - vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + 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) - c_from_active_power_flow = constraint(core, c_from_active_power_flow_polar(b, p[b.t_idx], - vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + o = objective(c, g.cost1 * pg[g.i]^2 + g.cost2 * pg[g.i] + g.cost3 for g in data.gen) - c_from_reactive_power_flow = constraint(core, c_from_reactive_power_flow_polar(b, q[b.t_idx], - vm[b.f_bus],vm[b.t_bus],va[b.f_bus],va[b.t_bus]) for b in data.branch) + # Reference bus angle ------------------------------------------------------ + c1 = constraint(c, va[i] for i in data.ref_buses) - c_phase_angle_diff = constraint( - core, - c_phase_angle_diff_polar(b,va[b.f_bus],va[b.t_bus]) for b in data.branch; - lcon = data.angmin, - ucon = data.angmax, + # 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), ) - c_active_power_balance = constraint(core, pd[b.i] + c_active_power_balance_demand_polar_nodemand(b, vm[b.i]) for b in data.bus) - - c_reactive_power_balance = constraint(core, qd[b.i] + c_reactive_power_balance_demand_polar_nodemand(b, vm[b.i]) for b in data.bus) + 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), + ) - c_active_power_balance_arcs = constraint!(core, c_active_power_balance, a.bus => p[a.i] for a in data.arc) - c_reactive_power_balance_arcs = constraint!(core, c_reactive_power_balance, a.bus => q[a.i] for a in data.arc) + 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), + ) - c_active_power_balance_gen = constraint!(core, c_active_power_balance, g.bus => -pg[g.i] for g in data.gen) - c_reactive_power_balance_gen = constraint!(core, c_reactive_power_balance, g.bus => -qg[g.i] for g in data.gen) + 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), + ) - c_from_thermal_limit = constraint( - core, c_thermal_limit(b,p[b.f_idx],q[b.f_idx]) for b in data.branch; - lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf), - ) + # Angle difference limits -------------------------------------------------- + constraint( + c, + va[b.f_bus] - va[b.t_bus] for b in data.branch; lcon = data.angmin, ucon = data.angmax, + ) - c_to_thermal_limit = constraint( - core, c_thermal_limit(b,p[b.t_idx],q[b.t_idx]) - for b in data.branch; + # 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), ) - - model =ExaModel(core; kwargs...) - - vars = ( - va = va, - vm = vm, - pg = pg, - qg = qg, - p = p, - q = q - ) - - cons = ( - c_ref_angle = c_ref_angle, - c_to_active_power_flow = c_to_active_power_flow, - c_to_reactive_power_flow = c_to_reactive_power_flow, - c_from_active_power_flow = c_from_active_power_flow, - c_from_reactive_power_flow = c_from_reactive_power_flow, - c_phase_angle_diff = c_phase_angle_diff, - c_active_power_balance = c_active_power_balance, - c_reactive_power_balance = c_reactive_power_balance, - c_from_thermal_limit = c_from_thermal_limit, - c_to_thermal_limit = c_to_thermal_limit + 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), ) - return model, vars, cons, (pd=pd, qd=qd) -end - -function create_parametric_ac_power_model( - filename; - backend = OpenCLBackend(), - T = Float64, - kwargs..., -) + # 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) - data, _ = parse_ac_power_data(filename) - data = convert_data(data, backend) + # 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 build_polar_opf_demand_parametric(data, backend = backend, T=T, kwargs...) + return ExaModel(c; prod = prod) end function create_power_models(backend = OpenCLBackend()) diff --git a/test/test_penalty.jl b/test/test_penalty.jl index a6e3f5c..641d6c4 100644 --- a/test/test_penalty.jl +++ b/test/test_penalty.jl @@ -4,6 +4,7 @@ using Lux.Training using MLUtils using Optimisers using CUDA +using Random dev_gpu = gpu_device() @@ -33,39 +34,41 @@ function feed_forward_builder( return Chain(dense_layers...) end - -model = create_parametric_ac_power_model("pglib_opf_case14_ieee.m"; backend = CUDABackend()) +# backend = CUDABackend() +batch_size = 32 +model = create_parametric_ac_power_model("pglib_opf_case14_ieee.m"; backend = CPU()) bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) function PenaltyLoss(model, ps, st, Θ) - X̂ , _ = model(Θ, ps, st) + X̂ , st_new = model(Θ, ps, st) - y = BNK.objective!(bm, X̂, Θ) + obj = BNK.objective!(bm, X̂, Θ) Vc, Vb = BNK.all_violations!(bm, X̂, Θ) - return sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb) + return sum(obj) + 1000 * sum(Vc) + 1000 * sum(Vb), st_new, (;obj=sum(obj), Vc=sum(Vc), Vb=sum(Vb)) end MOD = CUDA -batch_size = 32 dataset_size = 3200 +rng = Random.default_rng() nvar = model.meta.nvar ncon = model.meta.ncon nθ = length(model.θ) # X̂ = MOD.randn(nvar, dataset_size) +# Θ_train = randn(nθ, dataset_size) Θ_train = MOD.randn(nθ, dataset_size) lux_model = feed_forward_builder(nθ, nvar, [320, 320]) ps_model, st_model = Lux.setup(rng, lux_model) -X̂ , _ = model(Xtrain, ps_model, st_model) +X̂ , _ = lux_model(Θ_train, ps_model, st_model) -y = BNK.objective!(bm, X̂, Θ_train) -Vc, Vb = BNK.all_violations!(bm, X̂, Θ_train) +y = BNK.objective!(bm, X̂[:,1:32], Θ_train[:,1:32]) +Vc, Vb = BNK.all_violations!(bm, X̂[:,1:32], Θ_train[:,1:32]) -train_state = Training.TrainState(model, ps_model, st_model, Optimisers.Adam(1e-5)) +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 @@ -75,4 +78,4 @@ for (Θ) in data (Θ), # data train_state ) -end \ No newline at end of file +end From c34c2ef32f0a8ddee4d1d70fd1cf687e3195dce7 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 9 Jul 2025 17:41:28 -0400 Subject: [PATCH 04/13] update test --- test/power.jl | 4 +- test/runtests.jl | 9 ++++ test/test_penalty.jl | 105 +++++++++++++++++++++++++------------------ 3 files changed, 73 insertions(+), 45 deletions(-) diff --git a/test/power.jl b/test/power.jl index 21368e3..da79c73 100644 --- a/test/power.jl +++ b/test/power.jl @@ -207,8 +207,8 @@ function create_parametric_ac_power_model(filename::String = "pglib_opf_case14_i pg = variable(c, length(data.gen); lvar = data.pmin, uvar = data.pmax) qg = variable(c, length(data.gen); lvar = data.qmin, uvar = data.qmax) - pd = parameter(c, [b.pd for b in data.bus]) - qd = parameter(c, [b.qd for b in data.bus]) + @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) diff --git a/test/runtests.jl b/test/runtests.jl index ae48720..ca206e6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,6 +14,15 @@ using PGLib using LinearAlgebra using OpenCL, pocl_jll, AcceleratedKernels + +using Lux +using LuxCUDA +using Lux.Training +using MLUtils +using Optimisers +using CUDA +using Random + 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} diff --git a/test/test_penalty.jl b/test/test_penalty.jl index 641d6c4..9d8e09e 100644 --- a/test/test_penalty.jl +++ b/test/test_penalty.jl @@ -1,13 +1,3 @@ -using Lux -using LuxCUDA -using Lux.Training -using MLUtils -using Optimisers -using CUDA -using Random - -dev_gpu = gpu_device() - function feed_forward_builder( num_p::Integer, num_y::Integer, @@ -34,48 +24,77 @@ function feed_forward_builder( return Chain(dense_layers...) end -# backend = CUDABackend() -batch_size = 32 -model = create_parametric_ac_power_model("pglib_opf_case14_ieee.m"; backend = CPU()) -bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) -function PenaltyLoss(model, ps, st, Θ) - X̂ , st_new = model(Θ, ps, st) +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(), +) + model = create_parametric_ac_power_model(filename; backend = backend) + bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full)) + bm_all = BNK.BatchModel(model, dataset_size, config=BNK.BatchModelConfig(:full)) - obj = BNK.objective!(bm, X̂, Θ) - Vc, Vb = BNK.all_violations!(bm, X̂, Θ) + function PenaltyLoss(model, ps, st, Θ) + X̂ , st_new = model(Θ, ps, st) - return sum(obj) + 1000 * sum(Vc) + 1000 * sum(Vb), st_new, (;obj=sum(obj), Vc=sum(Vc), Vb=sum(Vb)) -end + obj = BNK.objective!(bm, X̂, Θ) + Vc, Vb = BNK.all_violations!(bm, X̂, Θ) -MOD = CUDA -dataset_size = 3200 -rng = Random.default_rng() + 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 = if typeof(backend) == CUDABackend + CUDA.randn(nθ, dataset_size) + else + randn(nθ, dataset_size) + end -nvar = model.meta.nvar -ncon = model.meta.ncon -nθ = length(model.θ) + lux_model = feed_forward_builder(nθ, nvar, [320, 320]) -# X̂ = MOD.randn(nvar, dataset_size) -# Θ_train = randn(nθ, dataset_size) -Θ_train = MOD.randn(nθ, dataset_size) + ps_model, st_model = Lux.setup(rng, lux_model) + X̂ , _ = lux_model(Θ_train, ps_model, st_model) -lux_model = feed_forward_builder(nθ, nvar, [320, 320]) + y = BNK.objective!(bm_all, X̂, Θ_train) -ps_model, st_model = Lux.setup(rng, lux_model) -X̂ , _ = lux_model(Θ_train, ps_model, st_model) + @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) -y = BNK.objective!(bm, X̂[:,1:32], Θ_train[:,1:32]) -Vc, Vb = BNK.all_violations!(bm, X̂[:,1:32], Θ_train[:,1:32]) + lagrangian_prev = sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb) -train_state = Training.TrainState(lux_model, ps_model, st_model, Optimisers.Adam(1e-5)) + 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 - ) + 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 = if haskey(ENV, "BNK_TEST_CUDA") + CUDABackend() + else + CPU() + end + + test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = gpu_device(), backend=backend) +end \ No newline at end of file From 4c87aee32bb12cabf58ceedea51c272f845cb17f Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 9 Jul 2025 17:42:44 -0400 Subject: [PATCH 05/13] run tests --- test/runtests.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index ca206e6..6fd3370 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,4 +46,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 From 42c5400bdbd5c79a51c13741821af35e92fa69b4 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 9 Jul 2025 17:50:48 -0400 Subject: [PATCH 06/13] add random tests --- Project.toml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 9d6239f..dd6612c 100644 --- a/Project.toml +++ b/Project.toml @@ -29,14 +29,15 @@ 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" -MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" -Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" [targets] -test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA", "MLUtils", "Optimisers"] +test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA", "MLUtils", "Optimisers", "Random"] From a530521542f14da9c1a289945e51a5f89ff3c7f2 Mon Sep 17 00:00:00 2001 From: Andrew Rosemberg Date: Wed, 9 Jul 2025 18:50:06 -0400 Subject: [PATCH 07/13] gou compatible --- test/test_penalty.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/test/test_penalty.jl b/test/test_penalty.jl index 9d8e09e..7b6ea75 100644 --- a/test/test_penalty.jl +++ b/test/test_penalty.jl @@ -48,15 +48,14 @@ function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = g ncon = model.meta.ncon nθ = length(model.θ) - Θ_train = if typeof(backend) == CUDABackend - CUDA.randn(nθ, dataset_size) - else - randn(nθ, dataset_size) - end + Θ_train = randn(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) From 75a3dc83a57683a0e845f2cfcad2c1db50042a0b Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 22:24:43 -0400 Subject: [PATCH 08/13] put slurm in global gi --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index e659031..230c5ed 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1 @@ /Manifest*.toml -*.out From 8e04da6b82b5ad4c1c1103da5646109ae6b192c4 Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 22:50:25 -0400 Subject: [PATCH 09/13] pass T for popf --- test/power.jl | 66 +++++++++++++++++++++++++-------------------------- 1 file changed, 33 insertions(+), 33 deletions(-) diff --git a/test/power.jl b/test/power.jl index da79c73..41fc980 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( @@ -192,14 +192,14 @@ end 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) + 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, Float64), 1.0), + start = fill!(similar(data.bus, T), 1.0), lvar = data.vmin, uvar = data.vmax, ) @@ -282,9 +282,9 @@ function create_parametric_ac_power_model(filename::String = "pglib_opf_case14_i return ExaModel(c; prod = prod) end -function create_power_models(backend = OpenCLBackend()) +function create_power_models(backend = OpenCLBackend(), T=Float64) models = ExaModel[] - push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend)) + push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend, T=T)) names = ["AC-OPF – IEEE-14"] return models, names end \ No newline at end of file From 4d36eae408818953ddace4916e16ac3fe9690ff7 Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 22:50:49 -0400 Subject: [PATCH 10/13] use float32 to match lux weights --- test/test_penalty.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/test/test_penalty.jl b/test/test_penalty.jl index 7b6ea75..7723f61 100644 --- a/test/test_penalty.jl +++ b/test/test_penalty.jl @@ -30,8 +30,9 @@ function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = g batch_size = 32, dataset_size = 3200, rng = Random.default_rng(), + T = Float64, ) - model = create_parametric_ac_power_model(filename; backend = backend) + 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)) @@ -48,7 +49,7 @@ function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = g ncon = model.meta.ncon nθ = length(model.θ) - Θ_train = randn(nθ, dataset_size) |> dev_gpu + Θ_train = randn(T, nθ, dataset_size) |> dev_gpu lux_model = feed_forward_builder(nθ, nvar, [320, 320]) @@ -89,11 +90,11 @@ function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = g end @testset "Penalty Training" begin - backend = if haskey(ENV, "BNK_TEST_CUDA") - CUDABackend() + backend, dev = if haskey(ENV, "BNK_TEST_CUDA") + CUDABackend(), gpu_device() else - CPU() + CPU(), cpu_device() end - test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = gpu_device(), backend=backend) + test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = dev, backend=backend, T=Float32) end \ No newline at end of file From a971cd96b51cb4c2d0b0fc858e1e2bdcdf0aa638 Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 22:51:12 -0400 Subject: [PATCH 11/13] cleanup imports --- test/runtests.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 6fd3370..1b1dfd1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,6 +22,7 @@ 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) @@ -33,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 From dfaf03cdcfd2264ff9e8876bb758dc0bfa58fc44 Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 23:03:25 -0400 Subject: [PATCH 12/13] typo --- test/power.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/power.jl b/test/power.jl index 41fc980..a3f6409 100644 --- a/test/power.jl +++ b/test/power.jl @@ -284,7 +284,7 @@ end function create_power_models(backend = OpenCLBackend(), T=Float64) models = ExaModel[] - push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend, T=T)) + push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend)) names = ["AC-OPF – IEEE-14"] return models, names end \ No newline at end of file From dd873c9fa979e721197e3cf20d3bd83ac32b1bfe Mon Sep 17 00:00:00 2001 From: Michael Klamkin Date: Wed, 9 Jul 2025 23:27:22 -0400 Subject: [PATCH 13/13] rm lux compat --- Project.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/Project.toml b/Project.toml index dd6612c..482f08d 100644 --- a/Project.toml +++ b/Project.toml @@ -17,8 +17,6 @@ BNKJuMP = "JuMP" [compat] ExaModels = "0.8.3" -Lux = "1.15.0" -LuxCUDA = "0.3.3" [extras] AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"