Merge pull request #132 from iwishiwasaneagle/perf-improvements

Improve performance of calculations

.github/workflows/CI.yml

@@ -26,6 +26,10 @@ jobs:
             pytest-flags: ""
             codecov: true
             platform: "ubuntu-latest"
+          - python-version: "3.11"
+            pytest-flags: ""
+            codecov: false
+            platform: "ubuntu-latest"
           - python-version: "3.10"
             pytest-flags: ""
             codecov: false
@@ -69,6 +73,7 @@ jobs:
 
     - name: Run tests
       run: |
+        NUMBA_DISABLE_JIT=${{ matrix.codecov }} \
         PYTHONPATH=tests \
         pytest tests \
         --cov-report=xml \

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "pybullet",
     "pydantic",
     "nptyping",
+    "numba"
 ]
 dynamic = ["version"]
 

src/jdrones/data_models.py

@@ -27,13 +27,14 @@ def __new__(cls, input_array=None):
         if input_array is None:
             obj = np.zeros(cls.k)
         else:
-            obj = np.asarray(input_array)
+            obj = np.array(input_array)
             if obj.shape != (cls.k,):
                 raise ValueError(f"Incorrect shape {obj.shape}")
 
-        if not np.can_cast(obj[0], DType):
-            raise ValueError(f"Incorrect dtype={obj.dtype}")
-        obj = obj.astype(DType)
+        if obj.dtype != DType:
+            if not np.can_cast(obj[0], DType):
+                raise ValueError(f"Incorrect dtype={obj.dtype}")
+            obj = obj.astype(DType)
 
         return obj.view(cls)
 
@@ -144,21 +145,32 @@ def prop_omega(self, prop_omega: VEC4):
     @classmethod
     def from_x(cls, x: LinearXAction):
         return cls(
-            np.concatenate(
-                [
-                    x[:3],
-                    (0, 0, 0, 0),
-                    x[6:9],
-                    x[3:6],
-                    x[9:12],
-                    (0, 0, 0, 0),
-                ],
-                dtype=DType,
-            )
+            [
+                x[0],
+                x[1],
+                x[2],
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+                x[6],
+                x[7],
+                x[8],
+                x[3],
+                x[4],
+                x[5],
+                x[9],
+                x[10],
+                x[11],
+                0.0,
+                0.0,
+                0.0,
+                0.0,
+            ],
         )
 
     def to_x(self) -> LinearXAction:
-        return np.concatenate([self.pos, self.vel, self.rpy, self.ang_vel])
+        return self[[0, 1, 2, 10, 11, 12, 7, 8, 9, 13, 14, 15]]
 
     def quat_rotation(self, quat: VEC4) -> "State":
         rotmat = quat_to_rotmat(quat)

src/jdrones/envs/base/lineardronenev.py

@@ -2,6 +2,7 @@
 #  SPDX-License-Identifier: GPL-3.0-only
 from typing import Tuple
 
+import numba
 import numpy as np
 from gymnasium import spaces
 from jdrones.data_models import State
@@ -84,6 +85,7 @@ def get_matrices(model: URDFModel):
         return A, B, C
 
     @staticmethod
+    @numba.njit
     def calc_dx(A, B, C, x, u):
         return (A @ x + B @ u + C.T).flatten()
 

src/jdrones/envs/base/nonlineardronenev.py

@@ -3,18 +3,18 @@
 import functools
 from typing import Tuple
 
+import numba
 import numpy as np
 from jdrones.data_models import State
 from jdrones.data_models import URDFModel
 from jdrones.envs.base.basedronenev import BaseDroneEnv
 from jdrones.transforms import euler_to_quat
 from jdrones.transforms import euler_to_rotmat
+from jdrones.types import DType
 from jdrones.types import MAT3X3
 from jdrones.types import PropellerAction
 from jdrones.types import VEC3
 
-UZ = np.array([0, 0, 1]).reshape((-1, 1))
-
 
 class NonlinearDynamicModelDroneEnv(BaseDroneEnv):
     """
@@ -49,11 +49,50 @@ def _get_cached_time_invariant_params(
         g: float
             acceleration due to gravity
         """
-        inertias = np.diag(np.array(model.I))
+        inertias = np.diag(np.array(model.I, dtype=DType))
         m = model.mass
         drag_coeffs = model.drag_coeffs
         g = model.g
