Test PETSc solver wrappers

python/test/unit/fem/test_petsc_solver_wrappers.py

@@ -0,0 +1,72 @@
+# Copyright (C) 2024 Jørgen S. Dokken
+#
+# This file is part of DOLFINx (https://www.fenicsproject.org)
+#
+# SPDX-License-Identifier:    LGPL-3.0-or-later
+"""Unit tests for high-level wrapper around PETSc for linear and non-linear problems"""
+
+from mpi4py import MPI
+
+import numpy as np
+import pytest
+
+import dolfinx
+import ufl
+
+
+@pytest.mark.petsc4py
+class TestPETScSolverWrappers:
+    @pytest.mark.parametrize(
+        "mode", [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet]
+    )
+    def test_compare_solvers(self, mode):
+        """Test that the wrapper for Linear problem and NonlinearProblem give the same result"""
+        from petsc4py import PETSc
+
+        import dolfinx.fem.petsc
+        import dolfinx.nls.petsc
+
+        msh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 12, 12, ghost_mode=mode)
+        V = dolfinx.fem.functionspace(msh, ("Lagrange", 1))
+        uh = dolfinx.fem.Function(V)
+        v = ufl.TestFunction(V)
+        x = ufl.SpatialCoordinate(msh)
+        f = x[0] * ufl.sin(x[1])
+        F = ufl.inner(uh, v) * ufl.dx - ufl.inner(f, v) * ufl.dx
+        u = ufl.TrialFunction(V)
+        a = ufl.replace(F, {uh: u})
+
+        sys = PETSc.Sys()
+        if MPI.COMM_WORLD.size == 1:
+            factor_type = "petsc"
+        elif sys.hasExternalPackage("superlu_dist"):
+            factor_type = "superlu_dist"
+        elif sys.hasExternalPackage("mumps"):
+            factor_type = "mumps"
+        else:
+            pytest.skip("No external solvers available in parallel")
+
+        petsc_options = {
+            "ksp_type": "preonly",
+            "pc_type": "lu",
+            "pc_factor_mat_solver_type": factor_type,
+        }
+        linear_problem = dolfinx.fem.petsc.LinearProblem(
+            ufl.lhs(a), ufl.rhs(a), petsc_options=petsc_options
+        )
+        u_lin = linear_problem.solve()
+
+        nonlinear_problem = dolfinx.fem.petsc.NonlinearProblem(F, uh)
+
+        solver = dolfinx.nls.petsc.NewtonSolver(msh.comm, nonlinear_problem)
+        ksp = solver.krylov_solver
+        ksp.setType("preonly")
+        ksp.getPC().setType("lu")
+        ksp.getPC().setFactorSolverType(factor_type)
+
+        eps = 100 * np.finfo(dolfinx.default_scalar_type).eps
+
+        solver.atol = eps
+        solver.rtol = eps
+        solver.solve(uh)
+        assert np.allclose(u_lin.x.array, uh.x.array, atol=eps, rtol=eps)