Skip to content

Commit b29456b

Browse files
andrewrosembergklamike
andrewrosemberg
and
klamike
authored
Lux integration test (#6)
* draft lux + bnk * updat ignore * running cpu * update test * run tests * add random tests * gou compatible * put slurm in global gi * pass T for popf * use float32 to match lux weights * cleanup imports * typo * rm lux compat --------- Co-authored-by: Michael Klamkin <klam@gatech.edu>
1 parent f144a80 commit b29456b

4 files changed

Lines changed: 241 additions & 36 deletions

File tree

‎Project.toml‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,17 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
2525
FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
2626
GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
2727
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
28+
Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
29+
LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
30+
MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
2831
OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2"
32+
Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
2933
PGLib = "07a8691f-3d11-4330-951b-3c50f98338be"
3034
PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655"
35+
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
3136
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
3237
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
3338
pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd"
3439

3540
[targets]
36-
test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels"]
41+
test = ["Test", "CUDA", "GPUArraysCore", "LinearAlgebra", "OpenCL", "pocl_jll", "AcceleratedKernels", "DifferentiationInterface", "FiniteDifferences", "Zygote", "PGLib", "PowerModels", "Lux", "LuxCUDA", "MLUtils", "Optimisers", "Random"]

‎test/power.jl‎

Lines changed: 123 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ end
1616

1717
_convert_array(data::N, backend) where {names,N<:NamedTuple{names}} =
1818
NamedTuple{names}(ExaModels.convert_array(d, backend) for d in data)
19-
function _parse_ac_data_raw(filename)
19+
function _parse_ac_data_raw(filename; T=Float64)
2020
ref = _build_power_ref(filename) # FIXME: only parse once
2121
arcdict = Dict(a => k for (k, a) in enumerate(ref[:arcs]))
2222
busdict = Dict(k => i for (i, (k, _)) in enumerate(ref[:bus]))
@@ -27,24 +27,24 @@ function _parse_ac_data_raw(filename)
2727
begin
2828
loads = [ref[:load][l] for l in ref[:bus_loads][k]]
2929
shunts = [ref[:shunt][s] for s in ref[:bus_shunts][k]]
30-
pd = sum(load["pd"] for load in loads; init = 0.0)
31-
qd = sum(load["qd"] for load in loads; init = 0.0)
32-
gs = sum(shunt["gs"] for shunt in shunts; init = 0.0)
33-
bs = sum(shunt["bs"] for shunt in shunts; init = 0.0)
30+
pd = T(sum(load["pd"] for load in loads; init = 0.0))
31+
qd = T(sum(load["qd"] for load in loads; init = 0.0))
32+
gs = T(sum(shunt["gs"] for shunt in shunts; init = 0.0))
33+
bs = T(sum(shunt["bs"] for shunt in shunts; init = 0.0))
3434
(i = busdict[k], pd = pd, gs = gs, qd = qd, bs = bs)
3535
end for (k, _) in ref[:bus]
3636
],
3737
gen = [
3838
(
3939
i = gendict[k],
40-
cost1 = v["cost"][1],
41-
cost2 = v["cost"][2],
42-
cost3 = v["cost"][3],
40+
cost1 = T(v["cost"][1]),
41+
cost2 = T(v["cost"][2]),
42+
cost3 = T(v["cost"][3]),
4343
bus = busdict[v["gen_bus"]],
4444
) for (k, v) in ref[:gen]
4545
],
4646
arc = [
47-
(i = k, rate_a = ref[:branch][l]["rate_a"], bus = busdict[i]) for
47+
(i = k, rate_a = T(ref[:branch][l]["rate_a"]), bus = busdict[i]) for
4848
(k, (l, i, _)) in enumerate(ref[:arcs])
4949
],
5050
branch = [
@@ -67,34 +67,34 @@ function _parse_ac_data_raw(filename)
6767
t_idx = t_idx,
6868
f_bus = busdict[branch["f_bus"]],
6969
t_bus = busdict[branch["t_bus"]],
70-
c1 = (-g * tr - b * ti) / ttm,
71-
c2 = (-b * tr + g * ti) / ttm,
72-
c3 = (-g * tr + b * ti) / ttm,
73-
c4 = (-b * tr - g * ti) / ttm,
74-
c5 = (g + g_fr) / ttm,
75-
c6 = (b + b_fr) / ttm,
76-
c7 = (g + g_to),
77-
c8 = (b + b_to),
78-
rate_a_sq = branch["rate_a"]^2,
70+
c1 = T((-g * tr - b * ti) / ttm),
71+
c2 = T((-b * tr + g * ti) / ttm),
72+
c3 = T((-g * tr + b * ti) / ttm),
73+
c4 = T((-b * tr - g * ti) / ttm),
74+
c5 = T((g + g_fr) / ttm),
75+
c6 = T((b + b_fr) / ttm),
76+
c7 = T((g + g_to)),
77+
c8 = T((b + b_to)),
78+
rate_a_sq = T(branch["rate_a"]^2),
7979
)
8080
end for (i, branch_raw) in ref[:branch]
8181
],
8282
ref_buses = [busdict[i] for (i, _) in ref[:ref_buses]],
83-
vmax = [v["vmax"] for (_, v) in ref[:bus]],
84-
vmin = [v["vmin"] for (_, v) in ref[:bus]],
85-
pmax = [v["pmax"] for (_, v) in ref[:gen]],
86-
pmin = [v["pmin"] for (_, v) in ref[:gen]],
87-
qmax = [v["qmax"] for (_, v) in ref[:gen]],
88-
qmin = [v["qmin"] for (_, v) in ref[:gen]],
89-
rate_a = [ref[:branch][l]["rate_a"] for (l, _, _) in ref[:arcs]],
90-
angmax = [b["angmax"] for (_, b) in ref[:branch]],
91-
angmin = [b["angmin"] for (_, b) in ref[:branch]],
83+
vmax = [T(v["vmax"]) for (_, v) in ref[:bus]],
84+
vmin = [T(v["vmin"]) for (_, v) in ref[:bus]],
85+
pmax = [T(v["pmax"]) for (_, v) in ref[:gen]],
86+
pmin = [T(v["pmin"]) for (_, v) in ref[:gen]],
87+
qmax = [T(v["qmax"]) for (_, v) in ref[:gen]],
88+
qmin = [T(v["qmin"]) for (_, v) in ref[:gen]],
89+
rate_a = [T(ref[:branch][l]["rate_a"]) for (l, _, _) in ref[:arcs]],
90+
angmax = [T(b["angmax"]) for (_, b) in ref[:branch]],
91+
angmin = [T(b["angmin"]) for (_, b) in ref[:branch]],
9292
)
9393
end
9494

