use package extension for Argos

exanauts · Jan 25, 2024 · e3a73e9 · e3a73e9
1 parent 02989e3
commit e3a73e9
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 32 deletions.
diff --git a/ext/api.jl b/ext/api.jl
@@ -7,44 +7,40 @@ end
 function Argos.run_opf_gpu(datafile::String, ::Argos.FullSpace; options...)
     flp = Argos.FullSpaceEvaluator(datafile; device=CUDABackend())
     model = Argos.OPFModel(Argos.bridge(flp))
-    ips = MadNLP.MadNLPSolver(
+    solver = MadNLP.MadNLPSolver(
         model;
+        kkt_system=MadNLP.SparseKKTSystem,
+        callback=MadNLP.SparseCallback,
         options...
     )
-    MadNLP.solve!(ips)
-    return ips
+    MadNLP.solve!(solver)
+    return solver
 end
 
 function Argos.run_opf_gpu(datafile::String, ::Argos.BieglerReduction; options...)
     flp = Argos.FullSpaceEvaluator(datafile; device=CUDABackend())
     model = Argos.OPFModel(Argos.bridge(flp))
-
-    madnlp_options = Dict{Symbol, Any}(options...)
-    # madnlp_options[:linear_solver] = LapackGPUSolver
-    opt_ipm, opt_linear, logger = MadNLP.load_options(; madnlp_options...)
-
     KKT = Argos.BieglerKKTSystem{Float64, CuVector{Int}, CuVector{Float64}, CuMatrix{Float64}}
-    ips = MadNLP.MadNLPSolver{Float64, KKT}(model, opt_ipm, opt_linear; logger=logger)
-    MadNLP.solve!(ips)
-    return ips
+    solver = MadNLP.MadNLPSolver(
+        model;
+        kkt_system=KKT,
+        callback=MadNLP.SparseCallback,
+        options...
+    )
+
+    MadNLP.solve!(solver)
+    return solver
 end
 
 function Argos.run_opf_gpu(datafile::String, ::Argos.DommelTinney; options...)
     flp = Argos.ReducedSpaceEvaluator(datafile; device=CUDABackend(), nbatch_hessian=256)
-    model = Argos.OPFModel(Argos.bridge(flp))
-
-    madnlp_options = Dict{Symbol, Any}(options...)
-    # madnlp_options[:linear_solver] = LapackGPUSolver
-    madnlp_options[:kkt_system] = MadNLP.DENSE_CONDENSED_KKT_SYSTEM
-    # madnlp_options[:inertia_correction_method] = MadNLP.INERTIA_FREE
-    madnlp_options[:lapack_algorithm] = MadNLP.CHOLESKY
-
-    opt_ipm, opt_linear, logger = MadNLP.load_options(; madnlp_options...)
-
-    QN = MadNLP.ExactHessian{Float64, CuVector{Float64}}
-    KKT = MadNLP.DenseCondensedKKTSystem{Float64, CuVector{Float64}, CuMatrix{Float64}, QN}
-    ips = MadNLP.MadNLPSolver{Float64, KKT}(model, opt_ipm, opt_linear; logger=logger)
-    MadNLP.solve!(ips)
-
-    return ips
+    model = Argos.OPFModel(flp)
+    solver = MadNLP.MadNLPSolver(
+        model;
+        kkt_system=MadNLP.DenseCondensedKKTSystem,
+        callback=MadNLP.DenseCallback,
+        options...
+    )
+    MadNLP.solve!(solver)
+    return solver
 end
diff --git a/test/Algorithms/MadNLP_wrapper.jl b/test/Algorithms/MadNLP_wrapper.jl
@@ -182,10 +182,10 @@ if CUDA.has_cuda_gpu()
         _test_results_match(solver_ref, solver_gpu; atol=1e-6)
     end
 
-    @testset "[CUDA] Solve OPF with $form" for (form, linear_solver_algo) in [
-        (Argos.FullSpace(), MadNLP.BUNCHKAUFMAN),
-        (Argos.BieglerReduction(), MadNLP.CHOLESKY),
-        (Argos.DommelTinney(), MadNLP.CHOLESKY),
+    @testset "[CUDA] Solve OPF with $form" for (form, linear_solver, linear_solver_algo) in [
+        (Argos.FullSpace(), LapackCPUSolver, MadNLP.BUNCHKAUFMAN),
+        (Argos.BieglerReduction(), LapackGPUSolver, MadNLP.CHOLESKY),
+        (Argos.DommelTinney(), LapackGPUSolver, MadNLP.CHOLESKY),
     ]
         case = "case9.m"
         datafile = joinpath(INSTANCES_DIR, case)
@@ -194,9 +194,10 @@ if CUDA.has_cuda_gpu()
             datafile,
             form;
             tol=1e-5,
-            linear_solver=LapackGPUSolver,
+            linear_solver=linear_solver,
             lapack_algorithm=linear_solver_algo,
             print_level=MadNLP.ERROR,
+            max_iter=100,
         )
         @test isa(solver, MadNLP.MadNLPSolver)
         @test solver.status == MadNLP.SOLVE_SUCCEEDED

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -20,6 +20,7 @@ const CASES = ["case9.m", "case30.m"]
 ARCHS = Any[(CPU(), Array, SparseMatrixCSC)]
 
 if has_cuda_gpu()
+    CUDA.allowscalar(false)
     CUDA_ARCH = (CUDABackend(), CuArray, nothing)
     push!(ARCHS, CUDA_ARCH)
     CUDA.versioninfo()