64b1e24467
Add a code generation pipeline that compiles Expr DAGs into flat Python functions, eliminating recursive tree-walk dispatch in the Newton-Raphson inner loop. Key changes: - Add to_code() method to all 11 Expr node types (expr.py) - New codegen.py module with CSE (common subexpression elimination), sparsity detection, and compile()/exec() compilation pipeline - Add ParamTable.env_ref() to avoid dict copies per iteration (params.py) - Newton and BFGS solvers accept pre-built jac_exprs and compiled_eval to avoid redundant diff/simplify and enable compiled evaluation - count_dof() and diagnostics accept pre-built jac_exprs - solver.py builds symbolic Jacobian once, compiles once, passes to all consumers (_monolithic_solve, count_dof, diagnostics) - Automatic fallback: if codegen fails, tree-walk eval is used Expected performance impact: - ~10-20x faster Jacobian evaluation (no recursive dispatch) - ~2-5x additional from CSE on quaternion-heavy systems - ~3x fewer entries evaluated via sparsity detection - Eliminates redundant diff().simplify() in DOF/diagnostics
313 lines
10 KiB
Python
313 lines
10 KiB
Python
"""Per-entity DOF diagnostics and overconstrained detection.
|
|
|
|
Provides per-body remaining degrees of freedom, human-readable free
|
|
motion labels, and redundant/conflicting constraint identification.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
|
|
from .entities import RigidBody
|
|
from .expr import Expr
|
|
from .params import ParamTable
|
|
|
|
# -- Per-entity DOF -----------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class EntityDOF:
|
|
"""DOF report for a single entity (rigid body)."""
|
|
|
|
entity_id: str
|
|
remaining_dof: int # 0 = well-constrained
|
|
free_motions: list[str] = field(default_factory=list)
|
|
|
|
|
|
def per_entity_dof(
|
|
residuals: list[Expr],
|
|
params: ParamTable,
|
|
bodies: dict[str, RigidBody],
|
|
rank_tol: float = 1e-8,
|
|
jac_exprs: "list[list[Expr]] | None" = None,
|
|
) -> list[EntityDOF]:
|
|
"""Compute remaining DOF for each non-grounded body.
|
|
|
|
For each body, extracts the Jacobian columns corresponding to its
|
|
7 parameters, performs SVD to find constrained directions, and
|
|
classifies null-space vectors as translations or rotations.
|
|
"""
|
|
free = params.free_names()
|
|
n_res = len(residuals)
|
|
env = params.get_env()
|
|
|
|
if n_res == 0:
|
|
# No constraints — every free body has 6 DOF
|
|
result = []
|
|
for pid, body in bodies.items():
|
|
if body.grounded:
|
|
continue
|
|
result.append(
|
|
EntityDOF(
|
|
entity_id=pid,
|
|
remaining_dof=6,
|
|
free_motions=[
|
|
"translation along X",
|
|
"translation along Y",
|
|
"translation along Z",
|
|
"rotation about X",
|
|
"rotation about Y",
|
|
"rotation about Z",
|
|
],
|
|
)
|
|
)
|
|
return result
|
|
|
|
# Build column index mapping: param_name -> column index in free list
|
|
free_index = {name: i for i, name in enumerate(free)}
|
|
|
|
# Build full Jacobian (for efficiency, compute once)
|
|
n_free = len(free)
|
|
J_full = np.empty((n_res, n_free))
|
|
if jac_exprs is not None:
|
|
for i in range(n_res):
|
|
for j in range(n_free):
|
|
J_full[i, j] = jac_exprs[i][j].eval(env)
|
|
else:
|
|
for i, r in enumerate(residuals):
|
|
for j, name in enumerate(free):
|
|
J_full[i, j] = r.diff(name).simplify().eval(env)
|
|
|
|
result = []
|
|
for pid, body in bodies.items():
|
|
if body.grounded:
|
|
continue
|
|
|
|
# Find column indices for this body's params
|
|
pfx = pid + "/"
|
|
body_param_names = [
|
|
pfx + "tx",
|
|
pfx + "ty",
|
|
pfx + "tz",
|
|
pfx + "qw",
|
|
pfx + "qx",
|
|
pfx + "qy",
|
|
pfx + "qz",
|
|
]
|
|
col_indices = [free_index[n] for n in body_param_names if n in free_index]
|
|
|
|
if not col_indices:
|
|
# All params fixed (shouldn't happen for non-grounded, but be safe)
|
|
result.append(EntityDOF(entity_id=pid, remaining_dof=0))
|
|
continue
|
|
|
|
# Extract submatrix: all residual rows, only this body's columns
|
|
J_sub = J_full[:, col_indices]
|
|
|
|
# SVD
|
|
U, sv, Vt = np.linalg.svd(J_sub, full_matrices=True)
|
|
constrained = int(np.sum(sv > rank_tol))
|
|
|
|
# Subtract 1 for the quaternion unit-norm constraint (already in residuals)
|
|
# The quat norm residual constrains 1 direction in the 7-D param space,
|
|
# so effective body DOF = 7 - 1 - constrained_by_other_constraints.
|
|
# But the quat norm IS one of the residual rows, so it's already counted
|
|
# in `constrained`. So: remaining = len(col_indices) - constrained
|
|
# But the quat norm takes 1 from 7 → 6 geometric DOF, and constrained
|
|
# includes the quat norm row. So remaining = 7 - constrained, which gives
|
|
# geometric remaining DOF directly.
|
|
remaining = len(col_indices) - constrained
|
|
|
|
# Classify null-space vectors as free motions
|
|
free_motions = []
|
|
if remaining > 0 and Vt.shape[0] > constrained:
|
|
null_space = Vt[constrained:] # rows = null vectors in param space
|
|
|
|
# Map column indices back to param types
|
|
param_types = []
|
|
for n in body_param_names:
|
|
if n in free_index:
|
|
if n.endswith(("/tx", "/ty", "/tz")):
|
|
param_types.append("t")
|
|
else:
|
|
param_types.append("q")
|
|
|
|
for null_vec in null_space:
|
|
label = _classify_motion(
|
|
null_vec, param_types, body_param_names, free_index
|
|
)
|
|
if label:
|
|
free_motions.append(label)
|
|
|
|
result.append(
|
|
EntityDOF(
|
|
entity_id=pid,
|
|
remaining_dof=remaining,
|
|
free_motions=free_motions,
|
|
)
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def _classify_motion(
|
|
null_vec: np.ndarray,
|
|
param_types: list[str],
|
|
body_param_names: list[str],
|
|
free_index: dict[str, int],
|
|
) -> str:
|
|
"""Classify a null-space vector as translation, rotation, or helical."""
|
|
# Split components into translation and rotation parts
|
|
trans_indices = [i for i, t in enumerate(param_types) if t == "t"]
|
|
rot_indices = [i for i, t in enumerate(param_types) if t == "q"]
|
|
|
|
trans_norm = np.linalg.norm(null_vec[trans_indices]) if trans_indices else 0.0
|
|
rot_norm = np.linalg.norm(null_vec[rot_indices]) if rot_indices else 0.0
|
|
|
|
total = trans_norm + rot_norm
|
|
if total < 1e-14:
|
|
return ""
|
|
|
|
trans_frac = trans_norm / total
|
|
rot_frac = rot_norm / total
|
|
|
|
# Determine dominant axis
|
|
if trans_frac > 0.8:
|
|
# Pure translation
|
|
axis = _dominant_axis(null_vec, trans_indices)
|
|
return f"translation along {axis}"
|
|
elif rot_frac > 0.8:
|
|
# Pure rotation
|
|
axis = _dominant_axis(null_vec, rot_indices)
|
|
return f"rotation about {axis}"
|
|
else:
|
|
# Mixed — helical
|
|
axis = _dominant_axis(null_vec, trans_indices)
|
|
return f"helical motion along {axis}"
|
|
|
|
|
|
def _dominant_axis(vec: np.ndarray, indices: list[int]) -> str:
|
|
"""Find the dominant axis (X/Y/Z) among the given component indices."""
|
|
if not indices:
|
|
return "?"
|
|
components = np.abs(vec[indices])
|
|
# Map to axis names — first 3 in group are X/Y/Z
|
|
axis_names = ["X", "Y", "Z"]
|
|
if len(components) >= 3:
|
|
idx = int(np.argmax(components[:3]))
|
|
return axis_names[idx]
|
|
elif len(components) == 1:
|
|
return axis_names[0]
|
|
else:
|
|
idx = int(np.argmax(components))
|
|
return axis_names[min(idx, 2)]
|
|
|
|
|
|
# -- Overconstrained detection ------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class ConstraintDiag:
|
|
"""Diagnostic for a single constraint."""
|
|
|
|
constraint_index: int
|
|
kind: str # "redundant" | "conflicting"
|
|
detail: str
|
|
|
|
|
|
def find_overconstrained(
|
|
residuals: list[Expr],
|
|
params: ParamTable,
|
|
residual_ranges: list[tuple[int, int, int]],
|
|
rank_tol: float = 1e-8,
|
|
jac_exprs: "list[list[Expr]] | None" = None,
|
|
) -> list[ConstraintDiag]:
|
|
"""Identify redundant and conflicting constraints.
|
|
|
|
Algorithm (following SolvSpace's FindWhichToRemoveToFixJacobian):
|
|
1. Build full Jacobian J, compute rank.
|
|
2. If rank == n_residuals, not overconstrained — return empty.
|
|
3. For each constraint: remove its rows, check if rank is preserved
|
|
→ if so, the constraint is **redundant**.
|
|
4. Compute left null space, project residual vector F → if a
|
|
constraint's residuals contribute to this projection, it is
|
|
**conflicting** (redundant + unsatisfied).
|
|
"""
|
|
free = params.free_names()
|
|
n_free = len(free)
|
|
n_res = len(residuals)
|
|
|
|
if n_free == 0 or n_res == 0:
|
|
return []
|
|
|
|
env = params.get_env()
|
|
|
|
# Build Jacobian and residual vector
|
|
J = np.empty((n_res, n_free))
|
|
r_vec = np.empty(n_res)
|
|
for i, r in enumerate(residuals):
|
|
r_vec[i] = r.eval(env)
|
|
if jac_exprs is not None:
|
|
for i in range(n_res):
|
|
for j in range(n_free):
|
|
J[i, j] = jac_exprs[i][j].eval(env)
|
|
else:
|
|
for i, r in enumerate(residuals):
|
|
for j, name in enumerate(free):
|
|
J[i, j] = r.diff(name).simplify().eval(env)
|
|
|
|
# Full rank
|
|
sv_full = np.linalg.svd(J, compute_uv=False)
|
|
full_rank = int(np.sum(sv_full > rank_tol))
|
|
|
|
if full_rank >= n_res:
|
|
return [] # not overconstrained
|
|
|
|
# Left null space: columns of U beyond rank
|
|
U, sv, Vt = np.linalg.svd(J, full_matrices=True)
|
|
left_null = U[:, full_rank:] # shape (n_res, n_res - rank)
|
|
|
|
# Project residual onto left null space
|
|
null_residual = left_null.T @ r_vec # shape (n_res - rank,)
|
|
residual_projection = left_null @ null_residual # back to residual space
|
|
|
|
diags = []
|
|
for start, end, c_idx in residual_ranges:
|
|
# Remove this constraint's rows and check rank
|
|
mask = np.ones(n_res, dtype=bool)
|
|
mask[start:end] = False
|
|
J_reduced = J[mask]
|
|
|
|
if J_reduced.shape[0] == 0:
|
|
continue
|
|
|
|
sv_reduced = np.linalg.svd(J_reduced, compute_uv=False)
|
|
reduced_rank = int(np.sum(sv_reduced > rank_tol))
|
|
|
|
if reduced_rank >= full_rank:
|
|
# Removing this constraint preserves rank → redundant
|
|
# Check if it's also conflicting (contributes to unsatisfied null projection)
|
|
constraint_proj = np.linalg.norm(residual_projection[start:end])
|
|
if constraint_proj > rank_tol:
|
|
kind = "conflicting"
|
|
detail = (
|
|
f"Constraint {c_idx} is conflicting (redundant and unsatisfied)"
|
|
)
|
|
else:
|
|
kind = "redundant"
|
|
detail = (
|
|
f"Constraint {c_idx} is redundant (can be removed without effect)"
|
|
)
|
|
diags.append(
|
|
ConstraintDiag(
|
|
constraint_index=c_idx,
|
|
kind=kind,
|
|
detail=detail,
|
|
)
|
|
)
|
|
|
|
return diags
|