Add non 2D paths wip [skip ci]

trananso · trananso · commit 43c8913200e8 · 2024-05-01T17:17:41.000-04:00
diff --git a/lib/matplotlib/path.py b/lib/matplotlib/path.py
@@ -97,13 +97,13 @@ class Path:
                              CLOSEPOLY: 1}
 
     def __init__(self, vertices, codes=None, _interpolation_steps=1,
-                 closed=False, readonly=False):
+                 closed=False, readonly=False, dims=2):
         """
         Create a new path with the given vertices and codes.
 
         Parameters
         ----------
-        vertices : (N, 2) array-like
+        vertices : (N, dims) array-like
             The path vertices, as an array, masked array or sequence of pairs.
             Masked values, if any, will be converted to NaNs, which are then
             handled correctly by the Agg PathIterator and other consumers of
@@ -125,9 +125,14 @@ def __init__(self, vertices, codes=None, _interpolation_steps=1,
         readonly : bool, optional
             Makes the path behave in an immutable way and sets the vertices
             and codes as read-only arrays.
+        dims : int, optional
         """
+        if dims <= 1:
+            raise ValueError("Path must be at least 2D")
+        self._dims = dims
+
         vertices = _to_unmasked_float_array(vertices)
-        _api.check_shape((None, 2), vertices=vertices)
+        _api.check_shape((None, dims), vertices=vertices)
 
         if codes is not None:
             codes = np.asarray(codes, self.code_type)
@@ -178,6 +183,7 @@ def _fast_from_codes_and_verts(cls, verts, codes, internals_from=None):
         pth._vertices = _to_unmasked_float_array(verts)
         pth._codes = codes
         pth._readonly = False
+        pth._dims = pth._vertices.shape[1]
         if internals_from is not None:
             pth._should_simplify = internals_from._should_simplify
             pth._simplify_threshold = internals_from._simplify_threshold
@@ -210,13 +216,15 @@ def _update_values(self):
 
     @property
     def vertices(self):
-        """The vertices of the `Path` as an (N, 2) array."""
+        """The vertices of the `Path` as an (N, dims) array."""
         return self._vertices
 
     @vertices.setter
     def vertices(self, vertices):
         if self._readonly:
             raise AttributeError("Can't set vertices on a readonly Path")
+        if not _api.check_shape((None, self._dims), vertices=vertices):
+            raise ValueError("Vertices shape does not match path dimensions")
         self._vertices = vertices
         self._update_values()
 
@@ -239,6 +247,26 @@ def codes(self, codes):
         self._codes = codes
         self._update_values()
 
+    @property
+    def dims(self):
+        """
+        The dimensions of vertices in the `Path`.
+        """
+        return self._dims
+
+    @dims.setter
+    def dims(self, dims):
+        if dims <= 2:
+            raise ValueError("Path must be at least 2D")
+
+        if dims < self._dims:
+            self._vertices = self._vertices[:, :dims]
+        elif dims > self._dims:
+            self._vertices = np.pad(self._vertices,
+                                    ((0, 0), (0, dims - self._dims)),
+                                    mode='constant', constant_values=np.nan)
+        self._dims = dims
+
     @property
     def simplify_threshold(self):
         """
@@ -298,17 +326,17 @@ def make_compound_path_from_polys(cls, XY):
 
         Parameters
         ----------
-        XY : (numpolys, numsides, 2) array
+        XY : (numpolys, numsides, dims) array
         """
         # for each poly: 1 for the MOVETO, (numsides-1) for the LINETO, 1 for
         # the CLOSEPOLY; the vert for the closepoly is ignored but we still
         # need it to keep the codes aligned with the vertices
-        numpolys, numsides, two = XY.shape
-        if two != 2:
-            raise ValueError("The third dimension of 'XY' must be 2")
+        numpolys, numsides,  dims = XY.shape
+        if dims < 2:
+            raise ValueError("The third dimension of 'XY' must be at least 2")
         stride = numsides + 1
         nverts = numpolys * stride
-        verts = np.zeros((nverts, 2))
+        verts = np.zeros((nverts, dims))
         codes = np.full(nverts, cls.LINETO, dtype=cls.code_type)
         codes[0::stride] = cls.MOVETO
         codes[numsides::stride] = cls.CLOSEPOLY
@@ -323,6 +351,9 @@ def make_compound_path(cls, *args):
         """
         if not args:
             return Path(np.empty([0, 2], dtype=np.float32))
+        if not all([path._dims != args[0]._dims for path in args]):
+            raise ValueError("Paths provided must be the same dimension")
+
         vertices = np.concatenate([path.vertices for path in args])
         codes = np.empty(len(vertices), dtype=cls.code_type)
         i = 0
@@ -338,6 +369,16 @@ def make_compound_path(cls, *args):
         not_stop_mask = codes != cls.STOP  # Remove STOPs, as internal STOPs are a bug.
         return cls(vertices[not_stop_mask], codes[not_stop_mask])
 
+    @classmethod
+    def _project_to_2d(cls, path, transform=None):
+        if transform is not None:
+            path = transform.transform_path(path)
+
+        if path._dims > 2:
+            path = Path._fast_from_codes_and_verts(path.vertices[:, 2], path.codes,
+                                                   path)
+        return path
+
     def __repr__(self):
         return f"Path({self.vertices!r}, {self.codes!r})"
 
@@ -478,6 +519,10 @@ def cleaned(self, transform=None, remove_nans=False, clip=None,
         --------
         Path.iter_segments : for details of the keyword arguments.
         """
