Merge pull request #304 from NeLy-EPFL/warp-typehints

Update type hints to be PEP 484 compliant

Author: sibocw

scripts/replay_behavior_gpu.py

@@ -79,10 +79,10 @@
 
 @wp.kernel
 def record_joint_angles_kernel(
-    qpos: wp.array2d(dtype=wp.float32),  # type: ignore  # (n_worlds, nq)
-    qpos_adrs: wp.array(dtype=wp.int32),  # type: ignore  # (n_jointdofs,)
-    step_counter: wp.array(dtype=wp.int32),  # type: ignore
-    recorded: wp.array3d(dtype=wp.float32),  # type: ignore  # (n_steps, n_worlds, n_dofs)
+    qpos: wp.array2d[float],  # (n_worlds, nq)
+    qpos_adrs: wp.array[int],  # (n_jointdofs,)
+    step_counter: wp.array[int],
+    recorded: wp.array3d[float],  # (n_steps, n_worlds, n_dofs)
 ):
     """Gather this step's joint angles into a pre-allocated, GPU-resident buffer.
 

src/flygym/compose/world/base_world.py

@@ -81,7 +81,7 @@ def mjcf_root(self) -> mj.MjSpec:
     def fly_lookup(self) -> dict[str, BaseFly]:
         """Lookup for `Fly` objects in the world, keyed by fly name."""
         return self._fly_lookup
-    
+
     @property
     def fly(self) -> BaseFly:
         """Get the single fly in the world.

src/flygym/simulation.py

@@ -11,6 +11,7 @@
 from flygym.compose.world import BaseWorld
 from flygym.rendering import Renderer
 from flygym.utils.profiling import print_perf_report
+from flygym.utils.typing import n_jointdofs, n_actuators, n_tendon_actuators
 
 
 class Simulation:
@@ -162,7 +163,7 @@ def render_as_needed_with_profile(self) -> bool:
             self._frames_rendered += 1
         return render_done
 
-    def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]:  # noqa: F821
+    def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]:
         """Get current joint angles ordered by the fly's skeleton.
 
         Args:
@@ -175,7 +176,7 @@ def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]:
         internal_ids = self._intern_qposadrs_by_fly[fly_name]
         return self.mj_data.qpos[internal_ids]
 
-    def get_joint_velocities(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]:  # noqa: F821
+    def get_joint_velocities(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]:
         """Get current joint angular velocities ordered by the fly's skeleton.
 
         Args:
@@ -216,7 +217,7 @@ def get_body_rotations(self, fly_name: str) -> Float[np.ndarray, "n_bodies 4"]:
 
     def get_actuator_forces(
         self, fly_name: str, actuator_type: ActuatorType
-    ) -> Float[np.ndarray, "n_actuators"]:  # noqa: F821
+    ) -> Float[np.ndarray, "n_actuators"]:
         """Get actuator forces for the given actuator type.
 
         Args:
@@ -354,7 +355,7 @@ def set_actuator_inputs(
         self,
         fly_name: str,
         actuator_type: ActuatorType,
-        inputs: Float[np.ndarray, "n_actuators"],  # noqa: F821
+        inputs: Float[np.ndarray, "n_actuators"],
     ) -> None:
         """Set control inputs for the given actuator type.
 
