Merge pull request #27 from GazzolaLab/feat/pose

[Feat] Pose

docs/api/geometry-composite.rst

@@ -12,3 +12,8 @@ Stacked geometry
 
 .. automodule:: bsr.geometry.composite.stack
    :members:
+
+Pose geometry
+-------------
+.. automodule:: bsr.geometry.composite.pose
+   :members:

src/bsr/__init__.py

@@ -11,7 +11,9 @@
 from .blender_commands.macros import (
     clear_materials,
     clear_mesh_objects,
+    deselect_all,
     scene_update,
+    select_camera,
 )
 from .frame import FrameManager
 from .geometry.composite.rod import Rod

src/bsr/blender_commands/macros.py

@@ -24,3 +24,13 @@ def clear_materials() -> None:
     # Clear existing materials in the scene
     for material in bpy.data.materials:
         bpy.data.materials.remove(material)
+
+
+def deselect_all() -> None:
+    # Deselect all objects
+    bpy.ops.object.select_all(action="DESELECT")
+
+
+def select_camera() -> None:
+    # Select the camera object
+    bpy.context.view_layer.objects.active = bpy.data.objects["Camera"]

src/bsr/frame.py

@@ -70,3 +70,31 @@ def set_frame_end(self, frame: Optional[int] = None) -> None:
                 isinstance(frame, int) and frame >= 0
             ), "frame must be a positive integer or 0"
         bpy.context.scene.frame_end = frame
+
+    def set_frame_rate(self, fps: int | float) -> None:
+        """
+        Set the frame rate of the scene.
+
+        Parameters
+        ----------
+        fps : float
+            The frame rate of the scene. (Frame per second)
+        """
+        assert isinstance(fps, (int, float)), "fps must be a number"
+        assert fps > 0, "fps must be a positive value"
+        bpy.context.scene.render.fps = int(fps)
+        bpy.context.scene.render.fps_base = int(fps) / fps
+
+    def get_frame_rate(self) -> float:
+        """
+        Get the frame rate of the scene.
+
+        Returns
+        -------
+        float
+            The frame rate of the scene. (Frame per second)
+        """
+        fps: float = (
+            bpy.context.scene.render.fps / bpy.context.scene.render.fps_base
+        )
+        return fps

src/bsr/geometry/composite/pose.py

@@ -0,0 +1,104 @@
+__doc__ = """
+Pose class for creating and updating poses in Blender
+"""
+__all__ = ["Pose"]
+
+import bpy
+from numpy.typing import NDArray
+
+from bsr.geometry.primitives.simple import Cylinder, Sphere
+from bsr.tools.keyframe_mixin import KeyFrameControlMixin
+
+
+class Pose(KeyFrameControlMixin):
+    """
+    Pose class for managing visualization and rendering in Blender
+
+    Parameters
+    ----------
+    position : NDArray
+        The position of pose. Expected shape is (n_dim,).
+        n_dim = 3
+    directors : NDArray
+        The directors of the pose. Expected shape is (n_dim, n_dim).
+        n_dim = 3
+
+    """
+
+    input_states = {"position", "directors"}
+
+    def __init__(
+        self,
+        position: NDArray,
+        directors: NDArray,
+        unit_length: float = 1.0,
+        thickness_ratio: float = 0.1,
+    ) -> None:
+        # create sphere and cylinder objects
+        self.spheres: list[Sphere] = []
+        self.cylinders: list[Cylinder] = []
+        self._bpy_objs: dict[str, bpy.types.Object] = {
+            "spheres": self.spheres,
+            "cylinders": self.cylinders,
+        }
+        self.__unit_length = unit_length
+        self.__ratio = thickness_ratio
+
+        self._build(position, directors)
+
+    @property
+    def object(self) -> dict[str, bpy.types.Object]:
+        """
+        Return the dictionary of Blender objects: spheres and cylinders
+        """
+        return self._bpy_objs
+
+    def _build(self, position: NDArray, directors: NDArray) -> None:
+        """
+        Build the pose object from the given position and directors
+        """
+        # create the sphere object at the position
+        sphere = Sphere(
+            position,
+            self.__unit_length * self.__ratio,
+        )
+        self.spheres.append(sphere)
+
+        # create cylinder and sphere objects for each director
+        for i in range(directors.shape[1]):
+            tip_position = position + directors[:, i] * self.__unit_length
+            cylinder = Cylinder(
+                position,
+                tip_position,
+                self.__unit_length * self.__ratio,
+            )
+            self.cylinders.append(cylinder)
+
+            sphere = Sphere(
+                tip_position,
+                self.__unit_length * self.__ratio,
+            )
+            self.spheres.append(sphere)
+
+    def update_states(self, position: NDArray, directors: NDArray) -> None:
+        """
+        Update the states of the pose object
+        """
+        self.spheres[0].update_states(position)
+
+        for i, cylinder in enumerate(self.cylinders):
+            tip_position = position + directors[:, i] * self.__unit_length
+            cylinder.update_states(position, tip_position)
+
+            sphere = self.spheres[i + 1]
+            sphere.update_states(tip_position)
+
+    def set_keyframe(self, keyframe: int) -> None:
+        """
+        Set the keyframe for the pose object
+        """
+        for shperes in self.spheres:
+            shperes.set_keyframe(keyframe)
+
+        for cylinder in self.cylinders:
+            cylinder.set_keyframe(keyframe)

tests/test_frame.py

@@ -0,0 +1,63 @@
+import math
+
+import bpy
+import pytest
+
+from bsr.frame import FrameManager
+
+
+class TestFrameManager:
+
+    def test_frame_manager_singleton(self):
+        assert FrameManager() is FrameManager()
+
+    def test_frame_manager_current_frame(self):
+        frame_manager = FrameManager()
+        assert frame_manager.current_frame == 0
+
+    def test_frame_manager_current_frame_setter(self):
+        frame_manager = FrameManager()
+        frame_manager.current_frame = 10
+        assert frame_manager.current_frame == 10
+
+    def test_frame_manager_current_frame_setter_with_wrong_frame(self):
+        frame_manager = FrameManager()
+        with pytest.raises(AssertionError):
+            frame_manager.current_frame = -1
+
+    def test_frame_manager_update(self):
+        frame_manager = FrameManager()
+        frame_manager.current_frame = 0
+        frame_manager.update(10)
+        assert frame_manager.current_frame == 10
+
+    def test_frame_manager_update_with_wrong_forwardframe(self):
+        frame_manager = FrameManager()
+        with pytest.raises(AssertionError):
+            frame_manager.update(-1)
+
+    def test_frame_manager_set_frame_end(self):
+        frame_manager = FrameManager()
+        frame_manager.set_frame_end(100)
+        assert bpy.context.scene.frame_end == 100
+
+    def test_frame_manager_set_frame_end_with_none(self):
+        frame_manager = FrameManager()
+        frame_manager.current_frame = 0
+        frame_manager.update(250)
+        frame_manager.set_frame_end()
+        assert bpy.context.scene.frame_end == 250
+
+    def test_frame_manager_set_frame_end_with_wrong_frame(self):
+        frame_manager = FrameManager()
+        with pytest.raises(AssertionError):
+            frame_manager.set_frame_end(-1)
+
+    def test_frame_manager_get_set_frame_rate(self):
+        frame_manager = FrameManager()
+        frame_manager.set_frame_rate(30)
+        assert frame_manager.get_frame_rate() == 30
+        frame_manager.set_frame_rate(29.97)
+        assert math.isclose(
+            frame_manager.get_frame_rate(), 29.97, abs_tol=0.001
+        )