add Topology.partition_interfaces

Author: joostvanzwieten

nutils/topology.py

@@ -1049,6 +1049,35 @@ def interfaces_spaces_unchecked(self, __spaces: FrozenSet[str]) -> 'Topology':
 
         raise NotImplementedError
 
+    def partition_interfaces(self, part_indices: Sequence[int]) -> 'Topology':
+        '''Return the interfaces between all parts of an element partition.
+
+        Given a partition of elements, this function returns the subset of
+        interfaces that face two different parts. The orientation of the
+        interfaces is arbitrary.
+
+        Parameters
+        ----------
+        part_indices : sequence or :class:`numpy.ndarray` of :class:`int`
+            For each element the index of the part the element belongs to.
+
+        Returns
+        -------
+        :class:`Topology`
+            The interfaces between all parts of the element partition, a subset
+            of :attr:`Topology.interfaces`.
+        '''
+
+        part_indices = function.Array.cast(part_indices)
+        if part_indices.dtype not in (bool, int):
+            raise ValueError(f'expected a sequence of integer part indices but got a sequence of type {part_indices.dtype}')
+        if part_indices.shape != (len(self),):
+            raise ValueError(f'expected a sequence of {len(self)} integer part indices but got an array with shape {part_indices.shape}')
+        interfaces = self.interfaces
+        part_indices = numpy.take(part_indices, self.f_index)
+        isnt_partition_interface = interfaces.elementwise_stack(part_indices == function.opposite(part_indices))
+        return interfaces.compress(~function.eval(isnt_partition_interface))
+
     def basis_discont(self, degree: int) -> function.Basis:
         'discontinuous shape functions'
 

tests/test_topology.py

@@ -174,6 +174,9 @@ def test_select(self):
             self.assertVertices(self.topo.select(((self.geom - center) * direction).sum()), desired_vertices)
             direction = numpy.roll(direction, shift=1)
 
+    def test_elementwise_stack(self):
+        self.assertEqual(function.eval(self.topo.elementwise_stack(self.topo.f_index)).tolist(), list(range(len(self.topo))))
+
 
 class ConformingTests:
 
@@ -456,8 +459,8 @@ class NewDisjointUnion(TestCase, CommonTests, ConformingTests):
 
     def setUp(self):
         super().setUp()
-        topo, self.geom = mesh.newrectilinear([8, 3], spaces='XY')
-        self.topo = topology.Topology.disjoint_union(topo.slice(slice(0, 3), 0), topo.slice(slice(4, 8), 0).slice(slice(0, 2), 1))
+        self.basetopo, self.geom = mesh.newrectilinear([8, 3], spaces='XY')
+        self.topo = topology.Topology.disjoint_union(self.basetopo.slice(slice(0, 3), 0), self.basetopo.slice(slice(4, 8), 0).slice(slice(0, 2), 1))
         self.desired_spaces = 'X', 'Y'
         self.desired_space_dims = 1, 1
         self.desired_ndims = 2
@@ -501,6 +504,12 @@ def test_trim(self):
         self.assertEqual(as_rounded_list(topo.trim(x-2.5, maxrefine=0).volume(x[None])), 0.5)
         self.assertEqual(as_rounded_list(topo.trim(0.5-x, maxrefine=0).volume(x[None])), 0.5)
 
+    def test_elementwise_stack(self):
+        self.assertEqual(
+            function.eval(self.topo.elementwise_stack(self.basetopo.f_index)).tolist(),
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 15, 16, 18, 19, 21, 22],
+        )
+
 
 class NewMul(TestCase, CommonTests, ConformingTests):
 
@@ -611,6 +620,18 @@ def test_basis(self):
         with self.assertRaises(ValueError):
             self.topo.basis('spline', degree=1, knotmultiplicities=['a', 'b'])
 
+    def test_elementwise_stack_mul(self):
+        desired = numpy.mgrid[:len(self.topo1), :len(self.topo2)].reshape(2, len(self.topo))
+        stack = self.topo.elementwise_stack(numpy.stack([self.topo1.f_index, self.topo2.f_index]))
+        self.assertEqual(function.eval(stack).tolist(), desired.T.tolist())
+        stack = self.topo.elementwise_stack(numpy.stack([self.topo1.f_index, self.topo2.f_index]), axis=1)
+        self.assertEqual(function.eval(stack).tolist(), desired.tolist())
+
+    def test_partition_interfaces(self):
+        centers = self.topo.partition_interfaces([0, 0, 1, 2, 1, 1]).sample('gauss', 0).eval(self.geom).tolist()
+        centers.sort()
+        self.assertAllAlmostEqual(centers, [[0.5, 2.0], [1.0, 0.5], [1.0, 1.5], [1.5, 1.0]])
+
 
 class NewWithGroupAliases(TestCase, CommonTests, ConformingTests):
 
@@ -1337,6 +1358,11 @@ def setUp(self):
         self.desired_references = [element.LineReference()**2]*4
         self.desired_vertices = [[[x, y] for x in X for y in Y] for X in pairwise(range(3)) for Y in pairwise(range(3))]
 
+    def test_partition_interfaces(self):
+        centers = self.topo.partition_interfaces([0, 0, 1, 2]).sample('gauss', 0).eval(self.geom).tolist()
+        centers.sort()
+        self.assertAllAlmostEqual(centers, [[1.0, 0.5], [1.0, 1.5], [1.5, 1.0]])
+
 
 class UnionTopology(TestCase, CommonTests, TransformChainsTests):