add Basis.discontinuous_at_partition_interfaces

This patch adds the `Basis.discontinuous_at_partition_interfaces` method which
makes the basis discontinuous at element partition interfaces by clipping each
basis function to each part on which the function has support.

nutils/function.py

@@ -17,6 +17,7 @@
 import numbers
 import inspect
 import fractions
+import treelog
 
 IntoArray = Union['Array', numpy.ndarray, bool, int, float, complex]
 Shape = Sequence[int]
@@ -2525,6 +2526,30 @@ def __getitem__(self, index: Any) -> Array:
         else:
             return super().__getitem__(index)
 
+    @treelog.withcontext
+    def discontinuous_at_partition_interfaces(self, part_indices: Sequence[int]):
+        '''Returns a basis that is discontinuous at element partition interfaces.
+
+        Given a partition of elements, this basis is made discontinuous at the
+        partition interfaces. All elements that have the same part index belong
+        to the same part.
+
+        The returned basis is formed by clipping each function of the basis to
+        each part individually and stacking all nonzero clipped functions. As a
+        consequence, if a basis function has support on three topologically
+        adjacent elements of which the first and the last element belong to one
+        part and the middle to another, this function will not be clipped to
+        each of the three elements individually, but to the first and the last
+        element and to the middle element.
+
+        Parameters
+        ----------
+        part_indices : sequence or :class:`numpy.ndarray` of :class:`int`
+            For each element the index of the part the element belongs to.
+        '''
+
+        return _DiscontinuousPartitionBasis(self, part_indices)
+
 
 class PlainBasis(Basis):
     '''A general purpose implementation of a :class:`Basis`.
@@ -2781,6 +2806,42 @@ def f_dofs_coeffs(self, index: evaluable.Array) -> Tuple[evaluable.Array, evalua
         return dofs, p_coeffs
 
 
+class _DiscontinuousPartitionBasis(Basis):
+
+    def __init__(self, parent: Basis, part_indices: Sequence[int]):
+        self._parent = parent
+
+        part_indices = numpy.array(part_indices).astype(int, casting='safe', copy=False)
+        if part_indices.shape != (parent.nelems,):
+            raise ValueError(f'expected a sequence of {self.nelems} integers but got an array with shape {part_indices.shape}')
+
+        # For each element we pair the parent dofs with the part indices and
+        # use that as partitioned dof. Then we renumber the partitioned dofs
+        # starting at 0, ordered by part index, then by parent dof.
+
+        ielem = evaluable.loop_index('ielem', parent.nelems)
+        parent_dofs, _ = parent.f_dofs_coeffs(ielem)
+        # Concatenate all parent dofs and stack all ndofs per element.
+        cc_parent_dofs = evaluable.loop_concatenate(parent_dofs, ielem)
+        cc_ndofs = evaluable.loop_concatenate(evaluable.insertaxis(parent_dofs.shape[0], 0, evaluable.asarray(1)), ielem)
+        cc_parent_dofs, cc_ndofs = evaluable.Tuple((cc_parent_dofs, cc_ndofs)).optimized_for_numpy.eval()
+        # Stack the part index for each element for each dof.
+        cc_part_indices = numpy.repeat(part_indices, cc_ndofs)
+        # Renumber and count all unique dofs.
+        unique_dofs, dofs = numpy.unique(numpy.stack([cc_part_indices, cc_parent_dofs], axis=1), axis=0, return_inverse=True)
+
+        self._dofs = evaluable.asarray(dofs)
+        self._ndofs = evaluable.asarray(cc_ndofs)
+        self._offsets = evaluable._SizesToOffsets(self._ndofs)
+
+        super().__init__(len(unique_dofs), parent.nelems, parent.index, parent.coords)
+
+    def f_dofs_coeffs(self, index):
+        dofs = evaluable.Take(self._dofs, evaluable.Range(evaluable.Take(self._ndofs, index)) + evaluable.Take(self._offsets, index))
+        _, coeffs = self._parent.f_dofs_coeffs(index)
+        return dofs, coeffs
+
+
 def Namespace(*args, **kwargs):
     from .expression_v1 import Namespace
     return Namespace(*args, **kwargs)

tests/test_function.py

@@ -1281,6 +1281,21 @@ def setUp(self):
         super().setUp()
 
 
+class DiscontinuousPartitionBasis(CommonBasis, TestCase):
+
+    def setUp(self):
+        self.checktransforms = transformseq.IndexTransforms(0, 4)
+        index, coords = self.mk_index_coords(0, self.checktransforms)
+        self.checkcoeffs = [[[1.], [2.], [3.], [4.]]] * 4
+        parent_dofs = [[0, 1, 3, 4], [1, 2, 4, 5], [3, 4, 6, 7], [4, 5, 7, 8]]
+        part_indices = [-1, 2, -1, 3]
+        self.basis = function.PlainBasis(self.checkcoeffs, parent_dofs, 9, index, coords).discontinuous_at_partition_interfaces(part_indices)
+        self.checkdofs = [[0, 1, 2, 3], [6, 7, 8, 9], [2, 3, 4, 5], [10, 11, 12, 13]]
+        self.checkndofs = 14
+        self.checkmasks = [[i in [0, 1, 7, 11] for i in range(self.checkndofs)]]
+        super().setUp()
+
+
 @parametrize
 class SurfaceGradient(TestCase):