support locate on 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. The Sample.zip method received a new optional ndims argument
to facilitate zipping of samples with varying dimensions.

nutils/sample.py

@@ -335,7 +335,7 @@ def take_elements(self, __indices: numpy.ndarray) -> 'Sample':
         else:
             return Sample.empty(self.spaces, self.ndims)
 
-    def zip(*samples: 'Sample') -> 'Sample':
+    def zip(*samples: 'Sample', ndims: Optional[int] = None) -> 'Sample':
         '''
         Join multiple samples, with identical point count but differing spaces, into
         a new sample with the same point count and the total of spaces. The result is
@@ -371,7 +371,12 @@ def zip(*samples: 'Sample') -> 'Sample':
         zipped : :class:`Sample`
         '''
 
-        return _Zip(*samples)
+        if ndims is None:
+            ndims = samples[0].ndims
+            if any(sample.ndims != ndims for sample in samples[1:]):
+                raise ValueError('zipped samples have varying dimensions; if this is correct please specify the resulting dimension via ndims=...')
+
+        return _Zip(samples, ndims)
 
 
 class _TransformChainsSample(Sample):
@@ -797,7 +802,7 @@ def basis(self, interpolation: str = 'none') -> Sample:
 
 class _Zip(Sample):
 
-    def __init__(self, *samples):
+    def __init__(self, samples, ndims):
         npoints = samples[0].npoints
         spaces = util.sum(sample.spaces for sample in samples)
         if not all(sample.npoints == npoints for sample in samples):
@@ -824,7 +829,7 @@ def __init__(self, *samples):
         self._ielems = tuple(types.arraydata(array) for array in numpy.unravel_index(flat_ielems, nelems))
         self._ilocals = tuple(types.arraydata(numpy.take(ilocal, self._indices, axis=0)) for ilocal in ilocals)
 
-        super().__init__(spaces=spaces, ndims=samples[0].ndims, nelems=self._sizes.shape[0], npoints=npoints)
+        super().__init__(spaces=spaces, ndims=ndims, nelems=self._sizes.shape[0], npoints=npoints)
 
     def getindex(self, ielem):
         return numpy.asarray(self._indices)[slice(*numpy.asarray(self._offsets)[ielem:ielem+2])]

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,18 @@ 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:])
+        return Sample.zip(sample1, sample2, ndims=self.ndims)
+
 
 class _Take(_TensorialTopology):
 
@@ -1534,76 +1621,6 @@ 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 _sample(self, ielems, coords, weights=None):
         index = numpy.argsort(ielems, kind='stable')
         sorted_ielems = ielems[index]
@@ -1727,6 +1744,9 @@ def basis_bernstein(self, degree):
 
     basis_std = basis_bernstein
 
+    def _lower_args(self, ielem, point):
+        return function.LowerArgs.for_space(self.space, (self.transforms, self.opposites), ielem, point)
+
 
 class LocateError(Exception):
     pass

tests/test_sample.py

@@ -268,7 +268,7 @@ def setUp(self):
         self.sampleY = topoY.sample('uniform', 3)
         self.slope = numpy.array([1, .5])  # geomX == geomY * slope
         self.sampleX = topoX.locate(self.geomX, self.sampleY.eval(self.geomY * self.slope), tol=1e-10)
-        self.stitched = self.sampleY.zip(self.sampleX)
+        self.stitched = self.sampleY.zip(self.sampleX, ndims=topoY.ndims)
 
     def test_eval(self):
         geomY, geomX = self.stitched.eval([self.geomY, self.geomX])
@@ -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))
@@ -290,7 +290,7 @@ def test_nested(self):
     def test_triplet(self):
         topoZ, geomZ = mesh.line(3, space='Z')
         sampleZ = topoZ.sample('uniform', 5)
-        triplet = Sample.zip(self.sampleY, self.sampleX, sampleZ)
+        triplet = Sample.zip(self.sampleY, self.sampleX, sampleZ, ndims=self.sampleY.ndims)
         geomX, geomY, geomZ = triplet.eval([self.geomX, self.geomY, geomZ])
         self.assertAllAlmostEqual(geomX, geomY[:, numpy.newaxis] * self.slope)
         self.assertAllAlmostEqual(geomY, geomZ / 3)

tests/test_topology.py

@@ -973,6 +973,8 @@ def test_missing_argument(self):
 
     @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))