Optimized :meth:.Axes.coords_to_point

manim/mobject/graphing/coordinate_systems.py

@@ -1995,40 +1995,42 @@ 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 every component in coords in the format [xi, yi, zi]?
+        # i.e. 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_components_xyz = False
+
+        # If coords is in the format ([[x1 y1 z1] [x2 y2 z2] ...]):
+        if coords.ndim == 3:
+            # Extract from original tuple: now the format is [[x1 y1 z1] [x2 y2 z2]]
+            coords = coords[0]
+
+            # If there's a single coord, extract it so that coords_to_point returns a single point
+            if coords.shape[0] == 1:
+                coords = coords[0]  # In this case, now coords = [x1 y1 z1]
+            else:
+                are_components_xyz = True
+                # Transform coords into the format [[x1 x2 ...] [y1 y2 ...] [z1 z2 ...]]
+                # for later processing.
+                coords = coords.T
+
+        # Now coords should be in the format [x y z], where each component is either a float or an ndarray
+        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
+
+        if are_components_xyz:
+            return points
+        return points.T  # Has no effect on a 1D array representing a single point
 
     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.