[SymForce] Set the default epsilon to a symbol for codegen

aaron-skydio · aaron-skydio · commit 559e521f190e · 2022-09-20T23:52:23.000-07:00
This PR does a couple things:

1) Add checks to `Codegen` for attempting to generate code with the
   default epsilon set to 0.  What to do when this happens is
   configurable in the `CodegenConfig`.  Options are warn (the default),
   raise, or proceed silently.
2) In the test mixin, set the behavior to raise.  This means if you
   create a codegen test that does not set the default epsilon to a
   symbol, you'll get an exception.  You could additionally override
   this by specifying in a particular `CodegenConfig` you create within
   the test `zero_epsilon_behavior=ALLOW`.
3) Set the default epsilon to the symbol for all codegen tests.  We've
   been impressively good at passing epsilons where required, so there
   are minimal changes to the generated code from this.
4) Add helpful error messages, and a `symforce.set_epsilon_to_invalid()`
   method that is helpful for tracking down exactly where an epsilon was
   missing in very large expressions.

Reviewers: dominic,nathan,harrison,hayk,bradley
Topic: sym-default-epsilon-symbol
Relative:
GitOrigin-RevId: 76573d6c36d2940a48ea8dadea942ff68b40a537
diff --git a/gen/cpp/sym/rot3.cc b/gen/cpp/sym/rot3.cc
diff --git a/gen/python/sym/rot3.py b/gen/python/sym/rot3.py
diff --git a/symforce/__init__.py b/symforce/__init__.py
@@ -315,3 +315,18 @@ def set_epsilon_to_zero() -> None:
     See `symforce.symbolic.epsilon` for more information.
     """
     _set_epsilon(0.0)
+
+
+def set_epsilon_to_invalid() -> None:
+    """
+    Set the default epsilon for SymForce to `None`.  Should not be used to actually create
+    expressions or generate code.
+
+    This is useful if you've forgotten to pass an epsilon somewhere, but are not sure where - using
+    this epsilon in an expression should throw a `TypeError` near the location where you forgot to
+    pass an epsilon.
+
+    This must be called before `symforce.symbolic` or other symbolic libraries have been imported.
+    See `symforce.symbolic.epsilon` for more information.
+    """
+    _set_epsilon(None)
diff --git a/symforce/cam/spherical_camera_cal.py b/symforce/cam/spherical_camera_cal.py
@@ -191,7 +191,7 @@ def pixel_from_camera_point(
 
         # compute theta
         xy_norm = point[:2, :].norm(epsilon)
-        theta = sf.atan2(xy_norm, point[2])
+        theta = sf.atan2(xy_norm, point[2], epsilon=0)
         is_valid = sf.Max(sf.sign(self.critical_theta - theta), 0)
 
         # clamp theta to critical_theta
diff --git a/symforce/codegen/backends/cpp/cpp_config.py b/symforce/codegen/backends/cpp/cpp_config.py
@@ -24,6 +24,7 @@ class CppConfig(CodegenConfig):
         use_eigen_types: Use eigen_lcm types for vectors instead of lists
         autoformat: Run a code formatter on the generated code
         cse_optimizations: Optimizations argument to pass to sf.cse
+        zero_epsilon_behavior: What should codegen do if a default epsilon is not set?
         support_complex: Generate code that can work with std::complex or with regular float types
         force_no_inline: Mark generated functions as `__attribute__((noinline))`
         zero_initialization_sparsity_threshold: Threshold between 0 and 1 for the sparsity below
diff --git a/symforce/codegen/backends/python/python_config.py b/symforce/codegen/backends/python/python_config.py
@@ -24,6 +24,7 @@ class PythonConfig(CodegenConfig):
         use_eigen_types: Use eigen_lcm types for vectors instead of lists
         autoformat: Run a code formatter on the generated code
         cse_optimizations: Optimizations argument to pass to sf.cse
+        zero_epsilon_behavior: What should codegen do if a default epsilon is not set?
         use_numba: Add the `@numba.njit` decorator to generated functions.  This will greatly
                    speed up functions by compiling them to machine code, but has large overhead
                    on the first call and some overhead on subsequent calls, so it should not be
diff --git a/symforce/codegen/codegen.py b/symforce/codegen/codegen.py
@@ -95,6 +95,39 @@ def __init__(
             docstring: The docstring to be used with the generated function
         """
 
