fix discrete_relative_degree==0

dynax/linearize.py

@@ -1,4 +1,4 @@
-"""Functions related to feedback linearization of nonlinear systems."""
+"""Functions related to input-output linearization of nonlinear systems."""
 
 from collections.abc import Callable
 from functools import partial
@@ -20,11 +20,23 @@
 )
 
 
-# TODO: make this a method of ControlAffine
 def relative_degree(
     sys: AbstractControlAffine, xs: Array, output: Optional[int] = None
 ) -> int:
-    """Estimate relative degree of system on region xs."""
+    """Estimate the relative degree of a single-output control-affine system.
+
+    Tests that Lie derivatives of the output function are zero exactly up to the
+    relative degree on the state space samples `xs`.
+
+    Args:
+        sys: Control affine system.
+        xs: Samples of the state space stacked along the first axis.
+        output: Index of the output function if the system has multiple outputs.
+
+    Returns:
+        Established relative degree of the system.
+
+    """
     if sys.n_inputs not in ["scalar", 1]:
         raise ValueError("System must be single input.")
     if output is None:
@@ -42,15 +54,16 @@ def relative_degree(
         else:
             LgLfn1h = lie_derivative(sys.g, lie_derivative(sys.f, h, n - 1))
             res = jax.vmap(LgLfn1h)(xs)
+
         if np.all(res == 0.0):
             continue
         elif np.all(res != 0.0):
             return n
-        else:
-            raise RuntimeError("sys has ill-defined relative degree.")
-    raise RuntimeError("Could not estimate relative degree. Increase max_reldeg.")
+
+    raise RuntimeError("sys has ill-defined relative degree.")
 
 
+# TODO: remove?
 def is_controllable(A, B) -> bool:
     """Test controllability of linear system."""
     n = A.shape[0]
@@ -155,28 +168,35 @@ def discrete_relative_degree(
     sys: AbstractSystem,
     xs: Array,
     us: Array,
-    max_reldeg=10,
     output: Optional[int] = None,
 ):
     """Estimate relative degree of discrete-time system on region xs.
 
     See :cite:p:`leeLinearizationNonlinearControl2022{def 7.7.}`.
 
     """
-    f = sys.vector_field
-    h = sys.output
+    if sys.n_inputs not in ["scalar", 1]:
+        raise ValueError("System must be single input.")
+    if output is None:
+        # Make sure system has single output
+        if sys.n_outputs not in ["scalar", 1]:
+            raise ValueError(f"Output is None, but system has {sys.n_outputs} outputs.")
+        h = sys.output
+    else:
+        h = lambda *args, **kwargs: sys.output(*args, **kwargs)[output]
 
-    y_depends_u = jax.grad(lambda n, x, u: h(propagate(f, n, x, u)), 2)
+    f = sys.vector_field
+    y = lambda n, x, u: h(propagate(f, n, x, u), u)
+    y_depends_u = jax.grad(y, 2)
 
-    for n in range(1, max_reldeg + 1):
+    max_reldeg = jnp.size(sys.initial_state)
+    for n in range(0, max_reldeg + 1):
         res = jax.vmap(partial(y_depends_u, n))(xs, us)
         if np.all(res == 0):
             continue
         elif np.all(res != 0):
             return n
-        else:
-            raise RuntimeError("sys has ill defined relative degree.")
-    raise RuntimeError("Could not estmate relative degree. Increase max_reldeg.")
+    raise RuntimeError("sys has ill defined relative degree.")
 
 
 def discrete_input_output_linearize(
@@ -226,7 +246,7 @@ def fn(u, args):
 
 
 class DiscreteLinearizingSystem(AbstractSystem, _CoupledSystemMixin):
-    r"""Dynamics computing linearizing feedback as output."""
+    """Dynamics computing linearizing feedback as output."""
 
     _v: Callable
 

tests/test_linearize.py

@@ -23,6 +23,23 @@
 tols = dict(rtol=1e-04, atol=1e-06)
 
 
+class Allpass(AbstractControlAffine):
+    initial_state = jnp.zeros(1)
+    n_inputs = "scalar"
+
+    def f(self, x):
+        return jnp.array(0.)
+
+    def g(self, x):
+        return jnp.array(0.)
+
+    def h(self, x):
+        return jnp.array(0.)
+
+    def i(self, x):
+        return jnp.array(1.)
+
+
 class SpringMassDamperWithOutput(AbstractControlAffine):
     m: float = 0.1
     r: float = 0.1
@@ -52,22 +69,6 @@ def test_relative_degree():
     sys = SpringMassDamperWithOutput(out=1)
     assert relative_degree(sys, xs) == 1
 
-    class Allpass(AbstractControlAffine):
-        initial_state = jnp.zeros(1)
-        n_inputs = "scalar"
-
-        def f(self, x):
-            return jnp.array(0.)
-
-        def g(self, x):
-            return jnp.array(0.)
-
-        def h(self, x):
-            return jnp.array(0.)
-
-        def i(self, x):
-            return jnp.array(1.)
-
     assert relative_degree(Allpass(), xs) == 0
 
 
@@ -78,12 +79,11 @@ def test_discrete_relative_degree():
     sys = SpringMassDamperWithOutput(out=0)
     assert discrete_relative_degree(sys, xs, us) == 2
 
-    with npt.assert_raises(RuntimeError):
-        discrete_relative_degree(sys, xs, us, max_reldeg=1)
-
     sys = SpringMassDamperWithOutput(out=1)
     assert discrete_relative_degree(sys, xs, us) == 1
 
+    assert discrete_relative_degree(Allpass(), xs, us) == 0
+
 
 def test_is_controllable():
     n = 3