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
55 lines
1.3 KiB
Python
55 lines
1.3 KiB
Python
"""Degrees-of-freedom counting via Jacobian rank."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
|
|
from .expr import Expr
|
|
from .params import ParamTable
|
|
|
|
|
|
def count_dof(
|
|
residuals: List[Expr],
|
|
params: ParamTable,
|
|
rank_tol: float = 1e-8,
|
|
jac_exprs: "List[List[Expr]] | None" = None,
|
|
) -> int:
|
|
"""Compute DOF = n_free_params - rank(Jacobian).
|
|
|
|
Evaluates the Jacobian numerically at the current parameter values
|
|
and computes its rank via SVD.
|
|
|
|
When *jac_exprs* is provided, reuses the pre-built symbolic
|
|
Jacobian instead of re-differentiating every residual.
|
|
"""
|
|
free = params.free_names()
|
|
n_free = len(free)
|
|
n_res = len(residuals)
|
|
|
|
if n_free == 0:
|
|
return 0
|
|
if n_res == 0:
|
|
return n_free
|
|
|
|
env = params.get_env()
|
|
|
|
J = 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[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)
|
|
|
|
if J.size == 0:
|
|
return n_free
|
|
|
|
sv = np.linalg.svd(J, compute_uv=False)
|
|
rank = int(np.sum(sv > rank_tol))
|
|
|
|
return n_free - rank
|