+        # Not implemented for non 2D
+        if self._dims != 2:
+            return self
+
         vertices, codes = _path.cleanup_path(
             self, transform, remove_nans, clip, snap, stroke_width, simplify,
             curves, sketch)
@@ -497,7 +542,7 @@ def transformed(self, transform):
             automatically update when the transform changes.
         """
         return Path(transform.transform(self.vertices), self.codes,
-                    self._interpolation_steps)
+                    self._interpolation_steps, dims=transform.output_dims)
 
     def contains_point(self, point, transform=None, radius=0.0):
         """
@@ -540,14 +585,22 @@ def contains_point(self, point, transform=None, radius=0.0):
         """
         if transform is not None:
             transform = transform.frozen()
+
+        # Transform the path, and then toss out the z dimension
+        if self.dims > 2:
+            pth = Path._project_to_2d(self, transform)
+            transform = None
+
         # `point_in_path` does not handle nonlinear transforms, so we
         # transform the path ourselves.  If *transform* is affine, letting
         # `point_in_path` handle the transform avoids allocating an extra
         # buffer.
-        if transform and not transform.is_affine:
+        elif transform and not transform.is_affine:
             self = transform.transform_path(self)
+            pth = self
             transform = None
-        return _path.point_in_path(point[0], point[1], radius, self, transform)
+
+        return _path.point_in_path(point[0], point[1], radius, pth, transform)
 
     def contains_points(self, points, transform=None, radius=0.0):
         """
@@ -590,7 +643,11 @@ def contains_points(self, points, transform=None, radius=0.0):
         """
         if transform is not None:
             transform = transform.frozen()
-        result = _path.points_in_path(points, radius, self, transform)
+        pth = self
+        if self._dims > 2:
+            pth = Path._project_to_2d(self, transform)
+            transform = None
+        result = _path.points_in_path(points, radius, pth, transform)
         return result.astype('bool')
 
     def contains_path(self, path, transform=None):
@@ -602,7 +659,15 @@ def contains_path(self, path, transform=None):
         """
         if transform is not None:
             transform = transform.frozen()
-        return _path.path_in_path(self, None, path, transform)
+
+        a_pth = Path._project_to_2d(self, None)
+        if path._dims > 2:
+            b_pth = Path._project_to_2d(path, transform)
+            transform = None
+        else:
+            b_pth = path
+
+        return _path.path_in_path(a_pth, None, b_pth, transform)
 
     def get_extents(self, transform=None, **kwargs):
         """
@@ -652,7 +717,9 @@ def intersects_path(self, other, filled=True):
         If *filled* is True, then this also returns True if one path completely
         encloses the other (i.e., the paths are treated as filled).
         """
-        return _path.path_intersects_path(self, other, filled)
+        a = Path._project_to_2d(self)
+        b = Path._project_to_2d(other)
+        return _path.path_intersects_path(a, b, filled)
 
     def intersects_bbox(self, bbox, filled=True):
         """
@@ -663,8 +730,9 @@ def intersects_bbox(self, bbox, filled=True):
 
         The bounding box is always considered filled.
         """
+        pth = Path._project_to_2d(self)
         return _path.path_intersects_rectangle(
-            self, bbox.x0, bbox.y0, bbox.x1, bbox.y1, filled)
+            pth, bbox.x0, bbox.y0, bbox.x1, bbox.y1, filled)
 
     def interpolated(self, steps):
         """
@@ -683,7 +751,7 @@ def interpolated(self, steps):
             new_codes[0::steps] = codes
         else:
             new_codes = None
-        return Path(vertices, new_codes)
+        return Path(vertices, new_codes, dims=vertices.shape[1])
 
     def to_polygons(self, transform=None, width=0, height=0, closed_only=True):
         """
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -1608,8 +1608,6 @@ def transform_path(self, path):
         In some cases, this transform may insert curves into the path
         that began as line segments.
         """
-        if self.input_dims != 2 or self.output_dims != 2:
-            raise NotImplementedError('Only defined in 2D')
         return self.transform_path_affine(self.transform_path_non_affine(path))
 
     def transform_path_affine(self, path):
@@ -1620,8 +1618,6 @@ def transform_path_affine(self, path):
         ``transform_path(path)`` is equivalent to
         ``transform_path_affine(transform_path_non_affine(values))``.
         """
-        if self.input_dims != 2 or self.output_dims != 2:
-            raise NotImplementedError('Only defined in 2D')
         return self.get_affine().transform_path_affine(path)
 
     def transform_path_non_affine(self, path):
@@ -1632,10 +1628,9 @@ def transform_path_non_affine(self, path):
         ``transform_path(path)`` is equivalent to
         ``transform_path_affine(transform_path_non_affine(values))``.
         """
-        if self.input_dims != 2 or self.output_dims != 2:
-            raise NotImplementedError('Only defined in 2D')
         x = self.transform_non_affine(path.vertices)
-        return Path._fast_from_codes_and_verts(x, path.codes, path)
+        return Path._fast_from_codes_and_verts(x, path.codes, path,
+                                               dims=self.output_dims)
 
     def transform_angles(self, angles, pts, radians=False, pushoff=1e-5):
         """
@@ -1817,7 +1812,7 @@ def transform_path(self, path):
     def transform_path_affine(self, path):
         # docstring inherited
         return Path(self.transform_affine(path.vertices),
-                    path.codes, path._interpolation_steps)
+                    path.codes, path._interpolation_steps, dims=self.output_dims)
 
     def transform_path_non_affine(self, path):
         # docstring inherited
