[BUG]: Contradictory locate_dofs_geometrical error message

[twmr]

How to reproduce the bug

When locate_dofs_geometrical is used for a function space that is a mixed-functionspace or a subspace of a mixed-functionspace, contradictory error messages are output.

Minimal Example (Python)

import basix
import numpy as np
import ufl
from dolfinx import fem, mesh
from mpi4py import MPI

msh = mesh.create_unit_interval(MPI.COMM_WORLD, nx=4)
P0 = ufl.FiniteElement("Lagrange", msh.ufl_cell(), 1)
P1 = P0
V = fem.FunctionSpace(msh, P0 * P1)
n = V.dofmap.index_map.size_global
print(f"{n=}")

fem.locate_dofs_geometrical(
    V,
    lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 1.0),
)
# raises
#    bcs_dofs1 = fem.locate_dofs_geometrical(
#  File "/home/thomas/miniconda/envs/fenicsx/lib/python3.10/site-packages/dolfinx/fem/bcs.py", line 54, in locate_dofs_geometrical
#    return _cpp.fem.locate_dofs_geometrical(V._cpp_object, marker)
# RuntimeError: Cannot locate dofs geometrically for mixed space. Use subspaces.

fem.locate_dofs_geometrical(
    V.sub(0),
    lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 1.0),
)
# raise
#  File "/home/thomas/miniconda/envs/fenicsx/lib/python3.10/site-packages/dolfinx/fem/bcs.py", line 54, in locate_dofs_geometrical
#    return _cpp.fem.locate_dofs_geometrical(V._cpp_object, marker)
#RuntimeError: Cannot tabulate coordinates for a FunctionSpace that is a subspace.

Output (Python)

No response

Version

0.5.1

DOLFINx git commit

No response

Installation

No response

Additional information

No response

[jorgensd]

The point is that when you work with mixed spaces, and want to locate dofs using their dof coordinates (the geometrical method), you need to use a collapsed sub space: i.e.:

dofs = fem.locate_dofs_geometrical(
    (V.sub(0), V.sub(0).collapse()[0]),
    lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 1.0),
)
print(dofs)

works.