Repository: pygfx/pygfx
Pull request #916 (2024-12-29 21:33:07)

Base AffineTransform state on position/scale/rotation instead of matrix
Changes:

* Base AffineTransform state on position/scale/rotation instead of matrix
* All ndarrays are now read only
* Reuse all local state ndarrays instead of constantly replacing them
* Remove performance workarounds to allow for performance testing

This yields really impressive performance increases (roughly 300-400%) for examples where transforms are updated every frame.
And we still have all our world space transform getters and setters with lazy evaluation & caching.
No need to call "update()" functions manually in pygfx (like in Three.js).
Best of both worlds! I'm very pleased with this outcome.

Future work (separate PR):
* I could add an option to AffineTransform to toggle between the matrix or position/scale/rotation as the source of truth. There are scenarios where that is preferred. See #920

Korijn: @panxinmiao with these changes I am getting ~130 FPS on my machine for the Skinned Mesh example:

![image](https://github.com/user-attachments/assets/b94ca3fa-a931-4b50-8f50-59a571bed2c8)

Is this "good enough" in terms of performance? :)
Korijn: I see I've managed to break something related to axis flips:

![image](https://github.com/user-attachments/assets/f6c578cf-1e61-4a32-9bcb-49a8e9204a3d)

Korijn: This is ready to merge from my end. So happy to reach this state, we now really have the public API that I initially envisioned when we started working on the new Transform API.

Patch:
diff --git a/examples/feature_demo/skinned_mesh.py b/examples/feature_demo/skinned_mesh.py
index 1041515..186eb79 100644
--- a/examples/feature_demo/skinned_mesh.py
+++ b/examples/feature_demo/skinned_mesh.py
@@ -67,8 +67,6 @@ def create_mesh(geometry, bones):
     skeleton = gfx.Skeleton(bones)
 
     mesh.add(bones[0])
-    # update the world matrix from root to leaf bone manually
-    bones[0].update_matrix_world()
     mesh.bind(skeleton)
 
     mesh.local.rotation = la.quat_from_euler((-math.pi / 2, 0, 0))
diff --git a/pygfx/objects/_more.py b/pygfx/objects/_more.py
index 17d76c8..7fb2667 100644
--- a/pygfx/objects/_more.py
+++ b/pygfx/objects/_more.py
@@ -4,7 +4,7 @@ import numpy as np
 from ._base import WorldObject
 from ..resources import Buffer
 from ..utils import unpack_bitfield, array_from_shadertype
-from ..utils.transform import AffineBase, callback
+from ..utils.transform import AffineBase
 from ..materials import BackgroundMaterial
 
 
@@ -434,12 +434,8 @@ class Text(WorldObject):
             name=name,
         )
 
-        # calling super from callback is possible, but slow so we register it as a second callback instead
-        self.world.on_update(super()._update_uniform_buffers)
-
-    @callback
     def _update_uniform_buffers(self, transform: AffineBase):
-        # super()._update_uniform_buffers(transform)
+        super()._update_uniform_buffers(transform)
         # When rendering in screen space, the world transform is used
         # to establish the point in the scene where the text is placed.
         # The only part of the local transform that is used is the
diff --git a/pygfx/objects/_skins.py b/pygfx/objects/_skins.py
index 608be60..b402de7 100644
--- a/pygfx/objects/_skins.py
+++ b/pygfx/objects/_skins.py
@@ -17,81 +17,10 @@ class Bone(WorldObject):
 
     """
 
-    class _Transform:
-        position: np.ndarray
-        rotation: np.ndarray
-        scale: np.ndarray
-        matrix: np.ndarray
-
     def __init__(self, name=""):
         super().__init__(name=name)
         self.visible = False
 
-        self._children: "List[Bone]" = []
-        self._parent: "Bone" = None
-
-        self.__matrix_world_needs_update = False
-
-        # compatible with cuurrent WorldObject RecursiveTransform system, but does nothing.
-        # temporary solution before we refactor the WorldObject transform system.
-        # See: https://github.com/pygfx/pygfx/pull/715#issuecomment-2053385803
-        self.world = Bone._Transform()
-        self.world.matrix = np.eye(4)
-        self.local = Bone._Transform()
-        self.local.position = np.zeros(3)
-        self.local.rotation = np.zeros(4)
-        self.local.scale = np.ones(3)
-        self.local.matrix = np.eye(4)
-
-    @property
-    def parent(self):
-        if self._parent is None:
-            return None
-        if isinstance(self._parent, Bone):
-            return self._parent
-        else:
-            return self._parent()
-
-    def update_matrix(self):
-        self.local.matrix = la.mat_compose(
-            self.local.position, self.local.rotation, self.local.scale
-        )
-        self.__matrix_world_needs_update = True
-
-    def update_matrix_world(self):
-        self.update_matrix()
-
-        if self.__matrix_world_needs_update:
-            if self.parent is not None:
-                self.world.matrix = self.parent.world.matrix @ self.local.matrix
-            else:
-                self.world.matrix = self.local.matrix
-            self.__matrix_world_needs_update = False
-
-        for child in self._children:
-            child.update_matrix_world()
-
-    def add(self, *bones: "Bone") -> "Bone":
-        for obj in bones:
-            if obj == self:
-                # can't add self as a child
-                continue
-
-            if obj and isinstance(obj, Bone):
-                if obj._parent is not None:
-                    obj._parent.remove(obj)
-
-                obj._parent = self
-                self._children.append(obj)
-            else:
-                pass
-                # Now the Bone class has specific logic, so we just pass if it's not a Bone
-
-        return self
-
-    def __repr__(self) -> str:
-        return f"Bone {self.name} {self.local.position} {self.local.rotation}\n"
-
 
 class Skeleton:
     """A skeleton.
@@ -173,13 +102,6 @@ class Skeleton:
         # See: https://github.com/pygfx/pygfx/pull/715#issuecomment-2046493145
         """Update the bone matrices buffer."""
 
-        # TODO: update bone matrices from root to leaf, it's a temporary solution.
-        # See: https://github.com/pygfx/pygfx/pull/715#issuecomment-2053385803
-        for bone in self.bones:
-            if bone.parent and isinstance(bone.parent, Bone):
-                continue
-            bone.update_matrix_world()
-
         for i, bone in enumerate(self.bones):
             self.bone_matrices_buffer.data[i]["bone_matrices"] = (
                 bone.world.matrix @ self.bone_inverses[i]
diff --git a/pygfx/utils/transform.py b/pygfx/utils/transform.py
index ac78160..f677500 100644
--- a/pygfx/utils/transform.py
+++ b/pygfx/utils/transform.py
@@ -4,7 +4,7 @@ from time import perf_counter_ns
 import weakref
 import functools
 
-from typing import Tuple, Union
+from typing import Tuple
 
 
 PRECISION_EPSILON = 1e-7
@@ -141,13 +141,19 @@ class AffineBase:
 
     last_modified: int
 
-    def __init__(self, *, reference_up=(0, 1, 0), is_camera_space=False):
+    def __init__(self, /, *, reference_up=(0, 1, 0), is_camera_space=False):
         self.update_callbacks = {}
+
         self.is_camera_space = int(is_camera_space)
-        self._scaling_signs = np.asarray([1, 1, 1], dtype=float)
 
         self._reference_up_provider = None
-        self._reference_up = la.vec_normalize(reference_up)
+        self._reference_up = la.vec_normalize(reference_up, dtype=float)
+        self._reference_up_view = self._reference_up.view()
+        self._reference_up_view.flags.writeable = False
+
+        self._scaling_signs = np.asarray([1, 1, 1], dtype=float)
+        self._scaling_signs_view = self._scaling_signs.view()
+        self._scaling_signs_view.flags.writeable = False
 
     def _set_reference_up_provider(self, provider: "RecursiveTransform"):
         self._reference_up_provider = provider
@@ -157,14 +163,14 @@ class AffineBase:
         """The zero-tilt reference vector used for the direction setters."""
         if self._reference_up_provider:
             return self._reference_up_provider._own_reference_up
-        return self._reference_up
+        return self._reference_up_view
 
     @reference_up.setter
     def reference_up(self, value):
         if self._reference_up_provider:
             self._reference_up_provider._own_reference_up = value
         else:
-            self._reference_up = la.vec_normalize(value)
+            self._reference_up[:] = la.vec_normalize(value)
 
     @property
     def matrix(self):
@@ -177,11 +183,13 @@ class AffineBase:
 
     @cached
     def _inverse_matrix(self) -> np.ndarray:
-        return mat_inv(self.matrix)
+        mat = mat_inv(self.matrix)
+        mat.flags.writeable = False
+        return mat
 
     @property
     def scaling_signs(self):
-        return self._scaling_signs
+        return self._scaling_signs_view
 
     @cached
     def _decomposed(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -193,9 +201,6 @@ class AffineBase:
             # any decomposed scaling is acceptable
             decomposed = la.mat_decompose(self.matrix)
 
-        # Stop the user from accidentally writing the temporary arrays
-        # that we return from these operations
-        # https://github.com/pygfx/pygfx/issues/651
         for m in decomposed:
             m.flags.writeable = False
 
@@ -211,9 +216,6 @@ class AffineBase:
         origin_target = la.vec_transform((0, 0, 0), self.matrix)
 
         directions = axes_target - origin_target
-        # Stop the user from accidentally writing the temporary arrays
-        # that we return from these operations
-        # https://github.com/pygfx/pygfx/issues/651
         directions.flags.writeable = False
         return directions
 
@@ -224,20 +226,12 @@ class AffineBase:
     @cached
     def _rotation_matrix(self):
         rotation = la.mat_from_quat(self._decomposed[1])
-
-        # Stop the user from accidentally writing the temporary arrays
-        # that we return from these operations
-        # https://github.com/pygfx/pygfx/issues/651
         rotation.flags.writeable = False
         return rotation
 
     @cached
     def _euler(self):
         euler = la.quat_to_euler(self._decomposed[1])
-
-        # Stop the user from accidentally writing the temporary arrays
-        # that we return from these operations
-        # https://github.com/pygfx/pygfx/issues/651
         euler.flags.writeable = False
         return euler
 
@@ -411,7 +405,7 @@ class AffineBase:
     @scale.setter
     def scale(self, value):
         m = la.mat_compose(self.position, self.rotation, value)
-        self._scaling_signs = np.sign(value)
+        self._scaling_signs[:] = np.sign(value)
         self.matrix = m
 
     @x.setter
@@ -505,7 +499,7 @@ class AffineBase:
         self.rotation = rotation
 
     def __array__(self, dtype=None):
-        return self.matrix.astype(dtype)
+        return self.matrix.astype(dtype, copy=False)
 
 
 class AffineTransform(AffineBase):
@@ -513,9 +507,6 @@ class AffineTransform(AffineBase):
 
     Parameters
     ----------
-    matrix : ndarray, [4, 4]
-        The affine matrix used to back this transform. If None, a new diagonal
-        matrix will be initialized.
     position : ndarray, [3]
         The position of this transform expressed in the target frame. This will
         overwrite the position component of ``matrix`` if present.
@@ -558,61 +549,115 @@ class AffineTransform(AffineBase):
 
     def __init__(
         self,
-        matrix=None,
+        position=(0, 0, 0),
+        rotation=(0, 0, 0, 1),
+        scale=(1, 1, 1),
         /,
         *,
-        position=None,
-        rotation=None,
-        scale=None,
         reference_up=(0, 1, 0),
         is_camera_space=False,
     ) -> None:
         super().__init__(reference_up=reference_up, is_camera_space=is_camera_space)
         self.last_modified = perf_counter_ns()
 
-        if matrix is None:
-            matrix = np.eye(4, dtype=float)
-
-        self.untracked_matrix = np.asarray(matrix)
+        self._position = np.asarray(position, dtype=float)
+        self._rotation = np.asarray(rotation, dtype=float)
+        self._scale = np.asarray(scale, dtype=float)
 
-        if position is not None:
-            self.position = position
-
-        if rotation is not None:
-            self.rotation = rotation
-
-        if scale is not None:
-            self.scale = scale
+        self._position_view = self._position.view()
+        self._position_view.flags.writeable = False
+        self._rotation_view = self._rotation.view()
+        self._rotation_view.flags.writeable = False
+        self._scale_view = self._scale.view()
+        self._scale_view.flags.writeable = False
 
     def flag_update(self):
         self.last_modified = perf_counter_ns()
         super().flag_update()
 
+    @property
+    def position(self) -> np.ndarray:
+        return self._position_view
+
+    @position.setter
+    def position(self, value):
+        self._position[:] = value
+        self.flag_update()
+
+    @property
+    def rotation(self) -> np.ndarray:
+        return self._rotation_view
+
+    @rotation.setter
+    def rotation(self, value):
+        self._rotation[:] = value
+        self.flag_update()
+
+    @property
+    def scale(self) -> np.ndarray:
+        return self._scale_view
+
+    @scale.setter
+    def scale(self, value):
+        self._scale[:] = value
+        self.flag_update()
+
+    @cached
+    def __scaling_signs(self):
+        signs = np.sign(self._scale)
+        signs.flags.writeable = False
+        return signs
+
+    @property
+    def scaling_signs(self):
+        return self.__scaling_signs
+
+    @cached
+    def _rotation_matrix(self):
+        rotation = la.mat_from_quat(self._rotation)
+        rotation.flags.writeable = False
+        return rotation
+
+    @cached
+    def _euler(self):
+        euler = la.quat_to_euler(self._rotation)
+        euler.flags.writeable = False
+        return euler
+
+    @cached
+    def _matrix(self):
+        mat = la.mat_compose(self._position, self._rotation, self._scale)
+        mat.flags.writeable = False
+        return mat
+
     @property
     def matrix(self):
-        view_array = self.untracked_matrix.view()
-        # Stop the user from accidentally writing the temporary arrays
-        # that we return from these operations
-        # https://github.com/pygfx/pygfx/issues/651
-        view_array.flags.writeable = False
-        return view_array
+        return self._matrix
 
     @matrix.setter
     def matrix(self, value):
-        self.untracked_matrix[:] = value
+        try:
+            # when the matrix is set manually
+            # try to maintain the most recent configured scaling signs
+            position, rotation, scale = la.mat_decompose(
+                value, scaling_signs=self.scaling_signs
+            )
+        except ValueError:
+            position, rotation, scale = la.mat_decompose(value)
+        self._position[:] = position
+        self._rotation[:] = rotation
+        self._scale[:] = scale
         self.flag_update()
 
     def __matmul__(self, other):
         if isinstance(other, AffineTransform):
-            return AffineTransform(
-                self.matrix @ other.matrix, is_camera_space=self.is_camera_space
-            )
+            matrix = self.matrix @ other.matrix
+            transform = AffineTransform(is_camera_space=self.is_camera_space)
+            transform.matrix = matrix
+            return transform
 
         return np.asarray(self) @ other
 
-    def __array__(self, dtype=None):
-        return self.untracked_matrix.astype(dtype, copy=False)
-
 
 class RecursiveTransform(AffineBase):
     """A transform that may be preceded by another transform.
@@ -650,7 +695,7 @@ class RecursiveTransform(AffineBase):
 
     Parameters
     ----------
-    matrix : AffineBase, ndarray [4, 4]
+    base : AffineBase
         The base transform that will be wrapped by this transform.
     parent : AffineBase, optional
         The parent transform that precedes the base transform.
@@ -678,7 +723,7 @@ class RecursiveTransform(AffineBase):
 
     def __init__(
         self,
-        matrix: Union[np.ndarray, AffineBase],
+        base: AffineBase,
         /,
         *,
         parent=None,
@@ -690,10 +735,7 @@ class RecursiveTransform(AffineBase):
         self.own = None
         self.last_modified = perf_counter_ns()
 
-        if isinstance(matrix, AffineBase):
-            self.own = matrix
-        else:
-            self.own = AffineTransform(matrix, is_camera_space=is_camera_space)
+        self.own = base
         self.own._set_reference_up_provider(self)
 
         if parent is None:
@@ -720,7 +762,9 @@ class RecursiveTransform(AffineBase):
     def __own_reference_up(self) -> np.ndarray:
         new_ref = la.vec_transform(self._reference_up, self._parent.inverse_matrix)
         origin = la.vec_transform((0, 0, 0), self._parent.inverse_matrix)
-        return la.vec_normalize(new_ref - origin)
+        vec = la.vec_normalize(new_ref - origin)
+        vec.flags.writeable = False
+        return vec
 
     @property
     def _own_reference_up(self):
@@ -730,7 +774,7 @@ class RecursiveTransform(AffineBase):
     def _own_reference_up(self, value):
         new_ref = la.vec_transform(value, self.parent.matrix)
         origin = self.parent.position
-        self._reference_up = la.vec_normalize(new_ref - origin)
+        self._reference_up[:] = la.vec_normalize(new_ref - origin)
         self.flag_update()
         # Note: since __own_reference_up is only used in a setter in AffineBase
         # we do not need to call flag_update() on self.own; all its state and cache
@@ -755,7 +799,9 @@ class RecursiveTransform(AffineBase):
 
     @cached
     def _matrix(self):
-        return self._parent.matrix @ self.own.matrix
+        mat = self._parent.matrix @ self.own.matrix
+        mat.flags.writeable = False
+        return mat
 
     @property
     def matrix(self):
@@ -773,4 +819,6 @@ class RecursiveTransform(AffineBase):
 
     @cached
     def scaling_signs(self):
-        return self._parent.scaling_signs * self.own.scaling_signs
+        signs = self._parent.scaling_signs * self.own.scaling_signs
+        signs.flags.writeable = False
+        return signs
diff --git a/pyproject.toml b/pyproject.toml
index 80191b7..73daf53 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ requires-python = ">= 3.9"
 dependencies = [
     "rendercanvas >= 2",
     "wgpu >=0.19.0,<0.20.0",
-    "pylinalg >=0.6.0,<0.7.0",
+    "pylinalg >=0.6.1,<0.7.0",
     "numpy",
     "freetype-py",
     "uharfbuzz",


Test patch:
diff --git a/tests/objects/test_world_object.py b/tests/objects/test_world_object.py
index f879061..880964f 100644
--- a/tests/objects/test_world_object.py
+++ b/tests/objects/test_world_object.py
@@ -439,9 +439,10 @@ def test_scaling_signs_manual_matrix():
     npt.assert_array_almost_equal(child.world.scale, expected)
 
     s2 = (-4, -4, 4)
+    expected2 = np.array(expected) * np.array(s2)
     child.local.scale = s2
     npt.assert_array_almost_equal(child.local.scale, s2)
-    npt.assert_array_almost_equal(child.world.scale, [-4, 8, 12])
+    npt.assert_array_almost_equal(child.world.scale, expected2)
 
 
 def test_rotation_derived():