Repository: projectmesa/mesa
Pull request #1926 (2023-12-23 19:52:12)

Reimplement schedulers to use AgentSet
This PR builds on #1916 and updates the schedulers to use the new AgentSet class and/or weakrefs. It also adds a new feature: schedulers take an optional agents keyword argument in `__init__` so you can spawn a scheduler with all agents. 

My aim with this PR is to be fully backward compatible but raise warnings when methods that have become redundant are being used. The current code passes all unit tests.

Still to do:

- [x] - update docs
- [x] - check type annotations
- [x] - double check to ensure backward compatible behavior
- [x] - code formatting
- [x] - add some additional tests



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) Report
Attention: `16 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`7bd42e7`)](https://app.codecov.io/gh/projectmesa/mesa/commit/7bd42e78516203083e576630af8a7f458bcd09c0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 79.14% compared to head [(`a06b099`)](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) 78.23%.

| [Files](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) | Patch % | Lines |
|---|---|---|
| [mesa/time.py](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa#diff-bWVzYS90aW1lLnB5) | 71.42% | [12 Missing and 4 partials :warning: ](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1926      +/-   ##
==========================================
- Coverage   79.14%   78.23%   -0.91%     
==========================================
  Files          15       15              
  Lines        1098     1112      +14     
  Branches      236      238       +2     
==========================================
+ Hits          869      870       +1     
- Misses        198      209      +11     
- Partials       31       33       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/projectmesa/mesa/pull/1926?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=projectmesa).

EwoutH: Looks like a nice start!

I think a big conceptual design decision we need to make is if we want to support multiple schedulers. Did we do so in the past?

If we only have one scheduler, we could simplify and assume a lot, for example that all agents are in the scheduler. We don't have to explicitly add or remove them, or even keep track of them in the scheduler. There would also be one single step/time clock. It would probably allow deprecating all schedulers, and moving a lot of the requesting of behavior to `AgentSet.do()`. It would be a flexible approach with lot possibilities, like described in #1912.

However, if we do want to support multiple schedulers, we need to leave a lot of stuff in the scheduler, at least the time of that scheduler and which agents are part of it. Personally, I can't try of that many cases you can't solve with the AgentSet approach, but there might be any.
quaquel: I have used multiple schedulers in the past (pre `RandomActivationByType`), and more generally, I like its flexibility. Also, this PR's only aim is to keep the current Schedulers up to date by using AgentSet. Discussing the future of schedulers is beyond the scope of this PR. 

As indicated in #1912, one possible direction could be to shift to a central EventList. This creates a central understanding of time. All existing functionality of the various schedulers could be rebuilt on such an event list. 
quaquel: I have reviewed the code several times, fixed the docs, and tested against various examples. This is ready for a proper check. @EwoutH, @Corvince, @jackiekazil, or @tpike3 do you have a chance to check this?
Corvince: Probably not in the coming days, but I really look forward to review this!
EwoutH: Will try to review it later today!

Patch:
diff --git a/.github/release.yml b/.github/release.yml
new file mode 100644
index 00000000..0277adf3
--- /dev/null
+++ b/.github/release.yml
@@ -0,0 +1,34 @@
+# Docs: https://docs.github.com/en/repositories/releasing-projects-on-github/automatically-generated-release-notes
+changelog:
+  exclude:
+    labels:
+      - ignore-for-release
+  categories:
+    - title: ⚠️ Breaking changes
+      labels:
+        - breaking
+    - title: 🧪 Experimental features
+      labels:
+        - experimental
+    - title: 🎉 New features added
+      labels:
+        - feature
+    - title: 🛠 Enhancements made
+      labels:
+        - enhancement
+    - title: 🐛 Bugs fixed
+      labels:
+        - bug
+    - title: 📜 Documentation improvements
+      labels:
+        - docs
+    - title: 🔧 Maintenance
+      labels:
+        - ci
+        - testing
+        - dependency
+        - maintenance
+        - packaging
+    - title: Other changes
+      labels:
+        - "*"
diff --git a/mesa/agent.py b/mesa/agent.py
index 0bb2f218..db452dbd 100644
--- a/mesa/agent.py
+++ b/mesa/agent.py
@@ -11,6 +11,7 @@ import contextlib
 import operator
 import warnings
 import weakref
+from collections import defaultdict
 from collections.abc import MutableSet, Sequence
 from random import Random
 
@@ -47,13 +48,26 @@ class Agent:
         self.pos: Position | None = None
 
         # register agent
-        self.model._agents[type(self)][self] = None
+        try:
+            self.model._agents[type(self)][self] = None
+        except AttributeError:
+            # model super has not been called
+            self.model._agents = defaultdict(dict)
+            self.model.agentset_experimental_warning_given = False
+
+            warnings.warn(
+                "The Mesa Model class wasn’t initialized. In the future, you need to explicitly initialize the Model by calling super().__init__() on initialization.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
 
     def remove(self) -> None:
         """Remove and delete the agent from the model."""
         with contextlib.suppress(KeyError):
             self.model._agents[type(self)].pop(self)
 
+
     def step(self) -> None:
         """A single step of the agent."""
 
@@ -91,6 +105,8 @@ class AgentSet(MutableSet, Sequence):
         which means that agents not referenced elsewhere in the program may be automatically removed from the AgentSet.
     """
 
+    agentset_experimental_warning_given = False
+
     def __init__(self, agents: Iterable[Agent], model: Model):
         """
         Initializes the AgentSet with a collection of agents and a reference to the model.
@@ -101,8 +117,8 @@ class AgentSet(MutableSet, Sequence):
         """
         self.model = model
 
-        if not self.model.agentset_experimental_warning_given:
-            self.model.agentset_experimental_warning_given = True
+        if not self.__class__.agentset_experimental_warning_given:
+            self.__class__.agentset_experimental_warning_given = True
             warnings.warn(
                 "The AgentSet is experimental. It may be changed or removed in any and all future releases, including patch releases.\n"
                 "We would love to hear what you think about this new feature. If you have any thoughts, share them with us here: https://github.com/projectmesa/mesa/discussions/1919",
diff --git a/mesa/model.py b/mesa/model.py
index 22e65448..ca0c653c 100644
--- a/mesa/model.py
+++ b/mesa/model.py
@@ -14,7 +14,7 @@ from collections import defaultdict
 # mypy
 from typing import Any
 
-from mesa.agent import AgentSet
+from mesa.agent import Agent, AgentSet
 from mesa.datacollection import DataCollector
 
 
@@ -83,9 +83,9 @@ class Model:
         """Return a list of different agent types."""
         return list(self._agents.keys())
 
-    def get_agents_of_type(self, agenttype: type) -> AgentSet:
+    def get_agents_of_type(self, agenttype: type[Agent]) -> AgentSet:
         """Retrieves an AgentSet containing all agents of the specified type."""
-        return AgentSet(self._agents[agenttype].values(), self)
+        return AgentSet(self._agents[agenttype].keys(), self)
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached. Overload as
diff --git a/mesa/space.py b/mesa/space.py
index 18cf7314..987afaf9 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -431,6 +431,69 @@ class _Grid:
         self.remove_agent(agent)
         self.place_agent(agent, pos)
 
+    def move_agent_to_one_of(
+        self,
+        agent: Agent,
+        pos: list[Coordinate],
+        selection: str = "random",
+        handle_empty: str | None = None,
+    ) -> None:
+        """
+        Move an agent to one of the given positions.
+
+        Args:
+            agent: Agent object to move. Assumed to have its current location stored in a 'pos' tuple.
+            pos: List of possible positions.
+            selection: String, either "random" (default) or "closest". If "closest" is selected and multiple
+                       cells are the same distance, one is chosen randomly.
+            handle_empty: String, either "warning", "error" or None (default). If "warning" or "error" is selected
+                          and no positions are given (an empty list), a warning or error is raised respectively.
+        """
+        # Only move agent if there are positions given (non-empty list)
+        if pos:
+            if selection == "random":
+                chosen_pos = agent.random.choice(pos)
+            elif selection == "closest":
+                current_pos = agent.pos
+                # Find the closest position without sorting all positions
+                closest_pos = None
+                min_distance = float("inf")
+                for p in pos:
+                    distance = self._distance_squared(p, current_pos)
+                    if distance < min_distance:
+                        min_distance = distance
+                        closest_pos = p
+                chosen_pos = closest_pos
+            else:
+                raise ValueError(
+                    f"Invalid selection method {selection}. Choose 'random' or 'closest'."
+                )
+            #  Move agent to chosen position
+            self.move_agent(agent, chosen_pos)
+
+        # If no positions are given, throw warning/error if selected
+        elif handle_empty == "warning":
+            warn(
+                f"No positions given, could not move agent {agent.unique_id}.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+        elif handle_empty == "error":
+            raise ValueError(
+                f"No positions given, could not move agent {agent.unique_id}."
+            )
+
+    def _distance_squared(self, pos1: Coordinate, pos2: Coordinate) -> float:
+        """
+        Calculate the squared Euclidean distance between two points for performance.
+        """
+        # Use squared Euclidean distance to avoid sqrt operation
+        dx, dy = abs(pos1[0] - pos2[0]), abs(pos1[1] - pos2[1])
+        if self.torus:
+            dx = min(dx, self.width - dx)
+            dy = min(dy, self.height - dy)
+        return dx**2 + dy**2
+
     def swap_pos(self, agent_a: Agent, agent_b: Agent) -> None:
         """Swap agents positions"""
         agents_no_pos = []
diff --git a/mesa/time.py b/mesa/time.py
index c4f055eb..7bd4ad1b 100644
--- a/mesa/time.py
+++ b/mesa/time.py
@@ -27,12 +27,12 @@ from __future__ import annotations
 
 import heapq
 import warnings
-from collections import defaultdict
+import weakref
 
 # mypy
-from typing import Union
+from typing import Iterable, Union
 
-from mesa.agent import Agent
+from mesa.agent import Agent, AgentSet
 from mesa.model import Model
 
 # BaseScheduler has a self.time of int, while
@@ -60,12 +60,25 @@ class BaseScheduler:
         - agents (property): Returns a list of all agent instances.
     """
 
-    def __init__(self, model: Model) -> None:
-        """Create a new, empty BaseScheduler."""
+    def __init__(self, model: Model, agents: Iterable[Agent] | None = None) -> None:
+        """Create a new BaseScheduler.
+
+        Args:
+            model (Model): The model to which the schedule belongs
+            agents (Iterable[Agent], None, optional): An iterable of agents who are controlled by the schedule
+
+        """
         self.model = model
         self.steps = 0
         self.time: TimeT = 0
-        self._agents: dict[int, Agent] = {}
+
+        if agents is None:
+            agents = []
+
+        self._agents: AgentSet = AgentSet(agents, model)
+
+        self._remove_warning_given = False
+        self._agents_key_warning_given = False
 
     def add(self, agent: Agent) -> None:
         """Add an Agent object to the schedule.
@@ -74,20 +87,23 @@ class BaseScheduler:
             agent: An Agent to be added to the schedule. NOTE: The agent must
             have a step() method.
         """
-        if agent.unique_id in self._agents:
-            raise Exception(
-                f"Agent with unique id {agent.unique_id!r} already added to scheduler"
-            )
 
-        self._agents[agent.unique_id] = agent
+        if agent not in self._agents:
+            self._agents.add(agent)
+        else:
+            raise ValueError("agent already added to scheduler")
 
     def remove(self, agent: Agent) -> None:
         """Remove all instances of a given agent from the schedule.
 
+        Note:
+            It is only necessary to explicitly remove agents from the schedule if
+            the agent is not removed from the model.
+
         Args:
             agent: An agent object.
         """
-        del self._agents[agent.unique_id]
+        self._agents.remove(agent)
 
     def step(self) -> None:
         """Execute the step of all the agents, one at a time."""
@@ -102,25 +118,31 @@ class BaseScheduler:
         return len(self._agents)
 
     @property
-    def agents(self) -> list[Agent]:
-        return list(self._agents.values())
+    def agents(self) -> AgentSet:
+        # a bit dirty, but returns a copy of the internal agent set
+        return self._agents.select()
 
     def get_agent_keys(self, shuffle: bool = False) -> list[int]:
         # To be able to remove and/or add agents during stepping
         # it's necessary to cast the keys view to a list.
-        agent_keys = list(self._agents.keys())
+
+        if not self._agents_key_warning_given:
+            self._agents_key_warning_given = True
+            warnings.warn(
+                "Because of the shift to using weakrefs, this method will be removed in a future version",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        agent_keys = [agent.unique_id for agent in self._agents]
         if shuffle:
             self.model.random.shuffle(agent_keys)
         return agent_keys
 
-    def do_each(self, method, agent_keys=None, shuffle=False):
-        if agent_keys is None:
-            agent_keys = self.get_agent_keys()
+    def do_each(self, method, shuffle=False):
         if shuffle:
-            self.model.random.shuffle(agent_keys)
-        for agent_key in agent_keys:
-            if agent_key in self._agents:
-                getattr(self._agents[agent_key], method)()
+            self.agents.shuffle(inplace=True)
+        self.agents.do(method)
 
 
 class RandomActivation(BaseScheduler):
@@ -202,6 +224,7 @@ class StagedActivation(BaseScheduler):
     def __init__(
         self,
         model: Model,
+        agents: Iterable[Agent] | None = None,
         stage_list: list[str] | None = None,
         shuffle: bool = False,
         shuffle_between_stages: bool = False,
@@ -209,15 +232,16 @@ class StagedActivation(BaseScheduler):
         """Create an empty Staged Activation schedule.
 
         Args:
-            - model: Model object associated with the schedule.
-            - stage_list: List of strings of names of stages to run, in the
+            model (Model): The model to which the schedule belongs
+            agents (Iterable[Agent], None, optional): An iterable of agents who are controlled by the schedule
+            stage_list (:obj:`list` of :obj:`str`): List of strings of names of stages to run, in the
                          order to run them in.
-            - shuffle: If True, shuffle the order of agents each step.
-            - shuffle_between_stages: If True, shuffle the agents after each
+            shuffle (bool, optional): If True, shuffle the order of agents each step.
+            shuffle_between_stages (bool, optional): If True, shuffle the agents after each
                                     stage; otherwise, only shuffle at the start
                                     of each step.
         """
-        super().__init__(model)
+        super().__init__(model, agents)
         self.stage_list = stage_list if stage_list else ["step"]
         self.shuffle = shuffle
         self.shuffle_between_stages = shuffle_between_stages
@@ -225,17 +249,14 @@ class StagedActivation(BaseScheduler):
 
     def step(self) -> None:
         """Executes all the stages for all agents."""
-        # To be able to remove and/or add agents during stepping
-        # it's necessary for the keys view to be a list.
-        agent_keys = self.get_agent_keys(self.shuffle)
+        shuffle = self.shuffle
         for stage in self.stage_list:
             if stage.startswith("model."):
                 getattr(self.model, stage[6:])()
             else:
-                self.do_each(stage, agent_keys=agent_keys)
-            # We recompute the keys because some agents might have been removed
-            # in the previous loop.
-            agent_keys = self.get_agent_keys(self.shuffle_between_stages)
+                self.do_each(stage, shuffle=shuffle)
+
+            shuffle = self.shuffle_between_stages
             self.time += self.stage_time
 
         self.steps += 1
@@ -265,9 +286,24 @@ class RandomActivationByType(BaseScheduler):
         - get_type_count: Returns the count of agents of a specific type.
     """
 
-    def __init__(self, model: Model) -> None:
-        super().__init__(model)
-        self.agents_by_type = defaultdict(dict)
+    def __init__(self, model: Model, agents: Iterable[Agent] | None = None) -> None:
+        super().__init__(model, agents)
+        """
+
+        Args:
+            model (Model): The model to which the schedule belongs
+            agents (Iterable[Agent], None, optional): An iterable of agents who are controlled by the schedule
+        """
+
+        # can't be a defaultdict because we need to pass model to AgentSet
+        self.agents_by_type: [type, AgentSet] = {}
+
+        if agents is not None:
+            for agent in agents:
+                try:
+                    self.agents_by_type[type(agent)].add(agent)
+                except KeyError:
+                    self.agents_by_type[type(agent)] = AgentSet([agent], self.model)
 
     def add(self, agent: Agent) -> None:
         """
@@ -277,17 +313,18 @@ class RandomActivationByType(BaseScheduler):
             agent: An Agent to be added to the schedule.
         """
         super().add(agent)
-        agent_class: type[Agent] = type(agent)
-        self.agents_by_type[agent_class][agent.unique_id] = agent
+
+        try:
+            self.agents_by_type[type(agent)].add(agent)
+        except KeyError:
+            self.agents_by_type[type(agent)] = AgentSet([agent], self.model)
 
     def remove(self, agent: Agent) -> None:
         """
         Remove all instances of a given agent from the schedule.
         """
-        del self._agents[agent.unique_id]
-
-        agent_class: type[Agent] = type(agent)
-        del self.agents_by_type[agent_class][agent.unique_id]
+        super().remove(agent)
+        self.agents_by_type[type(agent)].remove(agent)
 
     def step(self, shuffle_types: bool = True, shuffle_agents: bool = True) -> None:
         """
@@ -309,26 +346,25 @@ class RandomActivationByType(BaseScheduler):
         self.steps += 1
         self.time += 1
 
-    def step_type(self, type_class: type[Agent], shuffle_agents: bool = True) -> None:
+    def step_type(self, agenttype: type[Agent], shuffle_agents: bool = True) -> None:
         """
         Shuffle order and run all agents of a given type.
         This method is equivalent to the NetLogo 'ask [breed]...'.
 
         Args:
-            type_class: Class object of the type to run.
+            agenttype: Class object of the type to run.
         """
-        agent_keys: list[int] = list(self.agents_by_type[type_class].keys())
+        agents = self.agents_by_type[agenttype]
+
         if shuffle_agents:
-            self.model.random.shuffle(agent_keys)
-        for agent_key in agent_keys:
-            if agent_key in self.agents_by_type[type_class]:
-                self.agents_by_type[type_class][agent_key].step()
+            agents.shuffle(inplace=True)
+        agents.do("step")
 
-    def get_type_count(self, type_class: type[Agent]) -> int:
+    def get_type_count(self, agenttype: type[Agent]) -> int:
         """
         Returns the current number of agents of certain type in the queue.
         """
-        return len(self.agents_by_type[type_class])
+        return len(self.agents_by_type[agenttype])
 
 
 class DiscreteEventScheduler(BaseScheduler):
@@ -376,8 +412,15 @@ class DiscreteEventScheduler(BaseScheduler):
     """
 
     def __init__(self, model: Model, time_step: TimeT = 1) -> None:
+        """
+
+        Args:
+            model (Model): The model to which the schedule belongs
+            time_step (TimeT): The fixed time step between steps
+
+        """
         super().__init__(model)
-        self.event_queue: list[tuple[TimeT, float, Agent]] = []
+        self.event_queue: list[tuple[TimeT, float, weakref.ref]] = []
         self.time_step: TimeT = time_step  # Fixed time period for each step
 
         warnings.warn(
@@ -393,8 +436,13 @@ class DiscreteEventScheduler(BaseScheduler):
             raise ValueError(
                 f"Scheduled time ({time}) must be >= the current time ({self.time})"
             )
+        if agent not in self._agents:
+            raise ValueError(
+                "trying to schedule an event for agent which is not known to the scheduler"
+            )
+
         # Create an event, sorted first on time, secondary on a random value
-        event = (time, self.model.random.random(), agent)
+        event = (time, self.model.random.random(), weakref.ref(agent))
         heapq.heappush(self.event_queue, event)
 
     def schedule_in(self, delay: TimeT, agent: Agent) -> None:
@@ -411,12 +459,12 @@ class DiscreteEventScheduler(BaseScheduler):
         while self.event_queue and self.event_queue[0][0] <= end_time:
             # Get the next event (ignore the random value during unpacking)
             time, _, agent = heapq.heappop(self.event_queue)
+            agent = agent()  # unpack weakref
 
-            # Advance time to the event's time
-            self.time = time
-
-            # Execute the event
-            if agent.unique_id in self._agents:
+            if agent:
+                # Advance time to the event's time
+                self.time = time
+                # Execute the event
                 agent.step()
 
         # After processing events, advance time by the time_step


Test patch:
diff --git a/tests/test_space.py b/tests/test_space.py
index f8f2cc94..ab0e0da3 100644
--- a/tests/test_space.py
+++ b/tests/test_space.py
@@ -327,6 +327,33 @@ class TestSingleGrid(unittest.TestCase):
         assert self.space[initial_pos[0]][initial_pos[1]] is None
         assert self.space[final_pos[0]][final_pos[1]] == _agent
 
+    def test_move_agent_random_selection(self):
+        agent = self.agents[0]
+        possible_positions = [(10, 10), (20, 20), (30, 30)]
+        self.space.move_agent_to_one_of(agent, possible_positions, selection="random")
+        assert agent.pos in possible_positions
+
+    def test_move_agent_closest_selection(self):
+        agent = self.agents[0]
+        agent.pos = (5, 5)
+        possible_positions = [(6, 6), (10, 10), (20, 20)]
+        self.space.move_agent_to_one_of(agent, possible_positions, selection="closest")
+        assert agent.pos == (6, 6)
+
+    def test_move_agent_invalid_selection(self):
+        agent = self.agents[0]
+        possible_positions = [(10, 10), (20, 20), (30, 30)]
+        with self.assertRaises(ValueError):
+            self.space.move_agent_to_one_of(
+                agent, possible_positions, selection="invalid_option"
+            )
+
+    def test_distance_squared(self):
+        pos1 = (3, 4)
+        pos2 = (0, 0)
+        expected_distance_squared = 3**2 + 4**2
+        assert self.space._distance_squared(pos1, pos2) == expected_distance_squared
+
     def test_iter_cell_list_contents(self):
         """
         Test neighborhood retrieval
@@ -350,6 +377,68 @@ class TestSingleGrid(unittest.TestCase):
         assert len(cell_list_4) == 1
 
 
+class TestSingleGridTorus(unittest.TestCase):
+    def setUp(self):
+        self.space = SingleGrid(50, 50, True)  # Torus is True here
+        self.agents = []
+        for i, pos in enumerate(TEST_AGENTS_GRID):
+            a = MockAgent(i, None)
+            self.agents.append(a)
+            self.space.place_agent(a, pos)
+
+    def test_move_agent_random_selection(self):
+        agent = self.agents[0]
+        possible_positions = [(49, 49), (1, 1), (25, 25)]
+        self.space.move_agent_to_one_of(agent, possible_positions, selection="random")
+        assert agent.pos in possible_positions
+
+    def test_move_agent_closest_selection(self):
+        agent = self.agents[0]
+        agent.pos = (0, 0)
+        possible_positions = [(3, 3), (49, 49), (25, 25)]
+        self.space.move_agent_to_one_of(agent, possible_positions, selection="closest")
+        # Expecting (49, 49) to be the closest in a torus grid
+        assert agent.pos == (49, 49)
+
+    def test_move_agent_invalid_selection(self):
+        agent = self.agents[0]
+        possible_positions = [(10, 10), (20, 20), (30, 30)]
+        with self.assertRaises(ValueError):
+            self.space.move_agent_to_one_of(
+                agent, possible_positions, selection="invalid_option"
+            )
+
+    def test_move_agent_empty_list(self):
+        agent = self.agents[0]
+        possible_positions = []
+        agent.pos = (3, 3)
+        self.space.move_agent_to_one_of(agent, possible_positions, selection="random")
+        assert agent.pos == (3, 3)
+
+    def test_move_agent_empty_list_warning(self):
+        agent = self.agents[0]
+        possible_positions = []
+        # Should assert RuntimeWarning
+        with self.assertWarns(RuntimeWarning):
+            self.space.move_agent_to_one_of(
+                agent, possible_positions, selection="random", handle_empty="warning"
+            )
+
+    def test_move_agent_empty_list_error(self):
+        agent = self.agents[0]
+        possible_positions = []
+        with self.assertRaises(ValueError):
+            self.space.move_agent_to_one_of(
+                agent, possible_positions, selection="random", handle_empty="error"
+            )
+
+    def test_distance_squared_torus(self):
+        pos1 = (0, 0)
+        pos2 = (49, 49)
+        expected_distance_squared = 1**2 + 1**2  # In torus, these points are close
+        assert self.space._distance_squared(pos1, pos2) == expected_distance_squared
+
+
 class TestSingleNetworkGrid(unittest.TestCase):
     GRAPH_SIZE = 10
 
diff --git a/tests/test_time.py b/tests/test_time.py
index 5d6d66f9..a440d502 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -34,8 +34,7 @@ class MockAgent(Agent):
     def kill_other_agent(self):
         for agent in self.model.schedule.agents:
             if agent is not self:
-                self.model.schedule.remove(agent)
-                break
+                agent.remove()
 
     def stage_one(self):
         if self.model.enable_kill_other_agent:
@@ -78,7 +77,9 @@ class MockModel(Model):
         # Make scheduler
         if activation == STAGED:
             model_stages = ["stage_one", "model.model_stage", "stage_two"]
-            self.schedule = StagedActivation(self, model_stages, shuffle=shuffle)
+            self.schedule = StagedActivation(
+                self, stage_list=model_stages, shuffle=shuffle
+            )
         elif activation == RANDOM:
             self.schedule = RandomActivation(self)
         elif activation == SIMULTANEOUS:
@@ -140,9 +141,9 @@ class TestStagedActivation(TestCase):
         Test the staged activation can remove an agent
         """
         model = MockModel(shuffle=True)
-        agent_keys = list(model.schedule._agents.keys())
-        agent = model.schedule._agents[agent_keys[0]]
-        model.schedule.remove(agent)
+        agents = list(model.schedule._agents)
+        agent = agents[0]
+        model.schedule.remove(agents[0])
         assert agent not in model.schedule.agents
 
     def test_intrastep_remove(self):
@@ -166,6 +167,14 @@ class TestRandomActivation(TestCase):
     Test the random activation.
     """
 
+    def test_init(self):
+        model = Model()
+        agents = [MockAgent(model.next_id(), model) for _ in range(10)]
+
+        scheduler = RandomActivation(model, agents)
+        assert all(agent in scheduler.agents for agent in agents)
+
+
     def test_random_activation_step_shuffles(self):
         """
         Test the random activation step
@@ -205,6 +214,18 @@ class TestRandomActivation(TestCase):
         model.step()
         assert len(model.log) == 1
 
+    def test_get_agent_keys(self):
+        model = MockModel(activation=RANDOM)
+
+        keys = model.schedule.get_agent_keys()
+        agent_ids = [agent.unique_id for agent in model.agents]
+        assert all(entry_i == entry_j for entry_i, entry_j in zip(keys, agent_ids))
+
+        keys = model.schedule.get_agent_keys(shuffle=True)
+        agent_ids = {agent.unique_id for agent in model.agents}
+        assert all(entry in agent_ids for entry in keys)
+
+
 
 class TestSimultaneousActivation(TestCase):
     """
@@ -231,6 +252,15 @@ class TestRandomActivationByType(TestCase):
     does step for one type of agents, not the entire agents.
     """
 
+    def test_init(self):
+        model = Model()
+        agents = [MockAgent(model.next_id(), model) for _ in range(10)]
+        agents += [Agent(model.next_id(), model) for _ in range(10)]
+
+        scheduler = RandomActivationByType(model, agents)
+        assert all(agent in scheduler.agents for agent in agents)
+
+
     def test_random_activation_step_shuffles(self):
         """
         Test the random activation by type step
@@ -262,18 +292,30 @@ class TestRandomActivationByType(TestCase):
         # one step for each of 2 agents
         assert all(x == 1 for x in agent_steps)
 
-    def test_add_non_unique_ids(self):
+    def test_random_activation_counts(self):
         """
-        Test that adding agent with duplicate ids result in an error.
-        TODO: we need to run this test on all schedulers, not just
-        RandomActivationByType.
+        Test the random activation by type step causes each agent to step
         """
+
         model = MockModel(activation=RANDOM_BY_TYPE)
-        a = MockAgent(0, model)
-        b = MockAgent(0, model)
-        model.schedule.add(a)
-        with self.assertRaises(Exception):
-            model.schedule.add(b)
+
+        agent_types = model.agent_types
+        for agent_type in agent_types:
+            assert model.schedule.get_type_count(agent_type) == len(model.get_agents_of_type(agent_type))
+
+    # def test_add_non_unique_ids(self):
+    #     """
+    #     Test that adding agent with duplicate ids result in an error.
+    #     TODO: we need to run this test on all schedulers, not just
+    #     TODO:: identical IDs is something for the agent, not the scheduler and should be tested there
+    #     RandomActivationByType.
+    #     """
+    #     model = MockModel(activation=RANDOM_BY_TYPE)
+    #     a = MockAgent(0, model)
+    #     b = MockAgent(0, model)
+    #     model.schedule.add(a)
+    #     with self.assertRaises(Exception):
+    #         model.schedule.add(b)
 
 
 class TestDiscreteEventScheduler(TestCase):
@@ -297,19 +339,19 @@ class TestDiscreteEventScheduler(TestCase):
         self.assertEqual(len(self.scheduler.event_queue), 1)
         event_time, _, event_agent = self.scheduler.event_queue[0]
         self.assertEqual(event_time, 5)
-        self.assertEqual(event_agent, self.agent1)
+        self.assertEqual(event_agent(), self.agent1)
 
     def test_schedule_event_with_float_time(self):
         self.scheduler.schedule_event(5.5, self.agent1)
         self.assertEqual(len(self.scheduler.event_queue), 1)
         event_time, _, event_agent = self.scheduler.event_queue[0]
         self.assertEqual(event_time, 5.5)
-        self.assertEqual(event_agent, self.agent1)
+        self.assertEqual(event_agent(), self.agent1)
 
     def test_schedule_in(self):
         self.scheduler.schedule_in(3, self.agent2)
         _, _, event_agent = self.scheduler.event_queue[0]
-        self.assertEqual(event_agent, self.agent2)
+        self.assertEqual(event_agent(), self.agent2)
         self.assertEqual(self.scheduler.get_next_event_time(), self.scheduler.time + 3)
 
     # Testing Event Execution and Time Advancement
@@ -337,6 +379,11 @@ class TestDiscreteEventScheduler(TestCase):
         with self.assertRaises(ValueError):
             self.scheduler.schedule_event(-1, self.agent1)
 
+    def test_invalid_aget_time(self):
+        with self.assertRaises(ValueError):
+            agent3 = MockAgent(3, self.model)
+            self.scheduler.schedule_event(2, agent3)
+
     def test_immediate_event_execution(self):
         # Current time of the scheduler
         current_time = self.scheduler.time
@@ -366,7 +413,7 @@ class TestDiscreteEventScheduler(TestCase):
         self.assertEqual(len(self.scheduler.event_queue), 1)
         event_time, _, event_agent = self.scheduler.event_queue[0]
         self.assertEqual(event_time, self.scheduler.time)
-        self.assertEqual(event_agent, new_agent)
+        self.assertEqual(event_agent(), new_agent)
 
         # Step the scheduler and check if the agent's step method is executed
         self.scheduler.step()