From 2e72141c6b324d1522bb6b7f6c003c873fab2033 Mon Sep 17 00:00:00 2001
From: Elliott Sales de Andrade <quantum.analyst@gmail.com>
Date: Wed, 20 Dec 2023 16:54:55 -0500
Subject: [PATCH] Simplify affine_transform on C++ side

---
 lib/matplotlib/tests/test_transforms.py | 15 ++++++---------
 lib/matplotlib/transforms.py            | 11 ++++++++---
 src/_eigen.cpp                          |  6 ++----
 3 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/lib/matplotlib/tests/test_transforms.py b/lib/matplotlib/tests/test_transforms.py
index 31fcc1276a31..bb18ab6ee386 100644
--- a/lib/matplotlib/tests/test_transforms.py
+++ b/lib/matplotlib/tests/test_transforms.py
@@ -286,20 +286,17 @@ def test_translate_plus_other(self):
 
     def test_invalid_arguments(self):
         t = mtransforms.Affine2D()
-        # The wrong number of dimensions and a wrong shape with a possible number of
-        # dimensions are both caught by pybind11, which always raises the less precise
-        # RuntimeError.
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform(1)
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform([[[1]]])
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform([])
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform([1])
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform([[1]])
-        with pytest.raises(RuntimeError):
+        with pytest.raises(TypeError):
             t.transform([[1, 2, 3]])
 
     def test_copy(self):
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
index 99fa20a4a8c3..a70a8f02b068 100644
--- a/lib/matplotlib/transforms.py
+++ b/lib/matplotlib/transforms.py
@@ -1867,9 +1867,14 @@ def to_values(self):
     def transform_affine(self, values):
         mtx = self._get_eigen_matrix()
         if isinstance(values, np.ma.MaskedArray):
-            tpoints = mtx.affine_transform(values.data)
-            return np.ma.MaskedArray(tpoints, mask=np.ma.getmask(values))
-        return mtx.affine_transform(values)
+            tpoints = mtx.affine_transform(values.data.T).T
+            result = np.ma.MaskedArray(tpoints, mask=np.ma.getmask(values))
+        else:
+            values = np.asarray(values)
+            result = mtx.affine_transform(values.T).T
+        if len(values.shape) == 2:
+            result = np.atleast_2d(result)
+        return result
 
     if DEBUG:
         _transform_affine = transform_affine
diff --git a/src/_eigen.cpp b/src/_eigen.cpp
index 764f35685293..692118100749 100644
--- a/src/_eigen.cpp
+++ b/src/_eigen.cpp
@@ -128,10 +128,8 @@ PYBIND11_MODULE(_eigen, m) {
             }
         )
         .def("affine_transform",
-            [](const Eigen::Affine2d& self, py::array_t<double> vertices_arr) {
-                auto vertices = vertices_arr.attr("transpose")().cast<Eigen::Ref<const Eigen::Matrix2Xd>>();
-                auto result = py::cast(self * vertices, py::return_value_policy::move);
-                return result.attr("transpose")();
+            [](const Eigen::Affine2d& self, Eigen::Ref<const Eigen::Matrix2Xd> vertices) {
+                return self * vertices;
             }
         )