Complete geometry.py module: handling sphere and cylinder

src/bsr/geometry.py

@@ -1,38 +1,127 @@
+__doc__ = """
+This module provides a set of geometry-mesh interfaces for blender objects.
+"""
+__all__ = ["BlenderMeshInterfaceProtocol", "Sphere", "Cylinder"]
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ParamSpec,
+    Protocol,
+    Type,
+    TypedDict,
+    TypeVar,
+)
+from typing_extensions import Self
+
 import colorsys
 
 import bpy
 import numpy as np
 
+MeshDataType = dict[str, Any]
+
+S = TypeVar("S", bound="BlenderMeshInterfaceProtocol")
+P = ParamSpec("P")
+
+
+class BlenderMeshInterfaceProtocol(Protocol):
+    """
+    This protocol defines the interface for Blender mesh objects.
+    """
+
+    # TODO: For future implementation
+    # @property
+    # def data(self): ...
+
+    @property
+    def object(self) -> bpy.types.Object:
+        """Returns associated Blender object."""
+
+    @classmethod
+    def create(cls: Type[S], states: MeshDataType) -> S:
+        """Creates a new mesh object with the given states."""
+
+    def update_states(self, *args: Any) -> bpy.types.Object:
+        """Updates the mesh object with the given states."""
+
+    # def update_material(self, material) -> None: ...  # TODO: For future implementation
+
 
 class Sphere:
-    def __init__(self, location, radius=0.005):
-        self.obj = self.create_sphere(location, radius)
+    """
+    This class provides a mesh interface for Blender Sphere objects.
+    Sphere objects are created with the given position and radius.
 
-    def create_sphere(self, location, radius):
-        bpy.ops.mesh.primitive_uv_sphere_add(radius=radius, location=location)
-        return bpy.context.active_object
+    Parameters
+    ----------
+    position : np.ndarray
+        The position of the sphere object.
+    radius : float
+        The radius of the sphere object.
+    """
+
+    def __init__(self, position: np.ndarray, radius: float) -> None:
+        self._obj = self._create_sphere()
+        self.update_states(position, radius)
 
-    def update_position(self, location):
-        self.obj.location.z = location[2]
-        self.obj.location.y = location[1]
-        self.obj.location.x = location[0]
+    @classmethod
+    def create(cls, states: MeshDataType) -> "Sphere":
+        return cls(states["position"], states["radius"])
 
+    @property
+    def object(self) -> bpy.types.Object:
+        return self._obj
 
+    def update_states(
+        self, position: np.ndarray | None = None, radius: float | None = None
+    ) -> bpy.types.Object:
+        if position is not None:
+            self.object.location.x = position[0]
+            self.object.location.y = position[1]
+            self.object.location.z = position[2]
+        if radius is not None:
+            self.object.scale = (radius, radius, radius)
+        return self.object
+
+    def _create_sphere(self) -> bpy.types.Object:
+        """
+        Creates a new sphere object with the given position and radius.
+        """
+        bpy.ops.mesh.primitive_uv_sphere_add()
+        return bpy.context.active_object
+
+
+# FIXME: This class needs to be modified to conform to the BlenderMeshInterfaceProtocol
 class Cylinder:
-    def __init__(self, pos1, pos2):
-        self.obj = self.create_cylinder(pos1, pos2)
-        self.mat = bpy.data.materials.new(name="cyl_mat")
-        self.obj.active_material = self.mat
+    """
+    TODO: Add documentation
+    """
 
