Feature/inv kinematics

pyproject.toml

@@ -15,6 +15,7 @@ types-pyyaml = "^6.0.12.11"
 pydantic = "^2.3.0"
 ipython = "^8.14.0"
 networkx = "^3.1"
+pandas = "^2.2.2"
 
 [[tool.poetry.source]]
 name = "rocsys"

src/ktree/k_types.py

@@ -5,6 +5,7 @@
     BaseModel,
     ConfigDict,
     Field,
+    ValidationInfo,
     computed_field,
     field_validator,
     model_serializer,
@@ -35,10 +36,11 @@ class Vector(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
 
     vector: NDArray[np.float_] = Field(default=np.array([0.0, 0.0, 0.0]), min_length=3, max_length=3)
+    # mm: bool = Field(default=False)
 
-    def __init__(self, vector: NDArray[np.float_] | list[float] = np.array([0.0, 0.0, 0.0])) -> None:
-        super().__init__(**dict(vector=vector))
-        self.vector = _validate_list(self.vector)
+    # def __init__(self, vector: NDArray[np.float_] | list[float] = np.array([0.0, 0.0, 0.0])) -> None:
+    #     super().__init__(**dict(vector=vector))
+    #     self.vector = _validate_list(self.vector)
 
     # validators
     @field_validator("vector", mode="before")
@@ -103,9 +105,26 @@ def __str__(self) -> str:
     def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
             return np.allclose(self.vector, other.vector)
+        elif isinstance(other, np.ndarray):
+            if other.shape == (3,):
+                return np.allclose(self.vector, other)
+            else:
+                raise ValueError(f"Cannot compare Vector with {other}")
         else:
             raise NotImplementedError(f"Cannot compare Vector with {other}")
 
+    @staticmethod
+    def unit_x() -> "Vector":
+        return Vector(vector=np.array([1.0, 0.0, 0.0]))
+
+    @staticmethod
+    def unit_y() -> "Vector":
+        return Vector(vector=np.array([0.0, 1.0, 0.0]))
+
+    @staticmethod
+    def unit_z() -> "Vector":
+        return Vector(vector=np.array([0.0, 0.0, 1.0]))
+
 
 class JointType(str, Enum):
     FIXED = "fixed"
@@ -114,21 +133,34 @@ class JointType(str, Enum):
     SPATIAL = "spatial"
 
 
-class Rotation(BaseModel):
+class Quaternion(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    rpy: NDArray[np.float_] = Field(default=np.array([0.0, 0.0, 0.0]), min_length=3, max_length=3)
+    qx: float = Field(default=0.0)
+    qy: float = Field(default=0.0)
+    qz: float = Field(default=0.0)
+    qw: float = Field(default=1.0)
 
-    def __init__(self, rpy: NDArray[np.float_] | list[float] = np.array([0.0, 0.0, 0.0])) -> None:
-        super().__init__(**dict(rpy=rpy))
-        self.rpy = _validate_list(self.rpy)
+    @computed_field  # type: ignore[misc]
+    @property
+    def vector(self) -> NDArray:
+        return np.array([self.qx, self.qy, self.qz, self.qw])
+
+
+class AngleAxis(BaseModel):
+    axis: Vector = Field(default=Vector(), description="Axis of rotation")
+    angle: float = Field(default=0.0, description="Angle of rotation in radians")
+
+
+class Rotation(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
+
+    rpy: NDArray[np.float_] = Field(default=np.array([0.0, 0.0, 0.0]), min_length=3, max_length=3)
 
-    # validators
-    # _rpy_validator = field_validator("rpy", mode="before")(_validate_list)
     @field_validator("rpy", mode="before")
     @classmethod
-    def _rpy_validator(cls, v: NDArray[np.float_] | list[float]) -> NDArray:
-        return _validate_list(v)
+    def _rpy_validator(cls, rpy_value: NDArray[np.float_] | list[float], info: ValidationInfo) -> NDArray:
+        return _validate_list(rpy_value)
 
     @staticmethod
     def rot_x(angle: float) -> NDArray:
@@ -202,6 +234,45 @@ def matrix(self, matrix: NDArray) -> None:
 
         self.rpy = np.array([rx, ry, rz])
 
+    @computed_field  # type: ignore[misc]
+    @property
+    def quaternion(self) -> Quaternion:
+        cr = np.cos(self.rpy[0] * 0.5)
+        sr = np.sin(self.rpy[0] * 0.5)
+        cp = np.cos(self.rpy[1] * 0.5)
+        sp = np.sin(self.rpy[1] * 0.5)
+        cy = np.cos(self.rpy[2] * 0.5)
+        sy = np.sin(self.rpy[2] * 0.5)
+
+        return Quaternion(
+            qx=sr * cp * cy - cr * sp * sy,
+            qy=cr * sp * cy + sr * cp * sy,
+            qz=cr * cp * sy - sr * sp * cy,
+            qw=cr * cp * cy + sr * sp * sy,
+        )
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def axis_angle(self) -> AngleAxis:
+        """
+        Returns the axis and angle of rotation from quaternion.
+        """
+        q = self.quaternion.vector
+        angle = 2 * np.arccos(q[3])
+        s = np.sqrt(1 - q[3] ** 2)
+        if s < 1e-3:
+            return AngleAxis(axis=Vector(), angle=0.0)
+        axis = Vector(vector=q[:3] / s)
+        return AngleAxis(axis=axis, angle=angle)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def magnitude(self) -> float:
+        """
+        Returns the rotation magnitude based on the quaternion.
+        """
+        return 2 * np.arccos(self.quaternion.qw)
+
     def __mul__(self, other: Self | Vector | NDArray[np.float_]) -> "Vector | Rotation":
         if isinstance(other, Vector):
             return Vector(vector=self.matrix @ other.vector)
@@ -267,17 +338,155 @@ class JointAxis(str, Enum):
     Z = "z"
 
 
+class DHType(str, Enum):
+    STANDARD = "standard"
+    MODIFIED = "modified"
+    HAYATI = "hayati"
+
+
+class DHParameters(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)
+
+    a: float = Field(default=0.0)
+    alpha: float = Field(default=0.0)
+    d: float = Field(default=0.0)
+    theta: float = Field(default=0.0)
+    beta: float = Field(default=0.0)
+    dhtype: DHType = Field(default=DHType.MODIFIED)
+
+    @field_validator("a", "alpha", "d", "theta", "beta", mode="after")
+    @classmethod
+    def round_parameters(cls, v: float) -> float:
+        return round(v, 8)
+
+    @staticmethod
+    def from_matrix(matrix: NDArray, dhtype: DHType = DHType.MODIFIED) -> "DHParameters":
+        match dhtype:
+            case DHType.STANDARD:
+                return DHParameters(
+                    a=matrix[0, 3] / matrix[0, 0],
+                    alpha=np.arccos(matrix[2, 2]),
+                    d=matrix[2, 3],
+                    theta=np.arccos(matrix[0, 0]),
+                    dhtype=dhtype,
+                )
+            case DHType.MODIFIED:
+                return DHParameters(
+                    a=matrix[0, 3],
+                    alpha=np.arccos(matrix[2, 2]),
+                    d=matrix[2, 3] / matrix[2, 2],
+                    theta=np.arccos(matrix[0, 0]),
+                    dhtype=dhtype,
+                )
+            case DHType.HAYATI:
+                beta = np.arctan2(matrix[0, 2], matrix[2, 2])
+                return DHParameters(
+                    a=matrix[0, 3] / np.cos(beta),
+                    alpha=-np.arcsin(matrix[1, 2]),
+                    theta=np.arctan2(matrix[1, 0], matrix[1, 1]),
+                    beta=beta,
+                    dhtype=dhtype,
+                )
+            case _:
+                raise ValueError(f"Invalid DH type {type}")
+
+    def matrix(self) -> NDArray:
+        match self.dhtype:
+            case DHType.STANDARD:
+                return self._standard_matrix()
+            case DHType.MODIFIED:
+                return self._modified_matrix()
+            case DHType.HAYATI:
+                return self._hayati_matrix()
+            case _:
+                raise ValueError(f"Invalid DH type {type}")
+
+    def _standard_matrix(self) -> NDArray:
+        matrix = np.eye(4)
+        matrix[0, 0] = np.cos(self.theta)
+        matrix[0, 1] = -np.sin(self.theta) * np.cos(self.alpha)
+        matrix[0, 2] = np.sin(self.theta) * np.sin(self.alpha)
+        matrix[0, 3] = self.a * np.cos(self.theta)
+        matrix[1, 0] = np.sin(self.theta)
+        matrix[1, 1] = np.cos(self.theta) * np.cos(self.alpha)
+        matrix[1, 2] = -np.cos(self.theta) * np.sin(self.alpha)
+        matrix[1, 3] = self.a * np.sin(self.theta)
+        matrix[2, 1] = np.sin(self.alpha)
+        matrix[2, 2] = np.cos(self.alpha)
+        matrix[2, 3] = self.d
+
+        return matrix
+
+    def _modified_matrix(self) -> NDArray:
+        matrix = np.eye(4)
+        matrix[0, 0] = np.cos(self.theta)
+        matrix[0, 1] = -np.sin(self.theta)
+        matrix[0, 3] = self.a
+        matrix[1, 0] = np.sin(self.theta) * np.cos(self.alpha)
+        matrix[1, 1] = np.cos(self.theta) * np.cos(self.alpha)
+        matrix[1, 2] = -np.sin(self.alpha)
+        matrix[1, 3] = -self.d * np.sin(self.alpha)
+        matrix[2, 0] = np.sin(self.theta) * np.sin(self.alpha)
+        matrix[2, 1] = np.cos(self.theta) * np.sin(self.alpha)
+        matrix[2, 2] = np.cos(self.alpha)
+        matrix[2, 3] = self.d * np.cos(self.alpha)
+
+        return matrix
+
+    def _hayati_matrix(self) -> NDArray:
+        matrix = np.eye(4)
+        matrix[0, 0] = np.sin(self.alpha) * np.sin(self.beta) * np.sin(self.theta) + np.cos(self.beta) * np.cos(
+            self.theta
+        )
+        matrix[0, 1] = np.sin(self.alpha) * np.sin(self.beta) * np.cos(self.theta) - np.sin(self.theta) * np.cos(
+            self.beta
+        )
+        matrix[0, 2] = np.sin(self.beta) * np.cos(self.alpha)
+        matrix[0, 3] = self.a * np.cos(self.beta)
+        matrix[1, 0] = np.sin(self.theta) * np.cos(self.alpha)
+        matrix[1, 1] = np.cos(self.alpha) * np.cos(self.theta)
+        matrix[1, 2] = -np.sin(self.alpha)
+        matrix[2, 0] = np.sin(self.alpha) * np.sin(self.theta) * np.cos(self.beta) - np.sin(self.beta) * np.cos(
+            self.theta
+        )
+        matrix[2, 1] = np.sin(self.alpha) * np.cos(self.beta) * np.cos(self.theta) + np.sin(self.beta) * np.sin(
+            self.theta
+        )
+        matrix[2, 2] = np.cos(self.alpha) * np.cos(self.beta)
+        matrix[2, 3] = -self.a * np.sin(self.beta)
+
+        return matrix
+
+    def to_list(self) -> list[float]:
+        return [self.a, self.alpha, self.d, self.theta]
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, self.__class__):
+            return np.allclose(self.to_list(), other.to_list())
+        else:
+            raise NotImplementedError(f"Cannot compare DHParameters with {other}")
+
+    def __str__(self) -> str:
+        return (
+            f"a: {self.a:.3g} m, alpha: {self.alpha:.3g} rad, d: {self.d:.3g} m, theta: {self.theta:.3g} rad, beta:"
+            f" {self.beta:.3g} rad"
+        )
+
+
 class Joint(BaseModel):
     type: JointType = Field(default=JointType.FIXED, description="Degree of freedom type of the joint")
     axis: JointAxis | None = Field(
         default=None, description="If `type` is other than FIXED, axis of rotation or translation (x, y or z)"
     )
+    value: float | None = Field(default=None, description="Value of the joint in SI Units (meters or radians)")
 
     @model_validator(mode="after")  # type: ignore[misc]
     def _axis_validator(self) -> "Joint":
         match self.type:
             case JointType.FIXED | JointType.SPATIAL:
                 self.axis = None
+                self.value = None
+
         return self
 
     @computed_field  # type: ignore[misc]
@@ -318,14 +527,44 @@ class Transformation(BaseModel):
     )
     joint: Joint = Field(default=Joint(), description="Joint connecting parent and child")
 
+    @model_validator(mode="after")  # type: ignore[misc]
+    def _joint_validator(self) -> "Transformation":
+        match self.joint.type:
+            case JointType.REVOLUTE if self.joint.value is not None:
+                match self.joint.axis:
+                    case JointAxis.X:
+                        self.pose.rotation.rx = self.joint.value
+                    case JointAxis.Y:
+                        self.pose.rotation.ry = self.joint.value
+                    case JointAxis.Z:
+                        self.pose.rotation.rz = self.joint.value
+            case JointType.PRISMATIC if self.joint.value is not None:
+                match self.joint.axis:
+                    case JointAxis.X:
+                        self.pose.translation.x = self.joint.value
+                    case JointAxis.Y:
+                        self.pose.translation.y = self.joint.value
+                    case JointAxis.Z:
+                        self.pose.translation.z = self.joint.value
+        return self
+
     @field_validator("pose", mode="before")
     @classmethod
     def _pose_validator(cls, v: Pose | NDArray[np.float_] | list[float] | dict[str, float]) -> Pose:
         match v:
             case list() | np.ndarray():
                 return Pose.from_list(v)
             case dict():
-                pose_dict = dict()
+                pose_dict = dict(
+                    {
+                        X + M_SUFFIX: 0.0,
+                        Y + M_SUFFIX: 0.0,
+                        Z + M_SUFFIX: 0.0,
+                        RX + RAD_SUFFIX: 0.0,
+                        RY + RAD_SUFFIX: 0.0,
+                        RZ + RAD_SUFFIX: 0.0,
+                    }
+                )
                 for key, value in v.items():
                     if key.endswith(MM_SUFFIX):
                         pose_dict[key.replace(MM_SUFFIX, M_SUFFIX)] = value / 1000
@@ -378,6 +617,11 @@ def inv(self) -> "Transformation":
             pose=Pose(translation=Vector(vector=new_pos), rotation=new_rot),
             parent=self.child,
             child=self.parent,
+            joint=Joint(
+                type=self.joint.type,
+                axis=self.joint.axis,
+                value=None if self.joint.value is None else -self.joint.value,
+            ),
         )
 
     def reset_rotation(self) -> "Transformation":