Use more generic parameter loading

Author: ratheron

drone_models/core.py

@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+import tomllib
 import warnings
 from functools import partial, wraps
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, ParamSpec, Protocol, TypeVar, runtime_checkable
 
+import numpy as np
+
 if TYPE_CHECKING:
     from types import ModuleType
 
@@ -138,18 +142,27 @@ def decorator(fn: Callable[P, R]) -> Callable[P, R]:
     return decorator
 
 
-# def load_xml_params(drone_model: str) -> dict[str, Array]:
-#     """TODO.
+def load_params(physics: str, drone_model: str, xp: ModuleType | None = None) -> dict:
+    """TODO.
 
-#     Args:
-#         model (str): _description_
+    Args:
+        physics: _description_
+        drone_model: _description_
+        xp: The array API module to use. If not provided, numpy is used.
 
-#     Returns:
-#         dict[str, Array]: _description_
-#     """
-#     drone_path = Path(__file__).parents[1] / path
-#     # read in all parameters from xml
-#     params = ET.parse(drone_path).findall(".//custom/numeric")
-#     # create a dict from parameters containing array of floats
-#     params = {p.get("name"): np.array(list(map(float, p.get("data").split()))) for p in params}
-#     return params
+    Returns:
+        dict[str, Array]: _description_
+    """
+    xp = np if xp is None else xp
+    with open(Path(__file__).parent / "data/params.toml", "rb") as f:
+        physical_params = tomllib.load(f)
+    if drone_model not in physical_params:
+        raise KeyError(f"Drone model `{drone_model}` not found in data/params.toml")
+    with open(Path(__file__).parent / f"{physics}/params.toml", "rb") as f:
+        model_params = tomllib.load(f)
+    if drone_model not in model_params:
+        raise KeyError(f"Drone model `{drone_model}` not found in model params.toml")
+    params = physical_params[drone_model] | model_params[drone_model]
+    params["J_inv"] = np.linalg.inv(params["J"])
+    params = {k: xp.asarray(v) for k, v in params.items()}  # if k in fields
+    return params

drone_models/data/params.toml

@@ -0,0 +1,58 @@
+# This file contains all the **physical** parameters for all models.
+
+[cf2x_L250]
+mass = 0.033
+gravity_vec = [0.0, 0.0, -9.81]
+J = [
+    [1.68e-5, 0.0, 0.0],
+    [0.0, 1.68e-5, 0.0],
+    [0.0, 0.0, 2.98e-5]
+]
+KF = 3.16e-10  # TODO: Update with new fit
+KM = 7.94e-12
+L = 0.03253
+mixing_matrix = [
+    [-1.0, -1.0, 1.0, 1.0],
+    [-1.0,  1.0,  1.0, -1.0],
+    [-1.0, 1.0,  -1.0, 1.0]
+]
+# TODO: This is an arbitrary number that leads to a working simulation. It MUST be fitted and
+# updated as soon as possible.
+thrust_tau = 0.03
+
+[cf2x_P250]
+mass = 0.03454
+gravity_vec = [0.0, 0.0, -9.81]
+# J from IROS 2022
+J = [
+    [1.4e-5, 0.0,    0.0],
+    [0.0, 1.4e-5, 0.0],
+    [0.0, 0.0, 2.17e-5]
+]
+KF = 5.792654323957372e-10
+KM = 7.94e-12
+L = 0.03253
+mixing_matrix = [
+    [-1.0, -1.0, 1.0, 1.0],
+    [-1.0,  1.0,  1.0, -1.0],
+    [-1.0, 1.0,  -1.0, 1.0]
+]
+thrust_tau = 7.2333515424212695
+
+[cf2x_T350]
+mass = 0.03454
+gravity_vec = [0.0, 0.0, -9.81]
+J = [
+    [1.4e-5, 0.0,    0.0],
+    [0.0, 1.4e-5, 0.0],
+    [0.0, 0.0, 2.17e-5]
+]
+KF = 1.068340683985525e-09
+KM = 7.94e-12
+L = 0.03253
+mixing_matrix = [
+    [-1.0, -1.0, 1.0, 1.0],
+    [-1.0,  1.0,  1.0, -1.0],
+    [-1.0, 1.0,  -1.0, 1.0]
+]
+thrust_tau = 8.719306881898989

drone_models/so_rpy/params.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-import tomllib
-from pathlib import Path
 from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
 
+from drone_models.core import load_params
+
 if TYPE_CHECKING:
     from array_api_typing import Array
 
@@ -28,10 +28,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:
-            raise KeyError(f"Drone model `{drone_model}` not found in params.toml")
-        params = {k: np.asarray(v) for k, v in params[drone_model].items()}
-        params["J_inv"] = np.linalg.inv(params["J"])
+        params = load_params("so_rpy", drone_model)
+        fields = SoRpyParams.__annotations__.keys()
+        params = {k: np.asarray(v) for k, v in params.items() if k in fields}
         return SoRpyParams(**params)