Enable V.sub(i) for blocked real-spaces

[jorgensd]

Resolve #4259.
Tensor spaces, e.g. (2, 3), (3,4,2, ...) are flattened so that one can only do V.sub(i) not V.sub(i).sub(j).

Additionally this massively simplifies the dofmap constructor for real elements by using the standard element dof layout creator.

Adds test for symmetric real spaces as well.

[garth-wells]

It looks slightly odd having a loop over the value_size. I would have expected to see a recursive algorithm. Is the right fix elsewhere? Maybe avoiding premature flattening?

[jorgensd]

It looks slightly odd having a loop over the value_size. I would have expected to see a recursive algorithm. Is the right fix elsewhere? Maybe avoiding premature flattening?

We do the exact same thing for generalized blocked spaces:

from mpi4py import MPI
import basix.ufl
import dolfinx
import numpy as np
r_el = basix.ufl.element("Lagrange", "triangle", degree=1)
v_s = (2, 3)
b_el = basix.ufl.blocked_element(r_el, shape=v_s)


mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 1, 1)
V = dolfinx.fem.functionspace(mesh, b_el)
assert V.value_shape == v_s
for i in range(int(np.prod(v_s))):
    V.sub(i).collapse()
    print(f"Value shape of subspace {i}: {V.sub(i).value_shape}")
print(f"Value shape of blocked space: {V.value_shape}")

which yields:

Value shape of subspace 0: ()
Value shape of subspace 1: ()
Value shape of subspace 2: ()
Value shape of subspace 3: ()
Value shape of subspace 4: ()
Value shape of subspace 5: ()
Value shape of blocked space: (2, 3)

So this simply follows the standard that we use.
You can see this in the creation of the element dof layout:

dolfinx/cpp/dolfinx/fem/utils.h

Lines 325 to 348 in b0b4b19

	int bs = element.block_size();
	for (int i = 0; i < element.num_sub_elements(); ++i)
	{
	// The ith sub-element. For mixed elements this is subelements()[i]. For
	// blocked elements, the sub-element will always be the same, so we'll use
	// sub_elements()[0]
	std::shared_ptr<const fem::FiniteElement<T>> sub_e
	= element.sub_elements()[bs > 1 ? 0 : i];
	// In a mixed element DOFs are ordered element by element, so the offset to
	// the next sub-element is sub_e->space_dimension(). Blocked elements use
	// xxyyzz ordering, so the offset to the next sub-element is 1
	std::vector<int> parent_map_sub(sub_e->space_dimension(), offsets.back());
	for (std::size_t j = 0; j < parent_map_sub.size(); ++j)
	parent_map_sub[j] += bs * j;
	offsets.push_back(offsets.back() + (bs > 1 ? 1 : sub_e->space_dimension()));
	sub_doflayout.push_back(
	dolfinx::fem::create_element_dof_layout(*sub_e, parent_map_sub));
	}
	return ElementDofLayout(bs, element.entity_dofs(),
	element.entity_closure_dofs(), parent_map,
	sub_doflayout);

[michalhabera]

It is very easy to get confused with many notions of value_shape, value_size, num_sub_elements, and block_size. And for Real element many of them are the same, so that confuses code reader even more.

It seems to me that what you do is the right thing, with one small comment: should the construction of Real element only ever consider element.num_sub_elements and pass that in here

dolfinx/python/dolfinx/fem/function.py

Lines 687 to 692 in b0b4b19

	cpp_dofmap = _cpp.fem.build_real_element_dofmap(
	mesh.topology._cpp_object,
	element.basix_element.entity_dofs, # type: ignore
	element.basix_element.entity_closure_dofs, # type: ignore
	int(np.prod(element.value_shape)), # type: ignore
	)

?
Similarly then the C++ function would have int num_sub_elements and the loop over that. No value_shape and no value_size in these codepaths. Then you match the generic dofmap builder even more closely.

Maybe this is what confused @garth-wells ?

