# ----------------------------------------------------------------------------
# SymForce - Copyright 2022, Skydio, Inc.
# This source code is under the Apache 2.0 license found in the LICENSE file.
# ----------------------------------------------------------------------------
import functools
from symforce import codegen
from symforce import typing as T
from symforce.type_helpers import symbolic_inputs
_T = T.TypeVar("_T")
[docs]def symbolic_eval(func: T.Callable[..., _T]) -> _T:
"""
Build symbolic arguments for a function, and return the function evaluated on those arguments.
Useful for easily visualizing what expressions a symbolic function produces
Args:
func: A callable; args should have type annotations, and those types should be constructible
automatically with :func:`symforce.ops.storage_ops.StorageOps.symbolic`
Returns:
The outputs of ``func`` evaluated on the constructed symbolic args
See also:
:func:`symforce.type_helpers.symbolic_inputs`
"""
return func(**symbolic_inputs(func))
[docs]def lambdify(f: T.Callable, use_numba: bool = False) -> T.Callable:
"""
Convert a symbolic function to a numerical one. This is a thin wrapper around
:meth:`Codegen.function <symforce.codegen.codegen.Codegen.function>` provided for convenience.
Args:
f: A callable with symbolic inputs and outputs - see
:meth:`Codegen.function <symforce.codegen.codegen.Codegen.function>` for details
use_numba: If True, use Numba to compile the generated function. This can be much faster,
but has some limitations - see
:class:`codegen.PythonConfig <symforce.codegen.backends.python.python_config.PythonConfig>`
for details
Returns:
A numerical function equivalent to ``f``
See also:
:meth:`Codegen.lambdify <symforce.codegen.codegen.Codegen.lambdify>`
:meth:`Codegen.function <symforce.codegen.codegen.Codegen.function>`
:class:`codegen.PythonConfig <symforce.codegen.backends.python.python_config.PythonConfig>`
"""
codegen_obj = codegen.Codegen.function(f, config=codegen.PythonConfig(use_numba=use_numba))
return codegen_obj.lambdify()
[docs]def numbify(f: T.Callable) -> T.Callable:
"""
Shorthand for ``lambdify(f, use_numba=True)``
See also:
:func:`lambdify`
"""
return lambdify(f, use_numba=True)
SymbolicFunction = T.TypeVar("SymbolicFunction", bound=T.Callable)
[docs]def specialize_types(
f: SymbolicFunction, type_replacements: T.Mapping[T.Type, T.Type]
) -> SymbolicFunction:
"""
Specialize the type annotations on the given function, replacing any types in
``type_replacements``
For example, this can be used to take a symbolic function that accepts a generic type and
generate it for several concrete types::
def f(x: sf.CameraCal) -> sf.Scalar:
...
Codegen.function(specialize_types(f, {sf.CameraCal: sf.LinearCameraCal}), ...)
Codegen.function(specialize_types(f, {sf.CameraCal: sf.PolynomialCameraCal}), ...)
See also:
:func:`specialize_args`
"""
@functools.wraps(f)
def specialized_function(*args: T.Any, **kwargs: T.Any) -> T.Any:
return f(*args, **kwargs)
specialized_function.__annotations__ = f.__annotations__.copy()
for annotation, cls in specialized_function.__annotations__.items():
if cls in type_replacements:
specialized_function.__annotations__[annotation] = type_replacements[cls]
return T.cast(SymbolicFunction, specialized_function)
[docs]def specialize_args(
f: SymbolicFunction, arg_replacements: T.Mapping[str, T.Type]
) -> SymbolicFunction:
"""
Specialize the type annotations on the given function, replacing types for any arguments in
``arg_replacements``
For example, this can be used to take a symbolic function that accepts a generic type and
generate it for several concrete types::
def f(x: sf.CameraCal, y: sf.CameraCal) -> sf.Scalar:
...
Codegen.function(
specialize_types(f, {"x": sf.LinearCameraCal, "y": sf.PolynomialCameraCal}), ...
)
See also:
:func:`specialize_types`
"""
@functools.wraps(f)
def specialized_function(*args: T.Any, **kwargs: T.Any) -> T.Any:
return f(*args, **kwargs)
specialized_function.__annotations__ = f.__annotations__.copy()
for annotation in specialized_function.__annotations__:
if annotation in arg_replacements:
specialized_function.__annotations__[annotation] = arg_replacements[annotation]
return T.cast(SymbolicFunction, specialized_function)