Periodic vertexonlymesh

firedrake/mesh.py

@@ -2173,13 +2173,10 @@ def VertexOnlyMesh(mesh, vertexcoords, missing_points_behaviour=None):
     if mesh.coordinates.function_space().ufl_element().degree() > 1:
         raise NotImplementedError("Only straight edged meshes are supported")
 
-    if hasattr(mesh, "_made_from_coordinates") and mesh._made_from_coordinates:
-        raise NotImplementedError("Meshes made from coordinate fields are not yet supported")
-
     if pdim != gdim:
         raise ValueError(f"Mesh geometric dimension {gdim} must match point list dimension {pdim}")
 
-    swarm = _pic_swarm_in_plex(mesh.topology.topology_dm, vertexcoords, fields=[("parentcellnum", 1, IntType), ("refcoord", tdim, RealType)])
+    swarm = _pic_swarm_in_plex(mesh, vertexcoords)
 
     if missing_points_behaviour:
 
@@ -2217,8 +2214,6 @@ def compare_arrays(x, y, datatype):
             else:
                 raise ValueError("missing_points_behaviour must be None, 'error' or 'warn'")
 
-    dmcommon.label_pic_parent_cell_info(swarm, mesh)
-
     # Topology
     topology = VertexOnlyMeshTopology(swarm, mesh.topology, name="swarmmesh", reorder=False)
 
@@ -2254,7 +2249,7 @@ def compare_arrays(x, y, datatype):
     return vmesh
 
 
-def _pic_swarm_in_plex(plex, coords, fields=[]):
+def _pic_swarm_in_plex(parent_mesh, coords, fields=None):
     """
     Create a Particle In Cell (PIC) DMSwarm, immersed in a DMPlex
     at given point coordinates.
@@ -2313,6 +2308,18 @@ def _pic_swarm_in_plex(plex, coords, fields=[]):
 
     # Check coords
     coords = np.asarray(coords, dtype=RealType)
+
+    plex = parent_mesh.topology.topology_dm
+    tdim = parent_mesh.topological_dimension()
+    gdim = parent_mesh.geometric_dimension()
+
+    if fields is None:
+        fields = []
+    fields += [("parentcellnum", 1, IntType), ("refcoord", tdim, RealType)]
+
+    coords, reference_coords, parent_cell_nums = \
+        _parent_mesh_embedding(coords, parent_mesh)
+
     _, coordsdim = coords.shape
 
     # Create a DMSWARM
@@ -2339,6 +2346,9 @@ def _pic_swarm_in_plex(plex, coords, fields=[]):
         swarm.registerField(name, size, dtype=dtype)
     swarm.finalizeFieldRegister()
 
+    num_vertices = len(coords)
+    swarm.setLocalSizes(num_vertices, -1)
+
     # Note that no new fields can now be associated with the DMSWARM.
 
     # Add point coordinates - note we set redundant mode to False
@@ -2351,7 +2361,28 @@ def _pic_swarm_in_plex(plex, coords, fields=[]):
     # coordinates associated with them. The DMPlex cell id associated
     # with each PIC in the DMSwarm is accessed with the `DMSwarm_cellid`
     # field.
-    swarm.setPointCoordinates(coords, redundant=False, mode=PETSc.InsertMode.INSERT_VALUES)
+    # swarm.setPointCoordinates(coords, redundant=False, mode=PETSc.InsertMode.INSERT_VALUES)
+
+    # do this ourselves by just setting DMSwarmPICField_coor and DMSwarmPICField_cellid
+
+    # get fields - NOTE this isn't copied so make sure
+    # swarm.restoreField is called for each field too!
+    swarm_coords = swarm.getField("DMSwarmPIC_coor").reshape((num_vertices, gdim))
+    # Why are the next two different?
+    swarm_parent_cell_nums = swarm.getField("DMSwarm_cellid")
+    field_parent_cell_nums = swarm.getField("parentcellnum")
+    field_reference_coords = swarm.getField("refcoord").reshape((num_vertices, tdim))
+
+    swarm_coords[...] = coords
+    swarm_parent_cell_nums[...] = parent_cell_nums
+    field_parent_cell_nums[...] = parent_cell_nums
+    field_reference_coords[...] = reference_coords
+
+    # have to restore fields once accessed to allow access again
+    swarm.restoreField("refcoord")
+    swarm.restoreField("parentcellnum")
+    swarm.restoreField("DMSwarmPIC_coor")
+    swarm.restoreField("DMSwarm_cellid")
 
     # Remove PICs which have been placed into ghost cells of a distributed DMPlex
     dmcommon.remove_ghosts_pic(swarm, plex)
@@ -2366,6 +2397,44 @@ def _pic_swarm_in_plex(plex, coords, fields=[]):
     return swarm
 
 
+def _parent_mesh_embedding(vertex_coords, parent_mesh):
+
+    # Sort parallel later.
+
+    max_num_vertices = len(vertex_coords)
+    num_vertices = max_num_vertices
+
+    gdim = parent_mesh.geometric_dimension()
+    tdim = parent_mesh.topological_dimension()
+
+    # Create an out of mesh point to use in locate_cell when needed
+
+    out_of_mesh_point = np.full((1, gdim), np.inf)
+
+    # TODO: Eliminate points outside rank bounding box before calculating.
+
+    parent_cell_nums = np.empty(num_vertices, dtype=IntType)
+    reference_coords = np.empty((num_vertices, tdim), dtype=RealType)
+    valid = np.full(num_vertices, False)
+
+    for i in range(max_num_vertices):
+        if i < num_vertices:
+            parent_cell_num, reference_coord = \
+                parent_mesh.locate_cell_and_reference_coordinate(vertex_coords[i])
+            if parent_cell_num is not None:
+                valid[i] = True
+                parent_cell_nums[i] = parent_cell_num
+                reference_coords[i] = reference_coord
+        else:
+            parent_mesh.locate_cell(out_of_mesh_point)  # should return None
+
+    vertex_coords = np.compress(valid, vertex_coords, axis=0)
+    reference_coords = np.compress(valid, reference_coords, axis=0)
+    parent_cell_nums = np.compress(valid, parent_cell_nums, axis=0)
+
+    return vertex_coords, reference_coords, parent_cell_nums
+
+
 @PETSc.Log.EventDecorator()
 def SubDomainData(geometric_expr):
     """Creates a subdomain data object from a boolean-valued UFL expression.

