diff --git a/src/Evaluators/reduced_evaluator.jl b/src/Evaluators/reduced_evaluator.jl index 3019654d..5543c866 100644 --- a/src/Evaluators/reduced_evaluator.jl +++ b/src/Evaluators/reduced_evaluator.jl @@ -86,7 +86,6 @@ mutable struct ReducedSpaceEvaluator{T, VI, VT, MT, Jacx, Jacu, JacCons, Hess} < # Options linear_solver::LinearSolvers.AbstractLinearSolver powerflow_solver::AbstractNonLinearSolver - factorization::Union{Nothing, Factorization} has_jacobian::Bool update_jacobian::Bool has_hessian::Bool @@ -95,7 +94,7 @@ end function ReducedSpaceEvaluator( model, x, u; constraints=Function[voltage_magnitude_constraints, active_power_constraints, reactive_power_constraints], - linear_solver=DirectSolver(), + linear_solver=direct_linear_solver(model), powerflow_solver=NewtonRaphson(tol=1e-12), want_jacobian=true, want_hessian=true, @@ -143,7 +142,7 @@ function ReducedSpaceEvaluator( constraints, g_min, g_max, buffer, state_ad, obj_ad, cons_ad, cons_jac, hess_ad, - linear_solver, powerflow_solver, nothing, want_jacobian, false, want_hessian, + linear_solver, powerflow_solver, want_jacobian, false, want_hessian, ) end function ReducedSpaceEvaluator( @@ -222,18 +221,6 @@ end bounds(nlp::ReducedSpaceEvaluator, ::Variables) = (nlp.u_min, nlp.u_max) bounds(nlp::ReducedSpaceEvaluator, ::Constraints) = (nlp.g_min, nlp.g_max) -function _update_factorization!(nlp::ReducedSpaceEvaluator) - ∇gₓ = nlp.state_jacobian.x.J - if !isa(∇gₓ, SparseMatrixCSC) || isa(nlp.linear_solver, LinearSolvers.AbstractIterativeLinearSolver) - return - end - if !isnothing(nlp.factorization) - lu!(nlp.factorization, ∇gₓ) - else - nlp.factorization = lu(∇gₓ) - end -end - ## Callbacks function update!(nlp::ReducedSpaceEvaluator, u) jac_x = nlp.state_jacobian.x @@ -255,8 +242,6 @@ function update!(nlp::ReducedSpaceEvaluator, u) # Refresh values of active and reactive powers at generators update!(nlp.model, PS.Generators(), PS.ActivePower(), nlp.buffer) - # Update factorization if direct solver - _update_factorization!(nlp) # Evaluate Jacobian of power flow equation on current u AutoDiff.jacobian!(nlp.model, nlp.state_jacobian.u, nlp.buffer) # Specify that constraint's Jacobian is not up to date @@ -289,8 +274,7 @@ function _forward_solve!(nlp::ReducedSpaceEvaluator, y, x) ∇gₓ = nlp.state_jacobian.x.J LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gₓ, x) else - ∇gₓ = nlp.factorization - ldiv!(y, ∇gₓ, x) + LinearSolvers.ldiv!(nlp.linear_solver, y, x) end end @@ -307,8 +291,7 @@ function _backward_solve!(nlp::ReducedSpaceEvaluator, y::VT, x::VT) where {VT <: ∇gT = LinearSolvers.get_transpose(nlp.linear_solver, ∇gₓ) LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gT, x) else - ∇gf = nlp.factorization - LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gf', x) + LinearSolvers.rdiv!(nlp.linear_solver, y, x) end end @@ -394,11 +377,9 @@ function jacobian!(nlp::ReducedSpaceEvaluator, jac, u) m, nᵤ = size(Ju) ∇gᵤ = nlp.state_jacobian.u.J ∇gₓ = nlp.state_jacobian.x.J - # Compute factorization with UMFPACK - ∇gfac = nlp.factorization - # Compute adjoints all in once, using the same factorization + # Compute state sensitivities all in once μ = zeros(nₓ, nᵤ) - ldiv!(μ, ∇gfac, ∇gᵤ) + LinearSolvers.ldiv!(nlp.linear_solver, μ, ∇gᵤ) # Compute reduced Jacobian copy!(jac, Ju) mul!(jac, Jx, μ, -1.0, 1.0) @@ -447,10 +428,11 @@ end function jtprod!(nlp::ReducedSpaceEvaluator, jv, u, v) ∂obj = nlp.obj_stack + μ = nlp.buffer.balance jvx = ∂obj.stack.jvₓ ; fill!(jvx, 0) jvu = ∂obj.stack.jvᵤ ; fill!(jvu, 0) full_jtprod!(nlp, jvx, jvu, u, v) - reduced_gradient!(nlp, jv, jvx, jvu, jvx, u) + reduced_gradient!(nlp, jv, jvx, jvu, μ, u) end function ojtprod!(nlp::ReducedSpaceEvaluator, jv, u, σ, v) diff --git a/src/LinearSolvers/LinearSolvers.jl b/src/LinearSolvers/LinearSolvers.jl index 01495c18..ce673ca7 100644 --- a/src/LinearSolvers/LinearSolvers.jl +++ b/src/LinearSolvers/LinearSolvers.jl @@ -54,6 +54,7 @@ get_transpose(::AbstractLinearSolver, M::AbstractMatrix) = transpose(M) """ ldiv!(solver, x, A, y) + ldiv!(solver, x, y) * `solver::AbstractLinearSolver`: linear solver to solve the system * `x::AbstractVector`: Solution @@ -61,11 +62,28 @@ get_transpose(::AbstractLinearSolver, M::AbstractMatrix) = transpose(M) * `y::AbstractVector`: RHS Solve the linear system ``A x = y`` using the algorithm -specified in `solver`. +specified in `solver`. If `A` is not specified, the function +will used directly the factorization stored inplace. """ function ldiv! end +""" + rdiv!(solver, x, A, y) + rdiv!(solver, x, y) + +* `solver::AbstractLinearSolver`: linear solver to solve the system +* `x::AbstractVector`: Solution +* `A::AbstractMatrix`: Input matrix +* `y::AbstractVector`: RHS + +Solve the linear system ``A^⊤ x = y`` using the algorithm +specified in `solver`. If `A` is not specified, the function +will used directly the factorization stored inplace. + +""" +function rdiv! end + """ DirectSolver <: AbstractLinearSolver @@ -75,27 +93,43 @@ Solve linear system ``A x = y`` with direct linear algebra. * On CUDA GPU, `DirectSolver` redirects the resolution to the method `CUSOLVER.csrlsvqr` """ -struct DirectSolver <: AbstractLinearSolver end -DirectSolver(precond) = DirectSolver() -function ldiv!(::DirectSolver, - y::Vector, J::AbstractMatrix, x::Vector, -) - y .= J \ x +struct DirectSolver{Fac<:Union{Nothing, LinearAlgebra.Factorization}} <: AbstractLinearSolver + factorization::Fac +end + +DirectSolver(J::SparseMatrixCSC) = DirectSolver(lu(J)) +DirectSolver() = DirectSolver(nothing) +DirectSolver(precond::AbstractPreconditioner) = DirectSolver(nothing) + +# Reuse factorization in update +function ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::Vector, J::AbstractMatrix, x::Vector) + lu!(s.factorization, J) # Update factorization inplace + LinearAlgebra.ldiv!(y, s.factorization, x) # Forward-backward solve return 0 end -function ldiv!(::DirectSolver, - y::Vector, J::Factorization, x::Vector, -) - LinearAlgebra.ldiv!(y, J, x) +# Solve system Ax = y +function ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::AbstractArray, x::AbstractArray) + LinearAlgebra.ldiv!(y, s.factorization, x) # Forward-backward solve + return 0 +end +# Solve system A'x = y +function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::Array, x::Array) + LinearAlgebra.ldiv!(y, s.factorization', x) # Forward-backward solve return 0 end -function ldiv!(::DirectSolver, + +function ldiv!(::DirectSolver{Nothing}, y::Vector, J::AbstractMatrix, x::Vector) + y .= J \ x + return 0 +end + +function ldiv!(::DirectSolver{Nothing}, y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSR, x::CUDA.CuVector, ) CUSOLVER.csrlsvqr!(J, x, y, 1e-8, one(Cint), 'O') return 0 end -function ldiv!(::DirectSolver, +function ldiv!(::DirectSolver{Nothing}, y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSC, x::CUDA.CuVector, ) csclsvqr!(J, x, y, 1e-8, one(Cint), 'O') diff --git a/src/Polar/polar.jl b/src/Polar/polar.jl index b9a7d181..e3bb5c4a 100644 --- a/src/Polar/polar.jl +++ b/src/Polar/polar.jl @@ -268,6 +268,28 @@ function init_buffer!(form::PolarForm{T, IT, VT, MT}, buffer::PolarNetworkState) copyto!(buffer.qinj, qbus) end +function direct_linear_solver(polar::PolarForm) + is_cpu = isa(polar.device, KA.CPU) + if is_cpu + jac = jacobian_sparsity(polar, power_balance, State()) + return LinearSolvers.DirectSolver(jac) + else + # Factorization is not yet supported on the GPU + return LinearSolvers.DirectSolver(nothing) + end +end + +function build_preconditioner(polar::PolarForm; nblocks=-1) + jac = jacobian_sparsity(polar, power_balance, State()) + n = size(jac, 1) + npartitions = if nblocks > 0 + nblocks + else + div(n, 32) + end + return LinearSolvers.BlockJacobiPreconditioner(jac, npartitions, polar.device) +end + function Base.show(io::IO, polar::PolarForm) # Network characteristics nbus = PS.get(polar.network, PS.NumberOfBuses()) diff --git a/test/Evaluators/powerflow.jl b/test/Evaluators/powerflow.jl index 38047ac9..807849c5 100644 --- a/test/Evaluators/powerflow.jl +++ b/test/Evaluators/powerflow.jl @@ -4,6 +4,8 @@ using KernelAbstractions using Test import ExaPF: PowerSystem +import ExaPF: LinearSolvers +const LS = LinearSolvers @testset "Powerflow solver" begin datafile = joinpath(INSTANCES_DIR, "case14.raw") @@ -19,13 +21,13 @@ import ExaPF: PowerSystem @testset "Deport computation on device $device" for device in DEVICES polar = PolarForm(pf, device) - jac = ExaPF.jacobian_sparsity(polar, ExaPF.power_balance, State()) - precond = ExaPF.LinearSolvers.BlockJacobiPreconditioner(jac, npartitions, device) + precond = ExaPF.build_preconditioner(polar; nblocks=npartitions) # Retrieve initial state of network x0 = ExaPF.initial(polar, State()) uk = ExaPF.initial(polar, Control()) @testset "Powerflow solver $(LinSolver)" for LinSolver in ExaPF.list_solvers(device) + (LinSolver == LS.DirectSolver) && continue algo = LinSolver(precond) xk = copy(x0) nlp = ExaPF.ReducedSpaceEvaluator(