From 455b6318d94951fa8fb0709bb8e47be9478d3153 Mon Sep 17 00:00:00 2001 From: forbes-0023 Date: Mon, 2 Feb 2026 13:50:16 -0600 Subject: [PATCH] feat: port JacobianVerifier to solver/datagen/jacobian.py Port the constraint Jacobian builder and numerical rank verifier from data/synthetic/pebble-game.py. All 11 joint type builders, SVD rank computation, and incremental dependency detection. - Full type annotations (mypy strict) - Ruff lint and format clean - Re-exported from solver.datagen.__init__ Closes #3 --- solver/datagen/__init__.py | 2 + solver/datagen/jacobian.py | 517 +++++++++++++++++++++++++++++++++++++ 2 files changed, 519 insertions(+) create mode 100644 solver/datagen/jacobian.py diff --git a/solver/datagen/__init__.py b/solver/datagen/__init__.py index 6504939..a48a988 100644 --- a/solver/datagen/__init__.py +++ b/solver/datagen/__init__.py @@ -1,5 +1,6 @@ """Data generation utilities for assembly constraint training data.""" +from solver.datagen.jacobian import JacobianVerifier from solver.datagen.pebble_game import PebbleGame3D from solver.datagen.types import ( ConstraintAnalysis, @@ -11,6 +12,7 @@ from solver.datagen.types import ( __all__ = [ "ConstraintAnalysis", + "JacobianVerifier", "Joint", "JointType", "PebbleGame3D", diff --git a/solver/datagen/jacobian.py b/solver/datagen/jacobian.py new file mode 100644 index 0000000..09bb769 --- /dev/null +++ b/solver/datagen/jacobian.py @@ -0,0 +1,517 @@ +"""Numerical Jacobian rank verification for assembly constraint analysis. + +Builds the constraint Jacobian matrix and analyzes its numerical rank +to detect geometric degeneracies that the combinatorial pebble game +cannot identify (e.g., parallel revolute axes creating hidden dependencies). + +References: + - Chappuis, "Constraints Derivation for Rigid Body Simulation in 3D" +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from solver.datagen.types import Joint, JointType, RigidBody + +if TYPE_CHECKING: + from typing import Any + +__all__ = ["JacobianVerifier"] + + +class JacobianVerifier: + """Builds and analyzes the constraint Jacobian for numerical rank check. + + The pebble game gives a combinatorial *necessary* condition for + rigidity. However, geometric special cases (e.g., all revolute axes + parallel, creating a hidden dependency) require numerical verification. + + For each joint, we construct the constraint Jacobian rows that map + the 6n-dimensional generalized velocity vector to the constraint + violation rates. The rank of this Jacobian equals the number of + truly independent constraints. + + The generalized velocity vector for n bodies is:: + + v = [v1_x, v1_y, v1_z, w1_x, w1_y, w1_z, ..., vn_x, ..., wn_z] + + Each scalar constraint C_i contributes one row to J such that:: + + dC_i/dt = J_i @ v = 0 + """ + + def __init__(self, bodies: list[RigidBody]) -> None: + self.bodies = {b.body_id: b for b in bodies} + self.body_index = {b.body_id: i for i, b in enumerate(bodies)} + self.n_bodies = len(bodies) + self.jacobian_rows: list[np.ndarray] = [] + self.row_labels: list[dict[str, Any]] = [] + + def _body_cols(self, body_id: int) -> tuple[int, int]: + """Return the column range [start, end) for a body in J.""" + idx = self.body_index[body_id] + return idx * 6, (idx + 1) * 6 + + def add_joint_constraints(self, joint: Joint) -> int: + """Add Jacobian rows for all scalar constraints of a joint. + + Returns the number of rows added. + """ + builder = { + JointType.FIXED: self._build_fixed, + JointType.REVOLUTE: self._build_revolute, + JointType.CYLINDRICAL: self._build_cylindrical, + JointType.SLIDER: self._build_slider, + JointType.BALL: self._build_ball, + JointType.PLANAR: self._build_planar, + JointType.DISTANCE: self._build_distance, + JointType.PARALLEL: self._build_parallel, + JointType.PERPENDICULAR: self._build_perpendicular, + JointType.UNIVERSAL: self._build_universal, + JointType.SCREW: self._build_screw, + } + + rows_before = len(self.jacobian_rows) + builder[joint.joint_type](joint) + return len(self.jacobian_rows) - rows_before + + def _make_row(self) -> np.ndarray: + """Create a zero row of width 6*n_bodies.""" + return np.zeros(6 * self.n_bodies) + + def _skew(self, v: np.ndarray) -> np.ndarray: + """Skew-symmetric matrix for cross product: ``skew(v) @ w = v x w``.""" + return np.array( + [ + [0, -v[2], v[1]], + [v[2], 0, -v[0]], + [-v[1], v[0], 0], + ] + ) + + # --- Ball-and-socket (spherical) joint: 3 translation constraints --- + + def _build_ball(self, joint: Joint) -> None: + """Ball joint: coincident point constraint. + + ``C_trans = (x_b + R_b @ r_b) - (x_a + R_a @ r_a) = 0`` + (3 equations) + + Jacobian rows (for each of x, y, z): + body_a linear: -I + body_a angular: +skew(R_a @ r_a) + body_b linear: +I + body_b angular: -skew(R_b @ r_b) + """ + # Use anchor positions directly as world-frame offsets + r_a = joint.anchor_a - self.bodies[joint.body_a].position + r_b = joint.anchor_b - self.bodies[joint.body_b].position + + col_a_start, col_a_end = self._body_cols(joint.body_a) + col_b_start, col_b_end = self._body_cols(joint.body_b) + + for axis_idx in range(3): + row = self._make_row() + e = np.zeros(3) + e[axis_idx] = 1.0 + + row[col_a_start : col_a_start + 3] = -e + row[col_a_start + 3 : col_a_end] = np.cross(r_a, e) + + row[col_b_start : col_b_start + 3] = e + row[col_b_start + 3 : col_b_end] = -np.cross(r_b, e) + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "ball_translation", + "axis": axis_idx, + } + ) + + # --- Fixed joint: 3 translation + 3 rotation constraints --- + + def _build_fixed(self, joint: Joint) -> None: + """Fixed joint = ball joint + locked rotation. + + Translation part: same as ball joint (3 rows). + Rotation part: relative angular velocity must be zero (3 rows). + """ + self._build_ball(joint) + + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + for axis_idx in range(3): + row = self._make_row() + row[col_a_start + 3 + axis_idx] = -1.0 + row[col_b_start + 3 + axis_idx] = 1.0 + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "fixed_rotation", + "axis": axis_idx, + } + ) + + # --- Revolute (hinge) joint: 3 translation + 2 rotation constraints --- + + def _build_revolute(self, joint: Joint) -> None: + """Revolute joint: rotation only about one axis. + + Translation: same as ball (3 rows). + Rotation: relative angular velocity must be parallel to hinge axis. + """ + self._build_ball(joint) + + axis = joint.axis / np.linalg.norm(joint.axis) + t1, t2 = self._perpendicular_pair(axis) + + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + for i, t in enumerate((t1, t2)): + row = self._make_row() + row[col_a_start + 3 : col_a_start + 6] = -t + row[col_b_start + 3 : col_b_start + 6] = t + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "revolute_rotation", + "perp_axis": i, + } + ) + + # --- Cylindrical joint: 2 translation + 2 rotation constraints --- + + def _build_cylindrical(self, joint: Joint) -> None: + """Cylindrical joint: allows rotation + translation along one axis. + + Translation: constrain motion perpendicular to axis (2 rows). + Rotation: constrain rotation perpendicular to axis (2 rows). + """ + axis = joint.axis / np.linalg.norm(joint.axis) + t1, t2 = self._perpendicular_pair(axis) + + r_a = joint.anchor_a - self.bodies[joint.body_a].position + r_b = joint.anchor_b - self.bodies[joint.body_b].position + + col_a_start, col_a_end = self._body_cols(joint.body_a) + col_b_start, col_b_end = self._body_cols(joint.body_b) + + for i, t in enumerate((t1, t2)): + row = self._make_row() + row[col_a_start : col_a_start + 3] = -t + row[col_a_start + 3 : col_a_end] = np.cross(r_a, t) + row[col_b_start : col_b_start + 3] = t + row[col_b_start + 3 : col_b_end] = -np.cross(r_b, t) + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "cylindrical_translation", + "perp_axis": i, + } + ) + + for i, t in enumerate((t1, t2)): + row = self._make_row() + row[col_a_start + 3 : col_a_start + 6] = -t + row[col_b_start + 3 : col_b_start + 6] = t + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "cylindrical_rotation", + "perp_axis": i, + } + ) + + # --- Slider (prismatic) joint: 2 translation + 3 rotation constraints --- + + def _build_slider(self, joint: Joint) -> None: + """Slider/prismatic joint: translation along one axis only. + + Translation: perpendicular translation constrained (2 rows). + Rotation: all relative rotation constrained (3 rows). + """ + axis = joint.axis / np.linalg.norm(joint.axis) + t1, t2 = self._perpendicular_pair(axis) + + r_a = joint.anchor_a - self.bodies[joint.body_a].position + r_b = joint.anchor_b - self.bodies[joint.body_b].position + + col_a_start, col_a_end = self._body_cols(joint.body_a) + col_b_start, col_b_end = self._body_cols(joint.body_b) + + for i, t in enumerate((t1, t2)): + row = self._make_row() + row[col_a_start : col_a_start + 3] = -t + row[col_a_start + 3 : col_a_end] = np.cross(r_a, t) + row[col_b_start : col_b_start + 3] = t + row[col_b_start + 3 : col_b_end] = -np.cross(r_b, t) + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "slider_translation", + "perp_axis": i, + } + ) + + for axis_idx in range(3): + row = self._make_row() + row[col_a_start + 3 + axis_idx] = -1.0 + row[col_b_start + 3 + axis_idx] = 1.0 + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "slider_rotation", + "axis": axis_idx, + } + ) + + # --- Planar joint: 1 translation + 2 rotation constraints --- + + def _build_planar(self, joint: Joint) -> None: + """Planar joint: constrains to a plane. + + Translation: motion along plane normal constrained (1 row). + Rotation: rotation about axes in the plane constrained (2 rows). + """ + normal = joint.axis / np.linalg.norm(joint.axis) + t1, t2 = self._perpendicular_pair(normal) + + r_a = joint.anchor_a - self.bodies[joint.body_a].position + r_b = joint.anchor_b - self.bodies[joint.body_b].position + + col_a_start, col_a_end = self._body_cols(joint.body_a) + col_b_start, col_b_end = self._body_cols(joint.body_b) + + row = self._make_row() + row[col_a_start : col_a_start + 3] = -normal + row[col_a_start + 3 : col_a_end] = np.cross(r_a, normal) + row[col_b_start : col_b_start + 3] = normal + row[col_b_start + 3 : col_b_end] = -np.cross(r_b, normal) + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "planar_translation", + } + ) + + for i, t in enumerate((t1, t2)): + row = self._make_row() + row[col_a_start + 3 : col_a_start + 6] = -t + row[col_b_start + 3 : col_b_start + 6] = t + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "planar_rotation", + "perp_axis": i, + } + ) + + # --- Distance constraint: 1 scalar --- + + def _build_distance(self, joint: Joint) -> None: + """Distance constraint: ``||p_b - p_a|| = d``. + + Single row: ``direction . (v_b + w_b x r_b - v_a - w_a x r_a) = 0`` + where ``direction = normalized(p_b - p_a)``. + """ + p_a = joint.anchor_a + p_b = joint.anchor_b + diff = p_b - p_a + dist = np.linalg.norm(diff) + direction = np.array([1.0, 0.0, 0.0]) if dist < 1e-12 else diff / dist + + r_a = joint.anchor_a - self.bodies[joint.body_a].position + r_b = joint.anchor_b - self.bodies[joint.body_b].position + + col_a_start, col_a_end = self._body_cols(joint.body_a) + col_b_start, col_b_end = self._body_cols(joint.body_b) + + row = self._make_row() + row[col_a_start : col_a_start + 3] = -direction + row[col_a_start + 3 : col_a_end] = np.cross(r_a, direction) + row[col_b_start : col_b_start + 3] = direction + row[col_b_start + 3 : col_b_end] = -np.cross(r_b, direction) + + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "distance", + } + ) + + # --- Parallel constraint: 3 rotation constraints --- + + def _build_parallel(self, joint: Joint) -> None: + """Parallel: all relative rotation constrained (same as fixed rotation). + + In practice only 2 of 3 are independent for a single axis, but + we emit 3 and let the rank check sort it out. + """ + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + for axis_idx in range(3): + row = self._make_row() + row[col_a_start + 3 + axis_idx] = -1.0 + row[col_b_start + 3 + axis_idx] = 1.0 + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "parallel_rotation", + "axis": axis_idx, + } + ) + + # --- Perpendicular constraint: 1 angular --- + + def _build_perpendicular(self, joint: Joint) -> None: + """Perpendicular: single dot-product angular constraint.""" + axis = joint.axis / np.linalg.norm(joint.axis) + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + row = self._make_row() + row[col_a_start + 3 : col_a_start + 6] = -axis + row[col_b_start + 3 : col_b_start + 6] = axis + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "perpendicular", + } + ) + + # --- Universal (Cardan) joint: 3 translation + 1 rotation --- + + def _build_universal(self, joint: Joint) -> None: + """Universal joint: ball + one rotation constraint. + + Allows rotation about two axes, constrains rotation about the third. + """ + self._build_ball(joint) + + axis = joint.axis / np.linalg.norm(joint.axis) + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + row = self._make_row() + row[col_a_start + 3 : col_a_start + 6] = -axis + row[col_b_start + 3 : col_b_start + 6] = axis + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "universal_rotation", + } + ) + + # --- Screw (helical) joint: 2 translation + 2 rotation + 1 coupled --- + + def _build_screw(self, joint: Joint) -> None: + """Screw joint: coupled rotation-translation along axis. + + Like cylindrical but with a coupling constraint: + ``v_axial - pitch * w_axial = 0`` + """ + self._build_cylindrical(joint) + + axis = joint.axis / np.linalg.norm(joint.axis) + col_a_start, _ = self._body_cols(joint.body_a) + col_b_start, _ = self._body_cols(joint.body_b) + + row = self._make_row() + row[col_a_start : col_a_start + 3] = -axis + row[col_b_start : col_b_start + 3] = axis + row[col_a_start + 3 : col_a_start + 6] = joint.pitch * axis + row[col_b_start + 3 : col_b_start + 6] = -joint.pitch * axis + self.jacobian_rows.append(row) + self.row_labels.append( + { + "joint_id": joint.joint_id, + "type": "screw_coupling", + } + ) + + # --- Utilities --- + + def _perpendicular_pair(self, axis: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Generate two unit vectors perpendicular to *axis* and each other.""" + if abs(axis[0]) < 0.9: + t1 = np.cross(axis, np.array([1.0, 0, 0])) + else: + t1 = np.cross(axis, np.array([0, 1.0, 0])) + t1 /= np.linalg.norm(t1) + t2 = np.cross(axis, t1) + t2 /= np.linalg.norm(t2) + return t1, t2 + + def get_jacobian(self) -> np.ndarray: + """Return the full constraint Jacobian matrix.""" + if not self.jacobian_rows: + return np.zeros((0, 6 * self.n_bodies)) + return np.array(self.jacobian_rows) + + def numerical_rank(self, tol: float = 1e-8) -> int: + """Compute the numerical rank of the constraint Jacobian via SVD. + + This is the number of truly independent scalar constraints, + accounting for geometric degeneracies that the combinatorial + pebble game cannot detect. + """ + j = self.get_jacobian() + if j.size == 0: + return 0 + sv = np.linalg.svd(j, compute_uv=False) + return int(np.sum(sv > tol)) + + def find_dependencies(self, tol: float = 1e-8) -> list[int]: + """Identify which constraint rows are numerically dependent. + + Returns indices of rows that can be removed without changing + the Jacobian's rank. + """ + j = self.get_jacobian() + if j.size == 0: + return [] + + n_rows = j.shape[0] + dependent: list[int] = [] + + current = np.zeros((0, j.shape[1])) + current_rank = 0 + + for i in range(n_rows): + candidate = np.vstack([current, j[i : i + 1, :]]) if current.size else j[i : i + 1, :] + sv = np.linalg.svd(candidate, compute_uv=False) + new_rank = int(np.sum(sv > tol)) + + if new_rank > current_rank: + current = candidate + current_rank = new_rank + else: + dependent.append(i) + + return dependent