-    def create_cylinder(self, pos1, pos2):
-        depth, center, angles = self.calc_cyl_orientation(pos1, pos2)
-        bpy.ops.mesh.primitive_cylinder_add(
-            radius=0.005, depth=1, location=center
+    def __init__(
+        self, pos1: np.ndarray, pos2: np.ndarray, radius: float, depth: float
+    ):
+        self.pos1 = pos1
+        self.pos2 = pos2
+        self.obj = self._create_cylinder(pos1, pos2, radius, depth)
+
+    @classmethod
+    def create(cls, states: MeshDataType) -> "Cylinder":
+        return cls(
+            states["pos1"], states["pos2"], states["radius"], states["depth"]
         )
-        cylinder = bpy.context.active_object
-        cylinder.rotation_euler = (0, angles[1], angles[0])
-        cylinder.scale[2] = depth
-        return cylinder
+
+    @property
+    def object(self) -> bpy.types.Object:
+        return self._obj
+
+    def update_states(self, pos1, pos2, radius):
+        depth, center, angles = self.calc_cyl_orientation(pos1, pos2)
+        self.obj.location = center
+        self.obj.rotation_euler = (0, angles[1], angles[0])
+        self.obj.scale[2] = depth
+        self.obj.scale[0] = radius
+        self.obj.scale[1] = radius
 
     def calc_cyl_orientation(self, pos1, pos2):
         pos1 = np.array(pos1)
@@ -47,29 +136,29 @@ def calc_cyl_orientation(self, pos1, pos2):
         angles = np.array([phi, theta])
         return depth, center, angles
 
-    def update_position(self, pos1, pos2):
+    def _create_cylinder(
+        self, pos1: np.ndarray, pos2: np.ndarray, radius: float, depth: float
+    ) -> bpy.types.Object:
+        """
+        Creates a new cylinder object with the given end positions, radius, centerpoint and depth.
+        """
         depth, center, angles = self.calc_cyl_orientation(pos1, pos2)
-        self.obj.location = (center[0], center[1], center[2])
-        self.obj.rotation_euler = (0, angles[1], angles[0])
-        self.obj.scale[2] = depth
-
-        # computing deformation heat-map
-        max_def = 0.07
-
-        h = (
-            -np.sqrt(self.obj.location[0] ** 2 + self.obj.location[2] ** 2)
-            / max_def
-            + 240 / 360
-        )
-        v = (
-            np.sqrt(self.obj.location[0] ** 2 + self.obj.location[2] ** 2)
-            / max_def
-            * 0.5
-            + 0.5
+        bpy.ops.mesh.primitive_uv_cylinder_add(
+            radius=radius, location=center, depth=depth
         )
+        cylinder = bpy.context.active_object
+        cylinder.rotation_euler = (0, angles[1], angles[0])
+        cylinder.scale[2] = depth
+        return cylinder
 
-        r, g, b = colorsys.hsv_to_rgb(h, 1, v)
-        self.update_color(r, g, b, 1)
 
-    def update_color(self, r, g, b, a):
-        self.mat.diffuse_color = (r, g, b, a)
+if TYPE_CHECKING:
+    # This is required for explicit type-checking
+    data = {"position": np.array([0, 0, 0]), "radius": 1.0}
+    _: BlenderMeshInterfaceProtocol = Sphere.create(data)
+    data = {
+        "position_1": np.array([0, 0, 0]),
+        "position_2": np.array([1, 1, 1]),
+        "radius": 1.0,
+    }
+    _: BlenderMeshInterfaceProtocol = Cylinder.create(data)

src/bsr/macros.py

@@ -8,3 +8,13 @@ def clear_mesh_objects() -> None:
     bpy.ops.object.select_all(action="DESELECT")
     bpy.ops.object.select_by_type(type="MESH")
     bpy.ops.object.delete()
+
+
+def scene_update() -> None:
+    """
+    Update the scene
+
+    Used to update object's matrix_world after transformations
+    (https://blender.stackexchange.com/questions/27667/incorrect-matrix-world-after-transformation)
+    """
+    bpy.context.view_layer.update()

tests/geometry/test_primitive_geometry_mesh.py

@@ -0,0 +1,138 @@
+import numpy as np
+import pytest
+
+from bsr.geometry import Cylinder, Sphere
+from tests.geometry.utils import get_mesh_limit
+
+# Visual tolerance for the mesh limit
+_VISUAL_ATOL = 1e-7
+_VISUAL_RTOL = 1e-4
+
+
+@pytest.mark.parametrize(
+    "center",
+    [
+        np.array([10, 10, 10]),
+        np.array([10, 11, 10]),
+        np.array([10, 11, 11]),
+        np.array([11, 11, 11]),
+    ],
+)
+@pytest.mark.parametrize("radius", [1, 2, 3, 5.5])
+def test_sphere_radius_and_position(center, radius):
+    x_min, x_max = center[0] - radius, center[0] + radius
+    y_min, y_max = center[1] - radius, center[1] + radius
+    z_min, z_max = center[2] - radius, center[2] + radius
+
+    sphere = Sphere(position=center, radius=radius)
+
+    mesh_limit = get_mesh_limit(sphere)
+
+    np.testing.assert_allclose(
+        (x_min, x_max, y_min, y_max, z_min, z_max),
+        mesh_limit,
+        rtol=_VISUAL_RTOL,
+        atol=_VISUAL_ATOL,
+    )
+
+
+@pytest.mark.parametrize(
+    "position_one",
+    [
+        np.array([10, 10, 10]),
+        np.array([10, 11, -10]),
+        np.array([-10, 11, 11]),
+    ],
+)
+@pytest.mark.parametrize("length", [1, 10.5, -1, -10.5])
+@pytest.mark.parametrize("radius", [1, 3, 5.5])
+def test_x_cylinder_radius_and_positions(position_one, length, radius):
+    position_two = position_one + np.array([length, 0, 0])
+    y, z = position_one[1], position_one[2]
+
+    x_min, x_max = min(position_one[0], position_two[0]), max(
+        position_one[0], position_two[0]
+    )
+    y_min, y_max = y - radius, y + radius
+    z_min, z_max = z - radius, z + radius
+
+    cylinder = Cylinder(
+        position_1=position_one, position_2=position_two, radius=radius
+    )
+
+    mesh_limit = get_mesh_limit(cylinder)
+
+    np.testing.assert_allclose(
+        (x_min, x_max, y_min, y_max, z_min, z_max),
+        mesh_limit,
+        rtol=_VISUAL_RTOL,
+        atol=_VISUAL_ATOL,
+    )
+
+
+@pytest.mark.parametrize(
+    "position_one",
+    [
+        np.array([10, 10, 10]),
+        np.array([10, 11, -10]),
+        np.array([-10, 11, 11]),
+    ],
+)
+@pytest.mark.parametrize("length", [1, 10.5, -1, -10.5])
+@pytest.mark.parametrize("radius", [1, 3, 5.5])
+def test_y_cylinder_radius_and_positions(position_one, length, radius):
+    position_two = position_one + np.array([0, length, 0])
+    x, z = position_one[0], position_one[2]
+
+    x_min, x_max = x - radius, x + radius
+    y_min, y_max = min(position_one[1], position_two[1]), max(
+        position_one[1], position_two[1]
+    )
+    z_min, z_max = z - radius, z + radius
+
+    cylinder = Cylinder(
+        position_1=position_one, position_2=position_two, radius=radius
+    )
+
+    mesh_limit = get_mesh_limit(cylinder)
+
+    np.testing.assert_allclose(
+        (x_min, x_max, y_min, y_max, z_min, z_max),
+        mesh_limit,
+        rtol=_VISUAL_RTOL,
+        atol=_VISUAL_ATOL,
+    )
+
+
+@pytest.mark.parametrize(
+    "position_one",
+    [
+        np.array([10, 10, 10]),
+        np.array([10, 11, -10]),
+        np.array([-10, 11, 11]),
+    ],
+)
+@pytest.mark.parametrize("length", [1, 10.5, -1, -10.5])
+@pytest.mark.parametrize("radius", [1, 3, 5.5])
+def test_z_cylinder_radius_and_positions(position_one, length, radius):
+    position_two = position_one + np.array([0, 0, length])
+    x, y = position_one[0], position_one[1]
+
+    x_min, x_max = x - radius, x + radius
+    y_min, y_max = y - radius, y + radius
+    z_min, z_max = min(position_one[2], position_two[2]), max(
+        position_one[2], position_two[2]
+    )
+
+    cylinder = Cylinder(
+        position_1=position_one, position_2=position_two, radius=radius
+    )
+
+    mesh_limit = get_mesh_limit(cylinder)
+
+    np.testing.assert_allclose(
+        (x_min, x_max, y_min, y_max, z_min, z_max),
+        mesh_limit,
+        rtol=_VISUAL_RTOL,
+        atol=_VISUAL_ATOL,
+    )

tests/geometry/utils.py

@@ -0,0 +1,22 @@
+import bpy
+import numpy as np
+
+from bsr.geometry import BlenderMeshInterfaceProtocol
+from bsr.macros import scene_update
+
+
+def get_mesh_limit(interface: BlenderMeshInterfaceProtocol):
+    """(For testing) Given blender mesh object, return xyz limit"""
+
+    obj = interface.object
+    scene_update()
+
+    vertices_coords = []
+    for v in obj.data.vertices:
+        global_coord = obj.matrix_world @ v.co
+        vertices_coords.append(list(global_coord))
+    vertices_coords = np.array(vertices_coords)
+
+    x_min, y_min, z_min = np.min(vertices_coords, axis=0)
+    x_max, y_max, z_max = np.max(vertices_coords, axis=0)
+    return x_min, x_max, y_min, y_max, z_min, z_max