@@ -393,7 +394,7 @@ def set_leg_adhesion_states(
     def set_tendon_actuator_inputs(
         self,
         fly_name: str,
-        inputs: Float[np.ndarray, "n_tendon_actuators"],  # noqa: F821
+        inputs: Float[np.ndarray, "n_tendon_actuators"],
     ) -> None:
         """Set control inputs for tendon actuators.
 
@@ -770,13 +771,13 @@ def close(self):
         self.eye_renderer = None
         # Don't destruct self.retina and self.eye_renderer_scene_option: they can be
         # reused and retina init requires some IO ops.
-    
+
     @property
     def fly(self) -> BaseFly:
         """Return the single fly in the world, or raise an error if there are multiple."""
         return self.world.fly
-    
+
     @property
     def fly_lookup(self) -> dict[str, BaseFly]:
         """Return the fly lookup dictionary from the world."""
-        return self.world.fly_lookup
+        return self.world.fly_lookup

src/flygym/utils/typing.py

@@ -0,0 +1,19 @@
+"""Canonical axis-name bindings for jaxtyping shape hints.
+
+Referencing these (instead of bare string identifiers) lets pyflakes/ruff
+resolve the forward references inside shape strings like
+``Float[np.ndarray, "n_bodies"]`` instead of flagging them as undefined
+names (F821) — no lint ignores needed.
+"""
+
+from typing import TypeVar
+
+n_worlds = TypeVar("n_worlds")
+n_jointdofs = TypeVar("n_jointdofs")
+n_actuators = TypeVar("n_actuators")
+n_tendon_actuators = TypeVar("n_tendon_actuators")
+n_bodies = TypeVar("n_bodies")
+n_sites = TypeVar("n_sites")
+n_bodysegments = TypeVar("n_bodysegments")
+n_cameras = TypeVar("n_cameras")
+n_ommatidia = TypeVar("n_ommatidia")

src/flygym/warp/utils.py

@@ -7,9 +7,7 @@
 
 @wp.kernel
 def wp_gather_indexed_rows_3d(
-    src: wp.array3d(dtype=wp.float32),  # type: ignore
-    dst: wp.array3d(dtype=wp.float32),  # type: ignore
-    rows: wp.array(dtype=wp.int32),  # type: ignore
+    src: wp.array3d[float], dst: wp.array3d[float], rows: wp.array[int]
 ):
     """Gather specific rows (dim 1) from a 3D Warp array into a narrower destination.
 
@@ -30,9 +28,7 @@ def wp_gather_indexed_rows_3d(
 
 @wp.kernel
 def wp_gather_indexed_rows_vec3f(
-    src: wp.array2d(dtype=wp.vec3f),  # type: ignore
-    dst: wp.array3d(dtype=wp.float32),  # type: ignore
-    rows: wp.array(dtype=wp.int32),  # type: ignore
+    src: wp.array2d[wp.vec3], dst: wp.array3d[float], rows: wp.array[int]
 ):
     """Gather specific rows from a 2D ``vec3f`` array into a ``(n_worlds, n_rows_narrow, 3)``
     ``float32`` destination.
@@ -57,9 +53,7 @@ def wp_gather_indexed_rows_vec3f(
 
 @wp.kernel
 def wp_gather_indexed_rows_quatf(
-    src: wp.array2d(dtype=wp.quatf),  # type: ignore
-    dst: wp.array3d(dtype=wp.float32),  # type: ignore
-    rows: wp.array(dtype=wp.int32),  # type: ignore
+    src: wp.array2d[wp.quat], dst: wp.array3d[float], rows: wp.array[int]
 ):
     """Gather specific rows from a 2D ``quatf`` array into a ``(n_worlds, n_rows_narrow, 4)``
     ``float32`` destination.
@@ -85,9 +79,7 @@ def wp_gather_indexed_rows_quatf(
 
 @wp.kernel
 def wp_scatter_indexed_cols_2d(
-    src: wp.array2d(dtype=wp.float32),  # type: ignore
-    dst: wp.array2d(dtype=wp.float32),  # type: ignore
-    cols: wp.array(dtype=wp.int32),  # type: ignore
+    src: wp.array2d[float], dst: wp.array2d[float], cols: wp.array[int]
 ):
     """Scatter a 2D Warp array into specific columns of a wider destination array.
 
@@ -108,9 +100,7 @@ def wp_scatter_indexed_cols_2d(
 
 @wp.kernel
 def wp_gather_indexed_cols_2d(
-    src: wp.array2d(dtype=wp.float32),  # type: ignore
-    dst: wp.array2d(dtype=wp.float32),  # type: ignore
-    cols: wp.array(dtype=wp.int32),  # type: ignore
+    src: wp.array2d[float], dst: wp.array2d[float], cols: wp.array[int]
 ):
     """Gather specific columns from a 2D Warp array into a narrower destination array.
 
@@ -132,12 +122,12 @@ def wp_gather_indexed_cols_2d(
 @wp.kernel
 def unpack_rgb_kernel_selected_worlds_and_cameras(
     # In:
-    packed: wp.array2d(dtype=wp.uint32),  # type: ignore
-    rgb_adr: wp.array(dtype=int),  # type: ignore
-    worldids_to_render: wp.array(dtype=int),  # type: ignore
-    camids_to_render: wp.array(dtype=int),  # type: ignore
+    packed: wp.array2d[wp.uint32],
+    rgb_adr: wp.array[int],
+    worldids_to_render: wp.array[int],
+    camids_to_render: wp.array[int],
     # Out:
-    rgb_out: wp.array4d(dtype=wp.vec3),  # type: ignore
+    rgb_out: wp.array4d[wp.vec3],
 ):
     """Unpack ABGR uint32 packed pixel data into separate R, G, and B channels."""
     idx_within_worldids, idx_within_camids, pixelid = wp.tid()
@@ -156,9 +146,9 @@ def unpack_rgb_kernel_selected_worlds_and_cameras(
 
 def get_rgb_selected_worlds_and_cameras(
     rc: mjw.RenderContext,
-    worldids: wp.array(dtype=int),  # type: ignore
-    camids: wp.array(dtype=int),  # type: ignore
-    rgb_out: wp.array4d(dtype=wp.vec3),  # type: ignore
+    worldids: wp.array[int],
+    camids: wp.array[int],
+    rgb_out: wp.array4d[wp.vec3],
 ):
     """Get the RGB data output from the render context buffers for the selected worlds
     and cameras.

src/flygym_demo/benchmark/time_gpu_simulation.py

@@ -88,9 +88,9 @@ def make_target_angles_all_worlds(self, n_worlds: int, sim_steps: int):
 
 @wp.kernel
 def update_target_angles_kernel(
-    dof_angles_all_worlds_gpu: wp.array3d(dtype=wp.float32),  # type: ignore
-    step_counter_gpu: wp.array(dtype=wp.int32),  # type: ignore
-    curr_target_angles_gpu: wp.array2d(dtype=wp.float32),  # type: ignore
+    dof_angles_all_worlds_gpu: wp.array3d[float],
+    step_counter_gpu: wp.array[int],
+    curr_target_angles_gpu: wp.array2d[float],
 ):
     world_id, actuator_id = wp.tid()
     step = step_counter_gpu[0]
@@ -99,9 +99,7 @@ def update_target_angles_kernel(
 
 
 @wp.kernel
-def increment_counter_kernel(
-    step_counter_gpu: wp.array(dtype=wp.int32),  # type: ignore
-):
+def increment_counter_kernel(step_counter_gpu: wp.array[int]):
     step_counter_gpu[0] = step_counter_gpu[0] + 1