Factor repeated code into proj3d functions

scottshambaugh · scottshambaugh · commit a24160f2011a · 2026-03-04T16:04:05.000-07:00
diff --git a/lib/mpl_toolkits/mplot3d/art3d.py b/lib/mpl_toolkits/mplot3d/art3d.py
@@ -72,33 +72,6 @@ def get_dir_vector(zdir):
         raise ValueError("'x', 'y', 'z', None or vector of length 3 expected")
 
 
-def _apply_scale_transforms(xs, ys, zs, axes):
-    """
-    Apply axis scale transforms to 3D coordinates.
-
-    Transforms data coordinates to transformed coordinates (applying log,
-    symlog, etc.) for 3D projection. Preserves masked arrays.
-    """
-    def transform_coord(coord, axis):
-        coord = np.asanyarray(coord)
-        data = np.ma.getdata(coord).ravel()
-        return axis.get_transform().transform(data).reshape(coord.shape)
-
-    xs_scaled = transform_coord(xs, axes.xaxis)
-    ys_scaled = transform_coord(ys, axes.yaxis)
-    zs_scaled = transform_coord(zs, axes.zaxis)
-
-    # Preserve combined mask from any masked input
-    masks = [np.ma.getmask(a) for a in [xs, ys, zs]]
-    if any(m is not np.ma.nomask for m in masks):
-        combined = np.ma.mask_or(np.ma.mask_or(masks[0], masks[1]), masks[2])
-        xs_scaled = np.ma.array(xs_scaled, mask=combined)
-        ys_scaled = np.ma.array(ys_scaled, mask=combined)
-        zs_scaled = np.ma.array(zs_scaled, mask=combined)
-
-    return xs_scaled, ys_scaled, zs_scaled
-
-
 def _viewlim_mask(xs, ys, zs, axes):
     """
     Return the mask of the points outside the axes view limits.
@@ -225,16 +198,9 @@ def draw(self, renderer):
         else:
             pos3d = np.array([self._x, self._y, self._z], dtype=float)
 
-        # Apply scale transforms before projection
-        pos3d_scaled = np.array(_apply_scale_transforms(
-            pos3d[0], pos3d[1], pos3d[2], self.axes))
-        # Also scale the direction vector endpoint
         dir_end = pos3d + self._dir_vec
-        dir_end_scaled = np.array(_apply_scale_transforms(
-            dir_end[0], dir_end[1], dir_end[2], self.axes))
-
-        proj = proj3d._proj_trans_points(
-            [pos3d_scaled, dir_end_scaled], self.axes.M)
+        proj = proj3d._proj_trans_points_scaled(
+            [pos3d, dir_end], self.axes)
         dx = proj[0][1] - proj[0][0]
         dy = proj[1][1] - proj[1][0]
         angle = math.degrees(math.atan2(dy, dx))
@@ -370,11 +336,8 @@ def draw(self, renderer):
                                            dtype=float, mask=mask).filled(np.nan)
         else:
             xs3d, ys3d, zs3d = self._verts3d
-        # Apply scale transforms before projection
-        xs3d, ys3d, zs3d = _apply_scale_transforms(xs3d, ys3d, zs3d, self.axes)
-        xs, ys, zs, tis = proj3d._proj_transform_clip(xs3d, ys3d, zs3d,
-                                                      self.axes.M,
-                                                      self.axes._focal_length)
+        xs, ys, zs, tis = proj3d._proj_transform_clip_scaled(
+            xs3d, ys3d, zs3d, self.axes)
         self.set_data(xs, ys)
         super().draw(renderer)
         self.stale = False
@@ -465,13 +428,8 @@ def do_3d_projection(self):
             vs_list = [np.ma.array(vs, mask=np.broadcast_to(
                        _viewlim_mask(*vs.T, self.axes), vs.shape))
                        for vs in vs_list]
-        # Apply scale transforms before projection
-        xyzs_list = []
-        for vs in vs_list:
-            xs_scaled, ys_scaled, zs_scaled = _apply_scale_transforms(
-                vs[:, 0], vs[:, 1], vs[:, 2], self.axes)
-            xyzs_list.append(proj3d.proj_transform(
-                xs_scaled, ys_scaled, zs_scaled, self.axes.M))
+        xyzs_list = [proj3d._proj_transform_scaled(
+            vs[:, 0], vs[:, 1], vs[:, 2], self.axes) for vs in vs_list]
         self._paths = [mpath.Path(np.ma.column_stack([xs, ys]), cs)
                        for (xs, ys, _), (_, cs) in zip(xyzs_list, self._3dverts_codes)]
         zs = np.concatenate([zs for _, _, zs in xyzs_list])
@@ -561,13 +519,9 @@ def do_3d_projection(self):
                                                (*viewlim_mask.shape, 3))
                 mask = mask | viewlim_mask
 
-        # Apply scale transforms before projection
-        xs_scaled, ys_scaled, zs_scaled = _apply_scale_transforms(
-            segments[..., 0], segments[..., 1], segments[..., 2], self.axes)
-        segments_scaled = np.stack([xs_scaled, ys_scaled, zs_scaled], axis=-1)
-
-        xyzs = np.ma.array(proj3d._proj_transform_vectors(segments_scaled, self.axes.M),
-                           mask=mask)
+        xyzs = np.ma.array(
+            proj3d._proj_transform_vectors_scaled(segments, self.axes),
+            mask=mask)
         segments_2d = xyzs[..., 0:2]
         LineCollection.set_segments(self, segments_2d)
 
@@ -650,11 +604,8 @@ def do_3d_projection(self):
                                      dtype=float, mask=mask).filled(np.nan)
         else:
             xs, ys, zs = zip(*s)
-        # Apply scale transforms before projection
-        xs, ys, zs = _apply_scale_transforms(xs, ys, zs, self.axes)
-        vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
-                                                         self.axes.M,
-                                                         self.axes._focal_length)
+        vxs, vys, vzs, vis = proj3d._proj_transform_clip_scaled(
+            xs, ys, zs, self.axes)
         self._path2d = mpath.Path(np.ma.column_stack([vxs, vys]))
         return min(vzs)
 
@@ -714,11 +665,8 @@ def do_3d_projection(self):
                                      dtype=float, mask=mask).filled(np.nan)
         else:
             xs, ys, zs = zip(*s)
-        # Apply scale transforms before projection
-        xs, ys, zs = _apply_scale_transforms(xs, ys, zs, self.axes)
-        vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
-                                                         self.axes.M,
-                                                         self.axes._focal_length)
+        vxs, vys, vzs, vis = proj3d._proj_transform_clip_scaled(
+            xs, ys, zs, self.axes)
         self._path2d = mpath.Path(np.ma.column_stack([vxs, vys]), self._code3d)
         return min(vzs)
 
@@ -861,11 +809,8 @@ def do_3d_projection(self):
             xs, ys, zs = np.ma.array(self._offsets3d, mask=mask)
         else:
             xs, ys, zs = self._offsets3d
-        # Apply scale transforms before projection
-        xs, ys, zs = _apply_scale_transforms(xs, ys, zs, self.axes)
-        vxs, vys, vzs, vis = proj3d._proj_transform_clip(xs, ys, zs,
-                                                         self.axes.M,
-                                                         self.axes._focal_length)
+        vxs, vys, vzs, vis = proj3d._proj_transform_clip_scaled(
+            xs, ys, zs, self.axes)
         self._vzs = vzs
         if np.ma.isMA(vxs):
             super().set_offsets(np.ma.column_stack([vxs, vys]))
@@ -1081,11 +1026,8 @@ def do_3d_projection(self):
             xyzs = np.ma.array(self._offsets3d, mask=mask)
         else:
             xyzs = self._offsets3d
-        # Apply scale transforms before projection
-        xyzs_scaled = _apply_scale_transforms(*xyzs, self.axes)
-        vxs, vys, vzs, vis = proj3d._proj_transform_clip(*xyzs_scaled,
-                                                         self.axes.M,
-                                                         self.axes._focal_length)
+        vxs, vys, vzs, vis = proj3d._proj_transform_clip_scaled(
+            *xyzs, self.axes)
         self._data_scale = _get_data_scale(vxs, vys, vzs)
         # Sort the points based on z coordinates
         # Performance optimization: Create a sorted index array and reorder
@@ -1416,16 +1358,11 @@ def do_3d_projection(self):
         num_faces = len(self._faces)
         mask = self._invalid_vertices
 
-        # Apply scale transforms to faces before projection
-        xs_scaled, ys_scaled, zs_scaled = _apply_scale_transforms(
-            self._faces[..., 0], self._faces[..., 1], self._faces[..., 2],
-            self.axes)
-        faces_scaled = np.stack([xs_scaled, ys_scaled, zs_scaled], axis=-1)
-
         # Some faces might contain masked vertices, so we want to ignore any
         # errors that those might cause
         with np.errstate(invalid='ignore', divide='ignore'):
-            pfaces = proj3d._proj_transform_vectors(faces_scaled, self.axes.M)
+            pfaces = proj3d._proj_transform_vectors_scaled(
+                self._faces, self.axes)
 
         if self._axlim_clip:
             viewlim_mask = _viewlim_mask(self._faces[..., 0], self._faces[..., 1],
diff --git a/lib/mpl_toolkits/mplot3d/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py
@@ -131,6 +131,33 @@ def _ortho_transformation(zfront, zback):
     return proj_matrix
 
 
+def _apply_scale_transforms(xs, ys, zs, axes):
+    """
+    Apply axis scale transforms to 3D coordinates.
+
+    Transforms data coordinates to transformed coordinates (applying log,
+    symlog, etc.) for 3D projection. Preserves masked arrays.
+    """
+    def transform_coord(coord, axis):
+        coord = np.asanyarray(coord)
+        data = np.ma.getdata(coord).ravel()
+        return axis.get_transform().transform(data).reshape(coord.shape)
+
+    xs_scaled = transform_coord(xs, axes.xaxis)
+    ys_scaled = transform_coord(ys, axes.yaxis)
+    zs_scaled = transform_coord(zs, axes.zaxis)
+
+    # Preserve combined mask from any masked input
+    masks = [np.ma.getmask(a) for a in [xs, ys, zs]]
+    if any(m is not np.ma.nomask for m in masks):
+        combined = np.ma.mask_or(np.ma.mask_or(masks[0], masks[1]), masks[2])
+        xs_scaled = np.ma.array(xs_scaled, mask=combined)
+        ys_scaled = np.ma.array(ys_scaled, mask=combined)
+        zs_scaled = np.ma.array(zs_scaled, mask=combined)
+
+    return xs_scaled, ys_scaled, zs_scaled
+
+
 def _proj_transform_vec(vec, M):
     vecw = np.dot(M, vec.data)
     ts = vecw[0:3]/vecw[3]
@@ -234,3 +261,50 @@ def _proj_trans_points(points, M):
     points = np.asanyarray(points)
     xs, ys, zs = points[:, 0], points[:, 1], points[:, 2]
     return proj_transform(xs, ys, zs, M)
+
+
+def _proj_transform_clip_scaled(xs, ys, zs, axes):
+    """
+    Apply scale transforms and project with clipping.
+
+    Combines `_apply_scale_transforms` and `_proj_transform_clip` into a
+    single call. Returns txs, tys, tzs, tis.
+    """
+    xs, ys, zs = _apply_scale_transforms(xs, ys, zs, axes)
+    return _proj_transform_clip(xs, ys, zs, axes.M, axes._focal_length)
+
+
+def _proj_transform_vectors_scaled(vecs, axes):
+    """
+    Apply scale transforms and project vectors.
+
+    Combines `_apply_scale_transforms` and `_proj_transform_vectors` into a
+    single call. *vecs* has shape ``(..., 3)``.
+    """
+    xs, ys, zs = _apply_scale_transforms(
+        vecs[..., 0], vecs[..., 1], vecs[..., 2], axes)
+    vecs = np.stack([xs, ys, zs], axis=-1)
+    return _proj_transform_vectors(vecs, axes.M)
+
+
+def _proj_transform_scaled(xs, ys, zs, axes):
+    """
+    Apply scale transforms and project.
+
+    Combines `_apply_scale_transforms` and `proj_transform` into a single
+    call. Returns txs, tys, tzs.
+    """
+    xs, ys, zs = _apply_scale_transforms(xs, ys, zs, axes)
+    return proj_transform(xs, ys, zs, axes.M)
+
+
+def _proj_trans_points_scaled(points, axes):
+    """
+    Apply scale transforms and project points.
+
+    Combines `_apply_scale_transforms` and `_proj_trans_points` into a single
+    call.
+    """
+    points = np.asanyarray(points)
+    xs, ys, zs = points[:, 0], points[:, 1], points[:, 2]
+    return _proj_transform_scaled(xs, ys, zs, axes)
