fix basis multipatch topology (#765)

This PR fixes creating a basis on a `MultipatchTopology` with C^0
continuity at the patch interfaces. In addition, this PR deduplicates
code for merging dofs.

nutils/_util.py

@@ -9,6 +9,7 @@
 import numpy
 import collections.abc
 import inspect
+import itertools
 import functools
 import operator
 import numbers
@@ -18,6 +19,7 @@
 import contextlib
 import treelog
 import datetime
+from typing import Iterable, Sequence, Tuple
 
 supports_outdirfd = os.open in os.supports_dir_fd and os.listdir in os.supports_fd
 
@@ -766,4 +768,52 @@ def cli(f, *, argv=None):
     return f(**kwargs)
 
 
+def merge_index_map(nin: int, merge_sets: Iterable[Sequence[int]]) -> Tuple[numpy.ndarray, int]:
+    '''Returns an index map relating ``nin`` unmerged elements to ``nout`` merged elements.
+
+    The index map, an array of length ``nin``, satisfies the following conditions:
+
+    *   For every merge set in ``merge_sets``: for every pair of indices ``i``
+        and ``j`` in a merge set, ``index_map[i] == index_map[j]`` is true.
+
+        In code, the following is true:
+
+            all(index_map[i] == index_map[j] for i, *js in merge_sets for j in js)
+
+    *   Selecting the first occurences of indices in ``index_map`` gives the
+        sequence ``range(nout)``.
+
+    Args
+    ----
+    nin : :class:`int`
+        The number of elements before merging.
+    merge_sets : iterable of sequences of at least one :class:`int`
+        An iterable of merge sets, where each merge set lists the indices of
+        input elements that should be merged. Every merge set should have at
+        least one index.
+
+    Returns
+    -------
+    index_map : :class:`numpy.ndarray`
+        Index map with satisfying the above conditions.
+    nout : :class:`int`
+        The number of output indices.
+    '''
+
+    index_map = numpy.arange(nin)
+    def resolve(index):
+        parent = index_map[index]
+        while index != parent:
+            index = parent
+            parent = index_map[index]
+        return index
+    for merge_set in merge_sets:
+        resolved = list(map(resolve, merge_set))
+        index_map[resolved] = min(resolved)
+    new_indices = itertools.count()
+    for iin, ptr in enumerate(index_map):
+        index_map[iin] = next(new_indices) if iin == ptr else index_map[ptr]
+    return index_map, next(new_indices)
+
+
 # vim:sw=4:sts=4:et

nutils/topology.py

@@ -1647,29 +1647,16 @@ def _basis_c0_structured(self, name, degree):
 
         coeffs = [ref.get_poly_coeffs(name, degree=degree) for ref in self.references]
         offsets = numpy.cumsum([0] + [len(c) for c in coeffs])
-        dofmap = numpy.repeat(-1, offsets[-1])
-        for ielem, ioppelems in enumerate(self.connectivity):
-            for iedge, jelem in enumerate(ioppelems):  # loop over element neighbors and merge dofs
-                if jelem < ielem:
-                    continue  # either there is no neighbor along iedge or situation will be inspected from the other side
-                jedge = util.index(self.connectivity[jelem], ielem)
-                idofs = offsets[ielem] + self.references[ielem].get_edge_dofs(degree, iedge)
-                jdofs = offsets[jelem] + self.references[jelem].get_edge_dofs(degree, jedge)
-                for idof, jdof in zip(idofs, jdofs):
-                    while dofmap[idof] != -1:
-                        idof = dofmap[idof]
-                    while dofmap[jdof] != -1:
-                        jdof = dofmap[jdof]
-                    if idof != jdof:
-                        dofmap[max(idof, jdof)] = min(idof, jdof)  # create left-looking pointer
-        # assign dof numbers left-to-right
-        ndofs = 0
-        for i, n in enumerate(dofmap):
-            if n == -1:
-                dofmap[i] = ndofs
-                ndofs += 1
-            else:
-                dofmap[i] = dofmap[n]
+        # To merge matching dofs we loop over the connectivity table to find
+        # neighbouring elements (limited to jelem > ielem to consider every
+        # neighbour pair exactly once as well as ignore exterior boundaries)
+        # and mark the degrees of freedom on both sides to be equal.
+        dofmap, ndofs = util.merge_index_map(offsets[-1], (merge_set
+            for ielem, ioppelems in enumerate(self.connectivity)
+                for iedge, jelem in enumerate(ioppelems) if jelem >= ielem
+                    for merge_set in zip(
+                        offsets[ielem] + self.references[ielem].get_edge_dofs(degree, iedge),
+                        offsets[jelem] + self.references[jelem].get_edge_dofs(degree, util.index(self.connectivity[jelem], ielem)))))
 
         elem_slices = map(slice, offsets[:-1], offsets[1:])
         dofs = tuple(types.frozenarray(dofmap[s]) for s in elem_slices)
@@ -3135,15 +3122,9 @@ def basis_spline(self, degree, patchcontinuous=True, knotvalues=None, knotmultip
         if patchcontinuous:
             # build merge mapping: merge common boundary dofs (from low to high)
             pairs = itertools.chain(*(zip(*dofs) for dofs in commonboundarydofs.values() if len(dofs) > 1))
-            merge = numpy.arange(dofcount)
-            for dofs in sorted(pairs):
-                merge[list(dofs)] = merge[list(dofs)].min()
-            assert all(numpy.all(merge[a] == merge[b]) for a, *B in commonboundarydofs.values() for b in B), 'something went wrong is merging interface dofs; this should not have happened'
-            # build renumber mapping: renumber remaining dofs consecutively, starting at 0
-            remainder, renumber = numpy.unique(merge, return_inverse=True)
+            renumber, dofcount = util.merge_index_map(dofcount, pairs)
             # apply mappings
             dofmap = tuple(types.frozenarray(renumber[v], copy=False) for v in dofmap)
-            dofcount = len(remainder)
 
         return function.PlainBasis(coeffs, dofmap, dofcount, self.f_index, self.f_coords)
 

tests/test_topology.py

@@ -1145,6 +1145,18 @@ def test_connectivity(self):
             self.assertEqual(opp1, interfaces2.opposites[i2])
 
 
+class multipatch_misc(TestCase):
+
+    def test_basis_merge_dofs(self):
+        patches = (0, 4, 2, 5), (9, 1, 5, 3), (2, 5, 6, 7), (5, 3, 7, 8)
+        domain, geom = mesh.multipatch(patches=patches, nelems=1)
+        basis = domain.basis('spline', 1)
+        actual = domain.sample('bezier', 2).eval(basis)
+        desired = numpy.zeros((16, 10))
+        desired[numpy.arange(16), [0, 1, 2, 3, 4, 5, 3, 6, 2, 3, 7, 8, 3, 6, 8, 9]] = 1
+        numpy.testing.assert_allclose(actual, desired)
+
+
 class multipatch_errors(TestCase):
 
     def test_reverse(self):

tests/test_util.py

@@ -387,3 +387,26 @@ def test(self):
         self.assertEqual(util.deep_reduce(min, L), 1)
         self.assertEqual(util.deep_reduce(max, L), 7)
         self.assertEqual(util.deep_reduce(sum, L), 28)
+
+
+class merge_index_map(TestCase):
+
+    def test_empty_merge_sets(self):
+        index_map, count = util.merge_index_map(4, [])
+        self.assertEqual(index_map.tolist(), [0, 1, 2, 3])
+        self.assertEqual(count, 4)
+
+    def test_merge_set_one(self):
+        index_map, count = util.merge_index_map(4, [[1]])
+        self.assertEqual(index_map.tolist(), [0, 1, 2, 3])
+        self.assertEqual(count, 4)
+
+    def test_multihop1(self):
+        index_map, count = util.merge_index_map(4, [[0, 2], [1, 3], [0, 3]])
+        self.assertEqual(index_map.tolist(), [0, 0, 0, 0])
+        self.assertEqual(count, 1)
+
+    def test_multihop2(self):
+        index_map, count = util.merge_index_map(4, [[2, 3], [1, 2], [0, 3]])
+        self.assertEqual(index_map.tolist(), [0, 0, 0, 0])
+        self.assertEqual(count, 1)