Additional tidying, difficult numpy indexing rules

ngsPETSc/utils/firedrake/meshes.py

@@ -12,6 +12,12 @@
 import numpy as np
 from petsc4py import PETSc
 
+try:
+    from scipy.spatial.distance import cdist
+    HAVE_SCIPY = True
+except ImportError:
+    HAVE_SCIPY = False
+
 import netgen
 import netgen.meshing as ngm
 from netgen.meshing import MeshingParameters
@@ -70,6 +76,45 @@ def refineMarkedElements(self, mark):
     else:
         raise NotImplementedError("No implementation for dimension other than 2 and 3.")
 
+
+def _slow_cdist(XA, XB):
+        dist = np.zeros([len(XA), len(XB)])
+        for ii, a in enumerate(XA):
+            for jj, b in enumerate(XB):
+                dist[ii, jj] = np.linalg.norm(b - a)
+        return dist
+
+
+if not HAVE_SCIPY:
+    cdist = _slow_cdist
+
+
+def find_permutation(points_a, points_b, tol=1e-5):
+    """ Find all permutations between a list of two sets of points.
+
+    Given two numpy arrays of shape (ncells, npoints, dim) containing
+    floating point coordinates for each cell, determine each index
+    permutation that takes `points_a` to `points_b`. Ie:
+    ```
+    permutation = find_permutation(points_a, points_b)
+    assert np.allclose(points_a[permutation], points_b, rtol=0, atol=tol)
+    ```
+    """
+    if points_a.shape != points_b.shape:
+        raise ValueError("`points_a` and `points_b` must have the same shape.")
+
+    p = [np.where(cdist(a, b).T < tol)[1] for a, b in zip(points_a, points_b)]
+    try:
+        permutation = np.array(p, ndmin=2)
+    except ValueError:
+        raise ValueError("It was not possible to find a permutation for every cell within the provided tolerance")
+
+    if permutation.shape != points_a.shape[0:2]:
+        raise ValueError("It was not possible to find a permutation for every cell within the provided tolerance")
+
+    return permutation
+
+
 def curveField(self, order, tol=1e-8):
     '''
     This method returns a curved mesh as a Firedrake function.
@@ -78,21 +123,21 @@ def curveField(self, order, tol=1e-8):
 
     '''
     #Checking if the mesh is a surface mesh or two dimensional mesh
-    surf = len(self.netgen_mesh.Elements3D()) == 0 # REMOVE
     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()
 
     #Constructing mesh as a function
     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))
-    newFunctionCoordinates = fd.assemble(interpolate(self.coordinates, firedrake_space))
+    new_coordinates = fd.assemble(interpolate(self.coordinates, firedrake_space))
 
     #Computing reference points using fiat
-    fiat_element = newFunctionCoordinates.function_space().finat_element.fiat_equivalent
+    fiat_element = new_coordinates.function_space().finat_element.fiat_equivalent
     entity_ids = fiat_element.entity_dofs()
     nodes = fiat_element.dual_basis()
     ref = []
@@ -105,9 +150,8 @@ def curveField(self, order, tol=1e-8):
     reference_space_points = np.array(ref)
 
     #Mapping to the physical domain
-    els = {True: self.netgen_mesh.Elements2D, False: self.netgen_mesh.Elements3D} # REMOVE
-    physical_space_points = np.ndarray((ng_dimension, reference_space_points.shape[0], self.geometric_dimension()))
-    curved_space_points = np.ndarray((ng_dimension, reference_space_points.shape[0], self.geometric_dimension()))
+    physical_space_points = np.ndarray((ng_dimension, reference_space_points.shape[0], geom_dim))
+    curved_space_points = np.ndarray((ng_dimension, reference_space_points.shape[0], geom_dim))
 
     if self.comm.rank == 0:
         #Curving the mesh on rank 0
@@ -122,7 +166,7 @@ def curveField(self, order, tol=1e-8):
     physical_space_points = self.comm.bcast(physical_space_points, root=0)
     curved_space_points = self.comm.bcast(curved_space_points, root=0)
     curved = self.comm.bcast(curved, root=0)
-    cell_node_map = newFunctionCoordinates.cell_node_map()
+    cell_node_map = new_coordinates.cell_node_map()
 
     # Select only the points in curved cells
     physical_space_points = physical_space_points[curved]
@@ -137,68 +181,28 @@ def curveField(self, order, tol=1e-8):
     barycentres = barycentres[owned]
     ng_index = [idx for idx, o in zip(ng_index, owned) if o]
 
