Generalize Affine2DBase to AffineImmutable

trananso · trananso · commit 12e7d0103a28 · 2024-04-21T16:54:41.000-04:00
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -439,10 +439,10 @@ def draw_image(self, gc, x, y, im, transform=None):
         im : (N, M, 4) array of `numpy.uint8`
             An array of RGBA pixels.
 
-        transform : `~matplotlib.transforms.Affine2DBase`
+        transform : `~matplotlib.transforms.AffineImmutable`
             If and only if the concrete backend is written such that
             `option_scale_image` returns ``True``, an affine transformation
-            (i.e., an `.Affine2DBase`) *may* be passed to `draw_image`.  The
+            (i.e., an `.AffineImmutable`) *may* be passed to `draw_image`.  The
             translation vector of the transformation is given in physical units
             (i.e., dots or pixels). Note that the transformation does not
             override *x* and *y*, and has to be applied *before* translating
diff --git a/lib/matplotlib/backend_bases.pyi b/lib/matplotlib/backend_bases.pyi
@@ -91,7 +91,7 @@ class RendererBase:
         x: float,
         y: float,
         im: ArrayLike,
-        transform: transforms.Affine2DBase | None = ...,
+        transform: transforms.AffineImmutable | None = ...,
     ) -> None: ...
     def option_image_nocomposite(self) -> bool: ...
     def option_scale_image(self) -> bool: ...
diff --git a/lib/matplotlib/backends/backend_svg.py b/lib/matplotlib/backends/backend_svg.py
@@ -22,7 +22,7 @@
 from matplotlib.dates import UTC
 from matplotlib.path import Path
 from matplotlib import _path
-from matplotlib.transforms import Affine2D, Affine2DBase
+from matplotlib.transforms import Affine2D, AffineImmutable
 
 
 _log = logging.getLogger(__name__)
@@ -255,7 +255,7 @@ def _generate_transform(transform_list):
                 or type == 'translate' and value == (0, 0)
                 or type == 'rotate' and value == (0,)):
             continue
-        if type == 'matrix' and isinstance(value, Affine2DBase):
+        if type == 'matrix' and isinstance(value, AffineImmutable):
             value = value.to_values()
         parts.append('{}({})'.format(
             type, ' '.join(_short_float_fmt(x) for x in value)))
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -1062,10 +1062,10 @@ def get_path_collection_extents(
     master_transform : `~matplotlib.transforms.Transform`
         Global transformation applied to all paths.
     paths : list of `Path`
-    transforms : list of `~matplotlib.transforms.Affine2DBase`
+    transforms : list of `~matplotlib.transforms.AffineImmutable`
         If non-empty, this overrides *master_transform*.
     offsets : (N, 2) array-like
-    offset_transform : `~matplotlib.transforms.Affine2DBase`
+    offset_transform : `~matplotlib.transforms.AffineImmutable`
         Transform applied to the offsets before offsetting the path.
 
     Notes
diff --git a/lib/matplotlib/projections/polar.py b/lib/matplotlib/projections/polar.py
@@ -155,7 +155,7 @@ def inverted(self):
         )
 
 
-class PolarAffine(mtransforms.Affine2DBase):
+class PolarAffine(mtransforms.AffineImmutable):
     r"""
     The affine part of the polar projection.
 
@@ -181,7 +181,7 @@ def __init__(self, scale_transform, limits):
             View limits of the data. The only part of its bounds that is used
             is the y limits (for the radius limits).
         """
-        super().__init__()
+        super().__init__(dims=2)
         self._scale_transform = scale_transform
         self._limits = limits
         self.set_children(scale_transform, limits)
diff --git a/lib/matplotlib/projections/polar.pyi b/lib/matplotlib/projections/polar.pyi
@@ -23,7 +23,7 @@ class PolarTransform(mtransforms.Transform):
     ) -> None: ...
     def inverted(self) -> InvertedPolarTransform: ...
 
-class PolarAffine(mtransforms.Affine2DBase):
+class PolarAffine(mtransforms.AffineImmutable):
     def __init__(
         self, scale_transform: mtransforms.Transform, limits: mtransforms.BboxBase
     ) -> None: ...
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -1,13 +1,13 @@
 """
 Matplotlib includes a framework for arbitrary geometric
-transformations that is used determine the final position of all
+transformations that is used to determine the final position of all
 elements drawn on the canvas.
 
 Transforms are composed into trees of `TransformNode` objects
-whose actual value depends on their children.  When the contents of
-children change, their parents are automatically invalidated.  The
+whose actual value depends on their children. When the contents of
+children change, their parents are automatically invalidated. The
 next time an invalidated transform is accessed, it is recomputed to
-reflect those changes.  This invalidation/caching approach prevents
+reflect those changes. This invalidation/caching approach prevents
 unnecessary recomputations of transforms, and contributes to better
 interactive performance.
 
@@ -1372,7 +1372,7 @@ def _iter_break_from_left_to_right(self):
         This is equivalent to flattening the stack then yielding
         ``flat_stack[:i], flat_stack[i:]`` where i=0..(n-1).
         """
