From 8c657807c0aa9da9bed3c808712f36bae130ebb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= Date: Sat, 21 Sep 2024 19:27:08 +0200 Subject: [PATCH] Fix argument block-size in Expressions (#3429) * Block size of argument was not taken into account when argument has block size * Add test that segfaults on main * Remove debug print * Ruff format --- cpp/dolfinx/fem/Expression.h | 3 +-- python/test/unit/fem/test_expression.py | 36 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/cpp/dolfinx/fem/Expression.h b/cpp/dolfinx/fem/Expression.h index 5e68f6be44c..f519cca5cd2 100644 --- a/cpp/dolfinx/fem/Expression.h +++ b/cpp/dolfinx/fem/Expression.h @@ -199,7 +199,7 @@ class Expression num_argument_dofs = _argument_function_space->dofmap()->element_dof_layout().num_dofs(); auto element = _argument_function_space->element(); - + num_argument_dofs *= _argument_function_space->dofmap()->bs(); assert(element); if (element->needs_dof_transformations()) { @@ -244,7 +244,6 @@ class Expression std::ranges::fill(values_local, 0); _fn(values_local.data(), coeff_cell, constant_data.data(), coord_dofs.data(), entity_index, nullptr); - post_dof_transform(values_local, cell_info, e, size0); for (std::size_t j = 0; j < values_local.size(); ++j) values[e * vshape[1] + j] = values_local[j]; diff --git a/python/test/unit/fem/test_expression.py b/python/test/unit/fem/test_expression.py index 25c729c512b..55ea54f8b67 100644 --- a/python/test/unit/fem/test_expression.py +++ b/python/test/unit/fem/test_expression.py @@ -486,3 +486,39 @@ def test_facet_expression(dtype): exact_expr = 2 * (midpoint[1] + midpoint[0]) * np.dot(grad_u, exact_n) assert np.allclose(values, exact_expr, atol=atol) + + +def test_rank1_blocked(): + """ + Check that a test function with tensor shape is unrolled as + (num_cells, num_points, num_dofs, bs) when evaluated as an expression + """ + mesh = dolfinx.mesh.create_unit_square( + MPI.COMM_SELF, 1, 1, cell_type=dolfinx.mesh.CellType.quadrilateral + ) + value_shape = (3, 2) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, value_shape)) + v = ufl.TestFunction(V) + + points = np.array([[0.513, 0.317], [0.11, 0.38]], dtype=mesh.geometry.x.dtype) + expr = dolfinx.fem.Expression(v, points) + + values = expr.eval(mesh, np.array([0], dtype=np.int32))[0] + # Tabulate returns (num_derivatives, num_points, num_dofs, value_size) + ref_values = V.element.basix_element.tabulate(1, points)[0] + + num_points = points.shape[0] + num_dofs = V.dofmap.dof_layout.num_dofs + bs = V.dofmap.bs + assert len(values) == num_dofs * num_points * bs * bs + for i in range(num_points): + # Get basis functions for all blocks for ith point + point_values = values[i * num_dofs * bs * bs : (i + 1) * num_dofs * bs * bs] + for j in range(bs): + block_values = point_values[j * num_dofs * bs : (j + 1) * num_dofs * bs] + for k in range(bs): + # Test functions are evaluated as a vector function + if j != k: + np.testing.assert_allclose(block_values[k::bs], 0.0) + else: + np.testing.assert_allclose(block_values[k::bs], ref_values[i].flatten())