From 6d1fe2269bf2a3ca1d3a25c06dc1031276c87be7 Mon Sep 17 00:00:00 2001 From: fhchl Date: Sun, 21 Apr 2024 22:09:09 +0200 Subject: [PATCH] improve docs for linearize.py --- docs/source/conf.py | 2 + dynax/linearize.py | 161 +++++++++++++++++++++++++++------------- tests/test_linearize.py | 4 +- 3 files changed, 115 insertions(+), 52 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 02955d6..0410668 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -85,6 +85,8 @@ "jaxtyping": ("https://docs.kidger.site/jaxtyping/", None), "diffrax": ("https://docs.kidger.site/diffrax/", None), "equinox": ("https://docs.kidger.site/equinox/", None), + "optimistix": ("https://docs.kidger.site/optimistix/", None), + "lineax": ("https://docs.kidger.site/lineax/", None), "jax": ("https://jax.readthedocs.io/en/latest/", None), } diff --git a/dynax/linearize.py b/dynax/linearize.py index bb21e94..eb5592c 100644 --- a/dynax/linearize.py +++ b/dynax/linearize.py @@ -10,6 +10,7 @@ import optimistix as optx from jax import Array +from .custom_types import Scalar from .derivative import lie_derivative from .system import ( _CoupledSystemMixin, @@ -23,18 +24,18 @@ def relative_degree( sys: AbstractControlAffine, xs: Array, output: Optional[int] = None ) -> int: - """Estimate the relative degree of a single-output control-affine system. + """Estimate the relative degree of a SISO 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`. + Tests that the Lie derivatives of the output are zero exactly up to the + relative-degree order for each state in `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. + output: Optional index of the output if the system has multiple outputs. Returns: - Established relative degree of the system. + Estimated relative degree of the system. """ if sys.n_inputs not in ["scalar", 1]: @@ -70,7 +71,7 @@ def is_controllable(A, B) -> bool: contrmat = np.hstack([np.linalg.matrix_power(A, ni).dot(B) for ni in range(n)]) return np.linalg.matrix_rank(contrmat) == n - +# TODO: Adapt to general nonlinear reference system. def input_output_linearize( sys: AbstractControlAffine, reldeg: int, @@ -78,22 +79,24 @@ def input_output_linearize( output: Optional[int] = None, asymptotic: Optional[Sequence] = None, reg: Optional[float] = None, -) -> Callable[[Array, Array, float], Array]: - """Construct input-output linearizing feedback law. +) -> Callable[[Array, Array, float], Scalar]: + """Construct an input-output linearizing feedback law. Args: - sys: nonlinear model with single input - reldeg: relative degree of `sys` and lower bound of relative degree of `ref` - ref: target model with single input - output: specify linearizing output if systems have multiple outputs - asymptotic: If `None`, compute the exactly linearizing law. Otherwise, - a sequence of length `reldeg` defining the tracking behaviour. - reg: parameter that control the linearization effort. Only effective if - asymptotic is not None. - - Note: - Relative degree of `ref` must be same or higher than degree of `sys`. - Only single-input-single-output systems are currently supported. + sys: Control-affine system with well defined relative degree and single input + and output. + reldeg: Relative degree of `sys` and lower bound of relative degree of `ref`. + ref: Linear target system with single input and output. + output: Optional index of the output if the `sys` has multiple outputs. + asymptotic: If `None`, compute the exactly linearizing law. Otherwise, compute + an asymptotically linearizing law. Then `asymptotic` is interpreted as the + sequence of length `reldeg` of coefficients of the characteristic polynomial + of the tracking error system. + reg: Regularization parameter that controls the linearization effort. Only + effective if asymptotic is not `None`. + + Returns: + A feedback law `u = u(x, z, v)` that input-output linearizes the system. """ assert sys.n_inputs == ref.n_inputs, "systems have same input dimension" @@ -115,7 +118,7 @@ def input_output_linearize( if asymptotic is None: - def feedbacklaw(x: Array, z: Array, v: float) -> Array: + def feedbacklaw(x: Array, z: Array, v: float) -> Scalar: y_reldeg_ref = cAn.dot(z) + cAnm1b * v y_reldeg = Lfnh(x) out = (y_reldeg_ref - y_reldeg) / LgLfnm1h(x) @@ -136,7 +139,7 @@ def feedbacklaw(x: Array, z: Array, v: float) -> Array: cAis = [c.dot(np.linalg.matrix_power(A, i)) for i in range(reldeg)] Lfihs = [lie_derivative(sys.f, h, i) for i in range(reldeg)] - def feedbacklaw(x: Array, z: Array, v: float) -> Array: + def feedbacklaw(x: Array, z: Array, v: float) -> Scalar: y_reldeg_ref = cAn.dot(z) + cAnm1b * v y_reldeg = Lfnh(x) ae0s = jnp.array( @@ -156,12 +159,13 @@ def feedbacklaw(x: Array, z: Array, v: float) -> Array: return feedbacklaw -def propagate(f: Callable[[Array, float], Array], n: int, x: Array, u: float) -> Array: - """Propagates system n steps.""" - # TODO: replace by lax.scan - if n == 0: - return x - return propagate(f, n - 1, f(x, u), u) +def _propagate(f: Callable[[Array, float], Array], n: int, x: Array, u: float) -> Array: + # Propagates system for n <= discrete_relative_degree(sys) steps.""" + def fun(x, _): + return f(x, u), None + + xn, _ = jax.lax.scan(fun, x, jnp.arange(n)) + return xn def discrete_relative_degree( @@ -170,7 +174,20 @@ def discrete_relative_degree( us: Array, output: Optional[int] = None, ): - """Estimate relative degree of discrete-time system on region xs. + """Estimate the relative degree of a SISO discrete-time system. + + Tests that exactly the first relative-degree - 1 output samples are independent of + the input for each `(x, u)` for the initial state and input samples `(xs, us)`. In + this way, the discrete relative-degree can be interpreted as a system delay. + + Args: + sys: Concrete dynamical system. + xs: Initial state samples stacked along the first axis. + us: Initial input samples stacked along the first axis. + output: Optional index of the output if the system has multiple outputs. + + Returns: + The discrete-time relative degree of the system. See :cite:p:`leeLinearizationNonlinearControl2022{def 7.7.}`. @@ -186,7 +203,7 @@ def discrete_relative_degree( h = lambda *args, **kwargs: sys.output(*args, **kwargs)[output] f = sys.vector_field - y = lambda n, x, u: h(propagate(f, n, x, u), u) + y = lambda n, x, u: h(_propagate(f, n, x, u), u) y_depends_u = jax.grad(y, 2) max_reldeg = jnp.size(sys.initial_state) @@ -206,9 +223,27 @@ def discrete_input_output_linearize( output: Optional[int] = None, solver: Optional[optx.AbstractRootFinder] = None, ) -> Callable[[Array, Array, float, float], float]: - """Construct the input-output linearizing feedback for a discrete-time system.""" + """Construct the input-output linearizing feedback for a discrete-time system. + + This is similar to model-predictive control with a horizon of a single time + step and without constraints. The reference system can be nonlinear, in + which case the feedback law implements an exact tracking controller. - # Lee 2022, Chap. 7.4 + Args: + sys: Concrete dynamical system. + reldeg: Relative degree of `sys` and lower bound of relative degree of `ref`. + ref: Discrete-time reference system. + output: Optional index of the output if the `sys` has multiple outputs. + solver: Root finding algorithm to solve the feedback law. Defaults to + :py:class:`optimistix.Newton` with absolute and relative tolerance `1e-6`. + + Returns: + A feedback law :math:`u_n = u(x_n, z_n, v_n, u_{n-1})` that input-output + linearizes the system. + + See :cite:p:`leeLinearizationNonlinearControl2022{def 7.4.}`. + + """ f = lambda x, u: sys.vector_field(x, u) h = sys.output if sys.n_inputs != ref.n_inputs != 1: @@ -231,12 +266,12 @@ def y_reldeg_ref(z, v): B_reldeg = c.dot(np.linalg.matrix_power(A, reldeg - 1)).dot(b) return _output(A_reldeg.dot(z) + B_reldeg.dot(v)) else: - _output(ref.output(propagate(ref.vector_field, reldeg, z, v))) + _output(ref.output(_propagate(ref.vector_field, reldeg, z, v))) - def feedbacklaw(x: Array, z: Array, v: float, u_prev: float): - def fn(u, args): + def feedbacklaw(x: Array, z: Array, v: float, u_prev: float) -> float: + def fn(u, _): return ( - _output(h(propagate(f, reldeg, x, u))) - y_reldeg_ref(z, v) + _output(h(_propagate(f, reldeg, x, u))) - y_reldeg_ref(z, v) ).squeeze() u = optx.root_find(fn, solver, u_prev).value @@ -246,7 +281,26 @@ def fn(u, args): class DiscreteLinearizingSystem(AbstractSystem, _CoupledSystemMixin): - """Dynamics computing linearizing feedback as output.""" + r"""Coupled discrete-time system of dynamics, reference and linearizing feedback. + + .. math:: + + x_{n+1} &= f^{sys}(x_n, v_n) \\ + z_{n+1} &= f^{ref}(z_n, u_n) \\ + y_n &= v_n = v(x_n, z_n, u_n) + + where :math:`v` is such that :math:`y_n^{sys} = h^{sys}(x_n, u_n)` equals + :math:`y^{ref}_n = h^{ref}(z_n, u_n)`. + + Args: + sys: Concrete discrete-time system. + refsys: Reference system. + reldeg: Relative degree of `sys` and lower bound of relative degree of + `refsys`. + fb_kwargs: Additional keyword arguments passed to + :py:func:`discrete_input_output_linearize`. + + """ _v: Callable @@ -255,18 +309,18 @@ class DiscreteLinearizingSystem(AbstractSystem, _CoupledSystemMixin): def __init__( self, sys: AbstractSystem, - refsys: AbstractSystem, + ref: AbstractSystem, reldeg: int, **fb_kwargs, ): if sys.n_inputs != "scalar": raise ValueError("Only single input systems supported.") self._sys1 = sys - self._sys2 = refsys + self._sys2 = ref self.initial_state = jnp.append( self._pack_states(self._sys1.initial_state, self._sys2.initial_state), 0.0 ) - self._v = discrete_input_output_linearize(sys, reldeg, refsys, **fb_kwargs) + self._v = discrete_input_output_linearize(sys, reldeg, ref, **fb_kwargs) def vector_field(self, x, u=None, t=None): (x, z), v_last = self._unpack_states(x[:-1]), x[-1] @@ -277,23 +331,28 @@ def vector_field(self, x, u=None, t=None): def output(self, x, u=None, t=None): (x, z), v_last = self._unpack_states(x[:-1]), x[-1] - v = self._v(x, z, u, v_last) # FIXME: feedback law called twice + v = self._v(x, z, u, v_last) # NOTE: feedback law is computed twice return v class LinearizingSystem(DynamicStateFeedbackSystem): - r"""Coupled ODE of nonlinear dynamics, linear reference and io linearizing law. + r"""Coupled ODE of nonlinear dynamics, linear reference and linearizing feedback. .. math:: - ẋ &= f(x) + g(x)y \\ - ż &= Az + Bu \\ - y &= h(x, z, u) + ẋ &= f(x) + g(x)v \\ + ż &= Az + Bu \\ + y &= v = v(x, z, u) + + where :math:`v` is such that :math:`y^{sys} = h(x) + i(x)v` equals + :math:`y^{ref} = Cz + Du`. Args: - sys: nonlinear control affine system - refsys: linear reference system - reldeg: relative degree of sys and lower bound of relative degree of refsys + sys: Concrete control-affine system. + ref: Linear reference system. + reldeg: Relative degree of `sys` and lower bound of relative degree of `refsys`. + fb_kwargs: Additional keyword arguments passed to + :py:func:`input_output_linearize`. """ @@ -302,12 +361,12 @@ class LinearizingSystem(DynamicStateFeedbackSystem): def __init__( self, sys: AbstractControlAffine, - refsys: LinearSystem, + ref: LinearSystem, reldeg: int, **fb_kwargs, ): - v = input_output_linearize(sys, reldeg, refsys, **fb_kwargs) - super().__init__(sys, refsys, v) + v = input_output_linearize(sys, reldeg, ref, **fb_kwargs) + super().__init__(sys, ref, v) def output(self, x, u, t=None): x, z = self._unpack_states(x) diff --git a/tests/test_linearize.py b/tests/test_linearize.py index d75d36e..7881d1f 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -24,7 +24,7 @@ class Allpass(AbstractControlAffine): - initial_state = jnp.zeros(1) + initial_state = jnp.zeros(0) n_inputs = "scalar" def f(self, x): @@ -69,6 +69,7 @@ def test_relative_degree(): sys = SpringMassDamperWithOutput(out=1) assert relative_degree(sys, xs) == 1 + xs = np.random.normal(size=100) assert relative_degree(Allpass(), xs) == 0 @@ -82,6 +83,7 @@ def test_discrete_relative_degree(): sys = SpringMassDamperWithOutput(out=1) assert discrete_relative_degree(sys, xs, us) == 1 + xs = np.random.normal(size=100) assert discrete_relative_degree(Allpass(), xs, us) == 0