demo BDDC mixed poisson

Author: stefanozampini

python/demo/demo_mixed-poisson-bddc.py

@@ -0,0 +1,166 @@
+# ---
+# jupyter:
+#   jupytext:
+#     text_representation:
+#       extension: .py
+#       format_name: light
+#       format_version: '1.5'
+#       jupytext_version: 1.14.1
+# ---
+
+# # Mixed formulation of the Poisson equation solved with BDDC preconditioner # noqa
+#
+# This demo illustrates how to solve the Poisson equation using a mixed
+# (two-field) formulation and a monolithic Balancing Domain Decomposition by
+# Constraits (BDDC) preconditioner for the Krylov solver. For details, see
+# the [paper](https://doi.org/10.1137/16M1080653). For details on the
+# equations, see demo_mixed-poisson.py
+#
+# ```{admonition} Download sources
+# :class: download
+# * {download}`Python script <./demo_mixed-poisson-bddc.py>`
+# * {download}`Jupyter notebook <./demo_mixed-poisson-bddc.ipynb>`
+# ```
+
+#
+# ## Implementation
+#
+# Import the required modules:
+
+# +
+from mpi4py import MPI
+from petsc4py import PETSc
+
+import numpy as np
+
+import dolfinx
+import ufl
+from basix.ufl import element
+from dolfinx import fem, mesh
+from dolfinx.fem.petsc import discrete_gradient, interpolation_matrix
+from dolfinx.mesh import CellType, create_unit_square, GhostMode
+
+# Solution scalar (e.g., float32, complex128) and geometry (float32/64)
+# types
+dtype = dolfinx.default_scalar_type
+xdtype = dolfinx.default_real_type
+# -
+
+# Create a two-dimensional mesh. The iterative solver supports three-
+# dimensional meshes too. BDDC methods needs the linear operator
+# to be stored in unassembled format, i.e. assembled on each MPI process
+# associated with the subdomain part of the mesh, but not across processes
+# To this end, we require the mesh to not be distributed with overlap
+
+# +
+msh = create_unit_square(MPI.COMM_WORLD, 96, 96, CellType.triangle, ghost_mode=GhostMode.none, dtype=xdtype)
+# -
+
+# From now on, the script is the same as demo_mixed-poisson.py.
+
+# +
+k = 1
+V = fem.functionspace(msh, element("RT", msh.basix_cell(), k, dtype=xdtype))
+W = fem.functionspace(msh, element("Discontinuous Lagrange", msh.basix_cell(), k - 1, dtype=xdtype))
+Q = ufl.MixedFunctionSpace(V, W)
+
+sigma, u = ufl.TrialFunctions(Q)
+tau, v = ufl.TestFunctions(Q)
+
+x = ufl.SpatialCoordinate(msh)
+f = 10 * ufl.exp(-((x[0] - 0.5) * (x[0] - 0.5) + (x[1] - 0.5) * (x[1] - 0.5)) / 0.02)
+
+dx = ufl.Measure("dx", msh)
+
+a = ufl.extract_blocks(
+    ufl.inner(sigma, tau) * dx + ufl.inner(u, ufl.div(tau)) * dx + ufl.inner(ufl.div(sigma), v) * dx
+)
+L = [ufl.ZeroBaseForm((tau,)), -ufl.inner(f, v) * dx]
+
+fdim = msh.topology.dim - 1
+facets_top = mesh.locate_entities_boundary(msh, fdim, lambda x: np.isclose(x[1], 1.0))
+facets_bottom = mesh.locate_entities_boundary(msh, fdim, lambda x: np.isclose(x[1], 0.0))
+dofs_top = fem.locate_dofs_topological(V, fdim, facets_top)
+dofs_bottom = fem.locate_dofs_topological(V, fdim, facets_bottom)
+
+cells_top_ = mesh.compute_incident_entities(msh.topology, facets_top, fdim, fdim + 1)
+cells_bottom = mesh.compute_incident_entities(msh.topology, facets_bottom, fdim, fdim + 1)
+g = fem.Function(V, dtype=dtype)
+g.interpolate(lambda x: np.vstack((np.zeros_like(x[0]), np.sin(5 * x[0]))), cells0=cells_top_)
+g.interpolate(lambda x: np.vstack((np.zeros_like(x[0]), -np.sin(5 * x[0]))), cells0=cells_bottom)
+bcs = [fem.dirichletbc(g, dofs_top), fem.dirichletbc(g, dofs_bottom)]
+
+sigma, u = fem.Function(V, name="sigma", dtype=dtype), fem.Function(W, name="u", dtype=dtype)
+# -
+
+# We now create a linear problem solver for the mixed problem.
+# As we will use BDDC, we specify the matrix kind to be "is".
+# We also show a minimal customization for the method to be
+# applied monolithically on the symmetric indefinite linear
+# system arising from the mixed problem.
+# In particular, we need solvers for the subproblems that can
+# support the saddle point formulation.
+
+# +
+if PETSc.Sys().hasExternalPackage("umfpack"):
+    solver = "lu"
+    solver_type = "umfpack"
+elif PETSc.Sys().hasExternalPackage("mumps"):
+    solver = "cholesky"
+    solver_type = "mumps"
+elif PETSc.Sys().hasExternalPackage("superlu"):
+    solver = "lu"
+    solver_type = "superlu"
+elif PETSc.Sys().hasExternalPackage("umfpack"):
+    solver = "lu"
+    solver_type = "umfpack"
+else:
+    solver = "svd"
+    solver_type = "dummy"
+
+problem = fem.petsc.LinearProblem(
+    a,
+    L,
+    u=[sigma, u],
+    kind="is",
+    bcs=bcs,
+    petsc_options_prefix="demo_mixed_poisson_",
+    petsc_options={
+        "ksp_rtol": 1e-8,
+        "ksp_view": None,
+        "pc_type": "bddc",
+        "pc_bddc_use_local_mat_graph" : False,
+        "pc_bddc_benign_trick": None,
+        "pc_bddc_nonetflux": None,
+        "pc_bddc_detect_disconnected": None,
+        "pc_bddc_dirichlet_pc_type": solver,
+        "pc_bddc_dirichlet_pc_factor_mat_solver_type": solver_type,
+        "pc_bddc_neumann_pc_factor_mat_solver_type": solver_type,
+        "pc_bddc_neumann_pc_type": solver,
+        "pc_bddc_coarse_redundant_pc_factor_mat_solver_type": solver_type,
+        "pc_bddc_coarse_redundant_pc_type": solver,
+    },
+)
+# -
+
+# Solve and save the solution `u` in VTX format:
+
+# +
+ksp = problem.solver
+ksp.setMonitor(
+    lambda _, its, rnorm: PETSc.Sys.Print(f"Iteration: {its:>4d}, residual: {rnorm:.3e}")
+)
+
+problem.solve()
+converged_reason = problem.solver.getConvergedReason()
+assert converged_reason > 0, f"Krylov solver has not converged, reason: {converged_reason}."
+
+if dolfinx.has_adios2:
+    from dolfinx.io import VTXWriter
+
+    u.name = "u"
+    with VTXWriter(msh.comm, "output_mixed_poisson_bddc.bp", u, "bp4") as f:
+        f.write(0.0)
+else:
+    print("ADIOS2 required for VTX output.")
+# -