diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9fb7e78..667d63c 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/poetry.lock b/poetry.lock index 4e454b3..cd150fb 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1220,6 +1220,79 @@ files = [ {file = "mslex-1.3.0.tar.gz", hash = "sha256:641c887d1d3db610eee2af37a8e5abda3f70b3006cdfd2d0d29dc0d1ae28a85d"}, ] +[[package]] +name = "mypy" +version = "1.18.2" +description = "Optional static typing for Python" +optional = false +python-versions = ">=3.9" +groups = ["dev"] +files = [ + {file = "mypy-1.18.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c1eab0cf6294dafe397c261a75f96dc2c31bffe3b944faa24db5def4e2b0f77c"}, + {file = "mypy-1.18.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7a780ca61fc239e4865968ebc5240bb3bf610ef59ac398de9a7421b54e4a207e"}, + {file = "mypy-1.18.2-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:448acd386266989ef11662ce3c8011fd2a7b632e0ec7d61a98edd8e27472225b"}, + {file = "mypy-1.18.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f9e171c465ad3901dc652643ee4bffa8e9fef4d7d0eece23b428908c77a76a66"}, + {file = "mypy-1.18.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:592ec214750bc00741af1f80cbf96b5013d81486b7bb24cb052382c19e40b428"}, + {file = "mypy-1.18.2-cp310-cp310-win_amd64.whl", hash = "sha256:7fb95f97199ea11769ebe3638c29b550b5221e997c63b14ef93d2e971606ebed"}, + {file = "mypy-1.18.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:807d9315ab9d464125aa9fcf6d84fde6e1dc67da0b6f80e7405506b8ac72bc7f"}, + {file = "mypy-1.18.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:776bb00de1778caf4db739c6e83919c1d85a448f71979b6a0edd774ea8399341"}, + {file = "mypy-1.18.2-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1379451880512ffce14505493bd9fe469e0697543717298242574882cf8cdb8d"}, + {file = "mypy-1.18.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1331eb7fd110d60c24999893320967594ff84c38ac6d19e0a76c5fd809a84c86"}, + {file = "mypy-1.18.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:3ca30b50a51e7ba93b00422e486cbb124f1c56a535e20eff7b2d6ab72b3b2e37"}, + {file = "mypy-1.18.2-cp311-cp311-win_amd64.whl", hash = "sha256:664dc726e67fa54e14536f6e1224bcfce1d9e5ac02426d2326e2bb4e081d1ce8"}, + {file = "mypy-1.18.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:33eca32dd124b29400c31d7cf784e795b050ace0e1f91b8dc035672725617e34"}, + {file = "mypy-1.18.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a3c47adf30d65e89b2dcd2fa32f3aeb5e94ca970d2c15fcb25e297871c8e4764"}, + {file = "mypy-1.18.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5d6c838e831a062f5f29d11c9057c6009f60cb294fea33a98422688181fe2893"}, + {file = "mypy-1.18.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01199871b6110a2ce984bde85acd481232d17413868c9807e95c1b0739a58914"}, + {file = "mypy-1.18.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a2afc0fa0b0e91b4599ddfe0f91e2c26c2b5a5ab263737e998d6817874c5f7c8"}, + {file = "mypy-1.18.2-cp312-cp312-win_amd64.whl", hash = "sha256:d8068d0afe682c7c4897c0f7ce84ea77f6de953262b12d07038f4d296d547074"}, + {file = "mypy-1.18.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:07b8b0f580ca6d289e69209ec9d3911b4a26e5abfde32228a288eb79df129fcc"}, + {file = "mypy-1.18.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ed4482847168439651d3feee5833ccedbf6657e964572706a2adb1f7fa4dfe2e"}, + {file = "mypy-1.18.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3ad2afadd1e9fea5cf99a45a822346971ede8685cc581ed9cd4d42eaf940986"}, + {file = "mypy-1.18.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a431a6f1ef14cf8c144c6b14793a23ec4eae3db28277c358136e79d7d062f62d"}, + {file = "mypy-1.18.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7ab28cc197f1dd77a67e1c6f35cd1f8e8b73ed2217e4fc005f9e6a504e46e7ba"}, + {file = "mypy-1.18.2-cp313-cp313-win_amd64.whl", hash = "sha256:0e2785a84b34a72ba55fb5daf079a1003a34c05b22238da94fcae2bbe46f3544"}, + {file = "mypy-1.18.2-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:62f0e1e988ad41c2a110edde6c398383a889d95b36b3e60bcf155f5164c4fdce"}, + {file = "mypy-1.18.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8795a039bab805ff0c1dfdb8cd3344642c2b99b8e439d057aba30850b8d3423d"}, + {file = "mypy-1.18.2-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6ca1e64b24a700ab5ce10133f7ccd956a04715463d30498e64ea8715236f9c9c"}, + {file = "mypy-1.18.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d924eef3795cc89fecf6bedc6ed32b33ac13e8321344f6ddbf8ee89f706c05cb"}, + {file = "mypy-1.18.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20c02215a080e3a2be3aa50506c67242df1c151eaba0dcbc1e4e557922a26075"}, + {file = "mypy-1.18.2-cp314-cp314-win_amd64.whl", hash = "sha256:749b5f83198f1ca64345603118a6f01a4e99ad4bf9d103ddc5a3200cc4614adf"}, + {file = "mypy-1.18.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:25a9c8fb67b00599f839cf472713f54249a62efd53a54b565eb61956a7e3296b"}, + {file = "mypy-1.18.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c2b9c7e284ee20e7598d6f42e13ca40b4928e6957ed6813d1ab6348aa3f47133"}, + {file = "mypy-1.18.2-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d6985ed057513e344e43a26cc1cd815c7a94602fb6a3130a34798625bc2f07b6"}, + {file = "mypy-1.18.2-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22f27105f1525ec024b5c630c0b9f36d5c1cc4d447d61fe51ff4bd60633f47ac"}, + {file = "mypy-1.18.2-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:030c52d0ea8144e721e49b1f68391e39553d7451f0c3f8a7565b59e19fcb608b"}, + {file = "mypy-1.18.2-cp39-cp39-win_amd64.whl", hash = "sha256:aa5e07ac1a60a253445797e42b8b2963c9675563a94f11291ab40718b016a7a0"}, + {file = "mypy-1.18.2-py3-none-any.whl", hash = "sha256:22a1748707dd62b58d2ae53562ffc4d7f8bcc727e8ac7cbc69c053ddc874d47e"}, + {file = "mypy-1.18.2.tar.gz", hash = "sha256:06a398102a5f203d7477b2923dda3634c36727fa5c237d8f859ef90c42a9924b"}, +] + +[package.dependencies] +mypy_extensions = ">=1.0.0" +pathspec = ">=0.9.0" +tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} +typing_extensions = ">=4.6.0" + +[package.extras] +dmypy = ["psutil (>=4.0)"] +faster-cache = ["orjson"] +install-types = ["pip"] +mypyc = ["setuptools (>=50)"] +reports = ["lxml"] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +description = "Type system extensions for programs checked with the mypy type checker." +optional = false +python-versions = ">=3.8" +groups = ["dev"] +files = [ + {file = "mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505"}, + {file = "mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558"}, +] + [[package]] name = "nodeenv" version = "1.9.1" @@ -1407,6 +1480,18 @@ files = [ {file = "packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f"}, ] +[[package]] +name = "pathspec" +version = "0.12.1" +description = "Utility library for gitignore style pattern matching of file paths." +optional = false +python-versions = ">=3.8" +groups = ["dev"] +files = [ + {file = "pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08"}, + {file = "pathspec-0.12.1.tar.gz", hash = "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712"}, +] + [[package]] name = "pillow" version = "11.3.0" @@ -2041,7 +2126,6 @@ files = [ {file = "typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548"}, {file = "typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466"}, ] -markers = {dev = "python_version == \"3.10\""} [[package]] name = "urllib3" @@ -2177,4 +2261,4 @@ files = [ [metadata] lock-version = "2.1" python-versions = ">=3.10,<4.0" -content-hash = "5158b5e58ce2e24cb12be75c04037353e1e73e0320c7bdfd291235ae284633e3" +content-hash = "a0a15cf814745906fd8e2149793370860408c8ec2c88022ab9858c33ba795098" diff --git a/pyproject.toml b/pyproject.toml index 0fef565..c7c2f00 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/thants/envs/mono.py b/src/thants/envs/mono.py index a66b7e7..d781318 100644 --- a/src/thants/envs/mono.py +++ b/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 diff --git a/src/thants/envs/multi.py b/src/thants/envs/multi.py index a893d92..d71bc6f 100644 --- a/src/thants/envs/multi.py +++ b/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 diff --git a/src/thants/generators/colonies/utils.py b/src/thants/generators/colonies/utils.py index 7593858..c8ef9bf 100644 --- a/src/thants/generators/colonies/utils.py +++ b/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,11 +79,11 @@ 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) @@ -91,7 +91,7 @@ def init_colony( 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]: diff --git a/src/thants/generators/food.py b/src/thants/generators/food.py index 2ae276f..732c311 100644 --- a/src/thants/generators/food.py +++ b/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 diff --git a/src/thants/generators/terrain.py b/src/thants/generators/terrain.py index 0c05d70..27b27c4 100644 --- a/src/thants/generators/terrain.py +++ b/src/thants/generators/terrain.py @@ -1,6 +1,7 @@ import abc import chex +import jax import jax.numpy as jnp @@ -10,7 +11,7 @@ class TerrainGenerator(abc.ABC): """ @abc.abstractmethod - def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array: + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> jax.Array: """ Generate an array indicating passable/unpassable terrain @@ -33,7 +34,7 @@ class OpenTerrainGenerator(TerrainGenerator): Generates completely open terrain """ - def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array: + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> jax.Array: """ Generate an array indicating passable/unpassable terrain @@ -57,7 +58,7 @@ class BoundedTerrainGenerator(TerrainGenerator): Generate open terrain with an unpassable boundary """ - def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> chex.Array: + def __call__(self, dims: tuple[int, int], key: chex.PRNGKey) -> jax.Array: """ Generate an array indicating passable/unpassable terrain diff --git a/src/thants/observations.py b/src/thants/observations.py index 3141faf..5114455 100644 --- a/src/thants/observations.py +++ b/src/thants/observations.py @@ -40,16 +40,16 @@ def observations_from_state( idxs = jnp.indices((3, 3)) idxs = idxs.swapaxes(0, 2).reshape(9, 2) - 1 - def get_ant_view(arr: chex.Array, i: chex.Array, x: chex.Array) -> chex.Array: + def get_ant_view(arr: jax.Array, i: jax.Array, x: jax.Array) -> jax.Array: return arr.at[i[:, jnp.newaxis], x[:, 0], x[:, 1]].get() - def get_view(arr: chex.Array, x: chex.Array) -> chex.Array: + def get_view(arr: jax.Array, x: jax.Array) -> jax.Array: return arr.at[x[:, 0], x[:, 1]].get() - def get_signals(arr: chex.Array, i: int, x: chex.Array) -> chex.Array: + def get_signals(arr: jax.Array, i: jax.Array, x: jax.Array) -> jax.Array: return arr.at[i, jnp.arange(n_signals)[:, jnp.newaxis], x[:, 0], x[:, 1]].get() - def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: + def get_nest(i: chex.Numeric, arr: jax.Array, x: jax.Array) -> jax.Array: return arr.at[x[:, 0], x[:, 1]].get() == (i + 1) occupation = jnp.zeros((n_colonies, *dims), dtype=float) @@ -75,9 +75,8 @@ def get_nest(i: int, arr: chex.Array, x: chex.Array) -> chex.Array: float ) - boundaries = [0, *colony_sizes] - boundaries = np.array(boundaries) - boundaries = np.cumsum(boundaries) + colony_bins = np.array([0, *colony_sizes]) + boundaries = np.cumsum(colony_bins) observations = [ Observations( diff --git a/src/thants/py.typed b/src/thants/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/src/thants/rewards.py b/src/thants/rewards.py index 378ec4e..b0104be 100644 --- a/src/thants/rewards.py +++ b/src/thants/rewards.py @@ -1,7 +1,7 @@ import abc from typing import Optional, Sequence -import chex +import jax import jax.numpy as jnp import numpy as np @@ -9,12 +9,12 @@ def delivered_food( - nest: chex.Array, - pos: chex.Array, - carrying_before: chex.Array, - carrying_after: chex.Array, - colony_idxs: Optional[chex.Array] = None, -) -> chex.Array: + nest: jax.Array, + pos: jax.Array, + carrying_before: jax.Array, + carrying_after: jax.Array, + colony_idxs: Optional[jax.Array] = None, +) -> jax.Array: """ Calculate food deposited by individual ant on a nest @@ -49,9 +49,8 @@ def delivered_food( def _get_boundaries(colony_sizes: Sequence[int]) -> np.typing.NDArray: - boundaries = [0, *colony_sizes] - boundaries = np.array(boundaries) - boundaries = np.cumsum(boundaries) + colony_bounds = np.array([0, *colony_sizes]) + boundaries = np.cumsum(colony_bounds) return boundaries @@ -59,7 +58,7 @@ class RewardFn(abc.ABC): @abc.abstractmethod def __call__( self, colony_sizes: Sequence[int], old_state: State, new_state: State - ) -> Sequence[chex.Array]: + ) -> Sequence[jax.Array]: """ Generate individual ant rewards for each colony @@ -82,7 +81,7 @@ def __call__( class NullRewardFn(RewardFn): def __call__( self, colony_sizes: Sequence[int], old_state: State, new_state: State - ) -> Sequence[chex.Array]: + ) -> Sequence[jax.Array]: """ Assigns 0 reward to all agents @@ -106,7 +105,7 @@ def __call__( class DeliveredFoodRewards(RewardFn): def __call__( self, colony_sizes: Sequence[int], old_state: State, new_state: State - ) -> Sequence[chex.Array]: + ) -> Sequence[jax.Array]: """ Assigns rewards for ants depositing food on their own nest diff --git a/src/thants/signals.py b/src/thants/signals.py index a2165e5..cf3d2ad 100644 --- a/src/thants/signals.py +++ b/src/thants/signals.py @@ -1,6 +1,7 @@ import abc import chex +import jax import jax.numpy as jnp @@ -10,7 +11,7 @@ class SignalPropagator(abc.ABC): """ @abc.abstractmethod - def __call__(self, key: chex.PRNGKey, signals: chex.Array) -> chex.Array: + def __call__(self, key: chex.PRNGKey, signals: jax.Array) -> jax.Array: """ Generate updated signal state @@ -53,7 +54,7 @@ def __init__(self, decay_rate: float, dissipation_rate) -> None: self.dissipation_rate = dissipation_rate super().__init__() - def __call__(self, key: chex.PRNGKey, signals: chex.Array) -> chex.Array: + def __call__(self, key: chex.PRNGKey, signals: jax.Array) -> jax.Array: """ Generate updated signal state diff --git a/src/thants/steps.py b/src/thants/steps.py index a24c777..dee9a54 100644 --- a/src/thants/steps.py +++ b/src/thants/steps.py @@ -1,19 +1,19 @@ from typing import Sequence -import chex +import jax import jax.numpy as jnp from thants.types import Ants, Colonies, Colony, SignalActions def update_food( - food: chex.Array, - pos: chex.Array, - take: chex.Array, - deposit: chex.Array, - carrying: chex.Array, + food: jax.Array, + pos: jax.Array, + take: jax.Array, + deposit: jax.Array, + carrying: jax.Array, capacity: float, -) -> tuple[chex.Array, chex.Array]: +) -> tuple[jax.Array, jax.Array]: """ Update food piles due to ant actions @@ -51,8 +51,8 @@ def update_food( def update_positions( - dims: tuple[int, int], pos: chex.Array, terrain: chex.Array, updates: chex.Array -) -> chex.Array: + dims: tuple[int, int], pos: jax.Array, terrain: jax.Array, updates: jax.Array +) -> jax.Array: """ Update agent positions, checking for collisions @@ -85,7 +85,7 @@ def update_positions( return new_pos -def clear_nest(nests: chex.Array, food: chex.Array) -> chex.Array: +def clear_nest(nests: jax.Array, food: jax.Array) -> jax.Array: """ Clear food deposited on each colony nest @@ -146,11 +146,11 @@ def merge_colonies(colonies: Sequence[Colony]) -> Colonies: def deposit_signals( - signals: chex.Array, - pos: chex.Array, - colony_idx: chex.Array, + signals: jax.Array, + pos: jax.Array, + colony_idx: jax.Array, deposits: SignalActions, -) -> chex.Array: +) -> jax.Array: """ Deposit signals for the relevant colony and channel diff --git a/src/thants/types.py b/src/thants/types.py index 23fa2e8..bd13a8d 100644 --- a/src/thants/types.py +++ b/src/thants/types.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING import chex +import jax if TYPE_CHECKING: from dataclasses import dataclass @@ -18,16 +19,16 @@ class Ants: carrying: Amount of food being carried by ants """ - pos: chex.Array # (n_ants, 2) - health: chex.Array # (n_ants,) - carrying: chex.Array # (n_ants,) + pos: jax.Array # (n_ants, 2) + health: jax.Array # (n_ants,) + carrying: jax.Array # (n_ants,) @dataclass class Colony: ants: Ants - signals: chex.Array # [n-channels, *env-size] - nest: chex.Array # [*env-size] + signals: jax.Array # [n-channels, *env-size] + nest: jax.Array # [*env-size] @dataclass @@ -42,9 +43,9 @@ class Colonies: """ ants: Ants - colony_idx: chex.Array # [n-ants,] - signals: chex.Array # [n-colonies, n-channels, *env-size] - nests: chex.Array # [*env-size] + colony_idx: jax.Array # [n-ants,] + signals: jax.Array # [n-colonies, n-channels, *env-size] + nests: jax.Array # [*env-size] @dataclass @@ -62,8 +63,8 @@ class State: step: int key: chex.PRNGKey colonies: Colonies - food: chex.Array # [*env-size] - terrain: chex.Array # [*env-size] + food: jax.Array # [*env-size] + terrain: jax.Array # [*env-size] @dataclass @@ -75,8 +76,8 @@ class SignalActions: amount: Amount to deposit """ - channel: chex.Array # (n_ants,) - amount: chex.Array # (n_ants,) + channel: jax.Array # (n_ants,) + amount: jax.Array # (n_ants,) @dataclass @@ -90,9 +91,9 @@ class Actions: deposit_signals: Signal deposit channels and amounts """ - movements: chex.Array # (n_ants, 2) - take_food: chex.Array # (n_ants,) - deposit_food: chex.Array # (n_ants,) + movements: jax.Array # (n_ants, 2) + take_food: jax.Array # (n_ants,) + deposit_food: jax.Array # (n_ants,) deposit_signals: SignalActions @@ -109,12 +110,12 @@ class Observations: carrying: Food amount being held """ - ants: chex.Array # (n_ants, 9) - food: chex.Array # (n_ants, 9) - signals: chex.Array # (n_ants, n-channels, 9) - nest: chex.Array # (n_ants, 9) - terrain: chex.Array # (n_ants, 9) - carrying: chex.Array # (n_ants,) + ants: jax.Array # (n_ants, 9) + food: jax.Array # (n_ants, 9) + signals: jax.Array # (n_ants, n-channels, 9) + nest: jax.Array # (n_ants, 9) + terrain: jax.Array # (n_ants, 9) + carrying: jax.Array # (n_ants,) @dataclass @@ -127,6 +128,6 @@ class ColorScheme: terrain: Array of blocked/passable terrain cells """ - ants: chex.Array # (n-colonies, 4) - food: chex.Array # (4,) - terrain: chex.Array # (2, 4) + ants: jax.Array # (n-colonies, 4) + food: jax.Array # (4,) + terrain: jax.Array # (2, 4) diff --git a/src/thants/viewer.py b/src/thants/viewer.py index 83effa3..be59e4f 100644 --- a/src/thants/viewer.py +++ b/src/thants/viewer.py @@ -58,8 +58,8 @@ def get_color_scheme(color_sequence: str, n_colonies: int) -> ColorScheme: Environment visualisation color-scheme """ colors = color_sequences[color_sequence] - colors = jnp.array([(*i, 1.0) for i in colors[: 3 + n_colonies]]) - return ColorScheme(terrain=colors[:2], food=colors[2], ants=colors[3:]) + colors_arr = jnp.array([(*i, 1.0) for i in colors[: 3 + n_colonies]]) + return ColorScheme(terrain=colors_arr[:2], food=colors_arr[2], ants=colors_arr[3:]) def _draw_env( @@ -130,8 +130,9 @@ def render( RGB array if the render_mode is ``rgb_array`` """ self._clear_display() - dims = state.food.shape + dims = (state.food.shape[0], state.food.shape[1]) self._set_figure_size(dims) + fig, ax = self._get_fig_ax(padding=0.01) ax.clear() fig, ax = format_plot(fig, ax, dims) @@ -165,7 +166,7 @@ def animate( if not states: raise ValueError(f"The states argument has to be non-empty, got {states}.") - dims = states[0].food.shape + dims = (states[0].food.shape[0], states[0].food.shape[1]) self._set_figure_size(dims) fig, ax = self._get_fig_ax(name_suffix="_animation", show=False, padding=0.01)