Repository: projectmesa/mesa

Issue:
No public version of NetworkGrid._remove_agent, though code exists


Currently it is not possible to remove an Agent from a NetworkGrid using a **public** function. Currently, you can do the following:

```python
...
model.grid = NetworkGrid(G)
...
a = Agent(...)
model.grid._remove_agent(a, a.pos)
...
```
But this function is a private function, therefore it isn't shown in documentation or in your editor. I'd like to be able to remove agents like with the other space modules.


**Describe the solution you'd like**
I'd like a similar solution as used for the other space modules:

Solution in space.Grid:
```python
def remove_agent(self, agent: Agent) -> None:  # solution in space.Grid
        """ Remove the agent from the grid and set its pos variable to None. """
        pos = agent.pos
        self._remove_agent(pos, agent)
        agent.pos = None

```
Solution for NetworkGrid:
```python
def remove_agent(self, agent: Agent) -> None:  # New solution
        """ Remove the agent from the network and set its pos variable to None. """
        pos = agent.pos
        self._remove_agent(agent, pos)
        agent.pos = None

```

**Additional context**
The _remove_agent function in NetworkGrid
```python
def _remove_agent(self, agent: Agent, node_id: int) -> None:
        """ Remove an agent from a node. """

        self.G.nodes[node_id]["agent"].remove(agent)
```

Discussion:

Patch:
diff --git a/mesa/space.py b/mesa/space.py
index 37716c55..83998580 100644
--- a/mesa/space.py
+++ b/mesa/space.py
@@ -17,21 +17,8 @@ import itertools
 
 import numpy as np
 
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Set,
-    Sequence,
-    Tuple,
-    Union,
-    cast,
-    overload,
-)
-from .agent import Agent
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
+from mesa.agent import Agent
 
 Coordinate = Tuple[int, int]
 GridContent = Union[Optional[Agent], Set[Agent]]
@@ -113,77 +100,13 @@ class Grid:
         # Add all cells to the empties list.
         self.empties = set(itertools.product(*(range(self.width), range(self.height))))
 
-        # Neighborhood Cache
-        self._neighborhood_cache: Dict[Any, List[Coordinate]] = dict()
-
     @staticmethod
     def default_val() -> None:
         """ Default value for new cell elements. """
         return None
 
-    @overload
     def __getitem__(self, index: int) -> List[GridContent]:
-        ...
-
-    @overload
-    def __getitem__(
-        self, index: Tuple[Union[int, slice], Union[int, slice]]
-    ) -> Union[GridContent, List[GridContent]]:
-        ...
-
-    @overload
-    def __getitem__(self, index: Sequence[Coordinate]) -> List[GridContent]:
-        ...
-
-    def __getitem__(
-        self,
-        index: Union[
-            int, Sequence[Coordinate], Tuple[Union[int, slice], Union[int, slice]],
-        ],
-    ) -> Union[GridContent, List[GridContent]]:
-        """Access contents from the grid."""
-
-        if isinstance(index, int):
-            # grid[x]
-            return self.grid[index]
-
-        if isinstance(index[0], tuple):
-            # grid[(x1, y1), (x2, y2)]
-            index = cast(Sequence[Coordinate], index)
-
-            cells = []
-            for pos in index:
-                x1, y1 = self.torus_adj(pos)
-                cells.append(self.grid[x1][y1])
-            return cells
-
-        x, y = index
-
-        if isinstance(x, int) and isinstance(y, int):
-            # grid[x, y]
-            index = cast(Coordinate, index)
-            x, y = self.torus_adj(index)
-            return self.grid[x][y]
-
-        if isinstance(x, int):
-            # grid[x, :]
-            x, _ = self.torus_adj((x, 0))
-            x = slice(x, x + 1)
-
-        if isinstance(y, int):
-            # grid[:, y]
-            _, y = self.torus_adj((0, y))
-            y = slice(y, y + 1)
-
-        # grid[:, :]
-        x, y = (cast(slice, x), cast(slice, y))
-        cells = []
-        for rows in self.grid[x]:
-            for cell in rows[y]:
-                cells.append(cell)
-        return cells
-
-        raise IndexError
+        return self.grid[index]
 
     def __iter__(self) -> Iterator[GridContent]:
         """
@@ -209,7 +132,7 @@ class Grid:
                    diagonals) or Von Neumann (only up/down/left/right).
 
         """
-        neighborhood = self.get_neighborhood(pos, moore=moore)
+        neighborhood = self.iter_neighborhood(pos, moore=moore)
         return self.iter_cell_list_contents(neighborhood)
 
     def iter_neighborhood(
@@ -239,7 +162,31 @@ class Grid:
             including the center).
 
         """
-        yield from self.get_neighborhood(pos, moore, include_center, radius)
+        x, y = pos
+        coordinates: Set[Coordinate] = set()
+        for dy in range(-radius, radius + 1):
+            for dx in range(-radius, radius + 1):
+                if dx == 0 and dy == 0 and not include_center:
+                    continue
+                # Skip coordinates that are outside manhattan distance
+                if not moore and abs(dx) + abs(dy) > radius:
+                    continue
+                # Skip if not a torus and new coords out of bounds.
+                if not self.torus and (
+                    not (0 <= dx + x < self.width) or not (0 <= dy + y < self.height)
+                ):
+                    continue
+
+                px, py = self.torus_adj((x + dx, y + dy))
+
+                # Skip if new coords out of bounds.
+                if self.out_of_bounds((px, py)):
+                    continue
+
+                coords = (px, py)
+                if coords not in coordinates:
+                    coordinates.add(coords)
+                    yield coords
 
     def get_neighborhood(
         self,
@@ -267,35 +214,7 @@ class Grid:
             if not including the center).
 
         """
