add food and nestfinding behaviour on step function

agent.py

@@ -7,14 +7,15 @@ License: AGPL 3 (see end of file)
 (C) Alexander Bocken, Viviane Fahrni, Grace Kragho
 """
 import numpy as np
+import numpy.typing as npt
 from mesa.agent import Agent
 from mesa.space import Coordinate
-from typing import overload
-
 
 class RandomWalkerAnt(Agent):
     def __init__(self, unique_id, model, look_for_chemical=None,
-                 energy_0=1, chemical_drop_rate_0=1, sensitvity_0=1, alpha=0.5, drop_chemical=None) -> None:
+                 energy_0=1, chemical_drop_rate_0=1, sensitvity_0=0.1,
+                 alpha=0.5, drop_chemical=None,
+                 ) -> None:
         super().__init__(unique_id=unique_id, model=model)
 
         self._next_pos : None | Coordinate = None
@@ -28,19 +29,62 @@ class RandomWalkerAnt(Agent):
         self.alpha = alpha
 
 
-    def sensitvity_to_concentration(self, prop : float) -> float:
+    def sens_adj(self, props) -> npt.NDArray[np.float_] | float:
-        # TODO
+        # if props iterable create array, otherwise return single value
-        return prop
+        try:
+            iter(props)
+        except TypeError:
+            if props > self.sensitvity:
+                # TODO: nonlinear response
+                return props
+            else:
+                return 0
+
+        arr : list[float] = []
+        for prop in props:
+            arr.append(self.sens_adj(prop))
+        return np.array(arr)
+
 
     def step(self):
-        # follow positive gradient
+        # TODO: sensitvity decay
+
         if self.prev_pos is None:
             i = np.random.choice(range(6))
             self._next_pos = self.neighbors()[i]
             return
+
+        # Ants dropping A look for food
+        if self.drop_chemical == "A":
+            for neighbor in self.front_neighbors:
+                if self.model.grid.is_food(neighbor):
+                    self.drop_chemical = "B"
+                    self.prev_pos = neighbor
+                    self._next_pos = self.pos
+
+        # Ants dropping B look for nest
+        elif self.drop_chemical == "B":
+            for neighbor in self.front_neighbors:
+                if self.model.grid.is_nest(neighbor):
+                    self.look_for_chemical = "A" # Is this a correct interpretation?
+                    self.drop_chemical = "A"
+                    #TODO: Do we flip the ant here or reset prev pos?
+                    # For now, flip ant just like at food
+                    self.prev_pos = neighbor
+                    self._next_pos = self.pos
+
+                    for agent_id in self.model.get_unique_ids(self.model.num_new_recruits):
+                        agent = RandomWalkerAnt(unique_id=agent_id, model=self.model, look_for_chemical="B", drop_chemical="A")
+                        agent._next_pos = self.pos
+                        self.model.schedule.add(agent)
+                        self.model.grid.place_agent(agent, pos=neighbor)
+
+        # follow positive gradient
         if self.look_for_chemical is not None:
             front_concentration = [self.model.grid.fields[self.look_for_chemical][cell] for cell in self.front_neighbors ]
-            gradient = front_concentration - np.repeat(self.model.grid.fields[self.look_for_chemical][self.pos], 3)
+            front_concentration = self.sens_adj(front_concentration)
+            current_pos_concentration = self.sens_adj(self.model.grid.fields[self.look_for_chemical][self.pos])
+            gradient = front_concentration - np.repeat(current_pos_concentration, 3)
             index = np.argmax(gradient)
             if gradient[index] > 0:
                 self._next_pos = self.front_neighbors[index]

model.py

@@ -30,11 +30,12 @@ class ActiveWalkerModel(Model):
         self._unique_id_counter = -1
 
         self.max_steps = max_steps
-        self.nest_position : Coordinate = nest_position
+        self.grid.add_nest(nest_position)
         self.num_max_agents = num_max_agents
+        self.num_new_recruits = 5
 
-        self.decay_rates : dict[str, float] = {"A" :0.1,
+        self.decay_rates : dict[str, float] = {"A" :0.01,
-                                               "B": 0.1,
+                                               "B": 0.01,
                                                }
 
         for agent_id in self.get_unique_ids(num_initial_roamers):
@@ -42,6 +43,9 @@ class ActiveWalkerModel(Model):
             self.schedule.add(agent)
             self.grid.place_agent(agent, pos=nest_position)
 
+        for _ in range(5):
+            self.grid.add_food(5)
+
         self.datacollector = DataCollector(
                 model_reporters={},
                 agent_reporters={}

multihex.py

@@ -9,6 +9,7 @@ License: AGPL 3 (see end of file)
 (C) Alexander Bocken, Viviane Fahrni, Grace Kragho
 """
 
+from sys import dont_write_bytecode
 from mesa.space import HexGrid
 from mesa.agent import Agent
 import numpy as np
