Improved visuals

.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

README.md

@@ -85,6 +85,13 @@ env.animate(state_history, 100, "multi_colony.gif")
 
 ## Environment
 
+<div align="center">
+  <img src="https://github.com/zombie-einstein/thants/raw/main/.github/images/thants_env.gif" />
+  <br>
+  <em>A Thants environment with two competing colonies.</em>
+</div>
+<br>
+
 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

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)

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:])

src/thants/mono/viewer.py

@@ -1,4 +1,3 @@
-from functools import partial
 from typing import Optional, Sequence, Tuple
 
 import chex
@@ -10,34 +9,38 @@
 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]):
     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,

src/thants/multi/viewer.py

@@ -1,58 +1,51 @@
-from functools import partial
 from typing import Optional, Sequence, Tuple
 
 import chex
 import jax
 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]):
     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,