Refactoring

Adding Hinge joints to rigid body simulator

rainbow/simulators/prox_rigid_bodies/api.py

@@ -16,6 +16,7 @@
 import rainbow.math.coordsys as FRAME
 import rainbow.simulators.prox_rigid_bodies.steppers as STEPPERS
 from rainbow.simulators.prox_rigid_bodies.types import *
+import rainbow.math.matrix3 as M3
 
 
 def generate_unique_name(name: str) -> str:
@@ -77,60 +78,77 @@ def create_hinge(engine, hinge_name: str) -> None:
     engine.hinges[hinge_name] = hinge
 
 
-def set_hinge_sockets(engine, hinge_name: str,
-                      parent_name: str,
-                      r_parent: np.ndarray,
-                      q_parent: np.ndarray,
-                      child_name: str,
-                      r_child: np.ndarray,
-                      q_child: np.ndarray
-                      ) -> None:
-    """
-
-    :param engine:
-    :param hinge_name:
-    :param parent_name:
-    :param r_parent:
-    :param q_parent:
-    :param child_name:
-    :param r_child:
-    :param q_child:
+def set_hinge(
+        engine: Engine,
+        hinge_name: str,
+        parent_name: str,
+        child_name: str,
+        origin: np.ndarray,
+        axis: np.ndarray,
+        mode: str = "world"
+) -> None:
+    """
+    Set hinge joint parameters.
+
+    :param engine:        The engine that contains the hinge.
+    :param hinge_name:    The name of the hinge.
+    :param parent_name:   The name of the parent body (link).
+    :param child_name:    The name of the child body (link).
+    :param origin:        The position of the origin for the joint frame.
+    :param axis:          The joint axis of the hinge.
+    :param mode:          This controls the coordinate space in which axis and origin are given with respect
+                          to.
+                          It can be either "world", "body" frame of parent body, or "model" frame of parent.
     """
     if hinge_name in engine.hinges:
         hinge = engine.hinges[hinge_name]
     else:
-        raise RuntimeError("set_hinge_sockets() no such rigid body exist with that name")
+        raise RuntimeError("set_hinge() no such rigid body exist with that name")
     if parent_name in engine.bodies:
-        parent = engine.bodies[parent_name]
+        parent_link = engine.bodies[parent_name]
     else:
-        raise RuntimeError("set_hinge_sockets() no such rigid body exist with that name")
+        raise RuntimeError("set_hinge() no such rigid body exist with that name")
     if child_name in engine.bodies:
-        child = engine.bodies[child_name]
+        child_link = engine.bodies[child_name]
+    else:
+        raise RuntimeError("set_hinge() no such rigid body exist with that name")
+
+    o_world = None
+    s_world = None
+
+    if mode == "world":
+        o_world = np.copy(origin)
+        s_world = np.copy(axis)
+    elif mode == "body":
+        # Convert o and s from body frame to world frame
+        o_world = Q.rotate(parent_link.q, origin) + parent_link.r
+        s_world = Q.rotate(parent_link.q, axis)
+    elif mode == "model":
+        # Convert o and s from model frame to body frame
+        r_b2m = np.copy(parent_link.shape.r)
+        q_b2m = np.copy(parent_link.shape.q)
+        o_body = Q.rotate(Q.conjugate(q_b2m), origin - r_b2m)
+        s_body = Q.rotate(Q.conjugate(q_b2m), axis)
+        # Convert o and s from body frame to world frame
+        o_world = Q.rotate(parent_link.q, o_body) + parent_link.r
+        s_world = Q.rotate(parent_link.q, s_body)
     else:
-        raise RuntimeError("set_hinge_sockets() no such rigid body exist with that name")
+        raise ValueError(f"set_hinge() no such mode {mode} is supported")
 
-    socket_parent = FRAME.make(r_parent, q_parent)
-    socket_child = FRAME.make(r_child, q_child)
+    # Observe, we use convention that n (the local z-axis) will be the joint axis.
+    t, b, n = V3.make_orthonormal_vectors(s_world)
+    R_world = M3.make_from_cols(t, b, n)
+    q_world = Q.from_matrix(R_world)
 
-    # Currently, we assume that the socket joint frames live in the body frame coordinate systems
-    # of the rigid bodies they belong to.
-    #
-    # A socket is a coordinate mapping from joint frame space to body-space of the link.
-    #
-    # When rigging a simulation, it may be that the rigging person does not know the body
-    # frames.
-    # Instead, what is known is the model frames of the rigid bodies.
-    #
-    #  R_parent_bf2mf = np.copy(parent.shape.r)
-    #  q_parent_bf2mf = np.copy(parent.shape.q)
-    #
-    #
-    # Hence, we know (bf->mf) we are given (jf->mf) and we need to compute (jf->bf)
-    #
-    # Xjb = Xjm Xmb
-    #
-    hinge.set_parent_socket(parent, socket_parent)
-    hinge.set_parent_socket(child, socket_child)
+    q_parent_socket = Q.prod(Q.conjugate(parent_link.q), q_world)
+    q_child_socket = Q.prod(Q.conjugate(child_link.q), q_world)
+    r_parent_socket = Q.rotate(Q.conjugate(parent_link.q), o_world - parent_link.r)
+    r_child_socket = Q.rotate(Q.conjugate(child_link.q), o_world - child_link.r)
+
+    parent_socket = FRAME.make(r_parent_socket, q_parent_socket)
+    child_socket = FRAME.make(r_child_socket, q_child_socket)
+    hinge.set_parent_socket(parent_link, parent_socket)
+    hinge.set_child_socket(child_link, child_socket)
 
 
 def create_mesh(V: np.ndarray, T: np.ndarray) -> MESH.Mesh:
@@ -387,6 +405,21 @@ def set_position(engine, body_name: str, r: np.ndarray, use_model_frame=False) -
         body.r = np.copy(r)
 
 
+def get_position(engine, body_name: str) -> np.ndarray:
+    """
+    Retrieve the center of mass position for the rigid body.
+
+    :param engine:           The engine that stores the rigid body.
+    :param body_name:        The name of the rigid body.
+    :return:                 The center of mass position for the given rigid body.
+    """
+    if body_name in engine.bodies:
+        body = engine.bodies[body_name]
+    else:
+        raise RuntimeError("get_position() no such rigid body exist with that name")
+    return np.copy(body.r)
+
+
 def set_orientation(engine, body_name: str, q: np.ndarray, use_model_frame=False) -> None:
     """
     Set orientation of rigid body.
