From 2403bb18e6778885a653babea285c8cc1593afce Mon Sep 17 00:00:00 2001
From: Tycho Bogdanowitsch <tycho@stanford.edu>
Date: Sat, 9 Nov 2024 13:50:15 -0800
Subject: [PATCH] Implemented LQR

---
 simwise/equinoctial.py |  2 +-
 simwise/quaternion.py  | 43 ++++++++++++++++++++++++++++++++++++++++++
 simwise/simulation.py  | 18 ++++++++++++------
 3 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/simwise/equinoctial.py b/simwise/equinoctial.py
index ff19514..d9eba24 100644
--- a/simwise/equinoctial.py
+++ b/simwise/equinoctial.py
@@ -2,7 +2,7 @@
 
 def coe2mee(elements):
     """Transform classic (keplerian) orbital elements to modified equinoctial elements.
-
+sdajkshdkjsah
     Args:
         elements (np.ndarray): orbital elements of form [a, e, i, Ω, ω, θ]
     """
diff --git a/simwise/quaternion.py b/simwise/quaternion.py
index 076f5d8..b26102e 100644
--- a/simwise/quaternion.py
+++ b/simwise/quaternion.py
@@ -1,4 +1,5 @@
 import numpy as np
+from scipy.linalg import solve_continuous_are as ricatti_solver
 
 def normalize_quaternion(q):
     """Normalize a quaternion and force scalar positive.
@@ -166,3 +167,45 @@ def compute_control_torque(x, x_desired, K_p=1, K_d=1, tau_max=None):
             tau = tau_max * tau / np.linalg.norm(tau) 
 
     return tau
+def compute_control_torque_lqr(x,x_desired,Q=np.diag([1,1,1,1,1,1,1]),R=np.diag([1,1,1]),I=np.array([0.01461922201, 0.0412768466, 0.03235309961]), tau_max=None):
+    q = normalize_quaternion(x[:4]) # quaternion
+    q_d = normalize_quaternion(x_desired[:4])
+    q = q_d
+    q_0 = q[0]
+    q_1 = q[1]
+    q_2 = q[2]
+    q_3 = q[3]
+    
+    I_1 = I[0]
+    I_2 = I[1]
+    I_3 = I[2]
+
+    w_d = x_desired[4:]
+    w = x[4:] # omega
+    w = w_d
+    w_1 = w[0]
+    w_2 = w[1]
+    w_3 = w[2]
+    
+    A = 0.5 *np.array([[0,-w_1,-w_2,-w_3,-q_1,-q_2,-q_3],
+                 [w_1,0,w_3,-w_2,q_0,-q_3,q_2],
+                 [w_2,-w_3,0,w_1,q_3,q_0,-q_1],
+                 [w_3,w_2,-w_1,0,-q_2,q_1,q_0],
+                 [0,0,0,0,0,2*(w_3*(I_2-I_3))/I_1,2*(w_2*(I_2-I_3))/I_1],
+                 [0,0,0,0,2*(w_3*(I_3-I_1))/I_2,0,2*(w_1*(I_3-I_1))/I_2],
+                 [0,0,0,0,2*(w_2*(I_1-I_2))/I_3,2*(w_1*(I_1-I_2))/I_3,0]])
+    
+    B = np.array([[0,0,0],
+                 [0,0,0],
+                 [0,0,0],
+                 [0,0,0],
+                 [1/I_1,0,0],
+                 [0,1/I_2,0],
+                 [0,0,1/I_3]])
+    
+    P = np.array(ricatti_solver(A,B,Q,R)).T
+    K = np.matmul(np.matmul(np.linalg.inv(R),B.T),P)
+
+    u = -K@(x-x_desired) # control torque
+
+    return u
\ No newline at end of file
diff --git a/simwise/simulation.py b/simwise/simulation.py
index 875f5e6..43bcd81 100644
--- a/simwise/simulation.py
+++ b/simwise/simulation.py
@@ -1,6 +1,6 @@
 import numpy as np
 from scipy.integrate import solve_ivp
-from quaternion import quaternion2euler, quaternion_dynamics, compute_control_torque, angle_axis_between
+from quaternion import quaternion2euler, quaternion_dynamics, compute_control_torque, angle_axis_between, compute_control_torque_lqr
 from graph_utils import graph_euler, graph_vector_matplotlib, graph_quaternion, graph_quaternion_matplotlib
 import matplotlib.pyplot as plt
 from tqdm import tqdm
@@ -8,7 +8,7 @@
 if __name__ == "__main__":
     # Initial conditions
     q = np.array([1, 0, 0, 0])
-    Ω = np.array([0.0, 0.0, 0.1]) # [rad/s]
+    Ω = np.array([0.1, 0.1, 0.1]) # [rad/s]
     x0 = np.concatenate((q, Ω))
     
     # Desired state
@@ -24,15 +24,20 @@
     K_p = 0.0005
     K_d = 0.005
 
+    # LQR weighting matrices
+    Q=np.diag([1,1,1,1,1,1,1])
+    R=np.diag([1,1,1])
+
     # Noise parameter (standard deviation of gaussian) [Nm]
     tau_noise = 0.00000288
 
     # Maximum actuator torque [Nm]
     tau_max = 0.0032
 
-    dt = 1/60 # [sec]
-    t_end = 10 * 60 # [sec] 2 minutes
+    dt = 1 # [sec]
+    t_end = 5 * 60 # [sec] --> minutes
     epoch = 0
+
     num_points = int(t_end // dt)
 
     y = np.zeros((num_points, 7))
@@ -46,8 +51,9 @@
     for i in tqdm(range(num_points)):
         t = t_arr[i]
 
-        f = lambda t, x: quaternion_dynamics(x, dt, inertia, compute_control_torque(x, x_d, K_p, K_d, tau_max=tau_max), noise=tau_noise)
-        sol = solve_ivp(f, [t, t+dt], x, method='RK45')
+        f_lqr = lambda t, x: quaternion_dynamics(x, dt, inertia, compute_control_torque_lqr(x, x_d, Q, R, inertia, tau_max=tau_max), noise=tau_noise)
+        f_pid = lambda t, x: quaternion_dynamics(x, dt, inertia, compute_control_torque(x, x_d, K_p, K_d, tau_max=tau_max), noise=tau_noise)
+        sol = solve_ivp(f_pid, [t, t+dt], x, method='RK45')
         y[i] = sol.y[:, -1]
         x = y[i]
         e_angles[i] = quaternion2euler(y[i, :4])