tests/vertexonly/test_interpolation_from_parent.py

@@ -18,9 +18,7 @@
                                      # CalledProcessError is so the parallel tests correctly xfail
                                      marks=pytest.mark.xfail(raises=(subprocess.CalledProcessError, NotImplementedError),
                                                              reason="immersed parent meshes not supported")),
-                        pytest.param("periodicrectangle",
-                                     marks=pytest.mark.xfail(raises=(subprocess.CalledProcessError, NotImplementedError),
-                                                             reason="meshes made from coordinate fields are not supported")),
+                        pytest.param("periodicrectangle"),
                         pytest.param("shiftedmesh",
                                      marks=pytest.mark.skip(reason="meshes with modified coordinate fields are not supported"))],
                 ids=lambda x: f"{x}-mesh")
@@ -133,6 +131,13 @@ def test_scalar_function_interpolation(parentmesh, vertexcoords, fs):
     vm = VertexOnlyMesh(parentmesh, vertexcoords)
     vertexcoords = vm.coordinates.dat.data_ro.reshape(-1, parentmesh.geometric_dimension())
     fs_fam, fs_deg, fs_typ = fs
+    if (
+        parentmesh.coordinates.function_space().ufl_element().family()
+        == "Discontinuous Lagrange"
+        and fs_fam == "CG"
+    ):
+        pytest.skip(f"Interpolating into f{fs_fam} on a periodic mesh is not well-defined.")
+
     V = fs_typ(parentmesh, fs_fam, fs_deg)
     W = FunctionSpace(vm, "DG", 0)
     expr = reduce(add, SpatialCoordinate(parentmesh))
@@ -163,6 +168,11 @@ def test_vector_function_interpolation(parentmesh, vertexcoords, vfs):
     vfs_fam, vfs_deg, vfs_typ = vfs
     vm = VertexOnlyMesh(parentmesh, vertexcoords)
     vertexcoords = vm.coordinates.dat.data_ro
+    if (
+        parentmesh.coordinates.function_space().ufl_element().family()
+        == "Discontinuous Lagrange"
+    ):
+        pytest.skip(f"Interpolating into f{vfs_fam} on a periodic mesh is not well-defined.")
     V = vfs_typ(parentmesh, vfs_fam, vfs_deg)
     W = VectorFunctionSpace(vm, "DG", 0)
     expr = 2 * SpatialCoordinate(parentmesh)
@@ -200,6 +210,11 @@ def test_tensor_function_interpolation(parentmesh, vertexcoords, tfs):
     tfs_fam, tfs_deg, tfs_typ = tfs
     vm = VertexOnlyMesh(parentmesh, vertexcoords)
     vertexcoords = vm.coordinates.dat.data_ro
+    if (
+        parentmesh.coordinates.function_space().ufl_element().family()
+        == "Discontinuous Lagrange"
+    ):
+        pytest.skip(f"Interpolating into f{tfs_fam} on a periodic mesh is not well-defined.")
     V = tfs_typ(parentmesh, tfs_fam, tfs_deg)
     W = TensorFunctionSpace(vm, "DG", 0)
     x = SpatialCoordinate(parentmesh)

tests/vertexonly/test_swarm.py

@@ -42,7 +42,7 @@ def cell_midpoints(m):
                         "cube",
                         "tetrahedron",
                         pytest.param("immersedsphere", marks=pytest.mark.skip(reason="immersed parent meshes not supported and will segfault PETSc when creating the DMSwarm")),
-                        pytest.param("periodicrectangle", marks=pytest.mark.skip(reason="meshes made from coordinate fields are not supported")),
+                        pytest.param("periodicrectangle"),
                         pytest.param("shiftedmesh", marks=pytest.mark.skip(reason="meshes with modified coordinate fields are not supported"))])
 def parentmesh(request):
     if request.param == "interval":
@@ -78,7 +78,7 @@ def test_pic_swarm_in_plex(parentmesh):
     plex = parentmesh.topology.topology_dm
     from firedrake.petsc import PETSc
     fields = [("fieldA", 1, PETSc.IntType), ("fieldB", 2, PETSc.ScalarType)]
-    swarm = mesh._pic_swarm_in_plex(plex, inputpointcoords, fields=fields)
+    swarm = mesh._pic_swarm_in_plex(parentmesh, inputpointcoords, fields=fields)
     # Get point coords on current MPI rank
     localpointcoords = np.copy(swarm.getField("DMSwarmPIC_coor"))
     swarm.restoreField("DMSwarmPIC_coor")
@@ -103,14 +103,16 @@ def test_pic_swarm_in_plex(parentmesh):
         swarm.restoreField(name)
     # Check comm sizes match
     assert plex.comm.size == swarm.comm.size
+
     # Check coordinate list and parent cell indices match
     assert len(localpointcoords) == len(localparentcellindices)
     # check local points are found in list of input points
     for p in localpointcoords:
         assert np.any(np.isclose(p, inputpointcoords))
     # check local points are correct local points given mesh
     # partitioning (but don't require ordering to be maintained)
-    assert np.allclose(np.sort(inputlocalpointcoords), np.sort(localpointcoords))
+    assert np.allclose(np.sort(inputlocalpointcoords, axis=0),
+                       np.sort(localpointcoords, axis=0))
     # Check methods for checking number of points on current MPI rank
     assert len(localpointcoords) == swarm.getLocalSize()
     # Check there are as many local points as there are local cells

tests/vertexonly/test_vertex_only_fs.py

@@ -12,7 +12,7 @@
                         "cube",
                         "tetrahedron",
                         pytest.param("immersedsphere", marks=pytest.mark.xfail(reason="immersed parent meshes not supported")),
-                        pytest.param("periodicrectangle", marks=pytest.mark.xfail(reason="meshes made from coordinate fields are not supported")),
+                        pytest.param("periodicrectangle"),
                         pytest.param("shiftedmesh", marks=pytest.mark.skip(reason="meshes with modified coordinate fields are not supported"))])
 def parentmesh(request):
     if request.param == "interval":

tests/vertexonly/test_vertex_only_mesh_generation.py

@@ -42,7 +42,7 @@ def cell_midpoints(m):
                         "cube",
                         "tetrahedron",
                         pytest.param("immersedsphere", marks=pytest.mark.xfail(reason="immersed parent meshes not supported")),
-                        pytest.param("periodicrectangle", marks=pytest.mark.xfail(reason="meshes made from coordinate fields are not supported")),
+                        pytest.param("periodicrectangle"),
                         pytest.param("shiftedmesh", marks=pytest.mark.skip(reason="meshes with modified coordinate fields are not supported"))])
 def parentmesh(request):
     if request.param == "interval":