PyDiffGame 0.1.2__py3-none-any.whl → 2.0.0__py3-none-any.whl

PyDiffGame/examples/MassesWithSpringsComparison.py

@@ -1,218 +1,119 @@
-import numpy as np
-from typing import Sequence, Optional, Union
-
-from PyDiffGame.PyDiffGame import PyDiffGame
-from PyDiffGame.PyDiffGameLQRComparison import PyDiffGameLQRComparison
-from PyDiffGame.Objective import GameObjective, LQRObjective
-
-
-class MassesWithSpringsComparison(PyDiffGameLQRComparison):
-    def __init__(self,
-                 N: int,
-                 m: float,
-                 k: float,
-                 q: Union[float, Sequence[float], Sequence[Sequence[float]]],
-                 r: float,
-                 Ms: Optional[Sequence[np.array]] = None,
-                 x_0: Optional[np.array] = None,
-                 x_T: Optional[np.array] = None,
-                 T_f: Optional[float] = None,
-                 epsilon_x: Optional[float] = PyDiffGame.epsilon_x_default,
-                 epsilon_P: Optional[float] = PyDiffGame.epsilon_P_default,
-                 L: Optional[int] = PyDiffGame.L_default,
-                 eta: Optional[int] = PyDiffGame.eta_default):
-        I_N = np.eye(N)
-        Z_N = np.zeros((N, N))
-
-        M_masses = m * I_N
-        K = k * (2 * I_N - np.array([[int(abs(i - j) == 1) for j in range(N)] for i in range(N)]))
-        M_masses_inv = np.linalg.inv(M_masses)
+"""Coupled masses-and-springs benchmark: differential game vs. LQR.
 
-        M_inv_K = M_masses_inv @ K
+``N`` unit-coupled masses connected by springs form a ``2N``-dimensional linear
+system (positions and velocities).  The physical input space is decomposed along
+the *modal* directions of ``M^{-1}K`` so that each vibration mode becomes one
+player of a differential game, and the result is compared against a single
+LQR controller acting on the whole system.
 
-        if Ms is None:
-            eigenvectors = np.linalg.eig(M_inv_K)[1]
-            Ms = [eigenvector.reshape(1, N) for eigenvector in eigenvectors]
+Run directly to solve a small instance and report the costs::
 
-        A = np.block([[Z_N, I_N],
-                      [-M_inv_K, Z_N]])
-        B = np.block([[Z_N],
-                      [M_masses_inv]])
+    python -m PyDiffGame.examples.MassesWithSpringsComparison
+"""
 
-        Qs = [np.diag([0.0] * i + [q] + [0.0] * (N - 1) + [q] + [0.0] * (N - i - 1))
-              if isinstance(q, (int, float)) else
-              np.diag([0.0] * i + [q[i]] + [0.0] * (N - 1) + [q[i]] + [0.0] * (N - i - 1)) for i in range(N)]
+from __future__ import annotations
 
-        M = np.concatenate(Ms,
-                           axis=0)
+from collections.abc import Sequence
 
-        assert np.all(np.abs(np.linalg.inv(M) - M.T) < 10e-12)
-
-        Q_mat = np.kron(a=np.eye(2),
-                        b=M)
-
-        Qs = [Q_mat.T @ Q @ Q_mat for Q in Qs]
-
-        Rs = [np.array([r])] * N
-        R_lqr = 1 / 4 * r * I_N
-        Q_lqr = q * np.eye(2 * N) if isinstance(q, (int, float)) else np.diag(2 * q)
+import numpy as np
 
