[WIP] Continue fixing tests

amacati · amacati · commit 2275cce6e3b5 · 2025-08-12T19:05:21.000+02:00
diff --git a/drone_models/models/first_principles.py b/drone_models/models/first_principles.py
@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
-import array_api_extra as xpx
 import casadi as cs
 from array_api_compat import array_namespace
 from scipy.spatial.transform import Rotation as R
 
 import drone_models.models.symbols as symbols
 from drone_models.models.utils import supports
-from drone_models.transform import motor_force2rotor_speed
 from drone_models.utils import rotation
 
 if TYPE_CHECKING:
@@ -43,7 +42,7 @@ def dynamics(
         quat: Quaternion of the drone (xyzw).
         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.
+        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.
@@ -64,9 +63,10 @@ def 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.")
     else:
         rotor_vel_dot = (
-            1 / constants.ROTOR_TAU * (cmd - rotor_vel) - 1 / constants.ROTOR_D * rotor_vel**2
+            1 / constants.THRUST_TAU * (cmd - rotor_vel) - 1 / constants.KM * rotor_vel**2
         )
     # Creating force and torque vector
     forces_motor = constants.KF * rotor_vel**2
diff --git a/drone_models/models/identified/so_rpy.py b/drone_models/models/identified/so_rpy.py
@@ -38,16 +38,15 @@ def dynamics(
         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 rotor speed dynamics are calculated.
+            applied. If a value is given, the function raises an error.
         dist_f: Disturbance force (N) acting on the CoM.
         dist_t: Disturbance torque (Nm) acting on the CoM.
 
     Returns:
         The derivatives of all state variables.
     """
     xp = array_namespace(pos)
-    cmd_rotor_vel = cmd[..., -1]
-    cmd_f = xp.sum(constants.KF * cmd_rotor_vel**2, axis=-1)
+    cmd_f = cmd[..., -1]
     cmd_rpy = cmd[..., 0:3]
     rot = R.from_quat(quat)
     euler_angles = rot.as_euler("xyz")
diff --git a/drone_models/models/identified/so_rpy_rotor.py b/drone_models/models/identified/so_rpy_rotor.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 from array_api_compat import array_namespace
@@ -48,7 +49,7 @@ def dynamics(
     """
     xp = array_namespace(pos)
     cmd_f = cmd[..., -1]
-    cmd_rotor_vel = motor_force2rotor_speed(cmd_f, constants.KF)
+    cmd_rotor_vel = motor_force2rotor_speed(cmd_f / 4, constants.KF)
     cmd_rpy = cmd[..., 0:3]
     rot = R.from_quat(quat)
     euler_angles = rot.as_euler("xyz")
@@ -57,19 +58,11 @@ def dynamics(
     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.")
     else:
         rotor_vel_dot = (
-            1 / constants.ROTOR_TAU * (cmd_rotor_vel - rotor_vel)
-            - 1 / constants.ROTOR_D * rotor_vel**2
+            1 / constants.DI_D_ACC[2] * (cmd_rotor_vel - rotor_vel) - constants.KM * rotor_vel**2
         )
-
-    # Note: Due to the structure of the integrator, we split the commanded thrust into
-    # four equal parts and later apply the sum as total thrust again. Those four forces
-    # are not the true forces of the motors, but the sum is the true total thrust.
-    rotor_vel_dot = (
-        xp.asarray(1 / constants.DI_D_ACC[2] * (cmd_rotor_vel - rotor_vel))
-        - constants.KM * rotor_vel**2
-    )
     forces_motor = xp.sum(constants.KF * rotor_vel**2, axis=-1)
     forces_sum = xp.sum(forces_motor, axis=-1)
     thrust = constants.DI_D_ACC[0] + constants.DI_D_ACC[1] * forces_sum
diff --git a/drone_models/models/identified/so_rpy_rotor_drag.py b/drone_models/models/identified/so_rpy_rotor_drag.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 from array_api_compat import array_namespace
@@ -38,8 +39,8 @@ def dynamics(
         ang_vel: Angular velocity of the drone (rad/s).
         command: Roll pitch yaw (rad) and collective thrust (N) command.
         constants: Containing the constants of the drone.
-        rotor_vel: Thrust of the 4 motors in N. If None, the commanded thrust is directly applied.
-            If a value is given, thrust dynamics are calculated.
+        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 (N) acting on the CoM.
         dist_t: Disturbance torque (Nm) acting on the CoM.
 
@@ -48,16 +49,20 @@ def dynamics(
     """
     xp = array_namespace(pos)
     cmd_f = command[..., -1]
-    cmd_rotor_vel = motor_force2rotor_speed(cmd_f, constants.KF)
+    cmd_rotor_vel = motor_force2rotor_speed(cmd_f / 4, constants.KF)
     cmd_rpy = command[..., 0:3]
     rot = R.from_quat(quat)
     euler_angles = rot.as_euler("xyz")
     rpy_rates = rotation.ang_vel2rpy_rates(quat, ang_vel)
 
-    rotor_vel_dot = (
-        xp.asarray(1 / constants.DI_DD_ACC[2] * (cmd_rotor_vel / 4 - rotor_vel))
-        - constants.DI_DD_ACC[3] * rotor_vel**2
-    )
+    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.")
+    else:
+        rotor_vel_dot = (
+            1 / constants.DI_DD_ACC[2] * (cmd_rotor_vel - rotor_vel) - constants.KM * rotor_vel**2
+        )
     forces_motor = xp.sum(constants.KF * rotor_vel**2, axis=-1)
     forces_sum = xp.sum(forces_motor, axis=-1)
     thrust = constants.DI_DD_ACC[0] + constants.DI_DD_ACC[1] * forces_sum
diff --git a/drone_models/utils/rotation.py b/drone_models/utils/rotation.py
@@ -6,21 +6,52 @@
 
 from __future__ import annotations
 
-import os
-
-# TODO raise error if scipy is already loaded and the feature flag is not set
-os.environ["SCIPY_ARRAY_API"] = "1"  # Feature flag to activate Array API in scipy
 import re
 from typing import TYPE_CHECKING
 
 import casadi as cs
 from array_api_compat import array_namespace
+from jax.scipy.spatial.transform import Rotation as JR
 from scipy.spatial.transform import Rotation as R
 
 if TYPE_CHECKING:
     from array_api_strict import Array
 
 
+def from_quat(quat: Array, scalar_first: bool = False) -> R:
+    """Creates a rotation object compatible with the type of the given quat."""
+    if isinstance(quat, jp.ndarray):
+        if scalar_first:
+            raise ValueError("scalar_first is not supported by jax rotations")
+        return JR.from_quat(quat)
+    else:
+        return R.from_quat(quat, scalar_first=scalar_first)
+
+
+def from_rotvec(rotvec: Array, degrees: bool = False) -> R:
+    """Creates a rotation object compatible with the type of the given rotvec."""
+    if isinstance(rotvec, jp.ndarray):
+        return JR.from_rotvec(rotvec, degrees)
+    else:
+        return R.from_rotvec(rotvec, degrees)
+
+
+def from_euler(seq: str, angles: Array, degrees: bool = False) -> R:
+    """Creates a rotation object compatible with the type of the given angles."""
+    if isinstance(angles, jp.ndarray):
+        return JR.from_euler(seq, angles, degrees)
+    else:
+        return R.from_euler(seq, angles, degrees)
+
+
+def from_matrix(matrix: Array) -> R:
+    """Creates a rotation objecte compatible with the type of the given rotation matrix."""
+    if isinstance(matrix, jp.ndarray):
+        return JR.from_matrix(matrix)
+    else:
+        return R.from_matrix(matrix)
+
+
 def ang_vel2quat_dot(quat: Array, ang_vel: Array) -> Array:
     """Calculates the quaternion derivative based on an angular velocity."""
     xp = array_namespace(quat)
diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py
@@ -42,9 +42,9 @@ def create_rnd_commands(shape: tuple[int, ...] = (), dim: int = 4) -> Array:
 
 def skip_models_without_features(model: Callable, features: list[str]):
     """Skip the model if it does not have the required features."""
-    for feature in model_features(model):
-        if feature not in features:
-            pytest.skip(f"Model {model.__name__} does not have the required features.")
+    for feature in features:
+        if not model_features(model)[feature]:
+            pytest.skip(f"Model {model.__name__} does not have the feature '{feature}'.")
 
 
 @pytest.mark.unit
@@ -54,25 +54,76 @@ def test_model_features(model_name: str, model: Callable):
     assert hasattr(model, "__drone_model_features__"), (
         f"Model function {model_name} does not have __drone_model_features__ attribute"
     )
-    features = getattr(model, "__drone_model_features__")
+    features = model_features(model)
     assert isinstance(features, dict), (
-        f"__drone_model_features__ should be a dict, got {type(features)} for {model_name}"
+        f"model features should be a dict, got {type(features)} for {model_name}"
     )
     assert "rotor_dynamics" in features, (
-        f"__drone_model_features__ should contain 'rotor_dynamics' key for {model_name}"
+        f"model features should contain 'rotor_dynamics' key for {model_name}"
     )
 
 
 @pytest.mark.unit
 @pytest.mark.parametrize("model_name, model", available_models.items())
 @pytest.mark.parametrize("drone_name", Constants.available_configs)
-def test_model_single_input_no_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
+def test_model_single_no_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
     pos, quat, vel, ang_vel, _, _, _ = create_rnd_states()
-    commands = create_rnd_commands(dim=4)  # TODO make dependent on model
+    cmd = create_rnd_commands(dim=4)  # TODO make dependent on model
+    if model_features(model)["rotor_dynamics"]:
+        with pytest.warns(UserWarning, match="Rotor velocity is not provided"):
+            dpos, dquat, dvel, dang_vel, drotor_vel = model(
+                pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=None
+            )
+    else:
+        dpos, dquat, dvel, dang_vel, drotor_vel = model(
+            pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=None
+        )
+    assert drotor_vel is None, "Model should not return rotor velocities without rotor_vel input"
+    # Check if the output is on the correct device, has the correct type and shape
+    for dx, x in zip([dpos, dquat, dvel, dang_vel], [pos, quat, vel, ang_vel], strict=True):
+        assert isinstance(dx, type(x))
+        assert xp_device(dx) == xp_device(x)
+        assert dx.shape == x.shape
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("model_name, model", available_models.items())
+@pytest.mark.parametrize("drone_name", Constants.available_configs)
+def test_model_single_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
+    skip_models_without_features(model, ["rotor_dynamics"])
+
+    pos, quat, vel, ang_vel, rotor_vel, _, _ = create_rnd_states()
+    cmd = create_rnd_commands(dim=4)  # TODO make dependent on model
     dpos, dquat, dvel, dang_vel, drotor_vel = model(
-        pos, quat, vel, ang_vel, commands, Constants.from_config(drone_name, xp), rotor_vel=None
+        pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=rotor_vel
     )
     # Check if the output is on the correct device, has the correct type and shape
+    for dx, x in zip(
+        [dpos, dquat, dvel, dang_vel, drotor_vel], [pos, quat, vel, ang_vel, rotor_vel], strict=True
+    ):
+        assert isinstance(dx, type(x))
+        assert xp_device(dx) == xp_device(x)
+        assert dx.shape == x.shape
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("model_name, model", available_models.items())
+@pytest.mark.parametrize("drone_name", Constants.available_configs)
+def test_model_batched_no_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
+    batch_shape = (10,)
+    pos, quat, vel, ang_vel, _, _, _ = create_rnd_states(batch_shape)
+    cmd = create_rnd_commands(batch_shape, dim=4)  # TODO make dependent on model
+    if model_features(model)["rotor_dynamics"]:
+        with pytest.warns(UserWarning, match="Rotor velocity is not provided"):
+            dpos, dquat, dvel, dang_vel, drotor_vel = model(
+                pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=None
+            )
+    else:
+        dpos, dquat, dvel, dang_vel, drotor_vel = model(
+            pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=None
+        )
+    assert drotor_vel is None, "Model should not return rotor velocities without rotor_vel input"
+    # Check if the output is on the correct device, has the correct type and shape
     for dx, x in zip([dpos, dquat, dvel, dang_vel], [pos, quat, vel, ang_vel], strict=True):
         assert isinstance(dx, type(x))
         assert xp_device(dx) == xp_device(x)
@@ -82,19 +133,14 @@ def test_model_single_input_no_rotor_dynamics(model_name: str, model: Callable,
 @pytest.mark.unit
 @pytest.mark.parametrize("model_name, model", available_models.items())
 @pytest.mark.parametrize("drone_name", Constants.available_configs)
-def test_model_single_input_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
+def test_model_batched_rotor_dynamics(model_name: str, model: Callable, drone_name: str):
     skip_models_without_features(model, ["rotor_dynamics"])
 
-    pos, quat, vel, ang_vel, rotor_vel, _, _ = create_rnd_states()
-    commands = create_rnd_commands(dim=4)  # TODO make dependent on model
+    batch_shape = (10,)
+    pos, quat, vel, ang_vel, rotor_vel, _, _ = create_rnd_states(batch_shape)
+    cmd = create_rnd_commands(batch_shape, dim=4)  # TODO make dependent on model
     dpos, dquat, dvel, dang_vel, drotor_vel = model(
-        pos,
-        quat,
-        vel,
-        ang_vel,
-        commands,
-        Constants.from_config(drone_name, xp),
-        rotor_vel=rotor_vel,
+        pos, quat, vel, ang_vel, cmd, Constants.from_config(drone_name, xp), rotor_vel=rotor_vel
     )
     # Check if the output is on the correct device, has the correct type and shape
     for dx, x in zip(
