center of mass functions

pymotion/ops/center_of_mass.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+
+def human_center_of_mass(
+    joints_spine: np.array,
+    joints_left_arm: np.array,
+    joints_right_arm: np.array,
+    joints_left_leg: np.array,
+    joints_right_leg: np.array,
+) -> np.array:
+    """
+    Compute the center of mass of a human body definedy by the standard human weight distribution.
+    Each arm accounts for a 5% of the body weight, each leg accounts for a 15% of the body weight
+    and the spine accounts for a 60% of the body weight.
+
+    Parameters
+    ----------
+    joints_spine : np.array[..., n_joints, 3]
+        Joint positions of the spine.
+    joints_left_arm : np.array[..., n_joints, 3]
+        Joint positions of the left arm.
+    joints_right_arm : np.array[..., n_joints, 3]
+        Joint positions of the right arm.
+    joints_left_leg : np.array[..., n_joints, 3]
+        Joint positions of the left leg.
+    joints_right_leg : np.array[..., n_joints, 3]
+        Joint positions of the right leg.
+
+    Returns
+    -------
+    center_of_mass : np.array[..., 3]
+        Center of mass.
+    """
+    n_joints_spine = joints_spine.shape[-2]
+    n_joints_left_arm = joints_left_arm.shape[-2]
+    n_joints_right_arm = joints_right_arm.shape[-2]
+    n_joints_left_leg = joints_left_leg.shape[-2]
+    n_joints_right_leg = joints_right_leg.shape[-2]
+    weights = np.array(
+        [0.6 / n_joints_spine] * n_joints_spine
+        + [0.05 / n_joints_left_arm] * n_joints_left_arm
+        + [0.05 / n_joints_right_arm] * n_joints_right_arm
+        + [0.15 / n_joints_left_leg] * n_joints_left_leg
+        + [0.15 / n_joints_right_leg] * n_joints_right_leg
+    )
+    joints = np.concatenate(
+        [joints_spine, joints_left_arm, joints_right_arm, joints_left_leg, joints_right_leg], axis=-2
+    )
+    return center_of_mass(joints, weights)
+
+
+def center_of_mass(joints: np.array, weights: np.array) -> np.array:
+    """
+    Compute the center of mass of a set of joints.
+
+    Parameters
+    ----------
+    joints : np.array[..., n_joints, 3]
+        Joint positions.
+    weights : np.array[..., n_joints]
+        Weights of the joints. The weights should sum to 1 along the last dimension.
+
+    Returns
+    -------
+    center_of_mass : np.array[..., 3]
+        Center of mass.
+    """
+    return np.sum(joints * weights[..., np.newaxis], axis=-2)

pymotion/ops/center_of_mass_torch.py

