From c1216b25f71a446c5c121fe79a027ec35613c61a Mon Sep 17 00:00:00 2001
From: Gui-FernandesBR <guilherme_fernandes@usp.br>
Date: Mon, 22 Jan 2024 15:50:41 -0500
Subject: [PATCH] MNT: Add quaternion conversion functions to rocketpy.tools
 and update Flight

---
 .pylintrc                     |  1 +
 CHANGELOG.md                  |  1 +
 rocketpy/simulation/flight.py | 36 +++++++++-----------
 rocketpy/tools.py             | 64 +++++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 20 deletions(-)

diff --git a/.pylintrc b/.pylintrc
index 2611847f2..8cebb954c 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -436,6 +436,7 @@ disable=raw-checker-failed,
         no-else-return,
         inconsistent-return-statements,
         unspecified-encoding,
+        no-member, # because we use funcify_method decorator
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ad4ac536b..a3f73a2fe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,7 @@ straightforward as possible.
 - MNT: Add repr method to Parachute class [#490](https://github.com/RocketPy-Team/RocketPy/pull/490)
 - ENH: Function Reverse Arithmetic Priority [#488](https://github.com/RocketPy-Team/RocketPy/pull/488)
 - DOC: Update Header Related Docs
+- MNT: Encapsulate quaternion conversions [#537](https://github.com/RocketPy-Team/RocketPy/pull/537)
 
 ### Fixed
 
diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py
index 1c2549dfa..0b6bf3c10 100644
--- a/rocketpy/simulation/flight.py
+++ b/rocketpy/simulation/flight.py
@@ -11,7 +11,12 @@
 from ..mathutils.vector_matrix import Matrix, Vector
 from ..plots.flight_plots import _FlightPlots
 from ..prints.flight_prints import _FlightPrints
-from ..tools import find_closest
+from ..tools import (
+    find_closest,
+    quaternions_to_phi,
+    quaternions_to_precession,
+    quaternions_to_theta,
+)
 
 
 class Flight:
@@ -2295,33 +2300,24 @@ def lateral_attitude_angle(self):
     @funcify_method("Time (s)", "Precession Angle - ψ (°)", "spline", "constant")
     def psi(self):
         """Precession angle as a Function of time."""
-        psi = (180 / np.pi) * (
-            np.arctan2(self.e3[:, 1], self.e0[:, 1])
-            + np.arctan2(-self.e2[:, 1], -self.e1[:, 1])
-        )  # Precession angle
-        psi = np.column_stack([self.time, psi])  # Precession angle
-        return psi
+        psi = quaternions_to_precession(
+            self.e0.y_array, self.e1.y_array, self.e2.y_array, self.e3.y_array
+        )
+        return np.column_stack([self.time, psi])
 
     @funcify_method("Time (s)", "Spin Angle - φ (°)", "spline", "constant")
     def phi(self):
         """Spin angle as a Function of time."""
-        phi = (180 / np.pi) * (
-            np.arctan2(self.e3[:, 1], self.e0[:, 1])
-            - np.arctan2(-self.e2[:, 1], -self.e1[:, 1])
-        )  # Spin angle
-        phi = np.column_stack([self.time, phi])  # Spin angle
-        return phi
+        phi = quaternions_to_phi(
+            self.e0.y_array, self.e1.y_array, self.e2.y_array, self.e3.y_array
+        )
+        return np.column_stack([self.time, phi])
 
     @funcify_method("Time (s)", "Nutation Angle - θ (°)", "spline", "constant")
     def theta(self):
         """Nutation angle as a Function of time."""
-        theta = (
-            (180 / np.pi)
-            * 2
-            * np.arcsin(-((self.e1[:, 1] ** 2 + self.e2[:, 1] ** 2) ** 0.5))
-        )  # Nutation angle
-        theta = np.column_stack([self.time, theta])  # Nutation angle
-        return theta
+        theta = quaternions_to_theta(self.e1.y_array, self.e2.y_array)
+        return np.column_stack([self.time, theta])
 
     # Fluid Mechanics variables
     # Freestream Velocity
diff --git a/rocketpy/tools.py b/rocketpy/tools.py
index 6e765ca7a..21ac0006a 100644
--- a/rocketpy/tools.py
+++ b/rocketpy/tools.py
@@ -3,6 +3,7 @@
 import re
 from bisect import bisect_left
 
+import numpy as np
 import pytz
 from cftime import num2pydate
 from packaging import version as packaging_version
@@ -381,6 +382,69 @@ def check_requirement_version(module_name, version):
     return True
 
 
+# Flight
+
+
+def quaternions_to_precession(e0, e1, e2, e3):
+    """Calculates the Precession angle
+
+    Parameters
+    ----------
+    e0 : float
+        Euler parameter 0, must be between -1 and 1
+    e1 : float
+        Euler parameter 1, must be between -1 and 1
+    e2 : float
+        Euler parameter 2, must be between -1 and 1
+    e3 : float
+        Euler parameter 3, must be between -1 and 1
+    Returns
+    -------
+    float
+        Euler Precession angle in degrees
+    """
+    return (180 / np.pi) * (np.arctan2(e3, e0) + np.arctan2(-e2, -e1))
+
+
+def quaternions_to_phi(e0, e1, e2, e3):
+    """Calculates the Spin angle from quaternions.
+
+    Parameters
+    ----------
+    e0 : float
+        Euler parameter 0, must be between -1 and 1
+    e1 : float
+        Euler parameter 1, must be between -1 and 1
+    e2 : float
+        Euler parameter 2, must be between -1 and 1
+    e3 : float
+        Euler parameter 3, must be between -1 and 1
+
+    Returns
+    -------
+    float
+        Euler Spin angle in degrees
+    """
+    return (180 / np.pi) * (np.arctan2(e3, e0) - np.arctan2(-e2, -e1))
+
+
+def quaternions_to_theta(e1, e2):
+    """Calculates the Nutation angle from quaternions.
+
+    Parameters
+    ----------
+    e1 : float
+        Euler parameter 1, must be between -1 and 1
+    e2 : float
+        Euler parameter 2, must be between -1 and 1
+    Returns
+    -------
+    float
+        Euler Nutation angle in degrees
+    """
+    return (180 / np.pi) * 2 * np.arcsin(-((e1**2 + e2**2) ** 0.5))
+
+
 if __name__ == "__main__":
     import doctest