symforce-org / symforce

Fast symbolic computation, code generation, and nonlinear optimization for robotics
https://symforce.org
Apache License 2.0
1.44k stars 147 forks source link

Faster Python linearization functions #230

Open aaron-skydio opened 2 years ago

aaron-skydio commented 2 years ago

Is your feature request related to a problem? Please describe. SymForce can generate Python linearization functions for factors, which is how you can create factors on-the-fly for use with the optimizer from Python. SymForce can also optionally jit generated Python functions with Numba for improved performance.

Numba doesn't seem to like generated functions that take, e.g., geo types:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at /tmp/sf_codegen_residual_factor_k1qmivqr/python/symforce/sym_b5c9a81ad54e4690962994388d01caac/residual_factor.py (44)

File "../../../tmp/sf_codegen_residual_factor_k1qmivqr/python/symforce/sym_b5c9a81ad54e4690962994388d01caac/residual_factor.py", line 44:
<source missing, REPL/exec in use?> 

This error may have been caused by the following argument(s):
- argument 0: Cannot determine Numba type of <class 'sym.pose3.Pose3'>
- argument 1: Cannot determine Numba type of <class 'sym.pose3.Pose3'>
- argument 2: Cannot determine Numba type of <class 'sym.pose3.Pose3'>

It'd be awesome if we could jit "the majority of factors you'd like to create" for greatly improved linearization performance from Python for most problems.

Describe the solution you'd like Create a symbolic residual function that takes a geo or cam type, e.g.

def residual(pose: sf.Pose3, camera: sf.LinearCameraCal) -> sf.V3:
    return something

And generate it to a numba-compiled factor:

factor = Factor(["pose", "camera"], residual, codegen.PythonConfig(autoformat=False, use_numba=True))

Or even use numba here by default (if installed)