Fix some type hinting errors

Author: mstoelzle

examples/benchmark_planar_pcs.py

@@ -9,11 +9,11 @@
 import matplotlib.pyplot as plt
 import numpy as onp
 from pathlib import Path
-from typing import Callable, Dict, Optional, Literal
+from typing import Callable, Dict, Optional, Literal, Union
 
 import jsrm
 from jsrm import ode_factory
-from jsrm.systems import planar_pcs, planar_pcs_num
+from jsrm.systems import planar_pcs_num, planar_pcs_sym
 
 import time
 import pickle
@@ -25,21 +25,21 @@ def simulate_planar_pcs_value_eval(
     num_segments: int,
     type_of_derivation: Optional[Literal["symbolic", "numeric"]] = "symbolic",
     type_of_integration: Optional[Literal["gauss-legendre", "gauss-kronrad", "trapezoid"]] = "gauss-legendre",
-    param_integration: int = None,
+    param_integration: Optional[int] = None,
     type_of_jacobian: Optional[Literal["explicit", "autodiff"]] = "explicit",
-    robot_params: Dict[str, Array] = None,
-    strain_selector: Array = None,
-    q0: Array = None,
-    q_d0: Array = None,
+    robot_params: Optional[Dict[str, Array]] = None,
+    strain_selector: Optional[Array] = None,
+    q0: Optional[Array] = None,
+    q_d0: Optional[Array] = None,
     t: float = 1.0,
-    dt: float = None,
+    dt: Optional[float] = None,
     bool_print: bool = True,
     bool_plot: bool = True,
     bool_save_plot: bool = False,
     bool_save_video: bool = True,
     bool_save_res: bool = False,
-    results_path: str = None, 
-    results_path_extension: str = None
+    results_path: Optional[Union[str, Path]] = None, 
+    results_path_extension: Optional[str] = None
 ) -> Dict:
     """
     Simulate a planar PCS model. Save the video and figures.
@@ -81,7 +81,7 @@ def simulate_planar_pcs_value_eval(
             Defaults to None, which will use the default path.
         
     Returns:
-        TODO
+        simulation_dict (Dict): dictionary with the simulation results.
     """
 
     # ===================================================
