Move Firedrake code from ngsPETSc into Firedrake

.github/workflows/core.yml

@@ -194,7 +194,8 @@ jobs:
             : # because they rely on non-PyPI versions of petsc4py.
             pip install --no-build-isolation --no-deps \
               "$PETSC_DIR"/"$PETSC_ARCH"/externalpackages/git.slepc/src/binding/slepc4py
-            pip install --no-deps git+https://github.com/NGSolve/ngsPETSc.git netgen-mesher netgen-occt
+            : # UNDO ME
+            pip install --no-deps git+https://github.com/NGSolve/ngsPETSc.git@connorjward/remove-firedrake netgen-mesher netgen-occt
 
             : # We have to pass '--no-build-isolation' to use a custom petsc4py
             EXTRA_BUILD_ARGS='--no-isolation'

firedrake/mesh.py

@@ -2970,6 +2970,179 @@ def __iter__(self):
     def unique(self):
         return self
 
+    def refine_marked_elements(self, mark, netgen_flags=None):
+        """Refine a mesh using a DG0 marking function.
+
+        This method requires that the mesh has been constructed from a
+        netgen mesh.
+
+        :arg mark: the marking function which is a Firedrake DG0 function.
+        :arg netgen_flags: the dictionary of flags to be passed to ngsPETSc.
+
+        It includes the option:
+            - refine_faces, which is a boolean specifying if you want to refine faces.
+
+        """
+        import firedrake as fd
+
+        utils.check_netgen_installed()
+
+        if netgen_flags is None:
+            netgen_flags = {}
+        DistParams = self._distribution_parameters
+        els = {2: self.netgen_mesh.Elements2D, 3: self.netgen_mesh.Elements3D}
+        dim = self.geometric_dimension
+        refine_faces = netgen_flags.get("refine_faces", False)
+        if dim in [2, 3]:
+            with mark.dat.vec as marked:
+                marked0 = marked
+                getIdx = self._cell_numbering.getOffset
+                if self.sfBC is not None:
+                    sfBCInv = self.sfBC.createInverse()
+                    getIdx = lambda x: x
+                    _, marked0 = self.topology_dm.distributeField(sfBCInv,
+                                                                  self._cell_numbering,
+                                                                  marked)
+                if self.comm.Get_rank() == 0:
+                    mark = marked0.getArray()
+                    max_refs = np.max(mark)
+                    for _ in range(int(max_refs)):
+                        for i, el in enumerate(els[dim]()):
+                            if mark[getIdx(i)] > 0:
+                                el.refine = True
+                            else:
+                                el.refine = False
+                        if not refine_faces and dim == 3:
+                            self.netgen_mesh.Elements2D().NumPy()["refine"] = 0
+                        self.netgen_mesh.Refine(adaptive=True)
+                        mark = mark-np.ones(mark.shape)
+                    return fd.Mesh(self.netgen_mesh, distribution_parameters=DistParams, comm=self.comm)
+                return fd.Mesh(netgen.libngpy._meshing.Mesh(dim),
+                               distribution_parameters=DistParams, comm=self.comm)
+        else:
+            raise NotImplementedError("No implementation for dimension other than 2 and 3.")
+
+    @PETSc.Log.EventDecorator()
+    def curve_field(self, order, permutation_tol=1e-8, location_tol=1e-1, cg_field=False):
+        '''Return a function containing the curved coordinates of the mesh.
+
+        This method requires that the mesh has been constructed from a
+        netgen mesh.
+
+        :arg order: the order of the curved mesh.
+        :arg permutation_tol: tolerance used to construct the permutation of the reference element.
+        :arg location_tol: tolerance used to locate the cell a point belongs to.
+        :arg cg_field: return a CG function field representing the mesh, rather than the
+                       default DG field.
+
+        '''
+        import firedrake as fd
+
+        utils.check_netgen_installed()
+
+        from firedrake.netgen import find_permutation
+
+        # Check if the mesh is a surface mesh or two dimensional mesh
+        if len(self.netgen_mesh.Elements3D()) == 0:
+            ng_element = self.netgen_mesh.Elements2D
+        else:
+            ng_element = self.netgen_mesh.Elements3D
+        ng_dimension = len(ng_element())
+        geom_dim = self.geometric_dimension
+
+        # Construct the mesh as a Firedrake function
+        if cg_field:
+            firedrake_space = fd.VectorFunctionSpace(self, "CG", order)
+        else:
+            low_order_element = self.coordinates.function_space().ufl_element().sub_elements[0]
+            ufl_element = low_order_element.reconstruct(degree=order)
+            firedrake_space = fd.VectorFunctionSpace(self, fd.BrokenElement(ufl_element))
+        new_coordinates = fd.assemble(fd.interpolate(self.coordinates, firedrake_space))
+
+        # Compute reference points using fiat
+        fiat_element = new_coordinates.function_space().finat_element.fiat_equivalent
+        entity_ids = fiat_element.entity_dofs()
+        nodes = fiat_element.dual_basis()
+        ref = []
+        for dim in entity_ids:
+            for entity in entity_ids[dim]:
+                for dof in entity_ids[dim][entity]:
+                    # Assert singleton point for each node.
+                    pt, = nodes[dof].get_point_dict().keys()
+                    ref.append(pt)
+        reference_space_points = np.array(ref)
+
+        # Curve the mesh on rank 0 only
+        if self.comm.rank == 0:
+            # Construct numpy arrays for physical domain data
+            physical_space_points = np.zeros(
+                (ng_dimension, reference_space_points.shape[0], geom_dim)
+            )
+            curved_space_points = np.zeros(
+                (ng_dimension, reference_space_points.shape[0], geom_dim)
+            )
+            self.netgen_mesh.CalcElementMapping(reference_space_points, physical_space_points)
+            # NOTE: This will segfault!
+            self.netgen_mesh.Curve(order)
+            self.netgen_mesh.CalcElementMapping(reference_space_points, curved_space_points)
+            curved = ng_element().NumPy()["curved"]
+            # Broadcast a boolean array identifying curved cells
+            curved = self.comm.bcast(curved, root=0)
+            physical_space_points = physical_space_points[curved]
+            curved_space_points = curved_space_points[curved]
+        else:
+            curved = self.comm.bcast(None, root=0)
+            # Construct numpy arrays as buffers to receive physical domain data
+            ncurved = np.sum(curved)
+            physical_space_points = np.zeros(
+                (ncurved, reference_space_points.shape[0], geom_dim)
+            )
+            curved_space_points = np.zeros(
+                (ncurved, reference_space_points.shape[0], geom_dim)
+            )
+
+        # Broadcast curved cell point data
+        self.comm.Bcast(physical_space_points, root=0)
+        self.comm.Bcast(curved_space_points, root=0)
+        cell_node_map = new_coordinates.cell_node_map()
+
+        # Select only the points in curved cells
+        barycentres = np.average(physical_space_points, axis=1)
+        ng_index = [*map(lambda x: self.locate_cell(x, tolerance=location_tol), barycentres)]
+
+        # Select only the indices of points owned by this rank
+        owned = [(0 <= ii < len(cell_node_map.values)) if ii is not None else False for ii in ng_index]
+
+        # Select only the points owned by this rank
+        physical_space_points = physical_space_points[owned]
+        curved_space_points = curved_space_points[owned]
+        barycentres = barycentres[owned]
+        ng_index = [idx for idx, o in zip(ng_index, owned) if o]
+
+        # Get the PyOP2 indices corresponding to the netgen indices
+        pyop2_index = []
+        for ngidx in ng_index:
+            pyop2_index.extend(cell_node_map.values[ngidx])
+
+        # Find the correct coordinate permutation for each cell
+        # NB: Coordinates must be cast to real when running Firedrake in complex mode
+        permutation = find_permutation(
+            physical_space_points,
+            new_coordinates.dat.data[pyop2_index].reshape(
+                physical_space_points.shape
+            ).astype(np.float64, copy=False),
+            tol=permutation_tol
+        )
+
+        # Apply the permutation to each cell in turn
+        for ii, p in enumerate(curved_space_points):
+            curved_space_points[ii] = p[permutation[ii]]
+
+        # Assign the curved coordinates to the dat
+        new_coordinates.dat.data[pyop2_index] = curved_space_points.reshape(-1, geom_dim)
+
+        return new_coordinates
+
 
 @PETSc.Log.EventDecorator()
 def make_mesh_from_coordinates(coordinates, name, tolerance=0.5):
