import math
import numpy as np
import pygfx
from pylinalg import quat_from_vecs, vec_transform_quat
from ..utils.enums import RenderQueue
from ..utils import global_config
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
[docs]
class Grid(pygfx.Grid):
@property
def major_step(self) -> tuple[float, float]:
"""The step distance between the major grid lines."""
return self.material.major_step
@major_step.setter
def major_step(self, step: tuple[float, float]):
self.material.major_step = step
@property
def minor_step(self) -> tuple[float, float]:
"""The step distance between the minor grid lines."""
return self.material.minor_step
@minor_step.setter
def minor_step(self, step: tuple[float, float]):
self.material.minor_step = step
@property
def axis_thickness(self) -> float:
"""The thickness of the axis lines."""
return self.material.axis_thickness
@axis_thickness.setter
def axis_thickness(self, thickness: float):
self.material.axis_thickness = thickness
@property
def major_thickness(self) -> float:
"""The thickness of the major grid lines."""
return self.material.major_thickness
@major_thickness.setter
def major_thickness(self, thickness: float):
self.material.major_thickness = thickness
@property
def minor_thickness(self) -> float:
"""The thickness of the minor grid lines."""
return self.material.minor_thickness
@minor_thickness.setter
def minor_thickness(self, thickness: float):
self.material.minor_thickness = thickness
@property
def thickness_space(self) -> str:
"""The coordinate space in which the thicknesses are expressed.
See :obj:`pygfx.utils.enums.CoordSpace`:
"""
return self.material.thickness_space
@thickness_space.setter
def thickness_space(self, value: str):
self.material.thickness_space = value
@property
def axis_color(self) -> str:
"""The color of the axis lines."""
return self.material.axis_color
@axis_color.setter
def axis_color(self, color: str):
self.material.axis_color = color
@property
def major_color(self) -> str:
"""The color of the major grid lines."""
return self.material.major_color
@major_color.setter
def major_color(self, color: str):
self.material.major_color = color
@property
def minor_color(self) -> str:
"""The color of the minor grid lines."""
return self.material.minor_color
@minor_color.setter
def minor_color(self, color: str):
self.material.minor_color = color
@property
def infinite(self) -> bool:
"""Whether the grid is infinite.
If not infinite, the grid is 1x1 in world space, scaled, rotated, and
positioned with the object's transform.
(Infinite grids are not actually infinite. Rather they move along with
the camera, and are sized based on the distance between the camera and
the grid.)
"""
return self.material.infinite
@infinite.setter
def infinite(self, value: str):
self.material.infinite = value
[docs]
class Grids(pygfx.Group):
"""Just a class to make accessing the grids easier"""
def __init__(self, *, xy, xz, yz):
super().__init__()
self._xy = xy
self._xz = xz
self._yz = yz
self.add(xy, xz, yz)
@property
def xy(self) -> Grid:
"""xy grid"""
return self._xy
@property
def xz(self) -> Grid:
"""xz grid"""
return self._xz
@property
def yz(self) -> Grid:
"""yz grid"""
return self._yz
[docs]
class Ruler(pygfx.Ruler):
"""pygfx.Ruler subclass that adds a rotated axis label."""
def __init__(self, *, color="#fff", alpha_mode=None, render_queue=None, **kwargs):
super().__init__(
color=color, alpha_mode=alpha_mode, render_queue=render_queue, **kwargs
)
self._label = pygfx.Text(
screen_space=True,
anchor="middle-center",
font_size=20,
material=pygfx.TextMaterial(
color=color,
alpha_mode="auto",
render_queue=RenderQueue.overlay + 50,
aa=True,
),
)
self._label.visible = False
self.add(self._label)
self.text.material.outline_thickness = 0.5
@property
def label(self) -> pygfx.Text:
"""Axis label. Set text via ``label.set_text('label text')``"""
return self._label
@property
def color(self):
return self._text.material.color
@color.setter
def color(self, color):
self._text.material.color = color
self._line.material.color = color
self._points.material.edge_color = color
self._label.material.color = color
[docs]
def update(self, camera, canvas_size):
stats = super().update(camera, canvas_size)
self._update_label()
return stats
def _update_label(self):
# update the label position
t1, t2 = self._visible_part_coords
if t1 == t2:
self._label.visible = False
return
self._label.visible = True
mid_t = 0.5 * (t1 + t2)
self._label.local.position = (
self._start_pos * (1 - mid_t) + self._end_pos * mid_t
)
vec = self._visible_part_screen_vec
angle = math.atan2(vec[1], vec[0])
# the side of the line that the tick labels are on, as a screen space unit vector. this is
# the same rule that pygfx uses to anchor the tick labels themselves, and the screen vector
# already carries the camera scale, the viewport aspect and the ruler's orientation, so
# none of those need a case of their own.
if self.tick_side == "left":
px, py = -math.sin(angle), math.cos(angle)
else:
px, py = math.sin(angle), -math.cos(angle)
# a ruler that runs right to left, or top to bottom, on screen would render the label
# upside down, so turn it around. that turns the label's own axes around with it
upside_down = not (-0.5 * math.pi < angle <= 0.5 * math.pi)
if upside_down:
angle -= math.copysign(math.pi, angle)
# pylinalg uses [x, y, z, w] quaternion format
self._label.local.rotation = np.array(
[0.0, 0.0, math.sin(angle / 2), math.cos(angle / 2)]
)
# the label is rotated onto the line, so in its own frame the line runs along x and the
# ticks sit on one side of it, +y or -y. anchoring it to that side offsets it in screen
# pixels, which is what keeps the camera scale and the viewport aspect out of the placement
if (self.tick_side == "left") != upside_down:
anchor = "bottom-center"
else:
anchor = "top-center"
# max extent of the tick labels in that same perpendicular direction.
# tick labels are unrotated screen-space text, so we project their
# axis-aligned _rect onto (px, py) directly.
px_pos, px_neg = max(px, 0), min(px, 0)
py_pos, py_neg = max(py, 0), min(py, 0)
tick_extent_px = max(
(
px_pos * b._rect.right
+ px_neg * b._rect.left
+ py_pos * b._rect.top
+ py_neg * b._rect.bottom
for b in self.text._text_blocks
if b._rect.width > 0 or b._rect.height > 0
),
default=0.0,
)
# gap between the tick labels and the label. a text rect is tight on its left and right,
# but its top and bottom are the font's ascender and descender, which neither a tick
# number nor most labels reach. that padding already separates the two where the offset is
# vertical, so only add a gap to the extent that the offset is horizontal
gap_px = abs(px) * 0.5 * self._label.font_size
anchor_offset = max(tick_extent_px, 0.0) + gap_px
# both of these re-run the text layout, so only set them when they actually change
if self._label._anchor != anchor:
self._label.anchor = anchor
if self._label._anchor_offset != anchor_offset:
self._label.anchor_offset = anchor_offset
[docs]
@global_config.register
class Axes:
config = global_config.descriptor
@global_config.declare(
"intersection",
"tick_size",
"line_width",
"tick_marker",
"color",
"grids",
"grid_kwargs",
"auto_grid",
)
def __init__(
self,
plot_area,
intersection: tuple[int, int, int] | None = None,
tick_size: float = 8.0,
line_width: float = 2.0,
tick_marker: str = "tick",
color: str = "#fff",
x_kwargs: dict = None,
y_kwargs: dict = None,
z_kwargs: dict = None,
grids: bool = True,
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
x_kwargs = x_kwargs or {}
y_kwargs = y_kwargs or {}
z_kwargs = z_kwargs or {}
generic_kwargs = dict(
tick_size=tick_size,
line_width=line_width,
tick_marker=tick_marker, # 'tick' for both-sides, 'tick_left' or 'tick_right' for one-sided
color=color,
)
x_kwargs = dict(
tick_side="right",
**generic_kwargs,
**x_kwargs,
)
y_kwargs = dict(
tick_side="left",
**generic_kwargs,
**y_kwargs,
)
z_kwargs = dict(
tick_side="left",
**generic_kwargs,
**z_kwargs,
)
# create ruler for each dim
self._x = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **x_kwargs)
self._y = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **y_kwargs)
self._z = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **z_kwargs)
# We render the lines and ticks as solid, but enable aa for text for prettier glyphs
for ruler in self._x, self._y, self._z:
ruler.line.material.depth_compare = "<="
ruler.points.material.depth_compare = "<="
ruler.text.material.depth_compare = "<="
ruler.text.material.alpha_mode = "auto"
ruler.text.material.aa = True
ruler.label.material.depth_compare = "<="
self._offset = offset
# *MUST* instantiate some start and end positions for the rulers else kernel crashes immediately
# probably a WGPU rust panic
self.x.start_pos = 0, 0, 0
self.x.end_pos = 100, 0, 0
self.x.start_value = self.x.start_pos[0] - offset[0]
statsx = self.x.update(
self._plot_area.camera, self._plot_area.viewport.logical_size
)
self.y.start_pos = 0, 0, 0
self.y.end_pos = 0, 100, 0
self.y.start_value = self.y.start_pos[1] - offset[1]
statsy = self.y.update(
self._plot_area.camera, self._plot_area.viewport.logical_size
)
self.z.start_pos = 0, 0, 0
self.z.end_pos = 0, 0, 100
self.z.start_value = self.z.start_pos[1] - offset[2]
self.z.update(self._plot_area.camera, self._plot_area.viewport.logical_size)
# world object for the rulers + grids
self._world_object = pygfx.Group()
# add rulers
self.world_object.add(
self.x,
self.y,
self.z,
)
# set z ruler invisible for orthographic projections for now
if self._plot_area.camera.fov == 0:
# TODO: allow any orientation in the future even for orthographic projections
self.z.visible = False
if grid_kwargs is None:
grid_kwargs = dict()
# The grid is a bit weird, because it makes use of transparency to fade off in the distance.
# But w want it to write depth, so that objects that are drawn behind it are partually hidden.
# So we set alha_mode to 'auto'. We make it draw earlier than other 'auto' objects, under the
# assumption that most interesting stuff is in front of the grid, and artifacts behind the grid are less
# bad than those in front. Note that fully opaque objects blend perfectly fine with the grid. Artifacts
# should only emerge for objects that have semi-transparent fragments.
grid_kwargs = dict(
alpha_mode="auto",
render_queue=RenderQueue.auto + 50,
major_step=10,
minor_step=1,
thickness_space="screen",
major_thickness=2,
minor_thickness=0.5,
infinite=True,
**grid_kwargs,
)
if grids:
_grids = dict()
for plane in GRID_PLANES:
grid = Grid(
geometry=None,
material=pygfx.GridMaterial(**grid_kwargs),
orientation=plane,
visible=False,
)
_grids[plane] = grid
self._grids = Grids(**_grids)
self.world_object.add(self._grids)
if self._plot_area.camera.fov == 0:
# orthographic projection, place grids far away
self._grids.local.z = -1000
major_step_x, major_step_y = statsx["tick_step"], statsy["tick_step"]
self.grids.xy.material.major_step = major_step_x, major_step_y
self.grids.xy.material.minor_step = 0.2 * major_step_x, 0.2 * major_step_y
else:
self._grids = False
self._intersection = intersection
self._auto_grid = auto_grid
self._basis = None
self.basis = basis
self._last_state = self._get_view_state()
@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: Ruler = getattr(self, dim)
ruler.local.rotation = quat_from_vecs(cbasis, new_basis)
@property
def offset(self) -> np.ndarray:
"""offset of the axes"""
return self._offset
@offset.setter
def offset(self, value: np.ndarray):
self._offset = value
@property
def x(self) -> Ruler:
"""x axis ruler"""
return self._x
@property
def y(self) -> Ruler:
"""y axis ruler"""
return self._y
@property
def z(self) -> Ruler:
"""z axis ruler"""
return self._z
@property
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 color(self) -> pygfx.Color:
"""get or set a single color for all rulers"""
return self._x.color
@color.setter
def color(self, color: pygfx.Color | str):
for ruler in (self._x, self._y, self._z):
ruler.color = color
@property
def auto_grid(self) -> bool:
"""auto adjust the grid on each render cycle"""
return self._auto_grid
@auto_grid.setter
def auto_grid(self, value: bool):
self._auto_grid = value
@property
def visible(self) -> bool:
"""set visibility of all axes elements, rulers and grids"""
return self._world_object.visible
@visible.setter
def visible(self, value: bool):
self._world_object.visible = value
@property
def intersection(self) -> tuple[float, float, float] | None:
return self._intersection
@intersection.setter
def intersection(self, intersection: tuple[float, float, float] | None):
"""
intersection point of [x, y, z] rulers.
Set (0, 0, 0) for origin
Set to `None` to follow when panning through the scene with orthographic projection
"""
if intersection is None:
self._intersection = None
return
if len(intersection) != 3:
raise ValueError(
"intersection must be a float of 3 elements for [x, y, z] or `None`"
)
self._intersection = tuple(float(v) for v in intersection)
def _get_view_state(self) -> tuple:
viewport = self._plot_area.viewport
cam_matrix = self._plot_area.camera.camera_matrix.tobytes()
scale = self._plot_area.camera.local.scale.tobytes()
# the label margins are the other half of what places the rulers, and they are not known
# until the text has been laid out, which only happens once it has been drawn. tracking
# them here is what redoes the placement on the frame after that, and on any later change
# in the width of a tick label
return (
cam_matrix,
viewport.rect,
viewport.logical_size,
scale,
self._get_label_margins(),
)
[docs]
def update_using_bbox(self, bbox):
"""
Update the axes w.r.t. the given bbox
Parameters
----------
bbox: np.ndarray
array of shape [2, 3], [[xmin, ymin, zmin], [xmax, ymax, zmax]]
"""
# flip axes if camera scale is flipped
if self._plot_area.camera.local.scale_x < 0:
bbox[0, 0], bbox[1, 0] = bbox[1, 0], bbox[0, 0]
if self._plot_area.camera.local.scale_y < 0:
bbox[0, 1], bbox[1, 1] = bbox[1, 1], bbox[0, 1]
if self._plot_area.camera.local.scale_z < 0:
bbox[0, 2], bbox[1, 2] = bbox[1, 2], bbox[0, 2]
if self.intersection is None:
intersection = (0, 0, 0)
else:
intersection = self.intersection
self.update(bbox, intersection)
def _get_label_margins(self) -> tuple[float, float]:
"""
How far the x and y tick labels, plus their axis labels, reach from their ruler, in pixels
The fallbacks are for text that has not been laid out yet, which is the case until it has
been drawn once.
"""
x_blocks = [b for b in self.x.text._text_blocks if b._rect.height > 0]
x_margin = (
max(abs(b._rect.bottom) for b in x_blocks)
if x_blocks
else 1.5 * self.x.text.font_size
)
if self.x._label._text_blocks:
# the axis label starts at the tick margin, and its own body follows
x_margin += 1.5 * self.x._label.font_size
y_blocks = [b for b in self.y.text._text_blocks if b._rect.width > 0]
y_margin = (
max(abs(b._rect.left) for b in y_blocks)
if y_blocks
else 6 * self.y.text.font_size
)
if self.y._label._text_blocks:
y_margin += 1.5 * self.y._label.font_size
return x_margin, y_margin
[docs]
def update_using_camera(self):
"""
Update the axes w.r.t the current camera state
For orthographic projections of the xy plane, it will calculate the inverse projection
of the screen space onto world space to determine the current range of the world space
to set the rulers and ticks
For perspective projections it will just use the bbox of the scene to set the rulers
"""
if not self.visible:
return
state = self._get_view_state()
if state == self._last_state:
# no changes in the camera, the viewport rect, or the size of the labels
return
*_, (x_margin, y_margin) = state
if self._plot_area.camera.fov == 0:
xpos, ypos, width, height = self._plot_area.viewport.rect
# orthographic projection, get ranges using inverse
# get range of screen space by getting the corners
xmin, xmax = xpos, xpos + width
ymin, ymax = ypos + height, ypos
min_vals = self._plot_area.map_screen_to_world((xmin, ymin))
max_vals = self._plot_area.map_screen_to_world((xmax, ymax))
if min_vals is None or max_vals is None:
return
world_xmin, world_ymin, _ = min_vals
world_xmax, world_ymax, _ = max_vals
world_zmin, world_zmax = 0, 0
bbox = np.array(
[
[world_xmin, world_ymin, world_zmin],
[world_xmax, world_ymax, world_zmax],
]
)
else:
# set ruler start and end positions based on scene bbox
bbox = self._plot_area._fpl_graphics_scene.get_world_bounding_box()
if self.intersection is None:
if self._plot_area.camera.fov == 0:
# put the rulers in the bottom left corner, clear of their own labels
padding = 4
intersection = self._plot_area.map_screen_to_world(
(xpos + y_margin + padding, ypos + height - x_margin - padding)
)
else:
# force origin since None is not supported for Persepctive projections
self._intersection = (0, 0, 0)
intersection = self._intersection
else:
# axes intersect at the origin
intersection = self.intersection
self.update(bbox, intersection)
self._last_state = state
[docs]
def update(self, bbox, intersection):
"""
Update the axes using the given bbox and ruler intersection point
Parameters
----------
bbox: np.ndarray
array of shape [2, 3], [[xmin, ymin, zmin], [xmax, ymax, zmax]]
intersection: float, float, float
intersection point of the x, y, z ruler
"""
world_xmin, world_ymin, world_zmin = bbox[0]
world_xmax, world_ymax, world_zmax = bbox[1]
world_x_10, world_y_10, world_z_10 = intersection
# swap min and max for each dimension if necessary
if self._plot_area.camera.local.scale_y < 0:
world_ymin, world_ymax = world_ymax, world_ymin
self.y.tick_side = "right" # swap tick side
self.x.tick_side = "right"
else:
self.y.tick_side = "left"
self.x.tick_side = "right"
if self._plot_area.camera.local.scale_x < 0:
world_xmin, world_xmax = world_xmax, world_xmin
self.x.tick_side = "left"
self.x.start_pos = world_xmin, world_y_10, world_z_10
self.x.end_pos = world_xmax, world_y_10, world_z_10
self.x.start_value = self.x.start_pos[0] - self.offset[0]
statsx = self.x.update(
self._plot_area.camera, self._plot_area.viewport.logical_size
)
self.y.start_pos = world_x_10, world_ymin, world_z_10
self.y.end_pos = world_x_10, world_ymax, world_z_10
self.y.start_value = self.y.start_pos[1] - self.offset[1]
statsy = self.y.update(
self._plot_area.camera, self._plot_area.viewport.logical_size
)
if self._plot_area.camera.fov != 0:
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[2] - self.offset[2]
statsz = self.z.update(
self._plot_area.camera, self._plot_area.viewport.logical_size
)
major_step_z = statsz["tick_step"]
if self.grids:
if self.auto_grid:
major_step_x, major_step_y = statsx["tick_step"], statsy["tick_step"]
self.grids.xy.major_step = major_step_x, major_step_y
self.grids.xy.minor_step = 0.2 * major_step_x, 0.2 * major_step_y
if self._plot_area.camera.fov != 0:
self.grids.xz.major_step = major_step_x, major_step_z
self.grids.xz.minor_step = 0.2 * major_step_x, 0.2 * major_step_z
self.grids.yz.material.major_step = major_step_y, major_step_z
self.grids.yz.minor_step = 0.2 * major_step_y, 0.2 * major_step_z
