From a3fb774ac0ce98c3ee6d5cd0491936c54dace8bb Mon Sep 17 00:00:00 2001 From: Joost van Zwieten Date: Wed, 29 Nov 2023 13:10:46 +0100 Subject: [PATCH] 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 | 61 ++++++++++++++++++++++++++++++++++++++++++ tests/test_function.py | 15 +++++++++++ 2 files changed, 76 insertions(+) diff --git a/nutils/function.py b/nutils/function.py index 89e12381e..9649d8903 100644 --- a/nutils/function.py +++ b/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] @@ -2520,6 +2521,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`. @@ -2776,6 +2801,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) diff --git a/tests/test_function.py b/tests/test_function.py index 71fbe4942..8ca1e98dc 100644 --- a/tests/test_function.py +++ b/tests/test_function.py @@ -1288,6 +1288,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):