@@ -0,0 +1,69 @@
+import torch
+
+
+def human_center_of_mass(
+    joints_spine: torch.Tensor,
+    joints_left_arm: torch.Tensor,
+    joints_right_arm: torch.Tensor,
+    joints_left_leg: torch.Tensor,
+    joints_right_leg: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute the center of mass of a human body definedy by the standard human weight distribution.
+    Each arm accounts for a 5% of the body weight, each leg accounts for a 15% of the body weight
+    and the spine accounts for a 60% of the body weight.
+
+    Parameters
+    ----------
+    joints_spine : torch.Tensor[..., n_joints, 3]
+        Joint positions of the spine.
+    joints_left_arm : torch.Tensor[..., n_joints, 3]
+        Joint positions of the left arm.
+    joints_right_arm : torch.Tensor[..., n_joints, 3]
+        Joint positions of the right arm.
+    joints_left_leg : torch.Tensor[..., n_joints, 3]
+        Joint positions of the left leg.
+    joints_right_leg : torch.Tensor[..., n_joints, 3]
+        Joint positions of the right leg.
+
+    Returns
+    -------
+    center_of_mass : torch.Tensor[..., 3]
+        Center of mass.
+    """
+    n_joints_spine = joints_spine.shape[-2]
+    n_joints_left_arm = joints_left_arm.shape[-2]
+    n_joints_right_arm = joints_right_arm.shape[-2]
+    n_joints_left_leg = joints_left_leg.shape[-2]
+    n_joints_right_leg = joints_right_leg.shape[-2]
+    weights = torch.tensor(
+        [0.6 / n_joints_spine] * n_joints_spine
+        + [0.05 / n_joints_left_arm] * n_joints_left_arm
+        + [0.05 / n_joints_right_arm] * n_joints_right_arm
+        + [0.15 / n_joints_left_leg] * n_joints_left_leg
+        + [0.15 / n_joints_right_leg] * n_joints_right_leg,
+        device=joints_spine.device,
+    )
+    joints = torch.concatenate(
+        [joints_spine, joints_left_arm, joints_right_arm, joints_left_leg, joints_right_leg], axis=-2
+    )
+    return center_of_mass(joints, weights)
+
+
+def center_of_mass(joints: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+    """
+    Compute the center of mass of a set of joints.
+
+    Parameters
+    ----------
+    joints : torch.Tensor[..., n_joints, 3]
+        Joint positions.
+    weights : torch.Tensor[..., n_joints]
+        Weights of the joints. The weights should sum to 1 along the last dimension.
+
+    Returns
+    -------
+    center_of_mass : torch.Tensor[..., 3]
+        Center of mass.
+    """
+    return torch.sum(joints * weights[..., None], dim=-2)

pymotion/ops/tests/test_center_of_mass.py

@@ -0,0 +1,83 @@
+import numpy as np
+from pymotion.ops.center_of_mass import center_of_mass, human_center_of_mass
+from pymotion.ops.center_of_mass_torch import (
+    center_of_mass as center_of_mass_torch,
+    human_center_of_mass as human_center_of_mass_torch,
+)
+from numpy.testing import assert_allclose
+import torch
+
+
+class TestCoM:
+    atol = 1e-6
+    low_atol = 1e-3  # for those operations that are not as precise
+
+    def test_com(self):
+        assert_allclose(
+            center_of_mass(
+                np.array([[[0, 0, 0]]]),
+                np.array([[1]]),
+            ),
+            np.array([[0, 0, 0]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            center_of_mass(
+                np.array([[[0, 0, 0], [1, 0, 0]]]),
+                np.array([[0.5, 0.5]]),
+            ),
+            np.array([[0.5, 0, 0]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            center_of_mass(
+                np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]),
+                np.array([[0.3, 0.3, 0.4]]),
+            ),
+            np.array([[4.3, 5.3, 6.3]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            center_of_mass_torch(
+                torch.tensor([[[0, 0, 0]]]),
+                torch.tensor([[1]]),
+            ),
+            torch.tensor([[0, 0, 0]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            center_of_mass_torch(
+                torch.tensor([[[0, 0, 0], [1, 0, 0]]]),
+                torch.tensor([[0.5, 0.5]]),
+            ),
+            torch.tensor([[0.5, 0, 0]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            center_of_mass_torch(
+                torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]),
+                torch.tensor([[0.3, 0.3, 0.4]]),
+            ),
+            torch.tensor([[4.3, 5.3, 6.3]]),
+            atol=self.atol,
+        )
+        assert_allclose(
+            human_center_of_mass(
+                np.array([[[2, 1, 5]]]),
+                np.array([[[1, 3, 4]]]),
+                np.array([[[1, 1, 1]]]),
+                np.array([[[0, 7, 6]]]),
+                np.array([[[2, 8, 1]]]),
+            ),
+            np.array([[1.6, 3.05, 4.3]]),
+        )
+        assert_allclose(
+            human_center_of_mass_torch(
+                torch.tensor([[[2, 1, 5]]]),
+                torch.tensor([[[1, 3, 4]]]),
+                torch.tensor([[[1, 1, 1]]]),
+                torch.tensor([[[0, 7, 6]]]),
+                torch.tensor([[[2, 8, 1]]]),
+            ),
+            torch.tensor([[1.6, 3.05, 4.3]]),
+        )

pymotion/rotations/quat_torch.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import torch
 import torch.nn.functional as F
 import numpy as np
@@ -68,21 +69,15 @@ def from_euler(euler: torch.Tensor, order: np.array) -> torch.Tensor:
     }
     q0 = from_angle_axis(
         euler[..., 0:1],
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 0:1])
-        ).to(euler.device),
+        torch.from_numpy(np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 0:1])).to(euler.device),
     )
     q1 = from_angle_axis(
         euler[..., 1:2],
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 1:2])
-        ).to(euler.device),
+        torch.from_numpy(np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 1:2])).to(euler.device),
     )
     q2 = from_angle_axis(
         euler[..., 2:3],
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 2:3])
-        ).to(euler.device),
+        torch.from_numpy(np.apply_along_axis(lambda x: axis[x.item()], -1, order[..., 2:3])).to(euler.device),
     )
     return mul(q0, mul(q1, q2))
 
@@ -193,29 +188,17 @@ def to_euler(quaternions: torch.Tensor, order: np.array) -> torch.Tensor:
     angle_third = 0
 
     i = (
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 2:3])[
-                ..., np.newaxis
-            ]
-        )
+        torch.from_numpy(np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 2:3])[..., np.newaxis])
         .to(quaternions.device)
         .type(torch.long)
     )
     j = (
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 1:2])[
-                ..., np.newaxis
-            ]
-        )
+        torch.from_numpy(np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 1:2])[..., np.newaxis])
         .to(quaternions.device)
         .type(torch.long)
     )
     k = (
-        torch.from_numpy(
-            np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 0:1])[
-                ..., np.newaxis
-            ]
-        )
+        torch.from_numpy(np.apply_along_axis(lambda x: aux[x.item()], -1, order[..., 0:1])[..., np.newaxis])
         .to(quaternions.device)
         .type(torch.long)
     )
@@ -238,9 +221,7 @@ def to_euler(quaternions: torch.Tensor, order: np.array) -> torch.Tensor:
     )
 
     # compute second angle
-    euler[..., 1:2] = (2 * torch.arctan2(torch.hypot(c, d), torch.hypot(a, b))) - (
-        torch.pi / 2
-    )
+    euler[..., 1:2] = (2 * torch.arctan2(torch.hypot(c, d), torch.hypot(a, b))) - (torch.pi / 2)
 
     # compute first and third angle
     half_sum = torch.arctan2(b, a)
@@ -486,9 +467,7 @@ def unroll(quaternions: torch.Tensor, dim: int) -> torch.Tensor:
     return r
 
 
-def slerp(
-    q0: torch.Tensor, q1: torch.Tensor, t: float or torch.Tensor, shortest: bool = True
-) -> torch.Tensor:
+def slerp(q0: torch.Tensor, q1: torch.Tensor, t: float | torch.Tensor, shortest: bool = True) -> torch.Tensor:
     """
     Perform spherical linear interpolation (SLERP) between two unit quaternions.
 
@@ -524,9 +503,7 @@ def slerp(
     theta = theta_0 * t  # theta = angle between q0 vector and result
 
     q2 = q1 - q0 * dot
-    q2 /= torch.linalg.norm(
-        q2 + 0.000001, dim=-1, keepdim=True
-    )  # {q0, q2} is now an orthonormal basis
+    q2 /= torch.linalg.norm(q2 + 0.000001, dim=-1, keepdim=True)  # {q0, q2} is now an orthonormal basis
 
     return torch.cos(theta) * q0 + torch.sin(theta) * q2
 

pyproject.toml

@@ -10,7 +10,7 @@ exclude = ["*test*"]
 
 [project]
 name = "upc-pymotion"
-version = "0.1.3"
+version = "0.1.4"
 description = "A Python library for working with motion data in NumPy or PyTorch."
 readme = "README.md"
 authors = [{ name = "Jose Luis Ponton", email = "jose.luis.ponton@upc.edu" }]