Vectorization refactor

theseus/core/cost_function.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import abc
-from typing import Any, Iterable, List, Optional, Tuple, cast
+from typing import Any, List, Optional, Tuple, cast
 
 import torch
 import torch.autograd.functional as autogradF
@@ -17,19 +17,6 @@
 from .variable import Variable
 
 
-def _register_vars_in_list(
-    cost_fn: TheseusFunction, var_list_: Iterable[Variable], is_optim: bool = False
-):
-    for var_ in var_list_:
-        if hasattr(cost_fn, var_.name):
-            raise RuntimeError(f"Variable name {var_.name} is not allowed.")
-        setattr(cost_fn, var_.name, var_)
-        if is_optim:
-            cost_fn.register_optim_var(var_.name)
-        else:
-            cost_fn.register_aux_var(var_.name)
-
-
 # A cost function is defined by the variables interacting in it,
 # and a cost weight for weighting errors and jacobians
 # This is an abstract class for cost functions of different types.
@@ -128,8 +115,8 @@ def __init__(
             raise ValueError(
                 "AutodiffCostFunction must receive at least one optimization variable."
             )
-        _register_vars_in_list(self, optim_vars, is_optim=True)
-        _register_vars_in_list(self, aux_vars, is_optim=False)
+        self.register_vars(optim_vars, optim_vars=True)
+        self.register_vars(aux_vars, optim_vars=False)
 
         self._err_fn = err_fn
         self._dim = dim

theseus/core/theseus_function.py

@@ -64,6 +64,16 @@ def register_aux_vars(self, aux_var_names: Sequence[str]):
         for name in aux_var_names:
             self.register_aux_var(name)
 
+    def register_vars(self, vars: Iterable[Variable], optim_vars: bool = False):
+        for var_ in vars:
+            if hasattr(self, var_.name):
+                raise RuntimeError(f"Variable name {var_.name} is not allowed.")
+            setattr(self, var_.name, var_)
+            if optim_vars:
+                self.register_optim_var(var_.name)
+            else:
+                self.register_aux_var(var_.name)
+
     # Must copy everything
     @abc.abstractmethod
     def _copy_impl(self, new_name: Optional[str] = None) -> "TheseusFunction":

theseus/core/vectorizer.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from .cost_function import CostFunction, _register_vars_in_list
+from .cost_function import CostFunction
 from .objective import Objective
 from .variable import Variable
 
@@ -35,8 +35,8 @@ def __init__(self, cost_fn: CostFunction, name: Optional[str] = None):
             name = f"wrapper({cost_fn.name})"
         super().__init__(cost_fn.weight, name=name)
         self.cost_fn = cost_fn
-        _register_vars_in_list(self, cost_fn.optim_vars, is_optim=True)
-        _register_vars_in_list(self, cost_fn.aux_vars, is_optim=False)
+        self.register_vars(cost_fn.optim_vars, optim_vars=True)
+        self.register_vars(cost_fn.aux_vars, optim_vars=False)
         self._cached_error: Optional[torch.Tensor] = None
         self._cached_jacobians: Optional[List[torch.Tensor]] = None
 
@@ -74,10 +74,6 @@ def __repr__(self) -> str:
 
 # This class replaces the Objective's iterator for one that takes advantage of
 # cost function vectorization
-# TODO:
-#   - Tests to add:
-#         + Test that vectorization results in correct costs
-#         + Vectorize variable update after NLOPT step
 class Vectorize:
     _SHARED_TOKEN = "__shared__"
 
@@ -157,8 +153,12 @@ def _get_name(var_idx, name_set_):
         return info
 
     @staticmethod
-    def _get_all_vars(cf) -> List[Variable]:
-        return list(cf.optim_vars) + list(cf.aux_vars) + list(cf.weight.aux_vars)
+    def _get_all_vars(cf: CostFunction) -> List[Variable]:
+        return (
+            list(cf.optim_vars)
+            + list(cf.aux_vars)  # type:ignore
+            + list(cf.weight.aux_vars)  # type:ignore
+        )
 
     @staticmethod
     def _expand(tensor: torch.Tensor, size: int) -> torch.Tensor: