Files
solver/kindred_solver/solver.py
forbes-0023 92ae57751f feat(solver): graph decomposition for cluster-by-cluster solving (phase 3)
Add a Python decomposition layer using NetworkX that partitions the
constraint graph into biconnected components (rigid clusters), orders
them via a block-cut tree, and solves each cluster independently.
Articulation-point bodies propagate as boundary conditions between
clusters.

New module kindred_solver/decompose.py:
- DOF table mapping BaseJointKind to residual counts
- Constraint graph construction (nx.MultiGraph)
- Biconnected component detection + articulation points
- Block-cut tree solve ordering (root-first from grounded cluster)
- Cluster-by-cluster solver with boundary body fix/unfix cycling
- Pebble game integration for per-cluster rigidity classification

Changes to existing modules:
- params.py: add unfix() for boundary body cycling
- solver.py: extract _monolithic_solve(), add decomposition branch
  for assemblies with >= 8 free bodies

Performance: for k clusters of ~n/k params each, total cost drops
from O(n^3) to O(n^3/k^2).

220 tests passing (up from 207).
2026-02-20 22:19:35 -06:00

434 lines
12 KiB
Python

"""KindredSolver — IKCSolver implementation bridging KCSolve to the
expression-based Newton-Raphson solver."""
from __future__ import annotations
import kcsolve
from .bfgs import bfgs_solve
from .constraints import (
AngleConstraint,
BallConstraint,
CamConstraint,
CoincidentConstraint,
ConcentricConstraint,
ConstraintBase,
CylindricalConstraint,
DistanceCylSphConstraint,
DistancePointPointConstraint,
FixedConstraint,
GearConstraint,
LineInPlaneConstraint,
ParallelConstraint,
PerpendicularConstraint,
PlanarConstraint,
PointInPlaneConstraint,
PointOnLineConstraint,
RackPinionConstraint,
RevoluteConstraint,
ScrewConstraint,
SliderConstraint,
SlotConstraint,
TangentConstraint,
UniversalConstraint,
)
from .decompose import decompose, solve_decomposed
from .dof import count_dof
from .entities import RigidBody
from .newton import newton_solve
from .params import ParamTable
from .prepass import single_equation_pass, substitution_pass
# Assemblies with fewer free bodies than this use the monolithic path.
_DECOMPOSE_THRESHOLD = 8
# All BaseJointKind values this solver can handle
_SUPPORTED = {
# Phase 1
kcsolve.BaseJointKind.Coincident,
kcsolve.BaseJointKind.DistancePointPoint,
kcsolve.BaseJointKind.Fixed,
# Phase 2: point constraints
kcsolve.BaseJointKind.PointOnLine,
kcsolve.BaseJointKind.PointInPlane,
# Phase 2: orientation
kcsolve.BaseJointKind.Parallel,
kcsolve.BaseJointKind.Perpendicular,
kcsolve.BaseJointKind.Angle,
# Phase 2: axis/surface
kcsolve.BaseJointKind.Concentric,
kcsolve.BaseJointKind.Tangent,
kcsolve.BaseJointKind.Planar,
kcsolve.BaseJointKind.LineInPlane,
# Phase 2: kinematic joints
kcsolve.BaseJointKind.Ball,
kcsolve.BaseJointKind.Revolute,
kcsolve.BaseJointKind.Cylindrical,
kcsolve.BaseJointKind.Slider,
kcsolve.BaseJointKind.Screw,
kcsolve.BaseJointKind.Universal,
# Phase 2: mechanical
kcsolve.BaseJointKind.Gear,
kcsolve.BaseJointKind.RackPinion,
}
class KindredSolver(kcsolve.IKCSolver):
"""Expression-based Newton-Raphson constraint solver."""
def name(self):
return "Kindred (Newton-Raphson)"
def supported_joints(self):
return list(_SUPPORTED)
def solve(self, ctx):
params = ParamTable()
bodies = {} # part_id -> RigidBody
# 1. Build entities from parts
for part in ctx.parts:
pos = tuple(part.placement.position)
quat = tuple(part.placement.quaternion) # (w, x, y, z)
body = RigidBody(
part.id,
params,
position=pos,
quaternion=quat,
grounded=part.grounded,
)
bodies[part.id] = body
# 2. Build constraint residuals (track index mapping for decomposition)
all_residuals = []
constraint_objs = []
constraint_indices = [] # parallel to constraint_objs: index in ctx.constraints
for idx, c in enumerate(ctx.constraints):
if not c.activated:
continue
body_i = bodies.get(c.part_i)
body_j = bodies.get(c.part_j)
if body_i is None or body_j is None:
continue
marker_i_pos = tuple(c.marker_i.position)
marker_j_pos = tuple(c.marker_j.position)
obj = _build_constraint(
c.type,
body_i,
marker_i_pos,
body_j,
marker_j_pos,
c.marker_i,
c.marker_j,
c.params,
)
if obj is None:
continue
constraint_objs.append(obj)
constraint_indices.append(idx)
all_residuals.extend(obj.residuals())
# 3. Add quaternion normalization residuals for non-grounded bodies
quat_groups = []
for body in bodies.values():
if not body.grounded:
all_residuals.append(body.quat_norm_residual())
quat_groups.append(body.quat_param_names())
# 4. Pre-passes on full system
all_residuals = substitution_pass(all_residuals, params)
all_residuals = single_equation_pass(all_residuals, params)
# 5. Solve (decomposed for large assemblies, monolithic for small)
n_free_bodies = sum(1 for b in bodies.values() if not b.grounded)
if n_free_bodies >= _DECOMPOSE_THRESHOLD:
grounded_ids = {pid for pid, b in bodies.items() if b.grounded}
clusters = decompose(ctx.constraints, grounded_ids)
if len(clusters) > 1:
converged = solve_decomposed(
clusters,
bodies,
constraint_objs,
constraint_indices,
params,
)
else:
converged = _monolithic_solve(
all_residuals,
params,
quat_groups,
)
else:
converged = _monolithic_solve(all_residuals, params, quat_groups)
# 6. DOF
dof = count_dof(all_residuals, params)
# 7. Build result
result = kcsolve.SolveResult()
result.status = (
kcsolve.SolveStatus.Success if converged else kcsolve.SolveStatus.Failed
)
result.dof = dof
env = params.get_env()
placements = []
for body in bodies.values():
if body.grounded:
continue
pr = kcsolve.SolveResult.PartResult()
pr.id = body.part_id
pr.placement = kcsolve.Transform()
pr.placement.position = list(body.extract_position(env))
pr.placement.quaternion = list(body.extract_quaternion(env))
placements.append(pr)
result.placements = placements
return result
def is_deterministic(self):
return True
def _monolithic_solve(all_residuals, params, quat_groups):
"""Newton-Raphson solve with BFGS fallback on the full system."""
converged = newton_solve(
all_residuals,
params,
quat_groups=quat_groups,
max_iter=100,
tol=1e-10,
)
if not converged:
converged = bfgs_solve(
all_residuals,
params,
quat_groups=quat_groups,
max_iter=200,
tol=1e-10,
)
return converged
def _build_constraint(
kind,
body_i,
marker_i_pos,
body_j,
marker_j_pos,
marker_i,
marker_j,
c_params,
) -> ConstraintBase | None:
"""Create the appropriate constraint object from a BaseJointKind."""
marker_i_quat = tuple(marker_i.quaternion)
marker_j_quat = tuple(marker_j.quaternion)
# -- Phase 1 constraints --------------------------------------------------
if kind == kcsolve.BaseJointKind.Coincident:
return CoincidentConstraint(body_i, marker_i_pos, body_j, marker_j_pos)
if kind == kcsolve.BaseJointKind.DistancePointPoint:
distance = c_params[0] if c_params else 0.0
return DistancePointPointConstraint(
body_i,
marker_i_pos,
body_j,
marker_j_pos,
distance,
)
if kind == kcsolve.BaseJointKind.Fixed:
return FixedConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
# -- Phase 2: point constraints -------------------------------------------
if kind == kcsolve.BaseJointKind.PointOnLine:
return PointOnLineConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
if kind == kcsolve.BaseJointKind.PointInPlane:
offset = c_params[0] if c_params else 0.0
return PointInPlaneConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
offset=offset,
)
# -- Phase 2: orientation constraints -------------------------------------
if kind == kcsolve.BaseJointKind.Parallel:
return ParallelConstraint(body_i, marker_i_quat, body_j, marker_j_quat)
if kind == kcsolve.BaseJointKind.Perpendicular:
return PerpendicularConstraint(body_i, marker_i_quat, body_j, marker_j_quat)
if kind == kcsolve.BaseJointKind.Angle:
angle = c_params[0] if c_params else 0.0
return AngleConstraint(body_i, marker_i_quat, body_j, marker_j_quat, angle)
# -- Phase 2: axis/surface constraints ------------------------------------
if kind == kcsolve.BaseJointKind.Concentric:
distance = c_params[0] if c_params else 0.0
return ConcentricConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
distance=distance,
)
if kind == kcsolve.BaseJointKind.Tangent:
return TangentConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
if kind == kcsolve.BaseJointKind.Planar:
offset = c_params[0] if c_params else 0.0
return PlanarConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
offset=offset,
)
if kind == kcsolve.BaseJointKind.LineInPlane:
offset = c_params[0] if c_params else 0.0
return LineInPlaneConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
offset=offset,
)
# -- Phase 2: kinematic joints --------------------------------------------
if kind == kcsolve.BaseJointKind.Ball:
return BallConstraint(body_i, marker_i_pos, body_j, marker_j_pos)
if kind == kcsolve.BaseJointKind.Revolute:
return RevoluteConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
if kind == kcsolve.BaseJointKind.Cylindrical:
return CylindricalConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
if kind == kcsolve.BaseJointKind.Slider:
return SliderConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
if kind == kcsolve.BaseJointKind.Screw:
pitch = c_params[0] if c_params else 1.0
return ScrewConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
pitch=pitch,
)
if kind == kcsolve.BaseJointKind.Universal:
return UniversalConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
)
# -- Phase 2: mechanical constraints --------------------------------------
if kind == kcsolve.BaseJointKind.Gear:
radius_i = c_params[0] if len(c_params) > 0 else 1.0
radius_j = c_params[1] if len(c_params) > 1 else 1.0
return GearConstraint(
body_i,
marker_i_quat,
body_j,
marker_j_quat,
radius_i,
radius_j,
)
if kind == kcsolve.BaseJointKind.RackPinion:
pitch_radius = c_params[0] if c_params else 1.0
return RackPinionConstraint(
body_i,
marker_i_pos,
marker_i_quat,
body_j,
marker_j_pos,
marker_j_quat,
pitch_radius=pitch_radius,
)
# -- Stubs (accepted but produce no residuals) ----------------------------
if kind == kcsolve.BaseJointKind.Cam:
return CamConstraint()
if kind == kcsolve.BaseJointKind.Slot:
return SlotConstraint()
if kind == kcsolve.BaseJointKind.DistanceCylSph:
return DistanceCylSphConstraint()
return None