Simplify parametrize structure

Author: ratheron

drone_models/core.py

@@ -2,11 +2,12 @@
 
 from __future__ import annotations
 
+import inspect
 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
+from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar
 
 import numpy as np
 
@@ -49,16 +50,6 @@ def wrapper(
     return decorator
 
 
-model_parameter_registry: dict[str, type[ModelParams]] = {}
-
-
-def named_tuple2xp(params: ModelParams, xp: ModuleType, device: str | None = None) -> ModelParams:
-    """Convert a named tuple to an array API framework."""
-    return params.__class__(
-        **{k: xp.asarray(v, device=device) for k, v in params._asdict().items()}
-    )
-
-
 def parametrize(
     fn: Callable[P, R], drone_model: str, xp: ModuleType | None = None, device: str | None = None
 ) -> Callable[P, R]:
@@ -87,59 +78,28 @@ def parametrize(
     Returns:
         The parametrized controller function with all keyword argument only parameters filled in.
     """
-    model_id = fn.__module__ + "." + fn.__name__
     try:
-        params = model_parameter_registry[model_id].load(drone_model)
-        if xp is not None:  # Convert to any array API framework
-            params = named_tuple2xp(params, xp=xp, device=device)
+        xp = np if xp is None else xp
+        # physics = Path(sys.modules[fn.__module__].__file__).parent.name
+        physics = fn.__module__.split(".")[-2]
+        sig = inspect.signature(fn)
+        kwonly_params = [
+            name
+            for name, param in sig.parameters.items()
+            if param.kind == inspect.Parameter.KEYWORD_ONLY
+        ]
+        params = load_params(physics, drone_model)
+
+        params = {k: xp.asarray(v, device=device) for k, v in params.items() if k in kwonly_params}
+        # if xp is not None:  # Convert to any array API framework
+        #     params = named_tuple2xp(params, xp=xp, device=device)
     except KeyError as e:
         raise KeyError(
-            f"Model `{model_id}` does not exist in the parameter registry for drone `{drone_model}`"
+            f"Model `{physics}` does not exist in the parameter registry for drone `{drone_model}`"
         ) from e
     except ValueError as e:
-        raise ValueError(f"Drone model `{drone_model}` not supported for `{model_id}`") from e
-    return partial(fn, **params._asdict())
-
-
-@runtime_checkable
-class ModelParams(Protocol):
-    """Protocol for model parameters."""
-
-    @staticmethod
-    def load(drone_model: str) -> ModelParams:
-        """Load the parameters from the config file."""
-
-    def _asdict(self) -> dict[str, Any]:
-        """Convert the parameters to a dictionary."""
-
-
-def register_model_parameters(
-    params: ModelParams | type[ModelParams],
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    """Register the default model parameters for this model.
-
-    Warning:
-        The model parameters **must** be a named tuple with a function `load` that takes in the
-        drone model name and returns an instance of itself, or a class that implements the
-        ModelParams protocol.
-
-    Args:
-        params: The model parameter type.
-
-    Returns:
-        A decorator function that registers the parameters and returns the function unchanged.
-    """
-    if not isinstance(params, ModelParams):
-        raise ValueError(f"{params} does not implement the ModelParams protocol")
-
-    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
-        controller_id = fn.__module__ + "." + fn.__name__
-        if controller_id in model_parameter_registry:
-            raise ValueError(f"Model `{controller_id}` already registered")
-        model_parameter_registry[controller_id] = params
-        return fn
-
-    return decorator
+        raise ValueError(f"Drone model `{drone_model}` not supported for `{physics}`") from e
+    return partial(fn, **params)
 
 
 def load_params(physics: str, drone_model: str, xp: ModuleType | None = None) -> dict:

drone_models/first_principles/model.py

@@ -9,15 +9,13 @@
 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.first_principles.params import FirstPrinciplesParams
+from drone_models.core import supports
 from drone_models.utils import rotation, to_xp
 
 if TYPE_CHECKING:
     from array_api_typing import Array
 
 
-@register_model_parameters(FirstPrinciplesParams)
 @supports(rotor_dynamics=True)
 def dynamics(
     pos: Array,
@@ -110,7 +108,6 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-@register_model_parameters(FirstPrinciplesParams)
 def symbolic_dynamics(
     model_rotor_vel: bool = True,
     model_dist_f: bool = False,

drone_models/first_principles/params.py

@@ -9,15 +9,13 @@
 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.params import SoRpyParams
+from drone_models.core import supports
 from drone_models.utils import rotation, to_xp
 
 if TYPE_CHECKING:
     from array_api_typing import Array
 
 
-@register_model_parameters(SoRpyParams)
 @supports(rotor_dynamics=False)
 def dynamics(
     pos: Array,
@@ -108,7 +106,6 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-@register_model_parameters(SoRpyParams)
 def symbolic_dynamics(
     model_rotor_vel: bool = False,
     model_dist_f: bool = False,

drone_models/so_rpy/model.py

@@ -10,16 +10,14 @@
 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.core import supports
 from drone_models.transform import motor_force2rotor_vel
 from drone_models.utils import rotation, to_xp
 
 if TYPE_CHECKING:
     from array_api_typing import Array
 
 
-@register_model_parameters(SoRpyRotorParams)
 @supports(rotor_dynamics=True)
 def dynamics(
     pos: Array,
@@ -128,7 +126,6 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-@register_model_parameters(SoRpyRotorParams)
 def symbolic_dynamics(
     model_rotor_vel: bool = False,
     model_dist_f: bool = False,

drone_models/so_rpy/params.py

@@ -10,15 +10,13 @@
 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_drag.params import SoRpyRotorDragParams
+from drone_models.core import supports
 from drone_models.utils import rotation, to_xp
 
 if TYPE_CHECKING:
     from array_api_typing import Array
 
 
-@register_model_parameters(SoRpyRotorDragParams)
 @supports(rotor_dynamics=True)
 def dynamics(
     pos: Array,
@@ -142,7 +140,6 @@ def dynamics(
     return pos_dot, quat_dot, vel_dot, ang_vel_dot, rotor_vel_dot
 
 
-@register_model_parameters(SoRpyRotorDragParams)
 def symbolic_dynamics(
     model_rotor_vel: bool = True,
     model_dist_f: bool = False,

drone_models/so_rpy_rotor/model.py

@@ -0,0 +1,19 @@
+"""Tests of the parametrization of the models."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import pytest
+
+from drone_models import available_models
+from drone_models.core import parametrize
+from drone_models.drones import available_drones
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("model_name, model", available_models.items())
+@pytest.mark.parametrize("drone_type", available_drones)
+def test_model_parametrization(model_name: str, model: Callable, drone_type: str):
+    """TODO."""
+    parametrize(model, drone_type)