Route affine transforms through Eigen

Author: QuLogic

lib/matplotlib/_eigen.pyi

@@ -19,3 +19,4 @@ class Affine2d:
     def scale(self, sx: float, sy: float) -> None: ...
     def skew(self, xShear: float, yShear: float) -> None: ...
     def translate(self, tx: float, ty: float) -> None: ...
+    def affine_transform(self, vertices: np.ndarray) -> np.ndarray: ...

lib/matplotlib/_path.pyi

@@ -1,8 +1,5 @@
 from collections.abc import Sequence
 
-import numpy as np
-
 from .transforms import BboxBase
 
-def affine_transform(points: np.ndarray, trans: np.ndarray) -> np.ndarray: ...
 def count_bboxes_overlapping_bbox(bbox: BboxBase, bboxes: Sequence[BboxBase]) -> int: ...

lib/matplotlib/tests/test_transforms.py

@@ -285,24 +285,22 @@ def test_translate_plus_other(self):
         assert_array_almost_equal(trans.transform(self.multiple_points),
                                   [[45, 49.75], [48.5, 51.5], [48, 48.75]])
 
-    def test_invalid_transform(self):
+    def test_invalid_arguments(self):
         t = mtransforms.Affine2D()
-        # There are two different exceptions, since the wrong number of
-        # dimensions is caught when constructing an array_view, and that
-        # raises a ValueError, and a wrong shape with a possible number
-        # of dimensions is caught by our CALL_CPP macro, which always
-        # raises the less precise RuntimeError.
-        with pytest.raises(ValueError):
+        # 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):
             t.transform(1)
-        with pytest.raises(ValueError):
+        with pytest.raises(RuntimeError):
             t.transform([[[1]]])
         with pytest.raises(RuntimeError):
             t.transform([])
         with pytest.raises(RuntimeError):
             t.transform([1])
-        with pytest.raises(ValueError):
+        with pytest.raises(RuntimeError):
             t.transform([[1]])
-        with pytest.raises(ValueError):
+        with pytest.raises(RuntimeError):
             t.transform([[1, 2, 3]])
 
     def test_copy(self):

lib/matplotlib/transforms.py

@@ -45,7 +45,7 @@
 from numpy.linalg import inv
 
 from matplotlib import _api, _eigen
-from matplotlib._path import affine_transform, count_bboxes_overlapping_bbox
+from matplotlib._path import count_bboxes_overlapping_bbox
 from .path import Path
 
 DEBUG = False
@@ -1798,6 +1798,15 @@ def __array__(self, *args, **kwargs):
         # optimises the access of the transform matrix vs. the superclass
         return self.get_matrix()
 
+    def _get_eigen_matrix(self):
+        if self._invalid:
+            mtx = self.get_matrix()
+        # Back-compat with external stuff.
+        if isinstance(eigen := getattr(self, '_mtx'), _eigen.Affine2d):
+            return eigen
+        else:
+            return _eigen.Affine2d(mtx)
+
     def __eq__(self, other):
         if getattr(other, "is_affine", False) and hasattr(other, "get_matrix"):
             return (self.get_matrix() == other.get_matrix()).all()
@@ -1870,11 +1879,11 @@ def to_values(self):
         return tuple(mtx[:2].swapaxes(0, 1).flat)
 
     def transform_affine(self, values):
-        mtx = self.get_matrix()
+        mtx = self._get_eigen_matrix()
         if isinstance(values, np.ma.MaskedArray):
-            tpoints = affine_transform(values.data, mtx)
+            tpoints = mtx.affine_transform(values.data)
             return np.ma.MaskedArray(tpoints, mask=np.ma.getmask(values))
-        return affine_transform(values, mtx)
+        return mtx.affine_transform(values)
 
     if DEBUG:
         _transform_affine = transform_affine

src/_eigen.cpp

@@ -109,5 +109,19 @@ PYBIND11_MODULE(_eigen, m) {
             self.pretranslate(Eigen::Vector2d(tx, ty));
         }, "tx"_a, "ty"_a,
         "Add a translation in place.")
+
+        .def("affine_transform",
+            [](const Eigen::Affine2d& self, Eigen::Ref<const Eigen::Vector2d> vertices) {
+                Eigen::Vector2d result = self * vertices;
+                return result;
+            }
+        )
+        .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")();
+            }
+        )
     ;
 }

