diff --git a/solver/datagen/analysis.py b/solver/datagen/analysis.py index f862802..6ad8491 100644 --- a/solver/datagen/analysis.py +++ b/solver/datagen/analysis.py @@ -78,7 +78,17 @@ def analyze_assembly( 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) + # Classify based on effective (adjusted) DOF, not raw pebble game output, + # because the virtual ground body skews the raw numbers. + redundant = pg.get_redundant_count() + if redundant > 0 and effective_internal_dof > 0: + combinatorial_classification = "mixed" + elif redundant > 0: + combinatorial_classification = "overconstrained" + elif effective_internal_dof > 0: + combinatorial_classification = "underconstrained" + else: + combinatorial_classification = "well-constrained" # --- Jacobian Verification --- verifier = JacobianVerifier(bodies) diff --git a/solver/datagen/pebble_game.py b/solver/datagen/pebble_game.py index 655cd13..e3bdbc3 100644 --- a/solver/datagen/pebble_game.py +++ b/solver/datagen/pebble_game.py @@ -79,7 +79,7 @@ class PebbleGame3D: self.add_body(joint.body_a) self.add_body(joint.body_b) - num_constraints = joint.joint_type.value + num_constraints = joint.joint_type.dof results: list[dict[str, Any]] = [] for i in range(num_constraints): diff --git a/solver/datagen/types.py b/solver/datagen/types.py index 754e5dc..f7b3e84 100644 --- a/solver/datagen/types.py +++ b/solver/datagen/types.py @@ -37,20 +37,27 @@ class JointType(enum.Enum): representation, each joint maps to a number of edges equal to the DOF it removes. - DOF removed = number of scalar constraint equations the joint imposes. + Values are ``(ordinal, dof_removed)`` tuples so that joint types + sharing the same DOF count remain distinct enum members. Use the + :attr:`dof` property to get the scalar constraint count. """ - FIXED = 6 # Locks all relative motion - REVOLUTE = 5 # Allows rotation about one axis - CYLINDRICAL = 4 # Allows rotation + translation along one axis - SLIDER = 5 # Allows translation along one axis (prismatic) - BALL = 3 # Allows rotation about a point (spherical) - PLANAR = 3 # Allows 2D translation + rotation normal to plane - SCREW = 5 # Coupled rotation-translation (helical) - UNIVERSAL = 4 # Two rotational DOF (Cardan/U-joint) - PARALLEL = 3 # Forces parallel orientation (3 rotation constraints) - PERPENDICULAR = 1 # Single angular constraint - DISTANCE = 1 # Single scalar distance constraint + FIXED = (0, 6) # Locks all relative motion + REVOLUTE = (1, 5) # Allows rotation about one axis + CYLINDRICAL = (2, 4) # Allows rotation + translation along one axis + SLIDER = (3, 5) # Allows translation along one axis (prismatic) + BALL = (4, 3) # Allows rotation about a point (spherical) + PLANAR = (5, 3) # Allows 2D translation + rotation normal to plane + SCREW = (6, 5) # Coupled rotation-translation (helical) + UNIVERSAL = (7, 4) # Two rotational DOF (Cardan/U-joint) + PARALLEL = (8, 3) # Forces parallel orientation (3 rotation constraints) + PERPENDICULAR = (9, 1) # Single angular constraint + DISTANCE = (10, 1) # Single scalar distance constraint + + @property + def dof(self) -> int: + """Number of scalar constraints (DOF removed) by this joint type.""" + return self.value[1] # --------------------------------------------------------------------------- diff --git a/tests/datagen/__init__.py b/tests/datagen/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/datagen/test_analysis.py b/tests/datagen/test_analysis.py new file mode 100644 index 0000000..8edf893 --- /dev/null +++ b/tests/datagen/test_analysis.py @@ -0,0 +1,240 @@ +"""Tests for solver.datagen.analysis -- combined analysis function.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from solver.datagen.analysis import analyze_assembly +from solver.datagen.types import ( + ConstraintAnalysis, + Joint, + JointType, + RigidBody, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _two_bodies() -> list[RigidBody]: + return [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + ] + + +def _triangle_bodies() -> list[RigidBody]: + return [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + RigidBody(2, position=np.array([1.0, 1.7, 0.0])), + ] + + +# --------------------------------------------------------------------------- +# Scenario 1: Two bodies + revolute (underconstrained, 1 internal DOF) +# --------------------------------------------------------------------------- + + +class TestTwoBodiesRevolute: + """Demo scenario 1: two bodies connected by a revolute joint.""" + + @pytest.fixture() + def result(self) -> ConstraintAnalysis: + bodies = _two_bodies() + joints = [ + Joint( + 0, + body_a=0, + body_b=1, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + ] + return analyze_assembly(bodies, joints, ground_body=0) + + def test_internal_dof(self, result: ConstraintAnalysis) -> None: + assert result.jacobian_internal_dof == 1 + + def test_not_rigid(self, result: ConstraintAnalysis) -> None: + assert not result.is_rigid + + def test_classification(self, result: ConstraintAnalysis) -> None: + assert result.combinatorial_classification == "underconstrained" + + def test_no_redundant(self, result: ConstraintAnalysis) -> None: + assert result.combinatorial_redundant == 0 + + +# --------------------------------------------------------------------------- +# Scenario 2: Two bodies + fixed (well-constrained, 0 internal DOF) +# --------------------------------------------------------------------------- + + +class TestTwoBodiesFixed: + """Demo scenario 2: two bodies connected by a fixed joint.""" + + @pytest.fixture() + def result(self) -> ConstraintAnalysis: + bodies = _two_bodies() + joints = [ + Joint( + 0, + body_a=0, + body_b=1, + joint_type=JointType.FIXED, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + ), + ] + return analyze_assembly(bodies, joints, ground_body=0) + + def test_internal_dof(self, result: ConstraintAnalysis) -> None: + assert result.jacobian_internal_dof == 0 + + def test_rigid(self, result: ConstraintAnalysis) -> None: + assert result.is_rigid + + def test_minimally_rigid(self, result: ConstraintAnalysis) -> None: + assert result.is_minimally_rigid + + def test_classification(self, result: ConstraintAnalysis) -> None: + assert result.combinatorial_classification == "well-constrained" + + +# --------------------------------------------------------------------------- +# Scenario 3: Triangle with revolute joints (overconstrained) +# --------------------------------------------------------------------------- + + +class TestTriangleRevolute: + """Demo scenario 3: triangle of 3 bodies + 3 revolute joints.""" + + @pytest.fixture() + def result(self) -> ConstraintAnalysis: + bodies = _triangle_bodies() + joints = [ + Joint( + 0, + body_a=0, + body_b=1, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + Joint( + 1, + body_a=1, + body_b=2, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.5, 0.85, 0.0]), + anchor_b=np.array([1.5, 0.85, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + Joint( + 2, + body_a=2, + body_b=0, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([0.5, 0.85, 0.0]), + anchor_b=np.array([0.5, 0.85, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + ] + return analyze_assembly(bodies, joints, ground_body=0) + + def test_has_redundant(self, result: ConstraintAnalysis) -> None: + assert result.combinatorial_redundant > 0 + + def test_classification(self, result: ConstraintAnalysis) -> None: + assert result.combinatorial_classification in ("overconstrained", "mixed") + + def test_rigid(self, result: ConstraintAnalysis) -> None: + assert result.is_rigid + + def test_numerically_dependent(self, result: ConstraintAnalysis) -> None: + assert len(result.numerically_dependent) > 0 + + +# --------------------------------------------------------------------------- +# Scenario 4: Parallel revolute axes (geometric degeneracy) +# --------------------------------------------------------------------------- + + +class TestParallelRevoluteAxes: + """Demo scenario 4: parallel revolute axes create geometric degeneracies.""" + + @pytest.fixture() + def result(self) -> ConstraintAnalysis: + bodies = [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + RigidBody(2, position=np.array([4.0, 0.0, 0.0])), + ] + joints = [ + Joint( + 0, + body_a=0, + body_b=1, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + Joint( + 1, + body_a=1, + body_b=2, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([3.0, 0.0, 0.0]), + anchor_b=np.array([3.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ), + ] + return analyze_assembly(bodies, joints, ground_body=0) + + def test_geometric_degeneracies_detected(self, result: ConstraintAnalysis) -> None: + """Parallel axes produce at least one geometric degeneracy.""" + assert result.geometric_degeneracies > 0 + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- + + +class TestNoJoints: + """Assembly with bodies but no joints.""" + + def test_all_dof_free(self) -> None: + bodies = _two_bodies() + result = analyze_assembly(bodies, [], ground_body=0) + # Body 1 is completely free (6 DOF), body 0 is grounded + assert result.jacobian_internal_dof > 0 + assert not result.is_rigid + + def test_ungrounded(self) -> None: + bodies = _two_bodies() + result = analyze_assembly(bodies, []) + assert result.combinatorial_classification == "underconstrained" + + +class TestReturnType: + """Verify the return object is a proper ConstraintAnalysis.""" + + def test_instance(self) -> None: + bodies = _two_bodies() + joints = [Joint(0, 0, 1, JointType.FIXED)] + result = analyze_assembly(bodies, joints) + assert isinstance(result, ConstraintAnalysis) + + def test_per_edge_results_populated(self) -> None: + bodies = _two_bodies() + joints = [Joint(0, 0, 1, JointType.REVOLUTE)] + result = analyze_assembly(bodies, joints) + assert len(result.per_edge_results) == 5 diff --git a/tests/datagen/test_generator.py b/tests/datagen/test_generator.py new file mode 100644 index 0000000..06c10a5 --- /dev/null +++ b/tests/datagen/test_generator.py @@ -0,0 +1,166 @@ +"""Tests for solver.datagen.generator -- synthetic assembly generation.""" + +from __future__ import annotations + +import pytest + +from solver.datagen.generator import SyntheticAssemblyGenerator +from solver.datagen.types import JointType + + +class TestChainAssembly: + """generate_chain_assembly produces valid underconstrained chains.""" + + def test_returns_three_tuple(self) -> None: + gen = SyntheticAssemblyGenerator(seed=0) + bodies, joints, _analysis = gen.generate_chain_assembly(4) + assert len(bodies) == 4 + assert len(joints) == 3 + + def test_chain_underconstrained(self) -> None: + gen = SyntheticAssemblyGenerator(seed=0) + _, _, analysis = gen.generate_chain_assembly(4) + assert analysis.combinatorial_classification == "underconstrained" + + def test_chain_body_ids(self) -> None: + gen = SyntheticAssemblyGenerator(seed=0) + bodies, _, _ = gen.generate_chain_assembly(5) + ids = [b.body_id for b in bodies] + assert ids == [0, 1, 2, 3, 4] + + def test_chain_joint_connectivity(self) -> None: + gen = SyntheticAssemblyGenerator(seed=0) + _, joints, _ = gen.generate_chain_assembly(4) + for i, j in enumerate(joints): + assert j.body_a == i + assert j.body_b == i + 1 + + def test_chain_custom_joint_type(self) -> None: + gen = SyntheticAssemblyGenerator(seed=0) + _, joints, _ = gen.generate_chain_assembly(3, joint_type=JointType.BALL) + assert all(j.joint_type is JointType.BALL for j in joints) + + +class TestRigidAssembly: + """generate_rigid_assembly produces rigid assemblies.""" + + def test_rigid(self) -> None: + gen = SyntheticAssemblyGenerator(seed=42) + _, _, analysis = gen.generate_rigid_assembly(4) + assert analysis.is_rigid + + def test_spanning_tree_structure(self) -> None: + """n bodies should have at least n-1 joints (spanning tree).""" + gen = SyntheticAssemblyGenerator(seed=42) + bodies, joints, _ = gen.generate_rigid_assembly(5) + assert len(joints) >= len(bodies) - 1 + + def test_deterministic(self) -> None: + """Same seed produces same results.""" + g1 = SyntheticAssemblyGenerator(seed=99) + g2 = SyntheticAssemblyGenerator(seed=99) + _, j1, a1 = g1.generate_rigid_assembly(4) + _, j2, a2 = g2.generate_rigid_assembly(4) + assert a1.jacobian_rank == a2.jacobian_rank + assert len(j1) == len(j2) + + +class TestOverconstrainedAssembly: + """generate_overconstrained_assembly adds redundant constraints.""" + + def test_has_redundant(self) -> None: + gen = SyntheticAssemblyGenerator(seed=42) + _, _, analysis = gen.generate_overconstrained_assembly(4, extra_joints=2) + assert analysis.combinatorial_redundant > 0 + + def test_extra_joints_added(self) -> None: + gen = SyntheticAssemblyGenerator(seed=42) + _, joints_base, _ = gen.generate_rigid_assembly(4) + + gen2 = SyntheticAssemblyGenerator(seed=42) + _, joints_over, _ = gen2.generate_overconstrained_assembly(4, extra_joints=3) + # Overconstrained has base joints + extra + assert len(joints_over) > len(joints_base) + + +class TestTrainingBatch: + """generate_training_batch produces well-structured examples.""" + + @pytest.fixture() + def batch(self) -> list[dict]: + gen = SyntheticAssemblyGenerator(seed=42) + return gen.generate_training_batch(batch_size=20, n_bodies_range=(3, 6)) + + def test_batch_size(self, batch: list[dict]) -> None: + assert len(batch) == 20 + + def test_example_keys(self, batch: list[dict]) -> None: + expected = { + "example_id", + "n_bodies", + "n_joints", + "body_positions", + "joints", + "joint_labels", + "assembly_classification", + "is_rigid", + "is_minimally_rigid", + "internal_dof", + "geometric_degeneracies", + } + for ex in batch: + assert set(ex.keys()) == expected + + def test_example_ids_sequential(self, batch: list[dict]) -> None: + ids = [ex["example_id"] for ex in batch] + assert ids == list(range(20)) + + def test_classification_distribution(self, batch: list[dict]) -> None: + """Batch should contain multiple classification types.""" + classes = {ex["assembly_classification"] for ex in batch} + # With the 3-way generator split we expect at least 2 types + assert len(classes) >= 2 + + def test_body_count_in_range(self, batch: list[dict]) -> None: + for ex in batch: + assert 3 <= ex["n_bodies"] <= 5 # range is [3, 6) + + def test_joint_labels_match_joints(self, batch: list[dict]) -> None: + for ex in batch: + label_jids = set(ex["joint_labels"].keys()) + joint_jids = {j["joint_id"] for j in ex["joints"]} + assert label_jids == joint_jids + + def test_joint_label_fields(self, batch: list[dict]) -> None: + expected_fields = { + "independent_constraints", + "redundant_constraints", + "total_constraints", + } + for ex in batch: + for label in ex["joint_labels"].values(): + assert set(label.keys()) == expected_fields + + def test_joint_label_consistency(self, batch: list[dict]) -> None: + """independent + redundant == total for every joint.""" + for ex in batch: + for label in ex["joint_labels"].values(): + total = label["independent_constraints"] + label["redundant_constraints"] + assert total == label["total_constraints"] + + +class TestSeedReproducibility: + """Different seeds produce different results.""" + + def test_different_seeds_differ(self) -> None: + g1 = SyntheticAssemblyGenerator(seed=1) + g2 = SyntheticAssemblyGenerator(seed=2) + b1 = g1.generate_training_batch(batch_size=5, n_bodies_range=(3, 6)) + b2 = g2.generate_training_batch(batch_size=5, n_bodies_range=(3, 6)) + # Very unlikely to be identical with different seeds + c1 = [ex["assembly_classification"] for ex in b1] + c2 = [ex["assembly_classification"] for ex in b2] + r1 = [ex["is_rigid"] for ex in b1] + r2 = [ex["is_rigid"] for ex in b2] + # At least one of these should differ (probabilistically certain) + assert c1 != c2 or r1 != r2 diff --git a/tests/datagen/test_jacobian.py b/tests/datagen/test_jacobian.py new file mode 100644 index 0000000..864331d --- /dev/null +++ b/tests/datagen/test_jacobian.py @@ -0,0 +1,267 @@ +"""Tests for solver.datagen.jacobian -- Jacobian rank verification.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from solver.datagen.jacobian import JacobianVerifier +from solver.datagen.types import Joint, JointType, RigidBody + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _two_bodies() -> list[RigidBody]: + return [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + ] + + +def _three_bodies() -> list[RigidBody]: + return [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + RigidBody(2, position=np.array([4.0, 0.0, 0.0])), + ] + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestJacobianShape: + """Verify Jacobian matrix dimensions for each joint type.""" + + @pytest.mark.parametrize( + "joint_type,expected_rows", + [ + (JointType.FIXED, 6), + (JointType.REVOLUTE, 5), + (JointType.CYLINDRICAL, 4), + (JointType.SLIDER, 5), + (JointType.BALL, 3), + (JointType.PLANAR, 3), + (JointType.SCREW, 5), + (JointType.UNIVERSAL, 4), + (JointType.PARALLEL, 3), + (JointType.PERPENDICULAR, 1), + (JointType.DISTANCE, 1), + ], + ) + def test_row_count(self, joint_type: JointType, expected_rows: int) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + joint = Joint( + joint_id=0, + body_a=0, + body_b=1, + joint_type=joint_type, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ) + n_added = v.add_joint_constraints(joint) + assert n_added == expected_rows + + j = v.get_jacobian() + assert j.shape == (expected_rows, 12) # 2 bodies * 6 cols + + +class TestNumericalRank: + """Numerical rank checks for known configurations.""" + + def test_fixed_joint_rank_six(self) -> None: + """Fixed joint between 2 bodies: rank = 6.""" + bodies = _two_bodies() + v = JacobianVerifier(bodies) + j = Joint( + joint_id=0, + body_a=0, + body_b=1, + joint_type=JointType.FIXED, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + ) + v.add_joint_constraints(j) + assert v.numerical_rank() == 6 + + def test_revolute_joint_rank_five(self) -> None: + """Revolute joint between 2 bodies: rank = 5.""" + bodies = _two_bodies() + v = JacobianVerifier(bodies) + j = Joint( + joint_id=0, + body_a=0, + body_b=1, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ) + v.add_joint_constraints(j) + assert v.numerical_rank() == 5 + + def test_ball_joint_rank_three(self) -> None: + """Ball joint between 2 bodies: rank = 3.""" + bodies = _two_bodies() + v = JacobianVerifier(bodies) + j = Joint( + joint_id=0, + body_a=0, + body_b=1, + joint_type=JointType.BALL, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + ) + v.add_joint_constraints(j) + assert v.numerical_rank() == 3 + + def test_empty_jacobian_rank_zero(self) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + assert v.numerical_rank() == 0 + + +class TestParallelAxesDegeneracy: + """Parallel revolute axes create geometric dependencies.""" + + def _four_body_loop(self) -> list[RigidBody]: + return [ + RigidBody(0, position=np.array([0.0, 0.0, 0.0])), + RigidBody(1, position=np.array([2.0, 0.0, 0.0])), + RigidBody(2, position=np.array([2.0, 2.0, 0.0])), + RigidBody(3, position=np.array([0.0, 2.0, 0.0])), + ] + + def _loop_joints(self, axes: list[np.ndarray]) -> list[Joint]: + pairs = [(0, 1, [1, 0, 0]), (1, 2, [2, 1, 0]), (2, 3, [1, 2, 0]), (3, 0, [0, 1, 0])] + return [ + Joint( + joint_id=i, + body_a=a, + body_b=b, + joint_type=JointType.REVOLUTE, + anchor_a=np.array(anc, dtype=float), + anchor_b=np.array(anc, dtype=float), + axis=axes[i], + ) + for i, (a, b, anc) in enumerate(pairs) + ] + + def test_parallel_has_lower_rank(self) -> None: + """4-body closed loop: all-parallel revolute axes produce lower + Jacobian rank than mixed axes due to geometric dependency.""" + bodies = self._four_body_loop() + z_axis = np.array([0.0, 0.0, 1.0]) + + # All axes parallel to Z + v_par = JacobianVerifier(bodies) + for j in self._loop_joints([z_axis] * 4): + v_par.add_joint_constraints(j) + rank_par = v_par.numerical_rank() + + # Mixed axes + mixed = [ + np.array([0.0, 0.0, 1.0]), + np.array([0.0, 1.0, 0.0]), + np.array([0.0, 0.0, 1.0]), + np.array([1.0, 0.0, 0.0]), + ] + v_mix = JacobianVerifier(bodies) + for j in self._loop_joints(mixed): + v_mix.add_joint_constraints(j) + rank_mix = v_mix.numerical_rank() + + assert rank_par < rank_mix + + +class TestFindDependencies: + """Dependency detection.""" + + def test_fixed_joint_no_dependencies(self) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + j = Joint( + joint_id=0, + body_a=0, + body_b=1, + joint_type=JointType.FIXED, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + ) + v.add_joint_constraints(j) + assert v.find_dependencies() == [] + + def test_duplicate_fixed_has_dependencies(self) -> None: + """Two fixed joints on same pair: second is fully dependent.""" + bodies = _two_bodies() + v = JacobianVerifier(bodies) + for jid in range(2): + v.add_joint_constraints( + Joint( + joint_id=jid, + body_a=0, + body_b=1, + joint_type=JointType.FIXED, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + ) + ) + deps = v.find_dependencies() + assert len(deps) == 6 # Second fixed joint entirely redundant + + def test_empty_no_dependencies(self) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + assert v.find_dependencies() == [] + + +class TestRowLabels: + """Row label metadata.""" + + def test_labels_match_rows(self) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + j = Joint( + joint_id=7, + body_a=0, + body_b=1, + joint_type=JointType.REVOLUTE, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([1.0, 0.0, 0.0]), + axis=np.array([0.0, 0.0, 1.0]), + ) + v.add_joint_constraints(j) + assert len(v.row_labels) == 5 + assert all(lab["joint_id"] == 7 for lab in v.row_labels) + + +class TestPerpendicularPair: + """Internal _perpendicular_pair utility.""" + + @pytest.mark.parametrize( + "axis", + [ + np.array([1.0, 0.0, 0.0]), + np.array([0.0, 1.0, 0.0]), + np.array([0.0, 0.0, 1.0]), + np.array([1.0, 1.0, 1.0]) / np.sqrt(3), + ], + ) + def test_orthonormal(self, axis: np.ndarray) -> None: + bodies = _two_bodies() + v = JacobianVerifier(bodies) + t1, t2 = v._perpendicular_pair(axis) + + # All unit length + np.testing.assert_allclose(np.linalg.norm(t1), 1.0, atol=1e-12) + np.testing.assert_allclose(np.linalg.norm(t2), 1.0, atol=1e-12) + + # Mutually perpendicular + np.testing.assert_allclose(np.dot(axis, t1), 0.0, atol=1e-12) + np.testing.assert_allclose(np.dot(axis, t2), 0.0, atol=1e-12) + np.testing.assert_allclose(np.dot(t1, t2), 0.0, atol=1e-12) diff --git a/tests/datagen/test_pebble_game.py b/tests/datagen/test_pebble_game.py new file mode 100644 index 0000000..da71eab --- /dev/null +++ b/tests/datagen/test_pebble_game.py @@ -0,0 +1,206 @@ +"""Tests for solver.datagen.pebble_game -- (6,6)-pebble game.""" + +from __future__ import annotations + +import itertools + +import numpy as np +import pytest + +from solver.datagen.pebble_game import PebbleGame3D +from solver.datagen.types import Joint, JointType + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _revolute(jid: int, a: int, b: int, axis: np.ndarray | None = None) -> Joint: + """Shorthand for a revolute joint between bodies *a* and *b*.""" + if axis is None: + axis = np.array([0.0, 0.0, 1.0]) + return Joint( + joint_id=jid, + body_a=a, + body_b=b, + joint_type=JointType.REVOLUTE, + axis=axis, + ) + + +def _fixed(jid: int, a: int, b: int) -> Joint: + return Joint(joint_id=jid, body_a=a, body_b=b, joint_type=JointType.FIXED) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestAddBody: + """Body registration basics.""" + + def test_single_body_six_pebbles(self) -> None: + pg = PebbleGame3D() + pg.add_body(0) + assert pg.state.free_pebbles[0] == 6 + + def test_duplicate_body_no_op(self) -> None: + pg = PebbleGame3D() + pg.add_body(0) + pg.add_body(0) + assert pg.state.free_pebbles[0] == 6 + + def test_multiple_bodies(self) -> None: + pg = PebbleGame3D() + for i in range(5): + pg.add_body(i) + assert pg.get_dof() == 30 # 5 * 6 + + +class TestAddJoint: + """Joint insertion and DOF accounting.""" + + def test_revolute_removes_five_dof(self) -> None: + pg = PebbleGame3D() + results = pg.add_joint(_revolute(0, 0, 1)) + assert len(results) == 5 # 5 scalar constraints + assert all(r["independent"] for r in results) + # 2 bodies * 6 = 12, minus 5 independent = 7 free pebbles + assert pg.get_dof() == 7 + + def test_fixed_removes_six_dof(self) -> None: + pg = PebbleGame3D() + results = pg.add_joint(_fixed(0, 0, 1)) + assert len(results) == 6 + assert all(r["independent"] for r in results) + assert pg.get_dof() == 6 + + def test_ball_removes_three_dof(self) -> None: + pg = PebbleGame3D() + j = Joint(joint_id=0, body_a=0, body_b=1, joint_type=JointType.BALL) + results = pg.add_joint(j) + assert len(results) == 3 + assert all(r["independent"] for r in results) + assert pg.get_dof() == 9 + + +class TestTwoBodiesRevolute: + """Two bodies connected by a revolute -- demo scenario 1.""" + + def test_internal_dof(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_revolute(0, 0, 1)) + # Total DOF = 7, internal = 7 - 6 = 1 + assert pg.get_internal_dof() == 1 + + def test_not_rigid(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_revolute(0, 0, 1)) + assert not pg.is_rigid() + + def test_classification(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_revolute(0, 0, 1)) + assert pg.classify_assembly() == "underconstrained" + + +class TestTwoBodiesFixed: + """Two bodies + fixed joint -- demo scenario 2.""" + + def test_zero_internal_dof(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_fixed(0, 0, 1)) + assert pg.get_internal_dof() == 0 + + def test_rigid(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_fixed(0, 0, 1)) + assert pg.is_rigid() + + def test_well_constrained(self) -> None: + pg = PebbleGame3D() + pg.add_joint(_fixed(0, 0, 1)) + assert pg.classify_assembly() == "well-constrained" + + +class TestTriangleRevolute: + """Triangle of 3 bodies with revolute joints -- demo scenario 3.""" + + @pytest.fixture() + def pg(self) -> PebbleGame3D: + pg = PebbleGame3D() + pg.add_joint(_revolute(0, 0, 1)) + pg.add_joint(_revolute(1, 1, 2)) + pg.add_joint(_revolute(2, 2, 0)) + return pg + + def test_has_redundant_edges(self, pg: PebbleGame3D) -> None: + assert pg.get_redundant_count() > 0 + + def test_classification_overconstrained(self, pg: PebbleGame3D) -> None: + # 15 constraints on 3 bodies (Maxwell: 6*3-6=12 needed) + assert pg.classify_assembly() in ("overconstrained", "mixed") + + def test_rigid(self, pg: PebbleGame3D) -> None: + assert pg.is_rigid() + + +class TestChainNotRigid: + """A serial chain of 4 bodies with revolute joints is never rigid.""" + + def test_chain_underconstrained(self) -> None: + pg = PebbleGame3D() + for i in range(3): + pg.add_joint(_revolute(i, i, i + 1)) + assert not pg.is_rigid() + assert pg.classify_assembly() == "underconstrained" + + def test_chain_internal_dof(self) -> None: + pg = PebbleGame3D() + for i in range(3): + pg.add_joint(_revolute(i, i, i + 1)) + # 4 bodies * 6 = 24, minus 15 independent = 9 free, internal = 3 + assert pg.get_internal_dof() == 3 + + +class TestEdgeResults: + """Result dicts returned by add_joint.""" + + def test_result_keys(self) -> None: + pg = PebbleGame3D() + results = pg.add_joint(_revolute(0, 0, 1)) + expected_keys = {"edge_id", "joint_id", "constraint_index", "independent", "dof_remaining"} + for r in results: + assert set(r.keys()) == expected_keys + + def test_edge_ids_sequential(self) -> None: + pg = PebbleGame3D() + r1 = pg.add_joint(_revolute(0, 0, 1)) + r2 = pg.add_joint(_revolute(1, 1, 2)) + all_ids = [r["edge_id"] for r in r1 + r2] + assert all_ids == list(range(10)) + + def test_dof_remaining_monotonic(self) -> None: + pg = PebbleGame3D() + results = pg.add_joint(_revolute(0, 0, 1)) + dofs = [r["dof_remaining"] for r in results] + # Should be non-increasing (each independent edge removes a pebble) + for a, b in itertools.pairwise(dofs): + assert a >= b + + +class TestGroundedClassification: + """classify_assembly with grounded=True.""" + + def test_grounded_baseline_zero(self) -> None: + """With grounded=True the baseline is 0 (not 6).""" + pg = PebbleGame3D() + pg.add_joint(_fixed(0, 0, 1)) + # Ungrounded: well-constrained (6 pebbles = baseline 6) + assert pg.classify_assembly(grounded=False) == "well-constrained" + # Grounded: the 6 remaining pebbles on body 1 exceed baseline 0, + # so the raw pebble game (without a virtual ground body) sees this + # as underconstrained. The analysis function handles this properly + # by adding a virtual ground body. + assert pg.classify_assembly(grounded=True) == "underconstrained" diff --git a/tests/datagen/test_types.py b/tests/datagen/test_types.py new file mode 100644 index 0000000..7667243 --- /dev/null +++ b/tests/datagen/test_types.py @@ -0,0 +1,163 @@ +"""Tests for solver.datagen.types -- shared data types.""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest + +from solver.datagen.types import ( + ConstraintAnalysis, + Joint, + JointType, + PebbleState, + RigidBody, +) + + +class TestJointType: + """JointType enum construction and DOF values.""" + + EXPECTED_DOF: ClassVar[dict[str, int]] = { + "FIXED": 6, + "REVOLUTE": 5, + "CYLINDRICAL": 4, + "SLIDER": 5, + "BALL": 3, + "PLANAR": 3, + "SCREW": 5, + "UNIVERSAL": 4, + "PARALLEL": 3, + "PERPENDICULAR": 1, + "DISTANCE": 1, + } + + def test_member_count(self) -> None: + assert len(JointType) == 11 + + @pytest.mark.parametrize("name,dof", EXPECTED_DOF.items()) + def test_dof_values(self, name: str, dof: int) -> None: + assert JointType[name].dof == dof + + def test_access_by_name(self) -> None: + assert JointType["REVOLUTE"] is JointType.REVOLUTE + + def test_value_is_tuple(self) -> None: + assert JointType.REVOLUTE.value == (1, 5) + assert JointType.REVOLUTE.dof == 5 + + +class TestRigidBody: + """RigidBody dataclass defaults and construction.""" + + def test_defaults(self) -> None: + body = RigidBody(body_id=0) + np.testing.assert_array_equal(body.position, np.zeros(3)) + np.testing.assert_array_equal(body.orientation, np.eye(3)) + assert body.local_anchors == {} + + def test_custom_position(self) -> None: + pos = np.array([1.0, 2.0, 3.0]) + body = RigidBody(body_id=7, position=pos) + np.testing.assert_array_equal(body.position, pos) + assert body.body_id == 7 + + def test_local_anchors_mutable(self) -> None: + body = RigidBody(body_id=0) + body.local_anchors["top"] = np.array([0.0, 0.0, 1.0]) + assert "top" in body.local_anchors + + def test_default_factory_isolation(self) -> None: + """Each instance gets its own default containers.""" + b1 = RigidBody(body_id=0) + b2 = RigidBody(body_id=1) + b1.local_anchors["x"] = np.zeros(3) + assert "x" not in b2.local_anchors + + +class TestJoint: + """Joint dataclass defaults and construction.""" + + def test_defaults(self) -> None: + j = Joint(joint_id=0, body_a=0, body_b=1, joint_type=JointType.REVOLUTE) + np.testing.assert_array_equal(j.anchor_a, np.zeros(3)) + np.testing.assert_array_equal(j.anchor_b, np.zeros(3)) + np.testing.assert_array_equal(j.axis, np.array([0.0, 0.0, 1.0])) + assert j.pitch == 0.0 + + def test_full_construction(self) -> None: + j = Joint( + joint_id=5, + body_a=2, + body_b=3, + joint_type=JointType.SCREW, + anchor_a=np.array([1.0, 0.0, 0.0]), + anchor_b=np.array([2.0, 0.0, 0.0]), + axis=np.array([1.0, 0.0, 0.0]), + pitch=0.5, + ) + assert j.joint_id == 5 + assert j.joint_type is JointType.SCREW + assert j.pitch == 0.5 + + +class TestPebbleState: + """PebbleState dataclass defaults.""" + + def test_defaults(self) -> None: + s = PebbleState() + assert s.free_pebbles == {} + assert s.directed_edges == {} + assert s.independent_edges == set() + assert s.redundant_edges == set() + assert s.incoming == {} + assert s.outgoing == {} + + def test_default_factory_isolation(self) -> None: + s1 = PebbleState() + s2 = PebbleState() + s1.free_pebbles[0] = 6 + assert 0 not in s2.free_pebbles + + +class TestConstraintAnalysis: + """ConstraintAnalysis dataclass construction.""" + + def test_construction(self) -> None: + ca = ConstraintAnalysis( + combinatorial_dof=6, + combinatorial_internal_dof=0, + combinatorial_redundant=0, + combinatorial_classification="well-constrained", + per_edge_results=[], + jacobian_rank=6, + jacobian_nullity=0, + jacobian_internal_dof=0, + numerically_dependent=[], + geometric_degeneracies=0, + is_rigid=True, + is_minimally_rigid=True, + ) + assert ca.is_rigid is True + assert ca.is_minimally_rigid is True + assert ca.combinatorial_classification == "well-constrained" + + def test_per_edge_results_typing(self) -> None: + """per_edge_results accepts list[dict[str, Any]].""" + ca = ConstraintAnalysis( + combinatorial_dof=7, + combinatorial_internal_dof=1, + combinatorial_redundant=0, + combinatorial_classification="underconstrained", + per_edge_results=[{"edge_id": 0, "independent": True}], + jacobian_rank=5, + jacobian_nullity=1, + jacobian_internal_dof=1, + numerically_dependent=[], + geometric_degeneracies=0, + is_rigid=False, + is_minimally_rigid=False, + ) + assert len(ca.per_edge_results) == 1 + assert ca.per_edge_results[0]["edge_id"] == 0