Add support for non-dense matrices

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

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"

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.

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

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