add Topology.partition_interfaces

nutils/function.py

@@ -97,13 +97,13 @@ def empty(cls, points_shape=()) -> 'LowerArgs':
         return cls(points_shape, ())
 
     @classmethod
-    def for_space(cls, space: str, transforms: Tuple[Transforms, ...], index: evaluable.Array, coordinates: evaluable.Array) -> 'LowerArgs':
+    def for_space(cls, space: str, transforms: Tuple[Transforms, ...], index: evaluable.Array, coordinates: Optional[evaluable.Array] = None) -> 'LowerArgs':
         if index.dtype != int or index.ndim != 0:
             raise ValueError('argument `index` must be a scalar, integer `nutils.evaluable.Array`')
         args = LowerArg(space, transforms[0], index, coordinates),
         if len(transforms) > 1:
             args += LowerArg('~' + space, transforms[-1], index, coordinates),
-        return cls(coordinates.shape[:-1], args)
+        return cls(() if coordinates is None else coordinates.shape[:-1], args)
 
     def __post_init__(self):
         assert isinstance(self.points_shape, tuple) and all(map(evaluable._isindex, self.points_shape))

nutils/topology.py

@@ -414,6 +414,77 @@
 
         raise NotImplementedError
 
+    def elementwise_stack(self, func: function.Array, *, axis: int = 0) -> function.Array:
+        '''Return the stack of the function bound to every element of this topology.
+
+        The order of the stack matches the order of the elements of this
+        topology. Therefor, for any :class:`Topology` ``topo``
+
+            topo.elementwise_stack(topo.f_index)
+
+        is guaranteed to be the same (after evaluation) as
+
+            numpy.arange(len(topo))
+
+        Parameters
+        ----------
+        func : :class:`nutils.function.Array`
+            The function to bind to elements and stack.
+        axis : :class:`int`
+            The axis in the result array along which the elements are stacked.
+
+        Returns
+        -------
+        :class:`nutils.function.Array`
+            The stack of function at the first element, the second, etc. until
+            the last.
+
+        Examples
+        --------
+
+        Given a ``base`` :class:`Topology` and a ``refinement``, with
+        :meth:`Topology.elementwise_stack` we can obtain the element index of
+        the ``base`` for each element in the ``refinement``:
+
+        >>> from nutils import function, mesh
+        >>> base, geom = mesh.line(4)
+        >>> refinement = base.refined_by([1, 3]).refined_by([5])
+        >>> function.eval(refinement.elementwise_stack(base.f_index))
+        array([0, 2, 1, 1, 3, 3, 3])
+
+        Note that refined elements are positioned at the end of the refined
+        topology.
+
+        Notes
+        -----
+
+        This method delegates the work to
+        :meth:`Topology.elementwise_stack_last_axis`. Implementations of
+        :class:`Topology` should implement the latter.
+        '''
+
+        stack = self.elementwise_stack_last_axis(func)
+        axis = numeric.normdim(stack.ndim, axis)
+        if axis != stack.ndim - 1:
+            stack = numpy.transpose(stack, (*range(axis), *range(axis + 1, stack.ndim), axis))
+        return stack
+
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        '''Return the stack along the last axis of the function bound to every element of this topology.
+
+        Parameters
+        ----------
+        func : :class:`nutils.function.Array`
+            The function to bind to elements and stack.
+
+        Returns
+        -------
+        :class:`nutils.function.Array`
+            The stack of function at element 0, element 1, etc along the last axis.
+        '''
+
+        raise NotImplementedError
+
     @single_or_multiple
     def integrate_elementwise(self, funcs: Iterable[function.Array], *, degree: int, asfunction: bool = False, ischeme: str = 'gauss', arguments: Optional[_ArgDict] = None) -> Union[List[numpy.ndarray], List[function.Array]]:
         'element-wise integration'
@@ -992,6 +1063,10 @@
 
     @cached_property
     def interfaces(self) -> 'Topology':
+        '''
+        :class:`Topology`:
+          The interfaces of this topology.
+        '''
         return self.interfaces_spaces(self.spaces)
 
     def interfaces_spaces(self, __spaces: Iterable[str]) -> 'Topology':