-        return inertias, np.linalg.inv(inertias), m, np.array(drag_coeffs), g
+        return (
+            inertias,
+            np.linalg.inv(inertias),
+            m,
+            np.array(drag_coeffs, dtype=DType),
+            g,
+        )
+
+    @staticmethod
+    @numba.njit(
+        [
+            "float64[:](float64[:],float64[:],float64[:],float64[:],"
+            "float64[:,:],float64[:,:],float64,float64[:],float64)",
+        ]
+    )
+    def _calc_dstate(
+        rpy,
+        vel,
+        ang_vel,
+        u_star,
+        inertias,
+        inv_intertias,
+        m,
+        drag_coeffs,
+        g,
+    ):
+        UZ = np.array([0, 0, 1], dtype=DType).reshape((-1, 1))
+        R_W_Q = euler_to_rotmat(rpy)
+        R_Q_W = np.transpose(R_W_Q)
+
+        body_vel = np.dot(R_W_Q, vel)
+        drag_force = -np.sign(vel) * np.dot(R_Q_W, drag_coeffs * np.square(body_vel))
+        dstate = np.zeros((20,), dtype=DType)
+        dstate[:3] = vel
+        dstate[7:10] = ang_vel
+        dstate[10:13] = (
+            -m * g * UZ.T + (R_W_Q @ UZ).T * u_star[3] + drag_force
+        ).flatten() / m
+        dstate[13:16] = np.dot(inv_intertias, u_star[0:3])
+        return dstate
 
     @classmethod
     def calc_dstate(cls, action: PropellerAction, state: State, model: URDFModel):
@@ -72,33 +111,14 @@ def calc_dstate(cls, action: PropellerAction, state: State, model: URDFModel):
         State
             The state derivative
         """
-        (
-            inertias,
-            inv_intertias,
-            m,
-            drag_coeffs,
-            g,
-        ) = cls._get_cached_time_invariant_params(model)
         u_star = model.rpm2rpyT(np.square(action))
-
-        R_W_Q = euler_to_rotmat(state.rpy)
-        R_Q_W = np.transpose(R_W_Q)
-
-        body_vel = np.dot(R_W_Q, state.vel)
-        drag_force = -np.sign(state.vel) * np.dot(
-            R_Q_W, drag_coeffs * np.square(body_vel)
+        return cls._calc_dstate(
+            state.rpy,
+            state.vel,
+            state.ang_vel,
+            u_star,
+            *cls._get_cached_time_invariant_params(model),
         )
-        dstate = np.concatenate(
-            [
-                state.vel,
-                (0, 0, 0, 0),
-                state.ang_vel,
-                (-m * g * UZ.T + (R_W_Q @ UZ).T * u_star[3] + drag_force).flatten() / m,
-                np.dot(inv_intertias, u_star[0:3]),
-                (0, 0, 0, 0),
-            ]
-        )
-        return dstate
 
     def step(self, action: PropellerAction) -> Tuple[State, float, bool, bool, dict]:
         """

src/jdrones/envs/dronemodels.py

@@ -1,21 +1,23 @@
 #  Copyright 2023 Jan-Hendrik Ewers
 #  SPDX-License-Identifier: GPL-3.0-only
-import functools
 from importlib.resources import files
 
+import numba
 import numpy as np
 from jdrones.data_models import URDFModel
+from jdrones.types import DType
 
 
-@functools.cache
+@numba.njit
 def droneplus_mixing_matrix(*, length, k_Q, k_T):
     return np.array(
         [
             [0, -k_T * length, 0, k_T * length],
             [-k_T * length, 0, k_T * length, 0],
             [k_Q, -k_Q, k_Q, -k_Q],
             [k_T, k_T, k_T, k_T],
-        ]
+        ],
+        dtype=DType,
     )
 
 

src/jdrones/transforms.py

@@ -1,12 +1,13 @@
 #  Copyright 2023 Jan-Hendrik Ewers
 #  SPDX-License-Identifier: GPL-3.0-only
+import numba
 import numpy as np
-import pybullet as p
 from jdrones.types import MAT3X3
 from jdrones.types import VEC3
 from jdrones.types import VEC4
 
 