@@ -253,13 +253,12 @@ def simulate_planar_pcs_value_eval(
             results_path = (results_path_parent / file_name).with_suffix(".pkl")
 
         if isinstance(results_path, str) or isinstance(results_path, Path):
-            results_path = Path(results_path)
-            if results_path.suffix != ".pkl":
+            results_path_obj = Path(results_path)
+            if results_path_obj.suffix != ".pkl":
                 raise ValueError(
-                    f"results_path must have the suffix .pkl, but got {results_path.suffix}"
+                    f"results_path must have the suffix .pkl, but got {results_path_obj.suffix}"
                 )
-            else:
-                results_path = Path(results_path)
+            results_path = results_path_obj
         else:
             raise TypeError(
                 f"results_path must be a string, but got {type(results_path).__name__}"
@@ -390,7 +389,7 @@ def draw_robot(
 
     if type_of_derivation == "symbolic":
         strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
-            planar_pcs.factory(sym_exp_filepath, strain_selector)
+            planar_pcs_sym.factory(sym_exp_filepath, strain_selector)
         )
     elif type_of_derivation == "numeric":
         strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
@@ -441,7 +440,7 @@ def draw_robot(
                 y0=x0, 
                 max_steps=None, 
                 saveat=SaveAt(ts=video_ts))
-    diffeqsolve_fn = jit(diffeqsolve_fn)
+    diffeqsolve_fn = jit(diffeqsolve_fn)  # type: ignore
 
     print("Solving the ODE ...")
     sol = diffeqsolve_fn()  
@@ -463,7 +462,7 @@ def draw_robot(
     forward_kinematics_fn_end_effector = partial(forward_kinematics_fn, robot_params, s=s_max)
 
     print("JIT-compiling the forward kinematics function...")
-    forward_kinematics_fn_end_effector = jit(forward_kinematics_fn_end_effector)
+    forward_kinematics_fn_end_effector = jit(forward_kinematics_fn_end_effector)  # type: ignore
     forward_kinematics_fn_end_effector = vmap(forward_kinematics_fn_end_effector)
 
     print("Computing the end-effector position along the trajectory...")
@@ -619,7 +618,7 @@ def draw_robot(
         )
 
         # Initialize the video
-        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # type: ignore
         video = cv2.VideoWriter(
             str(video_path),
             fourcc,
@@ -649,6 +648,7 @@ def draw_robot(
     # ===========================
     if bool_save_res:
         print("Saving the simulation results...")
+        assert results_path is not None, "results_path should not be None when saving results"
         with open(results_path, "wb") as f:
             pickle.dump(simulation_dict, f)
         print(f"Simulation results saved at {results_path} \n")
@@ -665,20 +665,20 @@ def simulate_planar_pcs_time_eval(
     num_segments: int,
     type_of_derivation: str = "symbolic",
     type_of_integration: str = "gauss-legendre",
-    param_integration: int = None,
+    param_integration: Optional[int] = None,
     type_of_jacobian: str = "explicit",
-    robot_params: Dict[str, Array] = None,
-    strain_selector: Array = None,
-    q0: Array = None,
-    q_d0: Array = None,
+    robot_params: Optional[Dict[str, Array]] = None,
+    strain_selector: Optional[Array] = None,
+    q0: Optional[Array] = None,
+    q_d0: Optional[Array] = None,
     t: float = 1.0,
-    dt: float = None,
+    dt: Optional[float] = None,
     bool_save_res: bool = False,
-    results_path: str = None, 
-    results_path_extension: str = None,
-    type_time = "once",
-    nb_eval : int = None,
-    nb_samples: int = None
+    results_path: Optional[Union[str, Path]] = None, 
+    results_path_extension: Optional[str] = None,
+    type_time: str = "once",
+    nb_eval: Optional[int] = None,
+    nb_samples: Optional[int] = None
 ) -> Dict:
     """
     Simulate a planar PCS model. Save the video and figures.
@@ -922,13 +922,12 @@ def simulate_planar_pcs_time_eval(
             results_path = (results_path_parent / file_name).with_suffix(".pkl")
 
         if isinstance(results_path, str) or isinstance(results_path, Path):
-            results_path = Path(results_path)
-            if results_path.suffix != ".pkl":
+            results_path_obj = Path(results_path)
+            if results_path_obj.suffix != ".pkl":
                 raise ValueError(
-                    f"results_path must have the suffix .pkl, but got {results_path.suffix}"
+                    f"results_path must have the suffix .pkl, but got {results_path_obj.suffix}"
                 )
-            else:
-                results_path = Path(results_path)
+            results_path = results_path_obj
         else:
             raise TypeError(
                 f"results_path must be a string, but got {type(results_path).__name__}"
@@ -982,7 +981,7 @@ def simulate_planar_pcs_time_eval(
     timer_start = time.time()
     if type_of_derivation == "symbolic":
         strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
-            planar_pcs.factory(sym_exp_filepath, strain_selector)
+            planar_pcs_sym.factory(sym_exp_filepath, strain_selector)
         )
     elif type_of_derivation == "numeric":
         strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
@@ -1330,7 +1329,7 @@ def time_diffeqsolve_over_time():
 
         # JIT the functions
         print("JIT-compiling the ODE function...")
-        diffeqsolve_fn = jit(diffeqsolve_fn)
+        diffeqsolve_fn = jit(diffeqsolve_fn)  # type: ignore
 
         # First evaluation of the ODE to trigger JIT compilation
         print("Solving the ODE for the first time (JIT-compilation)...")
@@ -1579,6 +1578,7 @@ def time_kinetic_energy_over_time():
     # ===========================
     if bool_save_res:
         print("Saving the simulation results...")
+        assert results_path is not None, "results_path should not be None when saving results"
         with open(results_path, "wb") as f:
             pickle.dump(simulation_dict, f)
         print(f"Simulation results saved at {results_path} \n")

src/jsrm/systems/planar_pcs_num.py

@@ -5,7 +5,7 @@
 from quadax import GaussKronrodRule
 
 import numpy as onp
-from typing import Callable, Dict, Tuple, Optional, Literal
+from typing import Callable, Dict, Tuple, Optional, Literal, Any
 
 from .utils import (
     compute_strain_basis,
@@ -33,25 +33,25 @@
 
 def factory(
     num_segments: int, 
-    strain_selector: Array=None,
+    strain_selector: Optional[Array] = None,
     xi_eq: Optional[Array] = None,
     stiffness_fn: Optional[Callable] = None,
     actuation_mapping_fn: Optional[Callable] = None,
     global_eps: float = jnp.finfo(jnp.float32).eps, 
     integration_type: Optional[Literal["gauss-legendre", "gauss-kronrad", "trapezoid"]] = "gauss-legendre",
-    param_integration: int = None,
+    param_integration: Optional[int] = None,
     jacobian_type: Optional[Literal["explicit", "autodiff"]] = "explicit"
-    ) -> Tuple[
+) -> Tuple[
         Array,
         Callable[
-            [Dict[str, Array], Array, Array, Optional[float]], 
+            [Dict[str, Array], Array, Array, float], 
             Array
         ],
         Callable[
-            [Dict[str, Array], Array, Array, Optional[float]],
+            [Dict[str, Array], Array, Array, float],
             Tuple[Array, Array, Array, Array, Array, Array],
         ],
-        Dict[str, Callable],
+        Dict[str, Callable[..., Any]],
     ]:
     """
     Factory function to create the forward kinematics function for a planar robot.
@@ -190,7 +190,7 @@ def stiffness_fn(
                 S = B_xi.T @ S @ B_xi
 
             return S
-    if not isinstance(stiffness_fn, Callable):
+    if not isinstance(stiffness_fn, callable):
         raise TypeError(f"stiffness_fn must be a callable, but got {type(stiffness_fn).__name__}")
 
     # Actuation mapping function
@@ -295,7 +295,7 @@ def apply_eps_to_bend_strains(
     def classify_segment(
         params: Dict[str, Array], 
         s: Array
-        ) -> Tuple[Array, Array]:
+        ) -> Tuple[Array, Array, Array]:
         """
         Classify the point along the robot to the corresponding segment.
 
@@ -306,6 +306,7 @@ def classify_segment(
         Returns: 
             segment_idx (Array): index of the segment where the point is located
             s_segment (Array): point coordinate along the segment in the interval [0, l_segment]
+            l_cum (Array): cumulative length of the segments starting with 0
         """
         l = params["l"]
 
@@ -501,7 +502,7 @@ def J_explicit_local(
         def J_i(
             tuple_J_prev: Array,
             i: int
-        ) -> Array:
+        ) -> Tuple[Tuple[Array, Array], Array]:
             J_prev_Lprev, _ = tuple_J_prev
 
             start_index = 3 * i
@@ -534,7 +535,8 @@ def J_i(
         _, J_array = lax.scan(
             f = J_i,
             init = tuple_J_0, 
-            xs = jnp.arange(1, num_segments))
+            xs = jnp.arange(1, num_segments)
+        )
 
         # Add the initial condition to the Jacobian array
         J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0)
@@ -597,7 +599,7 @@ def J_explicit_global(
         def J_i(
             tuple_J_prev: Array,
             i: int
-        ) -> Array:
+        ) -> Tuple[Tuple[Array, Array], Array]:
             J_prev_Lprev, _ = tuple_J_prev
 
             start_index = 3 * i
@@ -708,7 +710,7 @@ def J_Jd_explicit_local(
         xi_d: Array,
         s: Array, 
         eps: float = global_eps
-    ) -> Array:
+    ) -> Tuple[Array, Array]:
         """
         Compute the body-frame jacobian and its derivative with respect to the strain vector 
         at a given point s using explicit expression in SE(2).
@@ -753,7 +755,7 @@ def J_Jd_explicit_local(
         def J_i(
             tuple_J_prev: Array,
             i: int
-        ) -> Array:
+        ) -> Tuple[Tuple[Array, Array], Array]:
             J_prev_Lprev, _ = tuple_J_prev
 
             start_index = 3 * i
@@ -837,7 +839,7 @@ def J_Jd_explicit_global(
         xi_d: Array,
         s: Array,
         eps: float = global_eps
-    ) -> Array:
+    ) -> Tuple[Array, Array]:
         """
         Compute the inertial-frame jacobian and its derivative with respect to the strain vector 
         at a given point s using explicit expression in SE(2).
@@ -882,7 +884,7 @@ def J_Jd_explicit_global(
         def J_i(
             tuple_J_prev: Array,
             i: int
-        ) -> Array:
+        ) -> Tuple[Tuple[Array, Array], Array]:
             J_prev_Lprev, _ = tuple_J_prev
 
             start_index = 3 * i
@@ -915,7 +917,8 @@ def J_i(
         _, J_array = lax.scan(
             f = J_i,
             init = tuple_J_0, 
-            xs = jnp.arange(1, num_segments))
+            xs = jnp.arange(1, num_segments)
+        )
 
         # Add the initial condition to the Jacobian array
         J_array = jnp.concatenate([J_0_s[jnp.newaxis, ...], J_array], axis=0)
@@ -1705,7 +1708,7 @@ def operational_space_dynamical_matrices_fn(
 
         return Lambda, mu, J, J_d, JB_pinv
 
-    auxiliary_fns = {
+    auxiliary_fns: Dict[str, Callable[..., Any]] = {
         "apply_eps_to_bend_strains": apply_eps_to_bend_strains,
         "classify_segment": classify_segment,
         "stiffness_fn": stiffness_fn,
@@ -1717,4 +1720,4 @@ def operational_space_dynamical_matrices_fn(
         "operational_space_dynamical_matrices_fn": operational_space_dynamical_matrices_fn,
     }
 
-    return B_xi, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns
+    return B_xi, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns

tests/test_fwd_kine_eps_planar_pcs.py

@@ -5,7 +5,7 @@
 from matplotlib import rc
 rc('animation', html='html5')
 
-from jsrm.systems import planar_pcs, planar_pcs_num
+from jsrm.systems import planar_pcs_num, planar_pcs_sym
 from pathlib import Path
 import jsrm
 from tqdm import tqdm
@@ -33,7 +33,7 @@
 def get_fwd_kine_fn(jacobian_type):
     if jacobian_type == "symbolic":
         sym_exp_filepath = Path(jsrm.__file__).parent / "symbolic_expressions" / f"planar_pcs_ns-{num_segments}.dill"
-        _, fwd, _, _ = planar_pcs.factory(sym_exp_filepath, strain_selector)
+        _, fwd, _, _ = planar_pcs_sym.factory(sym_exp_filepath, strain_selector)
     else:
         _, fwd, _, _ = planar_pcs_num.factory(
             num_segments, strain_selector,

tests/test_jacobian_eps_planar_pcs.py

@@ -5,7 +5,7 @@
 from matplotlib import rc
 rc('animation', html='html5')
 
-from jsrm.systems import planar_pcs, planar_pcs_num
+from jsrm.systems import planar_pcs_num, planar_pcs_sym
 from pathlib import Path
 import jsrm
 from tqdm import tqdm
@@ -33,7 +33,7 @@
 def get_jacobian_fn(jacobian_type):
     if jacobian_type == "symbolic":
         sym_exp_filepath = Path(jsrm.__file__).parent / "symbolic_expressions" / f"planar_pcs_ns-{num_segments}.dill"
-        _, _, _, aux = planar_pcs.factory(sym_exp_filepath, strain_selector)
+        _, _, _, aux = planar_pcs_sym.factory(sym_exp_filepath, strain_selector)
     else:
         _, _, _, aux = planar_pcs_num.factory(
             num_segments, strain_selector,