Merge remote-tracking branch 'origin/main' into centralize-backend-specific-code

Author: hayk-skydio

symforce/codegen/codegen.py

@@ -421,8 +421,8 @@ def generate_function(
             output_dir=output_dir,
             lcm_type_dir=Path(types_codegen_data["lcm_type_dir"]),
             function_dir=out_function_dir,
-            python_types_dir=Path(lcm_data["python_types_dir"]),
-            cpp_types_dir=Path(lcm_data["cpp_types_dir"]),
+            python_types_dir=lcm_data["python_types_dir"],
+            cpp_types_dir=lcm_data["cpp_types_dir"],
             generated_files=[Path(v.output_path) for v in templates.items],
         )
 

symforce/codegen/codegen_util.py

@@ -13,7 +13,6 @@
 import importlib.abc
 import importlib.util
 import itertools
-import os
 from pathlib import Path
 import sympy
 import sys
@@ -610,23 +609,28 @@ def get_base_instance(obj: T.Sequence[T.Any]) -> T.Any:
 
 def generate_lcm_types(
     lcm_type_dir: T.Openable, lcm_files: T.Sequence[str], lcm_output_dir: T.Openable = None
-) -> T.Dict[str, str]:
+) -> T.Dict[str, Path]:
     """
     Generates the language-specific type files for all symforce generated ".lcm" files.
 
     Args:
         lcm_type_dir: Directory containing symforce-generated .lcm files
         lcm_files: List of .lcm files to process
     """
+    lcm_type_dir = Path(lcm_type_dir)
+
     if lcm_output_dir is None:
-        lcm_output_dir = os.path.join(lcm_type_dir, "..")
+        lcm_output_dir = lcm_type_dir / ".."
+    else:
+        lcm_output_dir = Path(lcm_output_dir)
 
-    python_types_dir = os.path.join(lcm_output_dir, "python")
-    cpp_types_dir = os.path.join(lcm_output_dir, "cpp", "lcmtypes")
-    lcm_include_dir = os.path.join("lcmtypes")
+    python_types_dir = lcm_output_dir / "python"
+    cpp_types_dir = lcm_output_dir / "cpp" / "lcmtypes"
+    lcm_include_dir = "lcmtypes"
 
     result = {"python_types_dir": python_types_dir, "cpp_types_dir": cpp_types_dir}
 
+    # TODO(brad, aaron): Do something reasonable with lcm_files other than returning early
     # If no LCM files provided, do nothing
     if not lcm_files:
         return result
