98051ba0c9
- Move existing OndselSolver, GNN ML layer, and tooling into GNN/ directory for integration in later phases - Add Create addon scaffold: package.xml, Init.py - Add expression DAG with eval, symbolic diff, simplification - Add parameter table with fixed/free variable tracking - Add quaternion rotation as polynomial Expr trees - Add RigidBody entity (7 DOF: position + unit quaternion) - Add constraint classes: Coincident, DistancePointPoint, Fixed - Add Newton-Raphson solver with symbolic Jacobian + numpy lstsq - Add pre-solve passes: substitution + single-equation - Add DOF counting via Jacobian SVD rank - Add KindredSolver IKCSolver bridge for kcsolve integration - Add 82 unit tests covering all modules Registers as 'kindred' solver via kcsolve.register_solver() when loaded by Create's addon_loader.
425 lines
10 KiB
Python
425 lines
10 KiB
Python
"""Immutable expression DAG with eval, symbolic differentiation, and simplification."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Dict
|
|
|
|
|
|
class Expr:
|
|
"""Base class for all expression nodes."""
|
|
|
|
__slots__ = ()
|
|
|
|
def eval(self, env: Dict[str, float]) -> float:
|
|
raise NotImplementedError
|
|
|
|
def diff(self, var: str) -> Expr:
|
|
raise NotImplementedError
|
|
|
|
def simplify(self) -> Expr:
|
|
return self
|
|
|
|
def vars(self) -> set[str]:
|
|
"""Return the set of variable names in this expression."""
|
|
raise NotImplementedError
|
|
|
|
# -- operator overloads --------------------------------------------------
|
|
|
|
def __add__(self, other):
|
|
return Add(self, _wrap(other))
|
|
|
|
def __radd__(self, other):
|
|
return Add(_wrap(other), self)
|
|
|
|
def __sub__(self, other):
|
|
return Sub(self, _wrap(other))
|
|
|
|
def __rsub__(self, other):
|
|
return Sub(_wrap(other), self)
|
|
|
|
def __mul__(self, other):
|
|
return Mul(self, _wrap(other))
|
|
|
|
def __rmul__(self, other):
|
|
return Mul(_wrap(other), self)
|
|
|
|
def __truediv__(self, other):
|
|
return Div(self, _wrap(other))
|
|
|
|
def __rtruediv__(self, other):
|
|
return Div(_wrap(other), self)
|
|
|
|
def __neg__(self):
|
|
return Neg(self)
|
|
|
|
def __pow__(self, other):
|
|
return Pow(self, _wrap(other))
|
|
|
|
def __rpow__(self, other):
|
|
return Pow(_wrap(other), self)
|
|
|
|
|
|
def _wrap(x) -> Expr:
|
|
"""Coerce a number to Const."""
|
|
if isinstance(x, Expr):
|
|
return x
|
|
if isinstance(x, (int, float)):
|
|
return Const(float(x))
|
|
raise TypeError(f"Cannot coerce {type(x).__name__} to Expr")
|
|
|
|
|
|
# -- leaf nodes ---------------------------------------------------------------
|
|
|
|
|
|
class Const(Expr):
|
|
__slots__ = ("value",)
|
|
|
|
def __init__(self, value: float):
|
|
self.value = float(value)
|
|
|
|
def eval(self, env):
|
|
return self.value
|
|
|
|
def diff(self, var):
|
|
return ZERO
|
|
|
|
def simplify(self):
|
|
return self
|
|
|
|
def vars(self):
|
|
return set()
|
|
|
|
def __repr__(self):
|
|
return f"Const({self.value})"
|
|
|
|
def __eq__(self, other):
|
|
return isinstance(other, Const) and self.value == other.value
|
|
|
|
def __hash__(self):
|
|
return hash(("Const", self.value))
|
|
|
|
|
|
class Var(Expr):
|
|
__slots__ = ("name",)
|
|
|
|
def __init__(self, name: str):
|
|
self.name = name
|
|
|
|
def eval(self, env):
|
|
return env[self.name]
|
|
|
|
def diff(self, var):
|
|
return ONE if var == self.name else ZERO
|
|
|
|
def simplify(self):
|
|
return self
|
|
|
|
def vars(self):
|
|
return {self.name}
|
|
|
|
def __repr__(self):
|
|
return f"Var({self.name!r})"
|
|
|
|
def __eq__(self, other):
|
|
return isinstance(other, Var) and self.name == other.name
|
|
|
|
def __hash__(self):
|
|
return hash(("Var", self.name))
|
|
|
|
|
|
# -- unary nodes --------------------------------------------------------------
|
|
|
|
|
|
class Neg(Expr):
|
|
__slots__ = ("child",)
|
|
|
|
def __init__(self, child: Expr):
|
|
self.child = child
|
|
|
|
def eval(self, env):
|
|
return -self.child.eval(env)
|
|
|
|
def diff(self, var):
|
|
return Neg(self.child.diff(var))
|
|
|
|
def simplify(self):
|
|
c = self.child.simplify()
|
|
if isinstance(c, Const):
|
|
return Const(-c.value)
|
|
if isinstance(c, Neg):
|
|
return c.child
|
|
return Neg(c)
|
|
|
|
def vars(self):
|
|
return self.child.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Neg({self.child!r})"
|
|
|
|
|
|
class Sin(Expr):
|
|
__slots__ = ("child",)
|
|
|
|
def __init__(self, child: Expr):
|
|
self.child = child
|
|
|
|
def eval(self, env):
|
|
return math.sin(self.child.eval(env))
|
|
|
|
def diff(self, var):
|
|
# d/dx sin(f) = cos(f) * f'
|
|
return Mul(Cos(self.child), self.child.diff(var))
|
|
|
|
def simplify(self):
|
|
c = self.child.simplify()
|
|
if isinstance(c, Const):
|
|
return Const(math.sin(c.value))
|
|
return Sin(c)
|
|
|
|
def vars(self):
|
|
return self.child.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Sin({self.child!r})"
|
|
|
|
|
|
class Cos(Expr):
|
|
__slots__ = ("child",)
|
|
|
|
def __init__(self, child: Expr):
|
|
self.child = child
|
|
|
|
def eval(self, env):
|
|
return math.cos(self.child.eval(env))
|
|
|
|
def diff(self, var):
|
|
# d/dx cos(f) = -sin(f) * f'
|
|
return Mul(Neg(Sin(self.child)), self.child.diff(var))
|
|
|
|
def simplify(self):
|
|
c = self.child.simplify()
|
|
if isinstance(c, Const):
|
|
return Const(math.cos(c.value))
|
|
return Cos(c)
|
|
|
|
def vars(self):
|
|
return self.child.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Cos({self.child!r})"
|
|
|
|
|
|
class Sqrt(Expr):
|
|
__slots__ = ("child",)
|
|
|
|
def __init__(self, child: Expr):
|
|
self.child = child
|
|
|
|
def eval(self, env):
|
|
return math.sqrt(self.child.eval(env))
|
|
|
|
def diff(self, var):
|
|
# d/dx sqrt(f) = f' / (2 * sqrt(f))
|
|
return Div(self.child.diff(var), Mul(Const(2.0), Sqrt(self.child)))
|
|
|
|
def simplify(self):
|
|
c = self.child.simplify()
|
|
if isinstance(c, Const) and c.value >= 0:
|
|
return Const(math.sqrt(c.value))
|
|
return Sqrt(c)
|
|
|
|
def vars(self):
|
|
return self.child.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Sqrt({self.child!r})"
|
|
|
|
|
|
# -- binary nodes -------------------------------------------------------------
|
|
|
|
|
|
class Add(Expr):
|
|
__slots__ = ("a", "b")
|
|
|
|
def __init__(self, a: Expr, b: Expr):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def eval(self, env):
|
|
return self.a.eval(env) + self.b.eval(env)
|
|
|
|
def diff(self, var):
|
|
return Add(self.a.diff(var), self.b.diff(var))
|
|
|
|
def simplify(self):
|
|
a = self.a.simplify()
|
|
b = self.b.simplify()
|
|
if isinstance(a, Const) and isinstance(b, Const):
|
|
return Const(a.value + b.value)
|
|
if isinstance(a, Const) and a.value == 0.0:
|
|
return b
|
|
if isinstance(b, Const) and b.value == 0.0:
|
|
return a
|
|
return Add(a, b)
|
|
|
|
def vars(self):
|
|
return self.a.vars() | self.b.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Add({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class Sub(Expr):
|
|
__slots__ = ("a", "b")
|
|
|
|
def __init__(self, a: Expr, b: Expr):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def eval(self, env):
|
|
return self.a.eval(env) - self.b.eval(env)
|
|
|
|
def diff(self, var):
|
|
return Sub(self.a.diff(var), self.b.diff(var))
|
|
|
|
def simplify(self):
|
|
a = self.a.simplify()
|
|
b = self.b.simplify()
|
|
if isinstance(a, Const) and isinstance(b, Const):
|
|
return Const(a.value - b.value)
|
|
if isinstance(b, Const) and b.value == 0.0:
|
|
return a
|
|
if isinstance(a, Const) and a.value == 0.0:
|
|
return Neg(b).simplify()
|
|
return Sub(a, b)
|
|
|
|
def vars(self):
|
|
return self.a.vars() | self.b.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Sub({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class Mul(Expr):
|
|
__slots__ = ("a", "b")
|
|
|
|
def __init__(self, a: Expr, b: Expr):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def eval(self, env):
|
|
return self.a.eval(env) * self.b.eval(env)
|
|
|
|
def diff(self, var):
|
|
# product rule: a'b + ab'
|
|
return Add(Mul(self.a.diff(var), self.b), Mul(self.a, self.b.diff(var)))
|
|
|
|
def simplify(self):
|
|
a = self.a.simplify()
|
|
b = self.b.simplify()
|
|
if isinstance(a, Const) and isinstance(b, Const):
|
|
return Const(a.value * b.value)
|
|
if isinstance(a, Const) and a.value == 0.0:
|
|
return ZERO
|
|
if isinstance(b, Const) and b.value == 0.0:
|
|
return ZERO
|
|
if isinstance(a, Const) and a.value == 1.0:
|
|
return b
|
|
if isinstance(b, Const) and b.value == 1.0:
|
|
return a
|
|
if isinstance(a, Const) and a.value == -1.0:
|
|
return Neg(b).simplify()
|
|
if isinstance(b, Const) and b.value == -1.0:
|
|
return Neg(a).simplify()
|
|
return Mul(a, b)
|
|
|
|
def vars(self):
|
|
return self.a.vars() | self.b.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Mul({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class Div(Expr):
|
|
__slots__ = ("a", "b")
|
|
|
|
def __init__(self, a: Expr, b: Expr):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def eval(self, env):
|
|
return self.a.eval(env) / self.b.eval(env)
|
|
|
|
def diff(self, var):
|
|
# quotient rule: (a'b - ab') / b^2
|
|
return Div(
|
|
Sub(Mul(self.a.diff(var), self.b), Mul(self.a, self.b.diff(var))),
|
|
Mul(self.b, self.b),
|
|
)
|
|
|
|
def simplify(self):
|
|
a = self.a.simplify()
|
|
b = self.b.simplify()
|
|
if isinstance(a, Const) and isinstance(b, Const) and b.value != 0.0:
|
|
return Const(a.value / b.value)
|
|
if isinstance(a, Const) and a.value == 0.0:
|
|
return ZERO
|
|
if isinstance(b, Const) and b.value == 1.0:
|
|
return a
|
|
return Div(a, b)
|
|
|
|
def vars(self):
|
|
return self.a.vars() | self.b.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Div({self.a!r}, {self.b!r})"
|
|
|
|
|
|
class Pow(Expr):
|
|
__slots__ = ("base", "exp")
|
|
|
|
def __init__(self, base: Expr, exp: Expr):
|
|
self.base = base
|
|
self.exp = exp
|
|
|
|
def eval(self, env):
|
|
return self.base.eval(env) ** self.exp.eval(env)
|
|
|
|
def diff(self, var):
|
|
# For constant exponent: d/dx f^n = n * f^(n-1) * f'
|
|
# General case: d/dx f^g = f^g * (g' * ln(f) + g * f'/f)
|
|
# Phase 1: only support constant exponent
|
|
if isinstance(self.exp, Const):
|
|
n = self.exp.value
|
|
return Mul(
|
|
Mul(Const(n), Pow(self.base, Const(n - 1.0))), self.base.diff(var)
|
|
)
|
|
raise NotImplementedError("diff of Pow with non-constant exponent")
|
|
|
|
def simplify(self):
|
|
base = self.base.simplify()
|
|
exp = self.exp.simplify()
|
|
if isinstance(base, Const) and isinstance(exp, Const):
|
|
return Const(base.value**exp.value)
|
|
if isinstance(exp, Const):
|
|
if exp.value == 0.0:
|
|
return ONE
|
|
if exp.value == 1.0:
|
|
return base
|
|
if exp.value == 2.0:
|
|
return Mul(base, base).simplify()
|
|
return Pow(base, exp)
|
|
|
|
def vars(self):
|
|
return self.base.vars() | self.exp.vars()
|
|
|
|
def __repr__(self):
|
|
return f"Pow({self.base!r}, {self.exp!r})"
|
|
|
|
|
|
# -- sentinels ----------------------------------------------------------------
|
|
|
|
ZERO = Const(0.0)
|
|
ONE = Const(1.0)
|