From dc2fe7974845f7823ec0fe2981cfff98311ef8bf Mon Sep 17 00:00:00 2001 From: fhchl Date: Sun, 3 Mar 2024 21:48:34 +0100 Subject: [PATCH] refac: make LinearSystem sublass of ControlAffine --- dynax/derivative.py | 1 - dynax/linearize.py | 50 +++++++++------- dynax/system.py | 129 ++++++++++++++++++++-------------------- tests/test_linearize.py | 2 +- 4 files changed, 94 insertions(+), 88 deletions(-) diff --git a/dynax/derivative.py b/dynax/derivative.py index 3c333c2..f7f796b 100644 --- a/dynax/derivative.py +++ b/dynax/derivative.py @@ -9,7 +9,6 @@ from jax.experimental.jet import jet -@lru_cache def lie_derivative(f, h, n=1): r"""Return n-th directional derivative of h along f. diff --git a/dynax/linearize.py b/dynax/linearize.py index e9958d4..c377fa6 100644 --- a/dynax/linearize.py +++ b/dynax/linearize.py @@ -21,24 +21,25 @@ # TODO: make this a method of ControlAffine -def relative_degree( - sys: ControlAffine, xs, max_reldeg=10, output: Optional[int] = None -) -> int: +def relative_degree(sys: ControlAffine, xs: Array, output: Optional[int] = None) -> int: """Estimate relative degree of system on region xs.""" - # TODO: when ControlAffine has y = h(x) + i(x)u, include test for n = 0, - # i.e. i(x) == 0 for all x in xs. - assert sys.n_inputs in ["scalar", 1] + 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 - msg = f"Output is None, but system has {sys.n_outputs} outputs." - assert sys.n_outputs in ["scalar", 1], msg + if sys.n_outputs not in ["scalar", 1]: + raise ValueError(f"Output is None, but system has {sys.n_outputs} outputs.") h = sys.h else: h = lambda *args, **kwargs: sys.h(*args, **kwargs)[output] - for n in range(1, max_reldeg + 1): - LgLfn1h = lie_derivative(sys.g, lie_derivative(sys.f, h, n - 1)) - res = jax.vmap(LgLfn1h)(xs) + max_reldeg = jnp.size(sys.initial_state) + for n in range(0, max_reldeg + 1): + if n == 0: + res = jax.vmap(sys.i)(xs) + 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): @@ -62,7 +63,7 @@ def input_output_linearize( output: Optional[int] = None, asymptotic: Optional[Sequence] = None, reg: Optional[float] = None, -) -> Callable[[Array, Array, float], float]: +) -> Callable[[Array, Array, float], Array]: """Construct input-output linearizing feedback law. Args: @@ -80,7 +81,7 @@ def input_output_linearize( Only single-input-single-output systems are currently supported. """ - assert sys.n_inputs == ref.n_inputs, "systems habe same input dimension" + assert sys.n_inputs == ref.n_inputs, "systems have same input dimension" assert sys.n_inputs in [1, "scalar"] if output is None: @@ -99,24 +100,28 @@ def input_output_linearize( if asymptotic is None: - def feedbacklaw(x: Array, z: Array, v: float) -> float: + def feedbacklaw(x: Array, z: Array, v: float) -> Array: y_reldeg_ref = cAn.dot(z) + cAnm1b * v y_reldeg = Lfnh(x) - return (y_reldeg_ref - y_reldeg) / LgLfnm1h(x) + out = (y_reldeg_ref - y_reldeg) / LgLfnm1h(x) + return out if sys.n_inputs != "scalar" else out.squeeze() else: - msg = f"asymptotic must be of length {reldeg=} but, {len(asymptotic)=}" - assert len(asymptotic) == reldeg, msg + if len(asymptotic) != reldeg: + raise ValueError( + f"asymptotic must be of length {reldeg=} but, {len(asymptotic)=}" + ) coeffs = np.concatenate(([1], asymptotic)) - msg = "Polynomial must be Hurwitz" - assert np.all(np.real(np.roots(coeffs)) <= 0) + if not np.all(np.real(np.roots(coeffs)) <= 0): + raise ValueError("Polynomial must be Hurwitz") + alphas = asymptotic 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) -> float: + def feedbacklaw(x: Array, z: Array, v: float) -> Array: y_reldeg_ref = cAn.dot(z) + cAnm1b * v y_reldeg = Lfnh(x) ae0s = jnp.array( @@ -127,10 +132,11 @@ def feedbacklaw(x: Array, z: Array, v: float) -> float: ) error = y_reldeg_ref - y_reldeg + jnp.sum(ae0s) if reg is None: - return error / LgLfnm1h(x) + out = error / LgLfnm1h(x) else: l = LgLfnm1h(x) - return error * l / (l + reg) + out = error * l / (l + reg) + return out if sys.n_inputs != "scalar" else out.squeeze() return feedbacklaw diff --git a/dynax/system.py b/dynax/system.py index 89d1451..4816f13 100644 --- a/dynax/system.py +++ b/dynax/system.py @@ -75,13 +75,13 @@ class DynamicalSystem(eqx.Module): ẋ &= f(x, u, t) \\ y &= h(x, u, t) - Subclasses must set values for attributes n_states, n_inputs, and implement the - `vector_field` method. Use the optional `output` method to describe measurent + Subclasses must set values for attributes n_states, n_inputs, and implement the + `vector_field` method. Use the optional `output` method to describe measurent equations. Otherwise, the total state is returned as output. In most cases, it is not needed to define a custom __init__ method, as `DynamicalSystem` is a dataclass. - + Example:: class IntegratorAndGain(DynamicalSystem): @@ -108,7 +108,7 @@ def __check_init__(self): if not hasattr(self, attr): raise AttributeError(f"Attribute '{attr}' not initialized.") - # Check that vector_field returns Arrays or scalars and not PyTrees + # Check that vector_field and output returns Arrays or scalars and not PyTrees x = self.initial_state u = jax.ShapeDtypeStruct(dim2shape(self.n_inputs), jnp.float64) try: @@ -118,12 +118,11 @@ def __check_init__(self): raise ValueError( "Can not evaluate output shapes. Check your definitions!" ) from e - if not isinstance(dx, jax.ShapeDtypeStruct): - raise ValueError( - f"vector_field must return arrays or scalars, not {type(dx)}" - ) - if not isinstance(y, jax.ShapeDtypeStruct): - raise ValueError(f"outpuut must return arrays or scalars, not {type(y)}") + for val, func in zip((dx, y), ("vector_field, output")): # noqa: B905 + if not isinstance(val, jax.ShapeDtypeStruct): + raise ValueError( + f"{func} must return arrays or scalars, not {type(val)}" + ) @abstractmethod def vector_field(self, x, u=None, t=None) -> Array: @@ -212,7 +211,44 @@ def linearize(self, x0=None, u0=None, t=None) -> "LinearSystem": # (x, pytree_interal_states_x) -class LinearSystem(DynamicalSystem): +class ControlAffine(DynamicalSystem): + r"""A control-affine dynamical system. + + .. math:: + + ẋ &= f(x) + g(x)u \\ + y &= h(x) + i(x)u + + """ + + @abstractmethod + def f(self, x: Array) -> Array: + pass + + @abstractmethod + def g(self, x: Array) -> Array: + pass + + def h(self, x: Array) -> Array: + return x + + def i(self, x: Array) -> Array: + return jnp.array(0.0) + + def vector_field(self, x, u=None, t=None): + out = self.f(x) + if u is not None: + out += self.g(x).dot(u) + return out + + def output(self, x, u=None, t=None): + out = self.h(x) + if u is not None: + out += self.i(x).dot(u) + return out + + +class LinearSystem(ControlAffine): r"""A linear, time-invariant dynamical system. .. math:: @@ -222,10 +258,6 @@ class LinearSystem(DynamicalSystem): """ - # TODO: could be subclass of control-affine? Two blocking problems: - # - may h depend on u? Needed for D. If so, then one could compute - # relative degree. - A: Array B: Array C: Array @@ -255,48 +287,17 @@ def n_inputs(self) -> int | Literal["scalar"]: # type: ignore return self.B.shape[1] raise ValueError("Dimension mismatch.") - def vector_field(self, x, u=None, t=None) -> Array: - out = self.A.dot(x) - if u is not None: - out += self.B.dot(u) - return out - - def output(self, x, u=None, t=None) -> Array: - out = self.C.dot(x) - if u is not None: - out += self.D.dot(u) - return out - - -# TODO: make abstract - - -class ControlAffine(DynamicalSystem): - r"""A control-affine dynamical system. - - .. math:: - - ẋ &= f(x) + g(x)u \\ - y &= h(x) - - """ - def f(self, x): - raise NotImplementedError + return self.A.dot(x) def g(self, x): - raise NotImplementedError + return self.B def h(self, x): - return x + return self.C.dot(x) - def vector_field(self, x, u=None, t=None): - if u is None: - u = 0 - return self.f(x) + self.g(x) * u - - def output(self, x, u=None, t=None): - return self.h(x) + def i(self, x): + return self.D class _CoupledSystemMixin(eqx.Module): @@ -334,7 +335,7 @@ class SeriesSystem(DynamicalSystem, _CoupledSystemMixin): y_2 &= h_2(x_2, y1, t) .. aafig:: - + +------+ +------+ u --+->+ sys1 +--y1->+ sys2 +--> y2 +------+ +------+ @@ -369,7 +370,7 @@ def output(self, x, u=None, t=None): class FeedbackSystem(DynamicalSystem, _CoupledSystemMixin): r"""Two systems connected via feedback. - + .. math:: ẋ_1 &= f_1(x_1, u + y_2, t) \\ @@ -378,7 +379,7 @@ class FeedbackSystem(DynamicalSystem, _CoupledSystemMixin): y_2 &= h_2(x_2, y_1, t) \\ .. aafig:: - + +------+ u --+->+ sys1 +--+-> y1 ^ +------+ | @@ -427,12 +428,12 @@ class StaticStateFeedbackSystem(DynamicalSystem): y &= h(x, u, t) .. aafig:: - + +-----+ u --+------------->+ sys +----> y - ^ +--+--+ - | | - | | x + ^ +--+--+ + | | + | | x | +--------+ | +--+ "v(x)" +<----+ +--------+ @@ -468,22 +469,22 @@ class DynamicStateFeedbackSystem(DynamicalSystem, _CoupledSystemMixin): r"""System with dynamic state-feedback. .. math:: - + ẋ_1 &= f_1(x_1, v(x_1, x_2, u), t) \\ ẋ_2 &= f_2(x_2, u, t) \\ y &= h_1(x_1, u, t) .. aafig:: - + +--------------+ +-----+ u -+->+ v(x1, x2, u) +--v->+ sys +-> y - | +-+-------+----+ +--+--+ + | +-+-------+----+ +--+--+ | ^ ^ | | | x2 | x1 | | | +-------------+ - | +------+ - +->+ sys2 | - +------+ + | +------+ + +->+ sys2 | + +------+ """ diff --git a/tests/test_linearize.py b/tests/test_linearize.py index df146be..5b69b50 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -55,7 +55,7 @@ def test_relative_degree(): def test_discrete_relative_degree(): xs = np.random.normal(size=(100, 2)) - us = np.random.normal(size=(100, 1)) + us = np.random.normal(size=(100)) sys = SpringMassDamperWithOutput(out=0) assert discrete_relative_degree(sys, xs, us) == 2