feat: wrap GridapPETSc

Author: avik-pal

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

ext/NonlinearSolveGridapPETScExt.jl

@@ -0,0 +1,9 @@
+module NonlinearSolveGridapPETScExt
+
+using Gridap: Gridap
+using GridapPETSc: GridapPETSc
+
+using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolve: NonlinearSolve, GridapPETScSNES
+
+end

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

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