diff --git a/solver/datagen/__init__.py b/solver/datagen/__init__.py index a48a988..7d64bc0 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.analysis import analyze_assembly from solver.datagen.jacobian import JacobianVerifier from solver.datagen.pebble_game import PebbleGame3D from solver.datagen.types import ( @@ -18,4 +19,5 @@ __all__ = [ "PebbleGame3D", "PebbleState", "RigidBody", + "analyze_assembly", ] diff --git a/solver/datagen/analysis.py b/solver/datagen/analysis.py new file mode 100644 index 0000000..f862802 --- /dev/null +++ b/solver/datagen/analysis.py @@ -0,0 +1,130 @@ +"""Combined pebble game + Jacobian verification analysis. + +Provides :func:`analyze_assembly`, the main entry point for full rigidity +analysis of an assembly using both combinatorial and numerical methods. +""" + +from __future__ import annotations + +import numpy as np + +from solver.datagen.jacobian import JacobianVerifier +from solver.datagen.pebble_game import PebbleGame3D +from solver.datagen.types import ( + ConstraintAnalysis, + Joint, + JointType, + RigidBody, +) + +__all__ = ["analyze_assembly"] + +_GROUND_ID = -1 + + +def analyze_assembly( + bodies: list[RigidBody], + joints: list[Joint], + ground_body: int | None = None, +) -> ConstraintAnalysis: + """Full rigidity analysis of an assembly using both methods. + + Args: + bodies: List of rigid bodies in the assembly. + joints: List of joints connecting bodies. + ground_body: If set, this body is fixed (adds 6 implicit constraints). + + Returns: + ConstraintAnalysis with combinatorial and numerical results. + """ + # --- Pebble Game --- + pg = PebbleGame3D() + all_edge_results = [] + + # Add a virtual ground body (id=-1) if grounding is requested. + # Grounding body X means adding a fixed joint between X and + # the virtual ground. This properly lets the pebble game account + # for the 6 removed DOF without breaking invariants. + if ground_body is not None: + pg.add_body(_GROUND_ID) + + for body in bodies: + pg.add_body(body.body_id) + + if ground_body is not None: + ground_joint = Joint( + joint_id=-1, + body_a=ground_body, + body_b=_GROUND_ID, + joint_type=JointType.FIXED, + anchor_a=bodies[0].position if bodies else np.zeros(3), + anchor_b=bodies[0].position if bodies else np.zeros(3), + ) + pg.add_joint(ground_joint) + # Don't include ground joint edges in the output labels + # (they're infrastructure, not user constraints) + + for joint in joints: + results = pg.add_joint(joint) + all_edge_results.extend(results) + + combinatorial_independent = len(pg.state.independent_edges) + grounded = ground_body is not None + + # The virtual ground body contributes 6 pebbles to the total. + # Subtract those from the reported DOF for user-facing numbers. + raw_dof = pg.get_dof() + ground_offset = 6 if grounded else 0 + effective_dof = raw_dof - ground_offset + effective_internal_dof = max(0, effective_dof - (0 if grounded else 6)) + + combinatorial_classification = pg.classify_assembly(grounded=grounded) + + # --- Jacobian Verification --- + verifier = JacobianVerifier(bodies) + + for joint in joints: + verifier.add_joint_constraints(joint) + + # If grounded, remove the ground body's columns (fix its DOF) + j = verifier.get_jacobian() + if ground_body is not None and j.size > 0: + idx = verifier.body_index[ground_body] + cols_to_remove = list(range(idx * 6, (idx + 1) * 6)) + j = np.delete(j, cols_to_remove, axis=1) + + if j.size > 0: + sv = np.linalg.svd(j, compute_uv=False) + jacobian_rank = int(np.sum(sv > 1e-8)) + else: + jacobian_rank = 0 + + n_cols = j.shape[1] if j.size > 0 else 6 * len(bodies) + jacobian_nullity = n_cols - jacobian_rank + + dependent = verifier.find_dependencies() + + # Adjust for ground + trivial_dof = 0 if ground_body is not None else 6 + jacobian_internal_dof = jacobian_nullity - trivial_dof + + geometric_degeneracies = max(0, combinatorial_independent - jacobian_rank) + + # Rigidity: numerically rigid if nullity == trivial DOF + is_rigid = jacobian_nullity <= trivial_dof + is_minimally_rigid = is_rigid and len(dependent) == 0 + + return ConstraintAnalysis( + combinatorial_dof=effective_dof, + combinatorial_internal_dof=effective_internal_dof, + combinatorial_redundant=pg.get_redundant_count(), + combinatorial_classification=combinatorial_classification, + per_edge_results=all_edge_results, + jacobian_rank=jacobian_rank, + jacobian_nullity=jacobian_nullity, + jacobian_internal_dof=max(0, jacobian_internal_dof), + numerically_dependent=dependent, + geometric_degeneracies=geometric_degeneracies, + is_rigid=is_rigid, + is_minimally_rigid=is_minimally_rigid, + )