Search Lagrangian given CLF. (#32)

compatible_clf_cbf/clf.py

@@ -13,6 +13,7 @@
     check_array_of_polynomials,
     get_polynomial_result,
     new_sos_polynomial,
+    solve_with_id,
 )
 
 
@@ -29,65 +30,65 @@ class ClfWoInputLimitLagrangian:
     # The Lagrangian
     rho_minus_V: sym.Polynomial
 
-    @classmethod
-    def construct(
-        cls,
+    def get_result(
+        self,
+        result: solvers.MathematicalProgramResult,
+        coefficient_tol: Optional[float],
+    ) -> Self:
+        dVdx_times_f_result = get_polynomial_result(
+            result, self.dVdx_times_f, coefficient_tol
+        )
+        dVdx_times_g_result = get_polynomial_result(
+            result, self.dVdx_times_g, coefficient_tol
+        )
+        rho_minus_V_result = get_polynomial_result(
+            result, self.rho_minus_V, coefficient_tol
+        )
+        return ClfWoInputLimitLagrangian(
+            dVdx_times_f=dVdx_times_f_result,
+            dVdx_times_g=dVdx_times_g_result,
+            rho_minus_V=rho_minus_V_result,
+        )
+
+
+@dataclass
+class ClfWoInputLimitLagrangianDegrees:
+    dVdx_times_f: int
+    dVdx_times_g: List[int]
+    rho_minus_V: int
+
+    def to_lagrangians(
+        self,
         prog: solvers.MathematicalProgram,
         x_set: sym.Variables,
-        dVdx_times_f_degree: int,
-        dVdx_times_g_degrees: List[int],
-        rho_minus_V_degree: int,
         x_equilibrium: np.ndarray,
-    ) -> Self:
+    ) -> ClfWoInputLimitLagrangian:
         """
         Constructs the Lagrangians as SOS polynomials.
         """
-        dVdx_times_f, _ = new_sos_polynomial(prog, x_set, dVdx_times_f_degree)
+        dVdx_times_f, _ = new_sos_polynomial(prog, x_set, self.dVdx_times_f)
         dVdx_times_g = np.array(
-            [
-                new_sos_polynomial(prog, x_set, degree)[0]
-                for degree in dVdx_times_g_degrees
-            ]
+            [prog.NewFreePolynomial(x_set, degree) for degree in self.dVdx_times_g]
         )
         # We know that V(x_equilibrium) = 0 and dVdx at x_equilibrium is also
         # 0. Hence by
         # -(1+λ₀(x))*(∂V/∂x*f(x)+κ*V(x))−λ₁(x)*∂V/∂x*g(x)−λ₂(x)*(ρ−V(x))
         # being sos, we know that λ₂(x_equilibrium) = 0.
         # When x_equilibrium = 0, this means that λ₂(0) = 0. Since λ₂(x) is
         # sos, we know that its monomial basis doesn't contain 1.
-        if rho_minus_V_degree == 0:
+        if self.rho_minus_V == 0:
             rho_minus_V = sym.Polynomial()
         else:
             if np.all(x_equilibrium == 0):
                 monomial_basis = sym.MonomialBasis(
-                    x_set, int(np.floor(rho_minus_V_degree / 2))
+                    x_set, int(np.floor(self.rho_minus_V / 2))
                 )
                 assert monomial_basis[-1].total_degree() == 0
                 rho_minus_V, _ = prog.NewSosPolynomial(monomial_basis[:-1])
             else:
-                rho_minus_V, _ = prog.NewSosPolynomial(x_set, rho_minus_V_degree)
+                rho_minus_V, _ = prog.NewSosPolynomial(x_set, self.rho_minus_V)
         return ClfWoInputLimitLagrangian(dVdx_times_f, dVdx_times_g, rho_minus_V)
 
-    def get_result(
-        self,
-        result: solvers.MathematicalProgramResult,
-        coefficient_tol: Optional[float],
-    ) -> Self:
-        dVdx_times_f_result = get_polynomial_result(
-            result, self.dVdx_times_f, coefficient_tol
-        )
-        dVdx_times_g_result = get_polynomial_result(
-            result, self.dVdx_times_g, coefficient_tol
-        )
-        rho_minus_V_result = get_polynomial_result(
-            result, self.rho_minus_V, coefficient_tol
-        )
-        return ClfWoInputLimitLagrangian(
-            dVdx_times_f=dVdx_times_f_result,
-            dVdx_times_g=dVdx_times_g_result,
-            rho_minus_V=rho_minus_V_result,
-        )
-
 
 @dataclass
 class ClfWInputLimitLagrangian:
