Fix Hypre auxiliary spaces (#4712)

* Fix Hypre auxiliary spaces

Author: pbrubeck

firedrake/preconditioners/hypre_ads.py

@@ -1,11 +1,12 @@
 from firedrake.preconditioners.base import PCBase
+from firedrake.preconditioners.fdm import tabulate_exterior_derivative
 from firedrake.petsc import PETSc
 from firedrake.function import Function
-from firedrake.ufl_expr import TestFunction
+from firedrake.ufl_expr import TrialFunction
 from firedrake.dmhooks import get_function_space
 from firedrake.preconditioners.hypre_ams import chop
 from firedrake.interpolation import interpolate
-from finat.ufl import VectorElement
+from finat.ufl import FiniteElement, TensorElement, VectorElement
 from ufl import grad, curl, SpatialCoordinate
 from pyop2.utils import as_tuple
 
@@ -27,16 +28,34 @@ def initialize(self, obj):
         if formdegree != 2 or degree != 1:
             raise ValueError("Hypre ADS requires lowest order RT elements! (not %s of degree %d)" % (family, degree))
 
-        P1 = V.reconstruct(family="Lagrange", degree=1)
-        NC1 = V.reconstruct(family="N1curl" if mesh.ufl_cell().is_simplex() else "NCE", degree=1)
+        # Get the auxiliary Nedelec and Lagrange spaces and the coordinate space
+        cell = V.ufl_element().cell
+        NC1_element = FiniteElement("N1curl" if cell.is_simplex() else "NCE", cell=cell, degree=1)
+        P1_element = FiniteElement("Lagrange", cell=cell, degree=1)
+        coords_element = VectorElement(P1_element, dim=mesh.geometric_dimension())
+        if V.shape:
+            NC1_element = TensorElement(NC1_element, shape=V.shape)
+            P1_element = TensorElement(P1_element, shape=V.shape)
+
+        NC1 = V.reconstruct(element=NC1_element)
+        P1 = V.reconstruct(element=P1_element)
+        VectorP1 = V.reconstruct(element=coords_element)
+
         G_callback = appctx.get("get_gradient", None)
         if G_callback is None:
-            G = chop(assemble(interpolate(grad(TestFunction(P1)), NC1)).petscmat)
+            try:
+                G = chop(assemble(interpolate(grad(TrialFunction(P1)), NC1)).petscmat)
+            except NotImplementedError:
+                # dual evaluation not yet implemented see https://github.com/firedrakeproject/fiat/issues/109
+                G = tabulate_exterior_derivative(P1, NC1)
         else:
             G = G_callback(P1, NC1)
         C_callback = appctx.get("get_curl", None)
         if C_callback is None:
-            C = chop(assemble(interpolate(curl(TestFunction(NC1)), V)).petscmat)
+            try:
+                C = chop(assemble(interpolate(curl(TrialFunction(NC1)), V)).petscmat)
+            except NotImplementedError:
+                C = tabulate_exterior_derivative(NC1, V)
         else:
             C = C_callback(NC1, V)
 
@@ -50,7 +69,6 @@ def initialize(self, obj):
         pc.setHYPREDiscreteGradient(G)
         pc.setHYPREDiscreteCurl(C)
 
-        VectorP1 = P1.reconstruct(element=VectorElement(P1.ufl_element()))
         coords = Function(VectorP1).interpolate(SpatialCoordinate(mesh))
         pc.setCoordinates(coords.dat.data_ro.copy())
 

firedrake/preconditioners/hypre_ams.py

@@ -1,13 +1,14 @@
 import petsctools
 from firedrake.preconditioners.base import PCBase
+from firedrake.preconditioners.fdm import tabulate_exterior_derivative
 from firedrake.petsc import PETSc
 from firedrake.function import Function
-from firedrake.ufl_expr import TestFunction
+from firedrake.ufl_expr import TrialFunction
 from firedrake.dmhooks import get_function_space
 from firedrake.utils import complex_mode
 from firedrake.interpolation import interpolate
 from ufl import grad, SpatialCoordinate
-from finat.ufl import VectorElement
+from finat.ufl import FiniteElement, TensorElement, VectorElement
 from pyop2.utils import as_tuple
 
 __all__ = ("HypreAMS",)
@@ -48,10 +49,21 @@ def initialize(self, obj):
         if formdegree != 1 or degree != 1:
             raise ValueError("Hypre AMS requires lowest order Nedelec elements! (not %s of degree %d)" % (family, degree))
 
-        P1 = V.reconstruct(family="Lagrange", degree=1)
+        # Get the auxiliary Lagrange space and the coordinate space
+        P1_element = FiniteElement("Lagrange", degree=1)
+        coords_element = VectorElement(P1_element, dim=mesh.geometric_dimension())
+        if V.shape:
+            P1_element = TensorElement(P1_element, shape=V.shape)
+        P1 = V.reconstruct(element=P1_element)
+        VectorP1 = V.reconstruct(element=coords_element)
+
         G_callback = appctx.get("get_gradient", None)
         if G_callback is None:
-            G = chop(assemble(interpolate(grad(TestFunction(P1)), V)).petscmat)
+            try:
+                G = chop(assemble(interpolate(grad(TrialFunction(P1)), V)).petscmat)
+            except NotImplementedError:
+                # dual evaluation not yet implemented see https://github.com/firedrakeproject/fiat/issues/109
+                G = tabulate_exterior_derivative(P1, V)
         else:
             G = G_callback(P1, V)
 
@@ -68,7 +80,6 @@ def initialize(self, obj):
         if zero_beta:
             pc.setHYPRESetBetaPoissonMatrix(None)
 
-        VectorP1 = P1.reconstruct(element=VectorElement(P1.ufl_element()))
         coords = Function(VectorP1).interpolate(SpatialCoordinate(mesh))
         pc.setCoordinates(coords.dat.data_ro.copy())
         pc.setFromOptions()

tests/firedrake/regression/test_hypre_ads.py

@@ -2,89 +2,47 @@
 from firedrake import *
 
 
-@pytest.mark.skiphypre
-@pytest.mark.skipcomplex
-def test_homogeneous_field_linear():
-    mesh = UnitCubeMesh(10, 10, 10)
-    V = FunctionSpace(mesh, "RT", 1)
-
-    u = TrialFunction(V)
-    v = TestFunction(V)
-
-    a = inner(div(u), div(v))*dx + inner(u, v)*dx
-    L = inner(Constant((1, 0.5, 4)), v)*dx
-
-    bc = DirichletBC(V, Constant((1, 0.5, 4)), (1, 2, 3, 4))
-
-    params = {'snes_type': 'ksponly',
-              'ksp_type': 'cg',
-              'ksp_max_it': '30',
-              'ksp_rtol': '1e-15',
-              'ksp_atol': '1e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.HypreADS',
-              }
-
-    u = Function(V)
-    solve(a == L, u, bc, solver_parameters=params)
-    assert (errornorm(Constant((1, 0.5, 4)), u, 'L2') < 1e-10)
+@pytest.fixture(params=["simplex", "hexahedron"])
+def V(request):
+    cell = request.param
+    if cell == "simplex":
+        mesh = UnitCubeMesh(10, 10, 10)
+        V = FunctionSpace(mesh, "RT", 1)
+    elif cell == "hexahedron":
+        mesh = ExtrudedMesh(UnitSquareMesh(10, 10, quadrilateral=True), 10)
+        V = FunctionSpace(mesh, "NCF", 1)
+    else:
+        raise ValueError(f"Unrecognized cell {cell}.")
+    return V
 
 
 @pytest.mark.skiphypre
 @pytest.mark.skipcomplex
-def test_homogeneous_field_matfree():
-    mesh = UnitCubeMesh(10, 10, 10)
-    V = FunctionSpace(mesh, "RT", 1)
-
+@pytest.mark.parametrize("mat_type,interface", [("aij", "linear"), ("matfree", "linear"), ("aij", "nonlinear")])
+def test_homogeneous_field(V, mat_type, interface):
     u = TrialFunction(V)
     v = TestFunction(V)
 
+    u_exact = Constant((1, 0.5, 4))
     a = inner(div(u), div(v))*dx + inner(u, v)*dx
-    L = inner(Constant((1, 0.5, 4)), v)*dx
+    L = inner(u_exact, v)*dx
 
-    bc = DirichletBC(V, Constant((1, 0.5, 4)), (1, 2, 3, 4))
+    bc = DirichletBC(V, u_exact, (1, 2, 3, 4))
 
-    params = {'snes_type': 'ksponly',
-              'mat_type': 'matfree',
-              'ksp_type': 'cg',
-              'ksp_max_it': '30',
-              'ksp_rtol': '1e-15',
-              'ksp_atol': '1e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.AssembledPC',
-              'assembled_pc_type': 'python',
-              'assembled_pc_python_type': 'firedrake.HypreADS',
-              }
+    params = {
+        'snes_type': 'ksponly',
+        'mat_type': mat_type,
+        'pmat_type': 'aij',
+        'ksp_type': 'cg',
+        'ksp_max_it': '30',
+        'ksp_rtol': '2e-15',
+        'pc_type': 'python',
+        'pc_python_type': 'firedrake.HypreADS',
+    }
 
     u = Function(V)
     solve(a == L, u, bc, solver_parameters=params)
-    assert (errornorm(Constant((1, 0.5, 4)), u, 'L2') < 1e-10)
-
-
-@pytest.mark.skiphypre
-@pytest.mark.skipcomplex
-def test_homogeneous_field_nonlinear():
-    mesh = UnitCubeMesh(10, 10, 10)
-    V = FunctionSpace(mesh, "RT", 1)
-
-    u = Function(V)
-    v = TestFunction(V)
-
-    F = inner(div(u), div(v))*dx + inner(u, v)*dx - inner(Constant((1, 0.5, 4)), v)*dx
-
-    bc = DirichletBC(V, Constant((1, 0.5, 4)), (1, 2, 3, 4))
-
-    params = {'snes_type': 'ksponly',
-              'ksp_type': 'cg',
-              'ksp_itmax': '30',
-              'ksp_rtol': '1e-15',
-              'ksp_atol': '1e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.HypreADS',
-              }
-
-    solve(F == 0, u, bc, solver_parameters=params)
-    assert (errornorm(Constant((1, 0.5, 4)), u, 'L2') < 1e-10)
+    assert (errornorm(u_exact, u, 'L2') < 1e-10)
 
 
 @pytest.mark.skiphypre

tests/firedrake/regression/test_hypre_ams.py

@@ -3,107 +3,60 @@
 from firedrake import *
 
 
-@pytest.mark.skiphypre
-@pytest.mark.skipcomplex
-def test_homogeneous_field_linear():
-    mesh = UnitCubeMesh(5, 5, 5)
-    V = FunctionSpace(mesh, "N1curl", 1)
-    V0 = VectorFunctionSpace(mesh, "DG", 0)
-
-    u = TrialFunction(V)
-    v = TestFunction(V)
-
-    a = inner(curl(u), curl(v))*dx
-    L = inner(Constant((0., 0., 0.)), v)*dx
-
-    x, y, z = SpatialCoordinate(mesh)
-    B0 = 1
-    constant_field = as_vector([-0.5*B0*(y - 0.5), 0.5*B0*(x - 0.5), 0])
-
-    bc = DirichletBC(V, constant_field, (1, 2, 3, 4))
-
-    params = {'snes_type': 'ksponly',
-              'ksp_type': 'cg',
-              'ksp_max_it': '30',
-              'ksp_rtol': '2e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.HypreAMS',
-              'pc_hypre_ams_zero_beta_poisson': True
-              }
-
-    A = Function(V)
-    solve(a == L, A, bc, solver_parameters=params)
-    B = project(curl(A), V0)
-    assert numpy.allclose(B.dat.data_ro, numpy.array((0., 0., 1.)), atol=1e-6)
+@pytest.fixture(params=["simplex", "hexahedron"])
+def V(request):
+    cell = request.param
+    if cell == "simplex":
+        mesh = UnitCubeMesh(5, 5, 5)
+        V = FunctionSpace(mesh, "N1curl", 1)
+    elif cell == "hexahedron":
+        mesh = ExtrudedMesh(UnitSquareMesh(5, 5, quadrilateral=True), 5)
+        V = FunctionSpace(mesh, "NCE", 1)
+    else:
+        raise ValueError(f"Unrecognized cell {cell}.")
+    return V
 
 
 @pytest.mark.skiphypre
 @pytest.mark.skipcomplex
-def test_homogeneous_field_matfree():
-    mesh = UnitCubeMesh(5, 5, 5)
-    V = FunctionSpace(mesh, "N1curl", 1)
-    V0 = VectorFunctionSpace(mesh, "DG", 0)
-
+@pytest.mark.parametrize("mat_type,interface", [("aij", "linear"), ("matfree", "linear"), ("aij", "nonlinear")])
+def test_homogeneous_field(V, mat_type, interface):
+    mesh = V.mesh()
     u = TrialFunction(V)
     v = TestFunction(V)
 
     a = inner(curl(u), curl(v))*dx
-    L = inner(Constant((0., 0., 0.)), v)*dx
+    L = inner(Constant([0, 0, 0]), v) * dx
 
     x, y, z = SpatialCoordinate(mesh)
     B0 = 1
     constant_field = as_vector([-0.5*B0*(y - 0.5), 0.5*B0*(x - 0.5), 0])
 
     bc = DirichletBC(V, constant_field, (1, 2, 3, 4))
 
-    params = {'snes_type': 'ksponly',
-              'mat_type': 'matfree',
-              'ksp_type': 'cg',
-              'ksp_max_it': '30',
-              'ksp_rtol': '2e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.AssembledPC',
-              'assembled_pc_type': 'python',
-              'assembled_pc_python_type': 'firedrake.HypreAMS',
-              'assembled_pc_hypre_ams_zero_beta_poisson': True
-              }
+    params = {
+        'snes_type': 'ksponly',
+        'mat_type': mat_type,
+        'pmat_type': 'aij',
+        'ksp_type': 'cg',
+        'ksp_max_it': '30',
+        'ksp_rtol': '2e-15',
+        'pc_type': 'python',
+        'pc_python_type': 'firedrake.HypreAMS',
+        'pc_hypre_ams_zero_beta_poisson': True
+    }
 
     A = Function(V)
-    solve(a == L, A, bc, solver_parameters=params)
-    B = project(curl(A), V0)
-    assert numpy.allclose(B.dat.data_ro, numpy.array((0., 0., 1.)), atol=1e-6)
-
+    if interface == "linear":
+        solve(a == L, A, bc, solver_parameters=params)
+    elif interface == "nonlinear":
+        F = action(a, A) - L
+        solve(F == 0, A, bc, solver_parameters=params)
+    else:
+        raise ValueError(f"Unrecognized interface {interface}.")
 
-@pytest.mark.skiphypre
-@pytest.mark.skipcomplex
-def test_homogeneous_field_nonlinear():
-    mesh = UnitCubeMesh(5, 5, 5)
-    V = FunctionSpace(mesh, "N1curl", 1)
     V0 = VectorFunctionSpace(mesh, "DG", 0)
-
-    u = Function(V)
-    v = TestFunction(V)
-
-    a = inner(curl(u), curl(v))*dx
-    L = inner(Constant((0., 0., 0.)), v)*dx
-
-    x, y, z = SpatialCoordinate(mesh)
-    B0 = 1
-    constant_field = as_vector([-0.5*B0*(y - 0.5), 0.5*B0*(x - 0.5), 0])
-
-    bc = DirichletBC(V, constant_field, (1, 2, 3, 4))
-
-    params = {'snes_type': 'ksponly',
-              'ksp_type': 'cg',
-              'ksp_itmax': '30',
-              'ksp_rtol': '2e-15',
-              'pc_type': 'python',
-              'pc_python_type': 'firedrake.HypreAMS',
-              'pc_hypre_ams_zero_beta_poisson': True
-              }
-
-    solve(a - L == 0, u, bc, solver_parameters=params)
-    B = project(curl(u), V0)
+    B = project(curl(A), V0)
     assert numpy.allclose(B.dat.data_ro, numpy.array((0., 0., 1.)), atol=1e-6)