-        state_variables_names = ['x_{' + str(i) + '}' for i in range(1, N + 1)] + \
-                                ['\\dot{x}_{' + str(i) + '}' for i in range(1, N + 1)]
-        args = {'A': A,
-                'B': B,
-                'x_0': x_0,
-                'x_T': x_T,
-                'T_f': T_f,
-                'state_variables_names': state_variables_names,
-                'epsilon_x': epsilon_x,
-                'epsilon_P': epsilon_P,
-                'L': L,
-                'eta': eta}
+from PyDiffGame.base import PyDiffGame
+from PyDiffGame.comparison import PyDiffGameLQRComparison
+from PyDiffGame.objective import GameObjective, LQRObjective
+from PyDiffGame.plotting import show
 
-        lqr_objective = [LQRObjective(Q=Q_lqr,
-                                      R_ii=R_lqr)]
-        game_objectives = [GameObjective(Q=Q,
-                                         R_ii=R,
-                                         M_i=M_i) for Q, R, M_i in zip(Qs, Rs, Ms)]
-        games_objectives = [lqr_objective,
-                            game_objectives]
 
-        self.figure_filename_generator = lambda g: ('LQR' if g.is_LQR() else f'{N}-players') + \
-                                                                             f'_large_{1 if q[0] > q[1] else 2}'
-
-        super().__init__(args=args,
-                         M=M,
-                         games_objectives=games_objectives,
-                         continuous=True)
-
-    @staticmethod
-    def are_all_orthogonal(vectors: Sequence[np.array]) -> bool:
-        """
-        Check if a set of vectors are all orthogonal to each other.
-
-        Parameters
-        ----------
-
-        vectors: sequence of vectors of len(n)
-            A sequence of numpy arrays representing the vectors
-
-        Returns
-        ----------
-
-        all_orthogonal: bool
-            True if all the vectors are orthogonal to each other, False otherwise
-        """
-
-        all_orthogonal = not any([any([abs(vectors[i] @ vectors[j].T) > 1e-10 for j in range(i + 1, len(vectors))])
-                                  for i in range(len(vectors))])
-
-        return all_orthogonal
-
-    @staticmethod
-    def are_all_unit_vectors(vectors: Sequence[np.array]) -> bool:
-        """
-        Check if a set of vectors are all of length 1.
-
-        Parameters
-        ----------
-
-        vectors: sequence of vectors of len(n)
-            A sequence of numpy arrays representing the vectors
-
-        Returns
-        ----------
-
-        all_orthogonal: bool
-            True if all the vectors are orthogonal to each other, False otherwise
-        """
-
-        all_unit_vectors = all([abs(np.linalg.norm(v) - 1) < 1e-10 for v in vectors])
-
-        return all_unit_vectors
-
-
-N2_output_variables_names = ['$\\frac{x_1 + x_2}{\\sqrt{2}}$',
-                             '$\\frac{x_2 - x_1}{\\sqrt{2}}$',
-                             '$\\frac{\\dot{x}_1 + \\dot{x}_2}{\\sqrt{2}}$',
-                             '$\\frac{\\dot{x}_2 - \\dot{x}_1}{\\sqrt{2}}$']
-
-
-def multiprocess_worker_function(N: int,
-                                 k: float,
-                                 m: float,
-                                 q: float,
-                                 r: float,
-                                 epsilon_x: float,
-                                 epsilon_P: float):
-    x_0 = np.array([10 * i for i in range(1, N + 1)] + [0] * N)
-    x_T = x_0 * 10 if N == 2 else np.array([(10 * i) ** 3 for i in range(1, N + 1)] + [0] * N)
-
-    output_variables_names = N2_output_variables_names \
-        if N == 2 else [f'$q_{i}$' for i in range(1, N + 1)] + ['$\\dot{q}_{' + str(i) + '}$' for i in range(1, N + 1)]
-
-    T_f = 25
-    masses_with_springs = MassesWithSpringsComparison(N=N,
-                                                      m=m,
-                                                      k=k,
-                                                      q=q,
-                                                      r=r,
-                                                      x_0=x_0,
-                                                      x_T=x_T,
-                                                      T_f=T_f,
-                                                      epsilon_x=epsilon_x,
-                                                      epsilon_P=epsilon_P)
-
-    masses_with_springs(plot_state_spaces=True,
-                        plot_Mx=True,
-                        output_variables_names=output_variables_names,
-                        save_figure=True,
-                        figure_filename=masses_with_springs.figure_filename_generator
-                        )
-
-
-if __name__ == '__main__':
-    d = 0.2
-    N = 8
-
-    Ns = [N]
-    ks = [10]
-    ms = [50]
-    rs = [1]
-    epsilon_xs = [10e-8]
-    epsilon_Ps = [10e-8]
-    qs = [[500, 2000], [500, 250]] if N == 2 else \
-        [[500 * (1 + d) ** i for i in range(N)], [500 * (1 - d) ** i for i in range(N)]]
-
-    # N, k, m, r, epsilon_x, epsilon_P = Ns[0], ks[0], ms[0], rs[0], epsilon_xs[0], epsilon_Ps[0]
-    # q = [500, 2000] if N == 2 else [500 * (1 + d) ** i for i in range(N)]
-    # x_0 = np.array([10 * i for i in range(1, N + 1)] + [0] * N)
-    # x_T = x_0 * 10 if N == 2 else np.array([(10 * i) ** 3 for i in range(1, N + 1)] + [0] * N)
-    #
-    # masses_with_springs = MassesWithSpringsComparison(N=N,
-    #                                                   m=m,
-    #                                                   k=k,
-    #                                                   q=q,
-    #                                                   r=r,
-    #                                                   x_0=x_0,
-    #                                                   x_T=x_T,
-    #                                                   T_f=25,
-    #                                                   epsilon_x=epsilon_x,
-    #                                                   epsilon_P=epsilon_P)
-    #
-    # output_variables_names = N2_output_variables_names \
-    #     if N == 2 else [f'$q_{i}$' for i in range(1, N + 1)] + ['$\\dot{q}_{' + str(i) + '}$' for i in range(1, N + 1)]
-    #
-    # masses_with_springs(plot_state_spaces=True,
-    #                     plot_Mx=True,
-    #                     output_variables_names=output_variables_names,
-    #                     save_figure=True,
-    #                     figure_filename=masses_with_springs.figure_filename_generator
-    #                     )
-
-    #
-
-    params = [Ns, ks, ms, qs, rs, epsilon_xs, epsilon_Ps]
-    PyDiffGameLQRComparison.run_multiprocess(multiprocess_worker_function=multiprocess_worker_function,
-                                             values=params)
+class MassesWithSpringsComparison(PyDiffGameLQRComparison):
+    """Compare a modal differential game with an LQR for ``N`` masses on springs."""
+
+    def __init__(
+        self,
+        N: int,
+        m: float,
+        k: float,
+        q: float | Sequence[float],
+        r: float,
+        x_0: np.ndarray | None = None,
+        x_T: np.ndarray | None = None,
+        T_f: float | None = None,
+        epsilon_x: float = PyDiffGame.epsilon_x_default,
+        epsilon_P: float = PyDiffGame.epsilon_P_default,
+        L: int = PyDiffGame.L_default,
+        eta: int = PyDiffGame.eta_default,
+    ) -> None:
+        I_N = np.eye(N)
+        Z_N = np.zeros((N, N))
 