+@numba.njit
 def quat_to_euler(quat: VEC4) -> VEC3:
     """
     PyBullet does a Body 3-2-1 sequence:
@@ -53,10 +54,20 @@ def euler_to_quat(roll, pitch, yaw):
         Euler angle (roll,pitch,yaw)
 
     """
-
-    return p.getEulerFromQuaternion(quat)
-
-
+    x, y, z, w = quat
+    sinr_cosp = 2 * (w * x + y * z)
+    cosr_cosp = 1 - 2 * (x * x + y * y)
+    roll = np.arctan2(sinr_cosp, cosr_cosp)
+    sinp = np.sqrt(1 + 2 * (w * y - x * z))
+    cosp = np.sqrt(1 - 2 * (w * y - x * z))
+    pitch = 2 * np.arctan2(sinp, cosp) - np.pi / 2
+    siny_cosp = 2 * (w * z + x * y)
+    cosy_cosp = 1 - 2 * (y * y + z * z)
+    yaw = np.arctan2(siny_cosp, cosy_cosp)
+    return roll, pitch, yaw
+
+
+@numba.njit
 def euler_to_quat(euler: VEC3) -> VEC4:
     """PyBullet does a Body 3-2-1 sequence:
 
@@ -96,12 +107,46 @@ def quat_to_euler(x, y, z, w):
         Quaternion (x,y,z,w)
 
     """
-    return p.getQuaternionFromEuler(euler)
+    roll, pitch, yaw = euler
 
+    cy = np.cos(yaw * 0.5)
+    sy = np.sin(yaw * 0.5)
+    cp = np.cos(pitch * 0.5)
+    sp = np.sin(pitch * 0.5)
+    cr = np.cos(roll * 0.5)
+    sr = np.sin(roll * 0.5)
 
-def quat_to_rotmat(quat: VEC4) -> MAT3X3:
-    return np.reshape(p.getMatrixFromQuaternion(quat), (3, 3))
+    w = cr * cp * cy + sr * sp * sy
+    x = sr * cp * cy - cr * sp * sy
+    y = cr * sp * cy + sr * cp * sy
+    z = cr * cp * sy - sr * sp * cy
+
+    return np.array([x, y, z, w])
 
 
+@numba.njit
+def quat_to_rotmat(quat: VEC4) -> MAT3X3:
+    x, y, z, w = quat
+    x2, y2, z2, w2 = quat * quat
+    xy, xz, xw, yz, yw, zw = (
+        quat[0] * quat[1],
+        quat[0] * quat[2],
+        quat[0] * quat[3],
+        quat[1] * quat[2],
+        quat[1] * quat[3],
+        quat[2] * quat[3],
+    )
+
+    rot_mat = np.array(
+        [
+            [1 - 2 * (y2 + z2), 2 * (xy - zw), 2 * (xz + yw)],
+            [2 * (xy + zw), 1 - 2 * (x2 + z2), 2 * (yz - xw)],
+            [2 * (xz - yw), 2 * (yz + xw), 1 - 2 * (x2 + y2)],
+        ]
+    )
+    return rot_mat
+
+
+@numba.njit
 def euler_to_rotmat(euler: VEC3) -> MAT3X3:
     return quat_to_rotmat(euler_to_quat(euler))

tests/test_data_models.py

@@ -4,6 +4,7 @@
 import pytest
 from jdrones.data_models import State
 from jdrones.data_models import STATE_ENUM
+from jdrones.types import DType
 from scipy.spatial.transform import Rotation as R
 
 
@@ -269,7 +270,7 @@ def test_x_to_state():
 )
 def test_state_quat_rotation(quat, act, exp):
     exp = State(exp)
-    rotated = State(act).quat_rotation(quat)
+    rotated = State(act).quat_rotation(np.array(quat, dtype=DType))
     assert np.allclose(rotated.pos, exp.pos)
     assert np.allclose(rotated.quat, exp.quat)
     assert np.allclose(

tests/test_transforms.py

@@ -46,6 +46,8 @@ def test_quat_euler_sanity_check(quat, euler):
     assert np.allclose(
         pybullet.getQuaternionFromEuler(pybullet.getEulerFromQuaternion(quat)), quat
     )
+    assert np.allclose(pybullet.getQuaternionFromEuler(euler), quat)
+    assert np.allclose(pybullet.getEulerFromQuaternion(quat), euler)
 
 
 @QUAT_EULER