+        if sf.epsilon() == 0:
+            warning_message = """
+                Generating code with epsilon set to 0 - This is dangerous!  You may get NaNs, Infs,
+                or numerically unstable results from calling generated functions near singularities.
+
+                In order to safely generate code, you should set epsilon to either a symbol
+                (recommended) or a small numerical value like `sf.numeric_epsilon`.  You should do
+                this before importing any other code from symforce, e.g. with
+
+                    import symforce
+                    symforce.set_epsilon_to_symbol()
+
+                or
+
+                    import symforce
+                    symforce.set_epsilon_to_number()
+
+                For more information on use of epsilon to prevent singularities, take a look at the
+                Epsilon Tutorial: https://symforce.org/tutorials/epsilon_tutorial.html
+                """
+            warning_message = textwrap.indent(textwrap.dedent(warning_message), "    ")
+
+            if config.zero_epsilon_behavior == codegen_config.ZeroEpsilonBehavior.FAIL:
+                raise ValueError(warning_message)
+            elif config.zero_epsilon_behavior == codegen_config.ZeroEpsilonBehavior.WARN:
+                logger.warning(warning_message)
+            elif config.zero_epsilon_behavior == codegen_config.ZeroEpsilonBehavior.ALLOW:
+                pass
+            else:
+                raise ValueError(
+                    f"Invalid config.zero_epsilon_behavior: {config.zero_epsilon_behavior}"
+                )
+
         self.name = name
 
         # Inputs and outputs must be Values objects
