From 536d99bada6ac6fb58da96b3dbd0bc60aeaf8cea Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 17 Mar 2025 13:50:01 -0400 Subject: [PATCH 1/4] chore: add sources --- Project.toml | 6 ++++++ lib/BracketingNonlinearSolve/Project.toml | 3 +++ lib/NonlinearSolveBase/Project.toml | 3 +++ lib/NonlinearSolveFirstOrder/Project.toml | 4 ++++ lib/NonlinearSolveHomotopyContinuation/Project.toml | 3 +++ lib/NonlinearSolveQuasiNewton/Project.toml | 3 +++ lib/NonlinearSolveSpectralMethods/Project.toml | 3 +++ lib/SCCNonlinearSolve/Project.toml | 3 +++ lib/SimpleNonlinearSolve/Project.toml | 3 +++ 9 files changed, 31 insertions(+) diff --git a/Project.toml b/Project.toml index 5c51b73d0..4d46dc060 100644 --- a/Project.toml +++ b/Project.toml @@ -153,3 +153,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"] + +[sources] +NonlinearSolveBase = {path = "lib/NonlinearSolveBase"} +NonlinearSolveFirstOrder = {path = "lib/NonlinearSolveFirstOrder"} +NonlinearSolveQuasiNewton = {path = "lib/NonlinearSolveQuasiNewton"} +NonlinearSolveSpectralMethods = {path = "lib/NonlinearSolveSpectralMethods"} diff --git a/lib/BracketingNonlinearSolve/Project.toml b/lib/BracketingNonlinearSolve/Project.toml index 16cade168..8da9d6e97 100644 --- a/lib/BracketingNonlinearSolve/Project.toml +++ b/lib/BracketingNonlinearSolve/Project.toml @@ -42,3 +42,6 @@ TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" [targets] test = ["Aqua", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "Test", "TestItemRunner"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index dcd34303f..5be96109e 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -92,3 +92,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Aqua", "BandedMatrices", "DiffEqBase", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "LinearAlgebra", "SparseArrays", "Test"] + +[sources] +SciMLJacobianOperators = {path = "../SciMLJacobianOperators"} diff --git a/lib/NonlinearSolveFirstOrder/Project.toml b/lib/NonlinearSolveFirstOrder/Project.toml index 1df55eaa8..9da1c6741 100644 --- a/lib/NonlinearSolveFirstOrder/Project.toml +++ b/lib/NonlinearSolveFirstOrder/Project.toml @@ -88,3 +88,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ForwardDiff", "Enzyme", "ExplicitImports", "Hwloc", "InteractiveUtils", "LineSearch", "LineSearches", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test", "Zygote"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} +SciMLJacobianOperators = {path = "../SciMLJacobianOperators"} diff --git a/lib/NonlinearSolveHomotopyContinuation/Project.toml b/lib/NonlinearSolveHomotopyContinuation/Project.toml index 700b9f2ee..3f0cfc6ae 100644 --- a/lib/NonlinearSolveHomotopyContinuation/Project.toml +++ b/lib/NonlinearSolveHomotopyContinuation/Project.toml @@ -44,3 +44,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Aqua", "Test", "NonlinearSolve", "Enzyme", "NaNMath"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} diff --git a/lib/NonlinearSolveQuasiNewton/Project.toml b/lib/NonlinearSolveQuasiNewton/Project.toml index 7ebb52f88..68f52a574 100644 --- a/lib/NonlinearSolveQuasiNewton/Project.toml +++ b/lib/NonlinearSolveQuasiNewton/Project.toml @@ -80,3 +80,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] test = ["ADTypes", "Aqua", "BenchmarkTools", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LineSearch", "LineSearches", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "StaticArrays", "Test", "Zygote"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} diff --git a/lib/NonlinearSolveSpectralMethods/Project.toml b/lib/NonlinearSolveSpectralMethods/Project.toml index 53f568bc0..b0a8fd858 100644 --- a/lib/NonlinearSolveSpectralMethods/Project.toml +++ b/lib/NonlinearSolveSpectralMethods/Project.toml @@ -59,3 +59,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Aqua", "BenchmarkTools", "ExplicitImports", "Hwloc", "InteractiveUtils", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "StaticArrays", "Test"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} diff --git a/lib/SCCNonlinearSolve/Project.toml b/lib/SCCNonlinearSolve/Project.toml index 389323487..f237f2f6e 100644 --- a/lib/SCCNonlinearSolve/Project.toml +++ b/lib/SCCNonlinearSolve/Project.toml @@ -46,3 +46,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Aqua", "BenchmarkTools", "ExplicitImports", "Hwloc", "InteractiveUtils", "NonlinearSolveFirstOrder", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "StaticArrays", "Test"] + +[sources] +NonlinearSolveFirstOrder = {path = "../NonlinearSolveFirstOrder"} diff --git a/lib/SimpleNonlinearSolve/Project.toml b/lib/SimpleNonlinearSolve/Project.toml index 2e379dc16..10856caa3 100644 --- a/lib/SimpleNonlinearSolve/Project.toml +++ b/lib/SimpleNonlinearSolve/Project.toml @@ -91,3 +91,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] test = ["Aqua", "DiffEqBase", "Enzyme", "ExplicitImports", "InteractiveUtils", "NonlinearProblemLibrary", "Pkg", "PolyesterForwardDiff", "Random", "ReverseDiff", "StaticArrays", "Test", "TestItemRunner", "Tracker", "Zygote"] + +[sources] +NonlinearSolveBase = {path = "../NonlinearSolveBase"} From 3a211ad7e6c974400b3c91dfcdd390f7cb244b56 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 17 Mar 2025 19:08:41 -0400 Subject: [PATCH 2/4] feat: wrap GridapPETSc --- Project.toml | 15 +++++++++------ ext/NonlinearSolveGridapPETScExt.jl | 9 +++++++++ ext/NonlinearSolvePETScExt.jl | 9 ++++++++- src/extension_algs.jl | 13 +++++++++++++ 4 files changed, 39 insertions(+), 7 deletions(-) create mode 100644 ext/NonlinearSolveGridapPETScExt.jl diff --git a/Project.toml b/Project.toml index 4d46dc060..b61f48d18 100644 --- a/Project.toml +++ b/Project.toml @@ -34,6 +34,8 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [weakdeps] FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce" FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176" +Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e" +GridapPETSc = "bcdc36c2-0c3e-11ea-095a-c9dadae499f1" LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9" @@ -45,9 +47,16 @@ SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4" SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412" Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4" +[sources] +NonlinearSolveBase = {path = "lib/NonlinearSolveBase"} +NonlinearSolveFirstOrder = {path = "lib/NonlinearSolveFirstOrder"} +NonlinearSolveQuasiNewton = {path = "lib/NonlinearSolveQuasiNewton"} +NonlinearSolveSpectralMethods = {path = "lib/NonlinearSolveSpectralMethods"} + [extensions] NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt" NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration" +NonlinearSolveGridapPETScExt = ["Gridap", "GridapPETSc"] NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim" NonlinearSolveMINPACKExt = "MINPACK" NonlinearSolveNLSolversExt = "NLSolvers" @@ -153,9 +162,3 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"] - -[sources] -NonlinearSolveBase = {path = "lib/NonlinearSolveBase"} -NonlinearSolveFirstOrder = {path = "lib/NonlinearSolveFirstOrder"} -NonlinearSolveQuasiNewton = {path = "lib/NonlinearSolveQuasiNewton"} -NonlinearSolveSpectralMethods = {path = "lib/NonlinearSolveSpectralMethods"} diff --git a/ext/NonlinearSolveGridapPETScExt.jl b/ext/NonlinearSolveGridapPETScExt.jl new file mode 100644 index 000000000..8a5827798 --- /dev/null +++ b/ext/NonlinearSolveGridapPETScExt.jl @@ -0,0 +1,9 @@ +module NonlinearSolveGridapPETScExt + +using Gridap: Gridap +using GridapPETSc: GridapPETSc + +using NonlinearSolveBase: NonlinearSolveBase +using NonlinearSolve: NonlinearSolve, GridapPETScSNES + +end \ No newline at end of file diff --git a/ext/NonlinearSolvePETScExt.jl b/ext/NonlinearSolvePETScExt.jl index 7b36cecce..25c73b28a 100644 --- a/ext/NonlinearSolvePETScExt.jl +++ b/ext/NonlinearSolvePETScExt.jl @@ -17,6 +17,11 @@ function SciMLBase.__solve( maxiters = 1000, alias_u0::Bool = false, termination_condition = nothing, show_trace::Val = Val(false), kwargs... ) + if !MPI.Initialized() + @warn "MPI not initialized. Initializing MPI with MPI.Init()." maxlog = 1 + MPI.Init() + end + # XXX: https://petsc.org/release/manualpages/SNES/SNESSetConvergenceTest/ NonlinearSolveBase.assert_extension_supported_termination_condition( termination_condition, alg; abs_norm_supported = false @@ -68,8 +73,10 @@ function SciMLBase.__solve( PETSc.setfunction!(snes, f!, PETSc.VecSeq(zero(u0))) njac = Ref{Int}(-1) - if alg.autodiff !== missing || prob.f.jac !== nothing + # `missing` -> let PETSc compute the Jacobian using finite differences + if alg.autodiff !== missing autodiff = alg.autodiff === missing ? nothing : alg.autodiff + if prob.u0 isa Number jac! = NonlinearSolveBase.construct_extension_jac( prob, alg, prob.u0, prob.u0; autodiff diff --git a/src/extension_algs.jl b/src/extension_algs.jl index 1fad79839..e0631a3f5 100644 --- a/src/extension_algs.jl +++ b/src/extension_algs.jl @@ -488,3 +488,16 @@ function PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, end return PETScSNES(petsclib, mpi_comm, autodiff, kwargs) end + +# TODO: Docs +@concrete struct GridapPETScSNES <: AbstractNonlinearSolveAlgorithm + autodiff + snes_options +end + +function GridapPETScSNES(; autodiff = nothing, kwargs...) + if Base.get_extension(@__MODULE__, :NonlinearSolveGridapPETScExt) === nothing + error("`GridapPETScSNES` requires `GridapPETSc.jl` to be loaded") + end + return GridapPETScSNES(autodiff, kwargs) +end From c2a3ab45e67e5ed4e1cf88748bd373343d3bba39 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 17 Mar 2025 21:54:33 -0400 Subject: [PATCH 3/4] fix: working prototype of GridapPETSc --- ext/NonlinearSolveGridapPETScExt.jl | 118 +++++++++++++++++++++++++++- ext/NonlinearSolvePETScExt.jl | 5 +- src/NonlinearSolve.jl | 2 +- 3 files changed, 120 insertions(+), 5 deletions(-) diff --git a/ext/NonlinearSolveGridapPETScExt.jl b/ext/NonlinearSolveGridapPETScExt.jl index 8a5827798..e24215cad 100644 --- a/ext/NonlinearSolveGridapPETScExt.jl +++ b/ext/NonlinearSolveGridapPETScExt.jl @@ -1,9 +1,125 @@ module NonlinearSolveGridapPETScExt -using Gridap: Gridap +using Gridap: Gridap, Algebra using GridapPETSc: GridapPETSc using NonlinearSolveBase: NonlinearSolveBase using NonlinearSolve: NonlinearSolve, GridapPETScSNES +using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode + +using ConcreteStructs: @concrete +using FastClosures: @closure + +@concrete struct NonlinearSolveOperator <: Algebra.NonlinearOperator + f! + jac! + initial_guess_cache + resid_prototype + jacobian_prototype +end + +function Algebra.residual!(b::AbstractVector, op::NonlinearSolveOperator, x::AbstractVector) + op.f!(b, x) +end + +function Algebra.jacobian!( + A::AbstractMatrix, op::NonlinearSolveOperator, x::AbstractVector +) + op.jac!(A, x) +end + +function Algebra.zero_initial_guess(op::NonlinearSolveOperator) + fill!(op.initial_guess_cache, 0) + return op.initial_guess_cache +end + +function Algebra.allocate_residual(op::NonlinearSolveOperator, ::AbstractVector) + fill!(op.resid_prototype, 0) + return op.resid_prototype +end + +function Algebra.allocate_jacobian(op::NonlinearSolveOperator, ::AbstractVector) + fill!(op.jacobian_prototype, 0) + return op.jacobian_prototype +end + +# TODO: Later we should just wrap `Gridap` generally and pass in `PETSc` as the solver +function SciMLBase.__solve( + prob::NonlinearProblem, alg::GridapPETScSNES, args...; + abstol = nothing, reltol = nothing, + maxiters = 1000, alias_u0::Bool = false, termination_condition = nothing, + show_trace::Val = Val(false), kwargs... +) + # XXX: https://petsc.org/release/manualpages/SNES/SNESSetConvergenceTest/ + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg; abs_norm_supported = false + ) + + f_wrapped!, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0 + ) + T = eltype(u0) + + abstol = NonlinearSolveBase.get_tolerance(abstol, T) + reltol = NonlinearSolveBase.get_tolerance(reltol, T) + + nf = Ref{Int}(0) + + f! = @closure (fx, x) -> begin + nf[] += 1 + f_wrapped!(fx, x) + return fx + end + + if prob.u0 isa Number + jac! = NonlinearSolveBase.construct_extension_jac( + prob, alg, prob.u0, prob.u0; alg.autodiff + ) + J_init = zeros(T, 1, 1) + else + jac!, J_init = NonlinearSolveBase.construct_extension_jac( + prob, alg, u0, resid; alg.autodiff, initial_jacobian = Val(true) + ) + end + + njac = Ref{Int}(-1) + jac_fn! = @closure (J, x) -> begin + njac[] += 1 + jac!(J, x) + return J + end + + nop = NonlinearSolveOperator(f!, jac_fn!, u0, resid, J_init) + + petsc_args = [ + "-snes_rtol", string(reltol), "-snes_atol", string(abstol), + "-snes_max_it", string(maxiters) + ] + for (k, v) in pairs(alg.snes_options) + push!(petsc_args, "-$(k)") + push!(petsc_args, string(v)) + end + show_trace isa Val{true} && push!(petsc_args, "-snes_monitor") + + @show petsc_args + + # TODO: We can reuse the cache returned from this function + sol_u = GridapPETSc.with(args = petsc_args) do + sol_u = copy(u0) + Algebra.solve!(sol_u, GridapPETSc.PETScNonlinearSolver(), nop) + return sol_u + end + + f_wrapped!(resid, sol_u) + u_res = prob.u0 isa Number ? sol_u[1] : sol_u + resid_res = prob.u0 isa Number ? resid[1] : resid + + objective = maximum(abs, resid) + retcode = ifelse(objective ≤ abstol, ReturnCode.Success, ReturnCode.Failure) + return SciMLBase.build_solution( + prob, alg, u_res, resid_res; + retcode, stats = SciMLBase.NLStats(nf[], njac[], -1, -1, -1) + ) +end end \ No newline at end of file diff --git a/ext/NonlinearSolvePETScExt.jl b/ext/NonlinearSolvePETScExt.jl index 25c73b28a..25a93cf98 100644 --- a/ext/NonlinearSolvePETScExt.jl +++ b/ext/NonlinearSolvePETScExt.jl @@ -18,7 +18,7 @@ function SciMLBase.__solve( show_trace::Val = Val(false), kwargs... ) if !MPI.Initialized() - @warn "MPI not initialized. Initializing MPI with MPI.Init()." maxlog = 1 + @warn "MPI not initialized. Initializing MPI with MPI.Init()." maxlog=1 MPI.Init() end @@ -132,8 +132,7 @@ function SciMLBase.__solve( retcode = ifelse(objective ≤ abstol, ReturnCode.Success, ReturnCode.Failure) return SciMLBase.build_solution( prob, alg, u_res, resid_res; - retcode, original = snes, - stats = SciMLBase.NLStats(nf[], njac[], -1, -1, -1) + retcode, stats = SciMLBase.NLStats(nf[], njac[], -1, -1, -1) ) end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index a1b759011..38288368f 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -121,6 +121,6 @@ export NonlinearSolvePolyAlgorithm, FastShortcutNonlinearPolyalg, FastShortcutNL # Extension Algorithms export LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL -export PETScSNES, CMINPACK +export PETScSNES, GridapPETScSNES, CMINPACK end From 3242b267aa456fed006587b9e41c511819f66ac7 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 17 Mar 2025 21:55:39 -0400 Subject: [PATCH 4/4] fix: remove printing --- ext/NonlinearSolveGridapPETScExt.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/ext/NonlinearSolveGridapPETScExt.jl b/ext/NonlinearSolveGridapPETScExt.jl index e24215cad..abd473015 100644 --- a/ext/NonlinearSolveGridapPETScExt.jl +++ b/ext/NonlinearSolveGridapPETScExt.jl @@ -101,8 +101,6 @@ function SciMLBase.__solve( end show_trace isa Val{true} && push!(petsc_args, "-snes_monitor") - @show petsc_args - # TODO: We can reuse the cache returned from this function sol_u = GridapPETSc.with(args = petsc_args) do sol_u = copy(u0)