Fix BDDCPC/FDMPC pyop3

firedrake/preconditioners/bddc.py

@@ -112,11 +112,11 @@ def initialize(self, pc):
             dir_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs]))
         neu_nodes = numpy.setdiff1d(boundary_nodes, dir_nodes)
 
-        dir_nodes = V.dof_dset.lgmap.apply(dir_nodes)
+        dir_nodes = V.lgmap().apply(dir_nodes)
         dir_bndr = PETSc.IS().createGeneral(dir_nodes, comm=pc.comm)
         bddcpc.setBDDCDirichletBoundaries(dir_bndr)
 
-        neu_nodes = V.dof_dset.lgmap.apply(neu_nodes)
+        neu_nodes = V.lgmap().apply(neu_nodes)
         neu_bndr = PETSc.IS().createGeneral(neu_nodes, comm=pc.comm)
         bddcpc.setBDDCNeumannBoundaries(neu_bndr)
 
@@ -238,7 +238,7 @@ def local_bc(bc, cellwise):
 
     def local_to_global_map(V, cellwise):
         u = Function(V)
-        u.dat.data_wo[:] = numpy.arange(*V.dof_dset.layout_vec.getOwnershipRange())
+        u.dat.data_wo[:] = numpy.arange(*V.template_vec.getOwnershipRange())
 
         Vsub = local_space(V, False)
         usub = Function(Vsub).assign(u)
@@ -278,7 +278,7 @@ def update():
 def get_restricted_dofs(V, domain):
     W = FunctionSpace(V.mesh(), V.ufl_element()[domain])
     indices = get_restriction_indices(V, W)
-    indices = V.dof_dset.lgmap.apply(indices)
+    indices = V.lgmap().apply(indices)
     return PETSc.IS().createGeneral(indices, comm=V.comm)
 
 
@@ -337,7 +337,7 @@ def get_primal_indices(V, primal_markers):
         else:
             raise ValueError(f"Expecting markers in either {V.ufl_element()} or DG(0).")
         primal_indices = numpy.flatnonzero(markers.dat.data >= 1E-12)
-        primal_indices += V.dof_dset.layout_vec.getOwnershipRange()[0]
+        primal_indices += V.template_vec.getOwnershipRange()[0]
     else:
         primal_indices = numpy.asarray(primal_markers, dtype=PETSc.IntType)
     return primal_indices

firedrake/preconditioners/fdm.py

@@ -16,7 +16,7 @@
 from firedrake.function import Function
 from firedrake.cofunction import Cofunction
 from firedrake.cython.dmcommon import get_preallocation
-from firedrake.parloops import par_loop
+from firedrake.parloops import par_loop, READ, WRITE
 from firedrake.ufl_expr import TestFunction, TestFunctions, TrialFunctions
 from firedrake.utils import IntType, ScalarType
 from firedrake.pack import pack
@@ -1631,14 +1631,14 @@ def broken_function(V, val):
     """Return a Function(V, val=val) interpolated onto the broken space."""
     W = V.broken_space()
     w = Function(W, dtype=val.dtype)
-    v = Function(V, val=val)
-    domain = "{[i]: 0 <= i < v.dofs}"
+    v = Function(V, val=val, dtype=val.dtype)
+    domain = "{[i,j]: 0 <= i < v.dofs and 0 <= j < %d}" % V.block_size
     instructions = """
-    for i
-        w[i] = v[i]
+    for i, j
+        w[i,j] = v[i,j]
     end
     """
-    par_loop((domain, instructions), ufl.dx, {'w': (w, op2.WRITE), 'v': (v, op2.READ)})
+    par_loop((domain, instructions), ufl.dx, {'w': (w, WRITE), 'v': (v, READ)})
     return w
 
 

tests/firedrake/regression/test_bddc.py

@@ -5,9 +5,6 @@
 from firedrake.petsc import DEFAULT_DIRECT_SOLVER
 
 
-pytest.skip(reason="pyop3 TODO", allow_module_level=True)
-
-
 @pytest.fixture
 def rg():
     return RandomGenerator(PCG64(seed=123456789))