diff --git a/.github/images/thants_env.gif b/.github/images/thants_env.gif
new file mode 100644
index 0000000..f371b6a
Binary files /dev/null and b/.github/images/thants_env.gif differ
diff --git a/.github/workflows/pre-merge.yaml b/.github/workflows/pre-merge.yaml
index bfcf48b..27ccace 100644
--- a/.github/workflows/pre-merge.yaml
+++ b/.github/workflows/pre-merge.yaml
@@ -33,18 +33,3 @@ jobs:
installer-parallel: true
- run: poetry install
- run: pytest -vv tests
- build-docs:
- if: github.event.pull_request.draft == false
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v5
- - uses: actions/setup-python@v5
- with:
- python-version: '3.10'
- - uses: snok/install-poetry@v1
- with:
- version: 1.5.1
- virtualenvs-create: false
- installer-parallel: true
- - run: poetry install
- - run: task docs
diff --git a/README.md b/README.md
index 2539eda..0933af2 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,13 @@ env.animate(state_history, 100, "multi_colony.gif")
## Environment
+
+
+
+
A Thants environment with two competing colonies.
+
+
The environment is modelled as a grid, wrapped at the boundaries. Ants (the agents)
occupy individual cells on the grid (and cannot overlap). Ants can pick up, carry,
and deposit food, or deposit persistent signals that can be observed by other ants
diff --git a/src/thants/common/types.py b/src/thants/common/types.py
index 397ced4..5da1151 100644
--- a/src/thants/common/types.py
+++ b/src/thants/common/types.py
@@ -79,3 +79,18 @@ class Observations:
nest: chex.Array # (n_ants, 9)
terrain: chex.Array # (n_ants, 9)
carrying: chex.Array # (n_ants,)
+
+
+@dataclass
+class ColorScheme:
+ """
+ Visualisation color-scheme, given as an array of rgb values
+
+ ants: Array of colors for each colony
+ food: Color to represent food
+ terrain: Array of blocked/passable terrain cells
+ """
+
+ ants: chex.Array # (n-colonies, 4)
+ food: chex.Array # (4,)
+ terrain: chex.Array # (2, 4)
diff --git a/src/thants/common/utils.py b/src/thants/common/utils.py
index 231b88a..719a070 100644
--- a/src/thants/common/utils.py
+++ b/src/thants/common/utils.py
@@ -2,10 +2,11 @@
import chex
import jax.numpy as jnp
+from matplotlib import color_sequences
from matplotlib.axes import Axes
from matplotlib.figure import Figure
-from thants.common.types import Ants, Colony
+from thants.common.types import Ants, Colony, ColorScheme
def get_rectangular_indices(rec_dims: tuple[int, int]) -> chex.Array:
@@ -109,3 +110,24 @@ def format_plot(
ax.set_ylim(-0.5, env_dims[0] - 0.5)
return fig, ax
+
+
+def get_color_scheme(color_sequence: str, n_colonies: int) -> ColorScheme:
+ """
+ Get a environment visualisation colour scheme from a matplotlib sequence
+
+ Parameters
+ ----------
+ color_sequence
+ Matplotlib color-sequence name
+ n_colonies
+ Number of colonies to visualise
+
+ Returns
+ -------
+ 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:])
diff --git a/src/thants/mono/viewer.py b/src/thants/mono/viewer.py
index 614bb8a..b66c263 100644
--- a/src/thants/mono/viewer.py
+++ b/src/thants/mono/viewer.py
@@ -1,4 +1,3 @@
-from functools import partial
from typing import Optional, Sequence, Tuple
import chex
@@ -10,27 +9,30 @@
from matplotlib.image import AxesImage
from numpy.typing import NDArray
-from thants.common.utils import format_plot
+from thants.common.types import ColorScheme
+from thants.common.utils import format_plot, get_color_scheme
from thants.mono.types import State
-def _draw_env(state: State) -> tuple[chex.Array, chex.Array]:
- terrain = state.terrain.astype(float)
- terrain = jnp.stack([terrain, terrain, terrain, jnp.ones_like(terrain)], axis=2)
- trans_colors = jnp.array([1.0, 0.0, 0.0, 0.5])
- nest = state.colony.nest[:, :, jnp.newaxis] * trans_colors[jnp.newaxis]
- return terrain, nest
+def _draw_env(
+ state: State, colors: ColorScheme
+) -> tuple[chex.Array, chex.Array, chex.Array]:
+ terrain = state.terrain.astype(int)
+ terrain = colors.terrain.at[terrain].get()
+ nest_colors = jnp.clip(colors.ants + 0.1, 0.0, 1.0)
+ nest = state.colony.nest[:, :, jnp.newaxis] * nest_colors[jnp.newaxis]
+ food = jnp.full((*state.food.shape, 4), colors.food)
+ return terrain, nest, food
-@partial(jax.jit, static_argnames="dims")
-def _draw_ants(dims: tuple[int, int], state: State) -> tuple[chex.Array, chex.Array]:
+@jax.jit
+def _draw_ants(state: State, colors: ColorScheme) -> chex.Array:
+ dims = state.food.shape
ants = jnp.zeros((*dims, 4))
- color = jnp.array([1.0, 0.0, 0.0, 1.0])
- ants = ants.at[state.colony.ants.pos[:, 0], state.colony.ants.pos[:, 1]].set(color)
- food = jnp.stack(
- [jnp.zeros(dims), jnp.ones(dims), jnp.zeros(dims), state.food], axis=2
+ ants = ants.at[state.colony.ants.pos[:, 0], state.colony.ants.pos[:, 1]].set(
+ colors.ants[0]
)
- return ants, food
+ return ants
class ThantsViewer(MatplotlibViewer[State]):
@@ -38,6 +40,7 @@ def __init__(
self,
name: str = "thants",
render_mode: str = "human",
+ color_sequence: str = "tab20",
) -> None:
"""
Thants environment visualiser using Matplotlib
@@ -48,7 +51,10 @@ def __init__(
Plot name, default ``thants``
render_mode
Default ``human``
+ color_sequence
+ Matplotlib colour sequence to sample from
"""
+ self.color_sequence = color_sequence
super().__init__(name, render_mode)
def _set_figure_size(self, dims: tuple[int, int]) -> None:
@@ -77,6 +83,7 @@ def render(
dims = state.food.shape
self._set_figure_size(dims)
fig, ax = self._get_fig_ax(padding=0.01)
+ ax.clear()
fig, ax = format_plot(fig, ax, dims)
self._draw(ax, state)
@@ -114,18 +121,12 @@ def animate(
fig, ax = format_plot(fig, ax, env_dims)
plt.close(fig=fig)
- terrain, nest = _draw_env(states[0])
- ax.imshow(terrain)
- ax.imshow(nest)
-
- ants, food = _draw_ants(env_dims, states[0])
- food_img = ax.imshow(food)
- ants_img = ax.imshow(ants)
+ colors, ants_img, food_img = self._draw(ax, states[0])
def make_frame(state: State) -> tuple[AxesImage, AxesImage]:
- step_ants, step_food = _draw_ants(env_dims, state)
- food_img.set_data(step_food)
+ step_ants = _draw_ants(state, colors)
ants_img.set_data(step_ants)
+ food_img.set_alpha(jnp.clip(state.food, 0.0, 1.0))
return ants_img, food_img
self._animation = matplotlib.animation.FuncAnimation(
@@ -141,17 +142,21 @@ def make_frame(state: State) -> tuple[AxesImage, AxesImage]:
return self._animation
- def _draw(self, ax: plt.Axes, state: State) -> None:
- ax.clear()
- env_dims = state.terrain.shape
+ def _draw(
+ self, ax: plt.Axes, state: State
+ ) -> tuple[ColorScheme, AxesImage, AxesImage]:
+ colors = get_color_scheme(self.color_sequence, 1)
+
+ terrain, nest, food = _draw_env(state, colors)
+ ants = _draw_ants(state, colors)
- terrain, nest = _draw_env(state)
ax.imshow(terrain)
ax.imshow(nest)
- ants, food = _draw_ants(env_dims, state)
- ax.imshow(food)
- ax.imshow(ants)
+ food_img = ax.imshow(food, alpha=state.food)
+ ants_img = ax.imshow(ants)
+
+ return colors, ants_img, food_img
def _get_fig_ax(
self,
diff --git a/src/thants/multi/viewer.py b/src/thants/multi/viewer.py
index da2ee77..41b6e00 100644
--- a/src/thants/multi/viewer.py
+++ b/src/thants/multi/viewer.py
@@ -1,4 +1,3 @@
-from functools import partial
from typing import Optional, Sequence, Tuple
import chex
@@ -6,45 +5,39 @@
import jax.numpy as jnp
import matplotlib.animation
import matplotlib.pyplot as plt
-import numpy as np
from jumanji.viewer import MatplotlibViewer
-from matplotlib import colormaps
from matplotlib.image import AxesImage
from numpy.typing import NDArray
-from thants.common.utils import format_plot
+from thants.common.types import ColorScheme
+from thants.common.utils import format_plot, get_color_scheme
from thants.multi.types import State
-def draw_env(state: State, colors: chex.Array) -> tuple[chex.Array, chex.Array]:
- terrain = state.terrain.astype(float)
- terrain = jnp.stack([terrain, terrain, terrain, jnp.ones_like(terrain)], axis=2)
-
- trans_colors = colors * jnp.array([[1.0, 1.0, 1.0, 0.5]])
-
+def _draw_env(
+ state: State, colors: ColorScheme
+) -> tuple[chex.Array, chex.Array, chex.Array]:
+ terrain = state.terrain.astype(int)
+ terrain = colors.terrain.at[terrain].get()
+ nest_colors = jnp.clip(colors.ants + 0.1, 0.0, 1.0)
nests = [
- colony.nest[:, :, jnp.newaxis] * trans_colors[i, jnp.newaxis]
+ colony.nest[:, :, jnp.newaxis] * nest_colors[i, jnp.newaxis]
for i, colony in enumerate(state.colonies)
]
nests = jnp.sum(jnp.stack(nests, axis=0), axis=0)
+ food = jnp.full((*state.food.shape, 4), colors.food)
+ return terrain, nests, food
- return terrain, nests
-
-@partial(jax.jit, static_argnames="dims")
-def draw_ants(
- dims: tuple[int, int], state: State, colors: chex.Array
-) -> tuple[chex.Array, chex.Array]:
+@jax.jit
+def _draw_ants(state: State, colors: ColorScheme) -> chex.Array:
+ dims = state.food.shape
ants = jnp.zeros((*dims, 4))
for i, colony in enumerate(state.colonies):
- ants = ants.at[colony.ants.pos[:, 0], colony.ants.pos[:, 1]].set(colors[i])
-
- food = jnp.stack(
- [jnp.zeros(dims), jnp.ones(dims), jnp.zeros(dims), state.food], axis=2
- )
+ ants = ants.at[colony.ants.pos[:, 0], colony.ants.pos[:, 1]].set(colors.ants[i])
- return ants, food
+ return ants
class ThantsMultiColonyViewer(MatplotlibViewer[State]):
@@ -52,7 +45,7 @@ def __init__(
self,
name: str = "thants",
render_mode: str = "human",
- colony_colormap: str = "plasma",
+ color_sequence: str = "tab20",
) -> None:
"""
Thants multi-colony environment visualiser using Matplotlib
@@ -63,15 +56,12 @@ def __init__(
Plot name, default ``thants``
render_mode
Default ``human``
- colony_colormap
-
+ color_sequence
+ Matplotlib colour sequence to sample from
"""
- self.cmap = colormaps[colony_colormap]
+ self.color_sequence = color_sequence
super().__init__(name, render_mode)
- def _get_colony_colors(self, n: int) -> chex.Array:
- return jnp.array(self.cmap(np.linspace(0, 1, n)))
-
def _set_figure_size(self, dims: tuple[int, int]) -> None:
longest = max(dims[0], dims[1])
f_dims = (10.0 * dims[1] / longest, 10.0 * dims[0] / longest)
@@ -98,6 +88,7 @@ def render(
dims = state.food.shape
self._set_figure_size(dims)
fig, ax = self._get_fig_ax(padding=0.01)
+ ax.clear()
fig, ax = format_plot(fig, ax, dims)
self._draw(ax, state)
@@ -136,20 +127,12 @@ def animate(
fig, ax = format_plot(fig, ax, dims)
plt.close(fig=fig)
- colors = self._get_colony_colors(len(states[0].colonies))
-
- terrain, nests = draw_env(states[0], colors)
- ants, food = draw_ants(dims, states[0], colors)
-
- ax.imshow(terrain, cmap="grey")
- ax.imshow(nests)
- food_img = ax.imshow(food)
- ants_img = ax.imshow(ants)
+ colors, ants_img, food_img = self._draw(ax, states[0])
def make_frame(state: State) -> tuple[AxesImage, AxesImage]:
- step_ants, step_food = draw_ants(dims, state, colors)
+ step_ants = _draw_ants(state, colors)
ants_img.set_data(step_ants)
- food_img.set_data(step_food)
+ food_img.set_alpha(jnp.clip(state.food, 0.0, 1.0))
return ants_img, food_img
self._animation = matplotlib.animation.FuncAnimation(
@@ -165,18 +148,20 @@ def make_frame(state: State) -> tuple[AxesImage, AxesImage]:
return self._animation
- def _draw(self, ax: plt.Axes, state: State) -> None:
- ax.clear()
- dims = state.food.shape
- colors = self._get_colony_colors(len(state.colonies))
+ def _draw(
+ self, ax: plt.Axes, state: State
+ ) -> tuple[ColorScheme, AxesImage, AxesImage]:
+ colors = get_color_scheme(self.color_sequence, len(state.colonies))
- terrain, nests = draw_env(state, colors)
- ants, food = draw_ants(dims, state, colors)
+ terrain, nests, food = _draw_env(state, colors)
+ ants = _draw_ants(state, colors)
- ax.imshow(terrain, cmap="grey")
+ ax.imshow(terrain)
ax.imshow(nests)
- ax.imshow(food)
- ax.imshow(ants)
+ food_img = ax.imshow(food, alpha=state.food)
+ ants_img = ax.imshow(ants)
+
+ return colors, ants_img, food_img
def _get_fig_ax(
self,