"""Tests for graph decomposition and cluster-by-cluster solving.""" import math import pytest from kindred_solver.constraints import ( CoincidentConstraint, CylindricalConstraint, FixedConstraint, ParallelConstraint, RevoluteConstraint, SliderConstraint, ) from kindred_solver.decompose import ( SolveCluster, _constraints_for_bodies, build_constraint_graph_simple, build_solve_order, classify_cluster_rigidity, find_clusters, residual_count_by_name, solve_decomposed, ) from kindred_solver.dof import count_dof from kindred_solver.entities import RigidBody from kindred_solver.newton import newton_solve from kindred_solver.params import ParamTable from kindred_solver.prepass import single_equation_pass, substitution_pass ID_QUAT = (1, 0, 0, 0) c45 = math.cos(math.pi / 4) s45 = math.sin(math.pi / 4) ROT_90Z = (c45, 0, 0, s45) # ============================================================================ # DOF table tests # ============================================================================ class TestResidualCount: def test_known_types(self): assert residual_count_by_name("Fixed") == 6 assert residual_count_by_name("Revolute") == 5 assert residual_count_by_name("Cylindrical") == 4 assert residual_count_by_name("Ball") == 3 assert residual_count_by_name("Coincident") == 3 assert residual_count_by_name("Parallel") == 2 assert residual_count_by_name("Perpendicular") == 1 assert residual_count_by_name("DistancePointPoint") == 1 def test_stubs_zero(self): assert residual_count_by_name("Cam") == 0 assert residual_count_by_name("Slot") == 0 def test_unknown_zero(self): assert residual_count_by_name("DoesNotExist") == 0 # ============================================================================ # Graph construction tests # ============================================================================ class TestBuildConstraintGraph: def test_simple_pair(self): """Two bodies, one Revolute constraint.""" G = build_constraint_graph_simple( [("A", "B", "Revolute", 0)], grounded={"A"}, ) assert set(G.nodes()) == {"A", "B"} assert G.number_of_edges() == 1 assert G.nodes["A"]["grounded"] is True assert G.nodes["B"]["grounded"] is False def test_chain(self): """Chain: A-B-C-D.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Revolute", 1), ("C", "D", "Revolute", 2), ], ) assert G.number_of_nodes() == 4 assert G.number_of_edges() == 3 def test_multi_edge(self): """Two constraints between same body pair → two edges.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("A", "B", "Parallel", 1), ], ) assert G.number_of_nodes() == 2 assert G.number_of_edges() == 2 def test_stub_excluded(self): """Stub constraints (Cam, Slot) excluded from graph.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Cam", 1), ], ) assert G.number_of_nodes() == 2 # C not added assert G.number_of_edges() == 1 # ============================================================================ # Biconnected decomposition tests # ============================================================================ class TestFindClusters: def test_chain_decomposes(self): """Chain A-B-C-D: 3 biconnected components, 2 articulation points.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Revolute", 1), ("C", "D", "Revolute", 2), ], ) clusters, artic = find_clusters(G) assert len(clusters) == 3 assert artic == {"B", "C"} def test_loop_single_cluster(self): """Loop A-B-C-A: single biconnected component, no articulation points.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Revolute", 1), ("C", "A", "Revolute", 2), ], ) clusters, artic = find_clusters(G) assert len(clusters) == 1 assert artic == set() def test_star_decomposes(self): """Star: center connected to 3 leaves. 3 biconnected components (each is center + one leaf). Center is the only articulation point. """ G = build_constraint_graph_simple( [ ("center", "L1", "Revolute", 0), ("center", "L2", "Revolute", 1), ("center", "L3", "Revolute", 2), ], ) clusters, artic = find_clusters(G) assert len(clusters) == 3 assert artic == {"center"} def test_single_edge(self): """Two nodes, one edge: single biconnected component.""" G = build_constraint_graph_simple( [("A", "B", "Fixed", 0)], ) clusters, artic = find_clusters(G) assert len(clusters) == 1 assert artic == set() def test_tree_with_loop(self): """A-B-C-D with a loop B-C-E-B. The loop {B, C, E} forms one biconnected component. A-B and C-D are separate biconnected components. Articulation points: B and C. """ G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Revolute", 1), ("C", "E", "Revolute", 2), ("E", "B", "Revolute", 3), ("C", "D", "Revolute", 4), ], ) clusters, artic = find_clusters(G) # 3 clusters: {A,B}, {B,C,E}, {C,D} assert len(clusters) == 3 assert artic == {"B", "C"} # ============================================================================ # Solve order tests # ============================================================================ class TestBuildSolveOrder: def test_chain_grounded_at_start(self): """Chain: ground-A-B-C. Root cluster is {ground,A}, solved first. Then {A,B}, then {B,C}. """ G = build_constraint_graph_simple( [ ("ground", "A", "Revolute", 0), ("A", "B", "Revolute", 1), ("B", "C", "Revolute", 2), ], grounded={"ground"}, ) clusters, artic = find_clusters(G) solve_order = build_solve_order(G, clusters, artic, {"ground"}) assert len(solve_order) == 3 # First cluster should contain ground (root) assert "ground" in solve_order[0].bodies assert solve_order[0].has_ground is True # Last cluster should be the leaf assert "C" in solve_order[-1].bodies def test_star_grounded_center(self): """Star with grounded center: root cluster first, then leaves.""" G = build_constraint_graph_simple( [ ("center", "L1", "Revolute", 0), ("center", "L2", "Revolute", 1), ("center", "L3", "Revolute", 2), ], grounded={"center"}, ) clusters, artic = find_clusters(G) solve_order = build_solve_order(G, clusters, artic, {"center"}) assert len(solve_order) == 3 # First cluster contains grounded center assert "center" in solve_order[0].bodies assert solve_order[0].has_ground is True # All clusters have center as boundary for sc in solve_order: assert "center" in sc.boundary_bodies or "center" in sc.bodies def test_single_cluster_no_boundary(self): """Single cluster: no boundary bodies.""" G = build_constraint_graph_simple( [("A", "B", "Revolute", 0)], grounded={"A"}, ) clusters, artic = find_clusters(G) solve_order = build_solve_order(G, clusters, artic, {"A"}) assert len(solve_order) == 1 assert solve_order[0].boundary_bodies == set() assert solve_order[0].has_ground is True def test_constraint_assignment(self): """Constraints are correctly assigned to clusters.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("B", "C", "Revolute", 1), ], ) clusters, artic = find_clusters(G) solve_order = build_solve_order(G, clusters, artic, set()) # Each cluster should have exactly 1 constraint all_indices = [] for sc in solve_order: all_indices.extend(sc.constraint_indices) assert sorted(all_indices) == [0, 1] # ============================================================================ # Integration: cluster solving with actual Newton solver # ============================================================================ def _monolithic_solve(pt, all_bodies, constraint_objs): """Monolithic solve for comparison.""" residuals = [] for c in constraint_objs: residuals.extend(c.residuals()) quat_groups = [] for body in all_bodies: if not body.grounded: residuals.append(body.quat_norm_residual()) quat_groups.append(body.quat_param_names()) residuals = substitution_pass(residuals, pt) residuals = single_equation_pass(residuals, pt) return newton_solve(residuals, pt, quat_groups=quat_groups, max_iter=100, tol=1e-10) class TestClusterSolve: def test_chain_solve_matches_monolithic(self): """Chain: ground → A → B. Decomposed solve matches monolithic. Two biconnected components: {ground,A} and {A,B}. A is the articulation point. """ # --- Decomposed solve --- pt = ParamTable() ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) body_a = RigidBody("A", pt, (3, 2, 0), ID_QUAT) body_b = RigidBody("B", pt, (8, 3, 0), ID_QUAT) bodies = {"ground": ground, "A": body_a, "B": body_b} c0 = RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, body_a, (0, 0, 0), ID_QUAT) c1 = RevoluteConstraint(body_a, (5, 0, 0), ID_QUAT, body_b, (0, 0, 0), ID_QUAT) constraint_objs = [c0, c1] constraint_indices = [0, 1] # Solve order: grounded cluster first, then outward. # {ground,A} first (A gets solved), then {A,B} with A as boundary. clusters = [ SolveCluster( bodies={"ground", "A"}, constraint_indices=[0], boundary_bodies={"A"}, has_ground=True, ), SolveCluster( bodies={"A", "B"}, constraint_indices=[1], boundary_bodies={"A"}, has_ground=False, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt, ) assert converged env = pt.get_env() pos_a = body_a.extract_position(env) pos_b = body_b.extract_position(env) # A should be at origin (revolute to ground origin) assert abs(pos_a[0]) < 1e-8 assert abs(pos_a[1]) < 1e-8 assert abs(pos_a[2]) < 1e-8 # B should be at (5,0,0) — revolute at A's (5,0,0) marker, B's origin assert abs(pos_b[0] - 5.0) < 1e-8 assert abs(pos_b[1]) < 1e-8 assert abs(pos_b[2]) < 1e-8 def test_star_solve(self): """Star: grounded center with 3 revolute arms. Each arm is a separate cluster. Center is articulation point. Solve center-containing clusters first, then arms. """ pt = ParamTable() center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True) arm1 = RigidBody("L1", pt, (5, 5, 0), ID_QUAT) arm2 = RigidBody("L2", pt, (5, -5, 0), ID_QUAT) arm3 = RigidBody("L3", pt, (-5, 5, 0), ID_QUAT) bodies = {"C": center, "L1": arm1, "L2": arm2, "L3": arm3} c0 = RevoluteConstraint(center, (3, 0, 0), ID_QUAT, arm1, (0, 0, 0), ID_QUAT) c1 = RevoluteConstraint(center, (0, 3, 0), ID_QUAT, arm2, (0, 0, 0), ID_QUAT) c2 = RevoluteConstraint(center, (-3, 0, 0), ID_QUAT, arm3, (0, 0, 0), ID_QUAT) constraint_objs = [c0, c1, c2] constraint_indices = [0, 1, 2] # Each arm cluster has center as boundary. Since center is grounded, # its params are already fixed. Solve order doesn't matter much. clusters = [ SolveCluster( bodies={"C", "L1"}, constraint_indices=[0], boundary_bodies={"C"}, has_ground=True, ), SolveCluster( bodies={"C", "L2"}, constraint_indices=[1], boundary_bodies={"C"}, has_ground=True, ), SolveCluster( bodies={"C", "L3"}, constraint_indices=[2], boundary_bodies={"C"}, has_ground=True, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt, ) assert converged env = pt.get_env() # Each arm should be at its marker position assert abs(arm1.extract_position(env)[0] - 3.0) < 1e-8 assert abs(arm2.extract_position(env)[1] - 3.0) < 1e-8 assert abs(arm3.extract_position(env)[0] + 3.0) < 1e-8 def test_single_cluster_solve(self): """Single cluster: decomposed solve behaves same as monolithic.""" pt = ParamTable() ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True) body = RigidBody("b", pt, (5, 5, 5), ID_QUAT) bodies = {"g": ground, "b": body} c0 = RevoluteConstraint(ground, (0, 0, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT) constraint_objs = [c0] constraint_indices = [0] clusters = [ SolveCluster( bodies={"g", "b"}, constraint_indices=[0], boundary_bodies=set(), has_ground=True, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt, ) assert converged env = pt.get_env() pos = body.extract_position(env) assert abs(pos[0]) < 1e-8 assert abs(pos[1]) < 1e-8 assert abs(pos[2]) < 1e-8 # ============================================================================ # Full decompose() pipeline tests # ============================================================================ class TestDecomposePipeline: """Test the full decompose() entry point (requires kcsolve for build_constraint_graph, so we test via build_constraint_graph_simple and manual cluster construction instead).""" def test_chain_topology(self): """Chain decomposition produces correct number of clusters.""" G = build_constraint_graph_simple( [ ("g", "A", "Revolute", 0), ("A", "B", "Revolute", 1), ("B", "C", "Revolute", 2), ("C", "D", "Revolute", 3), ], grounded={"g"}, ) clusters, artic = find_clusters(G) solve_order = build_solve_order(G, clusters, artic, {"g"}) # Chain of 5 nodes → 4 biconnected components assert len(solve_order) == 4 # First cluster should contain ground assert "g" in solve_order[0].bodies def test_disconnected_components(self): """Two disconnected sub-assemblies produce independent cluster groups.""" G = build_constraint_graph_simple( [ ("A", "B", "Revolute", 0), ("C", "D", "Revolute", 1), ], ) # Two connected components, each single-cluster components = list(nx.connected_components(G)) assert len(components) == 2 # Each component decomposes to 1 cluster for comp_nodes in components: sub = G.subgraph(comp_nodes).copy() clusters, artic = find_clusters(sub) assert len(clusters) == 1 assert artic == set() import networkx as nx from kindred_solver.bfgs import bfgs_solve # ============================================================================ # Integration: decomposed solve vs monolithic on larger assemblies # ============================================================================ class TestDecomposedVsMonolithic: """Verify decomposed solving produces equivalent results to monolithic.""" def _build_chain(self, n_bodies, pt, displaced=True): """Build a revolute chain: ground → B0 → B1 → ... → B(n-1). Each body is connected at (2,0,0) local → (0,0,0) local of next. If displaced, bodies start at slightly wrong positions. Returns (bodies_dict, constraint_objs, constraint_indices). """ ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) bodies = {"ground": ground} constraints = [] indices = [] prev_body = ground for i in range(n_bodies): name = f"B{i}" # Correct position would be ((i+1)*2, 0, 0), displace slightly x = (i + 1) * 2 + (0.5 if displaced else 0) y = 0.3 if displaced else 0 body = RigidBody(name, pt, (x, y, 0), ID_QUAT) bodies[name] = body c = RevoluteConstraint( prev_body, (2, 0, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT, ) constraints.append(c) indices.append(i) prev_body = body return bodies, constraints, indices def test_chain_10_bodies(self): """10-body revolute chain: decomposed solve converges and satisfies constraints. Revolute chains are under-constrained (each joint has 1 rotation DOF), so monolithic and decomposed may find different valid solutions. We verify convergence and that all constraint residuals are near zero. """ pt_dec = ParamTable() bodies_dec, constraints_dec, indices_dec = self._build_chain(10, pt_dec) # Build clusters manually (chain -> each link is a biconnected component) cluster_list = [] body_names = ["ground"] + [f"B{i}" for i in range(10)] for i in range(10): pair = {body_names[i], body_names[i + 1]} boundary = set() if i > 0: boundary.add(body_names[i]) # shared with previous cluster if i < 9: boundary.add(body_names[i + 1]) # shared with next cluster cluster_list.append( SolveCluster( bodies=pair, constraint_indices=[i], boundary_bodies=boundary, has_ground=(i == 0), ) ) converged = solve_decomposed( cluster_list, bodies_dec, constraints_dec, indices_dec, pt_dec, ) assert converged # Verify all constraint residuals are satisfied env = pt_dec.get_env() for c in constraints_dec: for r in c.residuals(): val = r.eval(env) assert abs(val) < 1e-8, f"Residual not satisfied: {val}" # Verify quat norms for name, body in bodies_dec.items(): if not body.grounded: qn = body.quat_norm_residual().eval(env) assert abs(qn) < 1e-8, f"{name} quat not normalized: {qn}" def test_star_5_arms(self): """Star with 5 revolute arms off a grounded center.""" pt = ParamTable() center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True) bodies = {"C": center} constraints = [] indices = [] clusters = [] import math as m for i in range(5): name = f"arm{i}" angle = 2 * m.pi * i / 5 marker_pos = (3 * m.cos(angle), 3 * m.sin(angle), 0) body = RigidBody( name, pt, (marker_pos[0] + 1, marker_pos[1] + 1, 0), ID_QUAT ) bodies[name] = body c = RevoluteConstraint( center, marker_pos, ID_QUAT, body, (0, 0, 0), ID_QUAT, ) constraints.append(c) indices.append(i) clusters.append( SolveCluster( bodies={"C", name}, constraint_indices=[i], boundary_bodies={"C"}, has_ground=True, ) ) converged = solve_decomposed(clusters, bodies, constraints, indices, pt) assert converged env = pt.get_env() for i in range(5): name = f"arm{i}" angle = 2 * m.pi * i / 5 expected_x = 3 * m.cos(angle) expected_y = 3 * m.sin(angle) pos = bodies[name].extract_position(env) assert abs(pos[0] - expected_x) < 1e-6, ( f"{name}: x={pos[0]:.4f} expected={expected_x:.4f}" ) assert abs(pos[1] - expected_y) < 1e-6, ( f"{name}: y={pos[1]:.4f} expected={expected_y:.4f}" ) assert abs(pos[2]) < 1e-6 def test_single_cluster_loop(self): """A loop assembly (single biconnected component) solves as one cluster.""" pt = ParamTable() ground = RigidBody("g", pt, (0, 0, 0), ID_QUAT, grounded=True) b1 = RigidBody("B1", pt, (2.5, 0.3, 0), ID_QUAT) b2 = RigidBody("B2", pt, (4.5, -0.2, 0), ID_QUAT) bodies = {"g": ground, "B1": b1, "B2": b2} c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT) c1 = RevoluteConstraint(b1, (2, 0, 0), ID_QUAT, b2, (0, 0, 0), ID_QUAT) c2 = RevoluteConstraint(b2, (2, 0, 0), ID_QUAT, ground, (6, 0, 0), ID_QUAT) constraint_objs = [c0, c1, c2] constraint_indices = [0, 1, 2] clusters = [ SolveCluster( bodies={"g", "B1", "B2"}, constraint_indices=[0, 1, 2], boundary_bodies=set(), has_ground=True, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt ) assert converged env = pt.get_env() for c in constraint_objs: for r in c.residuals(): assert abs(r.eval(env)) < 1e-8 def test_boundary_propagation_accuracy(self): """Verify boundary body values propagate correctly between clusters. Chain: ground → A → B → C. A is boundary between clusters 0-1. B is boundary between clusters 1-2. After decomposed solve, all revolute constraints should be satisfied. """ pt = ParamTable() ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) body_a = RigidBody("A", pt, (2.5, 0.5, 0), ID_QUAT) # displaced body_b = RigidBody("B", pt, (4.5, -0.5, 0), ID_QUAT) # displaced body_c = RigidBody("C", pt, (7.0, 1.0, 0), ID_QUAT) # displaced bodies = {"ground": ground, "A": body_a, "B": body_b, "C": body_c} c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, body_a, (0, 0, 0), ID_QUAT) c1 = RevoluteConstraint(body_a, (2, 0, 0), ID_QUAT, body_b, (0, 0, 0), ID_QUAT) c2 = RevoluteConstraint(body_b, (2, 0, 0), ID_QUAT, body_c, (0, 0, 0), ID_QUAT) constraint_objs = [c0, c1, c2] constraint_indices = [0, 1, 2] clusters = [ SolveCluster( bodies={"ground", "A"}, constraint_indices=[0], boundary_bodies={"A"}, has_ground=True, ), SolveCluster( bodies={"A", "B"}, constraint_indices=[1], boundary_bodies={"A", "B"}, has_ground=False, ), SolveCluster( bodies={"B", "C"}, constraint_indices=[2], boundary_bodies={"B"}, has_ground=False, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt ) assert converged env = pt.get_env() # A at (2,0,0), B at (4,0,0), C at (6,0,0) pos_a = body_a.extract_position(env) pos_b = body_b.extract_position(env) pos_c = body_c.extract_position(env) assert abs(pos_a[0] - 2.0) < 1e-6 assert abs(pos_a[1]) < 1e-6 assert abs(pos_b[0] - 4.0) < 1e-6 assert abs(pos_b[1]) < 1e-6 assert abs(pos_c[0] - 6.0) < 1e-6 assert abs(pos_c[1]) < 1e-6 # Verify all params are free after solve (unfixed correctly) for name in ["A", "B", "C"]: body = bodies[name] for pname in body._param_names: assert not pt.is_fixed(pname), f"{pname} still fixed after solve" # ============================================================================ # Pebble game rigidity classification # ============================================================================ class TestPebbleGameClassification: """Test classify_cluster_rigidity using the pebble game.""" def test_well_constrained_fixed_pair(self): """Two bodies with a Fixed joint (6 residuals) → well-constrained.""" G = build_constraint_graph_simple( [("ground", "B", "Fixed", 0)], grounded={"ground"}, ) cluster = SolveCluster( bodies={"ground", "B"}, constraint_indices=[0], boundary_bodies=set(), has_ground=True, ) result = classify_cluster_rigidity(cluster, G, {"ground"}) assert result == "well-constrained", f"Expected well-constrained, got {result}" def test_underconstrained_revolute_pair(self): """Ground + body with Revolute (5 residuals) → 1 DOF under-constrained.""" G = build_constraint_graph_simple( [("ground", "B", "Revolute", 0)], grounded={"ground"}, ) cluster = SolveCluster( bodies={"ground", "B"}, constraint_indices=[0], boundary_bodies=set(), has_ground=True, ) result = classify_cluster_rigidity(cluster, G, {"ground"}) assert result == "underconstrained", f"Expected underconstrained, got {result}" def test_overconstrained(self): """Two Fixed joints between same pair → overconstrained.""" G = build_constraint_graph_simple( [ ("ground", "B", "Fixed", 0), ("ground", "B", "Fixed", 1), ], grounded={"ground"}, ) cluster = SolveCluster( bodies={"ground", "B"}, constraint_indices=[0, 1], boundary_bodies=set(), has_ground=True, ) result = classify_cluster_rigidity(cluster, G, {"ground"}) assert result == "overconstrained", f"Expected overconstrained, got {result}" def test_manual_edge_types(self): """Parallel (2 residuals, manual edge) is correctly handled.""" G = build_constraint_graph_simple( [ ("ground", "B", "Coincident", 0), # 3 residuals ("ground", "B", "Parallel", 1), # 2 residuals (manual) ], grounded={"ground"}, ) # Total: 3 + 2 = 5 residuals (like Revolute) → underconstrained (1 DOF) cluster = SolveCluster( bodies={"ground", "B"}, constraint_indices=[0, 1], boundary_bodies=set(), has_ground=True, ) result = classify_cluster_rigidity(cluster, G, {"ground"}) assert result == "underconstrained", f"Expected underconstrained, got {result}" def test_no_ground(self): """Two ungrounded bodies with Fixed joint → well-constrained (6 DOF trivial).""" G = build_constraint_graph_simple( [("A", "B", "Fixed", 0)], ) cluster = SolveCluster( bodies={"A", "B"}, constraint_indices=[0], boundary_bodies=set(), has_ground=False, ) result = classify_cluster_rigidity(cluster, G, set()) # Two bodies, no ground: 12 DOF total, 6 from Fixed, 6 trivial → well-constrained assert result == "well-constrained", f"Expected well-constrained, got {result}" # ============================================================================ # Large assembly tests (20+ bodies) # ============================================================================ class TestLargeAssembly: """End-to-end tests with large synthetic assemblies.""" def test_chain_20_bodies(self): """20-body revolute chain: decomposed solve converges.""" pt = ParamTable() ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) bodies = {"ground": ground} constraints = [] indices = [] prev = ground for i in range(20): name = f"B{i}" body = RigidBody(name, pt, ((i + 1) * 2 + 0.3, 0.2, 0), ID_QUAT) bodies[name] = body c = RevoluteConstraint(prev, (2, 0, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT) constraints.append(c) indices.append(i) prev = body # Build clusters: chain of 20 biconnected components body_names = ["ground"] + [f"B{i}" for i in range(20)] cluster_list = [] for i in range(20): pair = {body_names[i], body_names[i + 1]} boundary = set() if i > 0: boundary.add(body_names[i]) if i < 19: boundary.add(body_names[i + 1]) cluster_list.append( SolveCluster( bodies=pair, constraint_indices=[i], boundary_bodies=boundary, has_ground=(i == 0), ) ) converged = solve_decomposed(cluster_list, bodies, constraints, indices, pt) assert converged env = pt.get_env() for c in constraints: for r in c.residuals(): assert abs(r.eval(env)) < 1e-8 def test_star_10_arms(self): """Star with 10 revolute arms: 10 independent clusters.""" pt = ParamTable() center = RigidBody("C", pt, (0, 0, 0), ID_QUAT, grounded=True) bodies = {"C": center} constraints = [] indices = [] clusters = [] for i in range(10): name = f"arm{i}" angle = 2 * math.pi * i / 10 mx = 3 * math.cos(angle) my = 3 * math.sin(angle) body = RigidBody(name, pt, (mx + 0.5, my + 0.5, 0), ID_QUAT) bodies[name] = body c = RevoluteConstraint( center, (mx, my, 0), ID_QUAT, body, (0, 0, 0), ID_QUAT ) constraints.append(c) indices.append(i) clusters.append( SolveCluster( bodies={"C", name}, constraint_indices=[i], boundary_bodies={"C"}, has_ground=True, ) ) converged = solve_decomposed(clusters, bodies, constraints, indices, pt) assert converged env = pt.get_env() for i in range(10): name = f"arm{i}" angle = 2 * math.pi * i / 10 pos = bodies[name].extract_position(env) assert abs(pos[0] - 3 * math.cos(angle)) < 1e-6 assert abs(pos[1] - 3 * math.sin(angle)) < 1e-6 def test_tree_assembly(self): """Tree: ground → A → (B, C), A → D. Mixed branching. Topology: ground-A is root cluster, A is articulation point. Three leaf clusters: {A,B}, {A,C}, {A,D}. """ pt = ParamTable() ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) a = RigidBody("A", pt, (2.5, 0.3, 0), ID_QUAT) b = RigidBody("B", pt, (4.5, 2.5, 0), ID_QUAT) c = RigidBody("C", pt, (4.5, -2.5, 0), ID_QUAT) d = RigidBody("D", pt, (4.5, 0.3, 0), ID_QUAT) bodies = {"ground": ground, "A": a, "B": b, "C": c, "D": d} c0 = RevoluteConstraint(ground, (2, 0, 0), ID_QUAT, a, (0, 0, 0), ID_QUAT) c1 = RevoluteConstraint(a, (2, 2, 0), ID_QUAT, b, (0, 0, 0), ID_QUAT) c2 = RevoluteConstraint(a, (2, -2, 0), ID_QUAT, c, (0, 0, 0), ID_QUAT) c3 = RevoluteConstraint(a, (2, 0, 0), ID_QUAT, d, (0, 0, 0), ID_QUAT) constraint_objs = [c0, c1, c2, c3] constraint_indices = [0, 1, 2, 3] # Root cluster: {ground, A}, then 3 leaf clusters clusters = [ SolveCluster( bodies={"ground", "A"}, constraint_indices=[0], boundary_bodies={"A"}, has_ground=True, ), SolveCluster( bodies={"A", "B"}, constraint_indices=[1], boundary_bodies={"A"}, has_ground=False, ), SolveCluster( bodies={"A", "C"}, constraint_indices=[2], boundary_bodies={"A"}, has_ground=False, ), SolveCluster( bodies={"A", "D"}, constraint_indices=[3], boundary_bodies={"A"}, has_ground=False, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt ) assert converged env = pt.get_env() # A should be at (2, 0, 0) pos_a = a.extract_position(env) assert abs(pos_a[0] - 2.0) < 1e-6 assert abs(pos_a[1]) < 1e-6 # B at A + marker (2,2,0) = (4,2,0) pos_b = b.extract_position(env) assert abs(pos_b[0] - 4.0) < 1e-6 assert abs(pos_b[1] - 2.0) < 1e-6 # C at A + marker (2,-2,0) = (4,-2,0) pos_c = c.extract_position(env) assert abs(pos_c[0] - 4.0) < 1e-6 assert abs(pos_c[1] + 2.0) < 1e-6 # D at A + marker (2,0,0) = (4,0,0) pos_d = d.extract_position(env) assert abs(pos_d[0] - 4.0) < 1e-6 assert abs(pos_d[1]) < 1e-6 def test_mixed_constraint_chain(self): """Chain with mixed constraint types: Fixed, Coincident, Cylindrical.""" pt = ParamTable() ground = RigidBody("ground", pt, (0, 0, 0), ID_QUAT, grounded=True) b1 = RigidBody("B1", pt, (2.5, 0.3, 0), ID_QUAT) b2 = RigidBody("B2", pt, (5.5, -0.2, 0), ID_QUAT) b3 = RigidBody("B3", pt, (8.5, 0.4, 0), ID_QUAT) bodies = {"ground": ground, "B1": b1, "B2": b2, "B3": b3} c0 = FixedConstraint(ground, (2, 0, 0), ID_QUAT, b1, (0, 0, 0), ID_QUAT) c1 = CoincidentConstraint(b1, (3, 0, 0), b2, (0, 0, 0)) c2 = CylindricalConstraint(b2, (3, 0, 0), ID_QUAT, b3, (0, 0, 0), ID_QUAT) constraint_objs = [c0, c1, c2] constraint_indices = [0, 1, 2] clusters = [ SolveCluster( bodies={"ground", "B1"}, constraint_indices=[0], boundary_bodies={"B1"}, has_ground=True, ), SolveCluster( bodies={"B1", "B2"}, constraint_indices=[1], boundary_bodies={"B1", "B2"}, has_ground=False, ), SolveCluster( bodies={"B2", "B3"}, constraint_indices=[2], boundary_bodies={"B2"}, has_ground=False, ), ] converged = solve_decomposed( clusters, bodies, constraint_objs, constraint_indices, pt ) assert converged env = pt.get_env() for c in constraint_objs: for r in c.residuals(): assert abs(r.eval(env)) < 1e-8