import numpy                as np
import matplotlib.pyplot    as plt

gamma = 0.02
v0    = 10.0

def system_dynamics(t, x, u):
    return np.array([x[1], u - gamma * x[1]**2])

def simulate_no_control(t_span, x0):
    t_values = []
    x_values = []

    x     = x0
    x_dot = np.zeros_like(x)

    dt    = 0.01
    t     = t_span[0]

    while t <= t_span[1]:

        x_dot = system_dynamics(t, x, u=0)

        x = x + x_dot * dt

        t_values.append(t)
        x_values.append(x)

        t += dt

    return np.array(t_values), np.array(x_values)

fig = plt.figure(figsize=(8, 6), dpi=100)
ax1 = fig.add_subplot(111)

t_values, x_values = simulate_no_control(t_span=[0, 10], x0=[0, 10])

ax1.plot(t_values, x_values[:, 0], label='Position')
ax1.plot(t_values, x_values[:, 1], label='Velocity')
ax1.set_xlabel('Time')
ax1.set_ylabel('State')
ax1.legend()
ax1.grid()
plt.show()

!pip install control
import control              as ct

Ki_frac_Kp = 4

num = np.array([1, Ki_frac_Kp])
den = np.array([1, 2*gamma*v0, 0])
G = ct.tf(num, den)

fig = plt.figure(figsize=(9,6), dpi=100)
ax1 = fig.add_subplot(111)
ct.rlocus(G, ax=ax1, \
          title=r"Root locus for $H(s)=G_\tau(s)G_\theta(s)$")

ax1.set_title(rf"$K_i$/$K_p$={Ki_frac_Kp}")
ax1.set_xlabel(r"Re$(\omega)$")
ax1.set_ylabel(r"Im$(\omega)$")
ax1.set_xticks(np.arange(-50, 11, 5))
ax1.set_yticks(np.arange(-30, 31, 5))
ax1.set_xlim([-15, 5])
ax1.set_ylim([-10, 10])
ax1.hlines(0, -15, 5, color="k", linewidth=0.5)
ax1.vlines(0, -10, 10, color="k", linewidth=0.5)
ax1.set_box_aspect(1)

plt.show()

Kp = 8

fig = plt.figure(figsize=(9,6), dpi=100)
ax1 = fig.add_subplot(111)
ct.rlocus(G, ax=ax1, \
          title=r"Root locus for $H(s)=G_\tau(s)G_\theta(s)$")

ax1.set_title(rf"$K_i$/$K_p$={Ki_frac_Kp}, $K_p$={Kp}")
ax1.set_xlabel(r"Re$(\omega)$")
ax1.set_ylabel(r"Im$(\omega)$")
ax1.set_xticks(np.arange(-50, 11, 5))
ax1.set_yticks(np.arange(-30, 31, 5))
ax1.set_xlim([-15, 5])
ax1.set_ylim([-10, 10])
ax1.hlines(0, -15, 5, color="k", linewidth=0.5)
ax1.vlines(0, -10, 10, color="k", linewidth=0.5)
ax1.set_box_aspect(1)

poles, _ = ct.rlocus(G, gains=Kp, plot=False)
ax1.scatter(poles.real, poles.imag, color="r", marker="x", zorder=10)

plt.show()

from re import U
Ki = Kp * Ki_frac_Kp

def control_input(x, x_target, var_cum, Kp, Ki):
    error = x_target - x
    var_cum += error * 0.01
    u = Kp * error + Ki * var_cum
    return u, var_cum

def simulate_PI_control(t_span, x_target, x0, Kp, Ki):
    t_values = []
    x_values = []
    u_values = []

    x     = x0
    x_dot = np.zeros_like(x)
    var_cum = 0

    dt    = 0.01
    t     = t_span[0]

    while t <= t_span[1]:

        u, var_cum = control_input(x[1], x_target[1], var_cum, Kp, Ki)

        x_dot = system_dynamics(t, x, u)

        x = x + x_dot * dt

        t_values.append(t)
        x_values.append(x)
        u_values.append(u)

        t += dt

    return np.array(t_values), np.array(x_values), np.array(u_values)

t_values, x_values, u_values = simulate_PI_control( \
          t_span=[0, 3], x_target=[0,10], x0=[0, 8], Kp=Kp, Ki=Ki)

fig = plt.figure(figsize=(9, 9), dpi=100)
ax1 = fig.add_subplot(211)
ax1.plot(t_values, x_values[:, 1], label='Velocity')
ax1.hlines(10, 0, 3, linestyle='--', label='Target velocity')
ax1.set_xlabel('time [s]')
ax1.set_ylabel('velocity [m/s]')
ax1.set_xlim([0,3])
ax1.grid()
ax1.legend()
ax2 = fig.add_subplot(212)
ax2.plot(t_values, u_values, label='Control')
ax2.set_xlabel('time [s]')
ax2.set_ylabel(r'acceleration command [m/s$^2$]')
ax2.set_xlim([0,3])
ax2.grid()
ax2.legend()
plt.show()