@@ -3199,6 +3372,8 @@ def Mesh(meshfile, **kwargs):
 
     utils._init()
 
+    from_netgen = netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh)
+
     # We don't need to worry about using a user comm in these cases as
     # they all immediately call a petsc4py which in turn uses a PETSc
     # internal comm
@@ -3207,16 +3382,15 @@ def Mesh(meshfile, **kwargs):
         plex = meshfile
         if MPI.Comm.Compare(user_comm, plex.comm.tompi4py()) not in {MPI.CONGRUENT, MPI.IDENT}:
             raise ValueError("Communicator used to create `plex` must be at least congruent to the communicator used to create the mesh")
-    elif netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh):
-        try:
-            from ngsPETSc import FiredrakeMesh
-        except ImportError:
-            raise ImportError("Unable to import ngsPETSc. Please ensure that ngsolve is installed and available to Firedrake.")
+    elif from_netgen:
+        from firedrake.netgen import FiredrakeMesh
+
         petsctools.cite("Betteridge2024")
         netgen_flags = kwargs.get("netgen_flags", {"quad": False, "transform": None, "purify_to_tets": False})
         netgen_firedrake_mesh = FiredrakeMesh(meshfile, netgen_flags, user_comm)
         plex = netgen_firedrake_mesh.meshMap.petscPlex
         plex.setName(_generate_default_mesh_topology_name(name))
