Merge pull request #460 from sblauth/dev/fieldsplit_fields

Support for nested fieldsplits

CHANGELOG.rst

@@ -12,6 +12,8 @@ of the maintenance releases, please take a look at
 
 * Add a wrapper for PETSc's SNES solver for nonlinear equations. This is used internally in cashocs whenever possible. For the solution of the state system, our own Newton solver is the default for backwards compatibility. Users can use the new SNES backend by specifying :ini:`backend = petsc` in the Section StateSystem of the configuration.
 
+* Allows nesting of PETSc Fieldsplit PCs with the command line option "pc_fieldsplit_%d_fields <a,b,...>, as explained at `<https://petsc.org/main/manualpages/PC/PCFieldSplitSetFields/>`_
+
 * Increase the precision of the Gmsh output from cashocs
 
 * Add mesh quality constraints for shape optimization: These ensure that the angles of the (solid) angles of triangles and tetrahedrons cannot fall below a specified threshold.
@@ -28,7 +30,7 @@ of the maintenance releases, please take a look at
 
   * Section MeshQualityConstraints
 
-    * This section includes parameters for the new mesh quality constraints for shape optimization. These are described in the documentation at https://cashocs.readthedocs.io/en/stable/user/demos/shape_optimization/doc_config/#section-meshqualityconstraints
+    * This section includes parameters for the new mesh quality constraints for shape optimization. These are described in the documentation at `<https://cashocs.readthedocs.io/en/stable/user/demos/shape_optimization/doc_config/#section-meshqualityconstraints>`_
 
 
 

cashocs/_utils/linalg.py

@@ -232,17 +232,49 @@ def setup_fieldsplit_preconditioner(
                     "problem to be solved is not a mixed one.",
                 )
 
-            pc = ksp.getPC()
-            pc.setType(PETSc.PC.Type.FIELDSPLIT)
-            idx = []
-            name = []
-            for i in range(function_space.num_sub_spaces()):
-                idx_i = PETSc.IS().createGeneral(function_space.sub(i).dofmap().dofs())
-                idx.append(idx_i)
-                name.append(f"{i:d}")
-            idx_tuples = zip(name, idx)
-
-            pc.setFieldSplitIS(*idx_tuples)
+            if not any(key.endswith("_fields") for key in options.keys()):
+                pc = ksp.getPC()
+                pc.setType(PETSc.PC.Type.FIELDSPLIT)
+                idx = []
+                name = []
+                for i in range(function_space.num_sub_spaces()):
+                    idx_i = PETSc.IS().createGeneral(
+                        function_space.sub(i).dofmap().dofs()
+                    )
+                    idx.append(idx_i)
+                    name.append(f"{i:d}")
+                idx_tuples = zip(name, idx)
+
+                pc.setFieldSplitIS(*idx_tuples)
+            else:
+                dof_total = function_space.dofmap().dofs()
+                offset = np.min(dof_total)
+
+                num_sub_spaces = function_space.num_sub_spaces()
+                dof_list = [
+                    np.array(function_space.sub(i).dofmap().dofs())
+                    for i in range(num_sub_spaces)
+                ]
+
+                section = PETSc.Section().create()
+                section.setNumFields(num_sub_spaces)
+
+                for i in range(num_sub_spaces):
+                    section.setFieldName(i, f"{i:d}")
+                    section.setFieldComponents(i, 1)
+                section.setChart(0, len(dof_total))
+                for field_idx, dofs in enumerate(dof_list):
+                    for i in dofs:
+                        section.setDof(i - offset, 1)
+                        section.setFieldDof(i - offset, field_idx, 1)
+                section.setUp()
+
+                dm = PETSc.DMShell().create()
+                dm.setDefaultSection(section)
+                dm.setUp()
+
+                ksp.setDM(dm)
+                ksp.setDMActive(False)
 
 
 def _initialize_comm(comm: Optional[MPI.Comm] = None) -> MPI.Comm:

demos/undocumented/solvers/fieldsplit_nesting.py

