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
This commit is contained in:
@@ -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",
|
||||
|
||||
517
solver/datagen/jacobian.py
Normal file
517
solver/datagen/jacobian.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user