Nested integrals (#914)

Add support for nested integrals over the same space.

Author: joostvanzwieten

nutils/SI.py

@@ -346,6 +346,7 @@ def r(dispatch_func):
     @register(function.jump)
     @register(function.kronecker)
     @register(function.linearize)
+    @register(function.swap_spaces)
     @register(function.opposite)
     @register(function.replace_arguments)
     @register(function.scatter)

nutils/__init__.py

@@ -1,4 +1,4 @@
 'Numerical Utilities for Finite Element Analysis'
 
-__version__ = version = '9a62'
+__version__ = version = '9a63'
 version_name = 'jook-sing'

nutils/function.py

@@ -91,6 +91,9 @@ def __init__(self, spaces: Tuple[str, ...], ndims: int, nelems: int, npoints: in
             The number of points.
         '''
 
+        for i, space in enumerate(spaces):
+            if space in spaces[i+1:]:
+                raise ValueError(f'All spaces in a `Sample` must be unique, but space {space} is repeated.')
         self.spaces = spaces
         self.ndims = ndims
         self.nelems = nelems
@@ -479,6 +482,21 @@ def zip(*samples: 'Sample') -> 'Sample':
 
         return _Zip(*samples)
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> 'Sample':
+        '''Return a :class:`Sample` with spaces renamed according to ``map``.
+
+        Args
+        ----
+        map : mapping of :class:`str` to :class:`str`
+            A mapping of old space to new space.
+
+        Returns
+        -------
+        renamed : :class:`Sample`
+        '''
+
+        raise NotImplementedError
+
 
 class _TransformChainsSample(Sample):
 
@@ -554,6 +572,9 @@ def get_evaluable_indices(self, ielem: evaluable.Array) -> evaluable.Array:
     def _bind(self, func: function.Array) -> function.Array:
         return _ConcatenatePoints(func, self)
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _DefaultIndex(map.get(self.space, self.space), self.transforms, self.points)
+
 
 class _CustomIndex(_TransformChainsSample):
 
@@ -578,6 +599,9 @@ def tri(self) -> numpy.ndarray:
     def hull(self) -> numpy.ndarray:
         return numpy.take(self._index, self._parent.hull)
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _CustomIndex(self._parent.rename_spaces(map), self._index)
+
 
 if os.environ.get('NUTILS_TENSORIAL', None) == 'test':  # pragma: nocover
 
@@ -627,7 +651,7 @@ def get_evaluable_weights(self, __ielem: evaluable.Array) -> evaluable.Array:
         return evaluable.Zeros((evaluable.constant(0),) * len(self.spaces), dtype=float)
 
     def get_lower_args(self, __ielem: evaluable.Array) -> function.LowerArgs:
-        return function.LowerArgs((), {}, {})
+        return function.LowerArgs.empty()
 
     def get_element_tri(self, ielem: int) -> numpy.ndarray:
         raise IndexError('index out of range')
@@ -647,6 +671,9 @@ def _bind(self, func: function.Array) -> function.Array:
     def basis(self, interpolation: str = 'none') -> function.Array:
         return function.zeros((0,), float)
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _Empty(tuple(map.get(space, space) for space in self.spaces), self.ndims)
+
 
 class _Add(_TensorialSample):
 
@@ -694,6 +721,9 @@ def _integral(self, func: function.Array) -> function.Array:
     def _bind(self, func: function.Array) -> function.Array:
         return numpy.concatenate([self._sample1._bind(func), self._sample2._bind(func)])
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _Add(self._sample1.rename_spaces(map), self._sample2.rename_spaces(map))
+
 
 def _simplex_strip(strip):
     # Helper function that creates simplices for an extruded simplex, with
@@ -807,7 +837,7 @@ def get_evaluable_weights(self, __ielem: evaluable.Array) -> evaluable.Array:
 
     def get_lower_args(self, __ielem: evaluable.Array) -> function.LowerArgs:
         ielem1, ielem2 = evaluable.divmod(__ielem, self._sample2.nelems)
-        return self._sample1.get_lower_args(ielem1) | self._sample2.get_lower_args(ielem2)
+        return self._sample1.get_lower_args(ielem1) * self._sample2.get_lower_args(ielem2)
 
     @property
     def _reversed_factors(self):
@@ -900,6 +930,9 @@ def basis(self, interpolation: str = 'none') -> Sample:
         assert basis1.ndim == basis2.ndim == 1
         return numpy.ravel(basis1[:, None] * basis2[None, :])
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _Mul(self._sample1.rename_spaces(map), self._sample2.rename_spaces(map))
+
 
 class _Zip(Sample):
 
@@ -940,15 +973,16 @@ def _getslice(self, ielem):
 
     def get_lower_args(self, __ielem: evaluable.Array) -> function.LowerArgs:
         points_shape = evaluable.Take(evaluable.Constant(self._sizes), __ielem),
-        coordinates = {}
-        transform_chains = {}
+        args = function.LowerArgs.empty(points_shape)
         for samplei, ielemsi, ilocalsi in zip(self._samples, self._ielems, self._ilocals):
-            argsi = samplei.get_lower_args(evaluable.Take(evaluable.Constant(ielemsi), __ielem))
             slicei = evaluable.Take(evaluable.Constant(ilocalsi), self._getslice(__ielem))
-            transform_chains.update(argsi.transform_chains)
-            for space, coords in argsi.coordinates.items():
-                coordinates[space] = evaluable.Transpose.to_end(evaluable.Take(evaluable._flat(evaluable.Transpose.from_end(coords, 0), ndim=2), slicei), 0)
-        return function.LowerArgs(points_shape, transform_chains, coordinates)
+            args += samplei \
+                .get_lower_args(evaluable.Take(evaluable.Constant(ielemsi), __ielem)) \
+                .map_coordinates(
+                    points_shape,
+                    lambda coords: evaluable.Transpose.to_end(evaluable.Take(evaluable._flat(evaluable.Transpose.from_end(coords, 0), ndim=2), slicei), 0),
+                )
+        return args
 
     def get_evaluable_indices(self, ielem):
         return evaluable.Take(evaluable.Constant(self._indices), self._getslice(ielem))
@@ -959,6 +993,9 @@ def get_evaluable_weights(self, ielem):
         weights = self._samples[0].get_evaluable_weights(ielem0)
         return evaluable._take(evaluable._flat(weights), slice0, axis=0)
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _Zip(*[smpl.rename_spaces(map) for smpl in self._samples])
+
 
 class _TakeElements(_TensorialSample):
 
@@ -1014,6 +1051,9 @@ def get_element_hull(self, __ielem: int) -> numpy.ndarray:
     def take_elements(self, __indices: numpy.ndarray) -> Sample:
         return self._parent.take_elements(numpy.take(self._indices, __indices))
 
+    def rename_spaces(self, map: Mapping[str, str], /) -> Sample:
+        return _TakeElements(self._parent.rename_spaces(map), self._indices)
+
 
 class _Integral(function.Array):
 
@@ -1023,9 +1063,9 @@ def __init__(self, integrand: function.Array, sample: Sample) -> None:
         super().__init__(shape=integrand.shape, dtype=float if integrand.dtype in (bool, int) else integrand.dtype, spaces=integrand.spaces - frozenset(sample.spaces), arguments=integrand.arguments)
 
     def lower(self, args: function.LowerArgs) -> evaluable.Array:
-        ielem = evaluable.loop_index('_sample_' + '_'.join(self._sample.spaces), self._sample.nelems)
+        ielem = evaluable.loop_index(f'_sample_{len(args.args)}', self._sample.nelems)
         weights = evaluable.astype(self._sample.get_evaluable_weights(ielem), self.dtype)
-        integrand = evaluable.astype(self._integrand.lower(args | self._sample.get_lower_args(ielem)), self.dtype)
+        integrand = evaluable.astype(self._integrand.lower(args * self._sample.get_lower_args(ielem)), self.dtype)
         elem_integral = evaluable.einsum('B,ABC->AC', weights, integrand, B=weights.ndim, C=self.ndim)
         return evaluable.loop_sum(elem_integral, ielem)
 
@@ -1039,8 +1079,8 @@ def __init__(self, func: function.Array, sample: _TransformChainsSample) -> None
 
     def lower(self, args: function.LowerArgs) -> evaluable.Array:
         axis = len(args.points_shape)
-        ielem = evaluable.loop_index('_sample_' + '_'.join(self._sample.spaces), self._sample.nelems)
-        args |= self._sample.get_lower_args(ielem)
+        ielem = evaluable.loop_index(f'_sample_{len(args.args)}', self._sample.nelems)
+        args *= self._sample.get_lower_args(ielem)
         func = self._func.lower(args)
         func = evaluable.Transpose.to_end(func, *range(axis, len(args.points_shape)))
         for i in range(len(args.points_shape) - axis - 1):
@@ -1073,7 +1113,8 @@ def __init__(self, sample: _TransformChainsSample, interpolation: str) -> None:
         super().__init__(shape=(sample.npoints,), dtype=float, spaces=frozenset({sample.space}), arguments={})
 
     def lower(self, args: function.LowerArgs) -> evaluable.Array:
-        aligned_space_coords = args.coordinates[self._sample.space]
+        arg = args[self._sample.space]
+        aligned_space_coords = arg.coordinates
         assert aligned_space_coords.ndim == len(args.points_shape) + 1
         space_coords, where = evaluable.unalign(aligned_space_coords)
         # Reinsert the coordinate axis, the last axis of `aligned_space_coords`, or
@@ -1085,9 +1126,8 @@ def lower(self, args: function.LowerArgs) -> evaluable.Array:
             space_coords = evaluable.Transpose(space_coords, numpy.argsort(where))
             where = tuple(sorted(where))
 
-        (chain, *_), tip_index = args.transform_chains[self._sample.space]
-        index = evaluable.TransformIndex(self._sample.transforms[0], chain, tip_index)
-        coords = evaluable.TransformCoords(self._sample.transforms[0], chain, tip_index, space_coords)
+        index = evaluable.TransformIndex(self._sample.transforms[0], arg.transforms, arg.index)
+        coords = evaluable.TransformCoords(self._sample.transforms[0], arg.transforms, arg.index, space_coords)
         expect = self._sample.points.get_evaluable_coords(index)
         sampled = evaluable.Sampled(coords, expect, self._interpolation)
         indices = self._sample.get_evaluable_indices(index)

nutils/sample.py

@@ -1388,7 +1388,7 @@ def _lower_args(self, ielem, point):
         ielem1, ielem2 = evaluable.divmod(ielem, len(self.topo2))
         largs1 = self.topo1._lower_args(ielem1, point[:self.topo1.ndims])
         largs2 = self.topo2._lower_args(ielem2, point[self.topo1.ndims:])
-        return largs1 | largs2
+        return largs1 * largs2
 
     def _sample(self, ielems, coords, weights=None):
         ielems1, ielems2 = divmod(ielems, len(self.topo2))

nutils/topology.py

@@ -418,15 +418,16 @@ def _check(name, op, n_op, *args):
 class Unlower(TestCase):
 
     def test(self):
-        e = evaluable.Argument('arg', tuple(map(evaluable.constant, (2, 3, 4, 5))), int)
+        shape = tuple(map(evaluable.constant, (2, 3, 4, 5)))
+        e = evaluable.Argument('arg', shape, int)
         arguments = types.frozendict({'arg': ((2, 3), int)})
-        f = function._Unlower(e, frozenset(), arguments, function.LowerArgs((2, 3), {}, {}))
+        f = function._Unlower(e, frozenset(), arguments, function.LowerArgs.empty(shape[:2]))
         self.assertEqual(f.shape, (4, 5))
         self.assertEqual(f.dtype, int)
         self.assertEqual(f.arguments, arguments)
-        self.assertEqual(f.lower(function.LowerArgs((2, 3), {}, {})), e)
+        self.assertEqual(f.lower(function.LowerArgs.empty(shape[:2])), e)
         with self.assertRaises(ValueError):
-            f.lower(function.LowerArgs((3, 4), {}, {}))
+            f.lower(function.LowerArgs.empty(shape[1:3]))
 
 
 class Custom(TestCase):
@@ -435,10 +436,10 @@ def assertEvalAlmostEqual(self, factual, fdesired, **args):
         with self.subTest('0d-points'):
             self.assertAllAlmostEqual(evaluable.eval_once(factual.as_evaluable_array, arguments=args), evaluable.eval_once(fdesired.as_evaluable_array, arguments=args))
         with self.subTest('1d-points'):
-            lower_args = function.LowerArgs((evaluable.asarray(5),), {}, {})
+            lower_args = function.LowerArgs.empty((evaluable.asarray(5),))
             self.assertAllAlmostEqual(evaluable.eval_once(factual.lower(lower_args), arguments=args), evaluable.eval_once(fdesired.lower(lower_args), arguments=args))
         with self.subTest('2d-points'):
-            lower_args = function.LowerArgs((evaluable.asarray(5), evaluable.asarray(6)), {}, {})
+            lower_args = function.LowerArgs.empty((evaluable.asarray(5), evaluable.asarray(6)))
             self.assertAllAlmostEqual(evaluable.eval_once(factual.lower(lower_args), arguments=args), evaluable.eval_once(fdesired.lower(lower_args), arguments=args))
 
     def assertMultipy(self, leftval, rightval):
@@ -1573,7 +1574,7 @@ def test_constant(self):
 
     def test_lower_spaces(self):
         topo, geom = mesh.rectilinear([3])
-        with self.assertRaisesRegex(ValueError, r'cannot lower function with spaces \(.+\) - did you forget integral or sample?'):
+        with self.assertRaisesRegex(KeyError, r'no such space: .* - did you forget integral or sample?'):
             function.factor(geom)
 
 
@@ -1602,3 +1603,21 @@ def test_conflict(self):
         g = x**2 * y2
         with self.assertRaisesRegex(ValueError, "inconsistent shapes for argument 'y'"):
             function.arguments_for(f, g)
+
+
+class swap_spaces(TestCase):
+
+    def test_different(self):
+        X, x = mesh.line(2, space='X')
+        self.assertEqual(X.f_index.spaces, frozenset('X'))
+        f = function.swap_spaces(X.f_index, 'X', 'Y')
+        self.assertEqual(f.spaces, frozenset('Y'))
+        with self.assertRaisesRegex(KeyError, 'no such space'):
+            X.sample('gauss', 0).eval(f)
+        self.assertEqual(X.sample('gauss', 0).rename_spaces({'X': 'Y'}).eval(f).tolist(), [0, 1])
+
+    def test_same(self):
+        X, x = mesh.line(2, space='X')
+        f = function.swap_spaces(X.f_index, 'X', 'X')
+        self.assertEqual(f.spaces, frozenset('X'))
+        self.assertEqual(X.sample('gauss', 0).eval(f).tolist(), [0, 1])

tests/test_function.py

@@ -63,16 +63,14 @@ def test_get_lower_args(self):
             self.assertEqual(tuple(n.__index__() for n in actual_shape(dict(ielem=ielem))), desired_shape)
             offset = 0
             for space, desired_chain, desired_point in zip(self.desired_spaces, desired_chains, desired_points):
-                (chain, *_), index = args.transform_chains[space]
-                self.assertEqual(chain[evaluable.eval_once(index, arguments=dict(ielem=ielem)).__index__()], desired_chain)
+                arg = args[space]
+                self.assertEqual(arg.transforms[evaluable.eval_once(arg.index, arguments=dict(ielem=ielem)).__index__()], desired_chain)
                 desired_coords = desired_point.coords
                 desired_coords = numpy.lib.stride_tricks.as_strided(desired_coords, shape=(*desired_shape, desired_point.ndims,), strides=(0,)*offset+desired_coords.strides[:-1]+(0,)*(len(args.points_shape)-offset-desired_coords.ndim+1)+desired_coords.strides[-1:])
-                actual_coords = evaluable.eval_once(args.coordinates[space], arguments=dict(ielem=ielem))
+                actual_coords = evaluable.eval_once(arg.coordinates, arguments=dict(ielem=ielem))
                 self.assertEqual(actual_coords.shape, desired_coords.shape)
                 self.assertAllAlmostEqual(actual_coords, desired_coords)
                 offset += desired_point.coords.ndim - 1
-        with self.assertRaisesRegex(ValueError, '^Nested'):
-            args | self.sample.get_lower_args(evaluable.InRange(evaluable.Argument('ielem2', (), int), evaluable.constant(self.desired_nelems)))
 
     @property
     def _desired_element_tri(self):
@@ -128,8 +126,8 @@ def test_take_elements_single(self):
             self.assertEqual(take.ndims, self.desired_ndims)
             args = take.get_lower_args(evaluable.InRange(evaluable.Argument('ielem', (), int), evaluable.constant(1)))
             for space, desired_chain in zip(self.desired_spaces, self.desired_transform_chains[ielem]):
-                (chain, *_), index = args.transform_chains[space]
-                self.assertEqual(chain[evaluable.eval_once(index, arguments=dict(ielem=0)).__index__()], desired_chain)
+                arg = args[space]
+                self.assertEqual(arg.transforms[evaluable.eval_once(arg.index, arguments=dict(ielem=0)).__index__()], desired_chain)
 
     def test_take_elements_empty(self):
         take = self.sample.take_elements(numpy.array([], int))
@@ -148,6 +146,11 @@ def test_asfunction(self):
         func = self.sample.asfunction(numpy.arange(self.sample.npoints))
         self.assertEqual(self.sample.bind(func).eval().tolist(), numpy.arange(self.desired_npoints).tolist())
 
+    def test_rename_spaces(self):
+        assert not any(space.endswith('!') for space in self.desired_spaces)
+        renamed = self.sample.rename_spaces({space: space + '!' for space in self.desired_spaces})
+        self.assertEqual(renamed.spaces, tuple(space + '!' for space in self.desired_spaces))
+
 
 class Empty(TestCase, Common):
 
@@ -280,8 +283,7 @@ def test_integrate(self):
         self.assertAlmostEqual(self.stitched.integrate(function.J(self.geomX)), 5/9)  # NOTE: != norm(slope)
 
     def test_nested(self):
-        with self.assertRaisesRegex(ValueError, 'Nested integrals or samples in the same space: X.*, Y.'):
-            self.stitched.integral(self.stitched.integral(1)).eval()
+        self.stitched.integral(self.stitched.integral(1)).eval()
         topoZ, geomZ = mesh.line(2, space='Z')
         inner = self.stitched.integral((geomZ - self.geomX) * function.J(self.geomY))
         outer = topoZ.integral(inner * function.J(geomZ), degree=2)
@@ -295,6 +297,10 @@ def test_triplet(self):
         self.assertAllAlmostEqual(geomX, geomY[:, numpy.newaxis] * self.slope)
         self.assertAllAlmostEqual(geomY, geomZ / 3)
 
+    def test_rename_spaces(self):
+        renamed = self.stitched.rename_spaces({'X': 'B', 'X0': 'B', 'X1': 'C', 'Y': 'A'})
+        self.assertEqual(renamed.spaces, ('A', 'B') if 'X' in self.sampleX.spaces else ('A', 'B', 'C'))
+
 
 Zip(etype='square')
 Zip(etype='triangle')
@@ -438,6 +444,10 @@ def test_basis(self):
 
 class Special(TestCase):
 
+    def test_init_repeated_spaces(self):
+        with self.assertRaisesRegex(ValueError, 'space a is repeated.$'):
+            Sample(('a', 'b', 'a'), 1, 1, 1)
+
     def test_add_different_spaces(self):
         class Dummy(Sample):
             pass
@@ -576,3 +586,12 @@ def test_empty(self):
         empty = function.zeros(shape, float)
         array = empty.eval(legacy=False)
         self.assertAllEqual(array, numpy.zeros((2, 3)))
+
+    def test_reuse_space(self):
+        X0, x0 = mesh.line(2, space='X')
+        X1, x1 = mesh.line(3, space='X')
+        self.assertAllAlmostEqual(
+            function.eval(X0.integral(X1.integral(function.J(x1), degree=3) * function.J(x0), degree=2)),
+            6,
+            places=14,
+        )