solver/kindred_solver/constraints.py

"""Constraint classes that produce residual Expr lists.

Each constraint takes two RigidBody entities and marker transforms,
then generates residual expressions that equal zero when satisfied.
"""

from __future__ import annotations

from typing import List

from .entities import RigidBody
from .expr import Const, Expr


class ConstraintBase:
    """Abstract constraint producing residual expressions."""

    def residuals(self) -> List[Expr]:
        raise NotImplementedError


class CoincidentConstraint(ConstraintBase):
    """Point-on-point: marker origin on body_i == marker origin on body_j.

    3 residuals: dx, dy, dz.
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos

    def residuals(self) -> List[Expr]:
        wx_i, wy_i, wz_i = self.body_i.world_point(*self.marker_i_pos)
        wx_j, wy_j, wz_j = self.body_j.world_point(*self.marker_j_pos)
        return [wx_i - wx_j, wy_i - wy_j, wz_i - wz_j]


class DistancePointPointConstraint(ConstraintBase):
    """Point-to-point distance: ||p_i - p_j||^2 = d^2.

    1 residual (squared form avoids sqrt in Jacobian).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        distance: float,
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_j_pos = marker_j_pos
        self.distance = distance

    def residuals(self) -> List[Expr]:
        wx_i, wy_i, wz_i = self.body_i.world_point(*self.marker_i_pos)
        wx_j, wy_j, wz_j = self.body_j.world_point(*self.marker_j_pos)
        dx = wx_i - wx_j
        dy = wy_i - wy_j
        dz = wz_i - wz_j
        d2 = dx * dx + dy * dy + dz * dz
        return [d2 - Const(self.distance * self.distance)]


class FixedConstraint(ConstraintBase):
    """Lock all 6 DOF of body_j relative to body_i.

    Position: 3 residuals (marker origins coincide).
    Orientation: 3 residuals from relative quaternion (imaginary part == 0
    when orientations match).  The quaternion normalization constraint
    handles the 4th equation.

    The reference relative quaternion q_ref = conj(q_i_marker) * q_j_marker
    is computed from the marker quaternions at constraint creation time.
    At solve time:  q_rel = conj(q_i_world * q_i_marker) * (q_j_world * q_j_marker)
    Residual = imaginary part of (conj(q_ref) * q_rel).
    """

    def __init__(
        self,
        body_i: RigidBody,
        marker_i_pos: tuple[float, float, float],
        marker_i_quat: tuple[float, float, float, float],
        body_j: RigidBody,
        marker_j_pos: tuple[float, float, float],
        marker_j_quat: tuple[float, float, float, float],
    ):
        self.body_i = body_i
        self.body_j = body_j
        self.marker_i_pos = marker_i_pos
        self.marker_i_quat = marker_i_quat
        self.marker_j_pos = marker_j_pos
        self.marker_j_quat = marker_j_quat

    def residuals(self) -> List[Expr]:
        # Position residuals — same as Coincident
        wx_i, wy_i, wz_i = self.body_i.world_point(*self.marker_i_pos)
        wx_j, wy_j, wz_j = self.body_j.world_point(*self.marker_j_pos)
        pos_res = [wx_i - wx_j, wy_i - wy_j, wz_i - wz_j]

        # Orientation residuals via quaternion error
        # q_i_total = q_body_i * q_marker_i  (as Expr)
        # q_j_total = q_body_j * q_marker_j
        # q_error = conj(q_i_total) * q_j_total
        # We want q_error ≈ identity → imaginary parts ≈ 0
        qi = _quat_mul_const(
            self.body_i.qw,
            self.body_i.qx,
            self.body_i.qy,
            self.body_i.qz,
            *self.marker_i_quat,
        )
        qj = _quat_mul_const(
            self.body_j.qw,
            self.body_j.qx,
            self.body_j.qy,
            self.body_j.qz,
            *self.marker_j_quat,
        )
        # conj(qi) * qj
        err = _quat_mul_expr(
            qi[0],
            -qi[1],
            -qi[2],
            -qi[3],  # conjugate
            qj[0],
            qj[1],
            qj[2],
            qj[3],
        )
        # err should be (1, 0, 0, 0) — residuals are the imaginary parts
        ori_res = [err[1], err[2], err[3]]

        return pos_res + ori_res


# -- quaternion multiplication helpers ----------------------------------------


def _quat_mul_const(
    aw: Expr,
    ax: Expr,
    ay: Expr,
    az: Expr,
    bw: float,
    bx: float,
    by: float,
    bz: float,
) -> tuple[Expr, Expr, Expr, Expr]:
    """Multiply Expr quaternion (a) by constant quaternion (b)."""
    cbw, cbx, cby, cbz = Const(bw), Const(bx), Const(by), Const(bz)
    return _quat_mul_expr(aw, ax, ay, az, cbw, cbx, cby, cbz)


def _quat_mul_expr(
    aw: Expr,
    ax: Expr,
    ay: Expr,
    az: Expr,
    bw: Expr,
    bx: Expr,
    by: Expr,
    bz: Expr,
) -> tuple[Expr, Expr, Expr, Expr]:
    """Hamilton product of two Expr quaternions."""
    rw = aw * bw - ax * bx - ay * by - az * bz
    rx = aw * bx + ax * bw + ay * bz - az * by
    ry = aw * by - ax * bz + ay * bw + az * bx
    rz = aw * bz + ax * by - ay * bx + az * bw
    return rw, rx, ry, rz