@@ -0,0 +1,59 @@
+from fenics import *
+
+import cashocs
+
+mesh, subdomains, boundaries, dx, ds, dS = cashocs.regular_mesh(16)
+v_elem = VectorElement("CG", mesh.ufl_cell(), 2)
+p_elem = FiniteElement("CG", mesh.ufl_cell(), 1)
+V = FunctionSpace(mesh, MixedElement(v_elem, p_elem, p_elem))
+# order of the MixedElement determines (single) prefixes for PETSc solver
+# velocity block gets 0, pressure 1, and temperature 2
+
+upT = Function(V)
+u, p, T = split(upT)
+v, q, S = TestFunctions(V)
+
+mu = 1.0 / (T + 1)
+
+F = (
+    mu * inner(grad(u), grad(v)) * dx
+    + dot(grad(u) * u, v) * dx
+    - p * div(v) * dx
+    - q * div(u) * dx
+    + dot(grad(T), grad(S)) * dx
+    + dot(u, grad(T)) * S * dx
+    - Constant(100.0) * S * dx
+)
+u_in = Expression(("4.0 * x[1] * (1.0 - x[1])", "0.0"), degree=2)
+bcs = cashocs.create_dirichlet_bcs(V.sub(0), u_in, boundaries, 1)
+bcs += cashocs.create_dirichlet_bcs(V.sub(0), Constant((0.0, 0.0)), boundaries, [3, 4])
+bcs += cashocs.create_dirichlet_bcs(V.sub(2), Constant(1.0), boundaries, [1, 3, 4])
+
+petsc_options = {
+    "snes_type": "newtonls",
+    "snes_monitor": None,
+    "snes_max_it": 7,
+    "ksp_type": "fgmres",
+    "ksp_rtol": 1e-1,
+    "pc_type": "fieldsplit",
+    "pc_fieldsplit_type": "multiplicative",
+    "pc_fieldsplit_0_fields": "0,1",  # will get prefix fieldsplit_0_
+    "pc_fieldsplit_1_fields": "2",  # will get prefix fieldsplit_2_
+    "fieldsplit_0_ksp_type": "fgmres",
+    "fieldsplit_0_ksp_rtol": 1e-1,
+    "fieldsplit_0_pc_type": "fieldsplit",  # sub fields get global (mixed) index
+    "fieldsplit_0_pc_fieldsplit_type": "schur",
+    "fieldsplit_0_pc_fieldsplit_schur_precondition": "selfp",
+    "fieldsplit_0_fieldsplit_0_ksp_type": "preonly",
+    "fieldsplit_0_fieldsplit_0_pc_type": "lu",
+    "fieldsplit_0_fieldsplit_1_ksp_type": "gmres",
+    "fieldsplit_0_fieldsplit_1_ksp_rtol": 1e-1,
+    "fieldsplit_0_fieldsplit_1_pc_type": "hypre",
+    "fieldsplit_0_fieldsplit_1_ksp_converged_reason": None,
+    "fieldsplit_2_ksp_type": "gmres",
+    "fieldsplit_2_ksp_rtol": 1e-1,
+    "fieldsplit_2_pc_type": "hypre",
+}
+
+cashocs.snes_solve(F, upT, bcs, petsc_options=petsc_options, max_iter=8)
+u, p, T = upT.split(True)

tests/test_fieldsplit.py

