Clean-up interface of hsa ode_factory

- Allow `u` to be also a normal argument (instead of only keyword arguments)
- Pass `phi` as `args` keyword when calling `diffeqsolve`

examples/simulate_planar_hsa.py

@@ -286,7 +286,7 @@ def chi2u(chi: Array) -> Array:
     x0 = x0.at[: q0.shape[0]].set(q0)  # set initial configuration
 
     ode_fn = planar_hsa.ode_factory(dynamical_matrices_fn, params)
-    ode_term = ODETerm(partial(ode_fn, u=phi))
+    ode_term = ODETerm(ode_fn)
 
     sol = diffeqsolve(
         ode_term,
@@ -295,6 +295,7 @@ def chi2u(chi: Array) -> Array:
         t1=ts[-1],
         dt0=dt,
         y0=x0,
+        args=phi,
         max_steps=None,
         saveat=SaveAt(ts=video_ts),
     )

src/jsrm/systems/planar_hsa.py

@@ -634,13 +634,12 @@ def ode_factory(
     num_rods = params["rout"].shape[0] * params["rout"].shape[1]
 
     @jit
-    def ode_fn(t: float, x: Array, *args, u: Array) -> Array:
+    def ode_fn(t: float, x: Array, u: Array) -> Array:
         """
         ODE of the dynamical Lagrangian system.
         Args:
             t: time
             x: state vector of shape (2 * n_q, )
-            args: additional arguments
             u: input to the system.
                 - if consider_underactuation_model is True, then this is an array of shape (n_phi) with
                     motor positions / twist angles of the proximal end of the rods