@@ -446,6 +479,21 @@ def set_orientation(engine, body_name: str, q: np.ndarray, use_model_frame=False
         body.q = Q.unit(q)
 
 
+def get_orientation(engine, body_name: str) -> np.ndarray:
+    """
+    Retrieve body frame orientation for the rigid body.
+
+    :param engine:           The engine that stores the rigid body.
+    :param body_name:        The name of the rigid body.
+    :return:                 The orientation of the body frame for the given rigid body.
+    """
+    if body_name in engine.bodies:
+        body = engine.bodies[body_name]
+    else:
+        raise RuntimeError("get_orientation() no such rigid body exist with that name")
+    return np.copy(body.q)
+
+
 def set_velocity(engine, body_name: str, v: np.ndarray) -> None:
     """
     Set the linear velocity of a rigid body.
@@ -462,6 +510,21 @@ def set_velocity(engine, body_name: str, v: np.ndarray) -> None:
     body.v = np.copy(v)
 
 
+def get_velocity(engine, body_name: str) -> np.ndarray:
+    """
+    Retrieve the linear velocity for the rigid body.
+
+    :param engine:           The engine that stores the rigid body.
+    :param body_name:        The name of the rigid body.
+    :return:                 The linear velocity for the given rigid body.
+    """
+    if body_name in engine.bodies:
+        body = engine.bodies[body_name]
+    else:
+        raise RuntimeError("get_velocity() no such rigid body exist with that name")
+    return np.copy(body.v)
+
+
 def set_spin(engine, body_name: str, w: np.ndarray) -> None:
     """
     Set the angular velocity of a rigid body.
@@ -478,6 +541,21 @@ def set_spin(engine, body_name: str, w: np.ndarray) -> None:
     body.w = np.copy(w)
 
 
+def get_spin(engine, body_name: str) -> np.ndarray:
+    """
+    Retrieve the angular velocity for the rigid body.
+
+    :param engine:           The engine that stores the rigid body.
+    :param body_name:        The name of the rigid body.
+    :return:                 The angular velocity for the given rigid body.
+    """
+    if body_name in engine.bodies:
+        body = engine.bodies[body_name]
+    else:
+        raise RuntimeError("get_spin() no such rigid body exist with that name")
+    return np.copy(body.w)
+
+
 def set_mass_properties(engine, body_name: str, density: float) -> None:
     """
     Set the mass properties of a given rigid body.

rainbow/simulators/prox_rigid_bodies/gauss_seidel_solver.py

@@ -11,7 +11,12 @@
 from rainbow.util.timer import Timer
 
 
-def solve(engine, Problems: list, performance_data: dict[str, any], profiling_on: bool, prefix: str) -> None:
+def solve(
+        engine, Problems: list,
+        performance_data: dict[str, any],
+        profiling_on: bool,
+        prefix: str
+) -> None:
     """
 
     :param engine:

rainbow/simulators/prox_rigid_bodies/procedural/__init__.py

@@ -19,6 +19,12 @@
 from .create_temple import create_temple
 from .create_tower import create_tower
 from .create_ground import create_ground
+from .create_box_stack import create_box_stack
+from .create_cube_hinge_chain import create_cube_hinge_chain
+from .write_xml import write_xml
+from .create_scene import create_scene
+from .create_scene import get_scene_names
+
 
 __all__ = ["create_arch",
            "create_chainmail",
@@ -35,5 +41,10 @@
            "create_poles",
            "create_rock_slide",
            "create_sandbox",
-           "create_temple"
+           "create_temple",
+           "create_box_stack",
+           "create_cube_hinge_chain",
+           "write_xml",
+           "create_scene",
+           "get_scene_names"
            ]

rainbow/simulators/prox_rigid_bodies/procedural/create_box_stack.py

@@ -0,0 +1,61 @@
+"""
+This script contains code to create a sandbox scene.
+"""
+import numpy as np
+
+import rainbow.math.vector3 as V3
+import rainbow.math.quaternion as Q
+import rainbow.geometry.surface_mesh as MESH
+import rainbow.simulators.prox_rigid_bodies.api as API
+from rainbow.simulators.prox_rigid_bodies.types import Engine
+
+
+def create_box_stack(
+        engine: Engine,
+        box_width: float,
+        box_height: float,
+        box_depth: float,
+        K_boxes: int,
+        density: float,
+        material_name: str,
+) -> list[str]:
+    """
+    This function creates a box stack scene.
+
+    :param engine:                  The engine that will be used to create the dry stone rigid bodies in.
+    :param box_width:               The width of the boxes.
+    :param box_height:              The height of the boxes.
+    :param box_depth:               The depth of the boxes.
+    :param K_boxes:                 The number of boxes that will be stacked.
+    :param density:                 The mass density to use for all the rigid bodies.
+    :param material_name:           The material name to use for all the rigid bodies that are created.
+    :return:                        A list with the names of all the rigid bodies that were created.
+    """
+
+    shape_name = API.generate_unique_name("box")
+    V, T = MESH.create_box(box_width, box_height, box_depth)
+    mesh = API.create_mesh(V, T)
+    API.create_shape(engine, shape_name, mesh)
+
+    body_names = []
+
+    for k in range(K_boxes):
+
+        body_name = API.generate_unique_name("body")
+        body_names.append(body_name)
+        API.create_rigid_body(engine, body_name)
+        API.connect_shape(engine, body_name, shape_name)
+
+        height = k*box_height + box_height/2.0
+        r = V3.make(0.0, height, 0.0)
+
+        radians = k*np.pi/4
+        q = Q.Ry(radians)
+
+        API.set_position(engine, body_name, r, True)
+        API.set_orientation(engine, body_name, q, True)
+        API.set_body_type(engine, body_name, "free")
+        API.set_body_material(engine, body_name, material_name)
+        API.set_mass_properties(engine, body_name, density)
+
+    return body_names