diff --git a/.github/workflows/fenicsx-refs.env b/.github/workflows/fenicsx-refs.env index 8e958e229c3..b378b8296e2 100644 --- a/.github/workflows/fenicsx-refs.env +++ b/.github/workflows/fenicsx-refs.env @@ -1,4 +1,4 @@ basix_ref=main -ufl_ref=main -ffcx_ref=main +ufl_ref=dokken/multiple_coordinate_elements +ffcx_ref=mscroggs/ufl_coordinate_elements diff --git a/python/demo/demo_mixed-topology.py b/python/demo/demo_mixed-topology.py index 292d7584060..2aede11035a 100644 --- a/python/demo/demo_mixed-topology.py +++ b/python/demo/demo_mixed-topology.py @@ -24,16 +24,17 @@ import basix import dolfinx.cpp as _cpp import ufl -from dolfinx.cpp.mesh import GhostMode, create_cell_partitioner, create_mesh +from dolfinx.cpp.mesh import GhostMode, create_cell_partitioner +from dolfinx.mesh import create_mesh from dolfinx.fem import ( FunctionSpace, assemble_matrix, assemble_vector, - coordinate_element, mixed_topology_form, ) from dolfinx.io.utils import cell_perm_vtk from dolfinx.mesh import CellType, Mesh +from dolfinx import fem if MPI.COMM_WORLD.size > 1: print("Not yet running in parallel") @@ -89,48 +90,32 @@ cells_np = [np.array(c) for c in cells] geomx = np.array(geom, dtype=np.float64) -hexahedron = coordinate_element(CellType.hexahedron, 1) -prism = coordinate_element(CellType.prism, 1) part = create_cell_partitioner(GhostMode.none) mesh = create_mesh( - MPI.COMM_WORLD, cells_np, [hexahedron._cpp_object, prism._cpp_object], geomx, part + MPI.COMM_WORLD, + cells_np, + geomx, + [ + basix.ufl.element("Lagrange", "hexahedron", 1, shape=(3,)), + basix.ufl.element("Lagrange", "prism", 1, shape=(3,)), + ], + part, ) -# Create elements and dofmaps for each cell type -elements = [ - basix.create_element(basix.ElementFamily.P, basix.CellType.hexahedron, 1), - basix.create_element(basix.ElementFamily.P, basix.CellType.prism, 1), -] -elements_cpp = [_cpp.fem.FiniteElement_float64(e._e, None, True) for e in elements] -# NOTE: Both dofmaps have the same IndexMap, but different cell_dofs -dofmaps = _cpp.fem.create_dofmaps(mesh.comm, mesh.topology, elements_cpp) - -# Create C++ function space -V_cpp = _cpp.fem.FunctionSpace_float64(mesh, elements_cpp, dofmaps) - -# Create forms for each cell type. -# FIXME This hack is required at the moment because UFL does not yet know -# about mixed topology meshes. -a = [] -L = [] -for i, cell_name in enumerate(["hexahedron", "prism"]): - print(f"Creating form for {cell_name}") - element = basix.ufl.wrap_element(elements[i]) - domain = ufl.Mesh(basix.ufl.element("Lagrange", cell_name, 1, shape=(3,))) - V = FunctionSpace(Mesh(mesh, domain), element, V_cpp) - u, v = ufl.TrialFunction(V), ufl.TestFunction(V) - k = 12.0 - x = ufl.SpatialCoordinate(domain) - a += [(ufl.inner(ufl.grad(u), ufl.grad(v)) - k**2 * u * v) * ufl.dx] - f = ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1]) - L += [f * v * ufl.dx] +V = fem.functionspace(mesh, ("Lagrange", 1)) + + +u, v = ufl.TrialFunction(V), ufl.TestFunction(V) +k = 12.0 +x = ufl.SpatialCoordinate(domain) +a = (ufl.inner(ufl.grad(u), ufl.grad(v)) - k**2 * u * v) * ufl.dx +f = ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1]) +L = f * v * ufl.dx # Compile the form -# FIXME: For the time being, since UFL doesn't understand mixed topology -# meshes, we have to call mixed_topology_form instead of form. -a_form = mixed_topology_form(a, dtype=np.float64) -L_form = mixed_topology_form(L, dtype=np.float64) +a_form = form(a, dtype=np.float64) +L_form = form(L, dtype=np.float64) # Assemble the matrix A = assemble_matrix(a_form) diff --git a/python/dolfinx/fem/function.py b/python/dolfinx/fem/function.py index 118633fc1b4..ac14b8d61fe 100644 --- a/python/dolfinx/fem/function.py +++ b/python/dolfinx/fem/function.py @@ -583,6 +583,7 @@ def functionspace( mesh: Mesh, element: typing.Union[ ufl.finiteelement.AbstractFiniteElement, + Sequence[ufl.finiteelement.AbstractFiniteElement], ElementMetaData, tuple[str, int], tuple[str, int, tuple], @@ -600,22 +601,46 @@ def functionspace( """ # Create UFL element dtype = mesh.geometry.x.dtype - try: - e = ElementMetaData(*element) # type: ignore - ufl_e = basix.ufl.element( - e.family, - mesh.basix_cell(), # type: ignore - e.degree, - shape=e.shape, - symmetry=e.symmetry, - dtype=dtype, - ) - except TypeError: - ufl_e = element # type: ignore + if len(mesh.topology._cpp_object.cell_types) > 1: + try: + e = ElementMetaData(*element) # type: ignore + ufl_e = [ + basix.ufl.element( + e.family, + getattr(basix.CellType, _cpp.mesh.to_string(cell)), + e.degree, + shape=e.shape, + symmetry=e.symmetry, + dtype=dtype, + ) + for cell in mesh.topology._cpp_object.cell_types + ] + except TypeError: + ufl_e = element # type: ignore + + # Check that element and mesh cell types match + for domain, element in zip(mesh.ufl_domain(), ufl_e): + if domain is None or element.cell != domain.ufl_cell(): + raise ValueError("Non-matching UFL cell and mesh cell shapes.") + else: + try: + ufl_e = basix.ufl.element( + e.family, + mesh.basix_cell(), # type: ignore + e.degree, + shape=e.shape, + symmetry=e.symmetry, + dtype=dtype, + ) + except TypeError: + ufl_e = element # type: ignore + + # Check that element and mesh cell types match + if ((domain := mesh.ufl_domain()) is None) or ufl_e.cell != domain.ufl_cell(): + raise ValueError("Non-matching UFL cell and mesh cell shapes.") - # Check that element and mesh cell types match - if ((domain := mesh.ufl_domain()) is None) or ufl_e.cell != domain.ufl_cell(): - raise ValueError("Non-matching UFL cell and mesh cell shapes.") + # TODO + # TODO # Create DOLFINx objects element = finiteelement(mesh.topology.cell_type, ufl_e, dtype) # type: ignore diff --git a/python/dolfinx/mesh.py b/python/dolfinx/mesh.py index b5086a0b86b..4c40a66a30e 100644 --- a/python/dolfinx/mesh.py +++ b/python/dolfinx/mesh.py @@ -289,7 +289,7 @@ class Mesh: def __init__( self, msh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64], - domain: typing.Optional[ufl.Mesh], + domain: typing.Optional[typing.Union[ufl.Mesh, Sequence[ufl.Mesh]]], ): """Initialize mesh from a C++ mesh. @@ -307,8 +307,13 @@ def __init__( self._topology = Topology(self._cpp_object.topology) self._geometry = Geometry(self._cpp_object.geometry) self._ufl_domain = domain - if self._ufl_domain is not None: - self._ufl_domain._ufl_cargo = self._cpp_object # type: ignore + try: + for d in self._ufl_domain: + if d is not None: + d._ufl_cargo = self._cpp_object # type: ignore + except TypeError: + if self._ufl_domain is not None: + self._ufl_domain._ufl_cargo = self._cpp_object # type: ignore @property def comm(self): @@ -680,15 +685,60 @@ def refine( return Mesh(mesh1, ufl_domain), parent_cell, parent_facet +def extract_cmap_and_domain( + e: typing.Union[ + ufl.Mesh, + basix.finite_element.FiniteElement, + basix.ufl._BasixElement, + _CoordinateElement, + ], + gdim: int, +) -> tuple[typing.Any, typing.Any]: + try: + # e is a UFL domain + e_ufl = e.ufl_coordinate_element() + return _coordinate_element(e_ufl.basix_element), e + except AttributeError: + pass + + try: + # e is a Basix 'UFL' element + domain = ufl.Mesh(e) + assert domain.geometric_dimension == gdim + return _coordinate_element(e.basix_element), domain + except (AttributeError, TypeError): + pass + try: + # e is a Basix element + # TODO: Resolve geometric dimension vs shape for manifolds + e_ufl = basix.ufl._BasixElement(e) # type: ignore + e_ufl = basix.ufl.blocked_element(e_ufl, shape=(gdim,)) + domain = ufl.Mesh(e_ufl) + assert domain.geometric_dimension == gdim + return _coordinate_element(e), domain # type: ignore + except (AttributeError, TypeError): + pass + # e is a CoordinateElement + return e, None + + def create_mesh( comm: _MPI.Comm, - cells: npt.NDArray[np.int64], + cells: typing.Union[ + npt.NDArray[np.int64], + typing.Sequence[npt.NDArray[np.int64]], + ], x: npt.NDArray[np.floating], e: typing.Union[ ufl.Mesh, basix.finite_element.FiniteElement, basix.ufl._BasixElement, _CoordinateElement, + Sequence[ + typing.Union[ + basix.finite_element.FiniteElement, basix.ufl._BasixElement, _CoordinateElement + ] + ], ], partitioner: typing.Optional[Callable] = None, ) -> Mesh: @@ -696,7 +746,7 @@ def create_mesh( Args: comm: MPI communicator to define the mesh on. - cells: Cells of the mesh. ``cells[i]`` are the 'nodes' of cell + cells: Cells of the mesh. ``cells[i]`` are the 'nodes' of cell ``i``. x: Mesh geometry ('node' coordinates), with shape ``(num_nodes, gdim)``. @@ -721,37 +771,34 @@ def create_mesh( else: gdim = x.shape[1] - dtype = None try: - # e is a UFL domain - e_ufl = e.ufl_coordinate_element() # type: ignore - cmap = _coordinate_element(e_ufl.basix_element) # type: ignore - domain = e - dtype = cmap.dtype - # TODO: Resolve UFL vs Basix geometric dimension issue - # assert domain.geometric_dimension() == gdim - except AttributeError: - try: - # e is a Basix 'UFL' element - cmap = _coordinate_element(e.basix_element) # type: ignore - domain = ufl.Mesh(e) - dtype = cmap.dtype - assert domain.geometric_dimension() == gdim - except AttributeError: - try: - # e is a Basix element - # TODO: Resolve geometric dimension vs shape for manifolds - cmap = _coordinate_element(e) # type: ignore - e_ufl = basix.ufl._BasixElement(e) # type: ignore - e_ufl = basix.ufl.blocked_element(e_ufl, shape=(gdim,)) - domain = ufl.Mesh(e_ufl) + # e and cells are Sequences + dtype = None + formatted_cells = [] + cmaps = [] + domains = [] + for i, c in zip(e, cells): + cmap, domain = extract_cmap_and_domain(i, gdim) + if dtype is None: dtype = cmap.dtype - assert domain.geometric_dimension() == gdim - except (AttributeError, TypeError): - # e is a CoordinateElement - cmap = e - domain = None - dtype = cmap.dtype # type: ignore + else: + assert dtype == cmap.dtype + c = np.asarray(c, dtype=np.int64, order="C") + formatted_cells.append(c) + cmaps.append(cmap._cpp_object) + domains.append(domain) + + x = np.asarray(x, dtype=dtype, order="C") + msh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64] = _cpp.mesh.create_mesh( + comm, formatted_cells, cmaps, x, partitioner + ) + return Mesh(msh, domains) + + except KeyboardInterrupt: + pass + + cmap, domain = extract_cmap_and_domain(e, gdim) + dtype = cmap.dtype x = np.asarray(x, dtype=dtype, order="C") cells = np.asarray(cells, dtype=np.int64, order="C") diff --git a/python/test/unit/mesh/test_mesh.py b/python/test/unit/mesh/test_mesh.py index 9c502dfa20f..6c6ef6360d0 100644 --- a/python/test/unit/mesh/test_mesh.py +++ b/python/test/unit/mesh/test_mesh.py @@ -248,7 +248,7 @@ def test_UFLCell(interval, square, rectangle, cube, box): def test_UFLDomain(interval, square, rectangle, cube, box): def _check_ufl_domain(mesh): domain = mesh.ufl_domain() - assert mesh.geometry.dim == domain.geometric_dimension() + assert mesh.geometry.dim == domain.geometric_dimension assert mesh.topology.dim == domain.topological_dimension() assert mesh.ufl_cell() == domain.ufl_cell()