533ca91774
Add 18 new constraint classes covering all BaseJointKind types from Types.h: - Point: PointOnLine (2r), PointInPlane (1r) - Orientation: Parallel (2r), Perpendicular (1r), Angle (1r) - Surface: Concentric (4r), Tangent (1r), Planar (3r), LineInPlane (2r) - Kinematic: Ball (3r), Revolute (5r), Cylindrical (4r), Slider (5r), Screw (5r), Universal (4r) - Mechanical: Gear (1r), RackPinion (1r) - Stubs: Cam, Slot, DistanceCylSph New modules: - geometry.py: marker axis extraction, vector ops (dot3, cross3, sub3), geometric primitives (point_plane_distance, point_line_perp_components) - bfgs.py: L-BFGS-B fallback solver via scipy for when Newton fails solver.py changes: - Wire all 20 supported types in _build_constraint() - BFGS fallback after Newton-Raphson in solve() 183 tests passing (up from 82), including: - DOF counting for every joint type - Solve convergence from displaced initial conditions - Multi-body mechanisms (four-bar linkage, slider-crank, revolute chain)
128 lines
3.4 KiB
Python
128 lines
3.4 KiB
Python
"""L-BFGS-B fallback solver for when Newton-Raphson fails to converge.
|
|
|
|
Minimizes f(x) = 0.5 * sum(r_i(x)^2) using scipy's L-BFGS-B with
|
|
analytic gradient from the Expr DAG's symbolic differentiation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
|
|
from .expr import Expr
|
|
from .params import ParamTable
|
|
|
|
try:
|
|
from scipy.optimize import minimize as _scipy_minimize
|
|
|
|
_HAS_SCIPY = True
|
|
except ImportError:
|
|
_HAS_SCIPY = False
|
|
|
|
|
|
def bfgs_solve(
|
|
residuals: List[Expr],
|
|
params: ParamTable,
|
|
quat_groups: List[tuple[str, str, str, str]] | None = None,
|
|
max_iter: int = 200,
|
|
tol: float = 1e-10,
|
|
) -> bool:
|
|
"""Solve ``residuals == 0`` by minimizing sum of squared residuals.
|
|
|
|
Falls back gracefully to False if scipy is not available.
|
|
|
|
Returns True if converged (||r|| < tol).
|
|
"""
|
|
if not _HAS_SCIPY:
|
|
return False
|
|
|
|
free = params.free_names()
|
|
n_free = len(free)
|
|
n_res = len(residuals)
|
|
|
|
if n_free == 0 or n_res == 0:
|
|
return True
|
|
|
|
# Build symbolic gradient expressions once: d(r_i)/d(x_j)
|
|
jac_exprs: List[List[Expr]] = []
|
|
for r in residuals:
|
|
row = []
|
|
for name in free:
|
|
row.append(r.diff(name).simplify())
|
|
jac_exprs.append(row)
|
|
|
|
def objective_and_grad(x_vec):
|
|
# Update params
|
|
params.set_free_vector(x_vec)
|
|
if quat_groups:
|
|
_renormalize_quats(params, quat_groups)
|
|
|
|
env = params.get_env()
|
|
|
|
# Evaluate residuals
|
|
r_vals = np.array([r.eval(env) for r in residuals])
|
|
f = 0.5 * np.dot(r_vals, r_vals)
|
|
|
|
# Evaluate Jacobian
|
|
J = np.empty((n_res, n_free))
|
|
for i in range(n_res):
|
|
for j in range(n_free):
|
|
J[i, j] = jac_exprs[i][j].eval(env)
|
|
|
|
# Gradient of f = sum(r_i * dr_i/dx_j) = J^T @ r
|
|
grad = J.T @ r_vals
|
|
|
|
return f, grad
|
|
|
|
x0 = params.get_free_vector().copy()
|
|
|
|
result = _scipy_minimize(
|
|
objective_and_grad,
|
|
x0,
|
|
method="L-BFGS-B",
|
|
jac=True,
|
|
options={"maxiter": max_iter, "ftol": tol * tol, "gtol": tol},
|
|
)
|
|
|
|
# Apply final result
|
|
params.set_free_vector(result.x)
|
|
if quat_groups:
|
|
_renormalize_quats(params, quat_groups)
|
|
|
|
# Check convergence on actual residual norm
|
|
env = params.get_env()
|
|
r_vals = np.array([r.eval(env) for r in residuals])
|
|
return bool(np.linalg.norm(r_vals) < tol)
|
|
|
|
|
|
def _renormalize_quats(
|
|
params: ParamTable,
|
|
groups: List[tuple[str, str, str, str]],
|
|
):
|
|
"""Project quaternion params back onto the unit sphere."""
|
|
for qw_name, qx_name, qy_name, qz_name in groups:
|
|
if (
|
|
params.is_fixed(qw_name)
|
|
and params.is_fixed(qx_name)
|
|
and params.is_fixed(qy_name)
|
|
and params.is_fixed(qz_name)
|
|
):
|
|
continue
|
|
w = params.get_value(qw_name)
|
|
x = params.get_value(qx_name)
|
|
y = params.get_value(qy_name)
|
|
z = params.get_value(qz_name)
|
|
norm = math.sqrt(w * w + x * x + y * y + z * z)
|
|
if norm < 1e-15:
|
|
params.set_value(qw_name, 1.0)
|
|
params.set_value(qx_name, 0.0)
|
|
params.set_value(qy_name, 0.0)
|
|
params.set_value(qz_name, 0.0)
|
|
else:
|
|
params.set_value(qw_name, w / norm)
|
|
params.set_value(qx_name, x / norm)
|
|
params.set_value(qy_name, y / norm)
|
|
params.set_value(qz_name, z / norm)
|