k

pyop2/types/dat.py

@@ -779,15 +779,16 @@ def _vec(self):
 
     @utils.cached_property
     def _data_filtered(self):
-        size = self.dataset.layout_vec.getSizes()
-        data = self._data[:size[0]]
+        size, _ = self.dataset.layout_vec.getSizes()
+        size //= self.dataset.layout_vec.block_size
+        data = self._data[:size]
         return np.empty_like(data)
 
     @utils.cached_property
     def _data_filter(self):
         lgmap = self.dataset.lgmap
         n = self.dataset.size
-        lgmap_owned = lgmap.indices[:n]
+        lgmap_owned = lgmap.block_indices[:n]
         return lgmap_owned >= 0
 
     @contextlib.contextmanager

pyop2/types/dataset.py

@@ -503,6 +503,8 @@ def lgmap(self):
             tmp_indices = np.searchsorted(current_offsets, l2g, side="right") - 1
             idx[:] = l2g[:] - current_offsets[tmp_indices] + \
                 all_field_offsets[tmp_indices] + all_local_offsets[tmp_indices]
+            # Explicitly set -1 for constrained DoFs.
+            idx[l2g < 0] = -1
             self.comm.Allgather(owned_sz, current_offsets[1:])
             all_local_offsets += current_offsets[1:]
             start += s.total_size * s.cdim

tests/firedrake/regression/test_restricted_function_space.py

@@ -299,3 +299,36 @@ def test_restricted_function_space_extrusion_poisson(ncells):
     sol = Function(V_res)
     solve(a == L, sol, bcs=[bc])
     assert assemble(inner(sol - exact, sol - exact) * dx)**0.5 < 1.e-15
+
+
+@pytest.mark.parallel(nprocs=4)
+@pytest.mark.parametrize("ncells", [2, 16])
+def test_restricted_function_space_extrusion_stokes(ncells):
+    mesh = UnitIntervalMesh(ncells)
+    extm = ExtrudedMesh(mesh, ncells)
+    # Solve reference problem.
+    V = VectorFunctionSpace(extm, "CG", 2)
+    Q = FunctionSpace(extm, "CG", 1)
+    W = V * Q
+    u, p = TrialFunctions(W)
+    v, q = TestFunctions(W)
+    a = inner(2 * sym(grad(u)), grad(v)) * dx - inner(p, div(v)) * dx + inner(div(u), q) * dx
+    L = inner(as_vector([1., 1.]), v) * dx
+    bc = DirichletBC(W.sub(0), as_vector([0., 0.]), [1, 2, "bottom"])
+    sol = Function(W)
+    solve(a == L, sol, bcs=[bc])
+    # Solve problem on restricted space.
+    V_res = RestrictedFunctionSpace(V, boundary_set=[1, 2, "bottom"])
+    W_res = V_res * Q
+    u_res, p = TrialFunctions(W_res)
+    v_res, q = TestFunctions(W_res)
+    a_res = inner(2 * sym(grad(u_res)), grad(v_res)) * dx - inner(p, div(v_res)) * dx + inner(div(u_res), q) * dx
+    L_res = inner(as_vector([1., 1.]), v_res) * dx
+    bc_res = DirichletBC(W_res.sub(0), as_vector([0., 0.]), [1, 2, "bottom"])
+    sol_res = Function(W_res)
+    solve(a_res == L_res, sol_res, bcs=[bc_res])
+    # Compare.
+    assert assemble(inner(sol_res - sol, sol_res - sol) * dx)**0.5 < 1.e-15
+    # -- Actually, the ordering is the same.
+    assert np.allclose(sol_res.subfunctions[0].dat.data_ro_with_halos, sol.subfunctions[0].dat.data_ro_with_halos)
+    assert np.allclose(sol_res.subfunctions[1].dat.data_ro_with_halos, sol.subfunctions[1].dat.data_ro_with_halos)