Enable OptimizationProblem to generate multiple factors

Currently the `OptimizationProblem` class can only generate one big
factor to compute a joint residual/jacobian; it does not allow for
certain objectives to be added/removed from the overall linearization
function dynamically. This PR makes it so that multiple factors can be
generated by on `OptimizationProblem`.

This is accomplished by setting the "factor" metadata field of the
residual blocks. Residual blocks with the same value for the "factor"
metadata field will be grouped into a single factor separate from the
"main" optimization problem factor (which will contain everything else).

Topic: optimization_problem_split_factor
Relative: fixed_size_optimizer
GitOrigin-RevId: 3c7bb8fed45f4a896b990de88078ede573afe46b

symforce/opt/optimization_problem.py

@@ -3,6 +3,7 @@
 # This source code is under the Apache 2.0 license found in the LICENSE file.
 # ----------------------------------------------------------------------------
 
+import functools
 import itertools
 
 import symforce.symbolic as sf
@@ -113,8 +114,9 @@ def generate(
         config: T.Optional[CppConfig] = None,
     ) -> None:
         """
-        Generate everything needed to optimize ``self`` in C++. This currently assumes there is
-        only one factor generated for the optimization problem.
+        Generate everything needed to optimize ``self`` in C++. This will typically only generate
+        one factor, but can be configured to generate multiple different factors by setting the
+        "factor_name" field of the residual blocks to be split into separate factors.
 
         Args:
             output_dir: Directory in which to output the generated files.
@@ -142,6 +144,37 @@ def generate(
                 )
             )
 
+    def residual_blocks_per_factor(self, default_name: str) -> Values:
+        """
+        Returns a Values with the residual blocks split by factor name.
+        """
+        factor_blocks = Values()
+        for block_name, block in self.residual_blocks.items_recursive():
+            factor_name = block.factor_name if block.factor_name is not None else default_name
+            if factor_name not in factor_blocks:
+                factor_blocks[factor_name] = Values()
+            factor_blocks[factor_name][block_name] = block
+        return factor_blocks
+
+    @staticmethod
+    def compute_jacobians(
+        keys: T.Sequence[str], inputs: Values, residual_blocks: Values
+    ) -> sf.Matrix:
+        """
+        Functor that computes the jacobians of the residual with respect to a set of keys
+
+        The set of keys is not known when make_symbolic_factors is called, because we may want
+        to create a :class:`.numeric_factor.NumericFactor` which computes derivatives with
+        respect to different sets of optimized variables.
+        """
+        jacobians = [
+            residual_block.compute_jacobians(
+                [inputs[key] for key in keys], residual_name=residual_name, key_names=keys
+            )
+            for residual_name, residual_block in residual_blocks.items_recursive()
+        ]
+        return sf.Matrix.block_matrix(jacobians)
+
     def make_symbolic_factors(
         self,
         name: str,
@@ -161,21 +194,8 @@ def make_symbolic_factors(
         """
         inputs = self.inputs.dataclasses_to_values()
 
-        def compute_jacobians(keys: T.Iterable[str]) -> sf.Matrix:
-            """
-            Functor that computes the jacobians of the residual with respect to a set of keys
-
-            The set of keys is not known when make_symbolic_factors is called, because we may want
-            to create a :class:`.numeric_factor.NumericFactor` which computes derivatives with
-            respect to different sets of optimized variables.
-            """
-            jacobians = [
-                residual_block.compute_jacobians(
-                    [inputs[key] for key in keys], residual_name=residual_key, key_names=keys
-                )
-                for residual_key, residual_block in self.residual_blocks.items_recursive()
-            ]
-            return sf.Matrix.block_matrix(jacobians)
+        # Split blocks by factor
+        factor_blocks = self.residual_blocks_per_factor(default_name=name)
 
         return [
             Factor.from_inputs_and_residual(
@@ -186,11 +206,18 @@ def compute_jacobians(keys: T.Iterable[str]) -> sf.Matrix:
                         for key, value in inputs.items_recursive()
                     }
                 ),
-                residual=sf.M(self.residuals.to_storage()),
+                residual=sf.M(
+                    ops.StorageOps.to_storage(
+                        [block.residual for block in residual_blocks.values_recursive()]
+                    )
+                ),
                 config=config,
-                custom_jacobian_func=compute_jacobians,
+                custom_jacobian_func=functools.partial(
+                    self.compute_jacobians, inputs=inputs, residual_blocks=residual_blocks
+                ),
                 name=name,
             )
+            for name, residual_blocks in factor_blocks.items()
         ]
 
     def make_numeric_factors(

symforce/opt/residual_block.py

@@ -22,6 +22,9 @@ class ResidualBlock:
     residual: sf.Matrix
     extra_values: T.Optional[T.Dataclass] = None
     metadata: T.Optional[T.Dict[str, T.Any]] = None
+    # When creating an `OptimizationProblem`, all residual blocks with the same factor_name will
+    # be split out into a separate factor from the main optimization problem factor.
+    factor_name: T.Optional[str] = None
 
     def compute_jacobians(
         self,
@@ -54,6 +57,13 @@ def set_metadata(self, key: str, value: T.Any) -> ResidualBlock:
             self.metadata[key] = value
         return self
 
+    def set_factor_name(self, name: T.Optional[str]) -> ResidualBlock:
+        """
+        Sets the factor name of the residual block and returns the residual block.
+        """
+        self.factor_name = name
+        return self
+
 
 @dataclass
 class ResidualBlockWithCustomJacobian(ResidualBlock):