[SymForce] Consolidate lists of types

There are now three lists, for the cam types, the geo types, and geo
types including internal types, on the geo and cam modules.

Also added some documentation on how to add a new geo/cam type, which
should be easier now.

Also fixes the issue where we didn't include cam types we used in
generated headers.

Topic: sym-type-lists
Reviewers: bradley,nathan,chao,ryan-b,hayk
GitOrigin-RevId: 973b30e4e553944ef2d9fcf2cef7dae6289de030

docs/development.rst

@@ -116,3 +116,23 @@ workflow is currently run manually on a commit, and produces a ``symforce-wheels
 wheels (and sdists) for distribution (e.g. on PyPI).  It doesn't upload them to PyPI - to do that
 (after verifying that the built wheels work as expected) you should download and unzip the archive,
 and upload to PyPI with ``python -m twine upload [--repository testpypi] --verbose *``.
+
+*************************************************
+Adding new types
+*************************************************
+
+To add a new geo or cam type to SymForce:
+
+#. Add a symbolic implementation of your type, to either the :mod:`symforce.geo` or
+   :mod:`symforce.cam` module.  Add an import of your type in the ``__init__.py`` file for the
+   module.
+#. For geo types, you should add it to the ``notebooks/storage_D_tangent.ipynb`` and
+   ``notebooks/tangent_D_storage.ipynb`` notebooks, and use the results there for your symbolic
+   implementation.
+#. Create a test of your symbolic type, for example ``test/geo_rot3_test.py`` or
+   ``test/cam_linear_test.py``.
+#. For geo types, register their numerical equivalents in ``ops/__init__.py``
+#. Add any custom methods you'd like on the runtime numerical classes to the corresponding file in
+   the ``custom_methods`` directory for each backend language
+#. For geo types, add them to the ``"Test implicit construction"`` and ``"Test lie group ops"`` test
+   cases in ``test/symforce_values_test.cc``

symforce/cam/__init__.py

@@ -15,3 +15,6 @@
 from .polynomial_camera_cal import PolynomialCameraCal
 from .posed_camera import PosedCamera
 from .spherical_camera_cal import SphericalCameraCal
+
+# Default generated cam types
+CAM_TYPES = tuple(sorted(CameraCal.__subclasses__(), key=lambda cls: cls.__name__))

symforce/codegen/backends/cpp/templates/function/FUNCTION.h.jinja

@@ -22,9 +22,9 @@
 #include <Eigen/Sparse>
 {% endif %}
 
-{% for typename in ('Rot3', 'Rot2', 'Pose3', 'Pose2', 'Unit3') %}
-{% if typename in spec.types_included %}
-#include <sym/{{ python_util.camelcase_to_snakecase(typename) }}.h>
+{% for cls in sf.GEO_TYPES + sf.CAM_TYPES %}
+{% if cls.__name__ in spec.types_included %}
+#include <sym/{{ python_util.camelcase_to_snakecase(cls.__name__) }}.h>
 {% endif %}
 {% endfor %}
 

symforce/codegen/backends/cpp/templates/type_ops.h.jinja

@@ -11,13 +11,8 @@
 // Import all the known types.
 #include <sym/ops/lie_group_ops.h>
 #include <sym/ops/storage_ops.h>
-#include <sym/pose2.h>
-#include <sym/pose3.h>
-#include <sym/rot2.h>
-#include <sym/rot3.h>
-#include <sym/unit3.h>
-{% for class_name in camera_cal_class_names %}
-#include <sym/{{ python_util.camelcase_to_snakecase(class_name) }}.h>
+{% for cls in sf.GEO_TYPES + sf.CAM_TYPES %}
+#include <sym/{{ python_util.camelcase_to_snakecase(cls.__name__) }}.h>
 {% endfor %}
 #include <sym/util/typedefs.h>
 
@@ -44,16 +39,10 @@ static constexpr const bool kIsSparseEigenType =
   template <typename Scalar, typename... Args>                         \
   auto name(const type_t type, Args&&... args) {                       \
     switch (type.value) {                                              \
-      case type_t::ROT2:                                               \
-        return func<sym::Rot2<Scalar>>(args...);                       \
-      case type_t::ROT3:                                               \
-        return func<sym::Rot3<Scalar>>(args...);                       \
-      case type_t::POSE2:                                              \
-        return func<sym::Pose2<Scalar>>(args...);                      \
-      case type_t::POSE3:                                              \
-        return func<sym::Pose3<Scalar>>(args...);                      \
-      case type_t::UNIT3:                                              \
-        return func<sym::Unit3<Scalar>>(args...);                      \
+      {% for cls in sf.GEO_TYPES %}
+      case type_t::{{ python_util.camelcase_to_screaming_snakecase(cls.__name__) }}: \
+        return func<sym::{{ cls.__name__ }}<Scalar>>(args...);           \
+      {% endfor %}
       case type_t::SCALAR:                                             \
         return func<Scalar>(args...);                                  \
         {% for i in range(1, 10) %}
@@ -66,9 +55,9 @@ static constexpr const bool kIsSparseEigenType =
         return func<Eigen::Matrix<Scalar, {{ i }}, {{ j }}>>(args...); \
         {% endfor %}
         {% endfor %}
-      {% for class_name in camera_cal_class_names %}
-      case type_t::{{ python_util.camelcase_to_screaming_snakecase(class_name) }}: \
-        return func<sym::{{ class_name }}<Scalar>>(args...);         \
+      {% for cls in sf.CAM_TYPES %}
+      case type_t::{{ python_util.camelcase_to_screaming_snakecase(cls.__name__) }}: \
+        return func<sym::{{ cls.__name__ }}<Scalar>>(args...);         \
       {% endfor %}
       default:                                                         \
         SYM_ASSERT(false);                                             \

symforce/codegen/cam_package_codegen.py

@@ -22,18 +22,6 @@
 from symforce.codegen.ops_codegen_util import make_group_ops_funcs
 from symforce.codegen.ops_codegen_util import make_lie_group_ops_funcs
 
-# Default cam types to generate
-DEFAULT_CAM_TYPES = sf.CameraCal.__subclasses__()
-
-
-def camera_cal_class_names() -> T.List[str]:
-    """
-    Returns a sorted list of the CameraCal subclass names.
-    """
-    class_names = [cam_cal_cls.__name__ for cam_cal_cls in sf.CameraCal.__subclasses__()]
-    class_names.sort()
-    return class_names
-
 
 def pixel_from_camera_point_with_jacobians(
     self: sf.CameraCal, point: sf.V3, epsilon: sf.Scalar
@@ -265,7 +253,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
 
         # Build up templates for each type
 
-        for cls in DEFAULT_CAM_TYPES:
+        for cls in sf.CAM_TYPES:
             data = cam_class_data(cls, config=config)
 
             for base_dir, relative_path in (
@@ -292,7 +280,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
             template_path=Path("geo_package", "__init__.py.jinja"),
             data=dict(
                 Codegen.common_data(),
-                all_types=list(geo_package_codegen.DEFAULT_GEO_TYPES) + list(DEFAULT_CAM_TYPES),
+                all_types=list(sf.GEO_TYPES) + list(sf.CAM_TYPES),
                 numeric_epsilon=sf.numeric_epsilon,
             ),
             config=config.render_template_config,
@@ -305,7 +293,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
                 output_path=output_dir / "tests" / name,
                 data=dict(
                     Codegen.common_data(),
-                    all_types=DEFAULT_CAM_TYPES,
+                    all_types=sf.CAM_TYPES,
                     cam_cal_from_points=cam_cal_from_points,
                     _DISTORTION_COEFF_VALS=_DISTORTION_COEFF_VALS,
                 ),
@@ -322,7 +310,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
         geo_package_codegen.generate(config=config, output_dir=output_dir)
 
         # Build up templates for each type
-        for cls in DEFAULT_CAM_TYPES:
+        for cls in sf.CAM_TYPES:
             data = cam_class_data(cls, config=config)
 
             for base_dir, relative_path in (
@@ -376,15 +364,15 @@ def supports_camera_ray_from_pixel(cls: T.Type) -> bool:
                 output_path=output_dir / "tests" / name,
                 data=dict(
                     Codegen.common_data(),
-                    all_types=DEFAULT_CAM_TYPES,
+                    all_types=sf.CAM_TYPES,
                     cpp_cam_types=[
                         f"sym::{cls.__name__}<{scalar}>"
-                        for cls in DEFAULT_CAM_TYPES
+                        for cls in sf.CAM_TYPES
                         for scalar in data["scalar_types"]
                     ],
                     fully_implemented_cpp_cam_types=[
                         f"sym::{cls.__name__}<{scalar}>"
-                        for cls in DEFAULT_CAM_TYPES
+                        for cls in sf.CAM_TYPES
                         for scalar in data["scalar_types"]
                         if supports_camera_ray_from_pixel(cls)
                     ],

symforce/codegen/codegen.py

@@ -360,6 +360,9 @@ def common_data() -> T.Dict[str, T.Any]:
         data["typing_util"] = typing_util
         data["lcm_type_t_include_dir"] = "<lcmtypes/sym/type_t.hpp>"
 
+        # TODO(aaron): Replace uses of members of sf above
+        data["sf"] = sf
+
         def is_symbolic(T: T.Any) -> bool:
             return isinstance(T, (sf.Expr, sf.Symbol))
 

symforce/codegen/geo_package_codegen.py

@@ -23,9 +23,6 @@
 from symforce.codegen.ops_codegen_util import make_group_ops_funcs
 from symforce.codegen.ops_codegen_util import make_lie_group_ops_funcs
 
-# Default geo types to generate
-DEFAULT_GEO_TYPES = (sf.Rot2, sf.Pose2, sf.Rot3, sf.Pose3, sf.Unit3)
-
 
 def geo_class_common_data(cls: T.Type, config: CodegenConfig) -> T.Dict[str, T.Any]:
     """
@@ -205,7 +202,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
 
         # Build up templates for each type
 
-        for cls in DEFAULT_GEO_TYPES:
+        for cls in sf.GEO_TYPES:
             data = geo_class_common_data(cls, config)
             data["matrix_type_aliases"] = matrix_type_aliases.get(cls, {})
             data["custom_generated_methods"] = custom_generated_methods.get(cls, {})
@@ -238,7 +235,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
             template_path=Path("geo_package", "__init__.py.jinja"),
             data=dict(
                 Codegen.common_data(),
-                all_types=DEFAULT_GEO_TYPES,
+                all_types=sf.GEO_TYPES,
                 numeric_epsilon=sf.numeric_epsilon,
             ),
             config=config.render_template_config,
@@ -249,7 +246,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
         for name in ("geo_package_python_test.py",):
             templates.add(
                 template_path=Path("tests", name + ".jinja"),
-                data=dict(Codegen.common_data(), all_types=DEFAULT_GEO_TYPES),
+                data=dict(Codegen.common_data(), all_types=sf.GEO_TYPES),
                 config=config.render_template_config,
                 output_path=output_dir / "tests" / name,
             )
@@ -263,7 +260,7 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
         logger.debug(f'Creating C++ package at: "{package_dir}"')
 
         # Build up templates for each type
-        for cls in DEFAULT_GEO_TYPES:
+        for cls in sf.GEO_TYPES:
             data = geo_class_common_data(cls, config)
             data["matrix_type_aliases"] = matrix_type_aliases.get(cls, {})
             data["custom_generated_methods"] = custom_generated_methods.get(cls, {})
@@ -308,10 +305,10 @@ def generate(config: CodegenConfig, output_dir: Path = None) -> Path:
                 output_path=output_dir / "tests" / name,
                 data=dict(
                     Codegen.common_data(),
-                    all_types=DEFAULT_GEO_TYPES,
+                    all_types=sf.GEO_TYPES,
                     cpp_geo_types=[
                         f"sym::{cls.__name__}<{scalar}>"
-                        for cls in DEFAULT_GEO_TYPES
+                        for cls in sf.GEO_TYPES
                         for scalar in data["scalar_types"]
                     ],
                     cpp_matrix_types=[

symforce/codegen/lcm_templates/symforce_types.lcm.jinja

@@ -14,11 +14,15 @@ enum type_t : int32_t {
   SCALAR = 1,
 
   // Geometry types
-  ROT2 = 2,
-  ROT3 = 3,
-  POSE2 = 4,
-  POSE3 = 5,
-  UNIT3 = 7,
+  {% set k = namespace(value=2) %}
+  {% for cls in sf.GEO_TYPES %}
+  {{ python_util.camelcase_to_screaming_snakecase(cls.__name__) }} = {{ k.value }},
+  {% set k.value = k.value + 1 %}
+  {% if k.value == 6 %}
+    {# Skip 6, used for DATABUFFER #}
+    {% set k.value = k.value + 1 %}
+  {% endif %}
+  {% endfor %}
 
   // DataBuffer
   DATABUFFER = 6,
@@ -40,8 +44,8 @@ enum type_t : int32_t {
   {% endfor %}
 
   // Camera calibrations
-  {% for name in camera_cal_class_names %}
-  {{ python_util.camelcase_to_screaming_snakecase(name) }} = {{ k.value }},
+  {% for cls in sf.CAM_TYPES %}
+  {{ python_util.camelcase_to_screaming_snakecase(cls.__name__) }} = {{ k.value }},
   {% set k.value = k.value + 1 %}
   {% endfor %}
 };

symforce/codegen/lcm_types_codegen.py

@@ -3,15 +3,13 @@
 # This source code is under the Apache 2.0 license found in the LICENSE file.
 # ----------------------------------------------------------------------------
 
+import symforce.symbolic as sf
 from symforce import python_util
 from symforce import typing as T
-from symforce.codegen import cam_package_codegen
 
 
 def lcm_symforce_types_data() -> T.Dict[str, T.Any]:
     """
     Returns data for template generation with ``lcm_templates/symforce_types.lcm.jinja``.
     """
-    return dict(
-        python_util=python_util, camera_cal_class_names=cam_package_codegen.camera_cal_class_names()
-    )
+    return dict(python_util=python_util, sf=sf)

symforce/codegen/sym_util_package_codegen.py

@@ -6,10 +6,10 @@
 import tempfile
 from pathlib import Path
 
+import symforce.symbolic as sf
 from symforce import codegen
 from symforce import logger
 from symforce import python_util
-from symforce.codegen import cam_package_codegen
 from symforce.codegen import template_util
 
 
@@ -43,7 +43,7 @@ def generate(config: codegen.CodegenConfig, output_dir: Path = None) -> Path:
             output_path=package_dir / "type_ops.h",
             data=dict(
                 python_util=python_util,
-                camera_cal_class_names=cam_package_codegen.camera_cal_class_names(),
+                sf=sf,
             ),
             config=config.render_template_config,
             template_dir=template_dir,

symforce/geo/__init__.py

@@ -25,3 +25,20 @@
 from .rot2 import Rot2
 from .rot3 import Rot3
 from .unit3 import Unit3
+
+# Default generated geo types
+GEO_TYPES = (
+    Rot2,
+    Rot3,
+    Pose2,
+    Pose3,
+    Unit3,
+)
+
+# All geo types, including GEO_TYPES plus types used for internal representations that are not
+# generated
+ALL_GEO_TYPES = GEO_TYPES + (
+    Quaternion,
+    DualQuaternion,
+    Complex,
+)