foodforthought-cli 0.2.8__py3-none-any.whl → 0.3.1__py3-none-any.whl

ate/urdf/motion_analysis.py

@@ -0,0 +1,396 @@
+"""
+Motion analysis for joint discovery from point cloud trajectories.
+
+This module provides tools for analyzing how links move over time:
+- Trajectory extraction from per-frame point clouds
+- Joint type classification (revolute vs prismatic)
+- Motion correlation for parent-child relationship discovery
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    from scipy.spatial.transform import Rotation
+    DEPS_AVAILABLE = True
+except ImportError:
+    DEPS_AVAILABLE = False
+    np = None
+    Rotation = None
+
+
+class MotionAnalysisError(Exception):
+    """Error during motion analysis."""
+    pass
+
+
+@dataclass
+class LinkTrajectory:
+    """Trajectory of a single link over time."""
+    name: str
+    frame_indices: List[int]
+    centroids: "np.ndarray"  # (T, 3) positions
+    orientations: Optional["np.ndarray"] = None  # (T, 3, 3) rotation matrices
+
+    @property
+    def n_frames(self) -> int:
+        return len(self.frame_indices)
+
+    def get_velocity(self) -> "np.ndarray":
+        """Compute velocity (change in position per frame)."""
+        if self.n_frames < 2:
+            return np.zeros((0, 3))
+        return np.diff(self.centroids, axis=0)
+
+    def get_total_displacement(self) -> float:
+        """Get total path length traveled."""
+        velocity = self.get_velocity()
+        return np.sum(np.linalg.norm(velocity, axis=1))
+
+
+@dataclass
+class RelativeMotion:
+    """Relative motion between two links."""
+    parent_name: str
+    child_name: str
+    frame_indices: List[int]
+    relative_positions: "np.ndarray"  # (T, 3) child position in parent frame
+    relative_rotations: "np.ndarray"  # (T, 3, 3) child rotation in parent frame
+
+    def analyze_joint_type(self) -> Tuple[str, float]:
+        """
+        Classify joint type based on relative motion patterns.
+
+        Returns:
+            Tuple of (joint_type, confidence)
+            joint_type: "revolute", "prismatic", or "fixed"
+        """
+        n = len(self.frame_indices)
+        if n < 3:
+            return "fixed", 1.0
+
+        # Compute translation and rotation magnitudes
+        translations = np.linalg.norm(self.relative_positions - self.relative_positions[0], axis=1)
+        max_translation = np.max(translations)
+
+        # Compute rotation angles from rotation matrices
+        angles = []
+        for R in self.relative_rotations:
+            R0 = self.relative_rotations[0]
+            R_rel = R @ R0.T
+            angle = np.arccos(np.clip((np.trace(R_rel) - 1) / 2, -1, 1))
+            angles.append(angle)
+        angles = np.array(angles)
+        max_rotation = np.max(angles)
+
+        # Classification thresholds
+        TRANSLATION_THRESHOLD = 0.01  # 1cm
+        ROTATION_THRESHOLD = 0.05  # ~3 degrees
+
+        if max_translation < TRANSLATION_THRESHOLD and max_rotation < ROTATION_THRESHOLD:
+            return "fixed", 0.9
+
+        # Ratio of rotation to translation determines joint type
+        # Revolute: mostly rotation, minimal translation
+        # Prismatic: mostly translation, minimal rotation
+        if max_rotation > ROTATION_THRESHOLD and max_translation < 0.05:
+            return "revolute", 0.8
+        elif max_translation > TRANSLATION_THRESHOLD and max_rotation < 0.1:
+            return "prismatic", 0.8
+        else:
+            # Mixed motion - default to revolute as most common
+            return "revolute", 0.6
+
+
+def extract_trajectories(
+    link_clouds: Dict[str, List["np.ndarray"]],
+    frame_indices: List[int],
+) -> Dict[str, LinkTrajectory]:
+    """
+    Extract link trajectories from per-frame point clouds.
+
+    Args:
+        link_clouds: Dict mapping link_name -> list of point clouds per frame
+        frame_indices: List of frame indices
+
+    Returns:
+        Dict mapping link_name -> LinkTrajectory
+    """
+    if not DEPS_AVAILABLE:
+        raise MotionAnalysisError("NumPy/SciPy not available")
+
+    trajectories = {}
+
+    for link_name, clouds in link_clouds.items():
+        if len(clouds) != len(frame_indices):
+            logger.warning(f"Mismatched cloud count for {link_name}: {len(clouds)} vs {len(frame_indices)}")
+            continue
+
+        # Compute centroids
+        centroids = []
+        valid_indices = []
+
+        for i, cloud in enumerate(clouds):
+            if cloud is not None and len(cloud) > 0:
+                centroid = np.mean(cloud, axis=0)
+                centroids.append(centroid)
+                valid_indices.append(frame_indices[i])
+
+        if len(centroids) < 2:
+            logger.warning(f"Insufficient data for trajectory of '{link_name}'")
+            continue
+
+        trajectories[link_name] = LinkTrajectory(
+            name=link_name,
+            frame_indices=valid_indices,
+            centroids=np.array(centroids),
+            orientations=None,  # Could be computed with PCA if needed
+        )
+
+    return trajectories
+
+
+def compute_relative_motion(
+    parent_traj: LinkTrajectory,
+    child_traj: LinkTrajectory,
+) -> RelativeMotion:
+    """
+    Compute relative motion of child link w.r.t. parent link.
+
+    Args:
+        parent_traj: Parent link trajectory
+        child_traj: Child link trajectory
+
+    Returns:
+        RelativeMotion describing child motion in parent frame
+    """
+    if not DEPS_AVAILABLE:
+        raise MotionAnalysisError("NumPy not available")
+
+    # Find common frames
+    parent_frames = set(parent_traj.frame_indices)
+    child_frames = set(child_traj.frame_indices)
+    common_frames = sorted(parent_frames & child_frames)
+
+    if len(common_frames) < 2:
+        raise MotionAnalysisError(
+            f"Insufficient common frames between {parent_traj.name} and {child_traj.name}"
+        )
+
+    # Get positions at common frames
+    parent_pos = []
+    child_pos = []
+
+    for frame in common_frames:
+        p_idx = parent_traj.frame_indices.index(frame)
+        c_idx = child_traj.frame_indices.index(frame)
+        parent_pos.append(parent_traj.centroids[p_idx])
+        child_pos.append(child_traj.centroids[c_idx])
+
+    parent_pos = np.array(parent_pos)
+    child_pos = np.array(child_pos)
+
+    # Compute relative positions (child - parent)
+    relative_positions = child_pos - parent_pos
+
+    # For now, assume identity rotations (orientation estimation requires more data)
+    n = len(common_frames)
+    relative_rotations = np.tile(np.eye(3), (n, 1, 1))
+
+    return RelativeMotion(
+        parent_name=parent_traj.name,
+        child_name=child_traj.name,
+        frame_indices=common_frames,
+        relative_positions=relative_positions,
+        relative_rotations=relative_rotations,
+    )
+
+
+def estimate_revolute_axis(
+    relative_motion: RelativeMotion,
+) -> Tuple["np.ndarray", "np.ndarray", float]:
+    """
+    Estimate revolute joint axis from relative motion.
+
+    Uses SVD to find the best-fit rotation axis that explains
+    the child link's motion relative to the parent.
+
+    Args:
+        relative_motion: RelativeMotion between parent and child
+
+    Returns:
+        Tuple of (axis, pivot_point, confidence):
+        - axis: (3,) unit vector of rotation axis
+        - pivot_point: (3,) point on the rotation axis
+        - confidence: Quality of axis estimation (0-1)
+    """
+    if not DEPS_AVAILABLE:
+        raise MotionAnalysisError("NumPy not available")
+
+    positions = relative_motion.relative_positions
+    n = len(positions)
+
+    if n < 3:
+        # Not enough data - return default Y axis
+        return np.array([0, 1, 0]), positions[0], 0.3
+
+    # Center the positions
+    mean_pos = np.mean(positions, axis=0)
+    centered = positions - mean_pos
+
+    # SVD to find principal components
+    U, S, Vt = np.linalg.svd(centered)
+
+    # For a revolute joint, motion is planar perpendicular to the axis
+    # The axis is the direction with minimum variance (last singular value)
+    axis = Vt[-1]  # Direction with smallest variance
+
+    # Ensure consistent axis direction (positive Z preferred)
+    if axis[2] < 0:
+        axis = -axis
+
+    # Pivot point is at the mean position
+    pivot = mean_pos
+
+    # Confidence based on how planar the motion is
+    # Ratio of smallest to middle singular value
+    if S[1] > 1e-6:
+        planarity = 1 - (S[2] / S[1])
+        confidence = np.clip(planarity, 0.3, 1.0)
+    else:
+        confidence = 0.5
+
+    return axis.astype(np.float64), pivot.astype(np.float64), confidence
+
+
+def estimate_prismatic_axis(
+    relative_motion: RelativeMotion,
+) -> Tuple["np.ndarray", float]:
+    """
+    Estimate prismatic joint axis from relative motion.
+
+    Uses SVD to find the dominant direction of translation.
+
+    Args:
+        relative_motion: RelativeMotion between parent and child
+
+    Returns:
+        Tuple of (axis, confidence):
+        - axis: (3,) unit vector of translation direction
+        - confidence: Quality of axis estimation (0-1)
+    """
+    if not DEPS_AVAILABLE:
+        raise MotionAnalysisError("NumPy not available")
+
+    positions = relative_motion.relative_positions
+    n = len(positions)
+
+    if n < 2:
+        return np.array([0, 0, 1]), 0.3
+
+    # Compute displacement from first position
+    displacements = positions - positions[0]
+
+    # SVD to find principal direction
+    U, S, Vt = np.linalg.svd(displacements)
+
+    # Dominant translation direction is first singular vector
+    axis = Vt[0]
+
+    # Ensure consistent direction
+    if np.sum(displacements[-1] * axis) < 0:
+        axis = -axis
+
+    # Confidence based on how linear the motion is
+    if S[0] > 1e-6:
+        linearity = 1 - (S[1] / S[0]) if len(S) > 1 else 1.0
+        confidence = np.clip(linearity, 0.3, 1.0)
+    else:
+        confidence = 0.5
+
+    return axis.astype(np.float64), confidence
+
+
+def estimate_joint_limits(
+    relative_motion: RelativeMotion,
+    joint_type: str,
+    axis: "np.ndarray",
+    pivot: Optional["np.ndarray"] = None,
+) -> Tuple[float, float]:
+    """
+    Estimate joint limits from observed motion range.
+
+    Args:
+        relative_motion: RelativeMotion data
+        joint_type: "revolute" or "prismatic"
+        axis: Joint axis
+        pivot: Pivot point (for revolute joints)
+
+    Returns:
+        Tuple of (lower_limit, upper_limit)
+    """
+    if not DEPS_AVAILABLE:
+        raise MotionAnalysisError("NumPy not available")
+
+    positions = relative_motion.relative_positions
+
+    if joint_type == "revolute":
+        # Project positions onto plane perpendicular to axis
+        # Then measure angles
+        if pivot is None:
+            pivot = np.mean(positions, axis=0)
+
+        relative = positions - pivot
+
+        # Project onto plane
+        axis_component = np.outer(np.dot(relative, axis), axis)
+        in_plane = relative - axis_component
+
+        # Compute angles (relative to first position)
+        if len(in_plane) < 2:
+            return -np.pi, np.pi
+
+        ref = in_plane[0]
+        ref_norm = np.linalg.norm(ref)
+        if ref_norm < 1e-6:
+            return -np.pi, np.pi
+
+        ref = ref / ref_norm
+        angles = []
+
+        for pos in in_plane:
+            pos_norm = np.linalg.norm(pos)
+            if pos_norm < 1e-6:
+                angles.append(0)
+                continue
+            pos = pos / pos_norm
+            angle = np.arctan2(np.dot(np.cross(ref, pos), axis), np.dot(ref, pos))
+            angles.append(angle)
+
+        angles = np.array(angles)
+
+        # Add some margin to observed range
+        margin = 0.1  # ~6 degrees
+        return float(np.min(angles) - margin), float(np.max(angles) + margin)
+
+    else:  # prismatic
+        # Project onto axis and measure displacement
+        displacements = np.dot(positions - positions[0], axis)
+        margin = 0.01  # 1cm
+        return float(np.min(displacements) - margin), float(np.max(displacements) + margin)
+
+
+__all__ = [
+    "MotionAnalysisError",
+    "LinkTrajectory",
+    "RelativeMotion",
+    "extract_trajectories",
+    "compute_relative_motion",
+    "estimate_revolute_axis",
+    "estimate_prismatic_axis",
+    "estimate_joint_limits",
+]