support locate for tensorial topologies

This patch add support for locate on tensorial topologies, by moving the locate
implementation from TransformChainsTopology to the Topology base class. This is
mode possible by adding a new method _lower_args, as well as an implementation
of _Mul._sample.

nutils/topology.py

@@ -785,6 +785,78 @@ def locate(self, geom, coords, *, tol=0, eps=0, maxiter=0, arguments=None, weigh
         located : :class:`nutils.sample.Sample`
         '''
 
+        if max(tol, eps) <= 0:
+            raise ValueError('locate requires either tol or eps to be strictly positive')
+        coords = numpy.asarray(coords, dtype=float)
+        if geom.ndim == 0:
+            geom = geom[_]
+            coords = coords[..., _]
+        if not geom.shape == coords.shape[1:] == (self.ndims,):
+            raise ValueError('invalid geometry or point shape for {}D topology'.format(self.ndims))
+        if skip_missing and weights is not None:
+            raise ValueError('weights and skip_missing are mutually exclusive')
+        arguments = dict(arguments or ())
+        centroids = self.sample('_centroid', None).eval(geom, **arguments)
+        assert len(centroids) == len(self)
+        ielems = parallel.shempty(len(coords), dtype=int)
+        points = parallel.shempty((len(coords), len(geom)), dtype=float)
+        _ielem = evaluable.Argument('_locate_ielem', shape=(), dtype=int)
+        _point = evaluable.Argument('_locate_point', shape=(evaluable.constant(self.ndims),))
+        egeom = geom.lower(self._lower_args(_ielem, _point))
+        xJ = evaluable.Tuple((egeom, evaluable.derivative(egeom, _point))).simplified
+        with parallel.ctxrange('locating', len(coords)) as ipoints:
+            for ipoint in ipoints:
+                xt = coords[ipoint]  # target
+                dist = numpy.linalg.norm(centroids - xt, axis=1)
+                for ielem in numpy.argsort(dist) if maxdist is None \
+                        else sorted((dist < maxdist).nonzero()[0], key=dist.__getitem__):
+                    ref = self.references[ielem]
+                    arguments['_locate_ielem'] = ielem
+                    arguments['_locate_point'] = p = numpy.array(ref.centroid)
+                    ex = ep = numpy.inf
+                    iiter = 0
+                    while ex > tol and ep > eps:  # newton loop
+                        if iiter > maxiter > 0:
+                            break  # maximum number of iterations reached
+                        iiter += 1
+                        xp, Jp = xJ.eval(**arguments)
+                        dx = xt - xp
+                        ex0 = ex
+                        ex = numpy.linalg.norm(dx)
+                        if ex >= ex0:
+                            break  # newton is diverging
+                        try:
+                            dp = numpy.linalg.solve(Jp, dx)
+                        except numpy.linalg.LinAlgError:
+                            break  # jacobian is singular
+                        ep = numpy.linalg.norm(dp)
+                        p += dp  # NOTE: modifies arguments['_locate_point'] in place
+                    else:
+                        if ref.inside(p, max(eps, ep)):
+                            ielems[ipoint] = ielem
+                            points[ipoint] = p
+                            break
+                else:
+                    ielems[ipoint] = -1 # mark point as missing
+                    if not skip_missing:
+                        # rather than raising LocateError here, which
+                        # parallel.fork will reraise as a general Exception if
+                        # ipoint was assigned to a child process, we fast
+                        # forward through the remaining points to abandon the
+                        # loop and subsequently raise from the main process.
+                        for ipoint in ipoints:
+                            pass
+        if -1 not in ielems: # all points are found
+            return self._sample(ielems, points, weights)
+        elif skip_missing: # not all points are found and that's ok, we just leave those out
+            return self._sample(ielems[ielems != -1], points[ielems != -1])
+        else: # not all points are found and that's an error
+            raise LocateError(f'failed to locate point: {coords[ielems==-1][0]}')
+
+    def _lower_args(self, ielem, point):
+        raise NotImplementedError
+
+    def _sample(self, ielems, coords, weights=None):
         raise NotImplementedError
 
     @cached_property
@@ -985,8 +1057,11 @@ def withgroups(self, vgroups: Mapping[str, Union[str, Topology]] = {}, bgroups:
         def indicator(self, subtopo: Union[str, Topology]) -> Topology:
             raise SkipTest('`{}` does not implement `Topology.indicator`'.format(type(self).__qualname__))
 
-        def locate(self, geom, coords, *, tol=0, eps=0, maxiter=0, arguments=None, weights=None, maxdist=None, ischeme=None, scale=None, skip_missing=False) -> Sample:
-            raise SkipTest('`{}` does not implement `Topology.locate`'.format(type(self).__qualname__))
+        def _lower_args(self, ielem, point):
+            raise SkipTest('`{}` does not implement `Topology._lower_args`'.format(type(self).__qualname__))
+
+        def _sample(self, ielems, coords, weights=None):
+            raise SkipTest('`{}` does not implement `Topology._sample`'.format(type(self).__qualname__))
 
 else:
     _TensorialTopology = Topology
@@ -1241,6 +1316,23 @@ def basis(self, name: str, degree: Union[int, Sequence[int]], **kwargs) -> funct
     def sample(self, ischeme: str, degree: int) -> Sample:
         return self.topo1.sample(ischeme, degree) * self.topo2.sample(ischeme, degree)
 
+    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
+
+    def _sample(self, ielems, coords, weights=None):
+        ielems1, ielems2 = divmod(ielems, len(self.topo2))
+        sample1 = self.topo1._sample(ielems1, coords[...,:self.topo1.ndims], weights)
+        sample2 = self.topo2._sample(ielems2, coords[...,self.topo1.ndims:])
+        # NOTE: zip creates a sample with ndims set equal to that of the first
+        # argument, assuming (without verifying) that all samples are of the
+        # same dimension. Here this assumption is incorrect, as ndims should be
+        # the sum of the two samples, but since zipped samples do not support
+        # tri and hull the mistake is without consequences.
+        return Sample.zip(sample1, sample2)
+
 
 class _Take(_TensorialTopology):
 
@@ -1534,75 +1626,8 @@ def indicator(self, subtopo):
         values[numpy.fromiter(map(self.transforms.index, subtopo.transforms), dtype=int)] = 1
         return function.get(values, 0, self.f_index)
 
-    @log.withcontext
-    def locate(self, geom, coords, *, tol=0, eps=0, maxiter=0, arguments=None, weights=None, maxdist=None, skip_missing=False):
-        if max(tol, eps) <= 0:
-            raise ValueError('locate requires either tol or eps to be strictly positive')
-        coords = numpy.asarray(coords, dtype=float)
-        if geom.ndim == 0:
-            geom = geom[_]
-            coords = coords[..., _]
-        if not geom.shape == coords.shape[1:] == (self.ndims,):
-            raise ValueError('invalid geometry or point shape for {}D topology'.format(self.ndims))
-        if skip_missing and weights is not None:
-            raise ValueError('weights and skip_missing are mutually exclusive')
-        arguments = dict(arguments or ())
-        centroids = self.sample('_centroid', None).eval(geom, **arguments)
-        assert len(centroids) == len(self)
-        ielems = parallel.shempty(len(coords), dtype=int)
-        points = parallel.shempty((len(coords), len(geom)), dtype=float)
-        _ielem = evaluable.Argument('_locate_ielem', shape=(), dtype=int)
-        _point = evaluable.Argument('_locate_point', shape=(evaluable.constant(self.ndims),))
-        egeom = geom.lower(function.LowerArgs.for_space(self.space, (self.transforms, self.opposites), _ielem, _point))
-        xJ = evaluable.Tuple((egeom, evaluable.derivative(egeom, _point))).simplified
-        with parallel.ctxrange('locating', len(coords)) as ipoints:
-            for ipoint in ipoints:
-                xt = coords[ipoint]  # target
-                dist = numpy.linalg.norm(centroids - xt, axis=1)
-                for ielem in numpy.argsort(dist) if maxdist is None \
-                        else sorted((dist < maxdist).nonzero()[0], key=dist.__getitem__):
-                    ref = self.references[ielem]
-                    arguments['_locate_ielem'] = ielem
-                    arguments['_locate_point'] = p = numpy.array(ref.centroid)
-                    ex = ep = numpy.inf
-                    iiter = 0
-                    while ex > tol and ep > eps:  # newton loop
-                        if iiter > maxiter > 0:
-                            break  # maximum number of iterations reached
-                        iiter += 1
-                        xp, Jp = xJ.eval(**arguments)
-                        dx = xt - xp
-                        ex0 = ex
-                        ex = numpy.linalg.norm(dx)
-                        if ex >= ex0:
-                            break  # newton is diverging
-                        try:
-                            dp = numpy.linalg.solve(Jp, dx)
-                        except numpy.linalg.LinAlgError:
-                            break  # jacobian is singular
-                        ep = numpy.linalg.norm(dp)
-                        p += dp  # NOTE: modifies arguments['_locate_point'] in place
-                    else:
-                        if ref.inside(p, max(eps, ep)):
-                            ielems[ipoint] = ielem
-                            points[ipoint] = p
-                            break
-                else:
-                    ielems[ipoint] = -1 # mark point as missing
-                    if not skip_missing:
-                        # rather than raising LocateError here, which
-                        # parallel.fork will reraise as a general Exception if
-                        # ipoint was assigned to a child process, we fast
-                        # forward through the remaining points to abandon the
-                        # loop and subsequently raise from the main process.
-                        for ipoint in ipoints:
-                            pass
-        if -1 not in ielems: # all points are found
-            return self._sample(ielems, points, weights)
-        elif skip_missing: # not all points are found and that's ok, we just leave those out
-            return self._sample(ielems[ielems != -1], points[ielems != -1])
-        else: # not all points are found and that's an error
-            raise LocateError(f'failed to locate point: {coords[ielems==-1][0]}')
+    def _lower_args(self, ielem, point):
+        return function.LowerArgs.for_space(self.space, (self.transforms, self.opposites), ielem, point)
 
     def _sample(self, ielems, coords, weights=None):
         index = numpy.argsort(ielems, kind='stable')

tests/test_sample.py

@@ -280,7 +280,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.'):
+        with self.assertRaisesRegex(ValueError, 'Nested integrals or samples in the same space: X.*, Y.'):
             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))

tests/test_topology.py

@@ -966,13 +966,22 @@ def test_boundary_scalar(self):
         located = sample.eval(self.geom[1], scale=.123)
         self.assertAllAlmostEqual(located, target)
 
+    def test_integrate(self):
+        target = numpy.array([(.2, .3), (.1, .9), (0, 1), (.1, .3)])
+        weights = numpy.array([.1, .2, .3, .4])
+        sample = self.domain.locate(self.geom, target, eps=1e-15, tol=1e-12, arguments=dict(scale=.123), weights=weights)
+        integrated = sample.integrate(self.geom, scale=.123)
+        self.assertAllAlmostEqual(integrated, weights @ target)
+
     def test_missing_argument(self):
         target = numpy.array([(.2, .3)])
         with self.assertRaises(Exception):
             self.domain.locate(self.geom, target, eps=1e-15, tol=1e-12)
 
     @parametrize.enable_if(lambda etype, mode, **kwargs: etype == 'square' and mode != 'nonlinear')
     def test_detect_linear(self):
+        if os.environ.get('NUTILS_TENSORIAL'):
+            self.skipTest('linear detection not supported yet in tensorial mode')
         target = numpy.array([(.2, .3)])
         with self.assertLogs('nutils', level='DEBUG') as cm:
             self.domain.locate(self.geom, target, eps=1e-15, tol=1e-12, arguments=dict(scale=.123))