(I've spent a good hour trying to understand some relations, and we should simplify some concepts in few places. Or at least explain them better.)

[jorgensd]

It is very easy to get confused with many notions of value_shape, value_size, num_sub_elements, and block_size. And for Real element many of them are the same, so that confuses code reader even more.

It seems to me that what you do is the right thing, with one small comment: should the construction of Real element only ever consider element.num_sub_elements and pass that in here

dolfinx/python/dolfinx/fem/function.py

Lines 687 to 692 in b0b4b19

	cpp_dofmap = _cpp.fem.build_real_element_dofmap(
	mesh.topology._cpp_object,
	element.basix_element.entity_dofs, # type: ignore
	element.basix_element.entity_closure_dofs, # type: ignore
	int(np.prod(element.value_shape)), # type: ignore
	)

?
Similarly then the C++ function would have int num_sub_elements and the loop over that. No value_shape and no value_size in these codepaths. Then you match the generic dofmap builder even more closely.
Maybe this is what confused @garth-wells ?

(I've spent a good hour trying to understand some relations, and we should simplify some concepts in few places. Or at least explain them better.)

I guess we can assert that element.num_sub_elements is equal to int(np.prod(element.value_shape)) as anything else would mean that something has gone horribly wrong in construction of the element?

[Copilot]

🟢 Ready to approve

The change is narrowly scoped to Real-element dof layout construction and is backed by a targeted regression test that exercises the reported failure mode.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

This PR fixes subspace extraction for blocked/vector-valued Real elements by ensuring the C++ Real-element dof layout exposes per-component sub-layouts, allowing FunctionSpace.sub(i) to work as expected (resolving #4259).

Changes:

• Build ElementDofLayout sub-layouts for Real elements based on value_size, so extract_sub_dofmap([i]) no longer fails with “Invalid component”.
• Add a unit test covering scalar, length-1, vector, and tensor-shaped Real elements, validating V.sub(i) and collapse mapping behavior.

File summaries

File	Description
cpp/dolfinx/fem/dofmapbuilder.cpp	Constructs per-component ElementDofLayout sub-layouts for Real elements so V.sub(i) is valid for blocked Real spaces.
python/test/unit/fem/test_real_space.py	Adds regression tests for subspace extraction/collapse on Real elements with various value_shapes.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Not ready to approve

The newly added symmetric-space test computes num_dofs as a float (via / 2 and np.sum), making assertions and downstream uses brittle; it should be computed as an integer.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (1)

python/test/unit/fem/test_real_space.py:171

• num_dofs is computed via / 2 and wrapped in np.sum, which produces a float (e.g. 6.0). This makes num_sub_elements a float and relies on implicit int/float equality in assertions and on NumPy accepting float stops in arange, which is brittle and can break if the expression ever stops being exactly representable as a float. Compute num_dofs as an integer instead.

    num_dofs = np.sum(vs[0] * (vs[0] + 1) / 2)
    num_sub_elements = num_dofs if vs[0] > 1 else 0
    assert V.num_sub_spaces == num_sub_elements
    assert V.dofmap.index_map.size_global * V.dofmap.index_map_bs == num_dofs

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟢 Ready to approve

The core API/binding updates are consistent across C++ and Python and the PR adds targeted regression tests; only a minor test robustness nit was identified.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (1)

python/test/unit/fem/test_real_space.py:170

• num_dofs is computed via NumPy sum and / 2, which yields a float. This makes num_sub_elements a float and relies on implicit float↔int comparisons and array sizing; it’s clearer and more robust to keep this as an int using integer arithmetic.

    num_dofs = np.sum(vs[0] * (vs[0] + 1) / 2)
    num_sub_elements = num_dofs if vs[0] > 1 else 0
    assert V.num_sub_spaces == num_sub_elements

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The new tests need minor formatter cleanup and the symmetric Real-space test should also exercise subspace extraction/collapse to fully cover the target regression.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (3)

python/test/unit/fem/test_real_space.py:158

• Extra blank line at the start of this function body is inconsistent with the rest of the file and will be removed by ruff format; drop it to keep the test file formatter-clean.

def test_symmetric(vs, ftype):

python/test/unit/fem/test_real_space.py:126

• Extra blank line at the start of this function body will be removed by ruff format (Black-style) and can cause formatting CI failures; remove the empty line so the file is already formatted.

This issue also appears on line 157 of the same file.

def test_real_sub_spaces(vs, ftype):

python/test/unit/fem/test_real_space.py:172

• This test validates symmetry of evaluation but does not exercise the behavior this PR targets (extracting/collapsing subspaces for blocked/symmetric real elements). Add a simple V.sub(i)/collapse() loop so symmetric real spaces also cover the #4259 regression path.

    assert V.num_sub_spaces == num_sub_elements
    assert V.dofmap.index_map.size_global * V.dofmap.index_map_bs == num_dofs

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The newly added symmetric Real-space test can error on MPI ranks with no local/ghost cells due to reshaping an empty Expression.eval result.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (1)

python/test/unit/fem/test_real_space.py:182

• expr.eval(mesh, cell) returns an empty array when cell is empty (on ranks with no local cells), so .reshape(vs) will raise. Guard the evaluation so it only runs when there is at least one local/ghost cell, and index into the (cell, point) axes before reshaping.

        np.zeros(1, dtype=np.int32)
        if mesh.topology.index_map(mesh.topology.dim).size_local > 0
        else np.zeros(0, dtype=np.int32)
    )
    values = expr.eval(mesh, cell).reshape(vs)

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The new tests should be adjusted to follow existing test conventions (formatting and scatter_forward() after coefficient writes) to avoid CI/parallel robustness issues.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (2)

python/test/unit/fem/test_real_space.py:166

• The real_element(...) calls in this block are long enough to be hard to read and are likely to exceed the repo’s Python formatting conventions (most similar calls in this file are wrapped). Wrapping them avoids formatter/CI churn and improves readability.

    if vs == ():
        with pytest.raises(ValueError):
            basix.ufl.real_element(mesh.basix_cell(), value_shape=vs, dtype=ftype, symmetry=True)
        return
    el = basix.ufl.real_element(mesh.basix_cell(), value_shape=vs, dtype=ftype, symmetry=True)

python/test/unit/fem/test_real_space.py:176

• After writing into u.x.array, the test should call u.x.scatter_forward() to ensure ghost values are consistent across ranks before Expression.eval reads the coefficient vector. This file already follows that pattern in test_complex_real_space.

    u = dolfinx.fem.Function(V, dtype=ftype)
    u.x.array[:] = np.arange(num_dofs, dtype=ftype)
    expr = dolfinx.fem.Expression(u, np.array([[0.0, 0.0]], dtype=ftype))

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[Copilot]

🟡 Changes recommended

The updated Python files include long lines that are likely to fail the repo’s enforced ruff formatting/line-length checks.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Review details

Suppressed comments (2)

python/test/unit/fem/test_real_space.py:166

• Two calls to basix.ufl.real_element are written on a single long line; this will likely violate the repo's ruff formatting/line-length constraints and makes the symmetry=True cases harder to read.

    if vs == ():
        with pytest.raises(ValueError):
            basix.ufl.real_element(mesh.basix_cell(), value_shape=vs, dtype=ftype, symmetry=True)
        return
    el = basix.ufl.real_element(mesh.basix_cell(), value_shape=vs, dtype=ftype, symmetry=True)

python/dolfinx/fem/function.py:709

• The real-element dofmap creation lines exceed typical ruff formatting line-length limits and reduce readability; ruff/black-style wrapping will keep the file formatted consistently.

    if ufl_e.is_real:
        dof_layout = _cpp.fem.create_element_dof_layout(dolfinx_element._cpp_object, [])
        cpp_dofmap = _cpp.fem.build_real_element_dofmap(mesh.topology._cpp_object, dof_layout)

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.

[michalhabera]

Much nicer with common create_element_dof_layout.