src/_path.h

@@ -699,63 +699,6 @@ clip_path_to_rect(PathIterator &path, agg::rect_d &rect, bool inside, std::vecto
     _finalize_polygon(results, 1);
 }
 
-template <class VerticesArray, class ResultArray>
-void affine_transform_2d(VerticesArray &vertices, agg::trans_affine &trans, ResultArray &result)
-{
-    if (vertices.size() != 0 && vertices.shape(1) != 2) {
-        throw std::runtime_error("Invalid vertices array.");
-    }
-
-    size_t n = vertices.shape(0);
-    double x;
-    double y;
-    double t0;
-    double t1;
-    double t;
-
-    for (size_t i = 0; i < n; ++i) {
-        x = vertices(i, 0);
-        y = vertices(i, 1);
-
-        t0 = trans.sx * x;
-        t1 = trans.shx * y;
-        t = t0 + t1 + trans.tx;
-        result(i, 0) = t;
-
-        t0 = trans.shy * x;
-        t1 = trans.sy * y;
-        t = t0 + t1 + trans.ty;
-        result(i, 1) = t;
-    }
-}
-
-template <class VerticesArray, class ResultArray>
-void affine_transform_1d(VerticesArray &vertices, agg::trans_affine &trans, ResultArray &result)
-{
-    if (vertices.shape(0) != 2) {
-        throw std::runtime_error("Invalid vertices array.");
-    }
-
-    double x;
-    double y;
-    double t0;
-    double t1;
-    double t;
-
-    x = vertices(0);
-    y = vertices(1);
-
-    t0 = trans.sx * x;
-    t1 = trans.shx * y;
-    t = t0 + t1 + trans.tx;
-    result(0) = t;
-
-    t0 = trans.shy * x;
-    t1 = trans.sy * y;
-    t = t0 + t1 + trans.ty;
-    result(1) = t;
-}
-
 template <class BBoxArray>
 int count_bboxes_overlapping_bbox(agg::rect_d &a, BBoxArray &bboxes)
 {

src/_path_wrapper.cpp

@@ -116,36 +116,6 @@ Py_clip_path_to_rect(mpl::PathIterator path, agg::rect_d rect, bool inside)
     return convert_polygon_vector(result);
 }
 
-static py::object
-Py_affine_transform(py::array_t<double, py::array::c_style | py::array::forcecast> vertices_arr,
-                    agg::trans_affine trans)
-{
-    if (vertices_arr.ndim() == 2) {
-        auto vertices = vertices_arr.unchecked<2>();
-
-        check_trailing_shape(vertices, "vertices", 2);
-
-        py::ssize_t dims[] = { vertices.shape(0), 2 };
-        py::array_t<double> result(dims);
-        auto result_mutable = result.mutable_unchecked<2>();
-
-        affine_transform_2d(vertices, trans, result_mutable);
-        return result;
-    } else if (vertices_arr.ndim() == 1) {
-        auto vertices = vertices_arr.unchecked<1>();
-
-        py::ssize_t dims[] = { vertices.shape(0) };
-        py::array_t<double> result(dims);
-        auto result_mutable = result.mutable_unchecked<1>();
-
-        affine_transform_1d(vertices, trans, result_mutable);
-        return result;
-    } else {
-        throw py::value_error(
-            "vertices must be 1D or 2D, not" + std::to_string(vertices_arr.ndim()) + "D");
-    }
-}
-
 static int
 Py_count_bboxes_overlapping_bbox(agg::rect_d bbox, py::array_t<double> bboxes_obj)
 {
@@ -326,8 +296,6 @@ PYBIND11_MODULE(_path, m, py::mod_gil_not_used())
           "path_a"_a, "trans_a"_a, "path_b"_a, "trans_b"_a);
     m.def("clip_path_to_rect", &Py_clip_path_to_rect,
           "path"_a, "rect"_a, "inside"_a);
-    m.def("affine_transform", &Py_affine_transform,
-          "points"_a, "trans"_a);
     m.def("count_bboxes_overlapping_bbox", &Py_count_bboxes_overlapping_bbox,
           "bbox"_a, "bboxes"_a);
     m.def("path_intersects_path", &Py_path_intersects_path,