diff --git a/nutils/function.py b/nutils/function.py index 8eceebe61..6231679df 100644 --- a/nutils/function.py +++ b/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)) diff --git a/nutils/topology.py b/nutils/topology.py index 0acf6b7b3..ad6283552 100644 --- a/nutils/topology.py +++ b/nutils/topology.py @@ -414,6 +414,77 @@ def sample(self, ischeme: str, degree: int) -> Sample: 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 @@ def boundary_spaces_unchecked(self, __spaces: FrozenSet[str]) -> 'Topology': @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 @@ 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' @@ -1184,6 +1288,10 @@ def basis_std(self, degree: int, *args, **kwargs) -> function.Array: 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 integrate_elementwise(self, funcs: Iterable[function.Array], *, degree: int, 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 @@ def _sample(self, ielems, coords, weights=None): # 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 __init__(self, parent: Topology, indices: numpy.ndarray) -> None: 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 _tensorial_bases(self, btype, **kwargs) -> function.Basis: 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 @@ def interfaces_spaces_unchecked(self, spaces: FrozenSet[str]) -> Topology: 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 @@ def sample(self, ischeme, degree): 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) diff --git a/tests/test_topology.py b/tests/test_topology.py index 0a903dfb3..02a2e549d 100644 --- a/tests/test_topology.py +++ b/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):