start porting to new drone model

Author: yufei4hua

drone_estimators/estimator.py

@@ -14,7 +14,7 @@
     ukf_predict_correct,
 )
 from drone_estimators.structs import SigmaPointsSettings, UKFData, UKFSettings
-from drone_estimators.utils.dynamics import get_dynamics, observation_function
+from drone_estimators.utils.dynamics import dynamics_function, observation_function
 
 if TYPE_CHECKING:
     from array_api_typing import Array
@@ -85,8 +85,9 @@ def __init__(
             estimate_dist_t: If the disturbance torques should be estimated, defaults to False.
             initial_obs: Optional, the initial observation of the environment's state. See the environment's observation space for details.
         """
-        fx = get_dynamics(model, config)
-        hx = observation_function
+        
+        fx = dynamics_function(model, config)
+        hx = observation_function()
         # fx = jax.jit(dynamics_numeric(model, config))
         # hx = jax.jit(observation_function)
 

drone_estimators/utils/dynamics.py

@@ -5,44 +5,21 @@
 from functools import partial
 from typing import TYPE_CHECKING
 
-from drone_models.so_rpy_rotor_drag import dynamics as so_rpy_rotor_drag_dynamics
-from drone_models.so_rpy_rotor_drag.params import SoRpyRotorDragParams
+from drone_models.core import parametrize
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
     from array_api_typing import Array
 
 
-def get_dynamics(
-    model_type: str, drone_model: str
-) -> Callable[
-    [Array, Array, Array, Array, Array, Array | None, Array | None, Array | None],
-    tuple[Array, Array, Array, Array, Array | None],
-]:
-    """Simplifies drone-model dynamics by adding all arguments for the given drone_model."""
-    match model_type:
-        case "so_rpy_rotor_drag_dynamics":
-            params = SoRpyRotorDragParams.load(drone_model)
-            return partial(
-                so_rpy_rotor_drag_dynamics,
-                mass=params.mass,
-                gravity_vec=params.gravity_vec,
-                J=params.J,
-                J_inv=params.J_inv,
-                KF=params.KF,
-                KM=params.KM,
-                thrust_time_coef=params.thrust_time_coef,
-                acc_coef=params.acc_coef,
-                cmd_f_coef=params.cmd_f_coef,
-                rpy_coef=params.rpy_coef,
-                rpy_rates_coef=params.rpy_rates_coef,
-                cmd_rpy_coef=params.cmd_rpy_coef,
-                drag_linear_coef=params.drag_linear_coef,
-                drag_square_coef=params.drag_square_coef,
-            )
-        case _:
-            raise NotImplementedError(f"Model type {model_type} not supported.")
+def dynamics_function(model:str, config:str):
+    """TODO."""
+    # Idea: 
+    # from drone_models.{model}.model import dynamics as fn
+    
+    # return parametrize(fn, config)
+    ...
 
 
 def observation_function(