Extend testing of quadrature element

python/test/unit/fem/test_quadrature_elements.py

@@ -125,8 +125,11 @@ def test_interpolation_blocked(degree):
 
 
 def extract_diagonal(mat):
+    num_rows = mat._cpp_object.index_map(0).size_local
+    num_cols = mat._cpp_object.index_map(1).size_local
+    assert num_rows == num_cols, "Matrix must be square"
     bs = mat.block_size[0]
-    diag = np.empty(len(mat.indices) * bs)
+    diag = np.empty(num_rows * bs, dtype=mat.data.dtype)
     for row, (start, end) in enumerate(zip(mat.indptr[:-1], mat.indptr[1:])):
         for i in range(start, end):
             if mat.indices[i] == row:
@@ -135,19 +138,23 @@ def extract_diagonal(mat):
     return diag
 
 
-@pytest.mark.skip_in_parallel
+@pytest.mark.parametrize("degree", range(1, 4))
 @pytest.mark.parametrize("shape", [(), (1,), (2,), (3,), (4,), (2, 2), (3, 3)])
-def test_vector_element(shape):
+def test_vector_element(shape, degree):
+    """
+    Compare assembly into a vector with quadrature elements with the diagonal of
+    an assembled mass matrix with the same quadrature element.
+    """
     msh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
 
     dx_m = ufl.Measure(
         "dx",
         domain=msh,
-        metadata={"quadrature_degree": 1, "quadrature_scheme": "default"},
+        metadata={"quadrature_degree": degree, "quadrature_scheme": "default"},
     )
 
     Qe = basix.ufl.quadrature_element(
-        msh.topology.cell_name(), value_shape=shape, scheme="default", degree=1
+        msh.topology.cell_name(), value_shape=shape, scheme="default", degree=degree
     )
     Quad = dolfinx.fem.functionspace(msh, Qe)
     q_ = ufl.TestFunction(Quad)
@@ -156,7 +163,50 @@ def test_vector_element(shape):
     one.x.array[:] = 1.0
     mass_L_form = dolfinx.fem.form(ufl.inner(one, q_) * dx_m)
     mass_v = dolfinx.fem.assemble_vector(mass_L_form)
+    mass_v.scatter_reverse(dolfinx.la.InsertMode.add)
+    mass_v.scatter_forward()
     mass_a_form = dolfinx.fem.form(ufl.inner(dq, q_) * dx_m)
     mass_A = dolfinx.fem.assemble_matrix(mass_a_form)
+    mass_A.scatter_reverse()
+    num_owned_dofs = Quad.dofmap.index_map.size_local * Quad.dofmap.index_map_bs
+    assert np.allclose(extract_diagonal(mass_A), mass_v.array[:num_owned_dofs])
+
+
+@pytest.mark.parametrize("degree", range(1, 4))
+def test_quadrature_assembly(degree):
+    """
+    Test quadrature element against assembly with spatial coordinate and a fixed quadrature rule
+    """
+    msh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 7)
+    dx_m = ufl.Measure(
+        "dx",
+        domain=msh,
+        metadata={"quadrature_degree": degree, "quadrature_scheme": "default"},
+    )
+
+    Qe = basix.ufl.quadrature_element(
+        msh.topology.cell_name(), value_shape=(), scheme="default", degree=degree
+    )
+    Quad = dolfinx.fem.functionspace(msh, Qe)
+
+    V = dolfinx.fem.functionspace(msh, ("Lagrange", 1))
+    v = ufl.TestFunction(V)
+
+    def f(x):
+        return 1 + x[0] + x[1] ** degree
+
+    q = dolfinx.fem.Function(Quad)
+    q.interpolate(f)
+
+    L = ufl.inner(q, v) * dx_m
+    b = dolfinx.fem.assemble_vector(dolfinx.fem.form(L))
+    b.scatter_reverse(dolfinx.la.InsertMode.add)
+    b.scatter_forward()
+
+    x = ufl.SpatialCoordinate(msh)
+    L_ref = ufl.inner(f(x), v) * dx_m
+    b_ref = dolfinx.fem.assemble_vector(dolfinx.fem.form(L_ref))
+    b_ref.scatter_reverse(dolfinx.la.InsertMode.add)
+    b_ref.scatter_forward()
 
-    assert np.allclose(extract_diagonal(mass_A), mass_v.array)
+    np.testing.assert_allclose(b.array, b_ref.array)