Optimized :meth:.Axes.coords_to_point

manim/mobject/graphing/coordinate_systems.py

@@ -1995,40 +1995,47 @@ def construct(self):
 
                     self.add(plane, dot_scene, ax, dot_axes, lines)
         """
+        coords = np.asarray(coords)
         origin = self.x_axis.number_to_point(
             self._origin_shift([self.x_axis.x_min, self.x_axis.x_max]),
         )
 
-        coords = np.asarray(coords)
-
-        # if called like coords_to_point(1, 2, 3), then coords is a 1x3 array
-        transposed = False
-        if coords.ndim == 1:
-            # original implementation of coords_to_point for performance in the legacy case
-            result = np.array(origin)
-            for axis, number in zip(self.get_axes(), coords):
-                result += axis.number_to_point(number) - origin
-            return result
-        # if called like coords_to_point([1, 2, 3],[4, 5, 6]), then it shall be used as [1,4], [2,5], [3,6] and return the points as ([x_0,x_1],[y_0,y_1],[z_0,z_1])
-        elif coords.ndim == 2:
-            coords = coords.T
-            transposed = True
-        # if called like coords_to_point(np.array([[1, 2, 3],[4,5,6]])), reduce dimension by 1
-        elif coords.ndim == 3:
-            coords = np.squeeze(coords)
-        # else the coords is a Nx1, Nx2, Nx3 array so we do not need to modify the array
-
-        points = origin + np.sum(
-            [
-                axis.number_to_point(number) - origin
-                for number, axis in zip(coords.T, self.get_axes())
-            ],
-            axis=0,
-        )
-        # if called with single coord, then return a point instead of a list of points
-        if transposed:
-            return points.T
-        return points
+        # Is coords in the format ([[x1 y1 z1] [x2 y2 z2] ...])? (True)
+        # Or is coords in the format (x, y, z) or ([x1 x2 ...], [y1 y2 ...], [z1 z2 ...])? (False)
+        # The latter is preferred.
+        are_coordinates_transposed = False
+
+        # If coords is in the format ([[x1 y1 z1] [x2 y2 z2] ...]):
+        if coords.ndim == 3:
+            # Extract from original tuple: now coords looks like [[x y z]] or [[x1 y1 z1] [x2 y2 z2] ...].
+            coords = coords[0]
+            # If there's a single coord (coords = [[x y z]]), extract it so that
+            # coords = [x y z] and coords_to_point returns a single point.
+            if coords.shape[0] == 1:
+                coords = coords[0]
+            # Else, if coords looks more like [[x1 y1 z1] [x2 y2 z2] ...], transform them (by
+            # transposing) into the format [[x1 x2 ...] [y1 y2 ...] [z1 z2 ...]] for later processing.
+            else:
+                coords = coords.T
+                are_coordinates_transposed = True
+        # Otherwise, coords already looked like (x, y, z) or ([x1 x2 ...], [y1 y2 ...], [z1 z2 ...]),
+        # so no further processing is needed.
+
+        # Now coords should either look like [x y z] or [[x1 x2 ...] [y1 y2 ...] [z1 z2 ...]],
+        # so it can be iterated directly. Each element is either a float representing a single
+        # coordinate, or a float ndarray of coordinates corresponding to a single axis.
+        # Although "points" and "nums" are in plural, there might be a single point or number.
+        points = self.x_axis.number_to_point(coords[0])
+        other_axes = self.axes.submobjects[1:]
+        for axis, nums in zip(other_axes, coords[1:]):
+            points += axis.number_to_point(nums) - origin
+
+        # Return points as is, except if coords originally looked like
+        # ([x1 x2 ...], [y1 y2 ...], [z1 z2 ...]), which is determined by the conditions below. In
+        # that case, the current implementation requires that the results have to be transposed.
+        if are_coordinates_transposed or points.ndim == 1:
+            return points
+        return points.T
 
     def point_to_coords(self, point: Sequence[float]) -> np.ndarray:
         """Accepts a point from the scene and returns its coordinates with respect to the axes.