diff --git a/perf/acopf/ACOPF.jl b/perf/acopf/ACOPF.jl new file mode 100644 index 0000000..5139959 --- /dev/null +++ b/perf/acopf/ACOPF.jl @@ -0,0 +1,21 @@ +# Vectorized AC-OPF on ArrayDiff — entry point. `include` this file. +# +# * `structured.jl` — GPU-friendly constant matrix types (gather/scatter and +# ELLPACK) that ArrayDiff keeps by reference on the tape. +# * `data.jl` — case9mod (JuMP tutorial) and PowerModels-derived data. +# * `solver.jl` — first-order augmented-Lagrangian solver (projected Adam), +# built on `eval_residual!` / `eval_residual_jtprod!`; storage-generic. +# * `models.jl` — the two model builders: `build_rect` (JuMP-tutorial form, +# rectangular voltages + Ybus) and `build_polar` (GenOpt/ExaModels form, +# polar voltages with sin/cos branch flows). + +import ArrayDiff +import JuMP +import LinearAlgebra +import MathOptInterface as MOI +import SparseArrays + +include("structured.jl") +include("data.jl") +include("solver.jl") +include("models.jl") diff --git a/perf/acopf/Project.toml b/perf/acopf/Project.toml new file mode 100644 index 0000000..2478617 --- /dev/null +++ b/perf/acopf/Project.toml @@ -0,0 +1,14 @@ +[deps] +ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/perf/acopf/README.md b/perf/acopf/README.md new file mode 100644 index 0000000..f9b69db --- /dev/null +++ b/perf/acopf/README.md @@ -0,0 +1,94 @@ +# Vectorized AC-OPF on ArrayDiff (first-order, GPU-ready) + +Two AC-OPF formulations written as *whole-vector* expressions +(matrix–vector products + broadcasts) so ArrayDiff's array tape evaluates +them with a handful of `mul!`/broadcast kernels — the same code path on CPU +(`Vector{Float64}`), JLArrays (GPU semantics without a GPU) and CUDA +(`CuVector{Float64}`): + +1. **rect** — the JuMP tutorial (`optimal_power_flow.jl`) form: rectangular + complex voltages, power balance `S_G - S_D = V .* conj(Y V)` split into + real/imaginary parts with the sparse bus admittance components `G`, `B`. +2. **polar** — the GenOpt/ExaModels (`examples/opf.jl`) form: polar voltages, + branch flows with `sin`/`cos` of angle differences, incidence + gather/scatter for the nodal balance. + +Since AC-OPF needs second-order info for interior-point methods but ArrayDiff +is first-order (by design, for now), the problems are solved with a +first-order augmented Lagrangian (`solver.jl`): equality residual groups +compiled to `eval_residual!`/`eval_residual_jtprod!`, inequalities converted +to equalities with box slacks, variable bounds by projection, inner loop = +projected Adam. Don't expect Ipopt-grade precision — the point is the +vectorized evaluation pipeline, not the outer optimizer. + +## Structured constants (new ArrayDiff feature exercised here) + +Constant arrays that are *not* dense `Array`s are now kept **by reference** +on the tape (`Expression.arrays`, `NODE_ARRAY_VALUE` leaf) instead of being +serialized; matmul nodes call `LinearAlgebra.mul!` directly on them. +`structured.jl` provides two purpose-built types whose products are pure +broadcasts (GPU-safe, no atomics, no CUSPARSE dependency): + +* `GatherMatrix` — one 1 per row: `A*x` is a gather, `A'*w` a padded + fixed-width gather-accumulate (max-degree many fused broadcasts; no + atomics, no scan). Encodes branch↔bus / gen↔bus incidence. +* `ELLMatrix` — padded fixed-width sparse rows (ELLPACK), transpose stored + explicitly. Encodes the admittance components `G`, `B`; better suited to + GPUs than `SparseMatrixCSC` for the short uniform rows of power networks. + +`SparseMatrixCSC` also works (it just goes through its own `mul!`). + +## Files + +| file | purpose | +|---|---| +| `structured.jl` | `GatherMatrix`, `ELLMatrix`, `map_storage` (device transfer) | +| `data.jl` | `case9mod()` (tutorial data, per-unit) and `parse_polar_case` (PowerModels) | +| `models.jl` | `build_rect`, `build_polar` → `ALProblem` | +| `solver.jl` | first-order AL (projected Adam), storage-generic | +| `reference.jl` | Ipopt references (scalar JuMP model / PowerModels) | +| `main.jl` | CPU end-to-end: solve both forms, compare with Ipopt | +| `check_derivatives.jl` | finite-difference checks of all residual groups | +| `gpu_check.jl` | run everything on `JLArray` with scalar indexing disallowed | +| `run_gpu.jl` | CUDA driver (needs a GPU machine; `Pkg.add("CUDA")` first) | + +## Running + +```sh +julia --project=. check_derivatives.jl # derivative correctness +julia --project=. main.jl # CPU solves vs Ipopt +julia --project=. gpu_check.jl # GPU-semantics via JLArrays +julia --project=. run_gpu.jl case9.m # on a machine with an NVIDIA GPU +``` + +The development container had no GPU: `run_gpu.jl` is untested on real +hardware, but `gpu_check.jl` runs the identical code paths under GPUArrays' +scalar-indexing ban, which catches the class of bugs that breaks CUDA runs +(it already caught one: `cumsum!` on vectors silently falls back to a scalar +loop on GPU arrays). + +## Results (CPU, this container) + +Objectives vs Ipopt with the default Adam + SPG-polish schedule: + +| problem | Ipopt | first-order AL | gap | max violation | +|---|---|---|---|---| +| rect case9mod | 3087.84 | 3088.09 | 0.008% | 3.5e-7 | +| polar case9 | 347.70 | 348.39 | 0.20% | 1.5e-7 | +| polar case14 | 8081.52 | 8092.90 | 0.14% | 1.8e-7 | +| polar case30 | 204.97 | 205.89 | 0.45% | 7.5e-7 | + +`bench_cpu.jl` (AL gradient = 8 residual forward+J'v passes + objective +gradient, polar form): + +| case | GatherMatrix | SparseMatrixCSC | +|---|---|---| +| case9 (9 buses) | 64 µs | 52 µs | +| case118 (118 buses) | 281 µs | 110 µs | +| case1354 (1354 buses) | 2.75 ms | 1.54 ms | + +On CPU, `SparseMatrixCSC`'s tight loops win — use `use_gather = false` +there. The structured types exist for the GPU, where kernel count, coalesced +access and the absence of row-pointer indirection/atomics matter; +`run_gpu.jl` times GatherMatrix against CUSPARSE CSR to check that claim on +real hardware. diff --git a/perf/acopf/bench_cpu.jl b/perf/acopf/bench_cpu.jl new file mode 100644 index 0000000..fc36dad --- /dev/null +++ b/perf/acopf/bench_cpu.jl @@ -0,0 +1,53 @@ +# CPU timing of one AL-gradient evaluation (all residual groups: forward + +# J'v, plus the objective gradient) for the polar form, comparing the +# structured GatherMatrix constants against SparseMatrixCSC. +# +# julia --project=. bench_cpu.jl [case118.m] + +include("ACOPF.jl") + +import Downloads +import Printf + +# Fetch a pglib-opf case into /tmp (cached across runs). +function pglib_case(name::AbstractString) + path = joinpath(tempdir(), name) + if !isfile(path) + Downloads.download( + "https://raw.githubusercontent.com/power-grid-lib/pglib-opf/dc6be4b2f85ca0e776952ec22cbd4c22396ea5a3/$name", + path, + ) + end + return path +end + +function bench_case(case) + path = startswith(case, "pglib") ? pglib_case(case) : matpower_case(case) + d = parse_polar_case(path) + println("$case: $(d.nbus) buses, $(d.nbranch) branches, $(d.ngen) gens") + for (label, kw) in [ + ("GatherMatrix", (use_gather = true,)), + ("SparseMatrixCSC", (use_gather = false,)), + ] + prob = build_polar(d; kw...) + st = ALState(prob) + al_gradient!(st, prob, 10.0) # compile + n = 1_000 + t = @elapsed for _ in 1:n + al_gradient!(st, prob, 10.0) + end + Printf.@printf(" %-16s %8.1f µs / AL gradient\n", label, 1e6 * t / n) + end + return +end + +for case in ( + isempty(ARGS) ? + [ + "case9.m", + "pglib_opf_case118_ieee.m", + "pglib_opf_case1354_pegase.m", + ] : ARGS +) + bench_case(case) +end diff --git a/perf/acopf/check_derivatives.jl b/perf/acopf/check_derivatives.jl new file mode 100644 index 0000000..9257b95 --- /dev/null +++ b/perf/acopf/check_derivatives.jl @@ -0,0 +1,86 @@ +# Finite-difference validation of every residual group and the objective +# gradient for both AC-OPF forms, at a random interior point. +# +# julia --project=. check_derivatives.jl + +include("ACOPF.jl") + +import Random +import Test + +function _rand_point(prob) + lb = Vector(prob.lb) + ub = Vector(prob.ub) + x = similar(lb) + for i in eachindex(x) + lo = isfinite(lb[i]) ? lb[i] : -0.5 + hi = isfinite(ub[i]) ? ub[i] : 0.5 + t = 0.3 + 0.4 * rand() + x[i] = lo + t * (hi - lo) + end + return x +end + +function check_group(grp, x; h = 1e-6, atol = 1e-5) + n = length(x) + m = grp.dim + F = zeros(m) + ArrayDiff.eval_residual!(grp.evaluator, F, x) + v = randn(m) + Jtv = zeros(n) + ArrayDiff.eval_residual_jtprod!(grp.evaluator, Jtv, x, v) + Fp, Fm = zeros(m), zeros(m) + Jtv_fd = map(1:n) do i + xp = copy(x) + xp[i] += h + xm = copy(x) + xm[i] -= h + ArrayDiff.eval_residual!(grp.evaluator, Fp, xp) + ArrayDiff.eval_residual!(grp.evaluator, Fm, xm) + return LinearAlgebra.dot(v, (Fp .- Fm) ./ (2h)) + end + err = maximum(abs, Jtv .- Jtv_fd) + Test.@test err < atol + return err +end + +function check_objective(prob, x; h = 1e-6, atol = 1e-5) + g = zero(x) + MOI.eval_objective_gradient(prob.objective, g, x) + g_fd = map(eachindex(x)) do i + xp = copy(x) + xp[i] += h + xm = copy(x) + xm[i] -= h + return ( + MOI.eval_objective(prob.objective, xp) - + MOI.eval_objective(prob.objective, xm) + ) / (2h) + end + err = maximum(abs, g .- g_fd) + Test.@test err < atol + return err +end + +function check_problem(name, prob) + println("── $name") + x = _rand_point(prob) + err = check_objective(prob, x) + println(" objective gradient: max err $err") + for grp in prob.groups + err = check_group(grp, x) + println(" $(grp.name): max J'v err $err") + end + return +end + +Random.seed!(1) +Test.@testset "AC-OPF derivative checks" begin + d1 = case9mod() + check_problem("rect / ELLMatrix", build_rect(d1; matrix = ELLMatrix)) + check_problem("rect / SparseMatrixCSC", build_rect(d1; matrix = identity)) + check_problem("rect / dense tape", build_rect(d1; matrix = Matrix)) + d2 = parse_polar_case(matpower_case("case9.m")) + check_problem("polar / GatherMatrix", build_polar(d2; use_gather = true)) + check_problem("polar / SparseMatrixCSC", build_polar(d2; use_gather = false)) +end diff --git a/perf/acopf/data.jl b/perf/acopf/data.jl new file mode 100644 index 0000000..ef2919e --- /dev/null +++ b/perf/acopf/data.jl @@ -0,0 +1,248 @@ +# Data preparation for the two vectorized AC-OPF forms. + +import LinearAlgebra +import PowerModels +import SparseArrays + +# ── Form 1: the JuMP tutorial's case9mod, rectangular voltages + Ybus ──────── +# +# Data copied from JuMP/docs/src/tutorials/applications/optimal_power_flow.jl +# but converted to per-unit so every decision variable is O(1) (essential for +# a first-order method). The objective is kept in dollars by rescaling the +# cost coefficients, so the reference value 3087.84 from the tutorial still +# applies. + +struct RectData + N::Int + G::SparseArrays.SparseMatrixCSC{Float64,Int} # real(Y) in per-unit + B::SparseArrays.SparseMatrixCSC{Float64,Int} # imag(Y) in per-unit + Pg_lb::Vector{Float64} # per-unit bounds on real generation + Pg_ub::Vector{Float64} + Qg_lb::Vector{Float64} + Qg_ub::Vector{Float64} + Pd::Vector{Float64} # per-unit demands + Qd::Vector{Float64} + vmin::Float64 + vmax::Float64 + c2::Vector{Float64} # $ / (p.u.)² etc. so the objective is in dollars + c1::Vector{Float64} + c0::Float64 +end + +function case9mod() + N = 9 + base_MVA = 100.0 + sv(I, V) = Vector(SparseArrays.sparsevec(I, Float64.(V), N)) + Pg_lb = sv([1, 2, 3], [10, 10, 10]) ./ base_MVA + Pg_ub = sv([1, 2, 3], [250, 300, 270]) ./ base_MVA + Qg_lb = sv([1, 2, 3], [-5, -5, -5]) ./ base_MVA + Qg_ub = sv([1, 2, 3], [300, 300, 300]) ./ base_MVA + Pd = sv([5, 7, 9], [54, 60, 75]) ./ base_MVA + Qd = sv([5, 7, 9], [18, 21, 30]) ./ base_MVA + branch = [ + (1, 4, 0.0, 0.0576, 0.0), + (4, 5, 0.017, 0.092, 0.158), + (6, 5, 0.039, 0.17, 0.358), + (3, 6, 0.0, 0.0586, 0.0), + (6, 7, 0.0119, 0.1008, 0.209), + (8, 7, 0.0085, 0.072, 0.149), + (2, 8, 0.0, 0.0625, 0.0), + (8, 9, 0.032, 0.161, 0.306), + (4, 9, 0.01, 0.085, 0.176), + ] + M = length(branch) + F = [b[1] for b in branch] + T = [b[2] for b in branch] + # Everything in per-unit here (the tutorial divides z by base_MVA to work + # in MW/MVar instead). + z = [b[3] + im * b[4] for b in branch] + A = + SparseArrays.sparse(F, 1:M, 1.0, N, M) + + SparseArrays.sparse(T, 1:M, -1.0, N, M) + Y_0 = A * SparseArrays.spdiagm(1 ./ z) * A' + y_sh = [im * b[5] / 2 for b in branch] + Y_sh = SparseArrays.spdiagm( + LinearAlgebra.diag(A * SparseArrays.spdiagm(y_sh) * A'), + ) + Y = Y_0 + Y_sh + # Tutorial costs are in MW: 0.11 P² + 5 P + 150 etc. With P = base * p: + c2 = [0.11, 0.085, 0.1225] .* base_MVA^2 + c1 = [5.0, 1.2, 1.0] .* base_MVA + c0 = 150.0 + 600.0 + 335.0 + # Pad cost vectors to bus length (non-generator buses have Pg = 0 anyway + # because their bounds are [0, 0]). + c2v = zeros(N) + c1v = zeros(N) + c2v[1:3] .= c2 + c1v[1:3] .= c1 + return RectData( + N, + SparseArrays.sparse(real.(Y)), + SparseArrays.sparse(imag.(Y)), + Pg_lb, + Pg_ub, + Qg_lb, + Qg_ub, + Pd, + Qd, + 0.9, + 1.1, + c2v, + c1v, + c0, + ) +end + +# ── Form 2: PowerModels-derived branch data (ExaModels' parametrization) ──── +# +# Same `c1..c8` constants as ExaModels' OPF example / GenOpt's examples/opf.jl, +# but assembled into flat vectors indexed by branch / bus / gen, ready for the +# vectorized model. + +struct PolarData + nbus::Int + ngen::Int + nbranch::Int + f_bus::Vector{Int} + t_bus::Vector{Int} + gen_bus::Vector{Int} + ref_buses::Vector{Int} + c1::Vector{Float64} + c2::Vector{Float64} + c3::Vector{Float64} + c4::Vector{Float64} + c5::Vector{Float64} + c6::Vector{Float64} + c7::Vector{Float64} + c8::Vector{Float64} + rate_a::Vector{Float64} + rate_a_sq::Vector{Float64} + pd::Vector{Float64} + qd::Vector{Float64} + gs::Vector{Float64} + bs::Vector{Float64} + vmin::Vector{Float64} + vmax::Vector{Float64} + pmin::Vector{Float64} + pmax::Vector{Float64} + qmin::Vector{Float64} + qmax::Vector{Float64} + cost1::Vector{Float64} + cost2::Vector{Float64} + cost3::Vector{Float64} +end + +function parse_polar_case(filename::AbstractString) + data = PowerModels.parse_file(filename) + PowerModels.standardize_cost_terms!(data; order = 2) + PowerModels.calc_thermal_limits!(data) + ref = PowerModels.build_ref(data)[:it][:pm][:nw][0] + busdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:bus])) + gendict = Dict(k => i for (i, (k, v)) in enumerate(ref[:gen])) + branchdict = Dict(k => i for (i, (k, v)) in enumerate(ref[:branch])) + nbus = length(ref[:bus]) + ngen = length(ref[:gen]) + nbranch = length(ref[:branch]) + pd = zeros(nbus) + qd = zeros(nbus) + gs = zeros(nbus) + bs = zeros(nbus) + vmin = zeros(nbus) + vmax = zeros(nbus) + for (k, v) in ref[:bus] + i = busdict[k] + loads = [ref[:load][l] for l in ref[:bus_loads][k]] + shunts = [ref[:shunt][s] for s in ref[:bus_shunts][k]] + pd[i] = sum(load["pd"] for load in loads; init = 0.0) + qd[i] = sum(load["qd"] for load in loads; init = 0.0) + gs[i] = sum(shunt["gs"] for shunt in shunts; init = 0.0) + bs[i] = sum(shunt["bs"] for shunt in shunts; init = 0.0) + vmin[i] = v["vmin"] + vmax[i] = v["vmax"] + end + gen_bus = zeros(Int, ngen) + pmin = zeros(ngen) + pmax = zeros(ngen) + qmin = zeros(ngen) + qmax = zeros(ngen) + cost1 = zeros(ngen) + cost2 = zeros(ngen) + cost3 = zeros(ngen) + for (k, v) in ref[:gen] + i = gendict[k] + gen_bus[i] = busdict[v["gen_bus"]] + pmin[i] = v["pmin"] + pmax[i] = v["pmax"] + qmin[i] = v["qmin"] + qmax[i] = v["qmax"] + cost1[i] = v["cost"][1] + cost2[i] = v["cost"][2] + cost3[i] = v["cost"][3] + end + f_bus = zeros(Int, nbranch) + t_bus = zeros(Int, nbranch) + c = [zeros(nbranch) for _ in 1:8] + rate_a = zeros(nbranch) + for (k, branch) in ref[:branch] + i = branchdict[k] + f_bus[i] = busdict[branch["f_bus"]] + t_bus[i] = busdict[branch["t_bus"]] + g, b = PowerModels.calc_branch_y(branch) + tr, ti = PowerModels.calc_branch_t(branch) + ttm = tr^2 + ti^2 + g_fr = branch["g_fr"] + b_fr = branch["b_fr"] + g_to = branch["g_to"] + b_to = branch["b_to"] + c[1][i] = (-g * tr - b * ti) / ttm + c[2][i] = (-b * tr + g * ti) / ttm + c[3][i] = (-g * tr + b * ti) / ttm + c[4][i] = (-b * tr - g * ti) / ttm + c[5][i] = (g + g_fr) / ttm + c[6][i] = (b + b_fr) / ttm + c[7][i] = g + g_to + c[8][i] = b + b_to + rate_a[i] = branch["rate_a"] + end + ref_buses = [busdict[k] for (k, _) in ref[:ref_buses]] + return PolarData( + nbus, + ngen, + nbranch, + f_bus, + t_bus, + gen_bus, + ref_buses, + c[1], + c[2], + c[3], + c[4], + c[5], + c[6], + c[7], + c[8], + rate_a, + rate_a .^ 2, + pd, + qd, + gs, + bs, + vmin, + vmax, + pmin, + pmax, + qmin, + qmax, + cost1, + cost2, + cost3, + ) +end + +matpower_case(name::AbstractString) = joinpath( + dirname(dirname(pathof(PowerModels))), + "test", + "data", + "matpower", + name, +) diff --git a/perf/acopf/gpu_check.jl b/perf/acopf/gpu_check.jl new file mode 100644 index 0000000..891928f --- /dev/null +++ b/perf/acopf/gpu_check.jl @@ -0,0 +1,87 @@ +# GPU-semantics validation without a GPU: run both forms with the tape and +# all solver state on `JLArrays.JLArray` (the GPUArrays.jl reference backend) +# with scalar indexing disallowed. Any operation that would break on CUDA +# (scalar getindex/setindex on device arrays) throws here. +# +# Evaluations (residuals, J'v, objective gradient) are compared point-wise +# against the CPU path at the same input — these are deterministic up to +# reduction order, so tolerances are tight. Full solves are only smoke-tested +# (thousands of Adam steps on a nonconvex problem amplify last-bit +# differences between CPU and GPU-style reductions, so trajectories are not +# bitwise comparable). +# +# julia --project=. gpu_check.jl + +include("ACOPF.jl") + +import GPUArraysCore +import JLArrays +import Random +import Test + +GPUArraysCore.allowscalar(false) + +const JLV = JLArrays.JLArray{Float64,1} + +jl_device(x::AbstractArray) = JLArrays.JLArray(x) + +function compare_evaluations(name, prob_cpu, prob_dev) + println("── $name: evaluation comparison") + Random.seed!(42) + x_cpu = clamp.(Vector(prob_cpu.x0) .+ 0.01 .* randn(length(prob_cpu.x0)), + Vector(prob_cpu.lb), Vector(prob_cpu.ub)) + x_dev = JLV(x_cpu) + g_cpu = zero(x_cpu) + g_dev = JLV(zero(x_cpu)) + MOI.eval_objective_gradient(prob_cpu.objective, g_cpu, x_cpu) + MOI.eval_objective_gradient(prob_dev.objective, g_dev, x_dev) + Test.@test MOI.eval_objective(prob_dev.objective, x_dev) ≈ + MOI.eval_objective(prob_cpu.objective, x_cpu) rtol = 1e-12 + Test.@test Vector(g_dev) ≈ g_cpu rtol = 1e-12 + for (gc, gd) in zip(prob_cpu.groups, prob_dev.groups) + F_cpu = zeros(gc.dim) + F_dev = JLV(zeros(gd.dim)) + ArrayDiff.eval_residual!(gc.evaluator, F_cpu, x_cpu) + ArrayDiff.eval_residual!(gd.evaluator, F_dev, x_dev) + Test.@test Vector(F_dev) ≈ F_cpu rtol = 1e-12 atol = 1e-14 + v = randn(gc.dim) + Jtv_cpu = zero(x_cpu) + Jtv_dev = JLV(zero(x_cpu)) + ArrayDiff.eval_residual_jtprod!(gc.evaluator, Jtv_cpu, x_cpu, v) + ArrayDiff.eval_residual_jtprod!(gd.evaluator, Jtv_dev, x_dev, JLV(v)) + Test.@test Vector(Jtv_dev) ≈ Jtv_cpu rtol = 1e-10 atol = 1e-12 + println(" $(gc.name): F and J'v match") + end + return +end + +function smoke_solve(name, prob_dev) + x, stats = solve!(prob_dev; outer = 5, inner = 500, polish = 2, verbose = false) + println("── $name: smoke solve obj $(stats.obj) viol $(stats.viol)") + Test.@test isfinite(stats.obj) + Test.@test stats.viol < 5e-2 + return +end + +Test.@testset "GPU-semantics (JLArrays) checks" begin + d1 = case9mod() + rect_cpu = build_rect(d1; matrix = ELLMatrix) + rect_dev = build_rect( + d1; + matrix = ELLMatrix, + mode = ArrayDiff.Mode{JLV}(), + device = jl_device, + ) + compare_evaluations("rect / ELLMatrix", rect_cpu, rect_dev) + smoke_solve("rect / ELLMatrix", rect_dev) + d2 = parse_polar_case(matpower_case("case9.m")) + polar_cpu = build_polar(d2; use_gather = true) + polar_dev = build_polar( + d2; + use_gather = true, + mode = ArrayDiff.Mode{JLV}(), + device = jl_device, + ) + compare_evaluations("polar / GatherMatrix", polar_cpu, polar_dev) + smoke_solve("polar / GatherMatrix", polar_dev) +end diff --git a/perf/acopf/main.jl b/perf/acopf/main.jl new file mode 100644 index 0000000..b276b6c --- /dev/null +++ b/perf/acopf/main.jl @@ -0,0 +1,58 @@ +# CPU end-to-end validation of both vectorized AC-OPF forms. +# +# julia --project=. main.jl +# +# Solves form 1 (rectangular, case9mod) and form 2 (polar, case9) with the +# first-order AL solver and compares objectives against Ipopt references. + +include("ACOPF.jl") +include("reference.jl") + +import Printf + +function run_rect(; matrix = ELLMatrix, kwargs...) + d = case9mod() + prob = build_rect(d; matrix) + x, stats = solve!(prob; kwargs...) + return prob, x, stats +end + +function run_polar(case::AbstractString = "case9.m"; use_gather = true, kwargs...) + d = parse_polar_case(matpower_case(case)) + prob = build_polar(d; use_gather) + x, stats = solve!(prob; kwargs...) + return d, prob, x, stats +end + +function run_all() + println("═"^70) + println("Form 1: rectangular voltages + Ybus (JuMP tutorial, case9mod)") + println("═"^70) + ref1 = rect_reference(case9mod()) + Printf.@printf("Ipopt reference objective: %.2f\n", ref1) + _, _, stats1 = run_rect() + Printf.@printf( + "first-order AL objective: %.2f (gap %.3f%%, viol %.2e)\n", + stats1.obj, + 100 * (stats1.obj - ref1) / abs(ref1), + stats1.viol, + ) + println() + println("═"^70) + println("Form 2: polar voltages, sin/cos branch flows (GenOpt, case9)") + println("═"^70) + ref2 = polar_reference(matpower_case("case9.m")) + Printf.@printf("Ipopt/PowerModels reference objective: %.2f\n", ref2) + _, _, _, stats2 = run_polar("case9.m") + Printf.@printf( + "first-order AL objective: %.2f (gap %.3f%%, viol %.2e)\n", + stats2.obj, + 100 * (stats2.obj - ref2) / abs(ref2), + stats2.viol, + ) + return (ref1, stats1, ref2, stats2) +end + +if abspath(PROGRAM_FILE) == @__FILE__ + run_all() +end diff --git a/perf/acopf/models.jl b/perf/acopf/models.jl new file mode 100644 index 0000000..6c0aa0a --- /dev/null +++ b/perf/acopf/models.jl @@ -0,0 +1,212 @@ +# Vectorized AC-OPF model builders. Both build a JuMP model whose variables +# are contiguous `ArrayDiff.ArrayOfVariables` blocks, express every constraint +# group as one whole-vector residual (matrix-vector products + broadcasts), +# and compile each group into an ArrayDiff residual evaluator. +# +# Inequalities are turned into equalities with box-constrained slacks so the +# first-order AL solver only sees equality residuals + variable bounds: +# vmin² ≤ |V|² ≤ vmax² → Vr² + Vi² - w = 0, w ∈ [vmin², vmax²] +# p² + q² ≤ rate² → p² + q² + s - rate² = 0, s ∈ [0, rate²] + +import ArrayDiff +import JuMP +import LinearAlgebra +import MathOptInterface as MOI +import SparseArrays + +storage_type(::ArrayDiff.Mode{S}) where {S} = S + +to_storage(mode, v::Vector{Float64}) = storage_type(mode)(v) + +# ── Form 1: rectangular voltages + bus admittance matrix (JuMP tutorial) ──── +# +# Variables x = [Vr; Vi; Pg; Qg; w], all length N. +# Residuals: +# rP = Pg - Pd - (Vr .* Ir + Vi .* Ii) = 0 with I = Y V: +# rQ = Qg - Qd - (Vi .* Ir - Vr .* Ii) = 0 Ir = G Vr - B Vi +# rW = Vr.^2 + Vi.^2 - w = 0 Ii = G Vi + B Vr +# Reference-bus convention of the tutorial: Vi[1] = 0 (bounds), Vr[1] ≥ 0. + +function build_rect( + d::RectData; + mode = ArrayDiff.Mode{Vector{Float64}}(), + # Transform applied to the constant admittance components, e.g. + # `ELLMatrix`, `identity` (SparseMatrixCSC), or `Matrix` (dense tape). + matrix = ELLMatrix, + device = identity, # storage transform for structured matrices +) + N = d.N + Gm = map_storage(device, matrix(d.G)) + Bm = map_storage(device, matrix(d.B)) + model = JuMP.Model() + JuMP.@variable(model, Vr[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Vi[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Pg[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Qg[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, w[1:N], container = ArrayDiff.ArrayOfVariables) + Ir = Gm * Vr .- Bm * Vi + Ii = Gm * Vi .+ Bm * Vr + rP = Pg .- d.Pd .- (Vr .* Ir .+ Vi .* Ii) + rQ = Qg .- d.Qd .- (Vi .* Ir .- Vr .* Ii) + rW = Vr .^ 2 .+ Vi .^ 2 .- w + obj = sum(d.c2 .* Pg .^ 2 .+ d.c1 .* Pg) + groups = ResidualGroup[ + ResidualGroup("P-balance", build_residual_evaluator(model, rP, mode), N), + ResidualGroup("Q-balance", build_residual_evaluator(model, rQ, mode), N), + ResidualGroup("Vmag", build_residual_evaluator(model, rW, mode), N), + ] + obj_ev = build_objective_evaluator(model, obj, mode) + inf = fill(Inf, N) + lb_Vr = fill(-d.vmax, N) + ub_Vr = fill(d.vmax, N) + lb_Vi = fill(-d.vmax, N) + ub_Vi = fill(d.vmax, N) + lb_Vr[1] = 0.0 # tutorial: real(V[1]) ≥ 0 + lb_Vi[1] = ub_Vi[1] = 0.0 # tutorial: imag(V[1]) == 0 + lb = vcat(lb_Vr, lb_Vi, d.Pg_lb, d.Qg_lb, fill(d.vmin^2, N)) + ub = vcat(ub_Vr, ub_Vi, d.Pg_ub, d.Qg_ub, fill(d.vmax^2, N)) + x0 = vcat( + ones(N), + zeros(N), + (d.Pg_lb .+ d.Pg_ub) ./ 2, + (d.Qg_lb .+ d.Qg_ub) ./ 2, + ones(N), + ) + return ALProblem( + obj_ev, + 1.0 / max(maximum(d.c2), 1.0), + d.c0, + groups, + to_storage(mode, lb), + to_storage(mode, ub), + to_storage(mode, x0), + ) +end + +# ── Form 2: polar voltages, branch flows with sin/cos (GenOpt/ExaModels) ──── +# +# Variables x = [va; vm; pg; qg; pf; pt; qf; qt; sf; st]. +# With gather matrices F (branch → from-bus) and T (branch → to-bus), +# Δ = F va - T va, vv = (F vm) .* (T vm), the residual groups are the +# vectorized versions of GenOpt/examples/opf.jl's constraints; the power +# balance uses the transposed gathers (segmented sums) F', T' over branch +# flows and C' over generator injections. + +function build_polar( + d::PolarData; + mode = ArrayDiff.Mode{Vector{Float64}}(), + device = identity, + use_gather::Bool = true, +) + nb, ng, nl = d.nbus, d.ngen, d.nbranch + make_gather = if use_gather + idx -> map_storage(device, GatherMatrix(idx, nb)) + else + # Sparse baseline (kept by reference on the tape as well). `device` + # may convert the `SparseMatrixCSC`, e.g. to a `CuSparseMatrixCSR`. + idx -> device( + SparseArrays.sparse(1:length(idx), idx, 1.0, length(idx), nb), + ) + end + Fg = make_gather(d.f_bus) + Tg = make_gather(d.t_bus) + Cg = make_gather(d.gen_bus) + Ft = LinearAlgebra.transpose(Fg) + Tt = LinearAlgebra.transpose(Tg) + Ct = LinearAlgebra.transpose(Cg) + model = JuMP.Model() + JuMP.@variable(model, va[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, vm[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pt[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qt[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, sf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, st[1:nl], container = ArrayDiff.ArrayOfVariables) + vaf = Fg * va + vat = Tg * va + vmf = Fg * vm + vmt = Tg * vm + Δ = vaf .- vat + cΔ = cos.(Δ) + sΔ = sin.(Δ) + vv = vmf .* vmt + # From-side flow definitions (GenOpt's first two branch constraints). + g1 = pf .- (d.c5 .* vmf .^ 2 .+ d.c3 .* (vv .* cΔ) .+ d.c4 .* (vv .* sΔ)) + g2 = + qf .+ d.c6 .* vmf .^ 2 .+ d.c4 .* (vv .* cΔ) .- d.c3 .* (vv .* sΔ) + # To-side flows; cos(-Δ) = cos(Δ), sin(-Δ) = -sin(Δ) folded in. + g3 = pt .- (d.c7 .* vmt .^ 2 .+ d.c1 .* (vv .* cΔ) .- d.c2 .* (vv .* sΔ)) + g4 = + qt .+ d.c8 .* vmt .^ 2 .+ d.c2 .* (vv .* cΔ) .+ d.c1 .* (vv .* sΔ) + # Nodal power balance (scatter branch flows and generation to buses). + g5 = Ft * pf .+ Tt * pt .- Ct * pg .+ d.pd .+ d.gs .* vm .^ 2 + g6 = Ft * qf .+ Tt * qt .- Ct * qg .+ d.qd .- d.bs .* vm .^ 2 + # Thermal limits |S|² ≤ rate² with slacks in [0, rate²]. + g7 = pf .^ 2 .+ qf .^ 2 .+ sf .- d.rate_a_sq + g8 = pt .^ 2 .+ qt .^ 2 .+ st .- d.rate_a_sq + obj = sum(d.cost1 .* pg .^ 2 .+ d.cost2 .* pg) + groups = ResidualGroup[ + ResidualGroup("pf-def", build_residual_evaluator(model, g1, mode), nl), + ResidualGroup("qf-def", build_residual_evaluator(model, g2, mode), nl), + ResidualGroup("pt-def", build_residual_evaluator(model, g3, mode), nl), + ResidualGroup("qt-def", build_residual_evaluator(model, g4, mode), nl), + ResidualGroup("P-bal", build_residual_evaluator(model, g5, mode), nb), + ResidualGroup("Q-bal", build_residual_evaluator(model, g6, mode), nb), + ResidualGroup("sf-lim", build_residual_evaluator(model, g7, mode), nl), + ResidualGroup("st-lim", build_residual_evaluator(model, g8, mode), nl), + ] + obj_ev = build_objective_evaluator(model, obj, mode) + lb_va = fill(-Inf, nb) + ub_va = fill(Inf, nb) + for r in d.ref_buses + lb_va[r] = ub_va[r] = 0.0 + end + lb = vcat( + lb_va, + d.vmin, + d.pmin, + d.qmin, + -d.rate_a, + -d.rate_a, + -d.rate_a, + -d.rate_a, + zeros(nl), + zeros(nl), + ) + ub = vcat( + ub_va, + d.vmax, + d.pmax, + d.qmax, + d.rate_a, + d.rate_a, + d.rate_a, + d.rate_a, + d.rate_a_sq, + d.rate_a_sq, + ) + x0 = vcat( + zeros(nb), + ones(nb), + (d.pmin .+ d.pmax) ./ 2, + (d.qmin .+ d.qmax) ./ 2, + zeros(nl), + zeros(nl), + zeros(nl), + zeros(nl), + copy(d.rate_a_sq), + copy(d.rate_a_sq), + ) + return ALProblem( + obj_ev, + 1.0 / max(maximum(d.cost1; init = 1.0), maximum(d.cost2; init = 1.0), 1.0), + sum(d.cost3), + groups, + to_storage(mode, lb), + to_storage(mode, ub), + to_storage(mode, x0), + ) +end diff --git a/perf/acopf/reference.jl b/perf/acopf/reference.jl new file mode 100644 index 0000000..ea2994a --- /dev/null +++ b/perf/acopf/reference.jl @@ -0,0 +1,71 @@ +# Second-order (Ipopt) reference solutions used to validate the first-order +# GPU-style solves. + +import Ipopt +import JuMP +import PowerModels + +""" + rect_reference(d::RectData) -> objective + +Solve the scalar version of form 1 (the JuMP tutorial's model, in per-unit) +with Ipopt. For `case9mod()` this reproduces the tutorial's 3087.84. +""" +function rect_reference(d::RectData) + N = d.N + model = JuMP.Model(Ipopt.Optimizer) + JuMP.set_silent(model) + JuMP.@variable(model, -d.vmax <= Vr[1:N] <= d.vmax, start = 1.0) + JuMP.@variable(model, -d.vmax <= Vi[1:N] <= d.vmax, start = 0.0) + JuMP.@variable(model, d.Pg_lb[i] <= Pg[i in 1:N] <= d.Pg_ub[i]) + JuMP.@variable(model, d.Qg_lb[i] <= Qg[i in 1:N] <= d.Qg_ub[i]) + JuMP.set_lower_bound(Vr[1], 0.0) + JuMP.set_lower_bound(Vi[1], 0.0) + JuMP.set_upper_bound(Vi[1], 0.0) + G, B = Matrix(d.G), Matrix(d.B) + JuMP.@expression( + model, + Ir[i = 1:N], + sum(G[i, j] * Vr[j] - B[i, j] * Vi[j] for j in 1:N) + ) + JuMP.@expression( + model, + Ii[i = 1:N], + sum(G[i, j] * Vi[j] + B[i, j] * Vr[j] for j in 1:N) + ) + JuMP.@constraint( + model, + [i = 1:N], + Pg[i] - d.Pd[i] == Vr[i] * Ir[i] + Vi[i] * Ii[i] + ) + JuMP.@constraint( + model, + [i = 1:N], + Qg[i] - d.Qd[i] == Vi[i] * Ir[i] - Vr[i] * Ii[i] + ) + JuMP.@constraint(model, [i = 1:N], Vr[i]^2 + Vi[i]^2 >= d.vmin^2) + JuMP.@constraint(model, [i = 1:N], Vr[i]^2 + Vi[i]^2 <= d.vmax^2) + JuMP.@objective( + model, + Min, + sum(d.c2[i] * Pg[i]^2 + d.c1[i] * Pg[i] for i in 1:N) + d.c0 + ) + JuMP.optimize!(model) + return JuMP.objective_value(model) +end + +""" + polar_reference(file) -> objective + +PowerModels' AC-OPF (polar) solved with Ipopt. Note: PowerModels includes +angle-difference bounds that the GenOpt/ExaModels form omits; they are +inactive at the optimum for the standard test cases used here. +""" +function polar_reference(file::AbstractString) + PowerModels.silence() + result = PowerModels.solve_ac_opf( + file, + JuMP.optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0), + ) + return result["objective"] +end diff --git a/perf/acopf/run_gpu.jl b/perf/acopf/run_gpu.jl new file mode 100644 index 0000000..2ab0718 --- /dev/null +++ b/perf/acopf/run_gpu.jl @@ -0,0 +1,104 @@ +# Solve both AC-OPF forms on an NVIDIA GPU (tape, solver state, and the +# structured constant matrices all live on the device). +# +# julia --project=. run_gpu.jl [case] (default: case9.m for polar) +# +# Requires a CUDA-capable machine: `import Pkg; Pkg.add("CUDA")` in this +# project first. The container this code was developed in has no GPU; this +# script was validated indirectly through gpu_check.jl (JLArrays enforce the +# same device-semantics restrictions as CUDA arrays). + +include("ACOPF.jl") +include("reference.jl") + +import CUDA +import Printf + +@assert CUDA.functional() "CUDA is not functional on this machine" + +CUDA.allowscalar(false) + +const CUV = CUDA.CuVector{Float64} + +cu_device(x::AbstractArray{Float64}) = CUDA.CuArray(x) +cu_device(x::AbstractArray{Int}) = CUDA.CuArray(x) +cu_device(x) = x + +function bench(f, name) + f() # warm up / compile + t = @elapsed f() + Printf.@printf("%-28s %10.3f s\n", name, t) + return +end + +function main(case = "case9.m") + println("═"^70) + println("Form 1 on GPU: rectangular + Ybus (ELLMatrix), case9mod") + println("═"^70) + d1 = case9mod() + prob1 = build_rect( + d1; + matrix = ELLMatrix, + mode = ArrayDiff.Mode{CUV}(), + device = cu_device, + ) + x1, stats1 = solve!(prob1) + ref1 = rect_reference(d1) + Printf.@printf( + "GPU objective %.2f Ipopt %.2f viol %.2e\n", + stats1.obj, + ref1, + stats1.viol, + ) + println() + println("═"^70) + println("Form 2 on GPU: polar sin/cos (GatherMatrix), $case") + println("═"^70) + d2 = parse_polar_case(matpower_case(case)) + prob2 = build_polar( + d2; + use_gather = true, + mode = ArrayDiff.Mode{CUV}(), + device = cu_device, + ) + x2, stats2 = solve!(prob2) + ref2 = polar_reference(matpower_case(case)) + Printf.@printf( + "GPU objective %.2f Ipopt %.2f viol %.2e\n", + stats2.obj, + ref2, + stats2.viol, + ) + # AL-gradient timing on the polar form: CPU (CSC) vs GPU with the + # structured GatherMatrix vs GPU with CUSPARSE CSR. + println() + println("AL gradient timing (1000 evaluations), polar form:") + prob_cpu = build_polar(d2; use_gather = false) + cusparse(x::SparseArrays.SparseMatrixCSC) = + CUDA.CUSPARSE.CuSparseMatrixCSR(x) + cusparse(x::AbstractArray) = CUDA.CuArray(x) + prob_csr = build_polar( + d2; + use_gather = false, + mode = ArrayDiff.Mode{CUV}(), + device = cusparse, + ) + st_cpu = ALState(prob_cpu) + st_gpu = ALState(prob2) + st_csr = ALState(prob_csr) + bench( + () -> foreach(_ -> al_gradient!(st_cpu, prob_cpu, 10.0), 1:1000), + "CPU (CSC)", + ) + bench( + () -> CUDA.@sync(foreach(_ -> al_gradient!(st_gpu, prob2, 10.0), 1:1000)), + "GPU (GatherMatrix)", + ) + bench( + () -> CUDA.@sync(foreach(_ -> al_gradient!(st_csr, prob_csr, 10.0), 1:1000)), + "GPU (CUSPARSE CSR)", + ) + return +end + +main(length(ARGS) >= 1 ? ARGS[1] : "case9.m") diff --git a/perf/acopf/solver.jl b/perf/acopf/solver.jl new file mode 100644 index 0000000..e395064 --- /dev/null +++ b/perf/acopf/solver.jl @@ -0,0 +1,350 @@ +# A first-order augmented-Lagrangian solver operating purely through +# ArrayDiff's residual API. Everything in the hot loop is either a +# whole-vector broadcast or an ArrayDiff evaluator call, so the same code +# runs with `Vector{Float64}` (CPU), `JLArray` (GPU-semantics check) and +# `CuVector{Float64}` (GPU) storage. +# +# min f(x) s.t. F_g(x) = 0 (g = 1..#groups), lb ≤ x ≤ ub +# +# Inequalities are pre-converted to equalities with box-constrained slacks by +# the model builders, so only equality residual groups appear here. Bounds +# are enforced by projection (clamp) inside a projected-Adam inner loop. + +import ArrayDiff +import JuMP +import LinearAlgebra +import MathOptInterface as MOI +import Printf + +struct ResidualGroup{E} + name::String + evaluator::E + dim::Int +end + +struct ALProblem{S<:AbstractVector{Float64},O} + objective::O # ArrayDiff.Evaluator or nothing + # The AL minimizes `objective_scale * f(x) + multiplier/penalty terms`; + # scale down large cost coefficients so the objective gradient is O(1) + # like the constraint gradients. Reported objective values are unscaled. + objective_scale::Float64 + # Constant term left out of the evaluator (keeps the tape free of scalar + # +'s, which are not GPU-safe); added back when reporting. + objective_offset::Float64 + groups::Vector{ResidualGroup} + lb::S + ub::S + x0::S +end + +n_variables(p::ALProblem) = length(p.x0) + +""" + build_residual_evaluator(jump_model, expr, mode) + +Compile the vector-valued array expression `expr` (built from the variables +of `jump_model`) into an `ArrayDiff.Evaluator` supporting `eval_residual!` +and `eval_residual_jtprod!` with tape storage given by `mode`. +""" +function build_residual_evaluator(jump_model, expr, mode) + ad = ArrayDiff.model(mode) + ArrayDiff.set_residual!(ad, JuMP.moi_function(expr)) + ev = ArrayDiff.Evaluator( + ad, + mode, + JuMP.index.(JuMP.all_variables(jump_model)), + ) + MOI.initialize(ev, Symbol[:Grad, :Jac, :JacVec]) + return ev +end + +function build_objective_evaluator(jump_model, expr, mode) + ad = ArrayDiff.model(mode) + MOI.Nonlinear.set_objective(ad, JuMP.moi_function(expr)) + ev = ArrayDiff.Evaluator( + ad, + mode, + JuMP.index.(JuMP.all_variables(jump_model)), + ) + MOI.initialize(ev, Symbol[:Grad]) + return ev +end + +# Buffers for one residual group, allocated with the problem's storage type. +struct GroupState{S} + F::S # residual values + λ::S # multipliers + v::S # λ + ρ F seed for J'v + Jtv::S # J' (λ + ρ F) +end + +mutable struct ALState{S} + x::S + g::S # gradient of the AL + m::S # Adam first moment / SPG previous x + v::S # Adam second moment / SPG previous g + step::S # work buffer + trial::S # SPG line-search trial point + groups::Vector{GroupState{S}} +end + +function ALState(p::ALProblem{S}) where {S} + n = n_variables(p) + zero_n = () -> fill!(similar(p.x0, n), 0.0) + groups = [ + GroupState{S}( + fill!(similar(p.x0, g.dim), 0.0), + fill!(similar(p.x0, g.dim), 0.0), + fill!(similar(p.x0, g.dim), 0.0), + zero_n(), + ) for g in p.groups + ] + x = copy(p.x0) + x .= clamp.(x, p.lb, p.ub) + return ALState{S}( + x, + zero_n(), + zero_n(), + zero_n(), + zero_n(), + zero_n(), + groups, + ) +end + +objective_value(p::ALProblem{S,Nothing}, x) where {S} = p.objective_offset + +function objective_value(p::ALProblem, x) + return MOI.eval_objective(p.objective, x) + p.objective_offset +end + +# ∇(AL) = scale⁻¹-free scaled objective gradient + Σ_g J_g' (λ_g + ρ F_g). +function al_gradient!(st::ALState, p::ALProblem, ρ::Float64) + if p.objective === nothing + fill!(st.g, 0.0) + else + MOI.eval_objective_gradient(p.objective, st.g, st.x) + st.g .*= p.objective_scale + end + for (grp, gs) in zip(p.groups, st.groups) + ArrayDiff.eval_residual!(grp.evaluator, gs.F, st.x) + gs.v .= gs.λ .+ ρ .* gs.F + ArrayDiff.eval_residual_jtprod!(grp.evaluator, gs.Jtv, st.x, gs.v) + st.g .+= gs.Jtv + end + return +end + +function residuals!(st::ALState, p::ALProblem) + for (grp, gs) in zip(p.groups, st.groups) + ArrayDiff.eval_residual!(grp.evaluator, gs.F, st.x) + end + return +end + +# AL value at `x` (with the objective already scaled): uses each group's F +# buffer as scratch. +function al_value(st::ALState, p::ALProblem, x, ρ::Float64) + val = + p.objective === nothing ? 0.0 : + p.objective_scale * MOI.eval_objective(p.objective, x) + for (grp, gs) in zip(p.groups, st.groups) + ArrayDiff.eval_residual!(grp.evaluator, gs.F, x) + val += LinearAlgebra.dot(gs.λ, gs.F) + (ρ / 2) * LinearAlgebra.dot(gs.F, gs.F) + end + return val +end + +# One inner minimization of the AL with a nonmonotone spectral projected +# gradient (Birgin–Martínez–Raydan SPG). Everything is a broadcast or a +# `dot`, so it is storage-generic like the rest of the solver. +function spg!( + st::ALState, + p::ALProblem, + ρ::Float64; + iters::Int, + tol::Float64, + memory::Int = 10, + α0::Float64 = 1.0, + αmin::Float64 = 1e-12, + αmax::Float64 = 1e12, + γ::Float64 = 1e-4, +) + x_prev, g_prev, d, trial = st.m, st.v, st.step, st.trial + φ = al_value(st, p, st.x, ρ) + al_gradient!(st, p, ρ) + recent = fill(φ, memory) + α = α0 + for it in 1:iters + # d = P(x - α g) - x + d .= clamp.(st.x .- α .* st.g, p.lb, p.ub) .- st.x + dnorm = maximum(abs, d) + if dnorm < tol + break + end + gd = LinearAlgebra.dot(st.g, d) + # Nonmonotone Armijo backtracking on λ ∈ (0, 1]. + φ_ref = maximum(recent) + λ = 1.0 + φ_new = φ + for _ in 1:30 + trial .= st.x .+ λ .* d + φ_new = al_value(st, p, trial, ρ) + if φ_new <= φ_ref + γ * λ * gd + break + end + λ /= 2 + end + x_prev .= st.x + g_prev .= st.g + st.x .= trial + φ = φ_new + recent[1+it%memory] = φ + al_gradient!(st, p, ρ) + # Barzilai–Borwein step for the next iteration. + x_prev .= st.x .- x_prev # s + g_prev .= st.g .- g_prev # y + sy = LinearAlgebra.dot(x_prev, g_prev) + ss = LinearAlgebra.dot(x_prev, x_prev) + α = sy > 0 ? clamp(ss / sy, αmin, αmax) : αmax + end + return +end + +max_violation(st::ALState) = + isempty(st.groups) ? 0.0 : + maximum(gs -> isempty(gs.F) ? 0.0 : maximum(abs, gs.F), st.groups) + +""" + solve!(p::ALProblem; kwargs...) -> (x, stats) + +Run the first-order augmented-Lagrangian loop. Key knobs: + +* `method`: `:adam` (default) or `:spg` for the inner minimization. Adam's + slow trajectories reach better local basins from a cold start on these + nonconvex problems; SPG converges faster but to nearer (often worse) + stationary points, so it is used by default only in the final `polish` + phase, which drives the violation down from the point Adam reached. +* `outer`, `inner`: number of outer AL updates and inner steps. +* `lr`: initial Adam step, decayed by `lr_decay` each outer iteration + (floored at `lr_min`). +* `ρ0`, `ρ_growth`, `ρmax`: penalty schedule. Multipliers are updated only + when the violation meets the current target `η` (LANCELOT-style + safeguard, tightened by `viol_target_ratio` after each accepted update); + otherwise `ρ` grows. +* `polish`, `polish_inner`: number of SPG polish rounds at the end. +* `tol_feas`: stop when the max violation falls below this. +""" +function solve!( + p::ALProblem; + method::Symbol = :adam, + outer::Int = 40, + inner::Int = 6_000, + inner_tol0::Float64 = 1e-3, + inner_tol_decay::Float64 = 0.5, + lr::Float64 = 5e-3, + lr_decay::Float64 = 0.93, + lr_min::Float64 = 1e-5, + ρ0::Float64 = 10.0, + ρ_growth::Float64 = 3.0, + ρmax::Float64 = 1e6, + viol_target_ratio::Float64 = 0.25, + tol_feas::Float64 = 1e-7, + β1::Float64 = 0.9, + β2::Float64 = 0.999, + ϵ::Float64 = 1e-8, + polish::Int = 8, + polish_inner::Int = 2_000, + verbose::Bool = true, +) + st = ALState(p) + ρ = ρ0 + η = 0.1 # violation target for accepting a multiplier update + α = lr + inner_tol = inner_tol0 + history = NamedTuple[] + for out in 1:outer + if method === :spg + spg!(st, p, ρ; iters = inner, tol = inner_tol) + inner_tol = max(inner_tol * inner_tol_decay, 1e-10) + elseif method === :adam + fill!(st.m, 0.0) + fill!(st.v, 0.0) + for it in 1:inner + al_gradient!(st, p, ρ) + st.m .= β1 .* st.m .+ (1 - β1) .* st.g + st.v .= β2 .* st.v .+ (1 - β2) .* st.g .* st.g + c1 = 1 - β1^it + c2 = 1 - β2^it + st.step .= (α / c1) .* st.m ./ (sqrt.(st.v ./ c2) .+ ϵ) + st.x .= clamp.(st.x .- st.step, p.lb, p.ub) + end + else + error("unknown method $method") + end + residuals!(st, p) + viol = max_violation(st) + obj = objective_value(p, st.x) + push!(history, (outer = out, viol = viol, obj = obj, ρ = ρ, lr = α)) + if verbose + Printf.@printf( + "outer %3d obj %14.6f viol %10.3e ρ %8.1e lr %8.1e\n", + out, + obj, + viol, + ρ, + α, + ) + end + if viol < tol_feas + break + end + # LANCELOT-style safeguarded update: only trust a first-order + # multiplier update when the violation met the current target η; + # otherwise increase the penalty and keep the multipliers. + if viol <= η + for gs in st.groups + gs.λ .+= ρ .* gs.F + end + η = max(η * viol_target_ratio, tol_feas / 10) + else + ρ = min(ρ * ρ_growth, ρmax) + end + α = max(α * lr_decay, lr_min) + end + # Final polish: from the (near-feasible) point Adam reached, a few SPG + # solves with multiplier updates drive the violation down much faster + # than more Adam steps would. SPG alone from a cold start tends to land + # in worse local basins, so it is only used here at the end. + for pol in 1:polish + residuals!(st, p) + if max_violation(st) < tol_feas + break + end + spg!(st, p, ρ; iters = polish_inner, tol = 1e-10) + residuals!(st, p) + viol = max_violation(st) + obj = objective_value(p, st.x) + push!(history, (outer = -pol, viol = viol, obj = obj, ρ = ρ, lr = 0.0)) + if verbose + Printf.@printf( + "polish %2d obj %14.6f viol %10.3e ρ %8.1e\n", + pol, + obj, + viol, + ρ, + ) + end + for gs in st.groups + gs.λ .+= ρ .* gs.F + end + ρ = min(ρ * ρ_growth, ρmax) + end + residuals!(st, p) + stats = ( + viol = max_violation(st), + obj = objective_value(p, st.x), + history = history, + ) + return st.x, stats +end diff --git a/perf/acopf/structured.jl b/perf/acopf/structured.jl new file mode 100644 index 0000000..34e01e6 --- /dev/null +++ b/perf/acopf/structured.jl @@ -0,0 +1,207 @@ +# Device-generic structured constant matrices for vectorized AC-OPF. +# +# Both types implement only what ArrayDiff's matmul nodes need — +# `LinearAlgebra.mul!(y, A, x)` and `LinearAlgebra.mul!(y, transpose(A), w)` — +# using nothing but broadcasts over (views of) their storage vectors, plus +# `cumsum!`. That makes the exact same code run on `Vector`, `JLArray` (the +# GPUArrays reference backend used to validate GPU semantics without a GPU), +# and `CuArray`, without any CUDA dependency here. +# +# * `GatherMatrix`: a 0/1 matrix with exactly one 1 per row. `A * x` is a pure +# gather (`y .= x[idx]`, one fused kernel) and `A' * w` is a segmented sum +# implemented with one gather, one `cumsum!` and one subtraction — no +# atomics, no CSR machinery. This encodes branch↔bus / gen↔bus incidence. +# * `ELLMatrix`: fixed-width (padded) sparse rows. `A * x` is `K` fused +# gather-multiply-accumulate broadcasts where `K` is the max number of +# nonzeros per row (≈ node degree in a power network, small and uniform). +# The transpose is stored explicitly as another `ELLMatrix`. This encodes +# the bus admittance matrix and beats `SparseMatrixCSC`-style SpMV on GPU +# for such uniformly-short rows because it needs no row-pointer indirection +# and its memory access is coalesced. + +import LinearAlgebra +import SparseArrays + +# ── GatherMatrix ───────────────────────────────────────────────────────────── + +""" + GatherMatrix(idx::Vector{Int}, ncol::Int) + +The `length(idx) × ncol` matrix `A` with `A[i, idx[i]] = 1` and zeros +elsewhere, so that `A * x = x[idx]` (a pure gather) and `A' * w` scatter-adds +`w` into the selected columns. + +The transpose product is computed *without atomics and without a scan*: +column `j` collects the rows `{i : idx[i] == j}`; those row lists are padded +to the maximum count `D` (the maximum node degree, small in power networks) +into an `ncol × D` index matrix `tcols` pointing into a `length(idx) + 1` +buffer whose last slot is a permanent zero (the padding target). `A' * w` is +then `D` fused gather-add broadcasts. +""" +struct GatherMatrix{ + Vi<:AbstractVector{Int}, + Mi<:AbstractMatrix{Int}, + Vf<:AbstractVector{Float64}, +} <: AbstractMatrix{Float64} + idx::Vi + ncol::Int + tcols::Mi # ncol × D indices into `buf` (m + 1 = zero pad) + buf::Vf # length(idx) + 1 workspace with buf[end] == 0 kept invariant +end + +function GatherMatrix(idx::AbstractVector{<:Integer}, ncol::Integer) + m = length(idx) + lists = [Int[] for _ in 1:ncol] + for (i, j) in enumerate(idx) + push!(lists[j], i) + end + D = maximum(length, lists; init = 0) + tcols = fill(m + 1, ncol, D) # pad → permanent zero slot + for j in 1:ncol, (k, i) in enumerate(lists[j]) + tcols[j, k] = i + end + return GatherMatrix(collect(Int, idx), Int(ncol), tcols, zeros(m + 1)) +end + +Base.size(A::GatherMatrix) = (length(A.idx), A.ncol) + +# Only for `Matrix(A)` in reference computations / display; evaluation never +# uses scalar `getindex`. +Base.getindex(A::GatherMatrix, i::Int, j::Int) = Float64(A.idx[i] == j) + +function LinearAlgebra.mul!( + y::AbstractVector, + A::GatherMatrix, + x::AbstractVector, +) + y .= view(x, A.idx) + return y +end + +# cumsum! on vectors silently takes Base's scalar pairwise path (GPUArrays has no generic scan), +# which is why we use padded gathers instead +function LinearAlgebra.mul!( + y::AbstractVector, + At::LinearAlgebra.Transpose{Float64,<:GatherMatrix}, + w::AbstractVector, +) + A = LinearAlgebra.parent(At) + D = size(A.tcols, 2) + if D == 0 + fill!(y, 0.0) + return y + end + view(A.buf, eachindex(A.idx)) .= w # buf[m + 1] stays 0 (padding target) + y .= view(A.buf, view(A.tcols, :, 1)) + for k in 2:D + y .+= view(A.buf, view(A.tcols, :, k)) + end + return y +end + +# ── ELLMatrix ──────────────────────────────────────────────────────────────── + +""" + ELLMatrix(S::SparseMatrixCSC) + +ELLPACK-format copy of `S`: row `i`'s nonzeros are `vals[i, k]` at columns +`cols[i, k]` for `k = 1:K`, padded with zero values pointing at column 1. +The transpose is stored eagerly in `at` so that reverse-mode `A' * w` +products use the same kernel. +""" +struct ELLMatrix{ + Mi<:AbstractMatrix{Int}, + Mf<:AbstractMatrix{Float64}, + TT, +} <: AbstractMatrix{Float64} + m::Int + n::Int + cols::Mi + vals::Mf + at::TT # ELLMatrix of the transpose, or `nothing` +end + +function ELLMatrix(S::SparseArrays.SparseMatrixCSC; with_transpose::Bool = true) + m, n = size(S) + I, J, V = SparseArrays.findnz(S) + counts = zeros(Int, m) + for i in I + counts[i] += 1 + end + K = maximum(counts; init = 0) + cols = ones(Int, m, K) + vals = zeros(m, K) + fill!(counts, 0) + for (i, j, v) in zip(I, J, V) + counts[i] += 1 + cols[i, counts[i]] = j + vals[i, counts[i]] = v + end + at = if with_transpose + ELLMatrix( + SparseArrays.sparse(LinearAlgebra.transpose(S)); + with_transpose = false, + ) + else + nothing + end + return ELLMatrix(m, n, cols, vals, at) +end + +Base.size(A::ELLMatrix) = (A.m, A.n) + +function Base.getindex(A::ELLMatrix, i::Int, j::Int) + acc = 0.0 + for k in 1:size(A.cols, 2) + if A.cols[i, k] == j + acc += A.vals[i, k] + end + end + return acc +end + +function LinearAlgebra.mul!(y::AbstractVector, A::ELLMatrix, x::AbstractVector) + K = size(A.cols, 2) + if K == 0 + fill!(y, 0.0) + return y + end + y .= view(A.vals, :, 1) .* view(x, view(A.cols, :, 1)) + for k in 2:K + y .+= view(A.vals, :, k) .* view(x, view(A.cols, :, k)) + end + return y +end + +function LinearAlgebra.mul!( + y::AbstractVector, + At::LinearAlgebra.Transpose{Float64,<:ELLMatrix}, + w::AbstractVector, +) + A = LinearAlgebra.parent(At) + A.at === nothing && error("ELLMatrix built without transpose") + return LinearAlgebra.mul!(y, A.at, w) +end + +# ── Device transfer ────────────────────────────────────────────────────────── + +""" + map_storage(f, A) + +Rebuild `A` with every storage array passed through `f` (for example +`f = CUDA.CuArray` or `f = JLArrays.JLArray`). `f = identity` is a no-op. +""" +map_storage(f, A::AbstractArray) = f(A) + +function map_storage(f, A::GatherMatrix) + return GatherMatrix(f(A.idx), A.ncol, f(A.tcols), f(A.buf)) +end + +function map_storage(f, A::ELLMatrix) + at = A.at === nothing ? nothing : map_storage(f, A.at) + return ELLMatrix(A.m, A.n, f(A.cols), f(A.vals), at) +end + +function map_storage(f, At::LinearAlgebra.Transpose) + return LinearAlgebra.transpose(map_storage(f, LinearAlgebra.parent(At))) +end diff --git a/perf/bench/Project.toml b/perf/bench/Project.toml new file mode 100644 index 0000000..bfa286a --- /dev/null +++ b/perf/bench/Project.toml @@ -0,0 +1,20 @@ +[deps] +ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" +GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde" +HypertextLiteral = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +PGLib = "07a8691f-3d11-4330-951b-3c50f98338be" +PlutoTeachingTools = "661c6b06-c737-4d37-b85c-46df65de6f69" +PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +ShortCodes = "f62ebe17-55c5-4640-972f-b59c0dd11ccf" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" diff --git a/perf/bench/bench_opf.jl b/perf/bench/bench_opf.jl new file mode 100644 index 0000000..02c5e98 --- /dev/null +++ b/perf/bench/bench_opf.jl @@ -0,0 +1,147 @@ +# Apples-to-apples AC-OPF derivative-timing harness. +# +# Three backends produce the *same* AC-OPF as an `NLPModels.AbstractNLPModel`: +# :exa — ExaModels' `ac_power_model` (docs/src/opf.jl) → ExaModel +# :genopt — GenOpt's model (examples/opf.jl) → ExaModels.ExaModel → ExaModel +# :arraydiff — this repo's vectorized ArrayDiffNLPModel (perf/percival) +# +# Because all three are `AbstractNLPModel`, the timing step is identical for +# each — same call, gradient vs gradient, residual vs residual. +# +# julia --project=perf/bench +# include("perf/bench/bench_opf.jl") +# m = opf_model(:arraydiff, "case14.m") # step 1: build (PowerModels-bundled) +# opf_timings(m) # step 2: time +# +# Case names ending up in PowerModels' bundled data ("case9.m"/"case14.m"/ +# "case30.m") load offline; pglib names ("pglib_opf_case118_ieee.m") download. +# +# Swap storage/backend for GPU (you take over here): pass `backend=CUDABackend()` +# to :exa/:genopt, or `mode=ArrayDiff.Mode{CuVector{Float64}}(), device=CuArray` +# to :arraydiff. + +import ExaModels +import GenOpt +import NLPModels +import Printf +using BenchmarkTools + +# `:arraydiff` builder + its deps (parse_polar_case, build_polar_nlp). +include(joinpath(@__DIR__, "..", "percival", "build_percival.jl")) + +import Downloads + +const PGLIB_URL = + "https://raw.githubusercontent.com/power-grid-lib/pglib-opf/dc6be4b2f85ca0e776952ec22cbd4c22396ea5a3" + +# Resolve a case name to a local `.m` file. Prefer the MATPOWER cases bundled +# with PowerModels (`matpower_case`, from data.jl) — e.g. "case9.m", "case14.m", +# "case30.m" — so no download is needed. Fall back to a cached pglib download +# for names PowerModels doesn't ship (e.g. "pglib_opf_case118_ieee.m"). +function case_path(name::AbstractString) + bundled = matpower_case(name) + isfile(bundled) && return bundled + path = joinpath(tempdir(), name) + isfile(path) || Downloads.download("$(PGLIB_URL)/$(name)", path) + return path +end + +# Substituted into GenOpt's example so it uses the resolved local file instead +# of re-downloading (see the :genopt builder). +_bench_skip_download(args...) = nothing + +# Load only the function definitions from a script `path` (everything before the +# `stop` marker), optionally applying `subs` string replacements first. Lets us +# reuse the real upstream opf.jl files without their download/solve tails. +function _include_prefix(path, stop; subs = ()) + src = read(path, String) + for (a, b) in subs + src = replace(src, a => b) + end + i = findfirst(stop, src) + i === nothing || (src = src[1:prevind(src, first(i))]) + return Base.include_string(Main, src, path) +end + +const EXA_OPF = joinpath(dirname(dirname(pathof(ExaModels))), "docs", "src", "opf.jl") + +# ── Builders ───────────────────────────────────────────────────────────────── + +function opf_model(::Val{:arraydiff}, name; kwargs...) + d = parse_polar_case(case_path(name)) + return build_polar_nlp(d; kwargs...)[1] +end + +function opf_model(::Val{:exa}, name; backend = nothing, T = Float64) + isdefined(Main, :ac_power_model) || + _include_prefix(EXA_OPF, "# We first download") + return Base.invokelatest(Main.ac_power_model, case_path(name); backend = backend, T = T) +end + +include(joinpath(dirname(dirname(pathof(GenOpt))), "examples", "opf", "model.jl")) + +function opf_model(::Val{:genopt}, name; backend = nothing) + return ExaModels.ExaModel(build_model(PGLib.pglib(name)); backend = backend) +end + +opf_model(backend::Symbol, name; kwargs...) = + opf_model(Val(backend), name; kwargs...) + +# ── Timing (uniform NLPModels API) ─────────────────────────────────────────── + +""" + opf_timings(nlp) -> NamedTuple + +Median wall-clock time of one evaluation of each first-order primitive, at the +model's starting point. `residual` here is the constraint vector `cons!`. +""" +function opf_timings(nlp) + n, m = nlp.meta.nvar, nlp.meta.ncon + x = copy(nlp.meta.x0) + g = similar(x, n) + c = similar(x, m) + v = fill!(similar(x, m), one(eltype(x))) + Jtv = similar(x, n) + return ( + nvar = n, + ncon = m, + objective = (@belapsed NLPModels.obj($nlp, $x)), + gradient = (@belapsed NLPModels.grad!($nlp, $x, $g)), + residual = (@belapsed NLPModels.cons!($nlp, $x, $c)), + jtprod = (@belapsed NLPModels.jtprod!($nlp, $x, $v, $Jtv)), + ) +end + +function print_timings(name, t) + Printf.@printf( + "%-10s nvar %6d ncon %6d | obj %.2e grad %.2e resid %.2e jtprod %.2e\n", + name, t.nvar, t.ncon, t.objective, t.gradient, t.residual, t.jtprod, + ) + return +end + +# Convenience: build + time each backend on one case. +function compare(name; backends = (:exa, :genopt, :arraydiff)) + ts = [opf_timings(opf_model(b, name)) for b in backends] + for (t, b) in zip(ts, backends) + print_timings(string(b), t) + end + return +end + +case = "pglib_opf_case14_ieee.m" + +exa = opf_model(:exa, case) +genopt = opf_model(:genopt, case) +arraydiff = opf_model(:arraydiff, case) +print_timings(case, opf_timings(exa)) + +case = "pglib_opf_case10000_goc.m" +exa = opf_model(:exa, case) +t = opf_timings(exa) +print_timings(case, t) +compare(case) + +import CUDA +case = "pglib_opf_case14_ieee.m" +arraydiff = opf_model(:arraydiff, case, mode=ArrayDiff.Mode{CUDA.CuVector{Float64}}(), device=CUDA.CuArray) diff --git a/perf/gpu_bench.jl b/perf/gpu_bench.jl new file mode 100644 index 0000000..89a5ea2 --- /dev/null +++ b/perf/gpu_bench.jl @@ -0,0 +1,4 @@ +include("arraydiff.jl") +T, h, d, n = Float32, 4096, 13, 178 +display(ArrayDiffNeural.neural(T, h, d, n; gpu = true)) +display(ArrayDiffNeural.profile_gpu()) diff --git a/perf/percival/Project.toml b/perf/percival/Project.toml new file mode 100644 index 0000000..b1afe36 --- /dev/null +++ b/perf/percival/Project.toml @@ -0,0 +1,20 @@ +[deps] +ADNLPModels = "54578032-b7ea-4c30-94aa-7cbd1cce6c9a" +ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +JSOSolvers = "10dff2fc-5484-5881-a0e0-c90441020f8a" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/perf/percival/README.md b/perf/percival/README.md new file mode 100644 index 0000000..c167e55 --- /dev/null +++ b/perf/percival/README.md @@ -0,0 +1,116 @@ +# AC-OPF through NLPModels solvers (Percival / MadNLP) with ArrayDiff + +Solve the two vectorized AC-OPF forms from `../acopf` through the **NLPModels** +interface, so any NLPModels solver can consume them, with ArrayDiff providing +the (first-order, vectorized, GPU-ready) derivatives. + +``` +JuMP model (ArrayOfVariables + array expressions) + │ build objective + constraint-group residual evaluators (ArrayDiff) + ▼ +ArrayDiffNLPModel <: NLPModels.AbstractNLPModel (adnlp.jl) + obj/grad ← objective evaluator (reverse mode) + cons/jprod/jtprod ← vectorized residual evaluators (one fused pass each) + jac_coord ← Jacobian materialized from jtprod (for interior-point) + ▼ +NLPModels solver: MadNLP (CompactLBFGS) │ Percival (AL) +``` + +`min f(x) s.t. c(x) = 0, l ≤ x ≤ u`. Inequalities (thermal limits, voltage +magnitude) become equalities with box-constrained slacks, so only equality +constraints + variable bounds remain. + +## Results (CPU, case objectives vs Ipopt) + +| case | solver | gap | feasibility | status | iters | +|---|---|---|---|---|---| +| rect case9mod | MadNLP/LBFGS | 0.0000% | 1.5e-15 | SOLVE_SUCCEEDED | 27 | +| rect case9mod | Percival/LBFGS+Tron | 0.0000% | 1.9e-10 | max_iter | 301 | +| polar case9 | MadNLP/LBFGS | 0.0000% | 1.4e-14 | SOLVE_SUCCEEDED | 31 | +| polar case9 | Percival/LBFGS+Tron | -0.008% | 5.4e-05 | stalled | 30 | +| polar case14 | MadNLP/LBFGS | 0.0000% | 4.9e-15 | SOLVE_SUCCEEDED | 42 | +| polar case14 | Percival/LBFGS+Tron | -2.5% | 6.8e-03 | max_time | 10 | +| polar case30 | MadNLP/LBFGS | 0.0000% | 1.7e-14 | SOLVE_SUCCEEDED | 31 | +| polar case30 | Percival/LBFGS+Tron | -43% | 8.0e-02 | stalled | 24 | + +**MadNLP in quasi-Newton mode (`hessian_approximation = CompactLBFGS`) is the +robust winner**: it converges tightly and reliably on every case using only +gradients + the constraint Jacobian (no exact Hessian). Percival (augmented +Lagrangian) is correct on the small cases but its first-order subproblem +(LBFGS-approximated AL Hessian) fails to certify KKT as the penalty grows and +degrades on larger cases within the iteration/time budget. + +## Files + +| file | purpose | +|---|---| +| `adnlp.jl` | `ArrayDiffNLPModel` — the NLPModels bridge (obj/grad/cons/jprod/jtprod/jac_coord) | +| `build_percival.jl` | build rect / polar AC-OPF as `ArrayDiffNLPModel`s (reuses `../acopf`) | +| `run_madnlp.jl` | MadNLP quasi-Newton driver | +| `run_percival.jl` | Percival driver (LBFGS subproblem modifier) | +| `gpu_percival.jl` | JLArray (GPU-semantics) driver + CPU-vs-device eval checks | +| `compare.jl` | side-by-side MadNLP vs Percival on all cases | + +## Percival GPU-compatibility fixes (in `~/.julia/dev/Percival`) + +Percival could not run on GPU-resident (`CuVector` / `JLArray`) models. Fixes: + +1. **`AugLagModel.store_Jv/store_Jtv`** were hardcoded `Vector{T}`; used in + `grad!` (`g .+= jtprod!(model, x, μc_y, store_Jtv)`) they forced a device + mismatch. Changed to the storage type `V`. +2. **`AugLagModel`'s meta** ran `findall`-based variable-bound analysis, which + scalar-indexes GPU arrays. Disabled (`variable_bounds_analysis = false`); + the subproblem solver only projects with `lvar/uvar`. +3. **Dispatch guards** (`percival(nlp)`, `percival(Val{:equ})`) called + `NLPModels.equality_constrained`, whose fallback scalar-iterates the bound + vectors. Added `_is_equality_constrained` which compares them on the host. +4. **`SPGSubSolver`** (`src/spg_subsolver.jl`): a spectral projected gradient + subproblem solver that is pure broadcasts — a GPU-compatible replacement for + the default `TronSolver`, whose projected-CG inner loop uses scalar `x[i]` + indexing that GPU arrays disallow. Selected with `subsolver = SPGSubSolver`; + being first-order it needs no quasi-Newton wrapper. + +With these, `percival(nlp; subsolver = SPGSubSolver)` runs the **entire AL loop +on `JLArray` storage with scalar indexing disallowed** — validated end-to-end +(no scalar-indexing errors). Convergence of the pure-first-order SPG path on +ill-conditioned AC-OPF is slow (as expected); the fixes make Percival +*GPU-runnable*, and for GPU production MadNLPGPU (below) is the stronger option. + +Tests: `Percival/test/spg_subsolver.jl` (in the Percival test suite). + +## MadNLP on GPU + +MadNLP consumes the same `ArrayDiffNLPModel` and has first-class GPU support via +**MadNLPGPU** (cuDSS/cuSOLVER KKT solves). It was not runnable in this +container (no CUDA; JLArray has no dense factorization backend), but the model +is device-generic, so on real hardware `MadNLPGPU.CUDSSSolver` + a CuArray tape +is the intended GPU path. `jac_coord!` currently materializes a dense Jacobian +(fine for these sizes); large GPU problems want a sparse/matrix-free assembly. + +## Full `JuMP → MOI → NLPModelsJuMP.Optimizer` routing (done) + +The complete loop now works — see `jump_acopf.jl` for the AC-OPF written with +`@constraint(model, expr in MOI.Zeros(n))` and solved with `optimize!` +(case9: 347.67 vs Ipopt 347.70). The pieces: + +1. **ArrayDiff JuMP layer** (`src/JuMP/moi_bridge.jl`): `build_constraint` for + `AbstractJuMPArray in MOI.Zeros/Nonnegatives/Nonpositives` keeps the + expression whole (no scalarization) as an `ArrayNonlinearFunction`; plus + `MOI.Utilities.canonicalize!` so MOI caches accept it. Inside the JuMP + macros, MutableArithmetics' fused ops are normalized (`add_mul`/`sub_mul` → + `+`/`-` broadcasts) and its generic matmul is routed back to the whole-array + product (`_MA.operate(*, ::AbstractMatrix{<:Real}, ::AbstractJuMPArray)`), + which would otherwise silently scalarize `Matrix * ArrayOfVariables` into + `Vector{AffExpr}`. +2. **NLPModelsJuMP ext** (`_try_array_nlp_model` hook + implementation): + collects the vector constraints into a constrained `ArrayDiffNLPModel` + (obj/grad + per-constraint vectorized cons/jprod/jtprod, dense `jac_coord` + for KKT solvers, no Hessian). Row bounds map from the set: `Zeros` → [0,0], + `Nonnegatives` → [0,∞), `Nonpositives` → (−∞,0]. +3. **NLS-path guard** (bug fix): `_try_nls_model` no longer fires on + constrained models — it used to silently drop the constraints when the + objective matched `sum((...)^2)`. + +Test: `ArrayDiff/test/NLPModelsJuMP.jl::test_vector_constraint_solve` (skips +itself until the updated NLPModelsJuMP `bl/arraydiff` branch is pushed, since +the test env pins the GitHub remote). diff --git a/perf/percival/adnlp.jl b/perf/percival/adnlp.jl new file mode 100644 index 0000000..419248d --- /dev/null +++ b/perf/percival/adnlp.jl @@ -0,0 +1,197 @@ +# Constrained NLPModel backed by ArrayDiff evaluators: +# * `obj` — scalar objective f(x), via reverse-mode gradient +# * `cons` — a list of vector residuals c_g(x), each evaluated with the +# vectorized residual API (eval_residual! / jtprod! / jprod!). They are +# presented to NLPModels as one stacked equality constraint c(x) = 0. +# +# Solving min f(x) s.t. c(x) = 0, l ≤ x ≤ u through the NLPModels +# interface lets any NLPModels solver (here Percival) consume it. Storage- +# generic: the vector type `V` follows the ArrayDiff tape (`Vector`, `JLArray`, +# `CuVector`), so the whole solve runs on-device. +# +# Keeping the groups separate (instead of one `vcat`-ed residual) avoids +# needing n-ary `vcat` on the tape and keeps each group a single fused +# vectorized pass — exactly the structure of the hand-written AL solver. +# +# Prototype: moves into the NLPModelsJuMP ArrayDiff extension once validated. + +import ArrayDiff +import MathOptInterface as MOI +import NLPModels + +mutable struct ArrayDiffNLPModel{T, V, Ro, Rc} <: NLPModels.AbstractNLPModel{T, V} + meta::NLPModels.NLPModelMeta{T, V} + counters::NLPModels.Counters + obj::ArrayDiff.Evaluator{T, Ro} + cons::Vector{ArrayDiff.Evaluator{T, Rc}} + offsets::Vector{Int} # group g occupies rows offsets[g]+1 : offsets[g+1] + Jtv_tmp::V # scratch for accumulating J'v across groups +end + +function ArrayDiffNLPModel( + obj::ArrayDiff.Evaluator{T, Ro}, + cons::Vector{ArrayDiff.Evaluator{T, Rc}}, + nvar::Int, + lvar::V, + uvar::V, + x0::V; + name::String = "ArrayDiffNLP", +) where {T, V, Ro, Rc} + dims = [ArrayDiff.residual_dimension(c) for c in cons] + offsets = cumsum(vcat(0, dims)) + ncon = offsets[end] + z = fill!(similar(x0, ncon), zero(T)) + meta = NLPModels.NLPModelMeta{T, V}( + nvar; + x0 = x0, + lvar = lvar, + uvar = uvar, + ncon = ncon, + lcon = z, + ucon = copy(z), + y0 = copy(z), + nnzh = 0, # no exact Hessian; quasi-Newton solvers supply their own + minimize = true, + islp = false, + name = name, + lin = Int[], + # `findall`-based bound analysis scalar-indexes GPU bound vectors; + # the solver only projects with lvar/uvar, so skip it. + variable_bounds_analysis = false, + constraint_bounds_analysis = false, + # No second-order info from ArrayDiff. + hess_available = false, + hprod_available = false, + ) + Jtv_tmp = fill!(similar(x0, nvar), zero(T)) + return ArrayDiffNLPModel{T, V, Ro, Rc}( + meta, + NLPModels.Counters(), + obj, + cons, + offsets, + Jtv_tmp, + ) +end + +_slice(nlp::ArrayDiffNLPModel, g::Int) = (nlp.offsets[g] + 1):nlp.offsets[g + 1] + +function NLPModels.obj(nlp::ArrayDiffNLPModel, x::AbstractVector) + NLPModels.increment!(nlp, :neval_obj) + return MOI.eval_objective(nlp.obj, x) +end + +function NLPModels.grad!(nlp::ArrayDiffNLPModel, x::AbstractVector, g::AbstractVector) + NLPModels.increment!(nlp, :neval_grad) + MOI.eval_objective_gradient(nlp.obj, g, x) + return g +end + +function NLPModels.cons!(nlp::ArrayDiffNLPModel, x::AbstractVector, c::AbstractVector) + NLPModels.increment!(nlp, :neval_cons) + for g in eachindex(nlp.cons) + ArrayDiff.eval_residual!(nlp.cons[g], view(c, _slice(nlp, g)), x) + end + return c +end + +# Stacked J = [J_1; …; J_G], so (Jv)_g = J_g v: write each group into its slice. +function NLPModels.jprod!( + nlp::ArrayDiffNLPModel, + x::AbstractVector, + v::AbstractVector, + Jv::AbstractVector, +) + NLPModels.increment!(nlp, :neval_jprod) + for g in eachindex(nlp.cons) + ArrayDiff.eval_residual_jprod!(nlp.cons[g], view(Jv, _slice(nlp, g)), x, v) + end + return Jv +end + +# J' v = Σ_g J_g' v_g (v_g the slice of v): accumulate group contributions. +function NLPModels.jtprod!( + nlp::ArrayDiffNLPModel, + x::AbstractVector, + v::AbstractVector, + Jtv::AbstractVector, +) + NLPModels.increment!(nlp, :neval_jtprod) + fill!(Jtv, zero(eltype(Jtv))) + for g in eachindex(nlp.cons) + vg = collect_slice(v, _slice(nlp, g)) + ArrayDiff.eval_residual_jtprod!(nlp.cons[g], nlp.Jtv_tmp, x, vg) + Jtv .+= nlp.Jtv_tmp + end + return Jtv +end + +# `eval_residual_jtprod!` seeds the residual root with `v`; a plain `view` is +# fine on CPU and GPU. Kept as a hook in case a contiguous copy is needed. +collect_slice(v::AbstractVector, r) = view(v, r) + +# ── Explicit Jacobian (for interior-point solvers like MadNLP) ──────────────── +# +# ArrayDiff is matrix-free (jprod/jtprod), but MadNLP's KKT system needs the +# constraint Jacobian as coordinates. We materialize it densely: row i of J is +# ∇c_i = J' e_i, one reverse pass per constraint row. Fine for the dense KKT +# path on modest problems; large/GPU problems would want a sparse assembly. +# Structure is emitted in column-major dense order (all rows for column 1, +# then column 2, ...), and `jac_coord!` fills the same order. + +function NLPModels.jac_structure!( + nlp::ArrayDiffNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + m, n = nlp.meta.ncon, nlp.meta.nvar + k = 0 + for j in 1:n, i in 1:m + k += 1 + rows[k] = i + cols[k] = j + end + return rows, cols +end + +function NLPModels.jac_coord!( + nlp::ArrayDiffNLPModel, + x::AbstractVector, + vals::AbstractVector, +) + NLPModels.increment!(nlp, :neval_jac) + m, n = nlp.meta.ncon, nlp.meta.nvar + ei = fill!(similar(x, m), zero(eltype(x))) + row = similar(x, n) + valsm = reshape(vals, m, n) # column-major: valsm[i, j] = J[i, j] + for i in 1:m + fill!(ei, zero(eltype(x))) + ei_view = view(ei, i:i) + ei_view .= one(eltype(x)) + NLPModels.jtprod!(nlp, x, ei, row) # row = ∇c_i = J[i, :] + valsm[i, :] .= row + end + return vals +end + +# No exact Hessian (`nnzh == 0`). Quasi-Newton solvers (MadNLP's CompactLBFGS, +# Percival's LBFGSModel) build their own approximation from gradients; these +# no-ops satisfy the interface without ArrayDiff ever computing second order. +function NLPModels.hess_structure!( + nlp::ArrayDiffNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + return rows, cols +end + +function NLPModels.hess_coord!( + nlp::ArrayDiffNLPModel, + x::AbstractVector, + y::AbstractVector, + vals::AbstractVector; + obj_weight = one(eltype(x)), +) + NLPModels.increment!(nlp, :neval_hess) + return vals +end diff --git a/perf/percival/build_percival.jl b/perf/percival/build_percival.jl new file mode 100644 index 0000000..d0c83a4 --- /dev/null +++ b/perf/percival/build_percival.jl @@ -0,0 +1,177 @@ +# Build the two AC-OPF forms as `ArrayDiffNLPModel`s (min f s.t. c(x)=0, +# l≤x≤u), reusing the data and structured-matrix helpers from ../acopf. +# Inequalities are converted to equalities with box-constrained slacks, so the +# model has only equality constraints + variable bounds — the form Percival's +# `Val{:equ}` path handles directly. + +import ArrayDiff +import JuMP +import LinearAlgebra +import MathOptInterface as MOI +import SparseArrays + +const ACOPF_DIR = joinpath(@__DIR__, "..", "acopf") +include(joinpath(ACOPF_DIR, "structured.jl")) +include(joinpath(ACOPF_DIR, "data.jl")) +include("adnlp.jl") + +storage_type(::ArrayDiff.Mode{S}) where {S} = S +to_storage(mode, v::Vector{Float64}) = storage_type(mode)(v) + +# Build an ArrayDiff residual evaluator for a vector array expression. +function residual_evaluator(model, expr, mode) + ad = ArrayDiff.model(mode) + ArrayDiff.set_residual!(ad, JuMP.moi_function(expr)) + ev = ArrayDiff.Evaluator(ad, mode, JuMP.index.(JuMP.all_variables(model))) + MOI.initialize(ev, Symbol[:Grad, :Jac, :JacVec]) + return ev +end + +function objective_evaluator(model, expr, mode) + ad = ArrayDiff.model(mode) + MOI.Nonlinear.set_objective(ad, JuMP.moi_function(expr)) + ev = ArrayDiff.Evaluator(ad, mode, JuMP.index.(JuMP.all_variables(model))) + MOI.initialize(ev, Symbol[:Grad]) + return ev +end + +# ── Form 1: rectangular voltages + Ybus ────────────────────────────────────── + +function build_rect_nlp( + d::RectData; + mode = ArrayDiff.Mode{Vector{Float64}}(), + matrix = ELLMatrix, + device = identity, +) + N = d.N + Gm = map_storage(device, matrix(d.G)) + Bm = map_storage(device, matrix(d.B)) + model = JuMP.Model() + JuMP.@variable(model, Vr[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Vi[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Pg[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, Qg[1:N], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, w[1:N], container = ArrayDiff.ArrayOfVariables) + Ir = Gm * Vr .- Bm * Vi + Ii = Gm * Vi .+ Bm * Vr + rP = Pg .- d.Pd .- (Vr .* Ir .+ Vi .* Ii) + rQ = Qg .- d.Qd .- (Vi .* Ir .- Vr .* Ii) + rW = Vr .^ 2 .+ Vi .^ 2 .- w + obj_expr = sum(d.c2 .* Pg .^ 2 .+ d.c1 .* Pg) + obj = objective_evaluator(model, obj_expr, mode) + cons = [ + residual_evaluator(model, rP, mode), + residual_evaluator(model, rQ, mode), + residual_evaluator(model, rW, mode), + ] + lb_Vr = fill(-d.vmax, N) + ub_Vr = fill(d.vmax, N) + lb_Vi = fill(-d.vmax, N) + ub_Vi = fill(d.vmax, N) + lb_Vr[1] = 0.0 + lb_Vi[1] = ub_Vi[1] = 0.0 + lb = vcat(lb_Vr, lb_Vi, d.Pg_lb, d.Qg_lb, fill(d.vmin^2, N)) + ub = vcat(ub_Vr, ub_Vi, d.Pg_ub, d.Qg_ub, fill(d.vmax^2, N)) + x0 = vcat( + ones(N), + zeros(N), + (d.Pg_lb .+ d.Pg_ub) ./ 2, + (d.Qg_lb .+ d.Qg_ub) ./ 2, + ones(N), + ) + nvar = 5N + return ArrayDiffNLPModel( + obj, + cons, + nvar, + to_storage(mode, lb), + to_storage(mode, ub), + to_storage(mode, x0); + name = "acopf-rect", + ), + d.c0 # objective offset (constant term), added back when reporting +end + +# ── Form 2: polar voltages, sin/cos branch flows ───────────────────────────── + +function build_polar_nlp( + d::PolarData; + mode = ArrayDiff.Mode{Vector{Float64}}(), + device = identity, + use_gather::Bool = true, +) + nb, ng, nl = d.nbus, d.ngen, d.nbranch + make_gather = if use_gather + idx -> map_storage(device, GatherMatrix(idx, nb)) + else + idx -> device( + SparseArrays.sparse(1:length(idx), idx, 1.0, length(idx), nb), + ) + end + Fg = make_gather(d.f_bus) + Tg = make_gather(d.t_bus) + Cg = make_gather(d.gen_bus) + Ft = LinearAlgebra.transpose(Fg) + Tt = LinearAlgebra.transpose(Tg) + Ct = LinearAlgebra.transpose(Cg) + model = JuMP.Model() + JuMP.@variable(model, va[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, vm[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pt[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qt[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, sf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, st[1:nl], container = ArrayDiff.ArrayOfVariables) + vmf = Fg * vm + vmt = Tg * vm + Δ = Fg * va .- Tg * va + cΔ = cos.(Δ) + sΔ = sin.(Δ) + vv = vmf .* vmt + g1 = pf .- (d.c5 .* vmf .^ 2 .+ d.c3 .* (vv .* cΔ) .+ d.c4 .* (vv .* sΔ)) + g2 = qf .+ d.c6 .* vmf .^ 2 .+ d.c4 .* (vv .* cΔ) .- d.c3 .* (vv .* sΔ) + g3 = pt .- (d.c7 .* vmt .^ 2 .+ d.c1 .* (vv .* cΔ) .- d.c2 .* (vv .* sΔ)) + g4 = qt .+ d.c8 .* vmt .^ 2 .+ d.c2 .* (vv .* cΔ) .+ d.c1 .* (vv .* sΔ) + g5 = Ft * pf .+ Tt * pt .- Ct * pg .+ d.pd .+ d.gs .* vm .^ 2 + g6 = Ft * qf .+ Tt * qt .- Ct * qg .+ d.qd .- d.bs .* vm .^ 2 + g7 = pf .^ 2 .+ qf .^ 2 .+ sf .- d.rate_a_sq + g8 = pt .^ 2 .+ qt .^ 2 .+ st .- d.rate_a_sq + obj_expr = sum(d.cost1 .* pg .^ 2 .+ d.cost2 .* pg) + obj = objective_evaluator(model, obj_expr, mode) + cons = [ + residual_evaluator(model, g, mode) for + g in (g1, g2, g3, g4, g5, g6, g7, g8) + ] + lb_va = fill(-Inf, nb) + ub_va = fill(Inf, nb) + for r in d.ref_buses + lb_va[r] = ub_va[r] = 0.0 + end + lb = vcat( + lb_va, d.vmin, d.pmin, d.qmin, + -d.rate_a, -d.rate_a, -d.rate_a, -d.rate_a, zeros(nl), zeros(nl), + ) + ub = vcat( + ub_va, d.vmax, d.pmax, d.qmax, + d.rate_a, d.rate_a, d.rate_a, d.rate_a, d.rate_a_sq, d.rate_a_sq, + ) + x0 = vcat( + zeros(nb), ones(nb), (d.pmin .+ d.pmax) ./ 2, (d.qmin .+ d.qmax) ./ 2, + zeros(nl), zeros(nl), zeros(nl), zeros(nl), + copy(d.rate_a_sq), copy(d.rate_a_sq), + ) + nvar = 2nb + 2ng + 6nl + return ArrayDiffNLPModel( + obj, + cons, + nvar, + to_storage(mode, lb), + to_storage(mode, ub), + to_storage(mode, x0); + name = "acopf-polar", + ), + sum(d.cost3) +end diff --git a/perf/percival/compare.jl b/perf/percival/compare.jl new file mode 100644 index 0000000..e531ea4 --- /dev/null +++ b/perf/percival/compare.jl @@ -0,0 +1,75 @@ +# Compare NLPModels solvers on the AC-OPF `ArrayDiffNLPModel` (vectorized +# ArrayDiff derivatives), against Ipopt references. Everything goes through the +# NLPModels interface: obj/grad from the objective evaluator, cons/jprod/jtprod +# from the vectorized residual evaluators, and (for MadNLP) an explicit +# Jacobian materialized from jtprod. +# +# * MadNLP — interior point + CompactLBFGS quasi-Newton (no exact Hessian). +# * Percival — augmented Lagrangian; LBFGSModel subproblems + TRON (CPU) or +# the new SPGSubSolver (GPU-compatible). + +include("build_percival.jl") +include(joinpath(ACOPF_DIR, "reference.jl")) + +import MadNLP +import NLPModelsModifiers +import Percival +import Printf + +function madnlp_solve(nlp; kwargs...) + return MadNLP.madnlp( + nlp; + hessian_approximation = MadNLP.CompactLBFGS, + print_level = MadNLP.ERROR, + kwargs..., + ) +end + +percival_lbfgs(nlp; kwargs...) = Percival.percival( + nlp; + subproblem_modifier = NLPModelsModifiers.LBFGSModel, + verbose = 0, + kwargs..., +) + +function line(name, obj, ref, feas, status, iters) + Printf.@printf( + " %-18s obj %12.2f gap %8.4f%% feas %9.2e %-16s it %d\n", + name, obj, 100 * (obj - ref) / abs(ref), feas, status, iters, + ) + return +end + +function compare_rect() + d = case9mod() + ref = rect_reference(d) + println("── rect (case9mod), Ipopt ref $(round(ref; digits=2))") + nlp, off = build_rect_nlp(d) + s = madnlp_solve(nlp; max_iter = 1000) + line("MadNLP/LBFGS", s.objective + off, ref, s.primal_feas, s.status, s.iter) + nlp, off = build_rect_nlp(d) + s = percival_lbfgs(nlp; max_iter = 300, max_time = 60.0) + line("Percival/LBFGS+Tron", s.objective + off, ref, s.primal_feas, s.status, s.iter) + return +end + +function compare_polar(case = "case9.m") + d = parse_polar_case(matpower_case(case)) + ref = polar_reference(matpower_case(case)) + println("── polar ($case), Ipopt ref $(round(ref; digits=2))") + nlp, off = build_polar_nlp(d) + s = madnlp_solve(nlp; max_iter = 1000) + line("MadNLP/LBFGS", s.objective + off, ref, s.primal_feas, s.status, s.iter) + nlp, off = build_polar_nlp(d) + s = percival_lbfgs(nlp; max_iter = 300, max_time = 60.0) + line("Percival/LBFGS+Tron", s.objective + off, ref, s.primal_feas, s.status, s.iter) + return +end + +function compare_all() + compare_rect() + compare_polar("case9.m") + compare_polar("case14.m") + compare_polar("case30.m") + return +end diff --git a/perf/percival/gpu_percival.jl b/perf/percival/gpu_percival.jl new file mode 100644 index 0000000..105ded5 --- /dev/null +++ b/perf/percival/gpu_percival.jl @@ -0,0 +1,57 @@ +# Run the Percival path with the tape + solver state on `JLArray` (GPUArrays +# reference backend) with scalar indexing disallowed — the same device +# semantics as CUDA. Validates the Percival GPU fix (AugLagModel storage type) +# and that TRON + LBFGSModel + AugLagModel are device-generic. + +include("build_percival.jl") + +import GPUArraysCore +import JLArrays +import NLPModels +import NLPModelsModifiers +import Percival +import Printf + +GPUArraysCore.allowscalar(false) +const JLV = JLArrays.JLArray{Float64, 1} +jl_device(x::AbstractArray) = JLArrays.JLArray(x) + +function solve_gpu(nlp; kwargs...) + return Percival.percival( + nlp; + subproblem_modifier = NLPModelsModifiers.LBFGSModel, + kwargs..., + ) +end + +# GPU-native path: the SPG subsolver (pure broadcasts) replaces TRON, and the +# AL subproblem is minimized directly (no quasi-Newton wrapper needed since +# SPG is first-order). +function solve_gpu_spg(nlp; kwargs...) + return Percival.percival(nlp; subsolver = Percival.SPGSubSolver, kwargs...) +end + +# Compare a handful of NLPModels evaluations CPU vs JLArray to prove the +# device path is numerically identical before trusting the full solve. +function check_evals(nlp_cpu, nlp_gpu) + x_cpu = Vector(nlp_cpu.meta.x0) .+ 0.01 + x_gpu = JLV(x_cpu) + g_cpu = similar(x_cpu) + g_gpu = JLV(zero(x_cpu)) + NLPModels.grad!(nlp_cpu, x_cpu, g_cpu) + NLPModels.grad!(nlp_gpu, x_gpu, g_gpu) + @assert Vector(g_gpu) ≈ g_cpu + c_cpu = zeros(nlp_cpu.meta.ncon) + c_gpu = JLV(zeros(nlp_gpu.meta.ncon)) + NLPModels.cons!(nlp_cpu, x_cpu, c_cpu) + NLPModels.cons!(nlp_gpu, x_gpu, c_gpu) + @assert Vector(c_gpu) ≈ c_cpu + v_cpu = ones(nlp_cpu.meta.ncon) + Jtv_cpu = similar(x_cpu) + Jtv_gpu = JLV(zero(x_cpu)) + NLPModels.jtprod!(nlp_cpu, x_cpu, v_cpu, Jtv_cpu) + NLPModels.jtprod!(nlp_gpu, x_gpu, JLV(v_cpu), Jtv_gpu) + @assert Vector(Jtv_gpu) ≈ Jtv_cpu + println(" evals CPU vs JLArray match (grad, cons, jtprod)") + return +end diff --git a/perf/percival/jump_acopf.jl b/perf/percival/jump_acopf.jl new file mode 100644 index 0000000..4978e80 --- /dev/null +++ b/perf/percival/jump_acopf.jl @@ -0,0 +1,123 @@ +# A complete polar AC-OPF written as a JuMP model, using ArrayDiff's vectorized +# array variables and the new vector `@constraint(..., expr in set)` support. +# Every power-flow equation is ONE vectorized constraint (an +# `ArrayNonlinearFunction in MOI.Zeros`), not `nbranch` scalar rows. +# +# julia --project=perf/percival +# include("perf/percival/jump_acopf.jl") +# import Percival, NLPModelsModifiers, NLPModelsJuMP +# model = build_jump_polar("case9.m"; slacks = true) # bundled case, offline +# set_optimizer(model, NLPModelsJuMP.Optimizer) +# set_attribute(model, "solver", +# nlp -> Percival.PercivalSolver(nlp; +# subproblem_modifier = NLPModelsModifiers.LBFGSModel)) +# set_attribute(model, "subproblem_modifier", NLPModelsModifiers.LBFGSModel) +# set_attribute(model, MOI.AutomaticDifferentiationBackend(), ArrayDiff.Mode()) +# optimize!(model) # case9: 347.67 vs Ipopt 347.70 (−0.008%) +# objective_value(model); value.(vm) +# +# NLPModelsJuMP (branch bl/arraydiff) collects the `ArrayNonlinearFunction in +# Zeros/Nonpositives` constraints into a constrained `ArrayDiffNLPModel` (one +# vectorized residual evaluator per constraint), which any NLPModels solver +# consumes. With Percival use `slacks = true` (its augmented Lagrangian wants +# equality constraints + bounds); `build_polar_nlp` (build_percival.jl) still +# assembles the same NLPModel directly without going through MOI. + +include("build_percival.jl") # parse_polar_case, structured matrices, matpower_case + +import JuMP +import MathOptInterface as MOI + +_case(name) = isfile(matpower_case(name)) ? matpower_case(name) : name + +# `slacks = true` converts the thermal-limit inequalities into equalities with +# box-bounded slack variables (`p² + q² + s = rate²`, `0 ≤ s ≤ rate²`), so the +# model is equality-constrained + bounds — the form Percival's augmented +# Lagrangian handles directly. `slacks = false` keeps them as `Nonpositives` +# inequalities (matching the ExaModels/GenOpt variable layout). +function build_jump_polar(name::AbstractString; slacks::Bool = false) + d = parse_polar_case(_case(name)) + nb, ng, nl = d.nbus, d.ngen, d.nbranch + Fg = GatherMatrix(d.f_bus, nb) + Tg = GatherMatrix(d.t_bus, nb) + Cg = GatherMatrix(d.gen_bus, nb) + Ft, Tt, Ct = transpose(Fg), transpose(Tg), transpose(Cg) + + model = JuMP.Model() + JuMP.@variable(model, va[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, vm[1:nb], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qg[1:ng], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, pt[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, qt[1:nl], container = ArrayDiff.ArrayOfVariables) + + # Bounds (ArrayOfVariables is indexable → set per block). + _set_bounds(vm, d.vmin, d.vmax) + _set_bounds(pg, d.pmin, d.pmax) + _set_bounds(qg, d.qmin, d.qmax) + _set_bounds(pf, -d.rate_a, d.rate_a) + _set_bounds(pt, -d.rate_a, d.rate_a) + _set_bounds(qf, -d.rate_a, d.rate_a) + _set_bounds(qt, -d.rate_a, d.rate_a) + for i in 1:nb # flat voltage start + JuMP.set_start_value(vm[i], 1.0) + end + for r in d.ref_buses # reference-bus angle fixed to 0 + JuMP.fix(va[r], 0.0) + end + + # Vectorized branch/bus quantities. + vmf, vmt = Fg * vm, Tg * vm + Δ = Fg * va .- Tg * va + cΔ, sΔ = cos.(Δ), sin.(Δ) + vv = vmf .* vmt + + # Objective: generation cost (scalar reduction of array expressions). + JuMP.@objective(model, Min, sum(d.cost1 .* pg .^ 2 .+ d.cost2 .* pg)) + + # Power-flow equalities — one vector constraint each. + JuMP.@constraint(model, pf .- (d.c5 .* vmf .^ 2 .+ d.c3 .* (vv .* cΔ) .+ d.c4 .* (vv .* sΔ)) in MOI.Zeros(nl)) + JuMP.@constraint(model, qf .+ d.c6 .* vmf .^ 2 .+ d.c4 .* (vv .* cΔ) .- d.c3 .* (vv .* sΔ) in MOI.Zeros(nl)) + JuMP.@constraint(model, pt .- (d.c7 .* vmt .^ 2 .+ d.c1 .* (vv .* cΔ) .- d.c2 .* (vv .* sΔ)) in MOI.Zeros(nl)) + JuMP.@constraint(model, qt .+ d.c8 .* vmt .^ 2 .+ d.c2 .* (vv .* cΔ) .+ d.c1 .* (vv .* sΔ) in MOI.Zeros(nl)) + # Nodal power balance (scatter branch flows / generation to buses). + JuMP.@constraint(model, Ft * pf .+ Tt * pt .- Ct * pg .+ d.pd .+ d.gs .* vm .^ 2 in MOI.Zeros(nb)) + JuMP.@constraint(model, Ft * qf .+ Tt * qt .- Ct * qg .+ d.qd .- d.bs .* vm .^ 2 in MOI.Zeros(nb)) + # Thermal limits |S|² ≤ rate² — one vector constraint each. + if slacks + JuMP.@variable(model, sf[1:nl], container = ArrayDiff.ArrayOfVariables) + JuMP.@variable(model, st[1:nl], container = ArrayDiff.ArrayOfVariables) + _set_bounds(sf, zeros(nl), d.rate_a_sq) + _set_bounds(st, zeros(nl), d.rate_a_sq) + for i in 1:nl # start on the constraint: p = q = 0 ⇒ s = rate² + JuMP.set_start_value(sf[i], d.rate_a_sq[i]) + JuMP.set_start_value(st[i], d.rate_a_sq[i]) + end + JuMP.@constraint(model, pf .^ 2 .+ qf .^ 2 .+ sf .- d.rate_a_sq in MOI.Zeros(nl)) + JuMP.@constraint(model, pt .^ 2 .+ qt .^ 2 .+ st .- d.rate_a_sq in MOI.Zeros(nl)) + else + JuMP.@constraint(model, pf .^ 2 .+ qf .^ 2 .- d.rate_a_sq in MOI.Nonpositives(nl)) + JuMP.@constraint(model, pt .^ 2 .+ qt .^ 2 .- d.rate_a_sq in MOI.Nonpositives(nl)) + end + return model +end + +function _set_bounds(v, lo, hi) + for i in eachindex(lo) + JuMP.set_lower_bound(v[i], lo[i]) + JuMP.set_upper_bound(v[i], hi[i]) + end + return +end + +# Show the vectorized constraints stored on the model's backend. +function describe_constraints(model) + b = JuMP.backend(model) + for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent()) + n = MOI.get(model, MOI.NumberOfConstraints{F,S}()) + println(" ", n, " × ", F, " in ", S) + end + return +end diff --git a/perf/percival/run_madnlp.jl b/perf/percival/run_madnlp.jl new file mode 100644 index 0000000..6468057 --- /dev/null +++ b/perf/percival/run_madnlp.jl @@ -0,0 +1,49 @@ +# Solve the AC-OPF ArrayDiffNLPModel with MadNLP in quasi-Newton mode +# (`hessian_approximation = CompactLBFGS`), so no exact Hessian is needed — +# ArrayDiff supplies the objective gradient and the constraint Jacobian +# (materialized via `jac_coord!`). MadNLP is an interior-point method (a +# different family than Percival's augmented Lagrangian), and on CPU it +# converges tightly. The same model type also feeds MadNLPGPU on real CUDA +# hardware (untested here — no GPU in this container). + +include("build_percival.jl") +include(joinpath(ACOPF_DIR, "reference.jl")) + +import MadNLP +import Printf + +function solve_madnlp(nlp; quasi_newton = true, print_level = MadNLP.INFO, kwargs...) + opts = Dict{Symbol, Any}(kwargs) + if quasi_newton + opts[:hessian_approximation] = MadNLP.CompactLBFGS + end + return MadNLP.madnlp(nlp; print_level = print_level, opts...) +end + +function madnlp_rect(; kwargs...) + d = case9mod() + nlp, off = build_rect_nlp(d) + stats = solve_madnlp(nlp; kwargs...) + return nlp, stats, off +end + +function madnlp_polar(case = "case9.m"; kwargs...) + d = parse_polar_case(matpower_case(case)) + nlp, off = build_polar_nlp(d) + stats = solve_madnlp(nlp; kwargs...) + return nlp, stats, off +end + +function report_madnlp(name, stats, off, ref) + obj = stats.objective + off + Printf.@printf( + "%-16s MadNLP obj %.2f (Ipopt %.2f, gap %.4f%%) status %s iters %d\n", + name, + obj, + ref, + 100 * (obj - ref) / abs(ref), + stats.status, + stats.iter, + ) + return obj +end diff --git a/perf/percival/run_percival.jl b/perf/percival/run_percival.jl new file mode 100644 index 0000000..f7fe713 --- /dev/null +++ b/perf/percival/run_percival.jl @@ -0,0 +1,49 @@ +# Solve both AC-OPF forms through NLPModels → Percival, using the vectorized +# ArrayDiff derivatives. Percival's TRON subproblem needs Hessian-vector +# products; ArrayDiff is first-order, so we wrap the AL subproblem with an +# LBFGS quasi-Newton model (`subproblem_modifier = LBFGSModel`) which builds +# its Hessian approximation from gradient differences only. + +include("build_percival.jl") +include(joinpath(ACOPF_DIR, "reference.jl")) + +import NLPModelsModifiers +import Percival +import Printf + +function solve_percival(nlp; kwargs...) + return Percival.percival( + nlp; + subproblem_modifier = NLPModelsModifiers.LBFGSModel, + kwargs..., + ) +end + +function run_rect(; verbose = 1, kwargs...) + d = case9mod() + nlp, offset = build_rect_nlp(d) + stats = solve_percival(nlp; verbose = verbose, kwargs...) + return nlp, stats, offset +end + +function run_polar(case = "case9.m"; verbose = 1, kwargs...) + d = parse_polar_case(matpower_case(case)) + nlp, offset = build_polar_nlp(d) + stats = solve_percival(nlp; verbose = verbose, kwargs...) + return nlp, stats, offset +end + +function report(name, stats, offset, ref) + obj = stats.objective + offset + Printf.@printf( + "%-14s Percival obj %.2f (Ipopt %.2f, gap %.3f%%) feas %.2e status %s iters %d\n", + name, + obj, + ref, + 100 * (obj - ref) / abs(ref), + stats.primal_feas, + stats.status, + stats.iter, + ) + return obj +end diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl index ce85a9a..06558bd 100644 --- a/src/JuMP/moi_bridge.jl +++ b/src/JuMP/moi_bridge.jl @@ -12,7 +12,7 @@ function JuMP.moi_function(x::GenericArrayExpr{V,N}) where {V,N} return ArrayNonlinearFunction{N}(x.head, args, x.size, x.broadcasted) end -JuMP.moi_function(x::Array{<:Real}) = x +JuMP.moi_function(x::AbstractArray{<:Real}) = x # ── Detect whether a JuMP expression contains array args ───────────────────── @@ -46,3 +46,49 @@ function JuMP.set_objective_function( model.is_model_dirty = true return end + +# ── Vector constraints over array expressions ──────────────────────────────── +# +# `@constraint(model, expr in set)` where `expr` is an `AbstractJuMPArray` +# (e.g. a vectorized residual `Pg .- Pd .- ...`) and `set` is a vector set +# (`MOI.Zeros`, `MOI.Nonnegatives`, `MOI.Nonpositives`). JuMP's default +# `VectorConstraint` scalarizes the function (`func[idx]` for each index), +# which our array expressions deliberately don't support. Instead we keep the +# expression whole: `moi_function` turns it into a single +# `ArrayNonlinearFunction`, preserving the vectorized structure end-to-end. +# +# Relies on JuMP #3451 (`moi_function`/`_is_real` over `AbstractArray`). + +struct _ArrayVectorConstraint{F<:AbstractJuMPArray,S<:MOI.AbstractVectorSet} <: + JuMP.AbstractConstraint + func::F + set::S +end + +function JuMP.build_constraint( + _error::Function, + func::AbstractJuMPArray, + set::MOI.AbstractVectorSet, +) + n = length(func) + if n != MOI.dimension(set) + _error( + "Dimension of the function ($n) does not match the dimension of " * + "the set ($(MOI.dimension(set))).", + ) + end + return _ArrayVectorConstraint(func, set) +end + +# `jump_function`/`moi_function`/`moi_set` fall back to the generic +# `AbstractConstraint` methods (which read `.func`/`.set`), so we only need the +# shape and the belongs-to-model check. `moi_function(::GenericArrayExpr)` +# already yields the `ArrayNonlinearFunction`. +JuMP.shape(::_ArrayVectorConstraint) = JuMP.VectorShape() + +# The array expression carries whole variable blocks; the per-scalar ownership +# check JuMP does for `Vector`-valued functions doesn't apply. Some JuMP +# versions check the constraint, others the function, so cover both. +JuMP.check_belongs_to_model(::AbstractJuMPArray, ::JuMP.AbstractModel) = nothing +JuMP.check_belongs_to_model(::_ArrayVectorConstraint, ::JuMP.AbstractModel) = + nothing diff --git a/src/JuMP/nlp_expr.jl b/src/JuMP/nlp_expr.jl index 8ce1ab2..d648328 100644 --- a/src/JuMP/nlp_expr.jl +++ b/src/JuMP/nlp_expr.jl @@ -23,4 +23,4 @@ JuMP.variable_ref_type(::Type{GenericArrayExpr{V,N}}) where {V,N} = V JuMP._is_real(::GenericArrayExpr) = true -JuMP._is_real(::Array{<:Real}) = true +JuMP._is_real(::AbstractArray{<:Real}) = true diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl index 04b5a13..2ec9f20 100644 --- a/src/JuMP/operators.jl +++ b/src/JuMP/operators.jl @@ -2,10 +2,10 @@ function _matmul(::Type{V}, A, B) where {V} return GenericMatrixExpr{V}(:*, Any[A, B], (size(A, 1), size(B, 2)), false) end -function Base.:(*)(A::AbstractJuMPMatrix, B::Matrix) +function Base.:(*)(A::AbstractJuMPMatrix, B::AbstractMatrix) return _matmul(JuMP.variable_ref_type(A), A, B) end -function Base.:(*)(A::Matrix, B::AbstractJuMPMatrix) +function Base.:(*)(A::AbstractMatrix, B::AbstractJuMPMatrix) return _matmul(JuMP.variable_ref_type(B), A, B) end function Base.:(*)(A::AbstractJuMPMatrix, B::AbstractJuMPMatrix) @@ -14,15 +14,26 @@ end # Matrix-vector products: output is a 1-D `GenericArrayExpr` of length # `size(A, 1)`. Allows users to write `W * x` for a vector variable `x`. +# The constant matrix `A` can be any `AbstractMatrix`: dense matrices are +# serialized on the AD tape while sparse/structured ones are kept by +# reference (see `NODE_ARRAY_VALUE`). function _matvec(::Type{V}, A, b) where {V} return GenericArrayExpr{V,1}(:*, Any[A, b], (size(A, 1),), false) end -function Base.:(*)(A::AbstractJuMPMatrix, b::Vector) +function Base.:(*)(A::AbstractJuMPMatrix, b::AbstractVector) return _matvec(JuMP.variable_ref_type(A), A, b) end -function Base.:(*)(A::Matrix, b::AbstractJuMPVector{T}) where {T} +function Base.:(*)(A::AbstractMatrix, b::AbstractJuMPVector{T}) where {T} + return _matvec(JuMP.variable_ref_type(b), A, b) +end + +# Disambiguate against `LinearAlgebra.:*(::Diagonal, ::AbstractVector)`. +function Base.:(*)( + A::LinearAlgebra.Diagonal, + b::AbstractJuMPVector{T}, +) where {T} return _matvec(JuMP.variable_ref_type(b), A, b) end @@ -30,6 +41,18 @@ function Base.:(*)(A::AbstractJuMPMatrix, b::AbstractJuMPVector{T}) where {T} return _matvec(JuMP.variable_ref_type(A), A, b) end +# The JuMP macros rewrite `A * x` into MutableArithmetics calls whose generic +# matrix-vector product would silently scalarize the (indexable) +# `ArrayOfVariables` into a `Vector{AffExpr}`, defeating the vectorization. +# Route it back to the whole-array product. +function JuMP._MA.operate( + ::typeof(*), + A::AbstractMatrix{<:Real}, + x::AbstractJuMPArray, +) + return A * x +end + function __broadcast( ::Type{V}, axes::NTuple{N,Base.OneTo{Int}}, @@ -40,6 +63,14 @@ function __broadcast( end function _broadcast(::Type{V}, op::Function, args...) where {V} + # The JuMP macros rewrite broadcasted `.+`/`.-` into MutableArithmetics' + # fused `add_mul`/`sub_mul` (`op(a, b, c...) = a ± b * c * ...`). The AD + # tape only knows the elementary operators, so normalize here: + if op === JuMP._MA.add_mul || op === JuMP._MA.sub_mul + base = args[1] + rest = length(args) == 2 ? args[2] : _broadcast(V, *, args[2:end]...) + return _broadcast(V, op === JuMP._MA.add_mul ? (+) : (-), base, rest) + end return __broadcast(V, Broadcast.combine_axes(args...), op, Any[args...]) end @@ -47,6 +78,12 @@ function Base.broadcasted(op::Function, x::AbstractJuMPArray) return _broadcast(JuMP.variable_ref_type(x), op, x) end +# `value.(x)` is a query, not an expression: return the solution values +# instead of building a `GenericArrayExpr`. (`collect` materializes the +# indexable `ArrayOfVariables` into a `Vector{VariableRef}` first.) +Base.broadcasted(::typeof(JuMP.value), x::ArrayOfVariables) = + JuMP.value.(collect(x)) + function Base.broadcasted(op::Function, x::AbstractJuMPArray, y::AbstractArray) return _broadcast(JuMP.variable_ref_type(x), op, x, y) end diff --git a/src/array_nonlinear_function.jl b/src/array_nonlinear_function.jl index 49cfce4..13584cb 100644 --- a/src/array_nonlinear_function.jl +++ b/src/array_nonlinear_function.jl @@ -58,6 +58,13 @@ function Base.copy(f::ArrayOfContiguousVariables{N}) where {N} return f # immutable end +# An `ArrayNonlinearFunction` is always in canonical form (there are no +# duplicate/zero terms to merge as there would be for an affine function), so +# `canonicalize!` — called by MOI when a constraint is added to a cache — is a +# no-op. Without this, adding an `ArrayNonlinearFunction`-in-set constraint to +# a model errors. +MOI.Utilities.canonicalize!(f::ArrayNonlinearFunction) = f + # map_indices: remap MOI.VariableIndex values during MOI.copy_to function MOI.Utilities.map_indices( index_map::F, diff --git a/src/graph_tools.jl b/src/graph_tools.jl index 4d20d47..558470e 100644 --- a/src/graph_tools.jl +++ b/src/graph_tools.jl @@ -33,6 +33,12 @@ field, which should be interpreted as follows: * `NODE_VALUE_BLOCK`: a contiguous block of constants. `index` is the start index in the `.values` field; the next `m * n` entries (column-major) are the block's data. Shape stored in `Expression.block_shapes`. + * `NODE_ARRAY_VALUE`: a constant `AbstractArray` stored by reference (not + serialized on the tape). `index` is the index into the `.arrays` field of + `Expression`. Shape stored in `Expression.block_shapes`. Used for + non-dense constants (for example sparse or structured matrices) so that + `LinearAlgebra.mul!` can dispatch on the array's own type instead of a + dense view of the tape. These nodes occupy no tape storage. """ @enum( NodeType, @@ -67,6 +73,8 @@ field, which should be interpreted as follows: NODE_VARIABLE_BLOCK, # Block-of-constants node. NODE_VALUE_BLOCK, + # Constant `AbstractArray` stored by reference in `Expression.arrays`. + NODE_ARRAY_VALUE, ) @enum(Linearity, CONSTANT, LINEAR, PIECEWISE_LINEAR, NONLINEAR) @@ -151,7 +159,9 @@ function _classify_linearity( if node.type == NODE_VARIABLE || node.type == NODE_VARIABLE_BLOCK linearity[k] = LINEAR continue - elseif node.type == NODE_VALUE || node.type == NODE_VALUE_BLOCK + elseif node.type == NODE_VALUE || + node.type == NODE_VALUE_BLOCK || + node.type == NODE_ARRAY_VALUE linearity[k] = CONSTANT continue elseif node.type == NODE_PARAMETER diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl index b0f54b9..3d72f57 100644 --- a/src/mathoptinterface_api.jl +++ b/src/mathoptinterface_api.jl @@ -227,12 +227,18 @@ function MOI.initialize( return end +# Read a single leading element without scalar indexing so this also works +# when the tape lives on a GPU array (GPUArrays disallows scalar `getindex` +# by default outside the REPL). +_get_first(v::Vector) = v[1] +_get_first(v::AbstractVector) = first(Vector(view(v, 1:1))) + function MOI.eval_objective(d::NLPEvaluator, x) if d.objective === nothing error("No nonlinear objective.") end _reverse_mode(d, x) - return something(d.objective).expr.forward_storage[1] + return _get_first(something(d.objective).expr.forward_storage) end function MOI.eval_objective_gradient(d::NLPEvaluator, g, x) diff --git a/src/parse_moi.jl b/src/parse_moi.jl index 4598058..5b98252 100644 --- a/src/parse_moi.jl +++ b/src/parse_moi.jl @@ -161,7 +161,7 @@ function _parse_moi_stack!( ::Vector{Tuple{Int,Any}}, ::Model, expr::Expression, - x::AbstractArray{<:Real}, + x::Array{<:Real}, parent_index::Int, ) # Emit a single value block. We push the flat values to @@ -175,3 +175,20 @@ function _parse_moi_stack!( expr.block_shapes[length(expr.nodes)] = collect(size(x)) return end + +function _parse_moi_stack!( + ::Vector{Tuple{Int,Any}}, + ::Model, + expr::Expression, + x::AbstractArray{<:Real}, + parent_index::Int, +) + # Non-dense constant array (sparse, structured, GPU-resident, ...): store + # it by reference in `expr.arrays` instead of serializing it on the tape. + # Evaluation will pass the array itself to `LinearAlgebra.mul!` so its + # specialized methods apply. The node occupies no tape storage. + push!(expr.arrays, x) + push!(expr.nodes, Node(NODE_ARRAY_VALUE, length(expr.arrays), parent_index)) + expr.block_shapes[length(expr.nodes)] = collect(size(x)) + return +end diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl index 7dd620e..a09a0ce 100644 --- a/src/reverse_mode.jl +++ b/src/reverse_mode.jl @@ -4,28 +4,68 @@ # Use of this source code is governed by an MIT-style license that can be found # in the LICENSE.md file or at https://opensource.org/licenses/MIT. +# `_reshape_call`-style kernels for a matmul node with a `NODE_ARRAY_VALUE` +# child: the constant array `A`/`B` is passed by reference (it has no tape +# storage) and the remaining argument comes as a tape view. +_mul_const_lhs!(A, out, x) = LinearAlgebra.mul!(out, A, x) +_mul_const_rhs!(B, out, x) = LinearAlgebra.mul!(out, x, B) +# Reverse: for `out = A * x` with constant `A`, `rev_x = A' * rev_out`; for +# `out = x * B` with constant `B`, `rev_x = rev_out * B'`. +function _mul_rev_rhs_const_lhs!(A, rev_x, rev_out) + return LinearAlgebra.mul!(rev_x, LinearAlgebra.transpose(A), rev_out) +end +function _mul_rev_lhs_const_rhs!(B, rev_x, rev_out) + return LinearAlgebra.mul!(rev_x, rev_out, LinearAlgebra.transpose(B)) +end + +function _is_constant_node(node::Node) + return node.type == NODE_VALUE_BLOCK || node.type == NODE_ARRAY_VALUE +end + # Reverse-mode contribution for a matmul node `k` with children `ix1`, `ix2`. # `f.sizes.ndims[k]` may be 1 (mat-vec) or 2 (mat-mat); `_reshape_call` picks # the right view type for each node and `LinearAlgebra.mul!` covers both # shape combinations. function _matmul_reverse!(f, k::Int, ix1::Int, ix2::Int) - if f.nodes[ix1].type != NODE_VALUE_BLOCK - _reshape_call( - f.forward_storage, - f.sizes, - (ix2,), - _matmul_reverse_outer, - (f.reverse_storage, f.sizes, true, ix1, k), - ) + if !_is_constant_node(f.nodes[ix1]) + if f.nodes[ix2].type == NODE_ARRAY_VALUE + B = f.const_arrays[f.nodes[ix2].index] + _reshape_call( + f.reverse_storage, + f.sizes, + (ix1, k), + _mul_rev_lhs_const_rhs!, + (B,), + ) + else + _reshape_call( + f.forward_storage, + f.sizes, + (ix2,), + _matmul_reverse_outer, + (f.reverse_storage, f.sizes, true, ix1, k), + ) + end end - if f.nodes[ix2].type != NODE_VALUE_BLOCK - _reshape_call( - f.forward_storage, - f.sizes, - (ix1,), - _matmul_reverse_outer, - (f.reverse_storage, f.sizes, false, ix2, k), - ) + if !_is_constant_node(f.nodes[ix2]) + if f.nodes[ix1].type == NODE_ARRAY_VALUE + A = f.const_arrays[f.nodes[ix1].index] + _reshape_call( + f.reverse_storage, + f.sizes, + (ix2, k), + _mul_rev_rhs_const_lhs!, + (A,), + ) + else + _reshape_call( + f.forward_storage, + f.sizes, + (ix1,), + _matmul_reverse_outer, + (f.reverse_storage, f.sizes, false, ix2, k), + ) + end end return end @@ -192,6 +232,9 @@ function _forward_eval( ) elseif node.type == NODE_VALUE_BLOCK # Pre-loaded into `forward_storage` at construction. + elseif node.type == NODE_ARRAY_VALUE + # Constant array kept by reference in `f.const_arrays`; it has no + # tape storage. Consumed directly by its parent (see `:*`). elseif node.type == NODE_SUBEXPRESSION f.forward_storage[j] = d.subexpression_forward_values[node.index] elseif node.type == NODE_PARAMETER @@ -235,14 +278,37 @@ function _forward_eval( # `_reshape_call` dispatches each node to the right view # type based on its `ndims`. `LinearAlgebra.mul!` then # picks the matching method — mat-mat for `ndims[k] == 2`, - # mat-vec for `ndims[k] == 1`. - _reshape_call( - f.forward_storage, - f.sizes, - (k, ix1, ix2), - LinearAlgebra.mul!, - (), - ) + # mat-vec for `ndims[k] == 1`. A `NODE_ARRAY_VALUE` child + # has no tape storage: pass the referenced array itself so + # `mul!` dispatches on its concrete (sparse/structured) + # type. + if f.nodes[ix1].type == NODE_ARRAY_VALUE + A = f.const_arrays[f.nodes[ix1].index] + _reshape_call( + f.forward_storage, + f.sizes, + (k, ix2), + _mul_const_lhs!, + (A,), + ) + elseif f.nodes[ix2].type == NODE_ARRAY_VALUE + B = f.const_arrays[f.nodes[ix2].index] + _reshape_call( + f.forward_storage, + f.sizes, + (k, ix1), + _mul_const_rhs!, + (B,), + ) + else + _reshape_call( + f.forward_storage, + f.sizes, + (k, ix1, ix2), + LinearAlgebra.mul!, + (), + ) + end # We deliberately don't write v1/v2 into partials_storage # here: the matmul reverse branch reads forward_storage # directly, so those writes were dead. @@ -593,6 +659,31 @@ function _forward_eval( partials = _view_linear(f.partials_storage, f.sizes, child_idx) out .= tanh.(inp) partials .= one(T) .- out .* out + elseif operators.univariate_operators[node.index] in + (:sin, :cos, :exp, :sqrt, :log) + # Whole-array broadcasts: a single fused kernel on GPU + # storage instead of one scalar `eval_univariate_function_and_ + # gradient` round-trip per element. + op_sym = operators.univariate_operators[node.index] + out = _view_linear(f.forward_storage, f.sizes, k) + inp = _view_linear(f.forward_storage, f.sizes, child_idx) + partials = _view_linear(f.partials_storage, f.sizes, child_idx) + if op_sym === :sin + out .= sin.(inp) + partials .= cos.(inp) + elseif op_sym === :cos + out .= cos.(inp) + partials .= .-sin.(inp) + elseif op_sym === :exp + out .= exp.(inp) + partials .= out + elseif op_sym === :sqrt + out .= sqrt.(inp) + partials .= inv.(2 .* out) + else # :log + out .= log.(inp) + partials .= inv.(inp) + end else for j in _eachindex(f.sizes, k) ret_f, ret_f′ = eval_univariate_function_and_gradient( diff --git a/src/sizes.jl b/src/sizes.jl index d168638..d52995d 100644 --- a/src/sizes.jl +++ b/src/sizes.jl @@ -580,7 +580,8 @@ function _infer_sizes( # by construction. if node.type == NODE_VARIABLE_BLOCK || node.type == NODE_VALUE_BLOCK || - node.type == NODE_MOI_VARIABLE_BLOCK + node.type == NODE_MOI_VARIABLE_BLOCK || + node.type == NODE_ARRAY_VALUE _add_size!(sizes, k, block_shapes[k]) continue end @@ -657,7 +658,11 @@ function _infer_sizes( end end for k in eachindex(nodes) - sizes.storage_offset[k+1] = sizes.storage_offset[k] + _length(sizes, k) + # `NODE_ARRAY_VALUE` constants live outside the tape (they're kept by + # reference in `const_arrays`), so they occupy no tape storage. Their + # shape is still recorded above for size inference of their parents. + len = nodes[k].type == NODE_ARRAY_VALUE ? 0 : _length(sizes, k) + sizes.storage_offset[k+1] = sizes.storage_offset[k] + len end return sizes end @@ -667,6 +672,11 @@ struct _SubexpressionStorage{T<:Real,S<:AbstractVector{T}} adj::SparseArrays.SparseMatrixCSC{Bool,Int} sizes::Sizes const_values::Vector{T} + # Constant arrays kept by reference (see `NODE_ARRAY_VALUE`), indexed by + # `node.index`. The element type is abstract on purpose: each entry can be + # a different sparse/structured/GPU matrix type and is only touched via + # dynamic dispatch to `LinearAlgebra.mul!` once per node per evaluation. + const_arrays::Vector{AbstractArray} forward_storage::S partials_storage::S reverse_storage::S @@ -682,6 +692,7 @@ struct _SubexpressionStorage{T<:Real,S<:AbstractVector{T}} nodes::Vector{Node}, adj::SparseArrays.SparseMatrixCSC{Bool,Int}, const_values::Vector{T}, + const_arrays::Vector{AbstractArray}, block_shapes::Dict{Int,Vector{Int}}, partials_storage_ϵ::Vector{Float64}, linearity::Linearity, @@ -715,6 +726,7 @@ struct _SubexpressionStorage{T<:Real,S<:AbstractVector{T}} adj, sizes, const_values, + const_arrays, forward_storage, fill!(S(undef, N), zero(T)), # partials_storage, fill!(S(undef, N), zero(T)), # reverse_storage, diff --git a/src/types.jl b/src/types.jl index 4e2ce6f..6c5b423 100644 --- a/src/types.jl +++ b/src/types.jl @@ -8,6 +8,7 @@ struct Expression nodes::Vector{Node} values::Vector{Float64} + arrays::Vector{AbstractArray} block_shapes::Dict{Int,Vector{Int}} end @@ -18,21 +19,31 @@ tree. `block_shapes[k]` is the shape of node `k` (as a `Vector{Int}` of dimensions: `[m]` for a 1D vector block, `[m, n]` for a 2D matrix block, `[m, n, p]` for a 3D tensor block, ...) when `nodes[k]` is one of `NODE_MOI_VARIABLE_BLOCK`, -`NODE_VARIABLE_BLOCK`, or `NODE_VALUE_BLOCK`. Block nodes are leaves that -stand in for an entire `prod(shape)`-element array of variables (or -constants), preserving contiguity end-to-end so the AD tape can be filled and -gathered with single contiguous bulk operations. +`NODE_VARIABLE_BLOCK`, `NODE_VALUE_BLOCK`, or `NODE_ARRAY_VALUE`. Block nodes +are leaves that stand in for an entire `prod(shape)`-element array of +variables (or constants), preserving contiguity end-to-end so the AD tape can +be filled and gathered with single contiguous bulk operations. + +`arrays` holds constant `AbstractArray`s that are *not* serialized on the +tape but kept by reference (see `NODE_ARRAY_VALUE`). Dense `Array` constants +are serialized in `values`; anything else (sparse or structured matrices, +GPU arrays, ...) lands here so operations like `LinearAlgebra.mul!` can +dispatch on the concrete array type. """ struct Expression{T} nodes::Vector{Node} values::Vector{T} + arrays::Vector{AbstractArray} block_shapes::Dict{Int,Vector{Int}} - Expression{T}() where {T} = new{T}(Node[], T[], Dict{Int,Vector{Int}}()) + function Expression{T}() where {T} + return new{T}(Node[], T[], AbstractArray[], Dict{Int,Vector{Int}}()) + end end function Base.:(==)(x::Expression, y::Expression) return x.nodes == y.nodes && x.values == y.values && + x.arrays == y.arrays && x.block_shapes == y.block_shapes end @@ -115,6 +126,7 @@ function _subexpression_and_linearity( nodes, adj, convert(Vector{eltype(S)}, expr.values), + expr.arrays, copy(expr.block_shapes), partials_storage_ϵ, linearity[1], diff --git a/test/JuMP.jl b/test/JuMP.jl index 8ab0620..29997fe 100644 --- a/test/JuMP.jl +++ b/test/JuMP.jl @@ -1294,6 +1294,38 @@ function test_matmul_both_variables_overwrites_reverse() return end +# `@constraint(model, expr in set)` over a vectorized array expression keeps the +# function whole as an `ArrayNonlinearFunction` (no scalarization). Relies on +# JuMP #3451 allowing `AbstractArray` in nonlinear expressions. +function test_vector_constraint() + n = 3 + model = Model() + @variable(model, x[1:n], container = ArrayDiff.ArrayOfVariables) + linear = x .- [1.0, 2.0, 3.0] + nonlinear = sin.(x) + c1 = @constraint(model, linear in MOI.Zeros(n)) + c2 = @constraint(model, nonlinear in MOI.Zeros(n)) + types = MOI.get(model, MOI.ListOfConstraintTypesPresent()) + @test (ArrayDiff.ArrayNonlinearFunction{1}, MOI.Zeros) in types + b = JuMP.backend(model) + F = ArrayDiff.ArrayNonlinearFunction{1} + @test MOI.get(b, MOI.NumberOfConstraints{F,MOI.Zeros}()) == 2 + ci = MOI.get(b, MOI.ListOfConstraintIndices{F,MOI.Zeros}()) + fs = [MOI.get(b, MOI.ConstraintFunction(), c) for c in ci] + # Each constraint kept its whole vectorized structure (not scalarized). + @test all(f isa F for f in fs) + @test all(f.size == (n,) for f in fs) + return +end + +# A vector set whose dimension disagrees with the expression is rejected. +function test_vector_constraint_dimension() + model = Model() + @variable(model, x[1:3], container = ArrayDiff.ArrayOfVariables) + @test_throws ErrorException @constraint(model, x .- 1.0 in MOI.Zeros(2)) + return +end + end # module TestJuMP.runtests() diff --git a/test/NLPModelsJuMP.jl b/test/NLPModelsJuMP.jl index 8e8dc65..e3dfd6e 100644 --- a/test/NLPModelsJuMP.jl +++ b/test/NLPModelsJuMP.jl @@ -7,6 +7,8 @@ using ArrayDiff import MathOptInterface as MOI import NLPModelsJuMP import JSOSolvers +import NLPModelsModifiers +import Percival function runtests() for name in names(@__MODULE__; all = true) @@ -60,6 +62,45 @@ function test_neural_tronls() return _test_neural_nlpmodels_jump(JSOSolvers.TronSolverNLS) end +# Constrained solve through the full pipeline: `@constraint(model, expr in +# MOI.Zeros(m))` over a vectorized array expression → NLPModelsJuMP collects it +# into a constrained `ArrayDiffNLPModel` (one residual evaluator per vector +# constraint) → Percival (augmented Lagrangian; LBFGS subproblems since +# ArrayDiff is first-order). +# +# Projection problem with an analytic solution: +# min ‖x − t‖² s.t. A x = b ⇒ x* = t − A' (A A')⁻¹ (A t − b) +function test_vector_constraint_solve() + if !isdefined(NLPModelsJuMP, :_try_array_nlp_model) + @info "NLPModelsJuMP has no `_try_array_nlp_model`; skipping the " * + "constrained test (needs the updated bl/arraydiff branch)." + return + end + model = Model(NLPModelsJuMP.Optimizer) + set_attribute( + model, + "solver", + nlp -> Percival.PercivalSolver( + nlp; + subproblem_modifier = NLPModelsModifiers.LBFGSModel, + ), + ) + set_attribute(model, "subproblem_modifier", NLPModelsModifiers.LBFGSModel) + set_attribute(model, MOI.AutomaticDifferentiationBackend(), ArrayDiff.Mode()) + set_silent(model) + @variable(model, x[1:3], container = ArrayDiff.ArrayOfVariables) + A = [1.0 1.0 1.0; 1.0 -1.0 0.0] + b = [1.0, 0.0] + t = [2.0, 1.0, 0.5] + @objective(model, Min, sum((x .- t) .^ 2)) + @constraint(model, A * x .- b in MOI.Zeros(2)) + optimize!(model) + xstar = t - A' * ((A * A') \ (A * t - b)) + @test isapprox(value.(x), xstar; atol = 1e-4) + @test isapprox(objective_value(model), sum((xstar .- t) .^ 2); atol = 1e-4) + return +end + end TestWithNLPModelsJuMP.runtests() diff --git a/test/Project.toml b/test/Project.toml index 263d24f..bdb3ca4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -12,9 +12,11 @@ MathOptAI = "e52c2cb8-508e-4e12-9dd2-9c4755b60e73" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" +NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd" ONNX = "d0dd6a25-fac6-55c0-abf7-829e0c774d20" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" +Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" @@ -30,3 +32,5 @@ NLopt = {rev = "bl/diff_backend", url = "https://github.com/jump-dev/NLopt.jl/"} [compat] NLPModels = "0.21.12" +# Registered Percival (<= 0.7.6) requires NLPModelsModifiers 0.7. +NLPModelsModifiers = "0.7" diff --git a/test/ReverseAD.jl b/test/ReverseAD.jl index eed7c10..8599592 100644 --- a/test/ReverseAD.jl +++ b/test/ReverseAD.jl @@ -566,6 +566,7 @@ function test_linearity() nodes, adj, convert(Vector{Float64}, expr.values), + expr.arrays, expr.block_shapes, Float64[], ret[1], diff --git a/test/StructuredConstants.jl b/test/StructuredConstants.jl new file mode 100644 index 0000000..e3738b9 --- /dev/null +++ b/test/StructuredConstants.jl @@ -0,0 +1,229 @@ +module TestStructuredConstants + +# Tests for `NODE_ARRAY_VALUE`: constant `AbstractArray`s that are not dense +# `Array`s are kept by reference in `Expression.arrays` instead of being +# serialized on the AD tape, and matmul nodes call `LinearAlgebra.mul!` +# directly on them so their specialized (sparse / structured) methods apply. + +using Test + +using JuMP +using ArrayDiff +import LinearAlgebra +import MathOptInterface as MOI +import SparseArrays + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$(name)", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end + return +end + +# A minimal custom matrix type: `y = x[idx]` as a linear operator. Each row +# has exactly one entry equal to 1, in column `idx[row]`. It only implements +# the two `mul!` methods ArrayDiff needs, so hitting any other code path +# (for example, an attempt to serialize it densely) would error. +struct SelectionMatrix{V<:AbstractVector{<:Integer}} <: AbstractMatrix{Float64} + idx::V + ncol::Int +end + +Base.size(A::SelectionMatrix) = (length(A.idx), A.ncol) + +# Dense copy for the reference computations of the tests. We deliberately +# don't implement `getindex`: evaluation must never index the matrix, so +# leaving it out proves that only the two `mul!` methods are used. +function _matrix(A::SelectionMatrix) + B = zeros(size(A)) + for (i, j) in enumerate(A.idx) + B[i, j] = 1.0 + end + return B +end + +function LinearAlgebra.mul!( + y::AbstractVector, + A::SelectionMatrix, + x::AbstractVector, +) + y .= view(x, A.idx) + return y +end + +function LinearAlgebra.mul!( + y::AbstractVector, + At::LinearAlgebra.Transpose{Float64,<:SelectionMatrix}, + w::AbstractVector, +) + A = parent(At) + fill!(y, 0.0) + for (i, j) in enumerate(A.idx) + y[j] += w[i] + end + return y +end + +function _gradient(model, obj, x) + mode = ArrayDiff.Mode{Vector{Float64}}() + ad = ArrayDiff.model(mode) + MOI.Nonlinear.set_objective(ad, JuMP.moi_function(obj)) + evaluator = MOI.Nonlinear.Evaluator( + ad, + mode, + JuMP.index.(JuMP.all_variables(model)), + ) + MOI.initialize(evaluator, [:Grad]) + val = MOI.eval_objective(evaluator, x) + g = zero(x) + MOI.eval_objective_gradient(evaluator, g, x) + return val, g +end + +function test_parse_sparse_constant_by_reference() + n = 4 + A = SparseArrays.sprand(n, n, 0.5) + LinearAlgebra.I + model = JuMP.Model() + @variable(model, x[1:n], container = ArrayDiff.ArrayOfVariables) + expr = A * x + f = JuMP.moi_function(sum(expr .^ 2)) + mode = ArrayDiff.Mode{Vector{Float64}}() + ad = ArrayDiff.model(mode) + MOI.Nonlinear.set_objective(ad, f) + # The sparse matrix must be stored by reference, not copied to `values`. + obj = ad.objective + @test length(obj.arrays) == 1 + @test obj.arrays[1] === A + @test any(node -> node.type == ArrayDiff.NODE_ARRAY_VALUE, obj.nodes) + return +end + +function test_sparse_matvec_lhs_gradient() + n = 5 + A = SparseArrays.sprand(n, n, 0.4) + LinearAlgebra.I + A_dense = Matrix(A) + x_val = collect(1.0:n) + model = JuMP.Model() + @variable(model, x[1:n], container = ArrayDiff.ArrayOfVariables) + obj_sparse = sum((A * x) .^ 2) + obj_dense = sum((A_dense * x) .^ 2) + val_s, g_s = _gradient(model, obj_sparse, x_val) + val_d, g_d = _gradient(model, obj_dense, x_val) + @test val_s ≈ val_d + @test g_s ≈ g_d + # Reference: ∇ sum((Ax).^2) = 2 A' A x + @test g_s ≈ 2 * A_dense' * (A_dense * x_val) + return +end + +function test_custom_selection_matrix_gradient() + n = 4 + idx = [2, 4, 1, 1, 3] + A = SelectionMatrix(idx, n) + A_dense = _matrix(A) + x_val = [0.5, -1.0, 2.0, 3.0] + c = [1.0, 2.0, 3.0, 4.0, 5.0] + model = JuMP.Model() + @variable(model, x[1:n], container = ArrayDiff.ArrayOfVariables) + obj_custom = sum(c .* (A * x) .^ 2) + obj_dense = sum(c .* (A_dense * x) .^ 2) + val_c, g_c = _gradient(model, obj_custom, x_val) + val_d, g_d = _gradient(model, obj_dense, x_val) + @test val_c ≈ val_d + @test g_c ≈ g_d + @test g_c ≈ 2 * A_dense' * (c .* (A_dense * x_val)) + return +end + +function test_sparse_matmat_rhs_gradient() + m, n = 3, 4 + B = SparseArrays.sprand(n, n, 0.5) + LinearAlgebra.I + B_dense = Matrix(B) + W_val = reshape(collect(1.0:(m*n)), m, n) + model = JuMP.Model() + @variable(model, W[1:m, 1:n], container = ArrayDiff.ArrayOfVariables) + obj_sparse = sum((W * B) .^ 2) + obj_dense = sum((W * B_dense) .^ 2) + x_val = vec(W_val) + val_s, g_s = _gradient(model, obj_sparse, x_val) + val_d, g_d = _gradient(model, obj_dense, x_val) + @test val_s ≈ val_d + @test g_s ≈ g_d + # Reference: ∇_W sum((WB).^2) = 2 (WB) B' + @test reshape(g_s, m, n) ≈ 2 * (W_val * B_dense) * B_dense' + return +end + +function test_sparse_residual_jtprod() + n, m = 6, 4 + A = SparseArrays.sprand(m, n, 0.5) + A_dense = Matrix(A) + b = collect(range(-1.0, 1.0; length = m)) + f_sparse = x -> A * x .+ b + f_dense = x -> A_dense * x .+ b + x_val = sin.(1:n) + v = cos.(1:m) + ev_s = ArrayDiff.evaluator(f_sparse, n) + ev_d = ArrayDiff.evaluator(f_dense, n) + F_s, F_d = zeros(m), zeros(m) + ArrayDiff.eval_residual!(ev_s, F_s, x_val) + ArrayDiff.eval_residual!(ev_d, F_d, x_val) + @test F_s ≈ F_d + @test F_s ≈ A_dense * x_val .+ b + Jtv_s, Jtv_d = zeros(n), zeros(n) + ArrayDiff.eval_residual_jtprod!(ev_s, Jtv_s, x_val, v) + ArrayDiff.eval_residual_jtprod!(ev_d, Jtv_d, x_val, v) + @test Jtv_s ≈ Jtv_d + @test Jtv_s ≈ A_dense' * v + return +end + +function test_sparse_with_nonlinear_chain() + # Mimics the AC-OPF polar structure: gather, sin/cos of differences, + # elementwise products, scatter back. + nbus, nbr = 4, 5 + from = [1, 1, 2, 3, 4] + to = [2, 3, 3, 4, 1] + F = SelectionMatrix(from, nbus) + T_ = SelectionMatrix(to, nbus) + F_dense, T_dense = _matrix(F), _matrix(T_) + c = [0.3, -0.5, 1.1, 0.7, -0.2] + build = (F1, T1) -> function (va) + d = F1 * va .- T1 * va + return c .* sin.(d) .+ cos.(d) + end + va_val = [0.0, 0.1, -0.2, 0.3] + v = collect(1.0:nbr) + ev_c = ArrayDiff.evaluator(build(F, T_), nbus) + ev_d = ArrayDiff.evaluator(build(F_dense, T_dense), nbus) + F_c, F_d = zeros(nbr), zeros(nbr) + ArrayDiff.eval_residual!(ev_c, F_c, va_val) + ArrayDiff.eval_residual!(ev_d, F_d, va_val) + d_val = va_val[from] .- va_val[to] + @test F_c ≈ F_d + @test F_c ≈ c .* sin.(d_val) .+ cos.(d_val) + Jtv_c, Jtv_d = zeros(nbus), zeros(nbus) + ArrayDiff.eval_residual_jtprod!(ev_c, Jtv_c, va_val, v) + ArrayDiff.eval_residual_jtprod!(ev_d, Jtv_d, va_val, v) + @test Jtv_c ≈ Jtv_d + # Finite-difference check of J'v. + h = 1e-7 + Jtv_fd = map(eachindex(va_val)) do i + va_p = copy(va_val) + va_p[i] += h + va_m = copy(va_val) + va_m[i] -= h + fp = va -> c .* sin.(va[from] .- va[to]) .+ cos.(va[from] .- va[to]) + return LinearAlgebra.dot(v, (fp(va_p) .- fp(va_m)) ./ (2h)) + end + @test Jtv_c ≈ Jtv_fd atol = 1e-6 + return +end + +end + +TestStructuredConstants.runtests() diff --git a/test/runtests.jl b/test/runtests.jl index 5cc99fa..3975a5d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,7 @@ include("ReverseAD.jl") include("ArrayDiff.jl") include("JuMP.jl") +include("StructuredConstants.jl") include("MathOptAI.jl") include("ONNXExt.jl") if VERSION >= v"1.11"