@@ -116,6 +117,24 @@ def get_result(
         )
 
 
+@dataclass
+class ClfWInputLimitLagrangianDegrees:
+    V_minus_rho: int
+    Vdot: List[int]
+
+    def to_lagrangians(
+        self,
+        prog: solvers.MathematicalProgram,
+        x_set: sym.Variables,
+        x_equilibrium: np.ndarray,
+    ) -> ClfWInputLimitLagrangian:
+        V_minus_rho, _ = new_sos_polynomial(prog, x_set, self.V_minus_rho)
+        Vdot = np.array(
+            [new_sos_polynomial(prog, x_set, degree) for degree in self.Vdot]
+        )
+        return ClfWInputLimitLagrangian(V_minus_rho, Vdot)
+
+
 class ControlLyapunov:
     """
     For a control affine system
@@ -187,6 +206,28 @@ def __init__(
             assert u_vertices.shape[1] == self.nu
         self.u_vertices = u_vertices
 
+    def search_lagrangian_given_clf(
+        self,
+        V: sym.Polynomial,
+        rho: float,
+        kappa: float,
+        lagrangian_degrees: Union[
+            ClfWInputLimitLagrangianDegrees, ClfWoInputLimitLagrangianDegrees
+        ],
+        solver_id: Optional[solvers.SolverId] = None,
+        solver_options: Optional[solvers.SolverOptions] = None,
+        coefficient_tol: Optional[float] = None,
+    ) -> Union[ClfWInputLimitLagrangian, ClfWoInputLimitLagrangian]:
+        prog = solvers.MathematicalProgram()
+        prog.AddIndeterminates(self.x_set)
+        lagrangians = lagrangian_degrees.to_lagrangians(
+            prog, self.x_set, self.x_equilibrium
+        )
+        self._add_clf_condition(prog, V, lagrangians, rho, kappa)
+        result = solve_with_id(prog, solver_id, solver_options)
+        assert result.is_success()
+        return lagrangians.get_result(result, coefficient_tol)
+
     def _add_clf_condition(
         self,
         prog: solvers.MathematicalProgram,

tests/test_clf.py

@@ -66,14 +66,14 @@ def test_add_clf_condition_wo_input_limit(self):
         prog.AddIndeterminates(dut.x_set)
         V = self.get_V_init()
 
-        lagrangians = mut.ClfWoInputLimitLagrangian.construct(
-            prog,
-            dut.x_set,
-            dVdx_times_f_degree=2,
-            dVdx_times_g_degrees=[4, 4],
-            rho_minus_V_degree=4,
-            x_equilibrium=dut.x_equilibrium,
+        lagrangian_degrees = mut.ClfWoInputLimitLagrangianDegrees(
+            dVdx_times_f=2, dVdx_times_g=[4, 4], rho_minus_V=4
         )
+
+        lagrangians = lagrangian_degrees.to_lagrangians(
+            prog, dut.x_set, dut.x_equilibrium
+        )
+
         # The Lagrangian for rho-V doesn't contain constant or linear terms.
         assert (
             sym.Monomial()
@@ -99,3 +99,26 @@ def test_add_clf_condition_wo_input_limit(self):
             - lagrangians.rho_minus_V * (rho - V)
         )
         assert condition.EqualTo(condition_expected)
+
+    def test_search_lagrangian_given_clf_wo_input_limit(self):
+        """
+        Test search_lagrangian_given_clf without input limits.
+        """
+        dut = mut.ControlLyapunov(
+            f=self.f,
+            g=self.g,
+            x=self.x,
+            x_equilibrium=np.zeros(self.nx),
+            u_vertices=None,
+        )
+        V = self.get_V_init()
+
+        lagrangians = dut.search_lagrangian_given_clf(
+            V,
+            rho=0.01,
+            kappa=0.001,
+            lagrangian_degrees=mut.ClfWoInputLimitLagrangianDegrees(
+                dVdx_times_f=2, dVdx_times_g=[3, 3], rho_minus_V=4
+            ),
+        )
+        assert isinstance(lagrangians, mut.ClfWoInputLimitLagrangian)