support locate on tensorial topologies

nutils/topology.py

@@ -785,6 +785,78 @@
         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 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 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 @@
         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))