9595
_parse_ac_data(filename) = _parse_ac_data_raw(filename)
96-
function _parse_ac_data(filename, backend)
97-
_convert_array(_parse_ac_data_raw(filename), backend)
96+
function _parse_ac_data(filename, backend; T=Float64)
97+
_convert_array(_parse_ac_data_raw(filename, T=T), backend)
9898
end
9999

100100
function create_ac_power_model(
@@ -188,9 +188,101 @@ function create_ac_power_model(
188188
return ExaModel(c; prod = prod)
189189
end
190190

191-
# TODO: parametric version
191+
# Parametric version
192192

193-
function create_power_models(backend = OpenCLBackend())
193+
function create_parametric_ac_power_model(filename::String = "pglib_opf_case14_ieee.m";
194+
prod::Bool = true, backend = OpenCLBackend(), T=Float64, kwargs...)
195+
data = _parse_ac_data(filename, backend, T=T)
196+
c = ExaCore(T; backend = backend)
197+
198+
va = variable(c, length(data.bus);)
199+
vm = variable(
200+
c,
201+
length(data.bus);
202+
start = fill!(similar(data.bus, T), 1.0),
203+
lvar = data.vmin,
204+
uvar = data.vmax,
205+
)
206+
207+
pg = variable(c, length(data.gen); lvar = data.pmin, uvar = data.pmax)
208+
qg = variable(c, length(data.gen); lvar = data.qmin, uvar = data.qmax)
209+
210+
@allowscalar pd = parameter(c, [b.pd for b in data.bus])
211+
@allowscalar qd = parameter(c, [b.qd for b in data.bus])
212+
213+
p = variable(c, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a)
214+
q = variable(c, length(data.arc); lvar = -data.rate_a, uvar = data.rate_a)
215+
216+
o = objective(c, g.cost1 * pg[g.i]^2 + g.cost2 * pg[g.i] + g.cost3 for g in data.gen)
217+
218+
# Reference bus angle ------------------------------------------------------
219+
c1 = constraint(c, va[i] for i in data.ref_buses)
220+
221+
# Branch power-flow equations ---------------------------------------------
222+
constraint(
223+
c,
224+
(p[b.f_idx] - b.c5 * vm[b.f_bus]^2 -
225+
b.c3 * (vm[b.f_bus] * vm[b.t_bus] * cos(va[b.f_bus] - va[b.t_bus])) -
226+
b.c4 * (vm[b.f_bus] * vm[b.t_bus] * sin(va[b.f_bus] - va[b.t_bus])) for
227+
b in data.branch),
228+
)
229+
230+
constraint(
231+
c,
232+
(q[b.f_idx] + b.c6 * vm[b.f_bus]^2 +
233+
b.c4 * (vm[b.f_bus] * vm[b.t_bus] * cos(va[b.f_bus] - va[b.t_bus])) -
234+
b.c3 * (vm[b.f_bus] * vm[b.t_bus] * sin(va[b.f_bus] - va[b.t_bus])) for
235+
b in data.branch),
236+
)
237+
238+
constraint(
239+
c,
240+
(p[b.t_idx] - b.c7 * vm[b.t_bus]^2 -
241+
b.c1 * (vm[b.t_bus] * vm[b.f_bus] * cos(va[b.t_bus] - va[b.f_bus])) -
242+
b.c2 * (vm[b.t_bus] * vm[b.f_bus] * sin(va[b.t_bus] - va[b.f_bus])) for
243+
b in data.branch),
244+
)
245+
246+
constraint(
247+
c,
248+
(q[b.t_idx] + b.c8 * vm[b.t_bus]^2 +
249+
b.c2 * (vm[b.t_bus] * vm[b.f_bus] * cos(va[b.t_bus] - va[b.f_bus])) -
250+
b.c1 * (vm[b.t_bus] * vm[b.f_bus] * sin(va[b.t_bus] - va[b.f_bus])) for
251+
b in data.branch),
252+
)
253+
254+
# Angle difference limits --------------------------------------------------
255+
constraint(
256+
c,
257+
va[b.f_bus] - va[b.t_bus] for b in data.branch; lcon = data.angmin, ucon = data.angmax,
258+
)
259+
260+
# Apparent power thermal limits -------------------------------------------
261+
constraint(
262+
c,
263+
p[b.f_idx]^2 + q[b.f_idx]^2 - b.rate_a_sq for b in data.branch;
264+
lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf),
265+
)
266+
constraint(
267+
c,
268+
p[b.t_idx]^2 + q[b.t_idx]^2 - b.rate_a_sq for b in data.branch;
269+
lcon = fill!(similar(data.branch, Float64, length(data.branch)), -Inf),
270+
)
271+
272+
# Power balance at each bus -----------------------------------------------
273+
load_balance_p = constraint(c, pd[b.i] + b.gs * vm[b.i]^2 for b in data.bus)
274+
load_balance_q = constraint(c, qd[b.i] - b.bs * vm[b.i]^2 for b in data.bus)
275+
276+
# Map arc & generator variables into the bus balance equations
277+
constraint!(c, load_balance_p, a.bus => p[a.i] for a in data.arc)
278+
constraint!(c, load_balance_q, a.bus => q[a.i] for a in data.arc)
279+
constraint!(c, load_balance_p, g.bus => -pg[g.i] for g in data.gen)
280+
constraint!(c, load_balance_q, g.bus => -qg[g.i] for g in data.gen)
281+
282+
return ExaModel(c; prod = prod)
283+
end
284+
285+
function create_power_models(backend = OpenCLBackend(), T=Float64)
194286
models = ExaModel[]
195287
push!(models, create_ac_power_model("pglib_opf_case14_ieee.m"; backend = backend))
196288
names = ["AC-OPF – IEEE-14"]

‎test/runtests.jl‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,16 @@ using PGLib
1414
using LinearAlgebra
1515

1616
using OpenCL, pocl_jll, AcceleratedKernels
17+
18+
using Lux
19+
using LuxCUDA
20+
using Lux.Training
21+
using MLUtils
22+
using Optimisers
23+
using CUDA
24+
using Random
25+
import GPUArraysCore: @allowscalar
26+
1727
ExaModels.convert_array(x, ::OpenCLBackend) = CLArray(x)
1828
ExaModels.sort!(array::CLArray; lt = isless) = AcceleratedKernels.sort!(array; lt=lt)
1929
function Base.findall(f::F, bitarray::CLArray) where {F<:Function}
@@ -24,10 +34,7 @@ function Base.findall(f::F, bitarray::CLArray) where {F<:Function}
2434
end
2535
Base.findall(bitarray::CLArray) = Base.findall(identity, bitarray)
2636

27-
import GPUArraysCore: @allowscalar
28-
2937
if haskey(ENV, "BNK_TEST_CUDA")
30-
using CUDA
3138
@info "CUDA detected"
3239
end
3340

@@ -37,4 +44,5 @@ include("power.jl")
3744
include("test_viols.jl")
3845
include("test_diff.jl")
3946
include("api.jl")
40-
include("config.jl")
47+
include("config.jl")
48+
include("test_penalty.jl")

‎test/test_penalty.jl‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
function feed_forward_builder(
2+
num_p::Integer,
3+
num_y::Integer,
4+
hidden_layers::AbstractVector{<:Integer};
5+
activation = relu,
6+
)
7+
"""
8+
Builds a Chain of Dense layers with Lux
9+
"""
10+
# Combine all layers: input size, hidden sizes, output size
11+
layer_sizes = [num_p; hidden_layers; num_y]
12+
13+
# Build up a list of Dense layers
14+
dense_layers = Any[]
15+
for i in 1:(length(layer_sizes)-1)
16+
if i < length(layer_sizes) - 1
17+
# Hidden layers with activation
18+
push!(dense_layers, Dense(layer_sizes[i], layer_sizes[i+1], activation))
19+
else
20+
# Final layer with no activation
21+
push!(dense_layers, Dense(layer_sizes[i], layer_sizes[i+1]))
22+
end
23+
end
24+
25+
return Chain(dense_layers...)
26+
end
27+
28+
function test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = gpu_device(),
29+
backend=CPU(),
30+
batch_size = 32,
31+
dataset_size = 3200,
32+
rng = Random.default_rng(),
33+
T = Float64,
34+
)
35+
model = create_parametric_ac_power_model(filename; backend = backend, T=T)
36+
bm = BNK.BatchModel(model, batch_size, config=BNK.BatchModelConfig(:full))
37+
bm_all = BNK.BatchModel(model, dataset_size, config=BNK.BatchModelConfig(:full))
38+
39+
function PenaltyLoss(model, ps, st, Θ)
40+
X̂ , st_new = model(Θ, ps, st)
41+
42+
obj = BNK.objective!(bm, X̂, Θ)
43+
Vc, Vb = BNK.all_violations!(bm, X̂, Θ)
44+
45+
return sum(obj) + 1000 * sum(Vc) + 1000 * sum(Vb), st_new, (;obj=sum(obj), Vc=sum(Vc), Vb=sum(Vb))
46+
end
47+
48+
nvar = model.meta.nvar
49+
ncon = model.meta.ncon
50+
nθ = length(model.θ)
51+
52+
Θ_train = randn(T, nθ, dataset_size) |> dev_gpu
53+
54+
lux_model = feed_forward_builder(nθ, nvar, [320, 320])
55+
56+
ps_model, st_model = Lux.setup(rng, lux_model)
57+
ps_model = ps_model |> dev_gpu
58+
st_model = st_model |> dev_gpu
59+
60+
X̂ , _ = lux_model(Θ_train, ps_model, st_model)
61+
62+
y = BNK.objective!(bm_all, X̂, Θ_train)
63+
64+
@test length(y) == dataset_size
65+
Vc, Vb = BNK.all_violations!(bm_all, X̂, Θ_train)
66+
@test size(Vc) == (ncon, dataset_size)
67+
@test size(Vb) == (nvar, dataset_size)
68+
69+
lagrangian_prev = sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb)
70+
71+
train_state = Training.TrainState(lux_model, ps_model, st_model, Optimisers.Adam(1e-5))
72+
73+
data = DataLoader((Θ_train); batchsize=batch_size, shuffle=true) .|> dev_gpu
74+
for (Θ) in data
75+
_, loss_val, stats, train_state = Training.single_train_step!(
76+
AutoZygote(), # AD backend
77+
PenaltyLoss,
78+
(Θ), # data
79+
train_state
80+
)
81+
end
82+
83+
X̂ , st_model = lux_model(Θ_train, ps_model, st_model)
84+
85+
y = BNK.objective!(bm_all, X̂, Θ_train)
86+
Vc, Vb = BNK.all_violations!(bm_all, X̂, Θ_train)
87+
lagrangian_new = sum(y) + 1000 * sum(Vc) + 1000 * sum(Vb)
88+
89+
@test lagrangian_new < lagrangian_prev
90+
end
91+
92+
@testset "Penalty Training" begin
93+
backend, dev = if haskey(ENV, "BNK_TEST_CUDA")
94+
CUDABackend(), gpu_device()
95+
else
96+
CPU(), cpu_device()
97+
end
98+
99+
test_penalty_training(; filename="pglib_opf_case14_ieee.m", dev_gpu = dev, backend=backend, T=Float32)
100+
end

0 commit comments

Comments
 (0)