test: add unit tests for datagen modules
Some checks failed
CI / lint (push) Has been cancelled
CI / type-check (push) Has been cancelled
CI / test (push) Has been cancelled

- test_types.py: JointType enum values/count, dataclass defaults/isolation
- test_pebble_game.py: DOF accounting, rigidity, classification, edge results
- test_jacobian.py: Jacobian shape per joint type, rank, parallel axis degeneracy
- test_analysis.py: demo scenarios (revolute, fixed, triangle, parallel axes)
- test_generator.py: chain/rigid/overconstrained generation, training batch

Bug fixes found during testing:
- JointType enum: duplicate int values caused aliasing (SLIDER=REVOLUTE etc).
  Changed to (ordinal, dof) tuple values with a .dof property.
- pebble_game.py: .value -> .dof for constraint count
- analysis.py: classify from effective DOF (not raw pebble game with virtual
  ground body skew)

105 tests, all passing.

Closes #6
This commit is contained in:
2026-02-02 14:08:22 -06:00
parent 831a10cdb4
commit dc742bfc82
9 changed files with 1073 additions and 14 deletions

View File

@@ -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)

View File

@@ -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):

View File

@@ -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]
# ---------------------------------------------------------------------------

View File

View File

@@ -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

View File

@@ -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

View File

@@ -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)

View File

@@ -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"

163
tests/datagen/test_types.py Normal file
View File

@@ -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