Support control_fn in ode_factory

src/jsrm/systems/planar_hsa.py

@@ -700,6 +700,7 @@ def operational_space_dynamical_matrices_fn(
 def ode_factory(
     dynamical_matrices_fn: Callable,
     params: Dict[str, Array],
+    control_fn: Optional[Callable] = None,
     consider_underactuation_model: bool = True,
     consider_hysteresis: bool = False,
 ) -> Callable[[float, Array], Array]:
@@ -718,6 +719,9 @@ def ode_factory(
             alpha_fn is a function to compute the actuation vector of shape (n_q). It has the following signature:
                 alpha_fn(phi) -> tau_q where phi is the twist angle vector of shape (n_phi, )
         params: Dictionary with robot parameters
+        control_fn: Callable that returns the forcing function of the form control_fn(t, x) -> u. If consider_underactuation_model is True,
+            then u is an array of shape (n_q, ) with the configuration-space torques. If consider_underactuation_model is False,
+            then u is an array of shape (n_phi, ) with the motor positions / twist angles of the proximal end of the rods.
         consider_underactuation_model: If True, the underactuation model is considered. Otherwise, the fully-actuated
             model is considered with the identity matrix as the actuation matrix.
         consider_hysteresis: If True, Bouc-Wen is used to model hysteresis. Otherwise, hysteresis will be neglected.
@@ -758,6 +762,9 @@ def ode_fn(t: float, x: Array, u: Array) -> Array:
             q, q_d = x[:n_q], x[n_q:]
             z = None
 
+        if control_fn is not None:
+            u = u + control_fn(t, x)
+
         if consider_underactuation_model is True:
             phi = u
             B, C, G, K, D, tau_q = dynamical_matrices_fn(params, q, q_d, z=z, phi=phi)