@@ -104,6 +105,54 @@ class MultiHexGridScalarFields(MultiHexGrid):
     def reset_field(self, key : str) -> None:
         self.fields[key] = np.zeros((self.width, self.height))
 
+    def is_food(self, pos):
+        assert('food' in self.fields.keys())
+        return bool(self.fields['food'][pos])
+
+    def add_food(self, size : int , pos=None):
+        """
+        Adds food source to grid.
+        Args:
+        pos (optional): if None, selects random place on grid which
+                        is not yet occupied by either a nest or another food source
+        size:           how much food should be added to field
+        """
+        assert('food' in self.fields.keys())
+        if pos is None:
+            def select_random_place():
+                i = np.random.randint(0, self.width)
+                j = np.random.randint(0, self.height)
+                return i,j
+            pos = select_random_place()
+            while(self.is_nest(pos) or self.is_food(pos)):
+                pos = select_random_place()
+
+        self.fields['food'][pos] = size
+
+    def is_nest(self, pos : Coordinate) -> bool:
+        assert('nests' in self.fields.keys())
+        return bool(self.fields['nests'][pos])
+
+    def add_nest(self, pos:None|Coordinate=None):
+        """
+        Adds nest to grid.
+        Args:
+        pos:    if None, selects random place on grid which
+                is not yet occupied by either a nest or another food source
+        """
+        assert('nests' in self.fields.keys())
+        if pos is None:
+            def select_random_place():
+                i = np.random.randint(0, self.width)
+                j = np.random.randint(0, self.height)
+                return i,j
+            pos = select_random_place()
+            while(self.is_nest(pos) or self.is_food(pos)):
+                pos = select_random_place()
+
+        self.fields['nests'][pos] = True
+
+
 """
 This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, version 3.
 

server.py

@@ -12,14 +12,15 @@ License: AGPL 3 (see end of file)
 """
 
 import numpy as np
-from mesa.visualization.modules import CanvasHexGrid, ChartModule, CanvasGrid
+from mesa.visualization.modules import CanvasHexGrid, ChartModule, CanvasGrid, TextElement
 from mesa.visualization.ModularVisualization import ModularServer
 from mesa.visualization.UserParam import UserSettableParameter
 from model import ActiveWalkerModel
 from collections import defaultdict
 
-def setup():
+def setup(params=None):
     # Set the model parameters
+    if params is None:
         params = {
                   "width": 50, "height": 50,
                   "num_max_agents" : 100,
@@ -59,11 +60,19 @@ def setup():
         level: level to calculate color between white and black (linearly)
         normalization: value for which we want full black color
         """
-        rgb = max(int(255 - level * 255 / normalization), 0)
+        return max(int(255 - level * 255 / normalization), 0)
-        mono = f"{rgb:0{2}x}" # hex value of rgb value with fixed length 2
+
-        return f"#{3*mono}"
 
     def portray_ant_density(model, pos):
+        if model.grid.is_nest(pos):
+            col = "red"
+        elif model.grid.is_food(pos):
+            col = "green"
+        else:
+            col = get_color(level=len(model.grid[pos]), normalization=5)
+            col = f"rgb({col}, {col}, {col})"
+
+
         return {
             "Shape": "hex",
             "r": 1,
@@ -71,10 +80,12 @@ def setup():
             "Layer": 0,
             "x": pos[0],
             "y": pos[1],
-            "Color": get_color(level=len(model.grid[pos]), normalization=5)
+            "Color":  col,
         }
 
     def portray_pheromone_density(model, pos):
+        col_a = get_color(level=model.grid.fields["A"][pos], normalization=3)
+        col_b = get_color(level=model.grid.fields["B"][pos], normalization=3)
         return {
             "Shape": "hex",
             "r": 1,
@@ -82,7 +93,7 @@ def setup():
             "Layer": 0,
             "x": pos[0],
             "y": pos[1],
-            "Color": get_color(level=model.grid.fields["A"][pos], normalization=3)
+            "Color": f"rgb({col_a}, {col_b}, 255)"
         }
 
 
@@ -92,7 +103,8 @@ def setup():
     pixel_ratio = 10
     grid_ants = CanvasHexGridMultiAgents(portray_ant_density, width, height, width*pixel_ratio, height*pixel_ratio)
     grid_pheromones = CanvasHexGridMultiAgents(portray_pheromone_density, width, height, width*pixel_ratio, height*pixel_ratio)
-    return ModularServer(ActiveWalkerModel, [grid_ants, grid_pheromones],
+    test_text = TextElement()
+    return ModularServer(ActiveWalkerModel, [lambda m: "<h3>Ant density</h3><h5>Nest: Red, Food: Green</h5>", grid_ants, lambda m: "<h3>Pheromone Density</h3><h5>Pheromone A: Cyan, Pheromone B: Pink</h5>", grid_pheromones],
                            "Active Random Walker Ants", params)
 
 if __name__ == "__main__":