diff --git a/python/demo/demo_axis.py b/python/demo/demo_axis.py index f1a8385fb8f..9d6ec7a4192 100644 --- a/python/demo/demo_axis.py +++ b/python/demo/demo_axis.py @@ -676,6 +676,7 @@ def create_eps_mu(pml, rho, eps_bkg, mu_bkg): }, ) Esh_m = problem.solve() + assert isinstance(Esh_m, fem.Function) assert problem.solver.getConvergedReason() > 0, "Solver did not converge!" # Scattered magnetic field diff --git a/python/demo/demo_biharmonic.py b/python/demo/demo_biharmonic.py index 4ed85951321..d7e307d151e 100644 --- a/python/demo/demo_biharmonic.py +++ b/python/demo/demo_biharmonic.py @@ -228,6 +228,7 @@ problem = LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) uh = problem.solve() +assert isinstance(uh, fem.Function) # The solution can be written to a {py:class}`XDMFFile # ` file visualization with ParaView or VisIt diff --git a/python/demo/demo_pml.py b/python/demo/demo_pml.py index 1a879837250..9ff717d5776 100644 --- a/python/demo/demo_pml.py +++ b/python/demo/demo_pml.py @@ -721,6 +721,7 @@ def create_eps_mu( }, ) Esh = problem.solve() +assert isinstance(Esh, fem.Function) assert problem.solver.getConvergedReason() > 0, "Solver did not converge!" # - diff --git a/python/demo/demo_poisson.py b/python/demo/demo_poisson.py index 0ee822c719d..bf744706a4e 100644 --- a/python/demo/demo_poisson.py +++ b/python/demo/demo_poisson.py @@ -158,6 +158,7 @@ # + problem = LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"}) uh = problem.solve() +assert isinstance(uh, fem.Function) # - # The solution can be written to a {py:class}`XDMFFile diff --git a/python/demo/demo_scattering_boundary_conditions.py b/python/demo/demo_scattering_boundary_conditions.py index 5332e45cb74..4036b45e585 100644 --- a/python/demo/demo_scattering_boundary_conditions.py +++ b/python/demo/demo_scattering_boundary_conditions.py @@ -640,6 +640,7 @@ def curl_2d(f: fem.Function): gdim = mesh_data.mesh.geometry.dim V_dg = fem.functionspace(mesh_data.mesh, ("Discontinuous Lagrange", degree, (gdim,))) Esh_dg = fem.Function(V_dg) +assert isinstance(Esh, fem.Function) Esh_dg.interpolate(Esh) with io.VTXWriter(mesh_data.mesh.comm, "Esh.bp", Esh_dg) as vtx: diff --git a/python/dolfinx/fem/petsc.py b/python/dolfinx/fem/petsc.py index e98e3abb390..59f8b05bb49 100644 --- a/python/dolfinx/fem/petsc.py +++ b/python/dolfinx/fem/petsc.py @@ -742,20 +742,22 @@ class LinearProblem: def __init__( self, - a: ufl.Form, - L: ufl.Form, + a: typing.Union[ufl.Form, Iterable[Iterable[ufl.Form]]], + L: typing.Union[ufl.Form, Iterable[ufl.Form]], bcs: typing.Optional[Iterable[DirichletBC]] = None, - u: typing.Optional[_Function] = None, + u: typing.Optional[typing.Union[_Function, Iterable[_Function]]] = None, petsc_options: typing.Optional[dict] = None, form_compiler_options: typing.Optional[dict] = None, jit_options: typing.Optional[dict] = None, - ): + kind: typing.Optional[typing.Union[str, Iterable[Iterable[str]]]] = None, + P: typing.Optional[typing.Union[ufl.Form, Iterable[Iterable[ufl.Form]]]] = None, + ) -> None: """Initialize solver for a linear variational problem. Args: - a: Bilinear UFL form, the left hand side of the + a: Bilinear UFL form or a rectangular array of bilinear forms, the left hand side of the variational problem. - L: Linear UFL form, the right hand side of the variational + L: Linear UFL form or a sequence of lienar forms, the right hand side of the variational problem. bcs: Sequecne of Dirichlet boundary conditions. u: Solution function. It is be created if not @@ -771,6 +773,8 @@ def __init__( code generated by FFCx. See `python/dolfinx/jit.py` for all available options. Takes priority over all other option values. + kind: The PETSc matrix and vector type. See :func:`create_matrix` for options. + P: UFL form or a rectangular array of bilinear forms used as a preconditioner. Example:: @@ -778,6 +782,9 @@ def __init__( "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}) + + problem = LinearProblem([[a00, a01], [None, a11]], [L0, L1], bcs=[bc0, bc1], + u=[uh0, uh1]) """ self._a = _create_form( a, @@ -785,28 +792,48 @@ def __init__( form_compiler_options=form_compiler_options, jit_options=jit_options, ) - self._A = create_matrix(self._a) self._L = _create_form( L, dtype=PETSc.ScalarType, form_compiler_options=form_compiler_options, jit_options=jit_options, ) - self._b = create_vector(self._L) + self._A = create_matrix(self._a, kind=kind) + self._preconditioner = _create_form( + P, + dtype=PETSc.ScalarType, + form_compiler_options=form_compiler_options, + jit_options=jit_options, + ) + self._P = ( + create_matrix(self._preconditioner, kind=kind) + if self._preconditioner is not None + else None + ) + + # For nest matrices kind can be a nested list. + kind = "nest" if self._A.getType() == PETSc.Mat.Type.NEST else kind + self._b = create_vector(self._L, kind=kind) + self._x = create_vector(self._L, kind=kind) if u is None: - # Extract function space from TrialFunction (which is at the - # end of the argument list as it is numbered as 1, while the - # Test function is numbered as 0) - self.u = _Function(a.arguments()[-1].ufl_function_space()) + try: + # Extract function space for unknown from the right hand side of the equation. + self._u = _Function(L.arguments()[0].ufl_function_space()) + except AttributeError: + self._u = [_Function(Li.arguments()[0].ufl_function_space()) for Li in L] else: - self.u = u + self._u = u - self._x = dolfinx.la.petsc.create_vector_wrap(self.u.x) self.bcs = bcs - self._solver = PETSc.KSP().create(self.u.function_space.mesh.comm) - self._solver.setOperators(self._A) + try: + comm = self._u.function_space.mesh.comm + except AttributeError: + comm = self._u[0].function_space.mesh.comm + + self._solver = PETSc.KSP().create(comm) + self._solver.setOperators(self._A, self._P) # Give PETSc solver options a unique prefix problem_prefix = f"dolfinx_solve_{id(self)}" @@ -833,7 +860,7 @@ def __del__(self): self._b.destroy() self._x.destroy() - def solve(self) -> _Function: + def solve(self) -> typing.Union[_Function, Iterable[_Function]]: """Solve the problem.""" # Assemble lhs @@ -841,25 +868,47 @@ def solve(self) -> _Function: assemble_matrix(self._A, self._a, bcs=self.bcs) self._A.assemble() + # Assemble preconditioner + if self._P is not None: + self._P.zeroEntries() + assemble_matrix(self._P, self._preconditioner, bcs=self.bcs) + self._P.assemble() + # Assemble rhs - with self._b.localForm() as b_loc: - b_loc.set(0) + if self._b.getType() == PETSc.Vec.Type.NEST: + for b_sub in self._b.getNestSubVecs(): + with b_sub.localForm() as b_local: + b_local.set(0.0) + else: + with self._b.localForm() as b_loc: + b_loc.set(0) assemble_vector(self._b, self._L) # Apply boundary conditions to the rhs if self.bcs is not None: - apply_lifting(self._b, [self._a], bcs=[self.bcs]) - self._b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) - for bc in self.bcs: - bc.set(self._b.array_w) + try: + apply_lifting(self._b, [self._a], bcs=[self.bcs]) + dolfinx.la.petsc._ghost_update( + self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE + ) + for bc in self.bcs: + bc.set(self._b.array_w) + except RuntimeError: + bcs1 = _bcs_by_block(_extract_spaces(self._a, 1), self.bcs) # type: ignore + apply_lifting(self._b, self._a, bcs=bcs1) # type: ignore + dolfinx.la.petsc._ghost_update( + self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE + ) + bcs0 = _bcs_by_block(_extract_spaces(self._L), self.bcs) # type: ignore + dolfinx.fem.petsc.set_bc(self._b, bcs0) else: - self._b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE) + dolfinx.la.petsc._ghost_update(self._b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE) # Solve linear system and update ghost values in the solution self._solver.solve(self._b, self._x) - self.u.x.scatter_forward() - - return self.u + dolfinx.la.petsc._ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) + dolfinx.fem.petsc.assign(self._x, self._u) + return self._u @property def L(self) -> Form: @@ -886,6 +935,11 @@ def solver(self) -> PETSc.KSP: """Linear solver object""" return self._solver + @property + def u(self) -> typing.Union[_Function, list[_Function]]: + """Solution function""" + return self._u + class NonlinearProblem: """Nonlinear problem class for solving the non-linear problems. diff --git a/python/dolfinx/la/petsc.py b/python/dolfinx/la/petsc.py index cae88c52b61..9e3db6f2522 100644 --- a/python/dolfinx/la/petsc.py +++ b/python/dolfinx/la/petsc.py @@ -29,6 +29,15 @@ __all__ = ["assign", "create_vector", "create_vector_wrap"] +def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): # type: ignore + """Helper function for ghost updating PETSc vectors""" + if x.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined] + for x_sub in x.getNestSubVecs(): + x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode) + else: + x.ghostUpdate(addv=insert_mode, mode=scatter_mode) + + def create_vector(index_map: IndexMap, bs: int) -> PETSc.Vec: # type: ignore[name-defined] """Create a distributed PETSc vector. diff --git a/python/test/unit/fem/test_petsc_solver_wrappers.py b/python/test/unit/fem/test_petsc_solver_wrappers.py index 3737c81430e..39cf6956494 100644 --- a/python/test/unit/fem/test_petsc_solver_wrappers.py +++ b/python/test/unit/fem/test_petsc_solver_wrappers.py @@ -1,4 +1,4 @@ -# Copyright (C) 2024 Jørgen S. Dokken +# Copyright (C) 2024-2025 Jørgen S. Dokken # # This file is part of DOLFINx (https://www.fenicsproject.org) # @@ -10,6 +10,7 @@ import numpy as np import pytest +import basix.ufl import dolfinx import ufl @@ -70,3 +71,98 @@ def test_compare_solvers(self, mode): solver.rtol = eps solver.solve(uh) assert np.allclose(u_lin.x.array, uh.x.array, atol=eps, rtol=eps) + + @pytest.mark.parametrize( + "mode", [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet] + ) + @pytest.mark.parametrize("kind", [None, "mpi", "nest", [["aij", None], [None, "baij"]]]) + def test_mixed_system(self, mode, kind): + from petsc4py import PETSc + + msh = dolfinx.mesh.create_unit_square( + MPI.COMM_WORLD, 12, 12, ghost_mode=mode, dtype=PETSc.RealType + ) + + def top_bc(x): + return np.isclose(x[1], 1.0) + + msh.topology.create_connectivity(msh.topology.dim - 1, msh.topology.dim) + bndry_facets = dolfinx.mesh.locate_entities_boundary(msh, msh.topology.dim - 1, top_bc) + + el_0 = basix.ufl.element("Lagrange", msh.basix_cell(), 1, dtype=PETSc.RealType) + el_1 = basix.ufl.element("Lagrange", msh.basix_cell(), 2, dtype=PETSc.RealType) + + if kind is None: + me = basix.ufl.mixed_element([el_0, el_1]) + W = dolfinx.fem.functionspace(msh, me) + V, _ = W.sub(0).collapse() + Q, _ = W.sub(1).collapse() + else: + V = dolfinx.fem.functionspace(msh, el_0) + Q = dolfinx.fem.functionspace(msh, el_1) + W = ufl.MixedFunctionSpace(V, Q) + + u, p = ufl.TrialFunctions(W) + v, q = ufl.TestFunctions(W) + + a00 = ufl.inner(u, v) * ufl.dx + a11 = ufl.inner(p, q) * ufl.dx + x = ufl.SpatialCoordinate(msh) + f = x[0] + 3 * x[1] + g = -(x[1] ** 2) + x[0] + L0 = ufl.inner(f, v) * ufl.dx + L1 = ufl.inner(g, q) * ufl.dx + + f_expr = dolfinx.fem.Expression(f, V.element.interpolation_points) + g_expr = dolfinx.fem.Expression(g, Q.element.interpolation_points) + u_bc = dolfinx.fem.Function(V) + u_bc.interpolate(f_expr) + p_bc = dolfinx.fem.Function(Q) + p_bc.interpolate(g_expr) + + if kind is None: + a = a00 + a11 + L = L0 + L1 + dofs_V = dolfinx.fem.locate_dofs_topological( + (W.sub(0), V), msh.topology.dim - 1, bndry_facets + ) + dofs_Q = dolfinx.fem.locate_dofs_topological( + (W.sub(1), Q), msh.topology.dim - 1, bndry_facets + ) + bcs = [ + dolfinx.fem.dirichletbc(u_bc, dofs_V, W.sub(0)), + dolfinx.fem.dirichletbc(p_bc, dofs_Q, W.sub(1)), + ] + else: + a = [[a00, None], [None, a11]] + L = [L0, L1] + dofs_V = dolfinx.fem.locate_dofs_topological(V, msh.topology.dim - 1, bndry_facets) + dofs_Q = dolfinx.fem.locate_dofs_topological(Q, msh.topology.dim - 1, bndry_facets) + bcs = [ + dolfinx.fem.dirichletbc(u_bc, dofs_V), + dolfinx.fem.dirichletbc(p_bc, dofs_Q), + ] + + petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True, + } + + problem = dolfinx.fem.petsc.LinearProblem( + a, L, bcs=bcs, petsc_options=petsc_options, kind=kind + ) + wh = problem.solve() + if kind is None: + uh, ph = wh.split() + else: + uh, ph = wh + error_uh = dolfinx.fem.form(ufl.inner(uh - f, uh - f) * ufl.dx) + error_ph = dolfinx.fem.form(ufl.inner(ph - g, ph - g) * ufl.dx) + local_uh_L2 = dolfinx.fem.assemble_scalar(error_uh) + local_ph_L2 = dolfinx.fem.assemble_scalar(error_ph) + global_uh_L2 = np.sqrt(msh.comm.allreduce(local_uh_L2, op=MPI.SUM)) + global_ph_L2 = np.sqrt(msh.comm.allreduce(local_ph_L2, op=MPI.SUM)) + tol = 500 * np.finfo(dolfinx.default_scalar_type).eps + assert global_uh_L2 < tol and global_ph_L2 < tol