diff --git a/firedrake/mesh.py b/firedrake/mesh.py index c8cd4fd0f4..27b9585fc3 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1431,6 +1431,341 @@ def cell_set(self): size = list(self._entity_classes[self.cell_dimension(), :]) return op2.Set(size, "Cells", comm=self.comm) + @cached_property + def cell_facet_neighbours(self): + """Map each cell to its neighbouring cells across local facets. + + Returns a :class:`pyop2.types.dat.Dat` with one entry per local facet of + each cell. For cell ``c`` and local facet ``i``, + ``cell_facet_neighbours[c][i]`` is the Firedrake cell number of the cell + adjacent to ``c`` across that facet, or ``-1`` if the facet lies on the + boundary. + """ + num_facets = self.ufl_cell().num_facets + + dset = op2.DataSet(self.cell_set, dim=num_facets) + + # Create a local numpy buffer to store the neighbours of each cell + cell_neighbours = np.full((self.num_cells(), num_facets), -1, dtype=IntType) + + # Populate the local buffer by iterating over the cells and querying the DMPlex for local cell information + plex = self.topology_dm + cStart, cEnd = plex.getHeightStratum(0) # range of DMPlex point numbers representing cells + + for c_plex_point in range(cStart, cEnd): + cid = self._cell_numbering.getOffset(c_plex_point) # get Firedrake cell ID of plex cell point number + + # Instead of iterating over the plex facet points, iterate over FIAT's local facet IDs + for lf in range(num_facets): + lf_plex_point = self._cell_facet_point(cid, lf) # get plex point number corresponding to facet lf + support = plex.getSupport(lf_plex_point) + if len(support) == 2: + # an interior facet has 2 adjacent cells + # so the neighbouring cell corresponds to the adjacent cell that's not the current cell + other_c_plex_point = support[0] if support[1] == c_plex_point else support[1] + cell_neighbours[cid, lf] = self._cell_numbering.getOffset(other_c_plex_point) + + cell_facet_neighbours = op2.Dat( + dset, + cell_neighbours, + dtype=IntType, + name="cell-facet-neighbours-dat" + ) + + # NOTE: A fresh Dat has its halo entries marked as invalid -> next with_halos access can trigger a collective exchange. + # An exchange can overrite a rank's data with the owning rank's local data which is wrong. + # We avoid that by marking halos as valid. + cell_facet_neighbours.halo_valid = True + + return cell_facet_neighbours + + # helper methods + + @cached_property + def _cell_closure_facet_offset(self): + """Returns the correct column in cell_closure for a given local facet.""" + facet_dim = self.cell_dimension() - 1 + topology = FIAT.ufc_cell(self.ufl_cell()).get_topology() + return sum(len(topology[d]) for d in range(facet_dim)) + + def _cell_facet_point(self, cell, local_facet): + """Returns the DMPlex point for a local facet of a cell.""" + return self.cell_closure[cell, self._cell_closure_facet_offset + local_facet] + + def _get_facet_embedding_maps(self): + """ + For every local facet `lf` of the reference cell, FIAT provides an entity transform + that maps facet-local reference coordinates (X_lf) into the cell's reference coordinates (X): + + X = A_lf @ X_lf + b_lf. + + Returns + ------- + tuple of lists (A, b, A_inv) each indexed by local facet ID. + (A[lf], b[lf]) defines the affine embedding of facet reference coordinates into the cell reference frame, + and A_inv[lf] defines the pseudoinverse for mapping cell coordinates on the facet back to facet-local coordinates. + """ + import FIAT + ref_cell = FIAT.ufc_cell(self.ufl_cell()) + facet_dim = ref_cell.get_spatial_dimension() - 1 + num_facets = len(ref_cell.get_topology()[facet_dim]) + + A = [] + b = [] + A_inv = [] + + if facet_dim == 0: + vertices = np.asarray(ref_cell.get_vertices()) + for lf in range(num_facets): + A_lf = np.empty((ref_cell.get_spatial_dimension(), 0), dtype=RealType) + b_lf = np.asarray(vertices[lf], dtype=RealType) + A.append(A_lf) + b.append(b_lf) + A_inv.append(np.empty((0, ref_cell.get_spatial_dimension()), dtype=RealType)) + return A, b, A_inv + + for lf in range(num_facets): + phi = ref_cell.get_entity_transform(facet_dim, lf) + facet_on_ref_cell = ref_cell.construct_subelement(facet_dim) + facet_verts = facet_on_ref_cell.get_vertices() + + # Get (A_lf, b_lf) by evaluating the transform at the facet vertices + mapped_facet_verts = np.array([phi(v) for v in facet_verts]) + b_lf = mapped_facet_verts[0] # first vertex is always the origin + A_lf = (mapped_facet_verts[1:facet_dim+1] - b_lf).T + A.append(A_lf) + b.append(b_lf) + A_inv.append(np.linalg.pinv(A_lf)) + + return (A, b, A_inv) + + def _get_facet_orientation_coord_maps(self): + """ + When two adjacent cells share a facet, the ordering of the facet + vertices (and hence facet-local reference coordinates) may differ between + the two cells. This difference is encoded in the local facet orientation. + + This method constructs, for every possible facet orientation `o`, an affine + map acting on facet-local reference coordinates: + + X_f_o = Q_o @ X_f_can + t_o, + + where X_f_can defines the facet-local coordinates when the facet is in the canonical orientation and + X_f_o defines the facet-local coordinates of the same point when the facet is in orientation `o`. + + The FIAT `make_entity_permutations_simplex` method can be used to derive vertex permutations + for each possible facet permutation. We convert each such permutation into the corresponding affine coordinate transform using + `FIAT.reference_element.make_affine_mapping`. + + For tensor-product facets (e.g. quadrilateral facets of a hexahedra), FIAT + returns orientation keys as tuples (eo, i0, ..., ik), which are converted + into Firedrake's integer orientation before storing. + + Returns + ------- + dict + Mapping from facet orientation integer to (Q, t), where Q is the matrix + and t is the translation vector acting on facet-local coordinates. + """ + # First, identify facet type (simplex or tensor-product) + import FIAT + ref_cell = FIAT.ufc_cell(self.ufl_cell()) + facet_dim = ref_cell.get_spatial_dimension() - 1 + + if facet_dim == 0: + return { + 0: ( + np.empty((0, 0), dtype=RealType), + np.empty((0,), dtype=RealType), + ) + } + + facet_ref = ref_cell.construct_subelement(facet_dim) + facet_verts = np.asarray(facet_ref.get_vertices()) # FIAT's canonical facet vertices + + # Second, build orientation -> vertex permutation map + if isinstance(facet_ref, FIAT.reference_element.SimplicialComplex): + # 2D: quads and triangles (facet is an interval) + # 3D: tetrahedra (facet is a triangle) + o_p_maps = FIAT.orientation_utils.make_entity_permutations_simplex(dim=facet_dim, npoints=2) + + elif isinstance(facet_ref, FIAT.reference_element.Hypercube): + # 3D: hexahedra (facet is a quad) + factors = facet_ref.product.cells # interval in a quad + factor_dims = [f.get_spatial_dimension() for f in factors] + + # NOTE: We assume below that all factors are identical (and have the same dim), + # otherwise, we need to build factor permutations for each possible dim. + same_factor_dims = all(d == factor_dims[0] for d in factor_dims) + if same_factor_dims: + factor_o_p_maps = FIAT.orientation_utils.make_entity_permutations_simplex(dim=factor_dims[0], npoints=2) + o_p_maps = FIAT.orientation_utils.make_entity_permutations_tensorproduct(factors, factor_dims, list([factor_o_p_maps]*len(factors))) + else: + raise NotImplementedError("Facet permutation maps not yet implemented for facet type %s" % type(facet_ref)) + + # Finally, convert each vertex permutation into an affine map + o_coord_maps = {} + for o, p in o_p_maps.items(): + permuted_facet_verts = facet_verts[p] + if facet_verts.shape[0] == facet_dim + 1: + # Simplex facet: use all vertices + xs = facet_verts + ys = permuted_facet_verts + o_int = o + else: + # Tensor-product facet: use a fixed simplex subset + subset = [0, 1, 2] # any non-collinear triple works on a non-degenerate quad + xs = facet_verts[subset] + ys = permuted_facet_verts[subset] + + # For a tensor-product facet FIAT returns orientation as a tuple (eo, io1, io2..., iok) + # so we need to convert it to an integer (facet orientations stored on the mesh are expressed as integers). + # The integer orientation is obtained as o_int = (2**dim) * eo + io (as written in cython.dmcommon.pyx _compute_orientation_interval_tensor_product) + # intrinsic orientation is a binary encoding (each factor contributes one bit) + # io is obtained by converting this encoding into an integer (with factor 0 being the most significant bit) + + eo = o[0] + bits = o[1:] + io = 0 + for bit in bits: + io = 2*io + bit + o_int = (2**len(bits))*eo + io + + Q, t = FIAT.reference_element.make_affine_mapping(xs, ys) + o_coord_maps[o_int] = (Q, t) + + return o_coord_maps + + @cached_property + def cell_facet_coord_transforms(self): + """Returns affine reference-coordinate transforms between neighbouring cells across a given facet. + + For each cell `c` and each local facet `i` of this cell, this method constructs the affine + map + + X' = A[c,i] @ X + b[c,i], + + where `X` are reference coordinates of a point lying on facet `i` in cell `c` + and `X'` are the reference coordinates of the same point expressed in the neighbouring cell + across facet `i`. + + Returns + ------- + tuple + (A_dat, b_dat) indexed by global cell ID and local facet ID in each cell. + The two `pyop2.Dat` objects store the matrix and translation vector of + each cell-to-cell reference coordinate transforms. + """ + num_facets = self.ufl_cell().num_facets + ref_dim = self.ufl_cell().topological_dimension + A_dset = op2.DataSet(self.cell_set, dim=(num_facets, ref_dim, ref_dim)) + b_dset = op2.DataSet(self.cell_set, dim=(num_facets, ref_dim)) + + # Create local numpy buffers + A_transform = np.full((self.num_cells(), num_facets, ref_dim, ref_dim), np.nan) + b_transform = np.full((self.num_cells(), num_facets, ref_dim), np.nan) + + # Populate the local buffers by iterating over interior facets + # and extracting cell adjacency information from the mesh topology directly + facet_cells = self.interior_facets.facet_cell # all local facets (owned + halos) + local_facets = self.interior_facets.local_facet_dat.data_ro_with_halos + + local_facet_orientations = self.interior_facets.local_facet_orientation_dat.data_ro_with_halos + A, b, A_inv = self._get_facet_embedding_maps() + o_coord_maps = self._get_facet_orientation_coord_maps() + + NODATAVAL = np.iinfo(local_facet_orientations.dtype).max + + for f_idx in range(facet_cells.shape[0]): + c0, c1 = facet_cells[f_idx] # adjacent cells + lf0, lf1 = local_facets[f_idx] # local facet ID in each adjacent cell + o0, o1 = local_facet_orientations[f_idx] # local orientation in each adjacent cell + if o0 == NODATAVAL or o1 == NODATAVAL: + # skip outer halo facets where adjacency data isn't available + continue + + # Compute coords. map from o0 -> o1 + Q0, t0 = o_coord_maps[o0] # map canonical orientation -> o0 + Q1, t1 = o_coord_maps[o1] # map canonical orientation -> o1 + + # inverse coord. map + # o0 -> canonical orientation + Q0_inv = np.linalg.inv(Q0) + t0_inv = -Q0_inv @ t0 + + # o1 -> canonical orientation + Q1_inv = np.linalg.inv(Q1) + t1_inv = -Q1_inv @ t1 + + # compose maps both ways + # o0 -> canonical orientation with canonical orientation -> o1 + Q01 = Q1 @ Q0_inv + t01 = Q1 @ t0_inv + t1 + + # o1 -> canonical orientation with canonical orientation -> o0 + Q10 = Q0 @ Q1_inv + t10 = Q0 @ t1_inv + t0 + + # Compute forward transform c0 -> c1 and store on c0 + A_transform[c0, lf0] = A[lf1] @ Q01 @ A_inv[lf0] + b_transform[c0, lf0] = b[lf1] + A[lf1] @ t01 - A_transform[c0, lf0] @ b[lf0] + + # Compute backward transform c1 -> c0 and store on c1 + A_transform[c1, lf1] = A[lf0] @ Q10 @ A_inv[lf1] + b_transform[c1, lf1] = b[lf0] + A[lf0] @ t10 - A_transform[c1, lf1] @ b[lf1] + + cell_facet_A = op2.Dat( + A_dset, + A_transform, + dtype=RealType, + name="cell-facet-coord-transforms-A-dat" + ) + + cell_facet_A.halo_valid = True + + cell_facet_b = op2.Dat( + b_dset, + b_transform, + dtype=RealType, + name="cell-facet-coord-transforms-b-dat" + ) + cell_facet_b.halo_valid = True + + return cell_facet_A, cell_facet_b + + @cached_property + def cell_facet_exterior_mask(self): + """Identify exterior facets of each local cell, including halo cells. + + Returns a boolean array of shape ``(num_cells_with_halos, num_facets)``. + The entry ``cell_facet_exterior_mask[c, i]`` is ``True`` if local facet + ``i`` of local cell ``c`` lies on the domain boundary, and ``False`` + otherwise. + """ + num_facets = self.ufl_cell().num_facets + mask = np.zeros((self.cell_set.total_size, num_facets), dtype=bool) + + plex = self.topology_dm + if plex.getStratumSize("exterior_facets", 1) > 0: + ext_plex_points = frozenset( + plex.getStratumIS("exterior_facets", 1).getIndices().tolist() + ) + else: + ext_plex_points = frozenset() + + cstart, cend = plex.getHeightStratum(0) + for c_plex in range(cstart, cend): + cid = self._cell_numbering.getOffset(c_plex) + + if cid < 0 or cid >= self.cell_set.total_size: + continue + + for lf in range(num_facets): + f_point = self._cell_facet_point(cid, lf) + mask[cid, lf] = f_point in ext_plex_points + + return mask + @PETSc.Log.EventDecorator() def _set_partitioner(self, plex, distribute, partitioner_type=None): """Set partitioner for (re)distributing underlying plex over comm. @@ -2414,6 +2749,8 @@ def __init__(self, coordinates): V = functionspaceimpl.WithGeometry(coordinates.function_space(), self) self._coordinates_function = function.Function(V, val=coordinates) + self._topology_version = 0 + def _ufl_signature_data_(self, *args, **kwargs): return (type(self), self.extruded, self.variable_layers, super()._ufl_signature_data_(*args, **kwargs)) diff --git a/tests/firedrake/regression/test_cell_facet_topology.py b/tests/firedrake/regression/test_cell_facet_topology.py new file mode 100644 index 0000000000..0d1fe8c492 --- /dev/null +++ b/tests/firedrake/regression/test_cell_facet_topology.py @@ -0,0 +1,211 @@ +import numpy as np +import pytest + +from firedrake import * + + +@pytest.fixture(params=[ + # No interior facets + pytest.param(lambda: UnitIntervalMesh(1), id="interval-1"), + pytest.param(lambda: UnitSquareMesh(1, 1), id="tri-square-1x1"), + pytest.param(lambda: UnitSquareMesh(1, 1, quadrilateral=True), id="quad-square-1x1"), + pytest.param(lambda: UnitCubeMesh(1, 1, 1), id="tet-cube-1x1x1"), + pytest.param(lambda: UnitCubeMesh(1, 1, 1, hexahedral=True), id="hex-cube-1x1x1"), + + # With interior facets + pytest.param(lambda: UnitIntervalMesh(2), id="interval-2"), + pytest.param(lambda: UnitSquareMesh(2, 1), id="tri-square-2x1"), + pytest.param(lambda: UnitSquareMesh(2, 1, quadrilateral=True), id="quad-square-2x1"), + pytest.param(lambda: UnitCubeMesh(2, 1, 1), id="tet-cube-2x1x1"), + pytest.param(lambda: UnitCubeMesh(2, 1, 1, hexahedral=True), id="hex-cube-2x1x1"), +]) +def mesh(request): + return request.param() + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_neighbours_are_valid(mesh): + topology = mesh.topology + + # Check that every cell has exactly one neighbour entry per local facet + owned_neighbours = topology.cell_facet_neighbours.data_ro + neighbours_with_halos = topology.cell_facet_neighbours.data_ro_with_halos + + assert owned_neighbours.shape == ( + topology.cell_set.size, + mesh.ufl_cell().num_facets, + ) + + assert neighbours_with_halos.shape == ( + topology.cell_set.total_size, + mesh.ufl_cell().num_facets, + ) + + for c, row in enumerate(owned_neighbours): + for n in row: + if n == -1: + continue + + # Check that every non-boundary neighbour has a valid cell number + assert 0 <= n < topology.cell_set.total_size + + # Check reciprocity: If c is the neighbour of n across a given facet, + # then n also appears as the neighbour of c + assert c in neighbours_with_halos[n] + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_neighbours_match_interior_facets(mesh): + topology = mesh.topology + + neighbours = topology.cell_facet_neighbours.data_ro_with_halos + + # Test only with owned facets while allowing either adjacent cell to be halo + nowned_facets = topology.interior_facets.set.size + facet_cells = topology.interior_facets.facet_cell[:nowned_facets] + local_facets = topology.interior_facets.local_facet_dat.data_ro[:nowned_facets] + + # Check that the neighbouring cell is attached to the right facet + for (c0, c1), (lf0, lf1) in zip(facet_cells, local_facets): + if c0 == -1 or c1 == -1: + # Skip interior facets for which the rank doesn't know both adjacent cells + continue + + assert 0 <= c0 < topology.cell_set.total_size + assert 0 <= c1 < topology.cell_set.total_size + + assert neighbours[c0, lf0] == c1 + assert neighbours[c1, lf1] == c0 + + f0 = topology._cell_facet_point(c0, lf0) + f1 = topology._cell_facet_point(c1, lf1) + + assert f0 == f1 + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_exterior_mask(mesh): + topology = mesh.topology + + mask = topology.cell_facet_exterior_mask + + expected = np.zeros( + (topology.cell_set.total_size, mesh.ufl_cell().num_facets), + dtype=bool, + ) + + facet_cells = np.asarray(topology.exterior_facets.facet_cell).reshape((-1, 1)) + local_facets = np.asarray( + topology.exterior_facets.local_facet_dat.data_ro_with_halos + ).reshape((-1, 1)) + + for (c,), (lf,) in zip(facet_cells, local_facets): + if c == -1: + continue + assert 0 <= c < topology.cell_set.total_size + expected[c, lf] = True + + assert mask.shape == expected.shape + assert np.array_equal(mask, expected) + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_coord_transforms_are_inverse_on_interior_facets(mesh): + """Checks that the forward and backward transforms are consistent for two adjacent cells over a fixed shared facet.""" + topology = mesh.topology + + A_dat, b_dat = topology.cell_facet_coord_transforms + A = A_dat.data_ro_with_halos + b = b_dat.data_ro_with_halos + + embed_A, embed_b, _ = topology._get_facet_embedding_maps() + + # Restrict to owned interior facets + nowned_facets = topology.interior_facets.set.size + facet_cells = topology.interior_facets.facet_cell[:nowned_facets] + local_facets = topology.interior_facets.local_facet_dat.data_ro[:nowned_facets] + + for (c0, c1), (lf0, lf1) in zip(facet_cells, local_facets): + if c0 == -1 or c1 == -1: + continue + + assert 0 <= c0 < topology.cell_set.total_size + assert 0 <= c1 < topology.cell_set.total_size + + assert np.isfinite(A[c0, lf0]).all() + assert np.isfinite(b[c0, lf0]).all() + assert np.isfinite(A[c1, lf1]).all() + assert np.isfinite(b[c1, lf1]).all() + + facet_dim = embed_A[lf0].shape[1] + points = [np.zeros(facet_dim)] + points += [np.eye(facet_dim)[i] for i in range(facet_dim)] + + for Xf in points: + x = embed_A[lf0] @ Xf + embed_b[lf0] + y = A[c0, lf0] @ x + b[c0, lf0] + x_back = A[c1, lf1] @ y + b[c1, lf1] + + assert np.allclose(x_back, x) + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_coord_transforms_map_to_neighbour_facet(mesh): + """Checks geometric consistency, that is the linear transform for a cell c on local facet lf + sends points on that facet to points on the corresponding facet of the neighbouring cell.""" + topology = mesh.topology + + A_dat, b_dat = topology.cell_facet_coord_transforms + A = A_dat.data_ro_with_halos + b = b_dat.data_ro_with_halos + + embed_A, embed_b, _ = topology._get_facet_embedding_maps() + + # Restrict to owned interior facets + nowned_facets = topology.interior_facets.set.size + facet_cells = topology.interior_facets.facet_cell[:nowned_facets] + local_facets = topology.interior_facets.local_facet_dat.data_ro[:nowned_facets] + + for (c0, c1), (lf0, lf1) in zip(facet_cells, local_facets): + if c0 == -1 or c1 == -1: + continue + + assert 0 <= c0 < topology.cell_set.total_size + assert 0 <= c1 < topology.cell_set.total_size + + facet_dim = embed_A[lf0].shape[1] + points = [np.zeros(facet_dim)] + points += [np.eye(facet_dim)[i] for i in range(facet_dim)] + + for Xf in points: + X = embed_A[lf0] @ Xf + embed_b[lf0] + Y = A[c0, lf0] @ X + b[c0, lf0] + + # Y should lie on neighbour local facet lf1. + Yf = np.linalg.pinv(embed_A[lf1]) @ (Y - embed_b[lf1]) + assert np.allclose(embed_A[lf1] @ Yf + embed_b[lf1], Y) + + facet_dim = embed_A[lf1].shape[1] + points = [np.zeros(facet_dim)] + points += [np.eye(facet_dim)[i] for i in range(facet_dim)] + + for Yf in points: + Y = embed_A[lf1] @ Yf + embed_b[lf1] + X = A[c1, lf1] @ Y + b[c1, lf1] + + Xf = np.linalg.pinv(embed_A[lf0]) @ (X - embed_b[lf0]) + assert np.allclose(embed_A[lf0] @ Xf + embed_b[lf0], X) + + +@pytest.mark.parallel([1, 3]) +def test_cell_facet_topology_parallel_smoke(): + mesh = UnitSquareMesh(4, 4) + topology = mesh.topology + + topology.cell_facet_neighbours.data_ro_with_halos + + A, b = topology.cell_facet_coord_transforms + A.data_ro_with_halos + b.data_ro_with_halos + + topology.cell_facet_exterior_mask