-        cache_key = (pos, moore, include_center, radius)
-        neighborhood = self._neighborhood_cache.get(cache_key, None)
-
-        if neighborhood is None:
-            coordinates: Set[Coordinate] = set()
-
-            x, y = pos
-            for dy in range(-radius, radius + 1):
-                for dx in range(-radius, radius + 1):
-                    if dx == 0 and dy == 0 and not include_center:
-                        continue
-                    # Skip coordinates that are outside manhattan distance
-                    if not moore and abs(dx) + abs(dy) > radius:
-                        continue
-
-                    coord = (x + dx, y + dy)
-
-                    if self.out_of_bounds(coord):
-                        # Skip if not a torus and new coords out of bounds.
-                        if not self.torus:
-                            continue
-                        coord = self.torus_adj(coord)
-
-                    coordinates.add(coord)
-
-            neighborhood = sorted(coordinates)
-            self._neighborhood_cache[cache_key] = neighborhood
-
-        return neighborhood
+        return list(self.iter_neighborhood(pos, moore, include_center, radius))
 
     def iter_neighbors(
         self,
@@ -321,8 +240,9 @@ class Grid:
             An iterator of non-None objects in the given neighborhood;
             at most 9 if Moore, 5 if Von-Neumann
             (8 and 4 if not including the center).
+
         """
-        neighborhood = self.get_neighborhood(pos, moore, include_center, radius)
+        neighborhood = self.iter_neighborhood(pos, moore, include_center, radius)
         return self.iter_cell_list_contents(neighborhood)
 
     def get_neighbors(
@@ -331,8 +251,8 @@ class Grid:
         moore: bool,
         include_center: bool = False,
         radius: int = 1,
-    ) -> List[GridContent]:
-        """ Return a list of neighbors to a certain point.
+    ) -> List[Coordinate]:
+        """Return a list of neighbors to a certain point.
 
         Args:
             pos: Coordinate tuple for the neighborhood to get.
@@ -360,7 +280,8 @@ class Grid:
         elif not self.torus:
             raise Exception("Point out of bounds, and space non-toroidal.")
         else:
-            return pos[0] % self.width, pos[1] % self.height
+            x, y = pos[0] % self.width, pos[1] % self.height
+        return x, y
 
     def out_of_bounds(self, pos: Coordinate) -> bool:
         """
@@ -382,7 +303,7 @@ class Grid:
             An iterator of the contents of the cells identified in cell_list
 
         """
-        return filter(None, (self.grid[x][y] for x, y in cell_list))
+        return (self[x][y] for x, y in cell_list if not self.is_cell_empty((x, y)))
 
     @accept_tuple_argument
     def get_cell_list_contents(
@@ -950,6 +871,12 @@ class NetworkGrid:
 
         self.G.nodes[node_id]["agent"].remove(agent)
 
+    def remove_agent(self, agent: Agent) -> None:
+        """ Remove the agent from the network and set its pos variable to None. """
+        pos = agent.pos
+        self._remove_agent(agent, pos)
+        agent.pos = None
+
     def is_cell_empty(self, node_id: int) -> bool:
         """ Returns a bool of the contents of a cell. """
         return not self.G.nodes[node_id]["agent"]


Test patch:
diff --git a/tests/test_grid.py b/tests/test_grid.py
index 1079fc06..f7be0f6f 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -431,31 +431,5 @@ class TestHexGridTorus(TestBaseGrid):
         assert len(neighborhood) == 6
 
 
-class TestIndexing:
-    # Create a grid where the content of each coordinate is a tuple of its coordinates
-    grid = Grid(3, 5, True)
-    for _, x, y in grid.coord_iter():
-        grid.grid[x][y] = (x, y)
-
-    def test_int(self):
-        assert self.grid[0][0] == (0, 0)
-
-    def test_tuple(self):
-        assert self.grid[1, 1] == (1, 1)
-
-    def test_list(self):
-        assert self.grid[(0, 0), (1, 1)] == [(0, 0), (1, 1)]
-        assert self.grid[(0, 0), (5, 3)] == [(0, 0), (2, 3)]
-
-    def test_torus(self):
-        assert self.grid[3, 5] == (0, 0)
-
-    def test_slice(self):
-        assert self.grid[:, 0] == [(0, 0), (1, 0), (2, 0)]
-        assert self.grid[::-1, 0] == [(2, 0), (1, 0), (0, 0)]
-        assert self.grid[1, :] == [(1, 0), (1, 1), (1, 2), (1, 3), (1, 4)]
-        assert self.grid[:, :] == [(x, y) for x in range(3) for y in range(5)]
-
-
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_space.py b/tests/test_space.py
index d9f10177..38520b35 100644
--- a/tests/test_space.py
+++ b/tests/test_space.py
@@ -357,6 +357,15 @@ class TestSingleNetworkGrid(unittest.TestCase):
         assert _agent not in self.space.G.nodes[initial_pos]["agent"]
         assert _agent in self.space.G.nodes[final_pos]["agent"]
 
+    def test_remove_agent(self):
+        for i, pos in enumerate(TEST_AGENTS_NETWORK_SINGLE):
+            a = self.agents[i]
+            assert a.pos == pos
+            assert a in self.space.G.nodes[pos]["agent"]
+            self.space.remove_agent(a)
+            assert a.pos is None
+            assert a not in self.space.G.nodes[pos]["agent"]
+
     def test_is_cell_empty(self):
         assert not self.space.is_cell_empty(0)
         assert self.space.is_cell_empty(TestSingleNetworkGrid.GRAPH_SIZE - 1)