Fix all numeric models

Author: amacati

drone_models/core.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from functools import partial, wraps
 from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Protocol, TypeVar, runtime_checkable
 
@@ -33,6 +34,8 @@ def wrapper(
         ) -> tuple[Array, Array, Array, Array, Array | None]:
             if not rotor_dynamics and rotor_vel is not None:
                 raise ValueError("Rotor dynamics not supported, but rotor_vel is provided.")
+            if rotor_dynamics and rotor_vel is None:
+                warnings.warn("Rotor velocity not provided, using commanded rotor velocity.")
             return fn(pos, quat, vel, ang_vel, cmd, rotor_vel, *args, **kwargs)
 
         wrapper.__drone_model_features__ = {"rotor_dynamics": rotor_dynamics}

drone_models/first_principles/__init__.py

@@ -1,3 +1,8 @@
+"""First principles model.
+
+TODO: Add description.
+"""
+
 from drone_models.first_principles.model import dynamics, dynamics_symbolic
 
 __all__ = ["dynamics", "dynamics_symbolic"]

drone_models/first_principles/model.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING
 
 import casadi as cs
@@ -30,15 +29,15 @@ def dynamics(
     dist_f: Array | None = None,
     dist_t: Array | None = None,
     *,
-    thrust_tau: float,
+    mass: float,
+    gravity_vec: Array,
+    J: Array,
+    J_inv: Array,
     KF: float,
     KM: float,
     L: float,
     mixing_matrix: Array,
-    gravity_vec: Array,
-    mass: float,
-    J: Array,
-    J_inv: Array,
+    thrust_tau: float,
 ) -> tuple[Array, Array, Array, Array, Array | None]:
     r"""First principles model for a quatrotor.
 
@@ -52,12 +51,22 @@ def dynamics(
         vel: Velocity of the drone (m/s).
         ang_vel: Angular velocity of the drone (rad/s).
         cmd: Motor speeds (rad/s).
-        constants: Containing the constants of the drone.
         rotor_vel: Angular velocity of the 4 motors (rad/s). If None, the commanded thrust is
             directly applied. If value is given, thrust dynamics are calculated.
         dist_f: Disturbance force acting on the CoM (N).
         dist_t: Disturbance torque acting on the CoM (Nm).
 
+        mass: Mass of the drone (kg).
+        gravity_vec: Gravity vector (m/s^2). We assume the gravity vector points downwards, e.g.
+            [0, 0, -9.81].
+        J: Inertia matrix (kg m^2).
+        J_inv: Inverse inertia matrix (1/kg m^2).
+        KF: Motor force constant (N/rad^2).
+        KM: Motor torque constant (Nm/rad^2).
+        L: Distance from the CoM to the motor (m).
+        mixing_matrix: Mixing matrix denoting the turn direction of the motors (4x3).
+        thrust_tau: Thrust time constant (s).
+
     .. math::
         \sum_{i=1}^{\\infty} x_{i} TODO
 
@@ -67,15 +76,13 @@ def dynamics(
         More information https://ahrs.readthedocs.io/en/latest/filters/angular.html
     """
     xp = array_namespace(pos)