-        yield IdentityTransform(), self
+        yield IdentityTransform(dims=self.input_dims), self
 
     @property
     def depth(self):
@@ -1578,7 +1578,7 @@ def transform_bbox(self, bbox):
 
     def get_affine(self):
         """Get the affine part of this transform."""
-        return IdentityTransform()
+        return IdentityTransform(dims=self.input_dims)
 
     def get_matrix(self):
         """Get the matrix for the affine part of this transform."""
@@ -1821,11 +1821,13 @@ def get_affine(self):
         return self
 
 
-class Affine2DBase(AffineBase):
+class AffineImmutable(AffineBase):
     """
-    The base class of all 2D affine transformations.
+    The base class of all affine transformations.
 
-    2D affine transformations are performed using a 3x3 numpy array::
+    Affine transformations for the n-th degree are performed using a
+    numpy array with shape (n+1, n+1). For example, 2D affine
+    transformations are performed using a 3x3 numpy array::
 
         a c e
         b d f
@@ -1835,34 +1837,65 @@ class Affine2DBase(AffineBase):
     affine transformation, use `Affine2D`.
 
     Subclasses of this class will generally only need to override a
-    constructor and `~.Transform.get_matrix` that generates a custom 3x3 matrix.
+    constructor and `~.Transform.get_matrix` that generates a custom matrix
+    with the appropriate shape.
     """
-    input_dims = 2
-    output_dims = 2
+    def __init__(self, *args, dims=2, **kwargs):
+        self.input_dims = dims
+        self.output_dims = dims
+        super().__init__(*args, **kwargs)
 
     def frozen(self):
         # docstring inherited
-        return Affine2D(self.get_matrix().copy())
+        return _affine_factory(self.get_matrix().copy(), self.input_dims)
 
     @property
     def is_separable(self):
         mtx = self.get_matrix()
-        return mtx[0, 1] == mtx[1, 0] == 0.0
+        separable = True
+        for i in range(self.input_dims):
+            for j in range(i+1, self.input_dims):
+                separable = separable and mtx[i, j] == 0.0
+                separable = separable and mtx[j, i] == 0.0
+        return separable
 
     def to_values(self):
         """
-        Return the values of the matrix as an ``(a, b, c, d, e, f)`` tuple.
+        Return the values of the matrix as a tuple.
         """
         mtx = self.get_matrix()
-        return tuple(mtx[:2].swapaxes(0, 1).flat)
+        return tuple(mtx[:self.input_dims].swapaxes(0, 1).flat)
+
+    def _affine_transform(self, vertices):
+        """
+        Default implementation of affine_transform if a C implementation isn't
+        available
+        """
+        mtx = self.get_matrix()
+        values = np.asanyarray(vertices)
+        # single point
+        if (values.shape == (self.input_dims,)):
+            return mtx.dot(vertices + [1]).tolist()[:self.input_dims]
+        # multiple points
+        if (values.shape[1] == self.input_dims):
+            homogeneous = np.hstack((values, np.ones((values.shape[0], 1))))
+            return np.dot(mtx, homogeneous.T).T[:, :-1]
+
+        raise TypeError("Dimensions of input must match the input dimensions of "
+                        "the transform")
 
     @_api.rename_parameter("3.8", "points", "values")
     def transform_affine(self, values):
         mtx = self.get_matrix()
+
+        # Default to python implementation if C implementation isn't available
+        transform_fn = (affine_transform if self.input_dims <= 2
+                        else self._affine_transform)
+
         if isinstance(values, np.ma.MaskedArray):
-            tpoints = affine_transform(values.data, mtx)
+            tpoints = transform_fn(values.data, mtx)
             return np.ma.MaskedArray(tpoints, mask=np.ma.getmask(values))
-        return affine_transform(values, mtx)
+        return transform_fn(values, mtx)
 
     if DEBUG:
         _transform_affine = transform_affine
@@ -1886,12 +1919,25 @@ def inverted(self):
             shorthand_name = None
             if self._shorthand_name:
                 shorthand_name = '(%s)-1' % self._shorthand_name
-            self._inverted = Affine2D(inv(mtx), shorthand_name=shorthand_name)
+            self._inverted = _affine_factory(inv(mtx), self.input_dims,
+                                             shorthand_name=shorthand_name)
             self._invalid = 0
         return self._inverted
 
 
-class Affine2D(Affine2DBase):
+@_api.deprecated("3.9", alternative="AffineImmutable")
+class Affine2DBase(AffineImmutable):
+    pass
+
+
+def _affine_factory(mtx, dims, *args, **kwargs):
+    if dims == 2:
+        return Affine2D(mtx, *args, **kwargs)
+    else:
+        return NotImplemented
+
+
+class Affine2D(AffineImmutable):
     """
     A mutable 2D affine transformation.
     """