@@ -0,0 +1,157 @@
+# Copyright (C) 2020-2024 Sebastian Blauth
+#
+# This file is part of cashocs.
+#
+# cashocs is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# cashocs is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with cashocs.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Tests for the linear solvers.
+
+"""
+
+from fenics import *
+import numpy as np
+
+import cashocs
+
+
+def test_fieldsplit_pc():
+    mesh, subdomains, boundaries, dx, ds, dS = cashocs.regular_mesh(8)
+    v_elem = VectorElement("CG", mesh.ufl_cell(), 2)
+    p_elem = FiniteElement("CG", mesh.ufl_cell(), 1)
+    V = FunctionSpace(mesh, MixedElement(v_elem, p_elem))
+
+    up = Function(V)
+    u, p = split(up)
+    v, q = TestFunctions(V)
+
+    F = inner(grad(u), grad(v)) * dx - p * div(v) * dx - q * div(u) * dx
+    bcs = cashocs.create_dirichlet_bcs(V.sub(0), Constant((1.0, 0.0)), boundaries, 1)
+    bcs += cashocs.create_dirichlet_bcs(
+        V.sub(0), Constant((0.0, 0.0)), boundaries, [3, 4]
+    )
+
+    ksp_options = {
+        "ksp_type": "fgmres",
+        "ksp_rtol": 1e-5,
+        "ksp_max_it": 1,
+        "ksp_monitor_true_residual": None,
+        "pc_type": "fieldsplit",
+        "pc_fieldsplit_type": "schur",
+        "pc_fieldsplit_schur_precondition": "selfp",
+        "fieldsplit_0_ksp_type": "preonly",
+        "fieldsplit_0_pc_type": "lu",
+        "fieldsplit_1_ksp_type": "gmres",
+        "fieldsplit_1_ksp_rtol": 1e-6,
+        "fieldsplit_1_ksp_max_it": 25,
+        "fieldsplit_1_pc_type": "hypre",
+        "fieldsplit_1_ksp_converged_reason": None,
+    }
+
+    cashocs.linear_solve(F, up, bcs, ksp_options=ksp_options)
+    assert True
+
+
+def test_fieldsplit_pc_section():
+    mesh, subdomains, boundaries, dx, ds, dS = cashocs.regular_mesh(8)
+    v_elem = VectorElement("CG", mesh.ufl_cell(), 2)
+    p_elem = FiniteElement("CG", mesh.ufl_cell(), 1)
+    V = FunctionSpace(mesh, MixedElement(p_elem, v_elem))
+
+    pu = Function(V)
+    p, u = split(pu)
+    q, v = TestFunctions(V)
+
+    F = inner(grad(u), grad(v)) * dx - p * div(v) * dx - q * div(u) * dx
+    bcs = cashocs.create_dirichlet_bcs(V.sub(1), Constant((1.0, 0.0)), boundaries, 1)
+    bcs += cashocs.create_dirichlet_bcs(
+        V.sub(1), Constant((0.0, 0.0)), boundaries, [3, 4]
+    )
+
+    ksp_options = {
+        "ksp_type": "fgmres",
+        "ksp_rtol": 1e-5,
+        "ksp_max_it": 1,
+        "pc_type": "fieldsplit",
+        "pc_fieldsplit_type": "schur",
+        "pc_fieldsplit_schur_precondition": "selfp",
+        "pc_fieldsplit_0_fields": "1",
+        "pc_fieldsplit_1_fields": "0",
+        "fieldsplit_0_ksp_type": "gmres",
+        "fieldsplit_0_pc_type": "hypre",
+        "fieldsplit_0_ksp_rtol": 1e-6,
+        "fieldsplit_0_ksp_max_it": 25,
+        "fieldsplit_0_ksp_converged_reason": None,
+        "fieldsplit_1_ksp_type": "preonly",
+        "fieldsplit_1_pc_type": "lu",
+    }
+
+    cashocs.linear_solve(F, pu, bcs, ksp_options=ksp_options)
+    assert True
+
+
+def test_fieldsplit_snes_nested():
+    mesh, subdomains, boundaries, dx, ds, dS = cashocs.regular_mesh(8)
+    v_elem = VectorElement("CG", mesh.ufl_cell(), 2)
+    p_elem = FiniteElement("CG", mesh.ufl_cell(), 1)
+    V = FunctionSpace(mesh, MixedElement(v_elem, p_elem, p_elem))
+
+    upT = Function(V)
+    u, p, T = split(upT)
+    v, q, S = TestFunctions(V)
+
+    mu = 1.0 / (T + 1)
+
+    F = (
+        mu * inner(grad(u), grad(v)) * dx
+        + dot(grad(u) * u, v) * dx
+        - p * div(v) * dx
+        - q * div(u) * dx
+        + dot(grad(T), grad(S)) * dx
+        + dot(u, grad(T)) * S * dx
+        - Constant(1.0) * S * dx
+    )
+    bcs = cashocs.create_dirichlet_bcs(V.sub(0), Constant((1.0, 0.0)), boundaries, 1)
+    bcs += cashocs.create_dirichlet_bcs(
+        V.sub(0), Constant((0.0, 0.0)), boundaries, [3, 4]
+    )
+    bcs += cashocs.create_dirichlet_bcs(V.sub(2), Constant(1.0), boundaries, [1, 3, 4])
+
+    petsc_options = {
+        "snes_type": "newtonls",
+        "snes_monitor": None,
+        "snes_max_it": 7,
+        "ksp_type": "fgmres",
+        "ksp_rtol": 1e-1,
+        "pc_type": "fieldsplit",
+        "pc_fieldsplit_type": "multiplicative",
+        "pc_fieldsplit_0_fields": "0,1",
+        "pc_fieldsplit_1_fields": "2",
+        "fieldsplit_0_ksp_type": "fgmres",
+        "fieldsplit_0_ksp_rtol": 1e-1,
+        "fieldsplit_0_pc_type": "fieldsplit",
+        "fieldsplit_0_pc_fieldsplit_type": "schur",
+        "fieldsplit_0_pc_fieldsplit_schur_precondition": "selfp",
+        "fieldsplit_0_fieldsplit_0_ksp_type": "preonly",
+        "fieldsplit_0_fieldsplit_0_pc_type": "lu",
+        "fieldsplit_0_fieldsplit_1_ksp_type": "gmres",
+        "fieldsplit_0_fieldsplit_1_ksp_rtol": 1e-1,
+        "fieldsplit_0_fieldsplit_1_pc_type": "hypre",
+        "fieldsplit_0_fieldsplit_1_ksp_converged_reason": None,
+        "fieldsplit_2_ksp_type": "gmres",
+        "fieldsplit_2_ksp_rtol": 1e-1,
+        "fieldsplit_2_pc_type": "hypre",
+    }
+
+    cashocs.snes_solve(F, upT, bcs, petsc_options=petsc_options, max_iter=8)
+    assert True