-    mass, gravity_vec, KF, KM, L, mixing_matrix, J, J_inv = to_xp(
-        mass, gravity_vec, KF, KM, L, mixing_matrix, J, J_inv, xp=xp, device=device(pos)
+    mass, gravity_vec, J, J_inv, KF, KM, L, mixing_matrix, thrust_tau = to_xp(
+        mass, gravity_vec, J, J_inv, KF, KM, L, mixing_matrix, thrust_tau, xp=xp, device=device(pos)
     )
     rot = R.from_quat(quat)
     # Rotor dynamics
     if rotor_vel is None:
-        rotor_vel_dot = None
-        rotor_vel = cmd
-        warnings.warn("Rotor velocity is not provided, using commanded rotor velocity directly.")
+        rotor_vel, rotor_vel_dot = cmd, None
     else:
         rotor_vel_dot = 1 / thrust_tau * (cmd - rotor_vel) - 1 / KM * rotor_vel**2
     # Creating force and torque vector
@@ -87,7 +94,7 @@ def dynamics(
     # Because there currently is no way to identify the z torque in relation to the thrust,
     # we rely on a old identified value that can compute rpm to torque.
     # force = kf * rpm², torque = km * rpm² => torque = km/kf*force TODO
-    torques_motor_vec = forces_motor @ mixing_matrix * xp.stack([L, L, KM / KF])
+    torque = forces_motor @ mixing_matrix * xp.stack([L, L, KM / KF])
 
     # Linear equation of motion
     forces_motor_vec_world = rot.apply(forces_motor_vec)
@@ -99,12 +106,11 @@ def dynamics(
     vel_dot = forces_sum / mass
 
     # Rotational equation of motion
-    torques_sum = torques_motor_vec
     if dist_t is not None:
-        torques_sum = torques_sum + rot.apply(dist_t, inverse=True)
+        torque = torque + rot.apply(dist_t, inverse=True)
     quat_dot = rotation.ang_vel2quat_dot(quat, ang_vel)
-    ang_vel_dot = J_inv @ (torques_sum - xp.linalg.cross(ang_vel, J @ ang_vel))
-
+    torque = torque - xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
+    ang_vel_dot = (J_inv @ torque[..., None])[..., 0]
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 

drone_models/first_principles/params.py

@@ -1,3 +1,5 @@
+"""First principles model parameters."""
+
 from __future__ import annotations
 
 import tomllib
@@ -11,6 +13,8 @@
 
 
 class FirstPrinciplesParams(NamedTuple):
+    """Parameters for the FirstPrinciples model."""
+
     thrust_tau: float
     KF: float
     KM: float
@@ -23,6 +27,7 @@ class FirstPrinciplesParams(NamedTuple):
 
     @staticmethod
     def load(drone_model: str) -> FirstPrinciplesParams:
+        """Load the parameters for the drone model from the params.toml file."""
         with open(Path(__file__).parent / "params.toml", "rb") as f:
             params = tomllib.load(f)
         if drone_model not in params:

drone_models/so_rpy/__init__.py

@@ -1,3 +1,8 @@
+"""SoRpy model.
+
+TODO: Add description.
+"""
+
 from drone_models.so_rpy.model import dynamics, dynamics_symbolic
 
 __all__ = ["dynamics", "dynamics_symbolic"]

drone_models/so_rpy/model.py

@@ -33,13 +33,13 @@ def dynamics(
     *,
     mass: float,
     gravity_vec: Array,
+    J: Array,
+    J_inv: Array,
     acc_coef: Array,
     cmd_f_coef: Array,
     rpy_coef: Array,
     rpy_rates_coef: Array,
     cmd_rpy_coef: Array,
-    J: Array,
-    J_inv: Array,
 ) -> tuple[Array, Array, Array, Array, Array | None]:
     """The fitted double integrator (DI) model with optional motor delay (D).
 
@@ -49,12 +49,21 @@ def dynamics(
         vel: Velocity of the drone (m/s).
         ang_vel: Angular velocity of the drone (rad/s).
         cmd: Roll pitch yaw (rad) and collective thrust (N) command.
-        constants: Containing the constants of the drone.
-        rotor_vel: Speed of the 4 motors (rad/s). If None, the commanded thrust is directly
-            applied. If a value is given, the function raises an error.
+        rotor_vel: Speed of the 4 motors (rad/s). Kept for compatibility with the model signature.
         dist_f: Disturbance force (N) acting on the CoM.
         dist_t: Disturbance torque (Nm) acting on the CoM.
 
+        mass: Mass of the drone (kg).
+        gravity_vec: Gravity vector (m/s^2). We assume the gravity vector points downwards, e.g.
+            [0, 0, -9.81].
+        J: Inertia matrix (kg m^2).
+        J_inv: Inverse inertia matrix (1/kg m^2).
+        acc_coef: Coefficient for the acceleration (1/s^2).
+        cmd_f_coef: Coefficient for the collective thrust (N/rad^2).
+        rpy_coef: Coefficient for the roll pitch yaw dynamics (1/s).
+        rpy_rates_coef: Coefficient for the roll pitch yaw rates dynamics (1/s^2).
+        cmd_rpy_coef: Coefficient for the roll pitch yaw command dynamics (1/s).
+
     Returns:
         The derivatives of all state variables.
     """
@@ -93,7 +102,8 @@ def dynamics(
         # adding torque
         torque = torque + rot.apply(dist_t, inverse=True)
         # back to angular acceleration
-        ang_vel_dot = J_inv @ (torque - xp.linalg.cross(ang_vel, J @ ang_vel))
+        torque = torque - xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
+        ang_vel_dot = (J_inv @ torque[..., None])[..., 0]
 
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 

drone_models/so_rpy/params.py

@@ -13,7 +13,7 @@
 
 
 class SoRpyParams(NamedTuple):
-    """TODO."""
+    """Parameters for the SoRpy model."""
 
     mass: float
     gravity_vec: Array
@@ -27,6 +27,7 @@ class SoRpyParams(NamedTuple):
 
     @staticmethod
     def load(drone_model: str) -> SoRpyParams:
+        """Load the parameters for the drone model from the params.toml file."""
         with open(Path(__file__).parent / "params.toml", "rb") as f:
             params = tomllib.load(f)
         if drone_model not in params:

drone_models/so_rpy_rotor/__init__.py

@@ -1,3 +1,8 @@
+"""SoRpyRotor model.
+
+TODO: Add description.
+"""
+
 from drone_models.so_rpy_rotor.model import dynamics, dynamics_symbolic
 
 __all__ = ["dynamics", "dynamics_symbolic"]

drone_models/so_rpy_rotor/model.py

@@ -2,18 +2,18 @@
 
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING
 
 import casadi as cs
 from array_api_compat import array_namespace
+from array_api_compat import device as xp_device
 from scipy.spatial.transform import Rotation as R
 
 import drone_models.symbols as symbols
 from drone_models.core import register_model_parameters, supports
 from drone_models.so_rpy_rotor.params import SoRpyRotorParams
 from drone_models.transform import motor_force2rotor_vel
-from drone_models.utils import rotation
+from drone_models.utils import rotation, to_xp
 
 if TYPE_CHECKING:
     from array_api_typing import Array
@@ -54,31 +54,52 @@ def dynamics(
         vel: Velocity of the drone (m/s).
         ang_vel: Angular velocity of the drone (rad/s).
         cmd: Roll pitch yaw (rad) and collective thrust (N) command.
-        constants: Containing the constants of the drone.
         rotor_vel: Speed of the 4 motors (rad/s). If None, the commanded thrust is directly
             applied (not recommended). If value is given, rotor dynamics are calculated.
         dist_f: Disturbance force acting on the CoM (N).
         dist_t: Disturbance torque acting on the CoM (Nm).
 
+        mass: Mass of the drone (kg).
+        gravity_vec: Gravity vector (m/s^2). We assume the gravity vector points downwards, e.g.
+            [0, 0, -9.81].
+        KF: Motor force constant (N/rad^2).
+        KM: Motor torque constant (Nm/rad^2).
+        J: Inertia matrix (kg m^2).
+        J_inv: Inverse inertia matrix (1/kg m^2).
+        rotor_coef: Coefficient for the rotor dynamics (1/s).
+        acc_coef: Coefficient for the acceleration (1/s^2).
+        cmd_f_coef: Coefficient for the collective thrust (N/rad^2).
+        rpy_coef: Coefficient for the roll pitch yaw dynamics (1/s).
+        rpy_rates_coef: Coefficient for the roll pitch yaw rates dynamics (1/s^2).
+        cmd_rpy_coef: Coefficient for the roll pitch yaw command dynamics (1/s).
+
     Returns:
         tuple[Array, Array, Array, Array, Array | None]: _description_
     """
     xp = array_namespace(pos)
+    # Convert constants to the correct framework and device
+    device = xp_device(pos)
+    mass, gravity_vec, KF, KM, J, J_inv = to_xp(
+        mass, gravity_vec, KF, KM, J, J_inv, xp=xp, device=device
+    )
+    rotor_coef, acc_coef, cmd_f_coef = to_xp(rotor_coef, acc_coef, cmd_f_coef, xp=xp, device=device)
+    rpy_coef, rpy_rates_coef, cmd_rpy_coef = to_xp(
+        rpy_coef, rpy_rates_coef, cmd_rpy_coef, xp=xp, device=device
+    )
+
     cmd_f = cmd[..., -1]
     cmd_rotor_vel = motor_force2rotor_vel(cmd_f / 4, KF)
     cmd_rpy = cmd[..., 0:3]
     rot = R.from_quat(quat)
     euler_angles = rot.as_euler("xyz")
 
     if rotor_vel is None:
-        rotor_vel_dot = None
-        rotor_vel = cmd_rotor_vel
-        warnings.warn("Rotor velocity is not provided, using commanded rotor velocity directly.")
+        rotor_vel, rotor_vel_dot = cmd_rotor_vel[..., None], None
     else:
         rotor_vel_dot = 1 / rotor_coef * (cmd_rotor_vel[..., None] - rotor_vel) - KM * rotor_vel**2
-    forces_motor = xp.sum(KF * rotor_vel**2, axis=-1)
-    forces_sum = xp.sum(forces_motor, axis=-1)
-    thrust = acc_coef + cmd_f_coef * forces_sum
+
+    forces_motor = KF * xp.sum(rotor_vel**2, axis=-1)
+    thrust = acc_coef + cmd_f_coef * forces_motor
 
     drone_z_axis = rot.as_matrix()[..., -1]
 
@@ -98,11 +119,11 @@ def dynamics(
         # adding torque disturbances to the state
         # angular acceleration can be converted to total torque given the inertia matrix
         torque = ang_vel_dot @ J.mT + xp.linalg.cross(ang_vel, ang_vel @ J.mT)
-
         # adding torque
         torque = torque + rot.apply(dist_t, inverse=True)  # TODO rotation into body frame
         # back to angular acceleration
-        ang_vel_dot = J_inv @ (torque - xp.linalg.cross(ang_vel, J @ ang_vel))
+        torque = torque - xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
+        ang_vel_dot = (J_inv @ torque[..., None])[..., 0]
 
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 

drone_models/so_rpy_rotor/params.py

@@ -13,7 +13,7 @@
 
 
 class SoRpyRotorParams(NamedTuple):
-    """TODO."""
+    """Parameters for the SoRpyRotor model."""
 
     mass: float
     gravity_vec: Array
@@ -30,6 +30,7 @@ class SoRpyRotorParams(NamedTuple):
 
     @staticmethod
     def load(drone_model: str) -> SoRpyRotorParams:
+        """Load the parameters for the drone model from the params.toml file."""
         with open(Path(__file__).parent / "params.toml", "rb") as f:
             params = tomllib.load(f)
         if drone_model not in params: