fastplotlib · kushalkolar · Aug 2, 2024 · Jul 29, 2024 · Jul 29, 2024 · Aug 2, 2024
@@ -1,10 +1,13 @@
 import numpy as np
 
 import pygfx
+from pylinalg import quat_from_vecs, vec_transform_quat
 
 
 GRID_PLANES = ["xy", "xz", "yz"]
 
+CANONICAL_BAIS = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+
 
 # very thin subclass that just adds GridMaterial properties to this world object for easier user control
 class Grid(pygfx.Grid):
@@ -252,6 +255,9 @@ def __init__(
         grid_kwargs: dict = None,
         auto_grid: bool = True,
         offset: np.ndarray = np.array([0.0, 0.0, 0.0]),
+        basis: np.ndarray = np.array(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
+        ),
     ):
         self._plot_area = plot_area
 
@@ -362,10 +368,28 @@ def __init__(
         self._intersection = intersection
         self._auto_grid = auto_grid
 
+        self._basis = None
+        self.basis = basis
+
     @property
     def world_object(self) -> pygfx.WorldObject:
         return self._world_object
 
+    @property
+    def basis(self) -> np.ndarray:
+        """get or set the basis, shape is [3, 3]"""
+        return self._basis
+
+    @basis.setter
+    def basis(self, basis: np.ndarray):
+        if basis.shape != (3, 3):
+            raise ValueError
+
+        # apply quaternion to each of x, y, z rulers
+        for dim, cbasis, new_basis in zip(["x", "y", "z"], CANONICAL_BAIS, basis):
+            ruler: pygfx.Ruler = getattr(self, dim)
+            ruler.local.rotation = quat_from_vecs(cbasis, new_basis)
+
     @property
     def offset(self) -> np.ndarray:
         """offset of the axes"""
@@ -395,6 +419,19 @@ def grids(self) -> Grids | bool:
         """grids for each plane: xy, xz, yz"""
         return self._grids
 
+    @property
+    def colors(self) -> tuple[pygfx.Color]:
+        return tuple(getattr(self, dim).line.material.color for dim in ["x", "y", "z"])
+
+    @colors.setter
+    def colors(self, colors: tuple[pygfx.Color | str]):
+        """get or set the colors for the x, y, and z rulers"""
+        if len(colors) != 3:
+            raise ValueError
+
+        for dim, color in zip(["x", "y", "z"], colors):
+            getattr(self, dim).line.material.color = color
+
     @property
     def auto_grid(self) -> bool:
         """auto adjust the grid on each render cycle"""
@@ -482,10 +519,31 @@ def update_using_camera(self):
             xpos, ypos, width, height = self._plot_area.get_rect()
             # orthographic projection, get ranges using inverse
 
-            # get range of screen space
+            # get range of screen space by getting the corners
             xmin, xmax = xpos, xpos + width
             ymin, ymax = ypos + height, ypos
 
+            # apply quaternion to account for rotation of axes
+            # xmin, _, _ = vec_transform_quat(
+            #     [xmin, ypos + height / 2, 0],
+            #     self.x.local.rotation
+            # )
+            #
+            # xmax, _, _ = vec_transform_quat(
+            #     [xmax, ypos + height / 2, 0],
+            #     self.x.local.rotation,
+            # )
+            #
+            # _, ymin, _ = vec_transform_quat(
+            #     [xpos + width / 2, ymin, 0],
+            #     self.y.local.rotation
+            # )
+            #
+            # _, ymax, _ = vec_transform_quat(
+            #     [xpos + width / 2, ymax, 0],
+            #     self.y.local.rotation
+            # )
+
             min_vals = self._plot_area.map_screen_to_world((xmin, ymin))
             max_vals = self._plot_area.map_screen_to_world((xmax, ymax))
 
@@ -578,7 +636,7 @@ def update(self, bbox, intersection):
             self.z.start_pos = world_x_10, world_y_10, world_zmin
             self.z.end_pos = world_x_10, world_y_10, world_zmax
 
-            self.z.start_value = self.z.start_pos[1] - self.offset[2]
+            self.z.start_value = self.z.start_pos[2] - self.offset[2]
             statsz = self.z.update(
                 self._plot_area.camera, self._plot_area.viewport.logical_size
             )