speed up points traversal

Author: mikedh

Dockerfile

@@ -24,7 +24,7 @@ RUN bash /tmp/conda.bash
 
 # install python requirements
 COPY . /tmp/trimesh
-RUN /home/user/conda/bin/pip install /tmp/trimesh[all]
+RUN /home/user/conda/bin/pip install /tmp/trimesh[all] pytest
 
 # add user python to path 
 ENV PATH="/home/user/conda/bin:$PATH"
@@ -34,3 +34,6 @@ ENV XVFB_WHD="1920x1080x24"\
     DISPLAY=":99" \
     LIBGL_ALWAYS_SOFTWARE="1" \
     GALLIUM_DRIVER="llvmpipe"
+
+# make sure build fails if tests are failing
+RUN pytest /tmp/trimesh/tests

tests/test_points.py

@@ -129,6 +129,9 @@ def test_tsp(self):
                     assert len(idx) == len(points)
                     assert len(dist) == len(points) - 1
 
+                    # shouldn't be any negative indexes
+                    assert (idx >= 0).all()
+
                     # make sure distances returned are correct
                     dist_check = g.np.linalg.norm(g.np.diff(points[idx], axis=0), axis=1)
                     assert g.np.allclose(dist_check, dist)

trimesh/points.py

@@ -227,11 +227,15 @@ def k_means(points, k, **kwargs):
 def tsp(points, start=0):
     """
     Find an ordering of points where each is visited and
-    the next point is the closest in euclidean distance.
+    the next point is the closest in euclidean distance,
+    and if there are multiple points with equal distance
+    go to an arbitrary one.
 
     Assumes every point is visitable from every other point,
     i.e. the travelling salesman problem on a fully connected
-    graph.
+    graph. It is not a MINIMUM traversal; rather it is a
+    "not totally goofy traversal, quickly." On random points
+    this traversal is often ~20x shorter than random ordering.
 
     Parameters
     ---------------
@@ -249,32 +253,41 @@ def tsp(points, start=0):
     """
     # points should be float
     points = np.asanyarray(points, dtype=np.float64)
+
+    if len(points.shape) != 2:
+        raise ValueError('points must be (n, dimension)!')
+
     # start should be an index
     start = int(start)
 
     # a mask of unvisited points by index
     unvisited = np.ones(len(points), dtype=np.bool)
     unvisited[start] = False
+
     # traversal of points by index
-    traversal = [start]
+    traversal = np.zeros(len(points), dtype=np.int64) - 1
+    traversal[0] = start
     # list of distances
-    distances = []
+    distances = np.zeros(len(points) - 1, dtype=np.float64)
     # a mask of indexes in order
     index_mask = np.arange(len(points), dtype=np.int64)
 
-    # bound our traversal in case it's dumb
-    for i in range(len(points) + 2):
-
-        # we should always exit via this break
-        if not unvisited.any():
-            break
+    # in the loop we want to call distances.sum(axis=1)
+    # a lot and it's actually kind of slow for "reasons"
+    # but dot products with ones are equivilant and roughly
+    # 2x faster
+    sum_ones = np.ones(points.shape[1])
 
-        # which point are we currently at
+    # loop through all points
+    for i in range(len(points) - 1):
+        # which point are we currently on
         current = points[traversal[i]]
 
         # do NlogN distance query
-        # use sum instead of np.linalg.norm as it is slightly faster
-        dist = ((points[unvisited] - current) ** 2).sum(axis=1) ** 0.5
+        # use dot instead of .sum(axis=1) or np.linalg.norm
+        # as it is much faster, also don't square root
+        dist = np.dot((points[unvisited] - current) ** 2,
+                      sum_ones)
 
         # minimum distance index
         min_index = dist.argmin()
@@ -283,13 +296,12 @@ def tsp(points, start=0):
         # update the mask
         unvisited[successor] = False
         # append the index to the traversal
-        traversal.append(successor)
+        traversal[i + 1] = successor
         # append the distance
-        distances.append(dist[min_index])
+        distances[i] = dist[min_index]
 
-    # make sure results are numpy arrays of correct dtype
-    traversal = np.array(traversal, dtype=np.int64)
-    distances = np.array(distances, dtype=np.float64)
+    # we were comparing squared distance so root result
+    distances **= 0.5
 
     return traversal, distances
 

trimesh/util.py

@@ -75,7 +75,9 @@ def unitize(vectors,
     if len(vectors.shape) == 2:
         # for (m, d) arrays take the per- row unit vector
         # using sqrt and avoiding exponents is slightly faster
-        norm = np.sqrt((vectors * vectors).sum(axis=1))
+        # also dot with ones is faser than .sum(axis=1)
+        norm = np.sqrt(np.dot(vectors * vectors,
+                              np.ones(vectors.shape[1])))
         # non-zero norms
         valid = norm > threshold
         # in-place reciprocal of nonzero norms

trimesh/version.py

@@ -1 +1 @@
-__version__ = '2.34.3'
+__version__ = '2.34.4'