-    # Do we want the min or max of these???
-    norms = np.linalg.norm(physical_space_points - newFunctionCoordinates.dat.data[cell_node_map.values[ng_index]], axis=2)
+    breakpoint()
 
     # PyOP2 index
     pyop2_index = []
     for ngidx in ng_index:
         pyop2_index.extend(cell_node_map.values[ngidx])
     np.array(pyop2_index)
 
-    for dim in range(self.geometric_dimension()):
-        newFunctionCoordinates.sub(dim).dat.data[pyop2_index] = curved_space_points[:,:,dim].flatten()
-
-    # ~ breakpoint()
-
-    # ~ for i in range(physical_space_points.shape[0]):
-        # ~ #Inefficent code but runs only on curved elements
-        # ~ if curved[i]:
-            # ~ pts = physical_space_points[i][0:reference_space_points.shape[0]]
-            # ~ bary = sum([np.array(pts[i]) for i in range(len(pts))])/len(pts)
-            # ~ Idx = self.locate_cell(bary)
-            # ~ isInMesh = (0<=Idx<len(cell_node_map.values)) if Idx is not None else False
-            # ~ #Check if element is shared across processes
-            # ~ shared = self.comm.gather(isInMesh, root=0)
-            # ~ shared = self.comm.bcast(shared, root=0) # `shared` isn't used in anything other than warning messages!?
-            # ~ # self.comm.Allgather(isInMesh, shared)
-            # ~ #Bend if not shared
-            # ~ if np.sum(shared) == 1:
-                # ~ if isInMesh:
-                    # ~ # This seems dodgy, I think we're just trying to ensure the points in the netgen cell are close to the coordinates in the dat
-                    # ~ p = [np.argmin(np.sum((pts - pt)**2, axis=1))
-                            # ~ for pt in newFunctionCoordinates.dat.data[cell_node_map.values[Idx]][0:reference_space_points.shape[0]]]
-                    # ~ curved_space_points[i] = curved_space_points[i][p]
-                    # ~ res = np.linalg.norm(pts[p] - newFunctionCoordinates.dat.data[cell_node_map.values[Idx]][0:reference_space_points.shape[0]])
-                    # ~ if res > tol:
-                        # ~ fd.logging.warning(f"[{self.comm.rank}] Not able to curve Firedrake element {Idx} ({i}) -- residual: {res}")
-                    # ~ else:
-                        # ~ for j, datIdx in enumerate(cell_node_map.values[Idx][0:reference_space_points.shape[0]]):
-                            # ~ for dim in range(self.geometric_dimension()):
-                                # ~ coo = curved_space_points[i][j][dim]
-                                # ~ newFunctionCoordinates.sub(dim).dat.data[datIdx] = coo
-            # ~ else:
-                # ~ if isInMesh:
-                    # ~ p = [np.argmin(np.sum((pts - pt)**2, axis=1))
-                            # ~ for pt in newFunctionCoordinates.dat.data[cell_node_map.values[Idx]][0:reference_space_points.shape[0]]]
-                    # ~ curved_space_points[i] = curved_space_points[i][p]
-                    # ~ res = np.linalg.norm(pts[p]-newFunctionCoordinates.dat.data[cell_node_map.values[Idx]][0:reference_space_points.shape[0]])
-                # ~ else:
-                    # ~ res = np.inf
-                # ~ res = self.comm.gather(res, root=0)
-                # ~ res = self.comm.bcast(res, root=0)
-                # ~ rank = np.argmin(res)
-                # ~ if self.comm.rank == rank:
-                    # ~ if res[rank] > tol:
-                        # ~ fd.logging.warning("[{}, {}] Not able to curve Firedrake element {} \
-                            # ~ ({}) -- residual: {}".format(self.comm.rank, shared, Idx,i, res))
-                    # ~ else:
-                        # ~ for j, datIdx in enumerate(cell_node_map.values[Idx][0:reference_space_points.shape[0]]):
-                            # ~ for dim in range(self.geometric_dimension()):
-                                # ~ coo = curved_space_points[i][j][dim]
-                                # ~ newFunctionCoordinates.sub(dim).dat.data[datIdx] = coo
-    # ~ breakpoint()
-    return newFunctionCoordinates
+    # Find the correct coordinate permutation for each cell
+    permutation = find_permutation(
+        physical_space_points,
+        new_coordinates.dat.data[pyop2_index].reshape(physical_space_points.shape)
+    )
+
+    # 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
 
 def splitToQuads(plex, dim, comm):
     '''