+
     else:
         basename, ext = os.path.splitext(meshfile)
         if ext.lower() in ['.e', '.exo']:
@@ -3245,11 +3419,11 @@ def Mesh(meshfile, **kwargs):
                             permutation_name=kwargs.get("permutation_name"),
                             submesh_parent=submesh_parent.topology if submesh_parent else None,
                             comm=user_comm)
-    if netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh):
-        netgen_firedrake_mesh.createFromTopology(topology, name=name, comm=user_comm)
-        mesh = netgen_firedrake_mesh.firedrakeMesh
-    else:
-        mesh = make_mesh_from_mesh_topology(topology, name)
+    mesh = make_mesh_from_mesh_topology(topology, name)
+
+    if from_netgen:
+        mesh.netgen_mesh = netgen_firedrake_mesh.meshMap.ngMesh
+
     mesh.submesh_parent = submesh_parent
     mesh._tolerance = tolerance
     return mesh

firedrake/mg/mesh.py

@@ -7,6 +7,7 @@
 import petsctools
 import firedrake
 import firedrake.cython.dmcommon as dmcommon
+from firedrake import utils
 from firedrake.utils import cached_property
 from firedrake.cython import mgimpl as impl
 from .utils import set_level
@@ -113,12 +114,9 @@ def MeshHierarchy(mesh, refinement_levels,
     """
 
     if (isinstance(netgen_flags, bool) and netgen_flags) or isinstance(netgen_flags, dict):
-        try:
-            from ngsPETSc import NetgenHierarchy
-        except ImportError:
-            raise ImportError("Unable to import netgen and ngsPETSc. Please ensure that netgen and ngsPETSc "
-                              "are installed and available to Firedrake (see "
-                              "https://www.firedrakeproject.org/install.html#netgen).")
+        utils.check_netgen_installed()
+        from firedrake.mg.netgen import NetgenHierarchy
+
         if hasattr(mesh, "netgen_mesh"):
             return NetgenHierarchy(mesh, refinement_levels, flags=netgen_flags)
         else: