LM performance improvements

.gitignore

@@ -38,6 +38,7 @@ docs/site/
 *.ipynb
 
 Manifest.toml
+scripts/benchmarks/*.csv
 .vscode
 *.h5
 data

scripts/benchmarks/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+PowerFlows = "94fada2c-fd9a-4e89-8d82-81405f5cb4f6"
+PowerSystemCaseBuilder = "f00506e0-b84f-492a-93c2-c0a9afc4364e"
+PowerSystems = "bcd98974-b02a-5e2f-9ee0-a103f5c450dd"

src/RobustHomotopy/robust_homotopy_method.jl

@@ -18,17 +18,25 @@ function _newton_power_flow(pf::ACPolarPowerFlow{<:RobustHomotopyPowerFlow},
     symbolic_factor!(hSolver, homHess.Hv)
 
     success = true
+    total_iters = 0
     while true # go onto next t_k even if search doesn't terminate within max iterations.
-        converged_t_k, _ = _second_order_newton(homHess, t_k, time_step, x, hSolver)
+        converged_t_k, iters = _second_order_newton(homHess, t_k, time_step, x, hSolver)
+        total_iters += iters
         if t_k == 1.0
             success = converged_t_k
             break
         end
         t_k = min(t_k + Δt_k, 1.0)
     end
+    r_L2 = norm(homHess.pfResidual.Rv, 2)
+    r_Linf = norm(homHess.pfResidual.Rv, Inf)
+    @info("Final residual size: $(r_L2) L2, $(r_Linf) L∞.")
     if !success
-        @warn "RobustHomotopyPowerFlow failed to find a solution"
+        @error(
+            "The RobustHomotopyPowerFlow solver failed to converge after $total_iters iterations."
+        )
     else
+        @info("The RobustHomotopyPowerFlow solver converged after $total_iters iterations.")
         if get_calculate_loss_factors(data)
             _calculate_loss_factors(data, homHess.J.Jv, time_step)
         end

src/definitions.jl

@@ -39,8 +39,9 @@ const DEFAULT_IWAMOTO_FALLBACK = true # when a trust region step is rejected, tr
 const PF_MAX_LOG = 10
 # only used for Levenberg-Maquardt
 const DEFAULT_λ_0 = 1e-5
-# input is mix of powers (100 MW), voltages (0.8-1.2), and angles (-π/4 to π/4).
-const DEFAULT_MAX_TEST_λs = 50 # give up after increasing damping factor 50 times.
+# Upper bound on the LM damping factor μ. μ only grows on rejected steps, so
+# hitting this cap is a divergence signal — the solver aborts with an error.
+const DEFAULT_μ_MAX = 1e8
 
 const DEFAULT_Δt_k = 0.2
 

src/levenberg-marquardt.jl

@@ -1,157 +1,232 @@
-"""Driver for the LevenbergMarquardtACPowerFlow method: sets up the data 
-structures (e.g. residual), runs the power flow method via calling `_run_power_flow_method` 
+"""Pre-allocated workspace for the Levenberg-Marquardt solver.
+
+Holds the augmented matrix `[J; √λ·I]` with a fixed sparsity pattern,
+a mapping to update its entries in-place, and a cached QR factorization."""
+mutable struct LMWorkspace
+    A::SparseMatrixCSC{Float64, Int64}
+    # PERF: if we're doing dense access, I'd expect a bitmask to be faster.
+    #       Or only store λ_diag_indices (sparse) and infer J entries as the complement.
+    # Indices into A.nzval for the J block entries (same order as J.Jv.nzval)
+    j_nzval_indices::Vector{Int}
+    # Indices into A.nzval for the √λ diagonal entries (length n)
+    λ_diag_indices::Vector{Int}
+    # PERF: have not investigated other factorization methods. SPQR does not allow 
+    #       for in-place updates or re-use of symbolic factorization, but with non
+    #       square matrices--J^T J form is more unstable--we don't have a lot of options.
+    # Cached QR factorization
+    F::SparseArrays.SPQR.QRSparse{Float64, Int64}
+    # Preallocated augmented RHS [-Rv; 0] (length m + n); bottom n stay zero.
+    b::Vector{Float64}
+end
+
+"""Build the augmented matrix `[J; I]` once, recording which `A.nzval` entries
+correspond to J values vs the λ diagonal."""
+function LMWorkspace(Jv::SparseMatrixCSC{Float64, J_INDEX_TYPE})
+    m, n = size(Jv)
+
+    # Convert J to Int64 indices for SPQR compatibility, then vcat with identity.
+    Jv64 = SparseMatrixCSC{Float64, Int64}(
+        Jv.m, Jv.n,
+        Vector{Int64}(Jv.colptr),
+        Vector{Int64}(Jv.rowval),
+        copy(Jv.nzval),
+    )
+    Iλ = sparse(Int64.(1:n), Int64.(1:n), ones(n), n, n)
+    A = vcat(Jv64, Iλ)
+
+    # Identify which A.nzval entries come from J vs the diagonal.
+    j_nzval_indices = Vector{Int}(undef, length(Jv.nzval))
+    λ_diag_indices = Vector{Int}(undef, n)
+
+    j_idx = 0
+    for col in 1:n
+        for a_idx in A.colptr[col]:(A.colptr[col + 1] - 1)
+            row = A.rowval[a_idx]
+            if row <= m
+                j_idx += 1
+                j_nzval_indices[j_idx] = a_idx
+            elseif row == m + col
+                λ_diag_indices[col] = a_idx
+            end
+        end
+    end
+    @assert j_idx == length(Jv.nzval) "Expected $(length(Jv.nzval)) J entries, found $j_idx"
+
+    b = zeros(m + n)
+    F = LinearAlgebra.qr(A)
+
+    return LMWorkspace(A, j_nzval_indices, λ_diag_indices, F, b)
+end
+
+"""Copy current Jacobian values into the augmented matrix."""
+function copy_jacobian!(ws::LMWorkspace, Jv::SparseMatrixCSC{Float64, J_INDEX_TYPE})
+    nzv = Jv.nzval
+    for (i, a_idx) in enumerate(ws.j_nzval_indices)
+        ws.A.nzval[a_idx] = nzv[i]
+    end
+    return
+end
+
+"""Update the √λ diagonal and re-factorize."""
+function update_lambda!(ws::LMWorkspace, λ::Float64)
+    sqrtλ = sqrt(λ)
+    for a_idx in ws.λ_diag_indices
+        ws.A.nzval[a_idx] = sqrtλ
+    end
+    ws.F = LinearAlgebra.qr(ws.A)
+    return
+end
+
+"""Driver for the LevenbergMarquardtACPowerFlow method: sets up the data
+structures (e.g. residual), runs the power flow method via calling `_run_power_flow_method`
 on them, then handles post-processing (e.g. loss factors)."""
 function _newton_power_flow(
-    pf::ACPolarPowerFlow{LevenbergMarquardtACPowerFlow},
+    pf::AbstractACPowerFlow{LevenbergMarquardtACPowerFlow},
     data::ACPowerFlowData,
     time_step::Int64;
     tol::Float64 = DEFAULT_NR_TOL,
     maxIterations::Int = DEFAULT_NR_MAX_ITER,
     validate_voltage_magnitudes::Bool = DEFAULT_VALIDATE_VOLTAGES,
     vm_validation_range::MinMax = DEFAULT_VALIDATION_RANGE,
     λ_0::Float64 = DEFAULT_λ_0,
-    maxTestλs::Int = DEFAULT_MAX_TEST_λs,
+    x0::Union{Vector{Float64}, Nothing} = nothing,
     _ignored...,
 )
+    init_kwargs = if isnothing(x0)
+        (; validate_voltage_magnitudes, vm_validation_range)
+    else
+        (; validate_voltage_magnitudes, vm_validation_range, x0)
+    end
     residual, J, x0 = initialize_power_flow_variables(
-        pf, data, time_step;
-        validate_voltage_magnitudes, vm_validation_range)
+        pf, data, time_step; init_kwargs...)
     converged = norm(residual.Rv, Inf) < tol
     i = 0
     if !converged
+        ws = LMWorkspace(J.Jv)
         converged, i = _run_power_flow_method(
             time_step,
             x0,
             residual,
-            J;
-            tol, maxIterations, λ_0, maxTestλs,
+            J,
+            ws;
+            tol, maxIterations, λ_0,
         )
     end
+    # x0 was mutated in place to the converged state by _run_power_flow_method
+    # (or is the already-converged initial state if the loop was skipped).
+    _finalize_formulation!(pf, data, x0, residual, time_step)
     return _finalize_power_flow(
         converged, i, "LevenbergMarquardtACPowerFlow", residual, data, J.Jv, time_step)
 end
 
 function _run_power_flow_method(
     time_step::Int,
     x::Vector{Float64},
-    residual::ACPowerFlowResidual,
-    J::ACPowerFlowJacobian;
+    residual::Union{ACPowerFlowResidual, ACRectangularCIResidual},
+    J::Union{ACPowerFlowJacobian, ACRectangularCIJacobian},
+    ws::LMWorkspace;
     maxIterations::Int = DEFAULT_NR_MAX_ITER,
     tol::Float64 = DEFAULT_NR_TOL,
     λ_0::Float64 = DEFAULT_λ_0,
-    maxTestλs::Int = DEFAULT_MAX_TEST_λs,
     _ignored...,
 )
-    λ::Float64 = λ_0
+    μ::Float64 = λ_0
+    λ::Float64 = 0.0
     i, converged = 0, false
     residual(x, time_step)
     resSize = dot(residual.Rv, residual.Rv)
     linf = norm(residual.Rv, Inf)
     @debug "initially: sum of squares $(siground(resSize)), L ∞ norm $(siground(linf)), λ = $λ"
-    while i < maxIterations && !converged && !isnan(λ)
-        λ = update_damping_factor!(x, residual, J, time_step, maxTestλs)
-        converged = !isnan(λ) && norm(residual.Rv, Inf) < tol
+    while i < maxIterations && !converged && isfinite(λ) && μ < DEFAULT_μ_MAX
+        λ, μ = update_damping_factor!(x, residual, J, μ, time_step, ws)
+        converged = isfinite(λ) && norm(residual.Rv, Inf) < tol
         i += 1
     end
-    if isnan(λ)
-        @error "λ is NaN"
-    elseif i == maxIterations
-        @error "The LevenbergMarquardtACPowerFlow solver didn't coverge in $maxIterations iterations."
+    if !converged
+        if !isfinite(λ)
+            @error "λ is not finite ($(λ))"
+        elseif μ >= DEFAULT_μ_MAX
+            @error "The LevenbergMarquardtACPowerFlow damping factor μ hit the cap (DEFAULT_μ_MAX=$(DEFAULT_μ_MAX)) after $i iterations; aborting (likely divergence)."
+        elseif i == maxIterations
+            @error "The LevenbergMarquardtACPowerFlow solver didn't coverge in $maxIterations iterations."
+        end
     end
 
     return converged, i
 end
 
-# LM implementation and parameters values largely based upon this paper:
-# # https://www.sciencedirect.com/science/article/pii/S0142061518336342
+# LM implementation based on standard Levenberg-Marquardt method.
+# See Nocedal & Wright (2006), sections 10.3 and 11.2.
+
+"""Compute one LM trial step. Assumes `residual` and `J` are already evaluated
+at `x` by the caller. Returns the gain ratio ρ."""
 function compute_error(
     x::Vector{Float64},
-    residual::ACPowerFlowResidual,
-    J::ACPowerFlowJacobian,
+    residual::Union{ACPowerFlowResidual, ACRectangularCIResidual},
+    J::Union{ACPowerFlowJacobian, ACRectangularCIJacobian},
     λ::Float64,
     time_step::Int,
     residualSize::Float64,
+    ws::LMWorkspace,
 )
-    residual(x, time_step) # M(x_c)
-    J(time_step)
-
-    n = size(J.Jv, 2)
-    Iλ = sparse(1:n, 1:n, sqrt(λ) .* ones(n), n, n) # less error-prone compared to A = vcat(J.Jv, sqrt(λ) * sparse(LinearAlgebra.I, size(J.Jv)))
-    A = [J.Jv; Iλ]
+    # Update augmented matrix with current J values and λ, then factorize once.
+    copy_jacobian!(ws, J.Jv)
+    update_lambda!(ws, λ)
 
-    b_x = vcat(-residual.Rv, zeros(size(J.Jv, 2)))
-    Δx = A \ b_x
+    m = length(residual.Rv)
+    @assert m == length(ws.b) - size(J.Jv, 2) "residual/J size mismatch vs preallocated LM buffer (m=$m, buf=$(length(ws.b)), n=$(size(J.Jv, 2)))"
+    @views ws.b[1:m] .= .-residual.Rv   # bottom n entries stay zero from construction
+    Δx = ws.F \ ws.b
 
     temp_x = residual.Rv .+ J.Jv * Δx
 
     x_trial = x .+ Δx
-    residual(x_trial, time_step) # M(y_c)
-
-    b_y = vcat(-residual.Rv, zeros(size(J.Jv, 2)))
-    Δy = A \ b_y
-    temp_y = residual.Rv .+ J.Jv * Δy
-
-    newResidualSize_y = dot(residual.Rv, residual.Rv)
-    residual(x_trial .+ Δy, time_step) # M(x_c+Δx+Δy)
+    residual(x_trial, time_step) # M(x_c + Δx)
     newResidualSize = dot(residual.Rv, residual.Rv)
 
-    b_z = vcat(-residual.Rv, zeros(size(J.Jv, 2)))
-    Δz = A \ b_z
-    temp_z = residual.Rv .+ J.Jv * Δz
-    newResidualSize_z = dot(residual.Rv, residual.Rv)
+    predicted_reduction = residualSize - dot(temp_x, temp_x)
+    actual_reduction = residualSize - newResidualSize
 
-    residual(x_trial .+ Δy .+ Δz, time_step) # M(x_c+Δx+Δy+Δz)
-    newResidualSize = dot(residual.Rv, residual.Rv)
+    # Guard against zero/negative predicted reduction.
+    if predicted_reduction <= 0.0 || !isfinite(predicted_reduction)
+        residual(x, time_step)
+        return 0.0
+    end
 
-    predicted_reduction = (
-        residualSize - dot(temp_x, temp_x) + newResidualSize_y - dot(temp_y, temp_y) +
-        newResidualSize_z - dot(temp_z, temp_z)
-    )
-    actual_reduction = (residualSize - newResidualSize)
     ρ = actual_reduction / predicted_reduction
 
     if ρ > 1e-4
-        x .+= (Δx .+ Δy .+ Δz)
+        x .+= Δx
+    else
+        # Bad step: restore data state to match x (not x_trial).
+        residual(x, time_step)
     end
 
     return ρ
 end
 
 function update_damping_factor!(
     x::Vector{Float64},
-    residual::ACPowerFlowResidual,
-    J::ACPowerFlowJacobian,
+    residual::Union{ACPowerFlowResidual, ACRectangularCIResidual},
+    J::Union{ACPowerFlowJacobian, ACRectangularCIJacobian},
+    μ::Float64,
     time_step::Int,
-    maxTestλs::Int,
+    ws::LMWorkspace,
 )
     residual(x, time_step)
     residualSize = dot(residual.Rv, residual.Rv)
     J(time_step)
 
-    # Now update \lambda
-    test_lambda::Int = 1
-    # https://www.sciencedirect.com/science/article/pii/S0142061518336342
-    λ = DEFAULT_λ_0 * sqrt(residualSize)
-    while test_lambda < maxTestλs
-        ρ = compute_error(x, residual, J, λ, time_step, residualSize)
-        if ρ > 0.75 # good step
-            factor = 4.0
-            λ_temp = λ / factor
-            λ = max(λ_temp, 1.0e-8)
-            break
-        elseif ρ >= 0.25 # okay step
-            break
-        else # bad step
-            factor = 4.0
-            λ *= factor
-        end
-
-        test_lambda += 1
-    end
-    if test_lambda == maxTestλs
-        @error "Unable to improve: gave up after increasing damping factor $maxTestλs times."
-        λ = NaN
+    λ = μ * sqrt(residualSize)
+    ρ = compute_error(x, residual, J, λ, time_step, residualSize, ws)
+    coef = 4.0
+    if ρ > 0.75
+        μ = max(μ / coef, 1e-8)
+    elseif ρ >= 0.25
+        # intentional no-op
+    else
+        μ = min(μ * coef, DEFAULT_μ_MAX)
     end
 
-    return λ
+    return (λ, μ)
 end

src/power_flow_method.jl

@@ -725,7 +725,7 @@ function _finalize_power_flow(
         end
         return true
     end
-    @error("The $solver_name solver failed to converge.")
+    @error("The $solver_name solver failed to converge after $i iterations.")
     return false
 end