@@ -1045,6 +1120,35 @@
 
         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'
 
@@ -1184,6 +1288,10 @@
     def sample(self, ischeme: str, degree: int) -> Sample:
         return Sample.empty(self.spaces, self.ndims)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        func = function.Array.cast(func)
+        return function.zeros((*func.shape, 0), dtype=func.dtype)
+
 
 class _DisjointUnion(_TensorialTopology):
 
@@ -1238,6 +1346,9 @@
     def sample(self, ischeme: str, degree: int) -> Sample:
         return self.topo1.sample(ischeme, degree) + self.topo2.sample(ischeme, degree)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        return numpy.concatenate([self.topo1.elementwise_stack_last_axis(func), self.topo2.elementwise_stack_last_axis(func)], axis=-1)
+
     def trim(self, levelset: function.Array, maxrefine: int, ndivisions: int = 8, name: str = 'trimmed', leveltopo: Optional[Topology] = None, *, arguments: Optional[_ArgDict] = None) -> Topology:
         if leveltopo is not None:
             return super().trim(levelset, maxrefine, ndivisions, name, leveltopo, arguments=arguments)
@@ -1401,6 +1512,10 @@
         # tri and hull the mistake is without consequences.
         return Sample.zip(sample1, sample2)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        func = function.Array.cast(func)
+        return numpy.reshape(self.topo2.elementwise_stack_last_axis(self.topo1.elementwise_stack_last_axis(func)), (*func.shape, len(self)))
+
 
 class _Take(_TensorialTopology):
 
@@ -1417,6 +1532,9 @@
     def sample(self, ischeme: str, degree: int) -> Sample:
         return self.parent.sample(ischeme, degree).take_elements(self.indices)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        return numpy.take(self.parent.elementwise_stack_last_axis(func), self.indices, axis=-1)
+
 
 class _WithGroupAliases(_TensorialTopology):
 
@@ -1463,6 +1581,9 @@
     def sample(self, ischeme: str, degree: int) -> Sample:
         return self.parent.sample(ischeme, degree)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        return self.parent.elementwise_stack_last_axis(func)
+
     def refine_spaces_unchecked(self, spaces: FrozenSet[str]) -> Topology:
         return _WithGroupAliases(self.parent.refine_spaces(spaces), self.vgroups, self.bgroups, self.igroups)
 
@@ -1482,6 +1603,22 @@
         return _WithGroupAliases(self.parent.interfaces_spaces_unchecked(spaces), self.igroups, types.frozendict({}), types.frozendict({}))
 
 
+class _ElementwiseStack(function.Array):
+
+    def __init__(self, func: function.Array, space: str, transforms: transformseq.Transforms, opposites: transformseq.Transforms):
+        self.func = func
+        self.space = space
+        self.transforms = transforms
+        self.opposites = opposites
+        self.nelems = len(self.transforms)
+        super().__init__((*func.shape, self.nelems), func.dtype, func.spaces - frozenset({space}), func.arguments)
+
+    def lower(self, args: function.LowerArgs) -> evaluable.Array:
+        index = evaluable.loop_index(f'_loop_{self.space}', evaluable.asarray(self.nelems))
+        func = self.func.lower(args * function.LowerArgs.for_space(self.space, (self.transforms, self.opposites), index))
+        return evaluable.loop_concatenate(evaluable.appendaxes(func, (evaluable.asarray(1),)), index)
+
+
 class TransformChainsTopology(Topology):
     'base class for topologies with transform chains'
 
@@ -1589,6 +1726,9 @@
             transforms += self.opposites,
         return Sample.new(self.space, transforms, points)
 
+    def elementwise_stack_last_axis(self, func: function.Array) -> function.Array:
+        return _ElementwiseStack(function.Array.cast(func), self.space, self.transforms, self.opposites)
+
     def _refined_by(self, refine):
         fine = self.refined.transforms
         indices0 = numpy.setdiff1d(numpy.arange(len(self)), refine, assume_unique=True)

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):
 
@@ -1349,6 +1370,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):