@@ -1906,10 +1952,9 @@ def __init__(self, matrix=None, **kwargs):
 
         If *matrix* is None, initialize with the identity transform.
         """
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         if matrix is None:
-            # A bit faster than np.identity(3).
-            matrix = IdentityTransform._mtx
+            matrix = np.identity(3)
         self._mtx = matrix.copy()
         self._invalid = 0
 
@@ -1967,18 +2012,20 @@ def set_matrix(self, mtx):
     def set(self, other):
         """
         Set this transformation from the frozen copy of another
-        `Affine2DBase` object.
+        2D `AffineImmutable` object.
         """
-        _api.check_isinstance(Affine2DBase, other=other)
+        _api.check_isinstance(AffineImmutable, other=other)
+        if (other.input_dims != 2):
+            raise TypeError("Mismatch between dimensions of AffineImmutable "
+                            "and Affine2D")
         self._mtx = other.get_matrix()
         self.invalidate()
 
     def clear(self):
         """
         Reset the underlying matrix to the identity transform.
         """
-        # A bit faster than np.identity(3).
-        self._mtx = IdentityTransform._mtx.copy()
+        self._mtx = np.identity(3)
         self.invalidate()
         return self
 
@@ -2113,12 +2160,14 @@ def skew_deg(self, xShear, yShear):
         return self.skew(math.radians(xShear), math.radians(yShear))
 
 
-class IdentityTransform(Affine2DBase):
+class IdentityTransform(AffineImmutable):
     """
     A special class that does one thing, the identity transform, in a
     fast way.
     """
-    _mtx = np.identity(3)
+    def __init__(self, *args, **kwargs):
+        super().__init__(self, *args, **kwargs)
+        self._mtx = np.identity(self.input_dims + 1)
 
     def frozen(self):
         # docstring inherited
@@ -2532,7 +2581,7 @@ def composite_transform_factory(a, b):
     return CompositeGenericTransform(a, b)
 
 
-class BboxTransform(Affine2DBase):
+class BboxTransform(AffineImmutable):
     """
     `BboxTransform` linearly transforms points from one `Bbox` to another.
     """
@@ -2546,7 +2595,7 @@ def __init__(self, boxin, boxout, **kwargs):
         """
         _api.check_isinstance(BboxBase, boxin=boxin, boxout=boxout)
 
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         self._boxin = boxin
         self._boxout = boxout
         self.set_children(boxin, boxout)
@@ -2574,7 +2623,7 @@ def get_matrix(self):
         return self._mtx
 
 
-class BboxTransformTo(Affine2DBase):
+class BboxTransformTo(AffineImmutable):
     """
     `BboxTransformTo` is a transformation that linearly transforms points from
     the unit bounding box to a given `Bbox`.
@@ -2589,7 +2638,7 @@ def __init__(self, boxout, **kwargs):
         """
         _api.check_isinstance(BboxBase, boxout=boxout)
 
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         self._boxout = boxout
         self.set_children(boxout)
         self._mtx = None
@@ -2633,7 +2682,7 @@ def get_matrix(self):
         return self._mtx
 
 
-class BboxTransformFrom(Affine2DBase):
+class BboxTransformFrom(AffineImmutable):
     """
     `BboxTransformFrom` linearly transforms points from a given `Bbox` to the
     unit bounding box.
@@ -2643,7 +2692,7 @@ class BboxTransformFrom(Affine2DBase):
     def __init__(self, boxin, **kwargs):
         _api.check_isinstance(BboxBase, boxin=boxin)
 
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         self._boxin = boxin
         self.set_children(boxin)
         self._mtx = None
@@ -2668,13 +2717,13 @@ def get_matrix(self):
         return self._mtx
 
 
-class ScaledTranslation(Affine2DBase):
+class ScaledTranslation(AffineImmutable):
     """
     A transformation that translates by *xt* and *yt*, after *xt* and *yt*
     have been transformed by *scale_trans*.
     """
     def __init__(self, xt, yt, scale_trans, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         self._t = (xt, yt)
         self._scale_trans = scale_trans
         self.set_children(scale_trans)
@@ -2686,15 +2735,14 @@ def __init__(self, xt, yt, scale_trans, **kwargs):
     def get_matrix(self):
         # docstring inherited
         if self._invalid:
-            # A bit faster than np.identity(3).
-            self._mtx = IdentityTransform._mtx.copy()
+            self._mtx = np.identity(3)
             self._mtx[:2, 2] = self._scale_trans.transform(self._t)
             self._invalid = 0
             self._inverted = None
         return self._mtx
 
 
-class AffineDeltaTransform(Affine2DBase):
+class AffineDeltaTransform(AffineImmutable):
     r"""
     A transform wrapper for transforming displacements between pairs of points.
 
@@ -2712,7 +2760,7 @@ class AffineDeltaTransform(Affine2DBase):
     """
 
     def __init__(self, transform, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(dims=2, **kwargs)
         self._base_transform = transform
 
     __str__ = _make_str_method("_base_transform")
diff --git a/lib/matplotlib/transforms.pyi b/lib/matplotlib/transforms.pyi
