Add MyPy checks and fix annotation

.pre-commit-config.yaml

@@ -32,3 +32,12 @@ repos:
     rev: 23.12.0
     hooks:
     -   id: black
+
+-   repo: local
+    hooks:
+      - id: mypy
+        name: "Static type checker"
+        entry: python -m mypy src/thants/
+        language: system
+        types: [ python ]
+        pass_filenames: false

pyproject.toml

@@ -34,6 +34,7 @@ jumanji = "^1.1.1"
 pytest = "^8.4.1"
 pre-commit = "^4.3.0"
 taskipy = "^1.14.1"
+mypy = "^1.18.2"
 
 [build-system]
 requires = ["poetry-core"]

src/thants/envs/mono.py

@@ -2,6 +2,7 @@
 from typing import Optional, Sequence
 
 import chex
+import jax
 from jumanji import Environment, specs
 from jumanji.types import TimeStep
 from jumanji.viewer import Viewer
@@ -77,10 +78,10 @@ def __init__(
             Amount of signal deposited by the deposit signal action
         """
         colony_generator = colony_generator or BasicColonyGenerator(25, 2, (5, 5))
-        colony_generator = SingleColonyWrapper(colony_generator)
+        wrapped_colony_generator = SingleColonyWrapper(colony_generator)
         self.env = Thants(
             dims=dims,
-            colonies_generator=colony_generator,
+            colonies_generator=wrapped_colony_generator,
             food_generator=food_generator,
             terrain_generator=terrain_generator,
             signal_dynamics=signal_dynamics,
@@ -111,8 +112,8 @@ def reset(self, key: chex.PRNGKey) -> tuple[State, TimeStep[Observations]]:
         state, timestep = self.env.reset(key)
         return state, timestep[0]
 
-    def step(
-        self, state: State, actions: chex.Array
+    def step(  # type: ignore[override]
+        self, state: State, actions: jax.Array
     ) -> tuple[State, TimeStep[Observations]]:
         """
         Update the state of the environment

src/thants/envs/multi.py

@@ -96,7 +96,7 @@ def __init__(
         self.signal_deposit_amount = signal_deposit_amount
         self.max_steps = max_steps
         self._colonies_generator = colonies_generator or DualBasicColoniesGenerator(
-            [25, 25], 2, (5, 5)
+            (25, 25), 2, (5, 5)
         )
         self._food_generator = food_generator or BasicFoodGenerator((5, 5), 100, 1.0)
         self._terrain_generator = terrain_generator or OpenTerrainGenerator()
@@ -107,7 +107,7 @@ def __init__(
         self._viewer = viewer or ThantsMultiColonyViewer()
         super().__init__()
 
-    def reset(
+    def reset(  # type: ignore[override]
         self, key: chex.PRNGKey
     ) -> tuple[State, Sequence[TimeStep[Observations]]]:
         """
@@ -126,13 +126,13 @@ def reset(
         """
         key, colony_key, food_key, terrain_key = jax.random.split(key, num=4)
         colonies = self._colonies_generator(self.dims, colony_key)
-        colonies = merge_colonies(colonies)
+        merged_colonies = merge_colonies(colonies)
         food = self._food_generator.init(self.dims, food_key)
         terrain = self._terrain_generator(self.dims, terrain_key)
         state = State(
             step=0,
             key=key,
-            colonies=colonies,
+            colonies=merged_colonies,
             food=food,
             terrain=terrain,
         )
@@ -143,8 +143,8 @@ def reset(
         ]
         return state, time_steps
 
-    def step(
-        self, state: State, actions: Sequence[chex.Array]
+    def step(  # type: ignore[override]
+        self, state: State, actions: Sequence[jax.Array]
     ) -> tuple[State, Sequence[TimeStep[Observations]]]:
         """
         Update the state of the environment
@@ -171,10 +171,10 @@ def step(
             Tuple containing new state and list of TimeSteps for each colony
         """
         key, food_key, signals_key = jax.random.split(state.key, num=3)
-        actions = jnp.concatenate(actions, axis=0)
+        all_actions = jnp.concatenate(actions, axis=0)
         # Unwrap actions
-        actions = derive_actions(
-            actions,
+        unwrapped_actions = derive_actions(
+            all_actions,
             take_food_amount=self.take_food_amount,
             deposit_food_amount=self.deposit_food_amount,
             signal_deposit_amount=self.signal_deposit_amount,
@@ -185,14 +185,14 @@ def step(
             self.dims,
             state.colonies.ants.pos,
             state.terrain,
-            actions.movements,
+            unwrapped_actions.movements,
         )
         # Pick up and drop-off food for each colony
         new_food, new_carrying = update_food(
             state.food,
             new_pos,
-            actions.take_food,
-            actions.deposit_food,
+            unwrapped_actions.take_food,
+            unwrapped_actions.deposit_food,
             state.colonies.ants.carrying,
             self.carry_capacity,
         )
@@ -202,7 +202,10 @@ def step(
         new_signals = self._signal_dynamics(signals_key, state.colonies.signals)
         # Deposit signals
         new_signals = deposit_signals(
-            new_signals, new_pos, state.colonies.colony_idx, actions.deposit_signals
+            new_signals,
+            new_pos,
+            state.colonies.colony_idx,
+            unwrapped_actions.deposit_signals,
         )
         # Clear food dropped on nests
         new_food = clear_nest(state.colonies.nests, new_food)
@@ -251,7 +254,9 @@ def num_actions(self) -> int:
         return 7 + self._colonies_generator.n_signals
 
     @cached_property
-    def observation_spec(self) -> Sequence[specs.Spec[Observations]]:
+    def observation_spec(  # type: ignore[override]
+        self,
+    ) -> Sequence[specs.Spec[Observations]]:
         """
         List of observation specifications for each colony
 
@@ -296,7 +301,7 @@ def action_spec(self) -> Sequence[specs.BoundedArray]:
         ]
 
     @cached_property
-    def reward_spec(self) -> Sequence[specs.Array]:
+    def reward_spec(self) -> Sequence[specs.Array]:  # type: ignore[override]
         """
         List of reward specifications for each colony
 

src/thants/generators/colonies/utils.py

@@ -2,13 +2,13 @@
 from dataclasses import dataclass
 from typing import Sequence
 
-import chex
+import jax
 import jax.numpy as jnp
 
 from thants.types import Ants, Colony
 
 
-def get_rectangular_indices(rec_dims: tuple[int, int]) -> chex.Array:
+def get_rectangular_indices(rec_dims: tuple[int, int]) -> jax.Array:
     """
     Get cell indices for a rectangle
 
@@ -65,12 +65,12 @@ def init_colony(
     """
     x0 = jnp.array(bounds.x0)
     x1 = jnp.array(bounds.x1)
-    dims = jnp.array(dims)
+    dims_arr = jnp.array(dims)
     centre = (x0 + ((x1 - x0) // 2))[jnp.newaxis]
     d = math.ceil(math.sqrt(n_agents))
     ant_pos = get_rectangular_indices((d, d))[:n_agents]
     ant_pos = ant_pos + centre - (jnp.array([[d, d]]) // 2)
-    ant_pos = ant_pos % dims
+    ant_pos = ant_pos % dims_arr
 
     ant_health = jnp.ones((n_agents,))
     ant_carrying = jnp.zeros((n_agents,))
@@ -79,19 +79,19 @@ def init_colony(
 
     nest_idxs = get_rectangular_indices(nest_dims)
     nest_idxs = nest_idxs + centre - jnp.array(nest_dims) // 2
-    nest_idxs = nest_idxs % dims
-    nest = jnp.zeros(dims, dtype=bool)
+    nest_idxs = nest_idxs % dims_arr
+    nest = jnp.zeros(dims_arr, dtype=bool)
     nest = nest.at[nest_idxs[:, 0], nest_idxs[:, 1]].set(True)
 
-    signals = jnp.zeros((n_signals, *dims))
+    signals = jnp.zeros((n_signals, *dims_arr))
 
     return Colony(ants=ants, signals=signals, nest=nest)
 
 
 def init_colonies(
     env_dims: tuple[int, int],
     nest_dims: tuple[int, int],
-    n_agents: Sequence[BBox],
+    n_agents: Sequence[int],
     n_signals: int,
     bounds: Sequence[BBox],
 ) -> Sequence[Colony]:

src/thants/generators/food.py

@@ -13,7 +13,7 @@ class FoodGenerator(abc.ABC):
     """
 
     @abc.abstractmethod
-    def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array:
+    def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> jax.Array:
         """
         Initialise environment food state
 
@@ -31,7 +31,7 @@ def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array:
         """
 
     @abc.abstractmethod
-    def update(self, key: chex.PRNGKey, step: int, food: chex.Array) -> chex.Array:
+    def update(self, key: chex.PRNGKey, step: int, food: jax.Array) -> jax.Array:
         """
         Update food state during simulation, e.g. drop more food
 
@@ -80,7 +80,7 @@ def __init__(
         self.drop_interval = drop_interval
         self.drop_amount = drop_amount
 
-    def _drop_food(self, key: chex.PRNGKey, food: chex.Array) -> chex.Array:
+    def _drop_food(self, key: chex.PRNGKey, food: jax.Array) -> jax.Array:
         """
         Place a new fixed size block of at a random location
 
@@ -104,7 +104,7 @@ def _drop_food(self, key: chex.PRNGKey, food: chex.Array) -> chex.Array:
         food = food.at[food_idxs[:, 0], food_idxs[:, 1]].add(self.drop_amount)
         return food
 
-    def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array:
+    def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> jax.Array:
         """
         Initialise environment food state
 
@@ -126,7 +126,7 @@ def init(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array:
         food = self._drop_food(key, food)
         return food
 
-    def update(self, key: chex.PRNGKey, step: int, food: chex.Array) -> chex.Array:
+    def update(self, key: chex.PRNGKey, step: int, food: jax.Array) -> jax.Array:
         """
         Drop rectangular blocks of food at random locations at fixed intervals