diff --git a/solver/datagen/__init__.py b/solver/datagen/__init__.py index a1d0c34..169907e 100644 --- a/solver/datagen/__init__.py +++ b/solver/datagen/__init__.py @@ -7,6 +7,7 @@ from solver.datagen.generator import ( SyntheticAssemblyGenerator, ) from solver.datagen.jacobian import JacobianVerifier +from solver.datagen.labeling import AssemblyLabels, label_assembly from solver.datagen.pebble_game import PebbleGame3D from solver.datagen.types import ( ConstraintAnalysis, @@ -18,6 +19,7 @@ from solver.datagen.types import ( __all__ = [ "COMPLEXITY_RANGES", + "AssemblyLabels", "AxisStrategy", "ConstraintAnalysis", "JacobianVerifier", @@ -28,4 +30,5 @@ __all__ = [ "RigidBody", "SyntheticAssemblyGenerator", "analyze_assembly", + "label_assembly", ] diff --git a/solver/datagen/generator.py b/solver/datagen/generator.py index 15df68f..ad753f9 100644 --- a/solver/datagen/generator.py +++ b/solver/datagen/generator.py @@ -13,6 +13,7 @@ import numpy as np from scipy.spatial.transform import Rotation from solver.datagen.analysis import analyze_assembly +from solver.datagen.labeling import label_assembly from solver.datagen.types import ( ConstraintAnalysis, Joint, @@ -839,6 +840,11 @@ class SyntheticAssemblyGenerator: ) gen_name = "mixed" + # Produce ground truth labels (includes ConstraintAnalysis) + ground = 0 if grounded else None + labels = label_assembly(bodies, joints, ground_body=ground) + analysis = labels.analysis + # Build per-joint labels from edge results joint_labels: dict[int, dict[str, int]] = {} for result in analysis.per_edge_results: @@ -875,6 +881,7 @@ class SyntheticAssemblyGenerator: for j in joints ], "joint_labels": joint_labels, + "labels": labels.to_dict(), "assembly_classification": (analysis.combinatorial_classification), "is_rigid": analysis.is_rigid, "is_minimally_rigid": analysis.is_minimally_rigid, diff --git a/solver/datagen/labeling.py b/solver/datagen/labeling.py new file mode 100644 index 0000000..e03de42 --- /dev/null +++ b/solver/datagen/labeling.py @@ -0,0 +1,394 @@ +"""Ground truth labeling pipeline for synthetic assemblies. + +Produces rich per-constraint, per-joint, per-body, and assembly-level +labels by running both the pebble game and Jacobian verification and +correlating their results. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +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, +) + +if TYPE_CHECKING: + from typing import Any + +__all__ = ["AssemblyLabels", "label_assembly"] + +_GROUND_ID = -1 +_SVD_TOL = 1e-8 + + +# --------------------------------------------------------------------------- +# Label dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class ConstraintLabel: + """Per scalar-constraint label combining both analysis methods.""" + + joint_id: int + constraint_idx: int + pebble_independent: bool + jacobian_independent: bool + + +@dataclass +class JointLabel: + """Aggregated constraint counts for a single joint.""" + + joint_id: int + independent_count: int + redundant_count: int + total: int + + +@dataclass +class BodyDofLabel: + """Per-body DOF signature from nullspace projection.""" + + body_id: int + translational_dof: int + rotational_dof: int + + +@dataclass +class AssemblyLabel: + """Assembly-wide summary label.""" + + classification: str + total_dof: int + redundant_count: int + is_rigid: bool + is_minimally_rigid: bool + has_degeneracy: bool + + +@dataclass +class AssemblyLabels: + """Complete ground truth labels for an assembly.""" + + per_constraint: list[ConstraintLabel] + per_joint: list[JointLabel] + per_body: list[BodyDofLabel] + assembly: AssemblyLabel + analysis: ConstraintAnalysis + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable dict.""" + return { + "per_constraint": [ + { + "joint_id": c.joint_id, + "constraint_idx": c.constraint_idx, + "pebble_independent": c.pebble_independent, + "jacobian_independent": c.jacobian_independent, + } + for c in self.per_constraint + ], + "per_joint": [ + { + "joint_id": j.joint_id, + "independent_count": j.independent_count, + "redundant_count": j.redundant_count, + "total": j.total, + } + for j in self.per_joint + ], + "per_body": [ + { + "body_id": b.body_id, + "translational_dof": b.translational_dof, + "rotational_dof": b.rotational_dof, + } + for b in self.per_body + ], + "assembly": { + "classification": self.assembly.classification, + "total_dof": self.assembly.total_dof, + "redundant_count": self.assembly.redundant_count, + "is_rigid": self.assembly.is_rigid, + "is_minimally_rigid": self.assembly.is_minimally_rigid, + "has_degeneracy": self.assembly.has_degeneracy, + }, + } + + +# --------------------------------------------------------------------------- +# Per-body DOF from nullspace projection +# --------------------------------------------------------------------------- + + +def _compute_per_body_dof( + j_reduced: np.ndarray, + body_ids: list[int], + ground_body: int | None, + body_index: dict[int, int], +) -> list[BodyDofLabel]: + """Compute translational and rotational DOF per body. + + Uses SVD nullspace projection: for each body, extract its + translational (3 cols) and rotational (3 cols) components + from the nullspace basis and compute ranks. + """ + # Build column index mapping for the reduced Jacobian + # (ground body columns have been removed) + col_map: dict[int, int] = {} + col_idx = 0 + for bid in body_ids: + if bid == ground_body: + continue + col_map[bid] = col_idx + col_idx += 1 + + results: list[BodyDofLabel] = [] + + if j_reduced.size == 0: + # No constraints — every body is fully free + for bid in body_ids: + if bid == ground_body: + results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0)) + else: + results.append(BodyDofLabel(body_id=bid, translational_dof=3, rotational_dof=3)) + return results + + # Full SVD to get nullspace + _u, s, vh = np.linalg.svd(j_reduced, full_matrices=True) + rank = int(np.sum(s > _SVD_TOL)) + n_cols = j_reduced.shape[1] + + if rank >= n_cols: + # Fully constrained — no nullspace + for bid in body_ids: + results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0)) + return results + + # Nullspace basis: rows of Vh beyond the rank + nullspace = vh[rank:] # shape: (n_cols - rank, n_cols) + + for bid in body_ids: + if bid == ground_body: + results.append(BodyDofLabel(body_id=bid, translational_dof=0, rotational_dof=0)) + continue + + idx = col_map[bid] + trans_cols = nullspace[:, idx * 6 : idx * 6 + 3] + rot_cols = nullspace[:, idx * 6 + 3 : idx * 6 + 6] + + # Rank of each block = DOF in that category + if trans_cols.size > 0: + sv_t = np.linalg.svd(trans_cols, compute_uv=False) + t_dof = int(np.sum(sv_t > _SVD_TOL)) + else: + t_dof = 0 + + if rot_cols.size > 0: + sv_r = np.linalg.svd(rot_cols, compute_uv=False) + r_dof = int(np.sum(sv_r > _SVD_TOL)) + else: + r_dof = 0 + + results.append(BodyDofLabel(body_id=bid, translational_dof=t_dof, rotational_dof=r_dof)) + + return results + + +# --------------------------------------------------------------------------- +# Main labeling function +# --------------------------------------------------------------------------- + + +def label_assembly( + bodies: list[RigidBody], + joints: list[Joint], + ground_body: int | None = None, +) -> AssemblyLabels: + """Produce complete ground truth labels for an assembly. + + Runs both the pebble game and Jacobian verification internally, + then correlates their results into per-constraint, per-joint, + per-body, and assembly-level labels. + + Args: + bodies: Rigid bodies in the assembly. + joints: Joints connecting the bodies. + ground_body: If set, this body is fixed to the world. + + Returns: + AssemblyLabels with full label set and embedded ConstraintAnalysis. + """ + # ---- Pebble Game ---- + pg = PebbleGame3D() + all_edge_results: list[dict[str, Any]] = [] + + 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) + + for joint in joints: + results = pg.add_joint(joint) + all_edge_results.extend(results) + + grounded = ground_body is not None + combinatorial_independent = len(pg.state.independent_edges) + 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)) + + redundant_count = pg.get_redundant_count() + if redundant_count > 0 and effective_internal_dof > 0: + classification = "mixed" + elif redundant_count > 0: + classification = "overconstrained" + elif effective_internal_dof > 0: + classification = "underconstrained" + else: + classification = "well-constrained" + + # ---- Jacobian Verification ---- + verifier = JacobianVerifier(bodies) + + for joint in joints: + verifier.add_joint_constraints(joint) + + j_full = verifier.get_jacobian() + j_reduced = j_full.copy() + if ground_body is not None and j_reduced.size > 0: + idx = verifier.body_index[ground_body] + cols_to_remove = list(range(idx * 6, (idx + 1) * 6)) + j_reduced = np.delete(j_reduced, cols_to_remove, axis=1) + + if j_reduced.size > 0: + sv = np.linalg.svd(j_reduced, compute_uv=False) + jacobian_rank = int(np.sum(sv > _SVD_TOL)) + else: + jacobian_rank = 0 + + n_cols = j_reduced.shape[1] if j_reduced.size > 0 else 6 * len(bodies) + jacobian_nullity = n_cols - jacobian_rank + dependent_rows = verifier.find_dependencies() + dependent_set = set(dependent_rows) + + trivial_dof = 0 if grounded else 6 + jacobian_internal_dof = jacobian_nullity - trivial_dof + geometric_degeneracies = max(0, combinatorial_independent - jacobian_rank) + is_rigid = jacobian_nullity <= trivial_dof + is_minimally_rigid = is_rigid and len(dependent_rows) == 0 + + # ---- Per-constraint labels ---- + # Map Jacobian rows to (joint_id, constraint_index). + # Rows are added contiguously per joint in the same order as joints. + row_to_joint: list[tuple[int, int]] = [] + for joint in joints: + dof = joint.joint_type.dof + for ci in range(dof): + row_to_joint.append((joint.joint_id, ci)) + + per_constraint: list[ConstraintLabel] = [] + for edge_idx, edge_result in enumerate(all_edge_results): + jid = edge_result["joint_id"] + ci = edge_result["constraint_index"] + pebble_indep = edge_result["independent"] + + # Find matching Jacobian row + jacobian_indep = True + if edge_idx < len(row_to_joint): + row_idx = edge_idx + jacobian_indep = row_idx not in dependent_set + + per_constraint.append( + ConstraintLabel( + joint_id=jid, + constraint_idx=ci, + pebble_independent=pebble_indep, + jacobian_independent=jacobian_indep, + ) + ) + + # ---- Per-joint labels ---- + joint_agg: dict[int, JointLabel] = {} + for cl in per_constraint: + if cl.joint_id not in joint_agg: + joint_agg[cl.joint_id] = JointLabel( + joint_id=cl.joint_id, + independent_count=0, + redundant_count=0, + total=0, + ) + jl = joint_agg[cl.joint_id] + jl.total += 1 + if cl.pebble_independent: + jl.independent_count += 1 + else: + jl.redundant_count += 1 + + per_joint = [joint_agg[j.joint_id] for j in joints if j.joint_id in joint_agg] + + # ---- Per-body DOF labels ---- + body_ids = [b.body_id for b in bodies] + per_body = _compute_per_body_dof( + j_reduced, + body_ids, + ground_body, + verifier.body_index, + ) + + # ---- Assembly label ---- + assembly_label = AssemblyLabel( + classification=classification, + total_dof=max(0, jacobian_internal_dof), + redundant_count=redundant_count, + is_rigid=is_rigid, + is_minimally_rigid=is_minimally_rigid, + has_degeneracy=geometric_degeneracies > 0, + ) + + # ---- ConstraintAnalysis (for backward compat) ---- + analysis = ConstraintAnalysis( + combinatorial_dof=effective_dof, + combinatorial_internal_dof=effective_internal_dof, + combinatorial_redundant=redundant_count, + combinatorial_classification=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_rows, + geometric_degeneracies=geometric_degeneracies, + is_rigid=is_rigid, + is_minimally_rigid=is_minimally_rigid, + ) + + return AssemblyLabels( + per_constraint=per_constraint, + per_joint=per_joint, + per_body=per_body, + assembly=assembly_label, + analysis=analysis, + ) diff --git a/tests/datagen/test_generator.py b/tests/datagen/test_generator.py index dc70baa..3b120fa 100644 --- a/tests/datagen/test_generator.py +++ b/tests/datagen/test_generator.py @@ -513,6 +513,7 @@ class TestTrainingBatch: "body_orientations", "joints", "joint_labels", + "labels", "assembly_classification", "is_rigid", "is_minimally_rigid", @@ -586,6 +587,17 @@ class TestTrainingBatch: assert len(o) == 3 assert len(o[0]) == 3 + def test_labels_structure(self) -> None: + """Each example has labels dict with expected sub-keys.""" + gen = SyntheticAssemblyGenerator(seed=42) + batch = gen.generate_training_batch(10) + for ex in batch: + labels = ex["labels"] + assert "per_constraint" in labels + assert "per_joint" in labels + assert "per_body" in labels + assert "assembly" in labels + def test_grounded_field_present(self) -> None: gen = SyntheticAssemblyGenerator(seed=42) batch = gen.generate_training_batch(10) diff --git a/tests/datagen/test_labeling.py b/tests/datagen/test_labeling.py new file mode 100644 index 0000000..a81e2e3 --- /dev/null +++ b/tests/datagen/test_labeling.py @@ -0,0 +1,346 @@ +"""Tests for solver.datagen.labeling -- ground truth labeling pipeline.""" + +from __future__ import annotations + +import json + +import numpy as np + +from solver.datagen.labeling import ( + label_assembly, +) +from solver.datagen.types import Joint, JointType, RigidBody + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_bodies(*positions: tuple[float, ...]) -> list[RigidBody]: + return [RigidBody(body_id=i, position=np.array(pos)) for i, pos in enumerate(positions)] + + +def _make_joint( + jid: int, + a: int, + b: int, + jtype: JointType, + axis: tuple[float, ...] = (0.0, 0.0, 1.0), +) -> Joint: + return Joint( + joint_id=jid, + body_a=a, + body_b=b, + joint_type=jtype, + anchor_a=np.zeros(3), + anchor_b=np.zeros(3), + axis=np.array(axis), + ) + + +# --------------------------------------------------------------------------- +# Per-constraint labels +# --------------------------------------------------------------------------- + + +class TestConstraintLabels: + """Per-constraint labels combine pebble game and Jacobian results.""" + + def test_fixed_joint_all_independent(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.FIXED)] + labels = label_assembly(bodies, joints, ground_body=0) + assert len(labels.per_constraint) == 6 + for cl in labels.per_constraint: + assert cl.pebble_independent is True + assert cl.jacobian_independent is True + + def test_revolute_joint_all_independent(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + assert len(labels.per_constraint) == 5 + for cl in labels.per_constraint: + assert cl.pebble_independent is True + assert cl.jacobian_independent is True + + def test_chain_constraint_count(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (4, 0, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + assert len(labels.per_constraint) == 10 # 5 + 5 + + def test_constraint_joint_ids(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (4, 0, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.BALL), + ] + labels = label_assembly(bodies, joints, ground_body=0) + j0_constraints = [c for c in labels.per_constraint if c.joint_id == 0] + j1_constraints = [c for c in labels.per_constraint if c.joint_id == 1] + assert len(j0_constraints) == 5 # revolute + assert len(j1_constraints) == 3 # ball + + def test_overconstrained_has_pebble_redundant(self) -> None: + """Triangle with revolute joints: some constraints redundant.""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (1, 2, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + _make_joint(2, 2, 0, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + pebble_redundant = sum(1 for c in labels.per_constraint if not c.pebble_independent) + assert pebble_redundant > 0 + + +# --------------------------------------------------------------------------- +# Per-joint labels +# --------------------------------------------------------------------------- + + +class TestJointLabels: + """Per-joint aggregated labels.""" + + def test_fixed_joint_counts(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.FIXED)] + labels = label_assembly(bodies, joints, ground_body=0) + assert len(labels.per_joint) == 1 + jl = labels.per_joint[0] + assert jl.joint_id == 0 + assert jl.independent_count == 6 + assert jl.redundant_count == 0 + assert jl.total == 6 + + def test_overconstrained_has_redundant_joints(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (1, 2, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + _make_joint(2, 2, 0, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + total_redundant = sum(jl.redundant_count for jl in labels.per_joint) + assert total_redundant > 0 + + def test_joint_total_equals_dof(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.BALL)] + labels = label_assembly(bodies, joints, ground_body=0) + jl = labels.per_joint[0] + assert jl.total == 3 # ball has 3 DOF + + +# --------------------------------------------------------------------------- +# Per-body DOF labels +# --------------------------------------------------------------------------- + + +class TestBodyDofLabels: + """Per-body DOF signatures from nullspace projection.""" + + def test_fixed_joint_grounded_both_zero(self) -> None: + """Two bodies + fixed joint + grounded: both fully constrained.""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.FIXED)] + labels = label_assembly(bodies, joints, ground_body=0) + for bl in labels.per_body: + assert bl.translational_dof == 0 + assert bl.rotational_dof == 0 + + def test_revolute_has_rotational_dof(self) -> None: + """Two bodies + revolute + grounded: body 1 has rotational DOF.""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + b1 = next(b for b in labels.per_body if b.body_id == 1) + # Revolute allows 1 rotation DOF + assert b1.rotational_dof >= 1 + + def test_dof_bounds(self) -> None: + """All DOF values should be in [0, 3].""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (4, 0, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + for bl in labels.per_body: + assert 0 <= bl.translational_dof <= 3 + assert 0 <= bl.rotational_dof <= 3 + + def test_floating_more_dof_than_grounded(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + grounded = label_assembly(bodies, joints, ground_body=0) + floating = label_assembly(bodies, joints, ground_body=None) + g_total = sum(b.translational_dof + b.rotational_dof for b in grounded.per_body) + f_total = sum(b.translational_dof + b.rotational_dof for b in floating.per_body) + assert f_total > g_total + + def test_grounded_body_zero_dof(self) -> None: + """The grounded body should have 0 DOF.""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + b0 = next(b for b in labels.per_body if b.body_id == 0) + assert b0.translational_dof == 0 + assert b0.rotational_dof == 0 + + def test_body_count_matches(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (4, 0, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.BALL), + ] + labels = label_assembly(bodies, joints, ground_body=0) + assert len(labels.per_body) == 3 + + +# --------------------------------------------------------------------------- +# Assembly label +# --------------------------------------------------------------------------- + + +class TestAssemblyLabel: + """Assembly-wide summary labels.""" + + def test_underconstrained_chain(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (4, 0, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + assert labels.assembly.classification == "underconstrained" + assert labels.assembly.is_rigid is False + + def test_well_constrained(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.FIXED)] + labels = label_assembly(bodies, joints, ground_body=0) + assert labels.assembly.classification == "well-constrained" + assert labels.assembly.is_rigid is True + assert labels.assembly.is_minimally_rigid is True + + def test_overconstrained(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (1, 2, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE), + _make_joint(1, 1, 2, JointType.REVOLUTE), + _make_joint(2, 2, 0, JointType.REVOLUTE), + ] + labels = label_assembly(bodies, joints, ground_body=0) + assert labels.assembly.redundant_count > 0 + + def test_has_degeneracy_with_parallel_axes(self) -> None: + """Parallel revolute axes in a loop create geometric degeneracy.""" + z_axis = (0.0, 0.0, 1.0) + bodies = _make_bodies((0, 0, 0), (2, 0, 0), (2, 2, 0), (0, 2, 0)) + joints = [ + _make_joint(0, 0, 1, JointType.REVOLUTE, axis=z_axis), + _make_joint(1, 1, 2, JointType.REVOLUTE, axis=z_axis), + _make_joint(2, 2, 3, JointType.REVOLUTE, axis=z_axis), + _make_joint(3, 3, 0, JointType.REVOLUTE, axis=z_axis), + ] + labels = label_assembly(bodies, joints, ground_body=0) + assert labels.assembly.has_degeneracy is True + + +# --------------------------------------------------------------------------- +# Serialization +# --------------------------------------------------------------------------- + + +class TestToDict: + """to_dict produces JSON-serializable output.""" + + def test_top_level_keys(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + d = labels.to_dict() + assert set(d.keys()) == { + "per_constraint", + "per_joint", + "per_body", + "assembly", + } + + def test_per_constraint_keys(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + d = labels.to_dict() + for item in d["per_constraint"]: + assert set(item.keys()) == { + "joint_id", + "constraint_idx", + "pebble_independent", + "jacobian_independent", + } + + def test_assembly_keys(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + d = labels.to_dict() + assert set(d["assembly"].keys()) == { + "classification", + "total_dof", + "redundant_count", + "is_rigid", + "is_minimally_rigid", + "has_degeneracy", + } + + def test_json_serializable(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + d = labels.to_dict() + # Should not raise + serialized = json.dumps(d) + assert isinstance(serialized, str) + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +class TestLabelAssemblyEdgeCases: + """Edge cases for label_assembly.""" + + def test_no_joints(self) -> None: + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + labels = label_assembly(bodies, [], ground_body=0) + assert len(labels.per_constraint) == 0 + assert len(labels.per_joint) == 0 + assert labels.assembly.classification == "underconstrained" + # Non-ground body should be fully free + b1 = next(b for b in labels.per_body if b.body_id == 1) + assert b1.translational_dof == 3 + assert b1.rotational_dof == 3 + + def test_no_joints_floating(self) -> None: + bodies = _make_bodies((0, 0, 0)) + labels = label_assembly(bodies, [], ground_body=None) + assert len(labels.per_body) == 1 + assert labels.per_body[0].translational_dof == 3 + assert labels.per_body[0].rotational_dof == 3 + + def test_analysis_embedded(self) -> None: + """AssemblyLabels.analysis should be a valid ConstraintAnalysis.""" + bodies = _make_bodies((0, 0, 0), (2, 0, 0)) + joints = [_make_joint(0, 0, 1, JointType.REVOLUTE)] + labels = label_assembly(bodies, joints, ground_body=0) + analysis = labels.analysis + assert hasattr(analysis, "combinatorial_classification") + assert hasattr(analysis, "jacobian_rank") + assert hasattr(analysis, "is_rigid")