+        mass_matrix = m * I_N
+        stiffness = k * (2 * I_N - np.array([[int(abs(i - j) == 1) for j in range(N)] for i in range(N)]))
+        mass_inv = np.linalg.inv(mass_matrix)
+        mass_inv_stiffness = mass_inv @ stiffness
+
+        # Modal directions: eigh gives orthonormal eigenvectors, so M is orthogonal.
+        _, eigenvectors = np.linalg.eigh(mass_inv_stiffness)
+        Ms = [eigenvectors[:, i].reshape(1, N) for i in range(N)]
+        M = np.concatenate(Ms, axis=0)
+        assert np.allclose(np.linalg.inv(M), M.T, atol=1e-9), "modal transform must be orthogonal"
+
+        A = np.block([[Z_N, I_N], [-mass_inv_stiffness, Z_N]])
+        B = np.block([[Z_N], [mass_inv]])
+
+        q_per_mode = [q] * N if isinstance(q, (int, float)) else list(q)
+        modal_to_state = np.kron(np.eye(2), M)
+        game_objectives = []
+        for i, (q_i, M_i) in enumerate(zip(q_per_mode, Ms)):
+            modal_weight = np.diag([0.0] * i + [q_i] + [0.0] * (N - 1) + [q_i] + [0.0] * (N - i - 1))
+            Q_i = modal_to_state.T @ modal_weight @ modal_to_state
+            game_objectives.append(GameObjective(Q=Q_i, R=r, M=M_i))
+
+        # The LQR baseline optimizes *exactly* the aggregate of the game's
+        # objectives, so it is the genuine monolithic optimum for a fair
+        # comparison: sum_i Q_i, and the matching physical input weight r * I
+        # (the per-player modal weight r maps to r * I in physical input space
+        # because the modal transform M is orthogonal). The game, being a
+        # decomposed design, can only match or exceed this optimal cost.
+        modal_weight_total = np.diag(q_per_mode + q_per_mode)
+        Q_lqr = modal_to_state.T @ modal_weight_total @ modal_to_state
+        lqr_objective = [LQRObjective(Q=Q_lqr, R=r * I_N)]
+
+        state_variables_names = [f"x_{{{i}}}" for i in range(1, N + 1)] + [
+            rf"\dot{{x}}_{{{i}}}" for i in range(1, N + 1)
+        ]
+
+        super().__init__(
+            A=A,
+            B=B,
+            games_objectives=[lqr_objective, game_objectives],
+            continuous=True,
+            x_0=x_0,
+            x_T=x_T,
+            T_f=T_f,
+            L=L,
+            eta=eta,
+            epsilon_x=epsilon_x,
+            epsilon_P=epsilon_P,
+            state_variables_names=state_variables_names,
+        )
+
+
+def main(N: int = 2, *, plot: bool = True, save_figure: bool = False) -> None:
+    """Solve a single masses-and-springs comparison and print the costs."""
+
+    m, k, r = 50.0, 10.0, 1.0
+    q = [500.0, 2000.0] if N == 2 else [500.0 * 1.2**i for i in range(N)]
+    x_0 = np.array([10.0 * i for i in range(1, N + 1)] + [0.0] * N)
+    x_T = x_0 * 10.0
+
+    comparison = MassesWithSpringsComparison(N=N, m=m, k=k, q=q, r=r, x_0=x_0, x_T=x_T, T_f=25.0, L=300)
+    comparison.run(plot_state_spaces=plot, save_figure=save_figure)
+
+    lqr_cost, game_cost = comparison.costs()
+    print(f"LQR cost   = {lqr_cost:.4g}")
+    print(f"Game cost  = {game_cost:.4g}")
+    print(f"All games controllable: {comparison.are_all_controllable()}")
+    if plot:
+        show()
+
+
+if __name__ == "__main__":
+    main()

