Skip to content

Commit

Permalink
improve docs for linearize.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@fhchl
fhchl committed Apr 21, 2024
1 parent 915b6c9 commit 6d1fe22
Show file tree
Hide file tree
Showing 3 changed files with 115 additions and 52 deletions.
2 changes: 2 additions & 0 deletions docs/source/conf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}

Expand Down
161 changes: 110 additions & 51 deletions dynax/linearize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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]:
Expand Down Expand Up @@ -70,30 +71,32 @@ 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,
ref: LinearSystem,
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"
Expand All @@ -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)
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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.}`.
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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]
Expand All @@ -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`.
"""

Expand All @@ -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)
Expand Down
4 changes: 3 additions & 1 deletion tests/test_linearize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@


class Allpass(AbstractControlAffine):
initial_state = jnp.zeros(1)
initial_state = jnp.zeros(0)
n_inputs = "scalar"

def f(self, x):
Expand Down Expand Up @@ -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


Expand All @@ -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


Expand Down

0 comments on commit 6d1fe22

Please sign in to comment.