Fix symbolic models

Author: amacati

drone_models/first_principles/__init__.py

@@ -3,6 +3,6 @@
 TODO: Add description.
 """
 
-from drone_models.first_principles.model import dynamics, dynamics_symbolic
+from drone_models.first_principles.model import dynamics, symbolic_dynamics
 
-__all__ = ["dynamics", "dynamics_symbolic"]
+__all__ = ["dynamics", "symbolic_dynamics"]

drone_models/first_principles/model.py

@@ -114,34 +114,35 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-def dynamics_symbolic(
-    calc_rotor_vel: bool = True,
-    calc_dist_f: bool = False,
-    calc_dist_t: bool = False,
+@register_model_parameters(FirstPrinciplesParams)
+def symbolic_dynamics(
+    model_rotor_vel: bool = True,
+    model_dist_f: bool = False,
+    model_dist_t: bool = False,
     *,
-    thrust_tau: float,
-    KF: float,
-    KM: float,
-    L: float,
-    sign_matrix: Array,
-    gravity_vec: Array,
     mass: float,
+    gravity_vec: Array,
     J: Array,
     J_inv: Array,
+    KF: float,
+    KM: float,
+    L: float,
+    mixing_matrix: Array,
+    thrust_tau: float,
 ) -> tuple[cs.MX, cs.MX, cs.MX, cs.MX]:
     """TODO take from numeric."""
     # States and Inputs
     X = cs.vertcat(symbols.pos, symbols.quat, symbols.vel, symbols.ang_vel)
-    if calc_rotor_vel:
+    if model_rotor_vel:
         X = cs.vertcat(X, symbols.rotor_vel)
-    if calc_dist_f:
+    if model_dist_f:
         X = cs.vertcat(X, symbols.dist_f)
-    if calc_dist_t:
+    if model_dist_t:
         X = cs.vertcat(X, symbols.dist_t)
     U = symbols.cmd_rotor_vel
 
     # Defining the dynamics function
-    if calc_rotor_vel:
+    if model_rotor_vel:
         # Thrust dynamics
         rotor_vel_dot = 1 / thrust_tau * (U - symbols.rotor_vel) - 1 / KM * symbols.rotor_vel**2
         forces_motor = KF * symbols.rotor_vel**2
@@ -150,12 +151,12 @@ def dynamics_symbolic(
 
     # Creating force and torque vector
     forces_motor_vec = cs.vertcat(0, 0, cs.sum1(forces_motor))
-    torques_motor_vec = forces_motor @ sign_matrix * cs.vertcat(L, L, KM / KF)
+    torques_motor_vec = (forces_motor.T @ mixing_matrix).T * cs.vertcat(L, L, KM / KF)
 
     # Linear equation of motion
     forces_motor_vec_world = symbols.rot @ forces_motor_vec
     forces_sum = forces_motor_vec_world + gravity_vec * mass
-    if calc_dist_f is True:
+    if model_dist_f:
         forces_sum = forces_sum + symbols.dist_f
 
     pos_dot = symbols.vel
@@ -167,11 +168,11 @@ def dynamics_symbolic(
     )
     quat_dot = 0.5 * (xi @ symbols.quat)
     torques_sum = torques_motor_vec
-    if calc_dist_t:
+    if model_dist_t:
         torques_sum = torques_sum + symbols.rot.T @ symbols.dist_t
     ang_vel_dot = J_inv @ (torques_sum - cs.cross(symbols.ang_vel, J @ symbols.ang_vel))
 
-    if calc_rotor_vel:
+    if model_rotor_vel:
         X_dot = cs.vertcat(pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot)
     else:
         X_dot = cs.vertcat(pos_dot, quat_dot, vel_dot, ang_vel_dot)

drone_models/so_rpy/__init__.py

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

drone_models/so_rpy/model.py

@@ -16,8 +16,6 @@
 if TYPE_CHECKING:
     from array_api_typing import Array
 
-    from drone_models.utils.constants import Constants
-
 
 @register_model_parameters(SoRpyParams)
 @supports(rotor_dynamics=False)
@@ -98,8 +96,10 @@ def dynamics(
     if dist_t is not None:
         # 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 = (J @ ang_vel_dot[..., None])[..., 0]
+        torque = torque + xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
+        # adding torque. TODO: This should be a linear transformation. Can't we just transform the
+        # disturbance torque to an ang_vel_dot summand directly?
         torque = torque + rot.apply(dist_t, inverse=True)
         # back to angular acceleration
         torque = torque - xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
@@ -108,38 +108,46 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-def dynamics_symbolic(
-    constants: Constants,
-    calc_rotor_vel: bool = False,
-    calc_dist_f: bool = False,
-    calc_dist_t: bool = False,
+@register_model_parameters(SoRpyParams)
+def symbolic_dynamics(
+    model_rotor_vel: bool = False,
+    model_dist_f: bool = False,
+    model_dist_t: bool = False,
+    *,
+    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,
 ) -> tuple[cs.MX, cs.MX, cs.MX, cs.MX]:
     """The fitted double integrator (DI) model with optional motor delay (D).
 
     TODO.
     """
     # States and Inputs
     X = cs.vertcat(symbols.pos, symbols.quat, symbols.vel, symbols.ang_vel)
-    if calc_rotor_vel:
+    if model_rotor_vel:
         X = cs.vertcat(X, symbols.rotor_vel)
-    if calc_dist_f:
+    if model_dist_f:
         X = cs.vertcat(X, symbols.dist_f)
-    if calc_dist_t:
+    if model_dist_t:
         X = cs.vertcat(X, symbols.dist_t)
     U = cs.vertcat(symbols.cmd_roll, symbols.cmd_pitch, symbols.cmd_yaw, symbols.cmd_thrust)
 
     # Defining the dynamics function
     # Creating force vector
-    forces_motor_vec = cs.vertcat(
-        0, 0, constants.DI_ACC[0] + constants.DI_ACC[1] * symbols.cmd_thrust
-    )
+    forces_motor_vec = cs.vertcat(0, 0, acc_coef + cmd_f_coef * symbols.cmd_thrust)
 
     # Linear equation of motion
     pos_dot = symbols.vel
-    vel_dot = symbols.rot @ forces_motor_vec / constants.MASS + constants.GRAVITY_VEC
-    if calc_dist_f:
+    vel_dot = symbols.rot @ forces_motor_vec / mass + gravity_vec
+    if model_dist_f:
         # Adding force disturbances to the state
-        vel_dot = vel_dot + symbols.dist_f / constants.MASS
+        vel_dot = vel_dot + symbols.dist_f / mass
 
     # Rotational equation of motion
     euler_angles = rotation.cs_quat2euler(symbols.quat)
@@ -150,24 +158,19 @@ def dynamics_symbolic(
     quat_dot = 0.5 * (xi @ symbols.quat)
     rpy_rates = rotation.cs_ang_vel2rpy_rates(symbols.quat, symbols.ang_vel)
     rpy_rates_dot = (
-        constants.DI_PARAMS[:, 0] * euler_angles
-        + constants.DI_PARAMS[:, 1] * rpy_rates
-        + constants.DI_PARAMS[:, 2]
-        * cs.vertcat(symbols.cmd_roll, symbols.cmd_pitch, symbols.cmd_yaw)
+        rpy_coef * euler_angles
+        + rpy_rates_coef * rpy_rates
+        + cmd_rpy_coef * cs.vertcat(symbols.cmd_roll, symbols.cmd_pitch, symbols.cmd_yaw)
     )
     ang_vel_dot = rotation.cs_rpy_rates_deriv2ang_vel_deriv(symbols.quat, rpy_rates, rpy_rates_dot)
-    if calc_dist_t:
+    if model_dist_t:
         # adding torque disturbances to the state
         # angular acceleration can be converted to total torque
-        torque = constants.J @ ang_vel_dot + cs.cross(
-            symbols.ang_vel, constants.J @ symbols.ang_vel
-        )
+        torque = J @ ang_vel_dot + cs.cross(symbols.ang_vel, J @ symbols.ang_vel)
         # adding torque
         torque = torque + symbols.rot.T @ symbols.dist_t
         # back to angular acceleration
-        ang_vel_dot = constants.J_INV @ (
-            torque - cs.cross(symbols.ang_vel, constants.J @ symbols.ang_vel)
-        )
+        ang_vel_dot = J_inv @ (torque - cs.cross(symbols.ang_vel, J @ symbols.ang_vel))
 
     X_dot = cs.vertcat(pos_dot, quat_dot, vel_dot, ang_vel_dot)
     Y = cs.vertcat(symbols.pos, symbols.quat)

drone_models/so_rpy_rotor/__init__.py

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

drone_models/so_rpy_rotor/model.py

@@ -18,8 +18,6 @@
 if TYPE_CHECKING:
     from array_api_typing import Array
 
-    from drone_models.utils.constants import Constants
-
 
 @register_model_parameters(SoRpyRotorParams)
 @supports(rotor_dynamics=True)
@@ -35,10 +33,10 @@ def dynamics(
     *,
     mass: float,
     gravity_vec: Array,
-    KF: Array,
-    KM: Array,
     J: Array,
     J_inv: Array,
+    KF: Array,
+    KM: Array,
     rotor_coef: Array,
     acc_coef: Array,
     cmd_f_coef: Array,
@@ -62,10 +60,10 @@ def dynamics(
         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).
+        KF: Motor force constant (N/rad^2).
+        KM: Motor torque constant (Nm/rad^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).
@@ -118,54 +116,71 @@ def dynamics(
     if dist_t is not None:
         # 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
+        torque = (J @ ang_vel_dot[..., None])[..., 0]
+        torque = torque + xp.linalg.cross(ang_vel, (J @ ang_vel[..., None])[..., 0])
+        # adding torque. TODO: This should be a linear transformation. Can't we just transform the
+        # disturbance torque to an ang_vel_dot summand directly?
+        torque = torque + rot.apply(dist_t, inverse=True)
         # back to angular acceleration
         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
 
 
-def dynamics_symbolic(
-    constants: Constants,
-    calc_rotor_vel: bool = False,
-    calc_dist_f: bool = False,
-    calc_dist_t: bool = False,
+@register_model_parameters(SoRpyRotorParams)
+def symbolic_dynamics(
+    model_rotor_vel: bool = False,
+    model_dist_f: bool = False,
+    model_dist_t: bool = False,
+    *,
+    mass: float,
+    gravity_vec: Array,
+    J: Array,
+    J_inv: Array,
+    KF: Array,
+    KM: Array,
+    rotor_coef: Array,
+    acc_coef: Array,
+    cmd_f_coef: Array,
+    rpy_coef: Array,
+    rpy_rates_coef: Array,
+    cmd_rpy_coef: Array,
 ) -> tuple[cs.MX, cs.MX, cs.MX, cs.MX]:
     """The fitted double integrator (DI) model with optional motor delay (D).
 
     TODO.
     """
     # States and Inputs
     X = cs.vertcat(symbols.pos, symbols.quat, symbols.vel, symbols.ang_vel)
-    if calc_rotor_vel:
+    if model_rotor_vel:
         X = cs.vertcat(X, symbols.rotor_vel)
-    if calc_dist_f:
+    if model_dist_f:
         X = cs.vertcat(X, symbols.dist_f)
-    if calc_dist_t:
+    if model_dist_t:
         X = cs.vertcat(X, symbols.dist_t)
     U = cs.vertcat(symbols.cmd_roll, symbols.cmd_pitch, symbols.cmd_yaw, symbols.cmd_thrust)
     cmd_rpy = cs.vertcat(symbols.cmd_roll, symbols.cmd_pitch, symbols.cmd_yaw)
 
     # Defining the dynamics function
-    if calc_rotor_vel:
+    if model_rotor_vel:
         # motor_force2rotor_vel
-        cmd_rotor_vel = cs.sqrt(symbols.cmd_thrust / 4 / constants.KF)
-        rotor_vel_dot = 1 / constants.DI_D_ACC[2] * (cmd_rotor_vel - symbols.rotor_vel)
-        thrust = constants.KF * cs.sum1(symbols.rotor_vel**2)
+        cmd_rotor_vel = cs.sqrt(symbols.cmd_thrust / 4 / KF)
+        rotor_vel_dot = (
+            1 / rotor_coef * (cmd_rotor_vel - symbols.rotor_vel) - KM * symbols.rotor_vel**2
+        )
+        forces_motor = KF * cs.sum1(symbols.rotor_vel**2)
     else:
-        thrust = symbols.cmd_thrust
+        forces_motor = symbols.cmd_thrust
     # Creating force vector
-    forces_motor_vec = cs.vertcat(0, 0, constants.DI_D_ACC[0] + constants.DI_D_ACC[1] * thrust)
+    forces_motor_vec = cs.vertcat(0, 0, acc_coef + cmd_f_coef * forces_motor)
 
     # Linear equation of motion
     pos_dot = symbols.vel
-    vel_dot = symbols.rot @ forces_motor_vec / constants.MASS + constants.GRAVITY_VEC
-    if calc_dist_f:
+    vel_dot = symbols.rot @ forces_motor_vec / mass + gravity_vec
+    if model_dist_f:
         # Adding force disturbances to the state
-        vel_dot = vel_dot + symbols.dist_f / constants.MASS
+        vel_dot = vel_dot + symbols.dist_f / mass
 
     # Rotational equation of motion
     euler_angles = rotation.cs_quat2euler(symbols.quat)
@@ -175,26 +190,18 @@ def dynamics_symbolic(
     )
     quat_dot = 0.5 * (xi @ symbols.quat)
     rpy_rates = rotation.cs_ang_vel2rpy_rates(symbols.quat, symbols.ang_vel)
-    rpy_rates_dot = (
-        constants.DI_D_PARAMS[:, 0] * euler_angles
-        + constants.DI_D_PARAMS[:, 1] * rpy_rates
-        + constants.DI_D_PARAMS[:, 2] * cmd_rpy
-    )
+    rpy_rates_dot = rpy_coef * euler_angles + rpy_rates_coef * rpy_rates + cmd_rpy_coef * cmd_rpy
     ang_vel_dot = rotation.cs_rpy_rates_deriv2ang_vel_deriv(symbols.quat, rpy_rates, rpy_rates_dot)
-    if calc_dist_t:
+    if model_dist_t:
         # adding torque disturbances to the state
         # angular acceleration can be converted to total torque
-        torque = constants.J @ ang_vel_dot + cs.cross(
-            symbols.ang_vel, constants.J @ symbols.ang_vel
-        )
+        torque = J @ ang_vel_dot + cs.cross(symbols.ang_vel, J @ symbols.ang_vel)
         # adding torque
         torque = torque + symbols.rot.T @ symbols.dist_t
         # back to angular acceleration
-        ang_vel_dot = constants.J_INV @ (
-            torque - cs.cross(symbols.ang_vel, constants.J @ symbols.ang_vel)
-        )
+        ang_vel_dot = J_inv @ (torque - cs.cross(symbols.ang_vel, J @ symbols.ang_vel))
 
-    if calc_rotor_vel:
+    if model_rotor_vel:
         X_dot = cs.vertcat(pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot)
     else:
         X_dot = cs.vertcat(pos_dot, quat_dot, vel_dot, ang_vel_dot)

drone_models/so_rpy_rotor_drag/__init__.py

@@ -3,6 +3,6 @@
 TODO: Add description.
 """
 
-from drone_models.so_rpy_rotor_drag.model import dynamics, dynamics_symbolic
+from drone_models.so_rpy_rotor_drag.model import dynamics, symbolic_dynamics
 
-__all__ = ["dynamics", "dynamics_symbolic"]
+__all__ = ["dynamics", "symbolic_dynamics"]