From 41cf96cbcb7022ba7c41f7a32047555c7253a3fe Mon Sep 17 00:00:00 2001 From: Hanjin Liu Date: Sun, 12 May 2024 16:58:54 +0900 Subject: [PATCH] fix matplotlib vectors 3d --- .../matplotlib/components3d/vectors3d.py | 20 +++++++++++---- whitecanvas/backend/mock/__init__.py | 1 + whitecanvas/backend/mock/layers.py | 25 +++++++++++++++++++ whitecanvas/canvas/canvas3d/_base.py | 23 +++++++++++++++++ whitecanvas/plot/_methods.py | 2 ++ 5 files changed, 66 insertions(+), 5 deletions(-) diff --git a/whitecanvas/backend/matplotlib/components3d/vectors3d.py b/whitecanvas/backend/matplotlib/components3d/vectors3d.py index f3bdee2c..53ce7af1 100644 --- a/whitecanvas/backend/matplotlib/components3d/vectors3d.py +++ b/whitecanvas/backend/matplotlib/components3d/vectors3d.py @@ -27,11 +27,21 @@ def __init__(self, x0, dx, y0, dy, z0, dz): self._3d_w = dz def do_3d_projection(self, renderer=None): - x, y, z = proj3d.proj_transform(self._3d_x, self._3d_y, self._3d_z, self.axes.M) - u, v, w = proj3d.proj_transform(self._3d_u, self._3d_v, self._3d_w, self.axes.M) - self.set_UVC(u, v) - self.set_offsets(np.column_stack([x, y])) - return np.min(z) + x0, y0, z0 = proj3d.proj_transform( + self._3d_x, + self._3d_y, + self._3d_z, + self.axes.M, + ) + x1, y1, z1 = proj3d.proj_transform( + self._3d_x + self._3d_u, + self._3d_y + self._3d_v, + self._3d_z + self._3d_w, + self.axes.M, + ) + self.set_UVC(x1 - x0, y1 - y0) + self.set_offsets(np.column_stack([x0, y0])) + return np.min(z0) def post_add(self, canvas: Canvas3D): self.transform = canvas._axes.transData diff --git a/whitecanvas/backend/mock/__init__.py b/whitecanvas/backend/mock/__init__.py index afe7e837..d7322c6e 100644 --- a/whitecanvas/backend/mock/__init__.py +++ b/whitecanvas/backend/mock/__init__.py @@ -7,4 +7,5 @@ MonoLine, MultiLine, Texts, + Vectors, ) diff --git a/whitecanvas/backend/mock/layers.py b/whitecanvas/backend/mock/layers.py index 7dc8efcd..2674e20f 100644 --- a/whitecanvas/backend/mock/layers.py +++ b/whitecanvas/backend/mock/layers.py @@ -2,6 +2,7 @@ import numpy as np from cmap import Colormap +from numpy.typing import NDArray from whitecanvas import protocols from whitecanvas.backend.mock._base import ( @@ -214,3 +215,27 @@ def _plt_get_text_fontfamily(self) -> str: def _plt_set_text_fontfamily(self, family: str): self._fontfamily = family + + +@protocols.check_protocol(protocols.VectorsProtocol) +class Vectors(MockHasData, MockHasEdges): + def __init__(self, x0, dx, y0, dy): + super().__init__((x0, dx, y0, dy)) + + def _plt_get_ndata(self) -> int: + return self._plt_get_data()[0].size + + def _plt_get_edge_color(self) -> NDArray[np.float32]: + if not hasattr(self, "_edge_color"): + ndata = self._plt_get_ndata() + self._edge_color = np.zeros((ndata, 4), dtype=np.float32) + return self._edge_color + + def _plt_set_edge_color(self, color: NDArray[np.float32]): + self._edge_color = as_color_array(color, self._plt_get_ndata()) + + def _plt_get_antialias(self) -> bool: + return self._antialias + + def _plt_set_antialias(self, antialias: bool): + self._antialias = antialias diff --git a/whitecanvas/canvas/canvas3d/_base.py b/whitecanvas/canvas/canvas3d/_base.py index d1807028..c1e0eb96 100644 --- a/whitecanvas/canvas/canvas3d/_base.py +++ b/whitecanvas/canvas/canvas3d/_base.py @@ -133,6 +133,7 @@ def native(self) -> Any: @property def aspect_locked(self) -> bool: + """Whether the aspect ratio is locked.""" return self._canvas()._plt_get_aspect_locked() @aspect_locked.setter @@ -440,6 +441,28 @@ def add_vectors( alpha: float = 1.0, antialias: bool = True, ) -> layer3d.Vectors3D: + """ + Add a 3D vector field. + + Parameters + ---------- + x, y, z : array-like + Base points of the vectors. + vx, vy, vz : array-like + Components of the vectors. + name : str, optional + Name of the layer. + color : color-like, optional + Color of the vectors. + width : float, optional + Width of the vectors. + style : str or LineStyle, optional + Style of the vectors. + alpha : float, default 1.0 + Alpha channel of the vectors. + antialias : bool, default True + Whether to use antialiasing. + """ name = self._coerce_name(name) color = self._generate_colors(color) width = theme._default("line.width", width) diff --git a/whitecanvas/plot/_methods.py b/whitecanvas/plot/_methods.py index e2edbd03..19c3a64d 100644 --- a/whitecanvas/plot/_methods.py +++ b/whitecanvas/plot/_methods.py @@ -41,6 +41,7 @@ def _inner(*args, **kwargs): kde = _CANVAS.add_kde rug = _CANVAS.add_rug text = _CANVAS.add_text + vectors = _CANVAS.add_vectors # categorical methods cat = _CANVAS.cat cat_x = _CANVAS.cat_x @@ -67,6 +68,7 @@ def _inner(*args, **kwargs): kde = _make_method("kde") rug = _make_method("rug") text = _make_method("text") + vectors = _make_method("vectors") cat = _make_method("cat", pref="") cat_x = _make_method("cat_x", pref="") cat_y = _make_method("cat_y", pref="")