@@ -111,9 +144,42 @@ def __init__(
         # All symbols in outputs must be present in inputs
         input_symbols_list = codegen_util.flat_symbols_from_values(inputs)
         input_symbols = set(input_symbols_list)
-        assert self.output_symbols.issubset(
-            input_symbols
-        ), f"A symbol in the output expression is missing from inputs. inputs={input_symbols}"
+        if not self.output_symbols.issubset(input_symbols):
+            missing_outputs = self.output_symbols - input_symbols
+            error_msg = textwrap.dedent(
+                f"""
+                A symbol in the output expression is missing from inputs
+
+                Inputs:
+                {input_symbols}
+
+                Missing symbols:
+                {self.output_symbols - input_symbols}
+                """
+            )
+
+            if sf.epsilon() in missing_outputs:
+                error_msg += textwrap.dedent(
+                    f"""
+                    One of the missing symbols is `{sf.epsilon()}`, which is the default epsilon -
+                    this typically means you called a function that requires an epsilon without
+                    passing a value.  You need to either pass 0 for epsilon if you'd like to use 0,
+                    pass through the symbol you're using for epsilon if it's not `{sf.epsilon()}`,
+                    or add `{sf.epsilon()}` as an input to your generated function.  You would do
+                    this either by adding an argument `{sf.epsilon()}: sf.Scalar` if using a
+                    symbolic function, or setting `inputs["{sf.epsilon()}"] = sf.Symbol("{sf.epsilon()}")`
+                    if using `inputs` and `outputs` `Values`.
+
+                    If you aren't sure where you may have forgotten to pass an epsilon, setting
+                    epsilon to invalid may be helpful. You should do this before importing any other
+                    code from symforce, e.g. with
+
+                        import symforce
+                        symforce.set_epsilon_to_invalid()
+                    """
+                )
+
+            raise ValueError(error_msg)
 
         # Names given by keys in inputs/outputs must be valid variable names
         # TODO(aaron): Also check recursively
diff --git a/symforce/codegen/codegen_config.py b/symforce/codegen/codegen_config.py
@@ -4,6 +4,8 @@
 # ----------------------------------------------------------------------------
 from abc import abstractmethod
 from dataclasses import dataclass
+from dataclasses import field
+from enum import Enum
 from pathlib import Path
 
 from sympy.printing.codeprinter import CodePrinter
@@ -13,6 +15,21 @@
 CURRENT_DIR = Path(__file__).parent
 
 
+class ZeroEpsilonBehavior(Enum):
+    """
+    Options for what to do when attempting to generate code with the default epsilon set to 0
+    """
+
+    FAIL = 0
+    WARN = 1
+    ALLOW = 2
+
+
+# Default for new codegen configs - this lets you modify the default for all configs, e.g. for all
+# codegen tests
+DEFAULT_ZERO_EPSILON_BEHAVIOR = ZeroEpsilonBehavior.WARN
+
+
 # TODO(hayk): Address this type ignore, which comes from having abstract methods on a dataclass.
 @dataclass  # type: ignore
 class CodegenConfig:
@@ -26,6 +43,7 @@ class CodegenConfig:
         use_eigen_types: Use eigen_lcm types for vectors instead of lists
         autoformat: Run a code formatter on the generated code
         cse_optimizations: Optimizations argument to pass to sf.cse
+        zero_epsilon_behavior: What should codegen do if a default epsilon is not set?
     """
 
     doc_comment_line_prefix: str
@@ -35,6 +53,9 @@ class CodegenConfig:
     cse_optimizations: T.Optional[
         T.Union[T.Literal["basic"], T.Sequence[T.Tuple[T.Callable, T.Callable]]]
     ] = None
+    zero_epsilon_behavior: ZeroEpsilonBehavior = field(
+        default_factory=lambda: DEFAULT_ZERO_EPSILON_BEHAVIOR
+    )
 
     @classmethod
     @abstractmethod
diff --git a/symforce/codegen/geo_package_codegen.py b/symforce/codegen/geo_package_codegen.py
@@ -99,7 +99,9 @@ def codegen_mul(group: T.Type, multiplicand_type: T.Type) -> Codegen:
             config=config,
         ),
         Codegen.function(
-            func=lambda self: sf.V3(self.to_yaw_pitch_roll()),
+            # TODO(aaron): We currently can't generate custom methods with defaults - fix this, and
+            # pass epsilon as an argument with a default
+            func=lambda self: sf.V3(self.to_yaw_pitch_roll(epsilon=0)),
             input_types=[sf.Rot3],
             name="to_yaw_pitch_roll",
             config=config,
diff --git a/symforce/internal/symbolic.py b/symforce/internal/symbolic.py
@@ -271,7 +271,7 @@ def foo(x: Scalar, epsilon: Scalar = sf.epsilon()) -> Scalar:
     in the SymForce docs here: https://symforce.org/tutorials/epsilon_tutorial.html
 
     For purely numerical code that just needs a good default numerical epsilon, see
-    `symforce.symbolic.numerical_epsilon`.
+    `symforce.symbolic.numeric_epsilon`.
 
     Returns: The current default epsilon.  This is typically some kind of "Scalar", like a float or
              a Symbol.
diff --git a/symforce/test_util/test_case.py b/symforce/test_util/test_case.py
@@ -48,8 +48,11 @@ def main(*args: T.Any, **kwargs: T.Any) -> None:
 
     def setUp(self) -> None:
         super().setUp()
+
+        # Set random seeds
         np.random.seed(42)
         random.seed(42)
+
         # Store verbosity flag so tests can use
         self.verbose = ("-v" in sys.argv) or ("--verbose" in sys.argv)
 
diff --git a/symforce/test_util/test_case_mixin.py b/symforce/test_util/test_case_mixin.py
@@ -17,6 +17,7 @@
 from symforce import logger
 from symforce import python_util
 from symforce import typing as T
+from symforce.codegen import codegen_config
 from symforce.ops import LieGroupOps
 from symforce.ops import StorageOps
 from symforce.ops import interfaces
@@ -147,6 +148,11 @@ def setUp(self) -> None:
         mode is on)
         """
         super().setUp()
+
+        # Set to fail on default epsilon == 0
+        codegen_config.DEFAULT_ZERO_EPSILON_BEHAVIOR = codegen_config.ZeroEpsilonBehavior.FAIL
+
+        # Storage for temporary output directories
         self.output_dirs: T.List[str] = []
 
     def tearDown(self) -> None:
diff --git a/test/symforce_benchmarks_codegen_test.py b/test/symforce_benchmarks_codegen_test.py
@@ -5,6 +5,10 @@
 
 from pathlib import Path
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.benchmarks.inverse_compose_jacobian import generate_inverse_compose_jacobian
 from symforce.benchmarks.matrix_multiplication import generate_matrix_multiplication_benchmark
 from symforce.test_util import TestCase
diff --git a/test/symforce_codegen_test.py b/test/symforce_codegen_test.py
@@ -18,6 +18,9 @@
 from scipy import sparse
 
 import symforce
+
+symforce.set_epsilon_to_symbol()
+
 import symforce.symbolic as sf
 from symforce import codegen
 from symforce import logger
@@ -719,22 +722,20 @@ def test_invalid_codegen_raises(self) -> None:
         y = sf.Symbol("y")
         inputs = Values(input=x)
         outputs = Values(output=x + y)
-        self.assertRaises(
-            AssertionError, codegen.Codegen, "test", inputs, outputs, codegen.CppConfig()
-        )
+        self.assertRaises(ValueError, codegen.Codegen, inputs, outputs, codegen.CppConfig(), "test")
 
         # Inputs have non-unique symbols
         inputs = Values(in_x=x, in_y=x)
         outputs = Values(out_x=x)
         self.assertRaises(
-            AssertionError, codegen.Codegen, "test", inputs, outputs, codegen.CppConfig()
+            AssertionError, codegen.Codegen, inputs, outputs, codegen.CppConfig(), "test"
         )
 
         # Inputs and outputs have non-unique keys
         inputs = Values(x=x)
         outputs = Values(x=x)
         self.assertRaises(
-            AssertionError, codegen.Codegen, "test", inputs, outputs, codegen.CppConfig()
+            AssertionError, codegen.Codegen, inputs, outputs, codegen.CppConfig(), "test"
         )
 
     def test_name_deduction(self) -> None:
diff --git a/test/symforce_cpp_code_printer_codegen_test.py b/test/symforce_cpp_code_printer_codegen_test.py
@@ -6,6 +6,9 @@
 import os
 
 import symforce
+
+symforce.set_epsilon_to_symbol()
+
 import symforce.symbolic as sf
 from symforce import codegen
 from symforce.test_util import TestCase
diff --git a/test/symforce_databuffer_codegen_test.py b/test/symforce_databuffer_codegen_test.py
@@ -7,6 +7,9 @@
 import numpy as np
 
 import symforce
+
+symforce.set_epsilon_to_symbol()
+
 import symforce.symbolic as sf
 from symforce import codegen
 from symforce import typing as T
diff --git a/test/symforce_examples_bundle_adjustment_fixed_size_codegen_test.py b/test/symforce_examples_bundle_adjustment_fixed_size_codegen_test.py
@@ -5,6 +5,10 @@
 
 import os
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.examples.bundle_adjustment_fixed_size.generate_fixed_problem import (
     FixedBundleAdjustmentProblem,
 )
diff --git a/test/symforce_examples_bundle_adjustment_in_the_large_codegen_test.py b/test/symforce_examples_bundle_adjustment_in_the_large_codegen_test.py
@@ -5,6 +5,10 @@
 
 from pathlib import Path
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.examples.bundle_adjustment_in_the_large.bundle_adjustment_in_the_large import generate
 from symforce.test_util import TestCase
 from symforce.test_util import sympy_only
diff --git a/test/symforce_examples_custom_factor_generation_codegen_test.py b/test/symforce_examples_custom_factor_generation_codegen_test.py
@@ -5,6 +5,10 @@
 
 import os
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.examples.custom_factor_generation import generate_factors
 from symforce.test_util import TestCase
 from symforce.test_util import symengine_only
diff --git a/test/symforce_examples_robot_2d_localization_codegen_test.py b/test/symforce_examples_robot_2d_localization_codegen_test.py
@@ -5,6 +5,10 @@
 
 from pathlib import Path
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.examples.robot_2d_localization import robot_2d_localization
 from symforce.test_util import TestCase
 from symforce.test_util import symengine_only
diff --git a/test/symforce_examples_robot_3d_localization_codegen_test.py b/test/symforce_examples_robot_3d_localization_codegen_test.py
@@ -5,6 +5,10 @@
 
 from pathlib import Path
 
+import symforce
+
+symforce.set_epsilon_to_symbol()
+
 from symforce.examples.robot_3d_localization.robot_3d_localization import generate
 from symforce.test_util import TestCase
 from symforce.test_util import symengine_only
diff --git a/test/symforce_gen_codegen_test.py b/test/symforce_gen_codegen_test.py
@@ -10,6 +10,9 @@
 import numpy as np
 
 import symforce
+
+symforce.set_epsilon_to_symbol()
+
 import symforce.symbolic as sf
 from symforce import codegen
 from symforce import ops
diff --git a/test/symforce_gnc_codegen_test.py b/test/symforce_gnc_codegen_test.py
@@ -6,6 +6,9 @@
 import os
 
 import symforce
+
+symforce.set_epsilon_to_symbol()
+
 import symforce.symbolic as sf
 from symforce import codegen
 from symforce import typing as T
diff --git a/test/symforce_lcm_codegen_test.py b/test/symforce_lcm_codegen_test.py
diff --git a/test/symforce_optimization_problem_test.py b/test/symforce_optimization_problem_test.py
diff --git a/test/symforce_py_factor_test.py b/test/symforce_py_factor_test.py
diff --git a/test/symforce_py_optimizer_test.py b/test/symforce_py_optimizer_test.py
diff --git a/test/symforce_similarity_index_test.py b/test/symforce_similarity_index_test.py
diff --git a/test/symforce_types_codegen_test.py b/test/symforce_types_codegen_test.py
diff --git a/test/symforce_values_codegen_test.py b/test/symforce_values_codegen_test.py