@@ -638,23 +642,23 @@ def generate_lcm_types(
     skymarshal.main(
         [SkymarshalPython, SkymarshalCpp],
         args=[
-            lcm_type_dir,
+            str(lcm_type_dir),
             "--python",
             "--python-path",
-            os.path.join(python_types_dir, "lcmtypes"),
+            str(python_types_dir / "lcmtypes"),
             "--python-namespace-packages",
             "--python-package-prefix",
             "lcmtypes",
             "--cpp",
             "--cpp-hpath",
-            cpp_types_dir,
+            str(cpp_types_dir),
             "--cpp-include",
             lcm_include_dir,
         ],
     )
 
     # Autoformat generated python files
-    format_util.format_py_dir(python_types_dir)
+    format_util.format_py_dir(str(python_types_dir))
 
     return result
 

symforce/codegen/types_package_codegen.py

@@ -160,7 +160,7 @@ def generate_types(
     if not using_external_templates:
         templates.render()
         lcm_data = codegen_util.generate_lcm_types(lcm_type_dir, lcm_files, lcm_bindings_output_dir)
-        codegen_data.update(lcm_data)
+        codegen_data.update({key: str(val) for key, val in lcm_data.items()})
 
     return codegen_data
 

symforce/examples/robot_3d_localization/plotting.py

@@ -101,7 +101,7 @@ def plot_solution(
     if show_iteration_text:
         text = ax.text(8, 7, 9, "-", color="black")
 
-    def update_plot(slider_value: np.float64, show_iteration_text: bool) -> None:
+    def update_plot(slider_value: np.float64) -> None:
         """
         Update the plot using the given iteration.
         """

symforce/opt/internal/factor_utils.h

@@ -174,10 +174,13 @@ Factor<Scalar> JacobianFixed(Functor func, const std::vector<Key>& keys_to_func,
   constexpr int M = JacobianMat::RowsAtCompileTime;
   constexpr int N = JacobianMat::ColsAtCompileTime;
 
+  // NOTE(harrison): accoring to c++ spec you shouldn't have to specify a default lambda capture
+  // and should only have to specify a capture for func. However, this is not handled correctly in
+  // older versions of clang, which is why we have one here.
   return Factor<Scalar>(
-      [func](const Values<Scalar>& values, const std::vector<index_entry_t>& keys_to_func,
-             VectorX<Scalar>* residual, MatrixX<Scalar>* jacobian, MatrixX<Scalar>* hessian,
-             VectorX<Scalar>* rhs) {
+      [=](const Values<Scalar>& values, const std::vector<index_entry_t>& keys_to_func,
+          VectorX<Scalar>* residual, MatrixX<Scalar>* jacobian, MatrixX<Scalar>* hessian,
+          VectorX<Scalar>* rhs) {
         SYM_ASSERT(residual != nullptr);
         Eigen::Matrix<Scalar, M, 1> residual_fixed;
 

symforce/pybind/cc_factor.cc

@@ -5,6 +5,7 @@
 
 #include "./cc_factor.h"
 
+#include <cstring>
 #include <functional>
 
 #include <Eigen/Dense>
@@ -35,67 +36,89 @@ namespace {
 using PyHessianFunc =
     std::function<py::tuple(const sym::Valuesd&, const std::vector<index_entry_t>&)>;
 
+/**
+ * If Matrix is Eigen::SparseMatrix<double> and matrix is not a scipy.sparse.csc_matrix, or
+ * if Matrix is any other type and matrix is a scipy.sparse.csc_matrix, throws a value error.
+ */
+template <typename Matrix>
+void ThrowIfSparsityMismatch(const py::object& matrix) {
+  if (std::strcmp(Py_TYPE(matrix.ptr())->tp_name, "csc_matrix") == 0) {
+    throw py::value_error("Non-sparse matrix expected, scipy.sparse.csc_matrix found instead.");
+  }
+}
+template <>
+void ThrowIfSparsityMismatch<Eigen::SparseMatrix<double>>(const py::object& matrix) {
+  if (!py::isinstance(matrix, py::module_::import("scipy.sparse").attr("csc_matrix"))) {
+    throw py::value_error(
+        fmt::format("scipy.sparse.csc_matrix expected, found {} instead.", py::type::of(matrix)));
+  }
+}
+
+template <typename Matrix>
 auto WrapPyHessianFunc(PyHessianFunc&& hessian_func) {
   using Vec = Eigen::VectorXd;
-  using Mat = Eigen::MatrixXd;
   return [hessian_func = std::move(hessian_func)](
              const sym::Valuesd& values, const std::vector<index_entry_t>& keys,
-             Vec* const residual, Mat* const jacobian, Mat* const hessian, Vec* const rhs) {
+             Vec* const residual, Matrix* const jacobian, Matrix* const hessian, Vec* const rhs) {
     const py::tuple out_tuple = hessian_func(values, keys);
     if (residual != nullptr) {
       *residual = py::cast<Vec>(out_tuple[0]);
     }
     if (jacobian != nullptr) {
-      *jacobian = py::cast<Mat>(out_tuple[1]);
+      ThrowIfSparsityMismatch<Matrix>(out_tuple[1]);
+      *jacobian = py::cast<Matrix>(out_tuple[1]);
     }
     if (hessian != nullptr) {
-      *hessian = py::cast<Mat>(out_tuple[2]);
+      ThrowIfSparsityMismatch<Matrix>(out_tuple[2]);
+      *hessian = py::cast<Matrix>(out_tuple[2]);
     }
     if (rhs != nullptr) {
       *rhs = py::cast<Vec>(out_tuple[3]);
     }
   };
 }
 
-sym::Factord MakeHessianFactorSeparateKeys(PyHessianFunc hessian_func,
-                                           const std::vector<sym::Key>& keys_to_func,
-                                           const std::vector<sym::Key>& keys_to_optimize) {
-  return sym::Factord(WrapPyHessianFunc(std::move(hessian_func)), keys_to_func, keys_to_optimize);
-}
-
-sym::Factord MakeHessianFactorCommonKeys(PyHessianFunc hessian_func,
-                                         const std::vector<sym::Key>& keys) {
-  return sym::Factord(WrapPyHessianFunc(std::move(hessian_func)), keys);
+template <typename... Keys>
+sym::Factord MakeHessianFactor(PyHessianFunc hessian_func, const std::vector<Keys>&... keys,
+                               bool sparse) {
+  if (sparse) {
+    return sym::Factord(WrapPyHessianFunc<Eigen::SparseMatrix<double>>(std::move(hessian_func)),
+                        keys...);
+  } else {
+    return sym::Factord(WrapPyHessianFunc<Eigen::MatrixXd>(std::move(hessian_func)), keys...);
+  }
 }
 
 using PyJacobianFunc =
     std::function<py::tuple(const sym::Valuesd&, const std::vector<index_entry_t>&)>;
 
-sym::Factord::DenseJacobianFunc WrapPyJacobianFunc(PyJacobianFunc&& jacobian_func) {
-  return sym::Factord::DenseJacobianFunc(
+template <typename Matrix>
+sym::Factord::JacobianFunc<Matrix> WrapPyJacobianFunc(PyJacobianFunc&& jacobian_func) {
+  return sym::Factord::JacobianFunc<Matrix>(
       [jacobian_func = std::move(jacobian_func)](
           const sym::Valuesd& values, const std::vector<index_entry_t>& keys,
-          Eigen::VectorXd* const residual, Eigen::MatrixXd* const jacobian) {
+          Eigen::VectorXd* const residual, Matrix* const jacobian) {
         const py::tuple out_tuple = jacobian_func(values, keys);
         if (residual != nullptr) {
           *residual = py::cast<Eigen::VectorXd>(out_tuple[0]);
         }
         if (jacobian != nullptr) {
-          *jacobian = py::cast<Eigen::MatrixXd>(out_tuple[1]);
+          ThrowIfSparsityMismatch<Matrix>(out_tuple[1]);
+          *jacobian = py::cast<Matrix>(out_tuple[1]);
         }
       });
 }
 
-sym::Factord MakeJacobianFactorSeparateKeys(PyJacobianFunc jacobian_func,
-                                            const std::vector<sym::Key>& keys_to_func,
-                                            const std::vector<sym::Key>& keys_to_optimize) {
-  return sym::Factord::Jacobian(WrapPyJacobianFunc(std::move(jacobian_func)), keys_to_func,
-                                keys_to_optimize);
-}
-
-sym::Factord MakeJacobianFactorCommonKeys(PyJacobianFunc jacobian_func,
-                                          const std::vector<sym::Key>& keys) {
-  return sym::Factord::Jacobian(WrapPyJacobianFunc(std::move(jacobian_func)), keys);
+template <typename... Keys>
+sym::Factord MakeJacobianFactor(PyJacobianFunc jacobian_func, const std::vector<Keys>&... keys,
+                                bool sparse) {
+  if (sparse) {
+    return sym::Factord::Jacobian(
+        WrapPyJacobianFunc<Eigen::SparseMatrix<double>>(std::move(jacobian_func)), keys...);
+  } else {
+    return sym::Factord::Jacobian(WrapPyJacobianFunc<Eigen::MatrixXd>(std::move(jacobian_func)),
+                                  keys...);
+  }
 }
 
 }  // namespace
@@ -112,35 +135,49 @@ void AddFactorWrapper(pybind11::module_ module) {
       point, generates a linear approximation to the residual function.
   )")
       // TODO(brad): Add wrapper of the constructor from SparseHessianFunc
-      .def(py::init(&MakeHessianFactorCommonKeys), py::arg("hessian_func"), py::arg("keys"), R"(
-              Create directly from a (dense) hessian functor. This is the lowest-level constructor.
+      .def(py::init(&MakeHessianFactor<sym::Key>), py::arg("hessian_func"), py::arg("keys"),
+           py::arg("sparse") = false, R"(
+              Create directly from a hessian functor. This is the lowest-level constructor.
 
               Args:
                 keys: The set of input arguments, in order, accepted by func.
+                sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+              Precondition:
+                The jacobian and hessian returned by hessian_func have type scipy.sparse.csc_matrix if and only if sparse = True.
            )")
-      .def(py::init(&MakeHessianFactorSeparateKeys), py::arg("hessian_func"),
-           py::arg("keys_to_func"), py::arg("keys_to_optimize"),
+      .def(py::init(&MakeHessianFactor<sym::Key, sym::Key>), py::arg("hessian_func"),
+           py::arg("keys_to_func"), py::arg("keys_to_optimize"), py::arg("sparse") = false,
            R"(
-              Create directly from a (sparse) hessian functor. This is the lowest-level constructor.
+              Create directly from a hessian functor. This is the lowest-level constructor.
 
               Args:
                 keys_to_func: The set of input arguments, in order, accepted by func.
                 keys_to_optimize: The set of input arguments that correspond to the derivative in func. Must be a subset of keys_to_func.
+                sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+              Precondition:
+                The jacobian and hessian returned by hessian_func have type scipy.sparse.csc_matrix if and only if sparse = True.
            )")
       .def("is_sparse", &sym::Factord::IsSparse,
            "Does this factor use a sparse jacobian/hessian matrix?")
-      .def_static("jacobian", &MakeJacobianFactorCommonKeys, py::arg("jacobian_func"),
-                  py::arg("keys"), R"(
+      .def_static("jacobian", &MakeJacobianFactor<sym::Key>, py::arg("jacobian_func"),
+                  py::arg("keys"), py::arg("sparse") = false, R"(
                     Create from a function that computes the jacobian. The hessian will be computed using the
                     Gauss Newton approximation:
                         H   = J.T * J
                         rhs = J.T * b
 
                     Args:
                       keys: The set of input arguments, in order, accepted by func.
+                      sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+                    Precondition:
+                      The jacobian returned by jacobian_func has type scipy.sparse.csc_matrix if and only if sparse = True.
                   )")
-      .def_static("jacobian", &MakeJacobianFactorSeparateKeys, py::arg("jacobian_func"),
-                  py::arg("keys_to_func"), py::arg("keys_to_optimize"), R"(
+      .def_static("jacobian", &MakeJacobianFactor<sym::Key, sym::Key>, py::arg("jacobian_func"),
+                  py::arg("keys_to_func"), py::arg("keys_to_optimize"), py::arg("sparse") = false,
+                  R"(
             Create from a function that computes the jacobian. The hessian will be computed using the
             Gauss Newton approximation:
                 H   = J.T * J
@@ -149,6 +186,10 @@ void AddFactorWrapper(pybind11::module_ module) {
             Args:
               keys_to_func: The set of input arguments, in order, accepted by func.
               keys_to_optimize: The set of input arguments that correspond to the derivative in func. Must be a subset of keys_to_func.
+              sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+              Precondition:
+                The jacobian returned by jacobian_func has type scipy.sparse.csc_matrix if and only if sparse = True.
           )")
       .def(
           "linearize",

symforce/pybind/cc_sym.pyi

@@ -57,27 +57,37 @@ class Factor:
         self,
         hessian_func: typing.Callable[[Values, typing.List[index_entry_t]], tuple],
         keys: typing.List[Key],
+        sparse: bool = False,
     ) -> None:
         """
-        Create directly from a (dense) hessian functor. This is the lowest-level constructor.
+        Create directly from a hessian functor. This is the lowest-level constructor.
 
         Args:
           keys: The set of input arguments, in order, accepted by func.
+          sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
 
+        Precondition:
+          The jacobian and hessian returned by hessian_func have type scipy.sparse.csc_matrix if and only if sparse = True.
 
 
-        Create directly from a (sparse) hessian functor. This is the lowest-level constructor.
+
+        Create directly from a hessian functor. This is the lowest-level constructor.
 
         Args:
           keys_to_func: The set of input arguments, in order, accepted by func.
           keys_to_optimize: The set of input arguments that correspond to the derivative in func. Must be a subset of keys_to_func.
+          sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+        Precondition:
+          The jacobian and hessian returned by hessian_func have type scipy.sparse.csc_matrix if and only if sparse = True.
         """
     @typing.overload
     def __init__(
         self,
         hessian_func: typing.Callable[[Values, typing.List[index_entry_t]], tuple],
         keys_to_func: typing.List[Key],
         keys_to_optimize: typing.List[Key],
+        sparse: bool = False,
     ) -> None: ...
     def __repr__(self) -> str: ...
     def all_keys(self) -> typing.List[Key]:
@@ -93,6 +103,7 @@ class Factor:
     def jacobian(
         jacobian_func: typing.Callable[[Values, typing.List[index_entry_t]], tuple],
         keys: typing.List[Key],
+        sparse: bool = False,
     ) -> Factor:
         """
                 Create from a function that computes the jacobian. The hessian will be computed using the
@@ -102,6 +113,10 @@ class Factor:
 
                 Args:
                   keys: The set of input arguments, in order, accepted by func.
+                  sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+                Precondition:
+                  The jacobian returned by jacobian_func has type scipy.sparse.csc_matrix if and only if sparse = True.
 
 
 
@@ -113,13 +128,18 @@ class Factor:
         Args:
           keys_to_func: The set of input arguments, in order, accepted by func.
           keys_to_optimize: The set of input arguments that correspond to the derivative in func. Must be a subset of keys_to_func.
+          sparse: Create a sparse factor if True, dense factor if false. Defaults to dense.
+
+          Precondition:
+            The jacobian returned by jacobian_func has type scipy.sparse.csc_matrix if and only if sparse = True.
         """
     @staticmethod
     @typing.overload
     def jacobian(
         jacobian_func: typing.Callable[[Values, typing.List[index_entry_t]], tuple],
         keys_to_func: typing.List[Key],
         keys_to_optimize: typing.List[Key],
+        sparse: bool = False,
     ) -> Factor: ...
     def linearize(self, arg0: Values) -> tuple:
         """