Enabled back ellipsoidal damping in LM with linear solvers support checks

theseus/optimizer/nonlinear/levenberg_marquardt.py

@@ -9,10 +9,20 @@
 
 from theseus.core import Objective
 from theseus.optimizer import Linearization
-from theseus.optimizer.linear import DenseSolver, LinearSolver
+from theseus.optimizer.linear import DenseSolver, LinearSolver, LUCudaSparseSolver
 
 from .nonlinear_least_squares import NonlinearLeastSquares
 
+_LM_ALLOWED_SOLVERS = [DenseSolver, LUCudaSparseSolver]
+
+
+def _check_ellipsoidal_damping_cls(linear_solver: LinearSolver):
+    good = False
+    for lsc in _LM_ALLOWED_SOLVERS:
+        if isinstance(linear_solver, lsc):
+            good = True
+    return good
+
 
 # See Nocedal and Wright, Numerical Optimization, pp. 258 - 261
 # https://www.csie.ntu.edu.tw/~r97002/temp/num_optimization.pdf
@@ -40,6 +50,7 @@ def __init__(
             max_iterations=max_iterations,
             step_size=step_size,
         )
+        self._allows_ellipsoidal = _check_ellipsoidal_damping_cls(self.linear_solver)
 
     def compute_delta(
         self,
@@ -48,15 +59,17 @@ def compute_delta(
         damping_eps: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
-        if ellipsoidal_damping:
-            raise NotImplementedError("Ellipsoidal damping is not currently supported.")
-        if ellipsoidal_damping and not isinstance(self.linear_solver, DenseSolver):
+
+        solvers_str = ",".join(c.__name__ for c in _LM_ALLOWED_SOLVERS)
+        if ellipsoidal_damping and not self._allows_ellipsoidal:
             raise NotImplementedError(
-                "Ellipsoidal damping is only supported when using DenseSolver."
+                f"Ellipsoidal damping is only supported by solvers with type "
+                f"[{solvers_str}]."
             )
-        if damping_eps and not isinstance(self.linear_solver, DenseSolver):
+        if damping_eps and not self._allows_ellipsoidal:
             raise NotImplementedError(
-                "damping eps is only supported when using DenseSolver."
+                f"damping eps is only supported by solvers with type "
+                f"[{solvers_str}]."
             )
         damping_eps = damping_eps or 1e-8
 

theseus/optimizer/nonlinear/tests/test_levenberg_marquardt.py

@@ -4,14 +4,24 @@
 # LICENSE file in the root directory of this source tree.
 
 import pytest  # noqa: F401
+import torch
 
 import theseus as th
 
 from .common import run_nonlinear_least_squares_check
 
 
+@pytest.fixture
+def mock_objective():
+    objective = th.Objective()
+    v1 = th.Vector(1, name="v1")
+    v2 = th.Vector(1, name="v2")
+    objective.add(th.eb.VariableDifference(v1, th.ScaleCostWeight(1.0), v2))
+    return objective
+
+
 def test_levenberg_marquartd():
-    for ellipsoidal_damping in [False]:
+    for ellipsoidal_damping in [True, False]:
         for damping in [0, 0.001, 0.01, 0.1]:
             run_nonlinear_least_squares_check(
                 th.LevenbergMarquardt,
@@ -22,3 +32,27 @@ def test_levenberg_marquartd():
                 },
                 singular_check=damping < 0.001,
             )
+
+
+def test_ellipsoidal_damping_compatibility(mock_objective):
+    mock_objective.update({"v1": torch.ones(1, 1), "v2": torch.zeros(1, 1)})
+    for lsc in [th.LUDenseSolver, th.CholeskyDenseSolver]:
+        optimizer = th.LevenbergMarquardt(mock_objective, lsc)
+        optimizer.optimize(ellipsoidal_damping=True)
+        optimizer.optimize(damping_eps=0.1)
+
+    for lsc in [th.CholmodSparseSolver]:
+        optimizer = th.LevenbergMarquardt(mock_objective, lsc)
+        with pytest.raises(RuntimeError):
+            optimizer.optimize(ellipsoidal_damping=True)
+        with pytest.raises(RuntimeError):
+            optimizer.optimize(damping_eps=0.1)
+
+
+@pytest.mark.cuda
+def test_ellipsoidal_damping_compatibility_cuda(mock_objective):
+    mock_objective.update({"v1": torch.ones(1, 1), "v2": torch.zeros(1, 1)})
+    for lsc in [th.LUCudaSparseSolver]:
+        optimizer = th.LevenbergMarquardt(mock_objective, lsc)
+        optimizer.optimize(ellipsoidal_damping=True)
+        optimizer.optimize(damping_eps=0.1)