PyDiffGame/examples/PVTOL.py

@@ -2,55 +2,61 @@
 # RMM, 14 Jan 03
 #
 # This file works through an LQR based design problem, using the
-# planar vertical takeoff and landing (PVTOLComparison.py) aircraft example from
+# planar vertical takeoff and landing (PVTOL) aircraft example from
 # Astrom and Murray, Chapter 5.  It is intended to demonstrate the
 # basic functionality of the python-control package.
 #
 
-import os
-import numpy as np
-import matplotlib.pyplot as plt  # MATLAB-like plotting functions
+from __future__ import annotations
+
 import control as ct
+import matplotlib.pyplot as plt  # MATLAB-like plotting functions
+import numpy as np
+
+from PyDiffGame.plotting import show
 
 #
 # System dynamics
 #
-# These are the dynamics for the PVTOLComparison.py system, written in state space
+# These are the dynamics for the PVTOL system, written in state space
 # form.
 #
 
 # System parameters
-m = 4       # mass of aircraft
+m = 4  # mass of aircraft
 J = 0.0475  # inertia around pitch axis
-r = 0.25    # distance to center of force
-g = 9.8     # gravitational constant
-c = 0.05    # damping factor (estimated)
+r = 0.25  # distance to center of force
+g = 9.8  # gravitational constant
+c = 0.05  # damping factor (estimated)
 
 # State space dynamics
 xe = [0, 0, 0, 0, 0, 0]  # equilibrium point of interest
 ue = [0, m * g]  # (note these are lists, not matrices)
 
-# This will involve re-working the subsequent equations as the shapes
-# See below.
-
 A = np.array(
-    [[0, 0, 0, 1, 0, 0],
-     [0, 0, 0, 0, 1, 0],
-     [0, 0, 0, 0, 0, 1],
-     [0, 0, (-ue[0]*np.sin(xe[2]) - ue[1]*np.cos(xe[2]))/m, -c/m, 0, 0],
-     [0, 0, (ue[0]*np.cos(xe[2]) - ue[1]*np.sin(xe[2]))/m, 0, -c/m, 0],
-     [0, 0, 0, 0, 0, 0]]
+    [
+        [0, 0, 0, 1, 0, 0],
+        [0, 0, 0, 0, 1, 0],
+        [0, 0, 0, 0, 0, 1],
+        [0, 0, (-ue[0] * np.sin(xe[2]) - ue[1] * np.cos(xe[2])) / m, -c / m, 0, 0],
+        [0, 0, (ue[0] * np.cos(xe[2]) - ue[1] * np.sin(xe[2])) / m, 0, -c / m, 0],
+        [0, 0, 0, 0, 0, 0],
+    ]
 )
 
 # Input matrix
 B = np.array(
-    [[0, 0], [0, 0], [0, 0],
-     [np.cos(xe[2])/m, -np.sin(xe[2])/m],
-     [np.sin(xe[2])/m, np.cos(xe[2])/m],
-     [r/J, 0]]
+    [
+        [0, 0],
+        [0, 0],
+        [0, 0],
+        [np.cos(xe[2]) / m, -np.sin(xe[2]) / m],
+        [np.sin(xe[2]) / m, np.cos(xe[2]) / m],
+        [r / J, 0],
+    ]
 )
 
-# Output matrix 
+# Output matrix
 C = np.array([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0]])
 D = np.array([[0, 0], [0, 0]])
 
@@ -58,7 +64,7 @@ D = np.array([[0, 0], [0, 0]])
 # Construct inputs and outputs corresponding to steps in xy position
 #
 # The vectors xd and yd correspond to the states that are the desired
-# equilibrium states for the system.  The matrices Cx and Cy are the 
+# equilibrium states for the system.  The matrices Cx and Cy are the
 # corresponding outputs.
 #
 # The way these vectors are used is to compute the closed loop system
@@ -67,7 +73,7 @@ D = np.array([[0, 0], [0, 0]])
 #   xdot = Ax + B u  =>  xdot = (A-BK)x + K xd
 #      u = -K(x - xd)       y = Cx
 #
-# The closed loop dynamics can be simulated using the "step" command, 
+# The closed loop dynamics can be simulated using the "step" command,
 # with K*xd as the input vector (assumes that the "input" is unit size,
 # so that xd corresponds to the desired steady state.
 #
@@ -78,11 +84,8 @@ yd = np.array([[0], [1], [0], [0], [0], [0]])
 #
 # Extract the relevant dynamics for use with SISO library
 #
-# The current python-control library only supports SISO transfer
-# functions, so we have to modify some parts of the original MATLAB
-# code to extract out SISO systems.  To do this, we define the 'lat' and
-# 'alt' index vectors to consist of the states that are are relevant
-# to the lateral (x) and vertical (y) dynamics.
+# We define the 'lat' and 'alt' index vectors to consist of the states
+# that are relevant to the lateral (x) and vertical (y) dynamics.
 #
 
 # Indices for the parts of the state that we want
@@ -100,123 +103,114 @@ By = B[np.ix_(alt, [1])]
 Cy = C[np.ix_([1], alt)]
 Dy = D[np.ix_([1], [1])]
 
-# Label the plot
-plt.clf()
-plt.suptitle("LQR controllers for vectored thrust aircraft (pvtol-lqr)")
-
-#
-# LQR design
-#
-
-# Start with a diagonal weighting
-Qx1 = np.diag([1, 1, 1, 1, 1, 1])
-Qu1a = np.diag([1, 1])
-K1a, X, E = ct.lqr(A, B, Qx1, Qu1a)
-
-# Close the loop: xdot = Ax - B K (x-xd)
-#
-# Note: python-control requires we do this 1 input at a time
-# H1a = ss(A-B*K1a, B*K1a*concatenate((xd, yd), axis=1), C, D);
-# (T, Y) = step_response(H1a, T=np.linspace(0,10,100));
-#
-
-# Step response for the first input
-H1ax = ct.ss(Ax - Bx @ K1a[np.ix_([0], lat)],
-             Bx @ K1a[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-Tx, Yx = ct.step_response(H1ax, T=np.linspace(0, 10, 100))
-
-# Step response for the second input
-H1ay = ct.ss(Ay - By @ K1a[np.ix_([1], alt)],
-             By @ K1a[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
-Ty, Yy = ct.step_response(H1ay, T=np.linspace(0, 10, 100))
-
-plt.subplot(221)
-plt.title("Identity weights")
-# plt.plot(T, Y[:,1, 1], '-', T, Y[:,2, 2], '--')
-plt.plot(Tx.T, Yx.T, '-', Ty.T, Yy.T, '--')
-plt.plot([0, 10], [1, 1], 'k-')
-
-plt.axis([0, 10, -0.1, 1.4])
-plt.ylabel('position')
-plt.legend(('x', 'y'), loc='lower right')
-
-# Look at different input weightings
-Qu1a = np.diag([1, 1])
-K1a, X, E = ct.lqr(A, B, Qx1, Qu1a)
-H1ax = ct.ss(Ax - Bx @ K1a[np.ix_([0], lat)],
-             Bx @ K1a[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-
-Qu1b = (40 ** 2)*np.diag([1, 1])
-K1b, X, E = ct.lqr(A, B, Qx1, Qu1b)
-H1bx = ct.ss(Ax - Bx @ K1b[np.ix_([0], lat)],
-             Bx @ K1b[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-
-Qu1c = (200 ** 2)*np.diag([1, 1])
-K1c, X, E = ct.lqr(A, B, Qx1, Qu1c)
-H1cx = ct.ss(Ax - Bx @ K1c[np.ix_([0], lat)],
-             Bx @ K1c[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-
-T1, Y1 = ct.step_response(H1ax, T=np.linspace(0, 10, 100))
-T2, Y2 = ct.step_response(H1bx, T=np.linspace(0, 10, 100))
-T3, Y3 = ct.step_response(H1cx, T=np.linspace(0, 10, 100))
-
-plt.subplot(222)
-plt.title("Effect of input weights")
-plt.plot(T1.T, Y1.T, 'b-')
-plt.plot(T2.T, Y2.T, 'b-')
-plt.plot(T3.T, Y3.T, 'b-')
-plt.plot([0, 10], [1, 1], 'k-')
-
-plt.axis([0, 10, -0.1, 1.4])
-
-# arcarrow([1.3, 0.8], [5, 0.45], -6)
-plt.text(5.3, 0.4, r'$\rho$')
-
-# Output weighting - change Qx to use outputs
-Qx2 = C.T @ C
-Qu2 = 0.1 * np.diag([1, 1])
-K2, X, E = ct.lqr(A, B, Qx2, Qu2)
-
-H2x = ct.ss(Ax - Bx @ K2[np.ix_([0], lat)],
-            Bx @ K2[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-H2y = ct.ss(Ay - By @ K2[np.ix_([1], alt)],
-            By @ K2[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
-
-plt.subplot(223)
-plt.title("Output weighting")
-T2x, Y2x = ct.step_response(H2x, T=np.linspace(0, 10, 100))
-T2y, Y2y = ct.step_response(H2y, T=np.linspace(0, 10, 100))
-plt.plot(T2x.T, Y2x.T, T2y.T, Y2y.T)
-plt.ylabel('position')
-plt.xlabel('time')
-plt.ylabel('position')
-plt.legend(('x', 'y'), loc='lower right')
-
-#
-# Physically motivated weighting
-#
-# Shoot for 1 cm error in x, 10 cm error in y.  Try to keep the angle
-# less than 5 degrees in making the adjustments.  Penalize side forces
-# due to loss in efficiency.
-#
 
-Qx3 = np.diag([100, 10, 2*np.pi/5, 0, 0, 0])
-Qu3 = 0.1*np.diag([1, 10])
-K3, X, E = ct.lqr(A, B, Qx3, Qu3)
-
-H3x = ct.ss(Ax - Bx @ K3[np.ix_([0], lat)],
-            Bx @ K3[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
-H3y = ct.ss(Ay - By @ K3[np.ix_([1], alt)],
-            By @ K3[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
-plt.subplot(224)
-# step_response(H3x, H3y, 10)
-T3x, Y3x = ct.step_response(H3x, T=np.linspace(0, 10, 100))
-T3y, Y3y = ct.step_response(H3y, T=np.linspace(0, 10, 100))
-plt.plot(T3x.T, Y3x.T, T3y.T, Y3y.T)
-plt.title("Physically motivated weights")
-plt.xlabel('time')
-plt.legend(('x', 'y'), loc='lower right')
-plt.tight_layout()
-
-if 'PYCONTROL_TEST_EXAMPLES' not in os.environ:
-    plt.show()
+def main() -> None:
+    """Run the LQR design study and draw the comparison plots."""
+
+    # Label the plot
+    plt.clf()
+    plt.suptitle("LQR controllers for vectored thrust aircraft (pvtol-lqr)")
+
+    #
+    # LQR design
+    #
+
+    # Start with a diagonal weighting
+    Qx1 = np.diag([1, 1, 1, 1, 1, 1])
+    Qu1a = np.diag([1, 1])
+    K1a, _X, _E = ct.lqr(A, B, Qx1, Qu1a)
+
+    # Close the loop: xdot = Ax - B K (x-xd)
+    #
+    # Note: python-control requires we do this 1 input at a time.
+
+    # Step response for the first input
+    H1ax = ct.ss(Ax - Bx @ K1a[np.ix_([0], lat)], Bx @ K1a[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+    Tx, Yx = ct.step_response(H1ax, T=np.linspace(0, 10, 100))
+
+    # Step response for the second input
+    H1ay = ct.ss(Ay - By @ K1a[np.ix_([1], alt)], By @ K1a[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
+    Ty, Yy = ct.step_response(H1ay, T=np.linspace(0, 10, 100))
+
+    plt.subplot(221)
+    plt.title("Identity weights")
+    plt.plot(Tx.T, Yx.T, "-", Ty.T, Yy.T, "--")
+    plt.plot([0, 10], [1, 1], "k-")
+
+    plt.axis([0, 10, -0.1, 1.4])
+    plt.ylabel("position")
+    plt.legend(("x", "y"), loc="lower right")
+
+    # Look at different input weightings
+    Qu1a = np.diag([1, 1])
+    K1a, _X, _E = ct.lqr(A, B, Qx1, Qu1a)
+    H1ax = ct.ss(Ax - Bx @ K1a[np.ix_([0], lat)], Bx @ K1a[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+
+    Qu1b = (40**2) * np.diag([1, 1])
+    K1b, _X, _E = ct.lqr(A, B, Qx1, Qu1b)
+    H1bx = ct.ss(Ax - Bx @ K1b[np.ix_([0], lat)], Bx @ K1b[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+
+    Qu1c = (200**2) * np.diag([1, 1])
+    K1c, _X, _E = ct.lqr(A, B, Qx1, Qu1c)
+    H1cx = ct.ss(Ax - Bx @ K1c[np.ix_([0], lat)], Bx @ K1c[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+
+    T1, Y1 = ct.step_response(H1ax, T=np.linspace(0, 10, 100))
+    T2, Y2 = ct.step_response(H1bx, T=np.linspace(0, 10, 100))
+    T3, Y3 = ct.step_response(H1cx, T=np.linspace(0, 10, 100))
+
+    plt.subplot(222)
+    plt.title("Effect of input weights")
+    plt.plot(T1.T, Y1.T, "b-")
+    plt.plot(T2.T, Y2.T, "b-")
+    plt.plot(T3.T, Y3.T, "b-")
+    plt.plot([0, 10], [1, 1], "k-")
+
+    plt.axis([0, 10, -0.1, 1.4])
+
+    plt.text(5.3, 0.4, r"$\rho$")
+
+    # Output weighting - change Qx to use outputs
+    Qx2 = C.T @ C
+    Qu2 = 0.1 * np.diag([1, 1])
+    K2, _X, _E = ct.lqr(A, B, Qx2, Qu2)
+
+    H2x = ct.ss(Ax - Bx @ K2[np.ix_([0], lat)], Bx @ K2[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+    H2y = ct.ss(Ay - By @ K2[np.ix_([1], alt)], By @ K2[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
+
+    plt.subplot(223)
+    plt.title("Output weighting")
+    T2x, Y2x = ct.step_response(H2x, T=np.linspace(0, 10, 100))
+    T2y, Y2y = ct.step_response(H2y, T=np.linspace(0, 10, 100))
+    plt.plot(T2x.T, Y2x.T, T2y.T, Y2y.T)
+    plt.xlabel("time")
+    plt.ylabel("position")
+    plt.legend(("x", "y"), loc="lower right")
+
+    #
+    # Physically motivated weighting
+    #
+    # Shoot for 1 cm error in x, 10 cm error in y.  Try to keep the angle
+    # less than 5 degrees in making the adjustments.  Penalize side forces
+    # due to loss in efficiency.
+    #
+
+    Qx3 = np.diag([100, 10, 2 * np.pi / 5, 0, 0, 0])
+    Qu3 = 0.1 * np.diag([1, 10])
+    K3, _X, _E = ct.lqr(A, B, Qx3, Qu3)
+
+    H3x = ct.ss(Ax - Bx @ K3[np.ix_([0], lat)], Bx @ K3[np.ix_([0], lat)] @ xd[lat, :], Cx, Dx)
+    H3y = ct.ss(Ay - By @ K3[np.ix_([1], alt)], By @ K3[np.ix_([1], alt)] @ yd[alt, :], Cy, Dy)
+    plt.subplot(224)
+    T3x, Y3x = ct.step_response(H3x, T=np.linspace(0, 10, 100))
+    T3y, Y3y = ct.step_response(H3y, T=np.linspace(0, 10, 100))
+    plt.plot(T3x.T, Y3x.T, T3y.T, Y3y.T)
+    plt.title("Physically motivated weights")
+    plt.xlabel("time")
+    plt.legend(("x", "y"), loc="lower right")
+    plt.tight_layout()
+
+    show()
+
+
+if __name__ == "__main__":
+    main()